This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 25e378b5b07 Prune unused passthrough columns from UNNEST output
(opt-in) (#18782)
25e378b5b07 is described below
commit 25e378b5b07198fa1ebf2d0671572c3d6bfd24c2
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Thu Jun 25 17:17:21 2026 +0200
Prune unused passthrough columns from UNNEST output (opt-in) (#18782)
---
.../MultiStageBrokerRequestHandler.java | 4 +
pinot-common/src/main/proto/plan.proto | 5 +
.../tests/custom/UnnestIntegrationTest.java | 38 ++++
.../org/apache/pinot/query/QueryEnvironment.java | 31 +++-
.../planner/logical/PinotLogicalQueryPlanner.java | 6 +-
.../planner/logical/RelToPlanNodeConverter.java | 195 +++++++++++++++++++++
.../pinot/query/planner/plannode/UnnestNode.java | 44 ++++-
.../query/planner/serde/PlanNodeDeserializer.java | 10 +-
.../query/planner/serde/PlanNodeSerializer.java | 4 +-
.../query/planner/plannode/UnnestNodeTest.java | 32 ++++
.../query/planner/serde/PlanNodeSerDeTest.java | 40 +++++
.../pinot/query/queries/UnnestSqlPlannerTest.java | 149 ++++++++++++++++
.../query/runtime/operator/UnnestOperator.java | 28 ++-
.../query/runtime/operator/UnnestOperatorTest.java | 113 ++++++++++++
.../apache/pinot/spi/utils/CommonConstants.java | 15 ++
15 files changed, 702 insertions(+), 12 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
index 725c0925bf7..d9d60bbaad4 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
@@ -558,6 +558,9 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
int sortExchangeCopyThreshold = _config.getProperty(
CommonConstants.Broker.CONFIG_OF_SORT_EXCHANGE_COPY_THRESHOLD,
CommonConstants.Broker.DEFAULT_SORT_EXCHANGE_COPY_THRESHOLD);
+ boolean defaultUnnestColumnPruning = _config.getProperty(
+ CommonConstants.Broker.CONFIG_OF_UNNEST_COLUMN_PRUNING,
+ CommonConstants.Broker.DEFAULT_UNNEST_COLUMN_PRUNING);
WorkerManager workerManager =
QueryOptionsUtils.isMultiClusterRoutingEnabled(queryOptions, false)
? _multiClusterWorkerManager : _workerManager;
return QueryEnvironment.configBuilder()
@@ -569,6 +572,7 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
.isNullHandlingEnabled(QueryOptionsUtils.isNullHandlingEnabled(queryOptions))
.defaultInferPartitionHint(inferPartitionHint)
.defaultUseSpools(defaultUseSpool)
+ .defaultUnnestColumnPruning(defaultUnnestColumnPruning)
.defaultUseLeafServerForIntermediateStage(defaultUseLeafServerForIntermediateStage)
.defaultEnableGroupTrim(defaultEnableGroupTrim)
.defaultEnableDynamicFilteringSemiJoin(defaultEnableDynamicFilteringSemiJoin)
diff --git a/pinot-common/src/main/proto/plan.proto
b/pinot-common/src/main/proto/plan.proto
index 352cd92f4c1..8c2917d0968 100644
--- a/pinot-common/src/main/proto/plan.proto
+++ b/pinot-common/src/main/proto/plan.proto
@@ -236,6 +236,11 @@ message UnnestNode {
repeated int32 elementIndexes = 3;
// Optional absolute output index of the ordinality column in the stage
schema.
optional int32 ordinalityIndex = 4;
+ // When prunedPassthrough is true, the input-row column indexes copied into
the output, in output order.
+ // When absent/empty (old brokers), the operator copies the whole input row
(legacy behavior).
+ repeated int32 passthroughInputIndexes = 5;
+ // Whether the output schema was pruned to only the input columns referenced
downstream. Absent => false => legacy.
+ bool prunedPassthrough = 6;
}
enum WindowFrameType {
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UnnestIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UnnestIntegrationTest.java
index 5e824af94f6..c93522a16e3 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UnnestIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UnnestIntegrationTest.java
@@ -26,6 +26,7 @@ import org.apache.avro.generic.GenericData;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.data.Schema;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
@@ -303,6 +304,43 @@ public class UnnestIntegrationTest extends
CustomDataQueryClusterIntegrationTest
assertEquals(sum, 6 * getCountStarResult());
}
+ @DataProvider(name = "columnPruningQueries")
+ public Object[][] columnPruningQueries() {
+ String table = getTableName();
+ return new Object[][]{
+ // single array, source array dropped
+ {String.format("SELECT intCol, u.elem FROM %s CROSS JOIN
UNNEST(stringArrayCol) AS u(elem) "
+ + "ORDER BY intCol, u.elem", table)},
+ // select only the unnested element (zero passthrough)
+ {String.format("SELECT u.elem FROM %s CROSS JOIN
UNNEST(stringArrayCol) AS u(elem) ORDER BY u.elem", table)},
+ // WITH ORDINALITY (ordinality index recomputed against pruned output)
+ {String.format("SELECT intCol, u.elem, u.idx FROM %s CROSS JOIN
UNNEST(stringArrayCol) WITH ORDINALITY "
+ + "AS u(elem, idx) ORDER BY intCol, u.idx", table)},
+ // multiple arrays, both source arrays dropped
+ {String.format("SELECT intCol, u.longValue, u.stringValue FROM %s "
+ + "CROSS JOIN UNNEST(longArrayCol, stringArrayCol) AS u(longValue,
stringValue) "
+ + "ORDER BY intCol, u.longValue", table)},
+ // source array also selected (must be retained)
+ {String.format("SELECT intCol, stringArrayCol, u.elem FROM %s CROSS
JOIN UNNEST(stringArrayCol) AS u(elem) "
+ + "ORDER BY intCol, u.elem", table)}
+ };
+ }
+
+ @Test(dataProvider = "columnPruningQueries")
+ public void testCrossJoinUnnestColumnPruningMatchesDefault(String select)
+ throws Exception {
+ // Pruning input/passthrough columns from the UNNEST output must not
change results. Verified end-to-end on the
+ // multi-stage engine (the only engine that supports UNNEST CROSS JOIN),
with the flag toggled per query.
+ setUseMultiStageQueryEngine(true);
+ JsonNode defaultRows = postQuery(select).get("resultTable").get("rows");
+ JsonNode prunedRows = postQuery("SET unnestColumnPruning=true; " +
select).get("resultTable").get("rows");
+
+ assertNotNull(defaultRows);
+ assertNotNull(prunedRows);
+ assertEquals(prunedRows.toString(), defaultRows.toString(),
+ "Column pruning must not change UNNEST results for: " + select);
+ }
+
@Override
public String getTableName() {
return DEFAULT_TABLE_NAME;
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
index 2fb657de1fd..5801a0b6e5e 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
@@ -536,7 +536,7 @@ public class QueryEnvironment {
return pinotDispatchPlanner.createDispatchableSubPlanV2(plan.getLeft(),
plan.getRight());
}
SubPlan plan = PinotLogicalQueryPlanner.makePlan(relRoot, tracker,
useSpools(plannerContext.getOptions()),
- _envConfig.defaultHashFunction());
+ _envConfig.defaultHashFunction(),
pruneUnnestColumns(plannerContext.getOptions()));
PinotDispatchPlanner pinotDispatchPlanner =
new PinotDispatchPlanner(plannerContext,
_envConfig.getWorkerManager(), _envConfig.getRequestId(),
_envConfig.getTableCache());
@@ -719,6 +719,23 @@ public class QueryEnvironment {
return Boolean.parseBoolean(optionValue);
}
+ /**
+ * Whether to prune unused input (passthrough) columns from the UNNEST
output. Defaults to the broker config
+ * {@link Config#defaultUnnestColumnPruning()} (itself {@code false} by
default) because a broker emitting the
+ * smaller schema cannot be honored by an un-upgraded server; enable only
once all servers support it. Can be
+ * overridden per query via {@link
CommonConstants.Broker.Request.QueryOptionKey#UNNEST_COLUMN_PRUNING}.
+ * <p>
+ * NOTE: This is a no-op under {@code usePhysicalOptimizer} (the v2 path
does not go through
+ * {@link org.apache.pinot.query.planner.logical.RelToPlanNodeConverter}).
+ */
+ public boolean pruneUnnestColumns(Map<String, String> options) {
+ String optionValue =
options.get(CommonConstants.Broker.Request.QueryOptionKey.UNNEST_COLUMN_PRUNING);
+ if (optionValue == null) {
+ return _envConfig.defaultUnnestColumnPruning();
+ }
+ return Boolean.parseBoolean(optionValue);
+ }
+
@Value.Immutable
public interface Config {
@@ -780,6 +797,18 @@ public class QueryEnvironment {
return CommonConstants.Broker.DEFAULT_OF_SPOOLS;
}
+ /**
+ * Whether to prune unused input (passthrough) columns from the UNNEST
output by default.
+ *
+ * This is treated as the default value for the broker and it is expected
to be obtained from a Pinot configuration.
+ * This default value can be always overridden at query level by the query
option
+ * {@link
CommonConstants.Broker.Request.QueryOptionKey#UNNEST_COLUMN_PRUNING}.
+ */
+ @Value.Default
+ default boolean defaultUnnestColumnPruning() {
+ return CommonConstants.Broker.DEFAULT_UNNEST_COLUMN_PRUNING;
+ }
+
/// Whether to only use servers for leaf stages as the workers for the
intermediate stages.
/// This is useful to control the fanout of the query and reduce data
shuffling.
@Value.Default
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java
index e1c38e462ec..eceb18102d9 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java
@@ -46,6 +46,7 @@ import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.MailboxReceiveNode;
import org.apache.pinot.query.planner.plannode.MailboxSendNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.spi.utils.CommonConstants;
/**
@@ -60,8 +61,9 @@ public class PinotLogicalQueryPlanner {
*/
public static SubPlan makePlan(RelRoot relRoot,
@Nullable TransformationTracker.Builder<PlanNode, RelNode> tracker,
boolean useSpools,
- String hashFunction) {
- PlanNode rootNode = new RelToPlanNodeConverter(tracker,
hashFunction).toPlanNode(relRoot.rel);
+ String hashFunction, boolean pruneUnnestColumns) {
+ PlanNode rootNode = new RelToPlanNodeConverter(tracker, hashFunction,
+ !CommonConstants.Helix.DEFAULT_ENABLE_CASE_INSENSITIVE,
pruneUnnestColumns).toPlanNode(relRoot.rel);
PlanFragment rootFragment = planNodeToPlanFragment(rootNode, tracker,
useSpools, hashFunction);
return new SubPlan(rootFragment,
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java
index c347e18a79d..66271d5c006 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java
@@ -21,6 +21,7 @@ package org.apache.pinot.query.planner.logical;
import com.google.common.base.Preconditions;
import com.google.common.collect.Sets;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.List;
import java.util.Set;
import javax.annotation.Nullable;
@@ -111,6 +112,8 @@ public final class RelToPlanNodeConverter {
private final TransformationTracker.Builder<PlanNode, RelNode> _tracker;
private final String _hashFunction;
private final boolean _caseSensitive;
+ // When true, UNNEST output is pruned to drop input (passthrough) columns
not referenced downstream. Default off.
+ private final boolean _pruneUnnestColumns;
public RelToPlanNodeConverter(@Nullable
TransformationTracker.Builder<PlanNode, RelNode> tracker,
String hashFunction) {
@@ -119,9 +122,15 @@ public final class RelToPlanNodeConverter {
public RelToPlanNodeConverter(@Nullable
TransformationTracker.Builder<PlanNode, RelNode> tracker,
String hashFunction, boolean caseSensitive) {
+ this(tracker, hashFunction, caseSensitive, false);
+ }
+
+ public RelToPlanNodeConverter(@Nullable
TransformationTracker.Builder<PlanNode, RelNode> tracker,
+ String hashFunction, boolean caseSensitive, boolean pruneUnnestColumns) {
_tracker = tracker;
_hashFunction = hashFunction;
_caseSensitive = caseSensitive;
+ _pruneUnnestColumns = pruneUnnestColumns;
}
/**
@@ -227,6 +236,9 @@ public final class RelToPlanNodeConverter {
convertInputs(node.getInputs()), arrayExprs, tableFunctionContext);
}
+ // NOTE: Besides the main dispatch, this is also invoked from
tryPruneUnnestPassthrough. Its result is always used
+ // (either directly or rewritten into a pruned node), so it must remain free
of non-idempotent side effects beyond
+ // converting/tracking its own subtree.
private BasePlanNode convertLogicalCorrelate(LogicalCorrelate node) {
// Pattern: Correlate(left, Uncollect(Project(correlatedFields...)))
RelNode right = node.getRight();
@@ -709,10 +721,193 @@ public final class RelToPlanNodeConverter {
}
private ProjectNode convertLogicalProject(LogicalProject node) {
+ if (_pruneUnnestColumns) {
+ ProjectNode pruned = tryPruneUnnestPassthrough(node);
+ if (pruned != null) {
+ return pruned;
+ }
+ }
return new ProjectNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()),
NodeHint.fromRelHints(node.getHints()),
convertInputs(node.getInputs()),
RexExpressionUtils.fromRexNodes(node.getProjects()));
}
+ /**
+ * When a Project sits directly above a CROSS JOIN UNNEST (a {@link
LogicalCorrelate} over {@link Uncollect}), prunes
+ * the input (passthrough) columns that the Project does not reference -
notably the unnested source array - from the
+ * {@link UnnestNode} output, so the operator does not copy them into every
exploded row. The Project's own
+ * {@code InputRef}s are remapped from the full correlate-output index space
to the pruned space.
+ * <p>
+ * Returns {@code null} (caller falls back to the default conversion,
preserving all passthrough columns) when the
+ * pattern is not recognized, the layout is not the standard one, or pruning
would be a no-op.
+ */
+ @Nullable
+ private ProjectNode tryPruneUnnestPassthrough(LogicalProject node) {
+ List<RelNode> inputs = node.getInputs();
+ if (inputs.size() != 1 || !(inputs.get(0) instanceof LogicalCorrelate)) {
+ return null;
+ }
+ LogicalCorrelate correlate = (LogicalCorrelate) inputs.get(0);
+ RelNode right = correlate.getRight();
+ // Only handle the UNNEST pattern, and only when no correlate-filter wraps
the Uncollect. With a wrapping filter,
+ // convertLogicalCorrelate emits a FilterNode, so the Project no longer
sits directly on the UnnestNode.
+ if (findUncollect(right) == null || findCorrelateFilter(right) != null) {
+ // Cheap structural check failed before any conversion: let the caller
fall back to the default conversion.
+ return null;
+ }
+ // Convert the correlate exactly once. From here we always return a
ProjectNode wrapping this converted input, so
+ // the caller never re-runs the (side-effecting, tracker-registering)
conversion of the left subtree.
+ BasePlanNode converted = convertLogicalCorrelate(correlate);
+ if (converted instanceof UnnestNode) {
+ ProjectNode pruned = buildPrunedProject(node, (UnnestNode) converted,
correlate);
+ if (pruned != null) {
+ return pruned;
+ }
+ }
+ // Not prunable (unexpected layout, no-op, or non-UnnestNode): wrap the
already-converted input in a Project.
+ if (_tracker != null) {
+ _tracker.trackCreation(correlate, converted);
+ }
+ return new ProjectNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()),
NodeHint.fromRelHints(node.getHints()),
+ new ArrayList<>(List.of(converted)),
RexExpressionUtils.fromRexNodes(node.getProjects()));
+ }
+
+ /**
+ * Builds the pruned Project-over-UnnestNode for the standard CROSS JOIN
UNNEST layout, or returns {@code null} when
+ * the layout is non-standard or pruning would be a no-op (all passthrough
columns are referenced downstream). On
+ * success it registers the new UnnestNode with the tracker against {@code
correlate}.
+ */
+ @Nullable
+ private ProjectNode buildPrunedProject(LogicalProject node, UnnestNode full,
RelNode correlate) {
+ if (full.isPrunedPassthrough()) {
+ return null;
+ }
+ DataSchema fullSchema = full.getDataSchema();
+ int outSize = fullSchema.size();
+ int numArrays = full.getArrayExprs().size();
+ boolean withOrdinality = full.isWithOrdinality();
+ int leftCount = outSize - numArrays - (withOrdinality ? 1 : 0);
+ if (leftCount <= 0) {
+ // No passthrough columns to prune.
+ return null;
+ }
+ // Only handle the clean layout produced by the standard path: passthrough
columns occupy [0, leftCount) and the
+ // element (then ordinality) columns occupy the trailing region. Fall back
otherwise.
+ List<Integer> elementIndexes = full.getElementIndexes();
+ if (elementIndexes.size() != numArrays) {
+ return null;
+ }
+ for (int k = 0; k < numArrays; k++) {
+ if (elementIndexes.get(k) != leftCount + k) {
+ return null;
+ }
+ }
+ if (withOrdinality && full.getOrdinalityIndex() != leftCount + numArrays) {
+ return null;
+ }
+ // Collect referenced passthrough (left) columns from the project
expressions.
+ List<RexExpression> projects =
RexExpressionUtils.fromRexNodes(node.getProjects());
+ boolean[] referenced = new boolean[leftCount];
+ for (RexExpression project : projects) {
+ collectReferencedColumns(project, leftCount, referenced);
+ }
+ List<Integer> retained = new ArrayList<>();
+ for (int i = 0; i < leftCount; i++) {
+ if (referenced[i]) {
+ retained.add(i);
+ }
+ }
+ if (retained.size() == leftCount) {
+ // All passthrough columns are used downstream; pruning would be a no-op.
+ return null;
+ }
+ int numRetained = retained.size();
+ // Build the old (full correlate output) -> new (pruned output) index map.
+ int[] oldToNew = new int[outSize];
+ Arrays.fill(oldToNew, -1);
+ for (int i = 0; i < numRetained; i++) {
+ oldToNew[retained.get(i)] = i;
+ }
+ for (int k = 0; k < numArrays; k++) {
+ oldToNew[leftCount + k] = numRetained + k;
+ }
+ if (withOrdinality) {
+ oldToNew[leftCount + numArrays] = numRetained + numArrays;
+ }
+ // Build the pruned output schema: retained passthrough columns, then
element columns, then ordinality.
+ int newSize = numRetained + numArrays + (withOrdinality ? 1 : 0);
+ String[] columnNames = new String[newSize];
+ ColumnDataType[] columnTypes = new ColumnDataType[newSize];
+ for (int i = 0; i < numRetained; i++) {
+ int oldIdx = retained.get(i);
+ columnNames[i] = fullSchema.getColumnName(oldIdx);
+ columnTypes[i] = fullSchema.getColumnDataType(oldIdx);
+ }
+ for (int k = 0; k < numArrays; k++) {
+ columnNames[numRetained + k] = fullSchema.getColumnName(leftCount + k);
+ columnTypes[numRetained + k] = fullSchema.getColumnDataType(leftCount +
k);
+ }
+ if (withOrdinality) {
+ columnNames[numRetained + numArrays] =
fullSchema.getColumnName(leftCount + numArrays);
+ columnTypes[numRetained + numArrays] =
fullSchema.getColumnDataType(leftCount + numArrays);
+ }
+ DataSchema prunedSchema = new DataSchema(columnNames, columnTypes);
+ // Recompute element/ordinality indexes against the pruned (smaller)
output.
+ List<Integer> newElementIndexes = new ArrayList<>(numArrays);
+ for (int k = 0; k < numArrays; k++) {
+ newElementIndexes.add(numRetained + k);
+ }
+ int newOrdinalityIndex = withOrdinality ? numRetained + numArrays :
UnnestNode.UNSPECIFIED_INDEX;
+ UnnestNode.TableFunctionContext context = new
UnnestNode.TableFunctionContext(withOrdinality, newElementIndexes,
+ newOrdinalityIndex, retained, true);
+ UnnestNode prunedUnnest = new UnnestNode(DEFAULT_STAGE_ID, prunedSchema,
full.getNodeHint(), full.getInputs(),
+ full.getArrayExprs(), context);
+ if (_tracker != null) {
+ _tracker.trackCreation(correlate, prunedUnnest);
+ }
+ // Remap the project's input refs from the full correlate-output space to
the pruned space.
+ List<RexExpression> remappedProjects = new ArrayList<>(projects.size());
+ for (RexExpression project : projects) {
+ remappedProjects.add(remapInputRefs(project, oldToNew));
+ }
+ return new ProjectNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()),
NodeHint.fromRelHints(node.getHints()),
+ new ArrayList<>(List.of(prunedUnnest)), remappedProjects);
+ }
+
+ private static void collectReferencedColumns(RexExpression expr, int
leftCount, boolean[] referenced) {
+ if (expr instanceof RexExpression.InputRef) {
+ int idx = ((RexExpression.InputRef) expr).getIndex();
+ if (idx >= 0 && idx < leftCount) {
+ referenced[idx] = true;
+ }
+ } else if (expr instanceof RexExpression.FunctionCall) {
+ for (RexExpression op : ((RexExpression.FunctionCall)
expr).getFunctionOperands()) {
+ collectReferencedColumns(op, leftCount, referenced);
+ }
+ }
+ }
+
+ private static RexExpression remapInputRefs(RexExpression expr, int[]
oldToNew) {
+ if (expr instanceof RexExpression.InputRef) {
+ int idx = ((RexExpression.InputRef) expr).getIndex();
+ int mapped = (idx >= 0 && idx < oldToNew.length) ? oldToNew[idx] : -1;
+ // mapped is always >= 0 here: passthrough columns referenced by the
project were retained, and element/
+ // ordinality columns are always mapped. Guard defensively to avoid
silently corrupting indexes.
+ Preconditions.checkState(mapped >= 0, "Unexpected unmapped InputRef
index %s while pruning UNNEST passthrough",
+ idx);
+ return new RexExpression.InputRef(mapped);
+ } else if (expr instanceof RexExpression.FunctionCall) {
+ RexExpression.FunctionCall fc = (RexExpression.FunctionCall) expr;
+ List<RexExpression> ops = fc.getFunctionOperands();
+ List<RexExpression> rewritten = new ArrayList<>(ops.size());
+ for (RexExpression op : ops) {
+ rewritten.add(remapInputRefs(op, oldToNew));
+ }
+ return new RexExpression.FunctionCall(fc.getDataType(),
fc.getFunctionName(), rewritten);
+ } else {
+ return expr;
+ }
+ }
+
private FilterNode convertLogicalFilter(LogicalFilter node) {
return new FilterNode(DEFAULT_STAGE_ID, toDataSchema(node.getRowType()),
NodeHint.fromRelHints(node.getHints()),
convertInputs(node.getInputs()),
RexExpressionUtils.fromRexNode(node.getCondition()));
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/UnnestNode.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/UnnestNode.java
index 1135de4312c..786659a208f 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/UnnestNode.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/UnnestNode.java
@@ -106,6 +106,22 @@ public class UnnestNode extends BasePlanNode {
return _tableFunctionContext.getOrdinalityIndex();
}
+ /**
+ * Input-row column indexes copied into the output, in output order. Only
meaningful when
+ * {@link #isPrunedPassthrough()} is {@code true}; otherwise empty and the
operator copies the whole input row.
+ */
+ public List<Integer> getPassthroughInputIndexes() {
+ return _tableFunctionContext.getPassthroughInputIndexes();
+ }
+
+ /**
+ * Whether the output schema has been pruned to only the input columns
referenced downstream (plus the element and
+ * ordinality columns). When {@code false}, the operator preserves legacy
behavior of copying the whole input row.
+ */
+ public boolean isPrunedPassthrough() {
+ return _tableFunctionContext.isPrunedPassthrough();
+ }
+
@Override
public String explain() {
return "UNNEST";
@@ -157,11 +173,21 @@ public class UnnestNode extends BasePlanNode {
private final boolean _withOrdinality;
private final List<Integer> _elementIndexes;
private final int _ordinalityIndex;
+ private final List<Integer> _passthroughInputIndexes;
+ private final boolean _prunedPassthrough;
public TableFunctionContext(boolean withOrdinality, List<Integer>
elementIndexes, int ordinalityIndex) {
+ // Legacy/default: no passthrough pruning. The operator copies the whole
input row into the output.
+ this(withOrdinality, elementIndexes, ordinalityIndex, List.of(), false);
+ }
+
+ public TableFunctionContext(boolean withOrdinality, List<Integer>
elementIndexes, int ordinalityIndex,
+ List<Integer> passthroughInputIndexes, boolean prunedPassthrough) {
_withOrdinality = withOrdinality;
_elementIndexes = List.copyOf(elementIndexes);
_ordinalityIndex = ordinalityIndex;
+ _passthroughInputIndexes = List.copyOf(passthroughInputIndexes);
+ _prunedPassthrough = prunedPassthrough;
}
public boolean isWithOrdinality() {
@@ -176,8 +202,17 @@ public class UnnestNode extends BasePlanNode {
return _ordinalityIndex;
}
+ public List<Integer> getPassthroughInputIndexes() {
+ return _passthroughInputIndexes;
+ }
+
+ public boolean isPrunedPassthrough() {
+ return _prunedPassthrough;
+ }
+
public TableFunctionContext copy() {
- return new TableFunctionContext(_withOrdinality, _elementIndexes,
_ordinalityIndex);
+ return new TableFunctionContext(_withOrdinality, _elementIndexes,
_ordinalityIndex, _passthroughInputIndexes,
+ _prunedPassthrough);
}
@Override
@@ -190,12 +225,15 @@ public class UnnestNode extends BasePlanNode {
}
TableFunctionContext that = (TableFunctionContext) o;
return _withOrdinality == that._withOrdinality && _ordinalityIndex ==
that._ordinalityIndex
- && Objects.equals(_elementIndexes, that._elementIndexes);
+ && _prunedPassthrough == that._prunedPassthrough
+ && Objects.equals(_elementIndexes, that._elementIndexes)
+ && Objects.equals(_passthroughInputIndexes,
that._passthroughInputIndexes);
}
@Override
public int hashCode() {
- return Objects.hash(_withOrdinality, _elementIndexes, _ordinalityIndex);
+ return Objects.hash(_withOrdinality, _elementIndexes, _ordinalityIndex,
_passthroughInputIndexes,
+ _prunedPassthrough);
}
}
}
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java
index 7f2144124d3..c91e46d8f5f 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java
@@ -240,8 +240,16 @@ public class PlanNodeDeserializer {
int ordIdx = protoUnnestNode.hasOrdinalityIndex() ?
protoUnnestNode.getOrdinalityIndex()
: UnnestNode.UNSPECIFIED_INDEX;
+ // Passthrough pruning metadata. Absent on plans produced by older
brokers, in which case prunedPassthrough is
+ // false and the operator copies the whole input row (legacy behavior).
+ List<Integer> passthroughInputIndexes = new ArrayList<>();
+ for (int idx : protoUnnestNode.getPassthroughInputIndexesList()) {
+ passthroughInputIndexes.add(idx);
+ }
+
UnnestNode.TableFunctionContext context =
- new
UnnestNode.TableFunctionContext(protoUnnestNode.getWithOrdinality(),
elementIndexes, ordIdx);
+ new
UnnestNode.TableFunctionContext(protoUnnestNode.getWithOrdinality(),
elementIndexes, ordIdx,
+ passthroughInputIndexes, protoUnnestNode.getPrunedPassthrough());
return new UnnestNode(protoNode.getStageId(),
extractDataSchema(protoNode), extractNodeHint(protoNode),
extractInputs(protoNode), arrayExprs, context);
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java
index 99665d5c29b..f2283ed09f5 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java
@@ -287,7 +287,9 @@ public class PlanNodeSerializer {
.addAllArrayExprs(convertExpressions(node.getArrayExprs()))
.setWithOrdinality(context.isWithOrdinality())
.addAllElementIndexes(context.getElementIndexes())
- .setOrdinalityIndex(context.getOrdinalityIndex());
+ .setOrdinalityIndex(context.getOrdinalityIndex())
+ .addAllPassthroughInputIndexes(context.getPassthroughInputIndexes())
+ .setPrunedPassthrough(context.isPrunedPassthrough());
builder.setUnnestNode(unnestNodeBuilder.build());
return null;
}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/plannode/UnnestNodeTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/plannode/UnnestNodeTest.java
index 0396d2dd284..6045cd8aad0 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/plannode/UnnestNodeTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/plannode/UnnestNodeTest.java
@@ -185,4 +185,36 @@ public class UnnestNodeTest {
Assert.assertEquals(node1.hashCode(), node2.hashCode());
Assert.assertNotEquals(node1, node3);
}
+
+ @Test
+ public void testDefaultPassthroughIsNotPruned() {
+ DataSchema dataSchema = new DataSchema(new String[]{"id", "elem"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT,
DataSchema.ColumnDataType.STRING});
+ UnnestNode node = new UnnestNode(0, dataSchema, PlanNode.NodeHint.EMPTY,
+ new ArrayList<>(), List.of(new RexExpression.InputRef(1)),
List.of("elem"), false, null);
+
+ // Legacy/default construction: no pruning, empty passthrough map.
+ Assert.assertFalse(node.isPrunedPassthrough());
+ Assert.assertTrue(node.getPassthroughInputIndexes().isEmpty());
+ }
+
+ @Test
+ public void testPrunedPassthroughMetadata() {
+ DataSchema dataSchema = new DataSchema(new String[]{"id", "elem"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT,
DataSchema.ColumnDataType.STRING});
+ UnnestNode.TableFunctionContext context =
+ new UnnestNode.TableFunctionContext(false, List.of(1),
UnnestNode.UNSPECIFIED_INDEX, List.of(0), true);
+ UnnestNode node = new UnnestNode(0, dataSchema, PlanNode.NodeHint.EMPTY,
+ new ArrayList<>(), List.of(new RexExpression.InputRef(1)), context);
+
+ Assert.assertTrue(node.isPrunedPassthrough());
+ Assert.assertEquals(node.getPassthroughInputIndexes(), List.of(0));
+
+ // A pruned node must not equal an otherwise-identical non-pruned node.
+ UnnestNode legacy = new UnnestNode(0, dataSchema, PlanNode.NodeHint.EMPTY,
+ new ArrayList<>(), List.of(new RexExpression.InputRef(1)),
+ new UnnestNode.TableFunctionContext(false, List.of(1),
UnnestNode.UNSPECIFIED_INDEX));
+ Assert.assertNotEquals(node, legacy);
+ Assert.assertNotEquals(node.hashCode(), legacy.hashCode());
+ }
}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java
index 8fc02e78122..98745869d70 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/PlanNodeSerDeTest.java
@@ -18,10 +18,16 @@
*/
package org.apache.pinot.query.planner.serde;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.query.QueryEnvironmentTestBase;
+import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.planner.plannode.UnnestNode;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
@@ -38,4 +44,38 @@ public class PlanNodeSerDeTest extends
QueryEnvironmentTestBase {
assertEquals(stagePlan, deserializedStagePlan);
}
}
+
+ @Test
+ public void testPrunedUnnestNodeSerDe() {
+ // Round-trips the passthrough-pruning wire fields
(passthroughInputIndexes, prunedPassthrough). A non-sequential
+ // index list plus WITH ORDINALITY exercise the proto repeated/bool fields
and ordering.
+ DataSchema dataSchema = new DataSchema(new String[]{"col0", "col2",
"elem", "ord"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING,
ColumnDataType.INT, ColumnDataType.INT});
+ UnnestNode.TableFunctionContext context =
+ new UnnestNode.TableFunctionContext(true, List.of(2), 3, List.of(0,
2), true);
+ UnnestNode node = new UnnestNode(1, dataSchema, PlanNode.NodeHint.EMPTY,
new ArrayList<>(),
+ List.of(new RexExpression.InputRef(1)), context);
+
+ PlanNode deserialized =
PlanNodeDeserializer.process(PlanNodeSerializer.process(node));
+ assertEquals(deserialized, node);
+ UnnestNode deserializedUnnest = (UnnestNode) deserialized;
+ assertEquals(deserializedUnnest.getPassthroughInputIndexes(), List.of(0,
2));
+ assertEquals(deserializedUnnest.isPrunedPassthrough(), true);
+ assertEquals(deserializedUnnest.getOrdinalityIndex(), 3);
+ }
+
+ @Test
+ public void testLegacyUnnestNodeSerDe() {
+ // A non-pruned UnnestNode must round-trip with prunedPassthrough=false
and an empty passthrough map (the wire
+ // default an old broker produces).
+ DataSchema dataSchema = new DataSchema(new String[]{"id", "arr", "elem"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT_ARRAY,
ColumnDataType.INT});
+ UnnestNode node = new UnnestNode(1, dataSchema, PlanNode.NodeHint.EMPTY,
new ArrayList<>(),
+ new RexExpression.InputRef(1), "elem", false, null);
+
+ UnnestNode deserialized = (UnnestNode)
PlanNodeDeserializer.process(PlanNodeSerializer.process(node));
+ assertEquals(deserialized, node);
+ assertEquals(deserialized.isPrunedPassthrough(), false);
+ assertEquals(deserialized.getPassthroughInputIndexes(), List.of());
+ }
}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/UnnestSqlPlannerTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/UnnestSqlPlannerTest.java
index 5a4d1d3d1df..58888ff4635 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/UnnestSqlPlannerTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/queries/UnnestSqlPlannerTest.java
@@ -21,6 +21,10 @@ package org.apache.pinot.query.queries;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import org.apache.pinot.common.config.provider.TableCache;
+import org.apache.pinot.core.routing.MockRoutingManagerFactory;
+import org.apache.pinot.core.routing.RoutingManager;
+import org.apache.pinot.query.QueryEnvironment;
import org.apache.pinot.query.QueryEnvironmentTestBase;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
@@ -29,6 +33,8 @@ import org.apache.pinot.query.planner.plannode.AggregateNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.planner.plannode.UnnestNode;
+import org.apache.pinot.query.routing.WorkerManager;
+import org.apache.pinot.spi.utils.CommonConstants;
import org.testng.Assert;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
@@ -85,6 +91,149 @@ public class UnnestSqlPlannerTest extends
QueryEnvironmentTestBase {
assertOrdinality(unnestNode, withOrdinality);
}
+ @Test
+ public void testUnnestColumnPruningDropsSourceArray() {
+ // With pruning enabled, the unnested source array (mcol1) must not be
carried in the UnnestNode output.
+ String sql = "SET unnestColumnPruning=true; "
+ + "SELECT e.col1, u.s FROM e CROSS JOIN UNNEST(e.mcol1) AS u(s)";
+ DispatchableSubPlan subPlan = _queryEnvironment.planQuery(sql);
+ List<UnnestNode> unnestNodes = findUnnestNodes(subPlan);
+ Assert.assertEquals(unnestNodes.size(), 1);
+ UnnestNode unnestNode = unnestNodes.get(0);
+
+ Assert.assertTrue(unnestNode.isPrunedPassthrough(), "UnnestNode should be
pruned");
+ List<String> columns =
Arrays.asList(unnestNode.getDataSchema().getColumnNames());
+ Assert.assertFalse(columns.contains("mcol1"),
+ "Source array column should be pruned from UnnestNode output, found: "
+ columns);
+ Assert.assertTrue(columns.contains("col1"), "Referenced passthrough column
col1 should be retained: " + columns);
+ Assert.assertEquals(unnestNode.getDataSchema().size(), 2, "Expected [col1,
element] only: " + columns);
+ // col1 is the only retained passthrough column (input index 0); the
element lands right after it.
+ Assert.assertEquals(unnestNode.getPassthroughInputIndexes(), List.of(0));
+ Assert.assertEquals(unnestNode.getElementIndexes(), List.of(1));
+ }
+
+ @Test
+ public void testUnnestColumnPruningDisabledByDefault() {
+ // Without the flag, behavior is unchanged: the source array is still part
of the UnnestNode output.
+ String sql = "SELECT e.col1, u.s FROM e CROSS JOIN UNNEST(e.mcol1) AS
u(s)";
+ DispatchableSubPlan subPlan = _queryEnvironment.planQuery(sql);
+ List<UnnestNode> unnestNodes = findUnnestNodes(subPlan);
+ Assert.assertEquals(unnestNodes.size(), 1);
+ UnnestNode unnestNode = unnestNodes.get(0);
+
+ Assert.assertFalse(unnestNode.isPrunedPassthrough(), "UnnestNode should
not be pruned by default");
+ List<String> columns =
Arrays.asList(unnestNode.getDataSchema().getColumnNames());
+ Assert.assertTrue(columns.contains("mcol1"),
+ "Without pruning, the source array column should remain in the
UnnestNode output: " + columns);
+ }
+
+ @Test
+ public void testUnnestColumnPruningRetainsSelectedSourceArray() {
+ // When the user also selects the source array, it must be retained even
with pruning enabled.
+ String sql = "SET unnestColumnPruning=true; "
+ + "SELECT e.col1, e.mcol1, u.s FROM e CROSS JOIN UNNEST(e.mcol1) AS
u(s)";
+ DispatchableSubPlan subPlan = _queryEnvironment.planQuery(sql);
+ List<UnnestNode> unnestNodes = findUnnestNodes(subPlan);
+ Assert.assertEquals(unnestNodes.size(), 1);
+ UnnestNode unnestNode = unnestNodes.get(0);
+
+ List<String> columns =
Arrays.asList(unnestNode.getDataSchema().getColumnNames());
+ Assert.assertTrue(columns.contains("mcol1"),
+ "Explicitly selected source array must be retained: " + columns);
+ Assert.assertTrue(columns.contains("col1"), columns.toString());
+ }
+
+ @Test
+ public void testUnnestColumnPruningWithOrdinality() {
+ // The ordinality index must be recomputed against the pruned (smaller)
output.
+ String sql = "SET unnestColumnPruning=true; "
+ + "SELECT e.col1, u.s, u.ord FROM e CROSS JOIN UNNEST(e.mcol1) WITH
ORDINALITY AS u(s, ord)";
+ DispatchableSubPlan subPlan = _queryEnvironment.planQuery(sql);
+ List<UnnestNode> unnestNodes = findUnnestNodes(subPlan);
+ Assert.assertEquals(unnestNodes.size(), 1);
+ UnnestNode unnestNode = unnestNodes.get(0);
+
+ Assert.assertTrue(unnestNode.isWithOrdinality());
+ Assert.assertTrue(unnestNode.isPrunedPassthrough());
+ List<String> columns =
Arrays.asList(unnestNode.getDataSchema().getColumnNames());
+ Assert.assertFalse(columns.contains("mcol1"), "Source array should be
pruned: " + columns);
+ // Retained passthrough = [col1] (1), element at index 1, ordinality at
index 2.
+ Assert.assertEquals(unnestNode.getPassthroughInputIndexes(), List.of(0));
+ Assert.assertEquals(unnestNode.getElementIndexes(), List.of(1));
+ Assert.assertEquals(unnestNode.getOrdinalityIndex(), 2);
+ Assert.assertEquals(unnestNode.getDataSchema().size(), 3);
+ }
+
+ @Test
+ public void testUnnestColumnPruningMultipleArrays() {
+ // Both source arrays must be dropped; element indexes recompute
contiguously after the retained passthrough.
+ String sql = "SET unnestColumnPruning=true; "
+ + "SELECT e.col1, u.longVal, u.stringVal FROM e CROSS JOIN
UNNEST(e.mcol2, e.mcol1) AS u(longVal, stringVal)";
+ DispatchableSubPlan subPlan = _queryEnvironment.planQuery(sql);
+ List<UnnestNode> unnestNodes = findUnnestNodes(subPlan);
+ Assert.assertEquals(unnestNodes.size(), 1);
+ UnnestNode unnestNode = unnestNodes.get(0);
+
+ Assert.assertTrue(unnestNode.isPrunedPassthrough());
+ List<String> columns =
Arrays.asList(unnestNode.getDataSchema().getColumnNames());
+ Assert.assertFalse(columns.contains("mcol1"), columns.toString());
+ Assert.assertFalse(columns.contains("mcol2"), columns.toString());
+ Assert.assertTrue(columns.contains("col1"), columns.toString());
+ Assert.assertEquals(unnestNode.getPassthroughInputIndexes(), List.of(0));
+ Assert.assertEquals(unnestNode.getElementIndexes(), List.of(1, 2));
+ Assert.assertEquals(unnestNode.getDataSchema().size(), 3);
+ }
+
+ @Test
+ public void testUnnestColumnPruningViaBrokerDefault() {
+ // With the broker-config default on (and no per-query SET), pruning still
applies.
+ QueryEnvironment env = buildQueryEnvironment(true);
+ DispatchableSubPlan subPlan =
+ env.planQuery("SELECT e.col1, u.s FROM e CROSS JOIN UNNEST(e.mcol1) AS
u(s)");
+ List<UnnestNode> unnestNodes = findUnnestNodes(subPlan);
+ Assert.assertEquals(unnestNodes.size(), 1);
+ Assert.assertTrue(unnestNodes.get(0).isPrunedPassthrough(),
+ "Broker-config default should enable pruning without a SET option");
+
+ // A per-query SET overrides the broker default back off.
+ DispatchableSubPlan overridden = env.planQuery(
+ "SET unnestColumnPruning=false; SELECT e.col1, u.s FROM e CROSS JOIN
UNNEST(e.mcol1) AS u(s)");
+
Assert.assertFalse(findUnnestNodes(overridden).get(0).isPrunedPassthrough(),
+ "Per-query SET should override the broker-config default");
+ }
+
+ private static QueryEnvironment buildQueryEnvironment(boolean
defaultUnnestColumnPruning) {
+ MockRoutingManagerFactory factory = new MockRoutingManagerFactory(1, 2);
+ TABLE_SCHEMAS.forEach((name, schema) -> factory.registerTable(schema,
name));
+ SERVER1_SEGMENTS.forEach((table, segments) -> segments.forEach(s ->
factory.registerSegment(1, table, s)));
+ SERVER2_SEGMENTS.forEach((table, segments) -> segments.forEach(s ->
factory.registerSegment(2, table, s)));
+ RoutingManager routingManager = factory.buildRoutingManager(null);
+ TableCache tableCache = factory.buildTableCache();
+ return new QueryEnvironment(QueryEnvironment.configBuilder()
+ .requestId(-1L)
+ .database(CommonConstants.DEFAULT_DATABASE)
+ .tableCache(tableCache)
+ .workerManager(new WorkerManager("Broker_localhost", "localhost", 3,
routingManager))
+ .defaultUnnestColumnPruning(defaultUnnestColumnPruning)
+ .build());
+ }
+
+ @Test
+ public void testUnnestColumnPruningToZeroPassthrough() {
+ // Selecting only the unnested element retains zero passthrough columns.
+ String sql = "SET unnestColumnPruning=true; "
+ + "SELECT u.s FROM e CROSS JOIN UNNEST(e.mcol1) AS u(s)";
+ DispatchableSubPlan subPlan = _queryEnvironment.planQuery(sql);
+ List<UnnestNode> unnestNodes = findUnnestNodes(subPlan);
+ Assert.assertEquals(unnestNodes.size(), 1);
+ UnnestNode unnestNode = unnestNodes.get(0);
+
+ Assert.assertTrue(unnestNode.isPrunedPassthrough());
+ Assert.assertTrue(unnestNode.getPassthroughInputIndexes().isEmpty());
+ Assert.assertEquals(unnestNode.getElementIndexes(), List.of(0));
+ Assert.assertEquals(unnestNode.getDataSchema().size(), 1);
+ }
+
@Test
public void testAggregateWithOrdinality() {
String sql = "SELECT SUM(w.ord) FROM e CROSS JOIN UNNEST(e.mcol1) WITH
ORDINALITY AS w(s, ord)";
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnnestOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnnestOperator.java
index 0594e82a72c..367a565cac9 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnnestOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/UnnestOperator.java
@@ -50,6 +50,8 @@ public class UnnestOperator extends MultiStageOperator {
private final boolean _withOrdinality;
private final List<Integer> _elementIndexes;
private final int _ordinalityIndex;
+ private final int[] _passthroughInputIndexes;
+ private final boolean _prunedPassthrough;
private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
private boolean _loggedElementOverflow;
@@ -66,6 +68,13 @@ public class UnnestOperator extends MultiStageOperator {
_withOrdinality = node.isWithOrdinality();
_elementIndexes = node.getElementIndexes();
_ordinalityIndex = node.getOrdinalityIndex();
+ // Resolve to a primitive array once so the per-row hot path avoids
List.get + unboxing.
+ List<Integer> passthroughInputIndexes = node.getPassthroughInputIndexes();
+ _passthroughInputIndexes = new int[passthroughInputIndexes.size()];
+ for (int i = 0; i < _passthroughInputIndexes.length; i++) {
+ _passthroughInputIndexes[i] = passthroughInputIndexes.get(i);
+ }
+ _prunedPassthrough = node.isPrunedPassthrough();
}
@Override
@@ -226,10 +235,21 @@ public class UnnestOperator extends MultiStageOperator {
private Object[] appendElements(Object[] inputRow, List<Object> elements,
int ordinality) {
int outSize = _resultSchema.size();
Object[] out = new Object[outSize];
- // Copy left columns at beginning
- System.arraycopy(inputRow, 0, out, 0, inputRow.length);
- // Determine positions for elements. Track next free slot after the copied
input row.
- int base = inputRow.length;
+ // Copy passthrough (left) columns at the beginning of the output row.
+ int base;
+ if (_prunedPassthrough) {
+ // Only the input columns referenced downstream are retained, in output
order.
+ int numPassthrough = _passthroughInputIndexes.length;
+ for (int o = 0; o < numPassthrough; o++) {
+ out[o] = inputRow[_passthroughInputIndexes[o]];
+ }
+ base = numPassthrough;
+ } else {
+ // Legacy behavior: copy the whole input row (including the unnested
source array).
+ System.arraycopy(inputRow, 0, out, 0, inputRow.length);
+ base = inputRow.length;
+ }
+ // Determine positions for elements. Track next free slot after the copied
passthrough columns.
int nextFreePos = base;
for (int i = 0; i < elements.size(); i++) {
int elemPos = -1;
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/UnnestOperatorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/UnnestOperatorTest.java
index 124e7f0cb42..b3d32ffc397 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/UnnestOperatorTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/UnnestOperatorTest.java
@@ -165,6 +165,119 @@ public class UnnestOperatorTest {
assertEquals(rows.get(1)[3], "y");
}
+ @Test
+ public void shouldPrunePassthroughColumnsWhenSourceArrayDropped() {
+ // Input keeps the source array, but the pruned output schema drops it:
only "id" passes through.
+ DataSchema inputSchema = new DataSchema(new String[]{"id", "arr"}, new
ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.INT_ARRAY
+ });
+ when(_input.nextBlock()).thenReturn(
+ OperatorTestUtil.block(inputSchema,
+ new Object[]{1, new int[]{10, 20}},
+ new Object[]{2, new int[]{30}}));
+
+ DataSchema resultSchema = new DataSchema(new String[]{"id", "elem"}, new
ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.INT
+ });
+ RexExpression arrayExpr = new RexExpression.InputRef(1);
+ // Passthrough only input column 0 ("id") to output 0; element lands at
output 1. prunedPassthrough = true.
+ UnnestNode.TableFunctionContext context =
+ new UnnestNode.TableFunctionContext(false, List.of(1),
UnnestNode.UNSPECIFIED_INDEX, List.of(0), true);
+ UnnestNode node = new UnnestNode(-1, resultSchema,
PlanNode.NodeHint.EMPTY, List.of(), List.of(arrayExpr), context);
+ UnnestOperator operator = new
UnnestOperator(OperatorTestUtil.getTracingContext(), _input, inputSchema, node);
+
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 3);
+ for (Object[] row : rows) {
+ assertEquals(row.length, 2); // source array is not carried
+ }
+ assertEquals(rows.get(0)[0], 1);
+ assertEquals(rows.get(0)[1], 10);
+ assertEquals(rows.get(1)[0], 1);
+ assertEquals(rows.get(1)[1], 20);
+ assertEquals(rows.get(2)[0], 2);
+ assertEquals(rows.get(2)[1], 30);
+ }
+
+ @Test
+ public void shouldPrunePassthroughToZeroColumns() {
+ // SELECT only the unnested element: no passthrough columns retained.
+ DataSchema inputSchema = new DataSchema(new String[]{"id", "arr"}, new
ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.INT_ARRAY
+ });
+ when(_input.nextBlock()).thenReturn(
+ OperatorTestUtil.block(inputSchema, new Object[]{1, new int[]{10,
20}}));
+
+ DataSchema resultSchema = new DataSchema(new String[]{"elem"}, new
ColumnDataType[]{ColumnDataType.INT});
+ RexExpression arrayExpr = new RexExpression.InputRef(1);
+ UnnestNode.TableFunctionContext context =
+ new UnnestNode.TableFunctionContext(false, List.of(0),
UnnestNode.UNSPECIFIED_INDEX, List.of(), true);
+ UnnestNode node = new UnnestNode(-1, resultSchema,
PlanNode.NodeHint.EMPTY, List.of(), List.of(arrayExpr), context);
+ UnnestOperator operator = new
UnnestOperator(OperatorTestUtil.getTracingContext(), _input, inputSchema, node);
+
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 2);
+ assertEquals(rows.get(0).length, 1);
+ assertEquals(rows.get(0)[0], 10);
+ assertEquals(rows.get(1)[0], 20);
+ }
+
+ @Test
+ public void shouldPrunePassthroughWithOrdinality() {
+ DataSchema inputSchema = new DataSchema(new String[]{"id", "arr"}, new
ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.INT_ARRAY
+ });
+ when(_input.nextBlock()).thenReturn(
+ OperatorTestUtil.block(inputSchema, new Object[]{1, new int[]{5, 6}}));
+
+ // Pruned output [id, elem, ord]; source array dropped. element at 1,
ordinality at 2.
+ DataSchema resultSchema = new DataSchema(new String[]{"id", "elem",
"ord"}, new ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.INT, ColumnDataType.INT
+ });
+ UnnestNode.TableFunctionContext context =
+ new UnnestNode.TableFunctionContext(true, List.of(1), 2, List.of(0),
true);
+ UnnestNode node = new UnnestNode(-1, resultSchema,
PlanNode.NodeHint.EMPTY, List.of(),
+ List.of(new RexExpression.InputRef(1)), context);
+ UnnestOperator operator = new
UnnestOperator(OperatorTestUtil.getTracingContext(), _input, inputSchema, node);
+
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 2);
+ assertEquals(rows.get(0).length, 3);
+ assertEquals(rows.get(0)[0], 1);
+ assertEquals(rows.get(0)[1], 5);
+ assertEquals(rows.get(0)[2], 1);
+ assertEquals(rows.get(1)[1], 6);
+ assertEquals(rows.get(1)[2], 2);
+ }
+
+ @Test
+ public void shouldPrunePassthroughWithMultipleArrays() {
+ DataSchema inputSchema = new DataSchema(new String[]{"id", "a1", "a2"},
new ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.INT_ARRAY,
ColumnDataType.STRING_ARRAY
+ });
+ when(_input.nextBlock()).thenReturn(
+ OperatorTestUtil.block(inputSchema, new Object[]{10, new int[]{1, 2},
new String[]{"x", "y"}}));
+
+ // Pruned output [id, v1, v2]; both source arrays dropped.
+ DataSchema resultSchema = new DataSchema(new String[]{"id", "v1", "v2"},
new ColumnDataType[]{
+ ColumnDataType.INT, ColumnDataType.INT, ColumnDataType.STRING
+ });
+ UnnestNode.TableFunctionContext context =
+ new UnnestNode.TableFunctionContext(false, List.of(1, 2),
UnnestNode.UNSPECIFIED_INDEX, List.of(0), true);
+ UnnestNode node = new UnnestNode(-1, resultSchema,
PlanNode.NodeHint.EMPTY, List.of(),
+ List.of(new RexExpression.InputRef(1), new RexExpression.InputRef(2)),
context);
+ UnnestOperator operator = new
UnnestOperator(OperatorTestUtil.getTracingContext(), _input, inputSchema, node);
+
+ List<Object[]> rows = ((MseBlock.Data)
operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(rows.size(), 2);
+ assertEquals(rows.get(0).length, 3);
+ assertEquals(rows.get(0)[0], 10);
+ assertEquals(rows.get(0)[1], 1);
+ assertEquals(rows.get(0)[2], "x");
+ assertEquals(rows.get(1)[1], 2);
+ assertEquals(rows.get(1)[2], "y");
+ }
+
@Test
public void shouldZipMultipleArraysIntoColumns() {
DataSchema inputSchema = new DataSchema(new String[]{"id", "a1", "a2"},
new ColumnDataType[]{
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 33f94493840..b81f4fb2828 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -552,6 +552,15 @@ public class CommonConstants {
public static final String CONFIG_OF_SPOOLS =
"pinot.broker.multistage.spools";
public static final boolean DEFAULT_OF_SPOOLS = false;
+ /**
+ * Whether the multistage query engine prunes unused input (passthrough)
columns - notably the unnested source
+ * array - from the UNNEST output by default. This value can always be
overridden by the
+ * {@link Request.QueryOptionKey#UNNEST_COLUMN_PRUNING} query option. Keep
it disabled until all servers support
+ * it (a broker emitting the smaller UNNEST schema cannot be honored by an
un-upgraded server).
+ */
+ public static final String CONFIG_OF_UNNEST_COLUMN_PRUNING =
"pinot.broker.multistage.unnest.column.pruning";
+ public static final boolean DEFAULT_UNNEST_COLUMN_PRUNING = false;
+
/// Whether to only use servers for leaf stages as the workers for the
intermediate stages.
/// This value can always be overridden by
[Request.QueryOptionKey#USE_LEAF_SERVER_FOR_INTERMEDIATE_STAGE].
public static final String
CONFIG_OF_USE_LEAF_SERVER_FOR_INTERMEDIATE_STAGE =
@@ -812,6 +821,12 @@ public class CommonConstants {
public static final String APPLICATION_NAME = "applicationName";
public static final String USE_SPOOLS = "useSpools";
public static final String USE_PHYSICAL_OPTIMIZER =
"usePhysicalOptimizer";
+ // When true, the multi-stage planner prunes input (passthrough)
columns - notably the unnested source array -
+ // from the UNNEST output when they are not referenced downstream,
avoiding copying them into every exploded
+ // row. Defaults to false: enabling it makes the broker emit a smaller
UNNEST output schema, which an
+ // un-upgraded server cannot honor, so only enable it once all servers
support it.
+ // NOTE: This is a no-op under usePhysicalOptimizer (the v2 path does
not go through RelToPlanNodeConverter).
+ public static final String UNNEST_COLUMN_PRUNING =
"unnestColumnPruning";
/**
* When set to true, the broker uses the long-lived {@code
SubmitWithStream} bidi RPC to dispatch the query,
* receiving stage stats out-of-band as {@code OpChainComplete}
messages instead of via mailbox EOS. The
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
