Add fast path for single partition output

pettyjamesm · pettyjamesm · commit 1e9537314d13 · 2025-10-08T09:21:23.000-04:00
Adds special case handling for single partition output in
PagePartitioner such that partition calculations can be skipped when all
rows will necessarily go to a single output partition.
diff --git a/core/trino-main/src/main/java/io/trino/operator/output/PagePartitioner.java b/core/trino-main/src/main/java/io/trino/operator/output/PagePartitioner.java
@@ -44,6 +44,7 @@
 import java.util.function.IntUnaryOperator;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static io.trino.execution.buffer.PageSplitterUtil.splitAndSerializePage;
 import static io.trino.spi.block.PageBuilderStatus.DEFAULT_MAX_PAGE_SIZE_IN_BYTES;
@@ -138,7 +139,16 @@ public void partitionPage(Page page, OperatorContext operatorContext)
         }
 
         int outputPositionCount = replicatesAnyRow && !hasAnyRowBeenReplicated ? page.getPositionCount() + positionsAppenders.length - 1 : page.getPositionCount();
-        if (page.getPositionCount() < partitionFunction.partitionCount() * COLUMNAR_STRATEGY_COEFFICIENT) {
+        if (positionsAppenders.length == 1) {
+            // single output partition, skip partition calculation and append the entire page to the output partition
+            checkState(partitionFunction.partitionCount() == 1, "partitionFunction must be single partition");
+            // Any output rows are by definition "replicated" when only a single output partition exists
+            if (replicatesAnyRow && !hasAnyRowBeenReplicated) {
+                hasAnyRowBeenReplicated = true;
+            }
+            positionsAppenders[0].appendToOutputPartition(page);
+        }
+        else if (page.getPositionCount() < partitionFunction.partitionCount() * COLUMNAR_STRATEGY_COEFFICIENT) {
             // Partition will have on average less than COLUMNAR_STRATEGY_COEFFICIENT rows.
             // Doing it column-wise would degrade performance, so we fall back to row-wise approach.
             // Performance degradation is the worst in case of skewed hash distribution when only small subset
diff --git a/core/trino-main/src/main/java/io/trino/operator/output/PositionsAppender.java b/core/trino-main/src/main/java/io/trino/operator/output/PositionsAppender.java
@@ -14,14 +14,22 @@
 package io.trino.operator.output;
 
 import io.trino.spi.block.Block;
+import io.trino.spi.block.DictionaryBlock;
 import io.trino.spi.block.ValueBlock;
 import it.unimi.dsi.fastutil.ints.IntArrayList;
 
 public sealed interface PositionsAppender
         permits RowPositionsAppender, TypedPositionsAppender
 {
+    /**
+     * Appends the positions from the list, in the specified order. Implementations are not permitted to modify
+     * the contents of the {@link IntArrayList} argument as it may be passed directly from {@link DictionaryBlock#getRawIds()}
+     * without a defensive copy.
+     */
     void append(IntArrayList positions, ValueBlock source);
 
+    void appendRange(ValueBlock block, int offset, int length);
+
     /**
      * Appends the specified value positionCount times.
      * The result is the same as with using {@link PositionsAppender#append(IntArrayList, ValueBlock)} with
diff --git a/core/trino-main/src/main/java/io/trino/operator/output/PositionsAppenderPageBuilder.java b/core/trino-main/src/main/java/io/trino/operator/output/PositionsAppenderPageBuilder.java
@@ -73,6 +73,17 @@ private PositionsAppenderPageBuilder(
         }
     }
 
+    public void appendToOutputPartition(Page page)
+    {
+        int positionCount = page.getPositionCount();
+        declarePositions(positionCount);
+
+        for (int channel = 0; channel < channelAppenders.length; channel++) {
+            Block block = page.getBlock(channel);
+            channelAppenders[channel].appendRange(block, 0, positionCount);
+        }
+    }
+
     public void appendToOutputPartition(Page page, IntArrayList positions)
     {
         declarePositions(positions.size());
diff --git a/core/trino-main/src/main/java/io/trino/operator/output/RowPositionsAppender.java b/core/trino-main/src/main/java/io/trino/operator/output/RowPositionsAppender.java
@@ -115,6 +115,47 @@ public void append(IntArrayList positions, ValueBlock block)
         resetSize();
     }
 
+    @Override
+    public void appendRange(ValueBlock block, int offset, int length)
+    {
+        checkArgument(block instanceof RowBlock, "Block must be instance of %s", RowBlock.class);
+        if (length == 0) {
+            return;
+        }
+
+        RowBlock sourceRowBlock = (RowBlock) block;
+        ensureCapacity(length);
+
+        Block[] rawFieldBlocks = sourceRowBlock.getRawFieldBlocks();
+        int startOffset = sourceRowBlock.getOffsetBase();
+
+        for (int i = 0; i < fieldAppenders.length; i++) {
+            fieldAppenders[i].appendRange(rawFieldBlocks[i], startOffset + offset, length);
+        }
+
+        boolean[] rawRowIsNull = sourceRowBlock.getRawRowIsNull();
+        if (rawRowIsNull != null) {
+            for (int i = 0; i < length; i++) {
+                boolean isNull = rawRowIsNull[startOffset + offset + i];
+                hasNullRow |= isNull;
+                hasNonNullRow |= !isNull;
+                if (hasNullRow & hasNonNullRow) {
+                    System.arraycopy(rawRowIsNull, startOffset + offset + i, rowIsNull, positionCount + i, length - i);
+                    break;
+                }
+                else {
+                    rowIsNull[positionCount + i] = isNull;
+                }
+            }
+        }
+        else {
+            hasNonNullRow = true;
+        }
+
+        positionCount += length;
+        resetSize();
+    }
+
     @Override
     public void appendRle(ValueBlock value, int rlePositionCount)
     {
diff --git a/core/trino-main/src/main/java/io/trino/operator/output/TypedPositionsAppender.java b/core/trino-main/src/main/java/io/trino/operator/output/TypedPositionsAppender.java
@@ -44,6 +44,12 @@ public void append(IntArrayList positions, ValueBlock block)
         blockBuilder.appendPositions(block, positions.elements(), 0, positions.size());
     }
 
+    @Override
+    public void appendRange(ValueBlock block, int offset, int length)
+    {
+        blockBuilder.appendRange(block, offset, length);
+    }
+
     @Override
     public void appendRle(ValueBlock block, int count)
     {
diff --git a/core/trino-main/src/main/java/io/trino/operator/output/UnnestingPositionsAppender.java b/core/trino-main/src/main/java/io/trino/operator/output/UnnestingPositionsAppender.java
@@ -23,6 +23,7 @@
 import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
 import jakarta.annotation.Nullable;
 
+import java.util.Arrays;
 import java.util.Optional;
 
 import static com.google.common.base.Preconditions.checkArgument;
@@ -32,6 +33,7 @@
 import static io.trino.operator.output.PositionsAppenderUtil.calculateBlockResetSize;
 import static io.trino.operator.output.PositionsAppenderUtil.calculateNewArraySize;
 import static java.lang.Math.max;
+import static java.util.Objects.checkFromIndexSize;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -68,6 +70,50 @@ public UnnestingPositionsAppender(PositionsAppender delegate, Optional<BlockPosi
         this.identicalOperator = identicalOperator.orElse(null);
     }
 
+    public void appendRange(Block source, int offset, int length)
+    {
+        if (length == 0) {
+            return;
+        }
+
+        switch (source) {
+            case RunLengthEncodedBlock rleBlock -> {
+                appendRle(rleBlock.getValue(), length);
+            }
+            case DictionaryBlock dictionaryBlock -> {
+                ValueBlock dictionary = dictionaryBlock.getDictionary();
+                if (state == State.UNINITIALIZED) {
+                    state = State.DICTIONARY;
+                    this.dictionary = dictionary;
+                    dictionaryIdsBuilder.appendRange(dictionaryBlock, offset, length);
+                }
+                else if (state == State.DICTIONARY && this.dictionary == dictionary) {
+                    dictionaryIdsBuilder.appendRange(dictionaryBlock, offset, length);
+                }
+                else {
+                    transitionToDirect();
+
+                    int[] rawIds = dictionaryBlock.getRawIds();
+                    int rawOffset = dictionaryBlock.getRawIdsOffset();
+                    checkFromIndexSize(rawOffset + offset, length, rawIds.length);
+                    IntArrayList positionsList;
+                    if (rawOffset + offset == 0) {
+                        // Fast path, no copy necessary
+                        positionsList = IntArrayList.wrap(rawIds, length);
+                    }
+                    else {
+                        positionsList = IntArrayList.wrap(Arrays.copyOfRange(rawIds, rawOffset + offset, rawOffset + offset + length));
+                    }
+                    delegate.append(positionsList, dictionary);
+                }
+            }
+            case ValueBlock valueBlock -> {
+                transitionToDirect();
+                delegate.appendRange(valueBlock, offset, length);
+            }
+        }
+    }
+
     public void append(IntArrayList positions, Block source)
     {
         if (positions.isEmpty()) {
@@ -275,6 +321,17 @@ public void appendPositions(IntArrayList positions, DictionaryBlock block)
             size += positions.size();
         }
 
+        public void appendRange(DictionaryBlock block, int offset, int length)
+        {
+            checkArgument(length > 0, "block has no positions");
+            checkFromIndexSize(offset, length, block.getPositionCount());
+            ensureCapacity(size + length);
+            int[] rawIds = block.getRawIds();
+            int rawIdsOffset = block.getRawIdsOffset();
+            System.arraycopy(rawIds, rawIdsOffset + offset, dictionaryIds, size, length);
+            size += length;
+        }
+
         public DictionaryIdsBuilder newBuilderLike()
         {
             if (size == 0) {
diff --git a/core/trino-main/src/test/java/io/trino/operator/output/TestPagePartitioner.java b/core/trino-main/src/test/java/io/trino/operator/output/TestPagePartitioner.java
diff --git a/core/trino-main/src/test/java/io/trino/operator/output/TestPositionsAppender.java b/core/trino-main/src/test/java/io/trino/operator/output/TestPositionsAppender.java

Original file line number	Diff line number	Diff line change
`@@ -73,6 +73,17 @@ private PositionsAppenderPageBuilder(`
`73`	`73`	`}`
`74`	`74`	`}`
`75`	`75`
	`76`	`+ public void appendToOutputPartition(Page page)`
	`77`	`+ {`
	`78`	`+ int positionCount = page.getPositionCount();`
	`79`	`+ declarePositions(positionCount);`
	`80`	`+`
	`81`	`+ for (int channel = 0; channel < channelAppenders.length; channel++) {`
	`82`	`+ Block block = page.getBlock(channel);`
	`83`	`+ channelAppenders[channel].appendRange(block, 0, positionCount);`
	`84`	`+ }`
	`85`	`+ }`
	`86`	`+`
`76`	`87`	`public void appendToOutputPartition(Page page, IntArrayList positions)`
`77`	`88`	`{`
`78`	`89`	`declarePositions(positions.size());`
Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,12 @@ public void append(IntArrayList positions, ValueBlock block)`
`44`	`44`	`blockBuilder.appendPositions(block, positions.elements(), 0, positions.size());`
`45`	`45`	`}`
`46`	`46`
	`47`	`+ @Override`
	`48`	`+ public void appendRange(ValueBlock block, int offset, int length)`
	`49`	`+ {`
	`50`	`+ blockBuilder.appendRange(block, offset, length);`
	`51`	`+ }`
	`52`	`+`
`47`	`53`	`@Override`
`48`	`54`	`public void appendRle(ValueBlock block, int count)`
`49`	`55`	`{`