Reduce contribution of RLE blocks size in page splitting

Consider only the underlying value block size for RLE blocks in
page splitting calculation. This prevents the splitting up of Pages
due to getSizeInBytes of RLE blocks being a factor of positions count

Author: raunaqmorarka

core/trino-main/src/main/java/io/trino/execution/buffer/PageSplitterUtil.java

@@ -16,6 +16,10 @@
 import com.google.common.collect.ImmutableList;
 import io.airlift.slice.Slice;
 import io.trino.spi.Page;
+import io.trino.spi.block.Block;
+import io.trino.spi.block.DictionaryBlock;
+import io.trino.spi.block.RunLengthEncodedBlock;
+import io.trino.spi.block.ValueBlock;
 
 import java.util.List;
 
@@ -51,21 +55,41 @@ private static List<Page> splitPage(Page page, long maxPageSizeInBytes, long pre
         // the recursion would only terminate when page.getPositionCount() == 1
         // and create potentially a large number of Page's of size 1. So we check here that
         // if the size of the page doesn't improve from the previous call we terminate the recursion.
-        if (page.getSizeInBytes() == previousPageSize || page.getSizeInBytes() <= maxPageSizeInBytes || page.getPositionCount() == 1) {
+        long currentPageSize = getPageSizeForSplit(page);
+        if (currentPageSize == previousPageSize || currentPageSize <= maxPageSizeInBytes || page.getPositionCount() == 1) {
             return ImmutableList.of(page);
         }
 
         ImmutableList.Builder<Page> outputPages = ImmutableList.builder();
-        long previousSize = page.getSizeInBytes();
         int positionCount = page.getPositionCount();
         int half = positionCount / 2;
 
         Page leftHalf = page.getRegion(0, half);
-        outputPages.addAll(splitPage(leftHalf, maxPageSizeInBytes, previousSize));
+        outputPages.addAll(splitPage(leftHalf, maxPageSizeInBytes, currentPageSize));
 
         Page rightHalf = page.getRegion(half, positionCount - half);
-        outputPages.addAll(splitPage(rightHalf, maxPageSizeInBytes, previousSize));
+        outputPages.addAll(splitPage(rightHalf, maxPageSizeInBytes, currentPageSize));
 
         return outputPages.build();
     }
+
+    /**
+     * Returns the size of the page in bytes.
+     * This is used to determine if the page should be split.
+     * This differs from {@link Page#getSizeInBytes()} in that it purposely calculates the size RLE blocks
+     * as the size of the underlying value block. This is because we want to avoid creating a large number
+     * of smaller Pages when the source has produced RLE-only Pages with large positions count.
+     */
+    private static long getPageSizeForSplit(Page page)
+    {
+        long size = 0;
+        for (int channel = 0; channel < page.getChannelCount(); channel++) {
+            Block block = page.getBlock(channel);
+            switch (block) {
+                case RunLengthEncodedBlock rleBlock -> size += rleBlock.getUnderlyingValueBlock().getSizeInBytes();
+                case DictionaryBlock _, ValueBlock _ -> size += block.getSizeInBytes();
+            }
+        }
+        return size;
+    }
 }

core/trino-main/src/test/java/io/trino/execution/TestPageSplitterUtil.java

@@ -15,18 +15,15 @@
 
 import com.google.common.collect.ImmutableList;
 import io.airlift.slice.Slice;
-import io.trino.spi.Experimental;
 import io.trino.spi.Page;
-import io.trino.spi.block.ByteArrayBlock;
-import io.trino.spi.block.ValueBlock;
+import io.trino.spi.block.Block;
+import io.trino.spi.block.RunLengthEncodedBlock;
 import io.trino.spi.block.VariableWidthBlockBuilder;
 import io.trino.spi.type.Type;
 import io.trino.testing.MaterializedResult;
 import org.junit.jupiter.api.Test;
 
 import java.util.List;
-import java.util.Optional;
-import java.util.function.ObjLongConsumer;
 
 import static io.airlift.slice.Slices.wrappedBuffer;
 import static io.trino.SequencePageBuilder.createSequencePage;
@@ -79,13 +76,13 @@ public void testSplitPageNonDecreasingPageSize()
     {
         int positionCount = 100;
         int maxPageSizeInBytes = 1;
+        List<Type> types = ImmutableList.of(VARCHAR);
 
         Slice expectedValue = wrappedBuffer("test".getBytes(UTF_8));
         VariableWidthBlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, 1, expectedValue.length());
-        for (int i = 0; i < positionCount; i++) {
-            blockBuilder.writeEntry(expectedValue);
-        }
-        Page initialPage = new Page(new FixedDataSizeBlock(expectedValue.length(), blockBuilder.buildValueBlock()));
+        blockBuilder.writeEntry(expectedValue);
+        Block rleBlock = RunLengthEncodedBlock.create(blockBuilder.build(), positionCount);
+        Page initialPage = new Page(rleBlock);
         List<Page> pages = splitPage(initialPage, maxPageSizeInBytes);
 
         // the page should only be split in half as the recursion should terminate
@@ -98,89 +95,8 @@ public void testSplitPageNonDecreasingPageSize()
         assertThat((int) first.getSizeInBytes()).isGreaterThan(maxPageSizeInBytes);
         assertThat((int) second.getSizeInBytes()).isGreaterThan(maxPageSizeInBytes);
         assertPositionCount(pages, positionCount);
-    }
-
-    // Fake block that has retains a fixed size when split
-    private record FixedDataSizeBlock(long fixedSize, ValueBlock delegate)
-            implements ValueBlock
-    {
-        @Override
-        public ValueBlock copyPositions(int[] positions, int offset, int length)
-        {
-            return new FixedDataSizeBlock(fixedSize, delegate.copyPositions(positions, offset, length));
-        }
-
-        @Override
-        public ValueBlock getRegion(int positionOffset, int length)
-        {
-            return new FixedDataSizeBlock(fixedSize, delegate.getRegion(positionOffset, length));
-        }
-
-        @Override
-        public ValueBlock copyRegion(int position, int length)
-        {
-            return new FixedDataSizeBlock(fixedSize, delegate.copyRegion(position, length));
-        }
-
-        @Override
-        public ValueBlock copyWithAppendedNull()
-        {
-            return new FixedDataSizeBlock(fixedSize, delegate.copyWithAppendedNull());
-        }
-
-        @Experimental(eta = "2025-01-01")
-        @Override
-        public Optional<ByteArrayBlock> getNulls()
-        {
-            return delegate.getNulls();
-        }
-
-        @Override
-        public ValueBlock getSingleValueBlock(int position)
-        {
-            return delegate.getSingleValueBlock(position);
-        }
-
-        @Override
-        public int getPositionCount()
-        {
-            return delegate.getPositionCount();
-        }
-
-        @Override
-        public long getSizeInBytes()
-        {
-            return fixedSize;
-        }
-
-        @Override
-        public long getRegionSizeInBytes(int position, int length)
-        {
-            return delegate.getRegionSizeInBytes(position, length);
-        }
-
-        @Override
-        public long getRetainedSizeInBytes()
-        {
-            return delegate.getRetainedSizeInBytes();
-        }
-
-        @Override
-        public long getEstimatedDataSizeForStats(int position)
-        {
-            return delegate.getEstimatedDataSizeForStats(position);
-        }
-
-        @Override
-        public void retainedBytesForEachPart(ObjLongConsumer<Object> consumer)
-        {
-            delegate.retainedBytesForEachPart(consumer);
-        }
-
-        @Override
-        public boolean isNull(int position)
-        {
-            return delegate.isNull(position);
-        }
+        MaterializedResult actual = toMaterializedResult(TEST_SESSION, types, pages);
+        MaterializedResult expected = toMaterializedResult(TEST_SESSION, types, ImmutableList.of(initialPage));
+        assertThat(actual).containsExactlyElementsOf(expected);
     }
 }