This is an automated email from the ASF dual-hosted git repository.
yashmayya 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 5b6c7a2983c Fix missing UNNEST input columns in multi-stage EXPLAIN
when asking servers (#19771)
5b6c7a2983c is described below
commit 5b6c7a2983c305c486901a9e1010a18608496d66
Author: Yash Mayya <[email protected]>
AuthorDate: Wed Oct 7 08:35:17 2026 -0700
Fix missing UNNEST input columns in multi-stage EXPLAIN when asking servers
(#19771)
---
.../planner/logical/PlanNodeToRelConverter.java | 47 +++++++++++++++++++
.../logical/PlanNodeToRelConverterTest.java | 52 ++++++++++++++++++++++
2 files changed, 99 insertions(+)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java
index 2ff2e9289bb..5c7dd4346ae 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java
@@ -25,6 +25,7 @@ import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelCollation;
import org.apache.calcite.rel.RelCollations;
@@ -39,10 +40,15 @@ import org.apache.calcite.rel.logical.LogicalIntersect;
import org.apache.calcite.rel.logical.LogicalMinus;
import org.apache.calcite.rel.logical.LogicalSort;
import org.apache.calcite.rel.logical.LogicalUnion;
+import org.apache.calcite.rel.logical.LogicalValues;
import org.apache.calcite.rel.logical.LogicalWindow;
import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexShuttle;
import org.apache.calcite.rex.RexWindowBound;
import org.apache.calcite.rex.RexWindowBounds;
import org.apache.calcite.rex.RexWindowExclusion;
@@ -502,6 +508,13 @@ public final class PlanNodeToRelConverter {
if (arrayProjects.isEmpty()) {
arrayProjects.add(_builder.field(0));
}
+ // CROSS JOIN UNNEST also outputs the input columns, which parent
nodes may reference
+ int numPassthrough =
+ node.getDataSchema().size() - arrayProjects.size() -
(node.isWithOrdinality() ? 1 : 0);
+ if (numPassthrough > 0) {
+ pushCorrelatedUncollect(node, input, arrayProjects);
+ return null;
+ }
_builder.project(arrayProjects);
RelNode project = _builder.build();
@@ -517,6 +530,40 @@ public final class PlanNodeToRelConverter {
return null;
}
+ /// Replaces the input on top of the builder stack with the logical plan
shape of CROSS JOIN UNNEST: a Correlate of
+ /// the input with an Uncollect that reads the arrays through the
correlation variable. Like the [UnnestNode], it
+ /// outputs the input columns followed by the element (and ordinality)
columns.
+ private void pushCorrelatedUncollect(UnnestNode node, RelNode input,
List<RexNode> arrayProjects) {
+ RexBuilder rexBuilder = _builder.getRexBuilder();
+ RexCorrelVariable correlVariable =
+ (RexCorrelVariable) rexBuilder.makeCorrel(input.getRowType(),
_builder.getCluster().createCorrel());
+ RexShuttle toCorrelated = new RexShuttle() {
+ @Override
+ public RexNode visitInputRef(RexInputRef inputRef) {
+ return rexBuilder.makeFieldAccess(correlVariable,
inputRef.getIndex());
+ }
+ };
+ List<String> inputNames = input.getRowType().getFieldNames();
+ List<String> arrayNames = new ArrayList<>(arrayProjects.size());
+ for (RexNode arrayProject : arrayProjects) {
+ arrayNames.add(
+ arrayProject instanceof RexInputRef ?
inputNames.get(((RexInputRef) arrayProject).getIndex()) : null);
+ }
+ List<RexNode> requiredColumns =
_builder.fields(RelOptUtil.InputFinder.bits(arrayProjects, null));
+ _builder.push(LogicalValues.createOneRow(_builder.getCluster()))
+ .project(toCorrelated.apply(arrayProjects), arrayNames)
+ .uncollect(List.of(), node.isWithOrdinality())
+ .correlate(JoinRelType.INNER, correlVariable.id, requiredColumns);
+
+ if (node.isPrunedPassthrough()) {
+ // Keep only the passthrough input columns, followed by the element
(and ordinality) columns
+ List<RexNode> fields = _builder.fields();
+ List<RexNode> projects = new
ArrayList<>(_builder.fields(node.getPassthroughInputIndexes()));
+ projects.addAll(fields.subList(inputNames.size(), fields.size()));
+ _builder.project(projects);
+ }
+ }
+
@Override
public Void visitExchange(ExchangeNode exchangeNode, Void context) {
visitChildren(exchangeNode);
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverterTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverterTest.java
index 646d2ce4911..8f9f63285bc 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverterTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverterTest.java
@@ -30,6 +30,7 @@ import org.apache.pinot.query.planner.plannode.AggregateNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.planner.plannode.ProjectNode;
+import org.apache.pinot.query.planner.plannode.UnnestNode;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
@@ -81,4 +82,55 @@ public class PlanNodeToRelConverterTest {
"LogicalAggregate(group=[{1, 2}], groups=[[{1, 2}, {1}, {}]],
agg#0=[MAX($3)])\n"
+ " Leaf\n");
}
+
+ @Test
+ public void testAggregateOverUnnest() {
+ // COUNT(*) ... CROSS JOIN UNNEST(b) AS u(elem) GROUP BY a, elem, where
the unnest output is [a, b, elem]
+ PlanNode unnest = new UnnestNode(0, new DataSchema(new String[]{"a", "b",
"elem"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING_ARRAY,
ColumnDataType.STRING}), null,
+ List.of(unnestLeaf()), List.of(new RexExpression.InputRef(1)),
+ new UnnestNode.TableFunctionContext(false, List.of(2),
UnnestNode.UNSPECIFIED_INDEX));
+ DataSchema aggregateSchema = new DataSchema(new String[]{"a", "elem",
"cnt"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING,
ColumnDataType.LONG});
+ RexExpression.FunctionCall count = new
RexExpression.FunctionCall(ColumnDataType.LONG, "COUNT", List.of());
+ PlanNode aggregate = new AggregateNode(0, aggregateSchema, null,
List.of(unnest), List.of(count), List.of(-1),
+ List.of(0, 2), AggregateNode.AggType.LEAF, false, null, 0);
+
+ RelBuilder relBuilder =
RelBuilder.create(Frameworks.newConfigBuilder().build());
+ RelNode relNode = PlanNodeToRelConverter.convert(relBuilder, aggregate);
+ assertEquals(RelOptUtil.toString(relNode),
+ "LogicalAggregate(group=[{0, 2}], agg#0=[COUNT()])\n"
+ + " LogicalCorrelate(correlation=[$cor0], joinType=[inner],
requiredColumns=[{1}])\n"
+ + " Leaf\n"
+ + " Uncollect\n"
+ + " LogicalProject(b=[$cor0.b])\n"
+ + " LogicalValues(tuples=[[{ 0 }]])\n");
+ }
+
+ @Test
+ public void testProjectOverPrunedUnnest() {
+ // SELECT a, elem, ord ... CROSS JOIN UNNEST(b) WITH ORDINALITY AS u(elem,
ord), where the unnest output drops b
+ PlanNode unnest = new UnnestNode(0, new DataSchema(new String[]{"a",
"elem", "ord"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING,
ColumnDataType.INT}), null,
+ List.of(unnestLeaf()), List.of(new RexExpression.InputRef(1)),
+ new UnnestNode.TableFunctionContext(true, List.of(1), 2, List.of(0),
true));
+ PlanNode project = new ProjectNode(0, unnest.getDataSchema(), null,
List.of(unnest),
+ List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1),
new RexExpression.InputRef(2)));
+
+ RelBuilder relBuilder =
RelBuilder.create(Frameworks.newConfigBuilder().build());
+ RelNode relNode = PlanNodeToRelConverter.convert(relBuilder, project);
+ assertEquals(RelOptUtil.toString(relNode),
+ "LogicalProject(a=[$0], b=[$2], ORDINALITY=[$3])\n"
+ + " LogicalCorrelate(correlation=[$cor0], joinType=[inner],
requiredColumns=[{1}])\n"
+ + " Leaf\n"
+ + " Uncollect(withOrdinality=[true])\n"
+ + " LogicalProject(b=[$cor0.b])\n"
+ + " LogicalValues(tuples=[[{ 0 }]])\n");
+ }
+
+ private static PlanNode unnestLeaf() {
+ DataSchema leafSchema = new DataSchema(new String[]{"a", "b"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING_ARRAY});
+ return new ExplainedNode(0, leafSchema, null, List.of(), "Leaf", Map.of());
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
