Introduce newBlockBuilderLike() with expectedEntries

cjjxfli · Oct 23, 2020 · 8ad5079 · 8ad5079
1 parent ee56436
commit 8ad5079
Show file tree

Hide file tree

Showing 14 changed files with 134 additions and 1 deletion.
diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/ArrayBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/ArrayBlockBuilder.java
@@ -26,6 +26,7 @@
 import static com.facebook.presto.common.block.BlockUtil.calculateBlockResetSize;
 import static io.airlift.slice.SizeOf.sizeOf;
 import static java.lang.Math.max;
+import static java.lang.Math.toIntExact;
 import static java.util.Objects.requireNonNull;
 
 public class ArrayBlockBuilder
@@ -292,6 +293,14 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         return new ArrayBlockBuilder(blockBuilderStatus, values.newBlockBuilderLike(blockBuilderStatus), newSize);
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        int newSize = max(calculateBlockResetSize(positionCount), expectedEntries);
+        int valueExpectedEntries = positionCount == 0 ? expectedEntries : toIntExact((long) offsets[positionCount] * newSize / positionCount);
+        return new ArrayBlockBuilder(blockBuilderStatus, values.newBlockBuilderLike(blockBuilderStatus, valueExpectedEntries), newSize);
+    }
+
     @Override
     public String toString()
     {

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/BlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/BlockBuilder.java
@@ -118,4 +118,9 @@ default BlockBuilder appendStructureInternal(Block block, int position)
      * Creates a new block builder of the same type based on the current usage statistics of this block builder.
      */
     BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus);
+
+    /**
+     * Creates a new block builder of the same type based on the expectedEntries and the current usage statistics of this block builder.
+     */
+    BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries);
 }
diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/ByteArrayBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/ByteArrayBlockBuilder.java
@@ -114,6 +114,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         return new ByteArrayBlockBuilder(blockBuilderStatus, calculateBlockResetSize(positionCount));
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        return new ByteArrayBlockBuilder(blockBuilderStatus, max(calculateBlockResetSize(positionCount), expectedEntries));
+    }
+
     private void growCapacity()
     {
         int newSize;

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/Int128ArrayBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/Int128ArrayBlockBuilder.java
@@ -133,6 +133,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         return new Int128ArrayBlockBuilder(blockBuilderStatus, calculateBlockResetSize(positionCount));
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        return new Int128ArrayBlockBuilder(blockBuilderStatus, max(calculateBlockResetSize(positionCount), expectedEntries));
+    }
+
     private void growCapacity()
     {
         int newSize;

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/IntArrayBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/IntArrayBlockBuilder.java
@@ -114,6 +114,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         return new IntArrayBlockBuilder(blockBuilderStatus, calculateBlockResetSize(positionCount));
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        return new IntArrayBlockBuilder(blockBuilderStatus, max(calculateBlockResetSize(positionCount), expectedEntries));
+    }
+
     private void growCapacity()
     {
         int newSize;

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/LongArrayBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/LongArrayBlockBuilder.java
@@ -56,7 +56,6 @@ public LongArrayBlockBuilder(@Nullable BlockBuilderStatus blockBuilderStatus, in
     {
         this.blockBuilderStatus = blockBuilderStatus;
         this.initialEntryCount = max(expectedEntries, 1);
-
         updateDataSize();
     }
 
@@ -115,6 +114,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         return new LongArrayBlockBuilder(blockBuilderStatus, calculateBlockResetSize(positionCount));
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        return new LongArrayBlockBuilder(blockBuilderStatus, max(calculateBlockResetSize(positionCount), expectedEntries));
+    }
+
     private void growCapacity()
     {
         int newSize;

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/MapBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/MapBlockBuilder.java
@@ -29,6 +29,8 @@
 import static com.facebook.presto.common.block.BlockUtil.calculateBlockResetSize;
 import static com.facebook.presto.common.block.MapBlock.createMapBlockInternal;
 import static io.airlift.slice.SizeOf.sizeOf;
+import static java.lang.Math.max;
+import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
@@ -455,6 +457,22 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
                 newNegativeOneFilledArray(newSize * HASH_MULTIPLIER));
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        int newSize = max(calculateBlockResetSize(getPositionCount()), expectedEntries);
+        int nestedExpectedEntries = positionCount == 0 ? expectedEntries : toIntExact((long) offsets[positionCount] * newSize / positionCount);
+        return new MapBlockBuilder(
+                keyBlockEquals,
+                keyBlockHashCode,
+                blockBuilderStatus,
+                keyBlockBuilder.newBlockBuilderLike(blockBuilderStatus, nestedExpectedEntries),
+                valueBlockBuilder.newBlockBuilderLike(blockBuilderStatus, nestedExpectedEntries),
+                new int[newSize + 1],
+                new boolean[newSize],
+                newNegativeOneFilledArray(newSize * HASH_MULTIPLIER));
+    }
+
     @Override
     protected void ensureHashTableLoaded(MethodHandle keyBlockHashCode) {}
 

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/RowBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/RowBlockBuilder.java
@@ -27,6 +27,8 @@
 import static com.facebook.presto.common.block.BlockUtil.calculateBlockResetSize;
 import static com.facebook.presto.common.block.RowBlock.createRowBlockInternal;
 import static io.airlift.slice.SizeOf.sizeOf;
+import static java.lang.Math.max;
+import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
@@ -298,4 +300,16 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         }
         return new RowBlockBuilder(blockBuilderStatus, newBlockBuilders, new int[newSize + 1], new boolean[newSize]);
     }
+
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        int newSize = max(calculateBlockResetSize(getPositionCount()), expectedEntries);
+        BlockBuilder[] newBlockBuilders = new BlockBuilder[numFields];
+        int nestedExpectedEntries = positionCount == 0 ? expectedEntries : toIntExact((long) fieldBlockOffsets[positionCount] * newSize / positionCount);
+        for (int i = 0; i < numFields; i++) {
+            newBlockBuilders[i] = fieldBlockBuilders[i].newBlockBuilderLike(blockBuilderStatus, nestedExpectedEntries);
+        }
+        return new RowBlockBuilder(blockBuilderStatus, newBlockBuilders, new int[newSize + 1], new boolean[newSize]);
+    }
 }
diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/ShortArrayBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/ShortArrayBlockBuilder.java
@@ -114,6 +114,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         return new ShortArrayBlockBuilder(blockBuilderStatus, calculateBlockResetSize(positionCount));
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        return new ShortArrayBlockBuilder(blockBuilderStatus, max(calculateBlockResetSize(positionCount), expectedEntries));
+    }
+
     private void growCapacity()
     {
         int newSize;

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/SingleArrayBlockWriter.java b/presto-common/src/main/java/com/facebook/presto/common/block/SingleArrayBlockWriter.java
@@ -167,6 +167,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public String toString()
     {

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/SingleMapBlockWriter.java b/presto-common/src/main/java/com/facebook/presto/common/block/SingleMapBlockWriter.java
@@ -249,6 +249,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public String toString()
     {

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/SingleRowBlockWriter.java b/presto-common/src/main/java/com/facebook/presto/common/block/SingleRowBlockWriter.java
@@ -229,6 +229,12 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public String toString()
     {

diff --git a/presto-common/src/main/java/com/facebook/presto/common/block/VariableWidthBlockBuilder.java b/presto-common/src/main/java/com/facebook/presto/common/block/VariableWidthBlockBuilder.java
@@ -44,7 +44,9 @@
 import static io.airlift.slice.SizeOf.SIZE_OF_LONG;
 import static io.airlift.slice.SizeOf.SIZE_OF_SHORT;
 import static io.airlift.slice.SizeOf.sizeOf;
+import static java.lang.Math.max;
 import static java.lang.Math.min;
+import static java.lang.Math.toIntExact;
 
 public class VariableWidthBlockBuilder
         extends AbstractVariableWidthBlock
@@ -381,6 +383,16 @@ public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus)
         return new VariableWidthBlockBuilder(blockBuilderStatus, calculateBlockResetSize(positions), calculateBlockResetBytes(currentSizeInBytes));
     }
 
+    @Override
+    public BlockBuilder newBlockBuilderLike(BlockBuilderStatus blockBuilderStatus, int expectedEntries)
+    {
+        int newSize = max(calculateBlockResetSize(positions), expectedEntries);
+        int currentSizeInBytes = offsets[positions];
+        return new VariableWidthBlockBuilder(blockBuilderStatus,
+                newSize,
+                max(calculateBlockResetBytes(currentSizeInBytes), positions == 0 ? currentSizeInBytes : toIntExact((long) currentSizeInBytes * newSize / positions)));
+    }
+
     private int getOffset(int position)
     {
         return offsets[position];

diff --git a/presto-main/src/test/java/com/facebook/presto/block/TestBlockBuilder.java b/presto-main/src/test/java/com/facebook/presto/block/TestBlockBuilder.java
@@ -23,9 +23,12 @@
 import org.testng.annotations.Test;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.IntStream;
 
 import static com.facebook.presto.block.BlockAssertions.assertBlockEquals;
+import static com.facebook.presto.common.block.ArrayBlock.fromElementBlock;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static org.testng.Assert.assertEquals;
@@ -80,6 +83,31 @@ public void testNewBlockBuilderLike()
             assertEquals(newPageBuilder.getBlockBuilder(i).getPositionCount(), 0);
             assertTrue(newPageBuilder.getBlockBuilder(i).getRetainedSizeInBytes() < pageBuilder.getBlockBuilder(i).getRetainedSizeInBytes());
         }
+
+        BlockBuilder newBigintBlockBuilder = bigintBlockBuilder.newBlockBuilderLike(null, 200);
+        assertEquals(newBigintBlockBuilder.getPositionCount(), 0);
+        assertEquals(newBigintBlockBuilder.getRetainedSizeInBytes(), 80);
+        newBigintBlockBuilder.writeLong(0);
+        assertEquals(newBigintBlockBuilder.getPositionCount(), 1);
+        // Reserved 200 longs and booleans for nulls array
+        assertEquals(newBigintBlockBuilder.getRetainedSizeInBytes(), 1880);
+
+        BlockBuilder newVarcharBlockBuilder = varcharBlockBuilder.newBlockBuilderLike(null, 200);
+        assertEquals(newVarcharBlockBuilder.getPositionCount(), 0);
+        assertEquals(newVarcharBlockBuilder.getRetainedSizeInBytes(), 164);
+        newVarcharBlockBuilder.writeLong(0);
+        newVarcharBlockBuilder.closeEntry();
+        assertEquals(newVarcharBlockBuilder.getPositionCount(), 1);
+        // Reserved 200 varchars of average length 5.9, and 201 ints for offsets and 200 booleans for nulls
+        assertEquals(newVarcharBlockBuilder.getRetainedSizeInBytes(), 2360);
+
+        BlockBuilder newArrayBlockBuilder = arrayBlockBuilder.newBlockBuilderLike(null, 200);
+        assertEquals(newArrayBlockBuilder.getPositionCount(), 0);
+        assertEquals(newArrayBlockBuilder.getRetainedSizeInBytes(), 248);
+        newArrayBlockBuilder.appendStructure(fromElementBlock(1, Optional.empty(), IntStream.range(0, 2).toArray(), newBigintBlockBuilder.build()));
+        assertEquals(newArrayBlockBuilder.getPositionCount(), 1);
+        // Reserved 200 ARRAY(ARRAY(BIGINT)), and 201 ints for offsets and 200 booleans for nulls
+        assertEquals(newArrayBlockBuilder.getRetainedSizeInBytes(), 5848);
     }
 
     @Test