Move ensureCapacity methods from BlockUtil to Arrays

snash4 · Oct 22, 2021 · 247881d · 247881d
1 parent 3541380
commit 247881d
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 80 deletions.
diff --git a/presto-common/src/main/java/com/facebook/presto/common/array/Arrays.java b/presto-common/src/main/java/com/facebook/presto/common/array/Arrays.java
@@ -37,11 +37,7 @@ public static boolean[] ensureCapacity(boolean[] buffer, int capacity)
 
     public static byte[] ensureCapacity(byte[] buffer, int capacity)
     {
-        if (buffer == null || buffer.length < capacity) {
-            return new byte[(int) (capacity * SMALL.expansionFactor)];
-        }
-
-        return buffer;
+        return ensureCapacity(buffer, capacity, SMALL, NONE);
     }
 
     public static byte[] ensureCapacity(byte[] buffer, int capacity, ExpansionFactor expansionFactor, ExpansionOption expansionOption)
@@ -92,6 +88,18 @@ else if (buffer.length < capacity) {
         return newBuffer;
     }
 
+    public static short[] ensureCapacity(short[] buffer, int capacity)
+    {
+        if (buffer == null) {
+            buffer = new short[capacity];
+        }
+        else if (buffer.length < capacity) {
+            buffer = java.util.Arrays.copyOf(buffer, capacity);
+        }
+
+        return buffer;
+    }
+
     public static int[] ensureCapacity(int[] buffer, int capacity)
     {
         return ensureCapacity(buffer, capacity, SMALL, NONE);
@@ -147,8 +155,26 @@ else if (buffer.length < capacity) {
 
     public static long[] ensureCapacity(long[] buffer, int capacity)
     {
-        if (buffer == null || buffer.length < capacity) {
-            return new long[(int) (capacity * SMALL.expansionFactor)];
+        return ensureCapacity(buffer, capacity, SMALL, NONE);
+    }
+
+    public static long[] ensureCapacity(long[] buffer, int capacity, ExpansionFactor expansionFactor, ExpansionOption expansionOption)
+    {
+        int newCapacity = (int) (capacity * expansionFactor.expansionFactor);
+
+        if (buffer == null) {
+            buffer = new long[newCapacity];
+        }
+        else if (buffer.length < capacity) {
+            if (expansionOption == PRESERVE) {
+                buffer = java.util.Arrays.copyOf(buffer, newCapacity);
+            }
+            else {
+                buffer = new long[newCapacity];
+            }
+        }
+        else if (expansionOption == INITIALIZE) {
+            java.util.Arrays.fill(buffer, 0L);
         }
 
         return buffer;

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/BlockUtil.java b/presto-common/src/main/java/com/facebook/presto/common/block/BlockUtil.java
@@ -254,70 +254,6 @@ static int[] appendNullToOffsetsArray(int[] offsets, int offsetBase, int positio
         return newOffsets;
     }
 
-    /**
-     * Returns a byte array of size capacity if the input buffer is null or smaller than the capacity.
-     * If the input buffer is not null, the original values in buffer will be preserved.
-     */
-    public static byte[] ensureCapacity(byte[] buffer, int capacity)
-    {
-        if (buffer == null) {
-            buffer = new byte[capacity];
-        }
-        else if (buffer.length < capacity) {
-            buffer = Arrays.copyOf(buffer, capacity);
-        }
-
-        return buffer;
-    }
-
-    /**
-     * Returns a short array of size capacity if the input buffer is null or smaller than the capacity.
-     * If the input buffer is not null, the original values in buffer will be preserved.
-     */
-    public static short[] ensureCapacity(short[] buffer, int capacity)
-    {
-        if (buffer == null) {
-            buffer = new short[capacity];
-        }
-        else if (buffer.length < capacity) {
-            buffer = Arrays.copyOf(buffer, capacity);
-        }
-
-        return buffer;
-    }
-
-    /**
-     * Returns an int array of size capacity if the input buffer is null or smaller than the capacity.
-     * If the input buffer is not null, the original values in buffer will be preserved.
-     */
-    public static int[] ensureCapacity(int[] buffer, int capacity)
-    {
-        if (buffer == null) {
-            buffer = new int[capacity];
-        }
-        else if (buffer.length < capacity) {
-            buffer = Arrays.copyOf(buffer, capacity);
-        }
-
-        return buffer;
-    }
-
-    /**
-     * Returns a long array of size capacity if the input buffer is null or smaller than the capacity.
-     * If the input buffer is not null, the original values in buffer will be preserved.
-     */
-    public static long[] ensureCapacity(long[] buffer, int capacity)
-    {
-        if (buffer == null) {
-            buffer = new long[capacity];
-        }
-        else if (buffer.length < capacity) {
-            buffer = Arrays.copyOf(buffer, capacity);
-        }
-
-        return buffer;
-    }
-
     public static int getNum128Integers(int length)
     {
         int num128Integers = length / SIZE_OF_LONG / 2;

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/ByteArrayBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/ByteArrayBlock.java
@@ -21,12 +21,14 @@
 import java.util.Optional;
 import java.util.function.BiConsumer;
 
+import static com.facebook.presto.common.array.Arrays.ExpansionFactor.SMALL;
+import static com.facebook.presto.common.array.Arrays.ExpansionOption.PRESERVE;
+import static com.facebook.presto.common.array.Arrays.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.appendNullToIsNullArray;
 import static com.facebook.presto.common.block.BlockUtil.checkArrayRange;
 import static com.facebook.presto.common.block.BlockUtil.checkValidRegion;
 import static com.facebook.presto.common.block.BlockUtil.compactArray;
 import static com.facebook.presto.common.block.BlockUtil.countUsedPositions;
-import static com.facebook.presto.common.block.BlockUtil.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.internalPositionInRange;
 import static io.airlift.slice.SizeOf.sizeOf;
 import static java.lang.String.format;
@@ -260,7 +262,7 @@ public byte getByteUnchecked(int internalPosition)
     public Block appendNull()
     {
         boolean[] newValueIsNull = appendNullToIsNullArray(valueIsNull, arrayOffset, positionCount);
-        byte[] newValues = ensureCapacity(values, arrayOffset + positionCount + 1);
+        byte[] newValues = ensureCapacity(values, arrayOffset + positionCount + 1, SMALL, PRESERVE);
 
         return new ByteArrayBlock(arrayOffset, positionCount + 1, newValueIsNull, newValues);
     }

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/Int128ArrayBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/Int128ArrayBlock.java
@@ -23,12 +23,14 @@
 import java.util.Optional;
 import java.util.function.BiConsumer;
 
+import static com.facebook.presto.common.array.Arrays.ExpansionFactor.SMALL;
+import static com.facebook.presto.common.array.Arrays.ExpansionOption.PRESERVE;
+import static com.facebook.presto.common.array.Arrays.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.appendNullToIsNullArray;
 import static com.facebook.presto.common.block.BlockUtil.checkArrayRange;
 import static com.facebook.presto.common.block.BlockUtil.checkValidRegion;
 import static com.facebook.presto.common.block.BlockUtil.compactArray;
 import static com.facebook.presto.common.block.BlockUtil.countUsedPositions;
-import static com.facebook.presto.common.block.BlockUtil.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.getNum128Integers;
 import static com.facebook.presto.common.block.BlockUtil.internalPositionInRange;
 import static io.airlift.slice.SizeOf.SIZE_OF_LONG;
@@ -356,7 +358,7 @@ public boolean isNullUnchecked(int internalPosition)
     public Block appendNull()
     {
         boolean[] newValueIsNull = appendNullToIsNullArray(valueIsNull, positionOffset, positionCount);
-        long[] newValues = ensureCapacity(values, (positionOffset + positionCount + 1) * 2);
+        long[] newValues = ensureCapacity(values, (positionOffset + positionCount + 1) * 2, SMALL, PRESERVE);
         return new Int128ArrayBlock(positionOffset, positionCount + 1, newValueIsNull, newValues);
     }
 }
diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/IntArrayBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/IntArrayBlock.java
@@ -23,12 +23,14 @@
 import java.util.Optional;
 import java.util.function.BiConsumer;
 
+import static com.facebook.presto.common.array.Arrays.ExpansionFactor.SMALL;
+import static com.facebook.presto.common.array.Arrays.ExpansionOption.PRESERVE;
+import static com.facebook.presto.common.array.Arrays.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.appendNullToIsNullArray;
 import static com.facebook.presto.common.block.BlockUtil.checkArrayRange;
 import static com.facebook.presto.common.block.BlockUtil.checkValidRegion;
 import static com.facebook.presto.common.block.BlockUtil.compactArray;
 import static com.facebook.presto.common.block.BlockUtil.countUsedPositions;
-import static com.facebook.presto.common.block.BlockUtil.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.internalPositionInRange;
 import static io.airlift.slice.SizeOf.sizeOf;
 import static java.lang.String.format;
@@ -262,7 +264,7 @@ public int getIntUnchecked(int internalPosition)
     public Block appendNull()
     {
         boolean[] newValueIsNull = appendNullToIsNullArray(valueIsNull, arrayOffset, positionCount);
-        int[] newValues = ensureCapacity(values, arrayOffset + positionCount + 1);
+        int[] newValues = ensureCapacity(values, arrayOffset + positionCount + 1, SMALL, PRESERVE);
 
         return new IntArrayBlock(arrayOffset, positionCount + 1, newValueIsNull, newValues);
     }

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/LongArrayBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/LongArrayBlock.java
@@ -21,12 +21,14 @@
 import java.util.Optional;
 import java.util.function.BiConsumer;
 
+import static com.facebook.presto.common.array.Arrays.ExpansionFactor.SMALL;
+import static com.facebook.presto.common.array.Arrays.ExpansionOption.PRESERVE;
+import static com.facebook.presto.common.array.Arrays.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.appendNullToIsNullArray;
 import static com.facebook.presto.common.block.BlockUtil.checkArrayRange;
 import static com.facebook.presto.common.block.BlockUtil.checkValidRegion;
 import static com.facebook.presto.common.block.BlockUtil.compactArray;
 import static com.facebook.presto.common.block.BlockUtil.countUsedPositions;
-import static com.facebook.presto.common.block.BlockUtil.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.internalPositionInRange;
 import static io.airlift.slice.SizeOf.sizeOf;
 import static java.lang.Math.toIntExact;
@@ -298,7 +300,7 @@ public long getLongUnchecked(int internalPosition)
     public Block appendNull()
     {
         boolean[] newValueIsNull = appendNullToIsNullArray(valueIsNull, arrayOffset, positionCount);
-        long[] newValues = ensureCapacity(values, arrayOffset + positionCount + 1);
+        long[] newValues = ensureCapacity(values, arrayOffset + positionCount + 1, SMALL, PRESERVE);
 
         return new LongArrayBlock(arrayOffset, positionCount + 1, newValueIsNull, newValues);
     }

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/ShortArrayBlock.java b/presto-common/src/main/java/com/facebook/presto/common/block/ShortArrayBlock.java
@@ -21,12 +21,12 @@
 import java.util.Optional;
 import java.util.function.BiConsumer;
 
+import static com.facebook.presto.common.array.Arrays.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.appendNullToIsNullArray;
 import static com.facebook.presto.common.block.BlockUtil.checkArrayRange;
 import static com.facebook.presto.common.block.BlockUtil.checkValidRegion;
 import static com.facebook.presto.common.block.BlockUtil.compactArray;
 import static com.facebook.presto.common.block.BlockUtil.countUsedPositions;
-import static com.facebook.presto.common.block.BlockUtil.ensureCapacity;
 import static com.facebook.presto.common.block.BlockUtil.internalPositionInRange;
 import static io.airlift.slice.SizeOf.sizeOf;
 import static java.lang.String.format;