Fix missing UNNEST input columns in multi-stage EXPLAIN when asking servers

[yashmayya]

PR flow

Fix for CROSS JOIN UNNEST EXPLAIN: build Correlate plan when input columns are passed through.

flowchart TD
  N0["#91;F1#93; visitUnnest#58; check numPassthrough#62;0 #40;added#41; #40;F1#41;"]:::stAdded
  N1["#91;F1#93; pushCorrelatedUnnest#58; create correl variable and RexShuttle #40;added#41; #40;F1#41;"]:::stAdded
  N2["#91;F1#93; pushCorrelatedUnnest#58; build arrayNames and required columns #40;added#41; #40;F1#41;"]:::stAdded
  N3["#91;F1#93; pushCorrelatedUnnest#58; push Values#44; project arrayProjects #40;via shuttle#41; and#44; #40;F1#41;"]:::stAdded
  N4["#91;F1#93; pushCorrelatedUnnest#58; correlate input with uncollect #40;added#41; #40;F1#41;"]:::stAdded
  N5["#91;F1#93; pushCorrelatedUnnest#58; handle pruned passthrough via project #40;added#41; #40;F1#41;"]:::stAdded
  N0 -->|"call pushCorrelatedUncollect #40;when numPassthrough#62;0#41;"| N1
  N1 -->|"sequential"| N2
  N2 -->|"sequential"| N3
  N3 -->|"sequential"| N4
  N4 -->|"sequential"| N5
  classDef stAdded fill:#dafbe1,stroke:#1a7f37,color:#1f2328,stroke-width:2px
  classDef stModified fill:#fff8c5,stroke:#9a6700,color:#1f2328,stroke-width:2px
  classDef stRemoved fill:#ffebe9,stroke:#cf222e,color:#1f2328,stroke-width:2px
  classDef stUnchanged fill:#f6f8fa,stroke:#656d76,color:#1f2328,stroke-width:1px

Loading

AI-generated · Green: added · Yellow: modified · Red: removed · Gray: existing

Diff evidence

• F1: pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java — before · after

• Regenerate PR flow

The bug

With explainAskingServers (or the broker config pinot.query.multistage.explain.include.segment.plan), the broker gets the plan of each leaf stage from the servers. Then PlanNodeToRelConverter converts the plan nodes back to Calcite nodes.

visitUnnest converts an UnnestNode to an Uncollect over a Project of the arrays. The output of this Uncollect has only the element (and ordinality) columns. But for CROSS JOIN UNNEST, the UnnestNode output has the input columns first, then the element columns. So a parent node in the leaf stage cannot find its input fields, and EXPLAIN shows it as UnknownAggregate or UnknownProject.

For example:

SELECT Carrier, u.a, COUNT(*)
FROM mytable CROSS JOIN UNNEST(DivAirports) AS u(a)
GROUP BY Carrier, u.a

EXPLAIN shows the leaf stage as:

PinotLogicalExchange(distribution=[hash[0, 1]])
  UnknownAggregate
    Uncollect
      LogicalProject(DivAirports=[$1])
        LeafStageCombineOperator(table=[mytable])

The broker logs IllegalArgumentException: field ordinal [2] out of range; input fields are: [DivAirports]. These queries have the same problem: GROUP BY only the element column, WITH ORDINALITY, ROLLUP, SELECT Carrier, u.a ... LIMIT 10, and queries with unnestColumnPruning=true.

The fix

If the UnnestNode output has input columns, visitUnnest now builds the same shape as the logical plan. This shape is a LogicalCorrelate of the input with an Uncollect that reads the arrays through the correlation variable. If the planner pruned some input columns (unnestColumnPruning), a Project keeps only the remaining input columns. A plain UNNEST without input columns does not change.

Only the EXPLAIN output changes. Query execution does not use this conversion.

Testing

• PlanNodeToRelConverterTest.testAggregateOverUnnest converts an aggregate over CROSS JOIN UNNEST. testProjectOverPrunedUnnest converts a project over a pruned UNNEST WITH ORDINALITY. Without this change, they fail with UnknownAggregate and UnknownProject.
• I ran EXPLAIN for 7 such queries in a temporary integration test with the broker config above. Before this change, all of them showed an Unknown node. After it, none do, and the leaf stages match the logical plan shape. For the query above:

PinotLogicalExchange(distribution=[hash[0, 1]])
  LogicalAggregate(group=[{0, 2}], agg#0=[COUNT()])
    LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{1}])
      LeafStageCombineOperator(table=[mytable])
      Uncollect
        LogicalProject(DivAirports=[$cor0.DivAirports])
          LogicalValues(tuples=[[{ 0 }]])

• All 2662 pinot-query-planner tests pass. spotless, checkstyle and license are clean.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.59259% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.46%. Comparing base (73645da) to head (f580422).
⚠️ Report is 5 commits behind head on master.

Files with missing lines	Patch %	Lines
.../query/planner/logical/PlanNodeToRelConverter.java	92.59%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19771      +/-   ##
============================================
+ Coverage     68.40%   68.46%   +0.06%     
  Complexity     1470     1470              
============================================
  Files          3521     3522       +1     
  Lines        229091   229204     +113     
  Branches      36346    36377      +31     
============================================
+ Hits         156707   156924     +217     
+ Misses        60104    60001     -103     
+ Partials      12280    12279       -1     

Flag	Coverage Δ
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	68.46% <92.59%> (+0.06%)	⬆️
lane-a	100.00% <ø> (ø)
lane-b	0.00% <ø> (ø)
temurin	68.46% <92.59%> (+0.06%)	⬆️
unittests	68.46% <92.59%> (+0.06%)	⬆️
unittests1	58.29% <92.59%> (+0.08%)	⬆️
unittests2	40.03% <0.00%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.