Reduce overheads related to DictionaryBlock#getSizeInBytes()

core/trino-main/src/main/java/io/trino/operator/GroupByIdBlock.java

@@ -19,6 +19,7 @@
 import org.openjdk.jol.info.ClassLayout;
 
 import java.util.List;
+import java.util.OptionalInt;
 import java.util.function.BiConsumer;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
@@ -64,9 +65,15 @@ public long getRegionSizeInBytes(int positionOffset, int length)
     }
 
     @Override
-    public long getPositionsSizeInBytes(boolean[] positions)
+    public OptionalInt fixedSizeInBytesPerPosition()
     {
-        return block.getPositionsSizeInBytes(positions);
+        return block.fixedSizeInBytesPerPosition();
+    }
+
+    @Override
+    public long getPositionsSizeInBytes(boolean[] positions, int selectedPositionCount)
+    {
+        return block.getPositionsSizeInBytes(positions, selectedPositionCount);
     }
 
     @Override

core/trino-main/src/test/java/io/trino/block/AbstractTestBlock.java

@@ -221,11 +221,11 @@ private void assertBlockSize(Block block)
 
         boolean[] positions = new boolean[block.getPositionCount()];
         fill(positions, 0, firstHalf.getPositionCount(), true);
-        assertEquals(block.getPositionsSizeInBytes(positions), expectedFirstHalfSize);
+        assertEquals(block.getPositionsSizeInBytes(positions, firstHalf.getPositionCount()), expectedFirstHalfSize);
         fill(positions, true);
-        assertEquals(block.getPositionsSizeInBytes(positions), expectedBlockSize);
+        assertEquals(block.getPositionsSizeInBytes(positions, positions.length), expectedBlockSize);
         fill(positions, 0, firstHalf.getPositionCount(), false);
-        assertEquals(block.getPositionsSizeInBytes(positions), expectedSecondHalfSize);
+        assertEquals(block.getPositionsSizeInBytes(positions, positions.length - firstHalf.getPositionCount()), expectedSecondHalfSize);
     }
 
     // expectedValueType is required since otherwise the expected value type is unknown when expectedValue is null.

core/trino-main/src/test/java/io/trino/block/TestDictionaryBlock.java

@@ -19,10 +19,15 @@
 import io.trino.spi.block.BlockBuilder;
 import io.trino.spi.block.DictionaryBlock;
 import io.trino.spi.block.DictionaryId;
+import io.trino.spi.block.IntArrayBlock;
 import io.trino.spi.block.VariableWidthBlock;
 import io.trino.spi.block.VariableWidthBlockBuilder;
 import org.testng.annotations.Test;
 
+import java.util.Arrays;
+import java.util.Optional;
+import java.util.stream.IntStream;
+
 import static io.airlift.slice.SizeOf.SIZE_OF_INT;
 import static io.airlift.testing.Assertions.assertInstanceOf;
 import static io.trino.block.BlockAssertions.createSlicesBlock;
@@ -252,12 +257,12 @@ public void testNestedCompact()
 
         assertEquals(
                 dictionary.getSizeInBytes(),
-                valuesBlock.getPositionsSizeInBytes(new boolean[] {true, false, true, false, false, false}) + 4 * Integer.BYTES);
+                valuesBlock.getPositionsSizeInBytes(new boolean[] {true, false, true, false, false, false}, 2) + 4 * Integer.BYTES);
         assertFalse(dictionary.isCompact());
 
         assertEquals(
                 dictionaryWithAllPositionsUsed.getSizeInBytes(),
-                valuesBlock.getPositionsSizeInBytes(new boolean[] {true, true, true, false, true, true}) + 6 * Integer.BYTES);
+                valuesBlock.getPositionsSizeInBytes(new boolean[] {true, true, true, false, true, true}, 5) + 6 * Integer.BYTES);
         // dictionary is not compact (even though all positions were used) because it's unnested
         assertFalse(dictionaryWithAllPositionsUsed.isCompact());
 
@@ -392,6 +397,114 @@ public void testEstimatedDataSizeForStats()
         }
     }
 
+    @Test
+    public void testNestedDictionarySizes()
+    {
+        // fixed width block
+        Block fixedWidthBlock = new IntArrayBlock(100, Optional.empty(), IntStream.range(0, 100).toArray());
+        assertDictionarySizeMethods(fixedWidthBlock);
+        assertDictionarySizeMethods(new DictionaryBlock(fixedWidthBlock, IntStream.range(0, 50).toArray()));
+        assertDictionarySizeMethods(
+                new DictionaryBlock(
+                        new DictionaryBlock(fixedWidthBlock, IntStream.range(0, 50).toArray()),
+                        IntStream.range(0, 10).toArray()));
+
+        // variable width block
+        Block variableWidthBlock = createSlicesBlock(createExpectedValues(100));
+        assertDictionarySizeMethods(variableWidthBlock);
+        assertDictionarySizeMethods(new DictionaryBlock(variableWidthBlock, IntStream.range(0, 50).toArray()));
+        assertDictionarySizeMethods(
+                new DictionaryBlock(
+                        new DictionaryBlock(variableWidthBlock, IntStream.range(0, 50).toArray()),
+                        IntStream.range(0, 10).toArray()));
+    }
+
+    private static void assertDictionarySizeMethods(Block block)
+    {
+        int positions = block.getPositionCount();
+
+        int[] allIds = IntStream.range(0, positions).toArray();
+        if (block instanceof DictionaryBlock) {
+            assertEquals(
+                    new DictionaryBlock(block, allIds).getSizeInBytes(),
+                    block.getSizeInBytes(),
+                    "nested dictionary size should not be counted");
+        }
+        else {
+            assertEquals(new DictionaryBlock(block, allIds).getSizeInBytes(), block.getSizeInBytes() + (Integer.BYTES * (long) positions));
+        }
+
+        if (positions > 0) {
+            int firstHalfLength = positions / 2;
+            int secondHalfLength = positions - firstHalfLength;
+            int[] firstHalfIds = IntStream.range(0, firstHalfLength).toArray();
+            int[] secondHalfIds = IntStream.range(firstHalfLength, positions).toArray();
+
+            boolean[] selectedPositions = new boolean[positions];
+            selectedPositions[0] = true;
+            if (block instanceof DictionaryBlock) {
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getPositionsSizeInBytes(selectedPositions, 1),
+                        block.getPositionsSizeInBytes(selectedPositions, 1),
+                        "nested dictionary blocks must not include nested id overhead");
+                assertEquals(
+                        new DictionaryBlock(block, new int[]{0}).getSizeInBytes(),
+                        block.getPositionsSizeInBytes(selectedPositions, 1),
+                        "nested dictionary blocks must not include nested id overhead");
+
+                Arrays.fill(selectedPositions, true);
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getPositionsSizeInBytes(selectedPositions, positions),
+                        block.getSizeInBytes(),
+                        "nested dictionary blocks must not include nested id overhead");
+
+                assertEquals(
+                        new DictionaryBlock(block, firstHalfIds).getSizeInBytes(),
+                        block.getRegionSizeInBytes(0, firstHalfLength),
+                        "nested dictionary blocks must not include nested id overhead");
+                assertEquals(
+                        new DictionaryBlock(block, secondHalfIds).getSizeInBytes(),
+                        block.getRegionSizeInBytes(firstHalfLength, secondHalfLength),
+                        "nested dictionary blocks must not include nested id overhead");
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getRegionSizeInBytes(0, firstHalfLength),
+                        block.getRegionSizeInBytes(0, firstHalfLength),
+                        "nested dictionary blocks must not include nested id overhead");
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getRegionSizeInBytes(firstHalfLength, secondHalfLength),
+                        block.getRegionSizeInBytes(firstHalfLength, secondHalfLength),
+                        "nested dictionary blocks must not include nested id overhead");
+            }
+            else {
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getPositionsSizeInBytes(selectedPositions, 1),
+                        block.getPositionsSizeInBytes(selectedPositions, 1) + Integer.BYTES);
+
+                assertEquals(
+                        new DictionaryBlock(block, new int[]{0}).getSizeInBytes(),
+                        block.getPositionsSizeInBytes(selectedPositions, 1) + Integer.BYTES);
+
+                Arrays.fill(selectedPositions, true);
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getPositionsSizeInBytes(selectedPositions, positions),
+                        block.getSizeInBytes() + (Integer.BYTES * (long) positions));
+
+                assertEquals(
+                        new DictionaryBlock(block, firstHalfIds).getSizeInBytes(),
+                        block.getRegionSizeInBytes(0, firstHalfLength) + (Integer.BYTES * (long) firstHalfLength));
+                assertEquals(
+                        new DictionaryBlock(block, secondHalfIds).getSizeInBytes(),
+                        block.getRegionSizeInBytes(firstHalfLength, secondHalfLength) + (Integer.BYTES * (long) secondHalfLength));
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getRegionSizeInBytes(0, firstHalfLength),
+                        block.getRegionSizeInBytes(0, firstHalfLength) + (Integer.BYTES * (long) firstHalfLength));
+                assertEquals(
+                        new DictionaryBlock(block, allIds).getRegionSizeInBytes(firstHalfLength, secondHalfLength),
+                        block.getRegionSizeInBytes(firstHalfLength, secondHalfLength) + (Integer.BYTES * (long) secondHalfLength));
+            }
+        }
+    }
+
     private static DictionaryBlock createDictionaryBlockWithUnreferencedKeys(Slice[] expectedValues, int positionCount)
     {
         // adds references to 0 and all odd indexes

core/trino-main/src/test/java/io/trino/block/TestRunLengthEncodedBlock.java

@@ -26,6 +26,7 @@
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
 
 public class TestRunLengthEncodedBlock
         extends AbstractTestBlock
@@ -61,6 +62,26 @@ private static BlockBuilder createBlockBuilder()
         return new VariableWidthBlockBuilder(null, 1, 1);
     }
 
+    @Test
+    public void testPositionsSizeInBytes()
+    {
+        Block valueBlock = createSingleValueBlock(createExpectedValue(10));
+        Block rleBlock = new RunLengthEncodedBlock(valueBlock, 10);
+        // Size in bytes is not fixed per position
+        assertTrue(rleBlock.fixedSizeInBytesPerPosition().isEmpty());
+        // Accepts specific position selection
+        boolean[] positions = new boolean[rleBlock.getPositionCount()];
+        positions[0] = true;
+        positions[1] = true;
+        assertEquals(rleBlock.getPositionsSizeInBytes(positions, 2), valueBlock.getSizeInBytes());
+        // Accepts null positions array with count only
+        assertEquals(rleBlock.getPositionsSizeInBytes(null, 2), valueBlock.getSizeInBytes());
+        // Always reports the same size in bytes regardless of positions
+        for (int positionCount = 0; positionCount < rleBlock.getPositionCount(); positionCount++) {
+            assertEquals(rleBlock.getPositionsSizeInBytes(null, positionCount), valueBlock.getSizeInBytes());
+        }
+    }
+
     @Test
     public void testBuildingFromLongArrayBlockBuilder()
     {