You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ya...@apache.org on 2026/10/07 15:35:25 UTC

(pinot) branch master updated: Fix missing UNNEST input columns in multi-stage EXPLAIN when asking servers (#19771)

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 <ya...@gmail.com>
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: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org