Fix missing GROUPING SETS in multi-stage EXPLAIN when asking servers

[yashmayya]

PR flow

Fixes missing GROUPING SETS in multi-stage EXPLAIN by updating PlanNodeToRelConverter to handle grouping sets. Supported by F1 changes in visitAggregate method.

flowchart TD
  N0["Get group keys #40;list#41; #40;F1#41;"]:::stAdded
  N1["Check has grouping sets#63; #40;F1#41;"]:::stAdded
  N2["Build group key with grouping sets #40;F1#41;"]:::stAdded
  N3["Build group key #40;plain#41; #40;F1#41;"]:::stAdded
  N0 -->|"next"| N1
  N1 -->|"if true"| N2
  N1 -->|"if false"| N3
  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

This PR depends on #19767, which depends on #19766. It contains their commits. Only the last commit is new.

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.

visitAggregate builds the group key only from AggregateNode.getGroupKeys() and ignores getGroupingSets(). So a GROUPING SETS / ROLLUP / CUBE aggregate that runs as a multi-stage operator in a leaf stage shows as a plain GROUP BY.

For example, with is_partitioned_by_group_by_keys, the ROLLUP aggregate of this query runs in the leaf stage, above the inner aggregate:

SELECT Carrier, Origin, MAX(v)
FROM (
  SELECT /*+ aggOptions(is_partitioned_by_group_by_keys='true') */ Carrier, Origin, Dest, AVG(DepDelay) AS v
  FROM mytable GROUP BY Carrier, Origin, Dest)
GROUP BY ROLLUP(Carrier, Origin)

EXPLAIN PLAN WITHOUT IMPLEMENTATION shows this aggregate as PinotLogicalAggregate(group=[{0, 1}], groups=[[{0, 1}, {0}, {}]], agg#0=[MAX($2)], aggType=[LEAF]). EXPLAIN with the plan from the servers shows the leaf stage as:

PinotLogicalExchange(distribution=[hash[0, 1, 2]])
  LogicalAggregate(group=[{0, 1}], agg#0=[MAX($2)])
    LogicalProject(Carrier=[$0], Origin=[$2], $f2=[DIVIDE($3, $4)])
      LeafStageCombineOperator(table=[mytable])

The fix

In visitAggregate, if the node has grouping sets, build the group key with RelBuilder.groupKey(ImmutableBitSet, Iterable<ImmutableBitSet>). Each grouping set holds indexes into the group keys, so the converter maps them to input fields first. Plain GROUP BY does not change.

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

Testing

• PlanNodeToRelConverterTest.testAggregateGroupingSets converts a ROLLUP aggregate whose group keys are not the first input fields. Without this change, it fails, because the plan has no groups.
• I ran EXPLAIN in a temporary integration test with the broker config above. Before this change, the leaf stage showed a plain GROUP BY. After it, the leaf stage shows the grouping sets:

PinotLogicalExchange(distribution=[hash[0, 1, 2]])
  LogicalAggregate(group=[{0, 1}], groups=[[{0, 1}, {0}, {}]], agg#0=[MAX($2)])
    LogicalProject(Carrier=[$0], Origin=[$2], $f2=[DIVIDE($3, $4)])
      LeafStageCombineOperator(table=[mytable])

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

Not in this PR

• If the single-stage engine runs the ROLLUP aggregate (for example, ROLLUP directly on a table), the server plan shows GroupBy(groupKeys=[[Carrier, Origin]], ...). GroupByOperator does not print the grouping sets.
• An aggregate above UNNEST in a leaf stage shows as UnknownAggregate, with or without grouping sets. visitUnnest keeps only the unnested columns, so the aggregate cannot find its input fields.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 68.40%. Comparing base (73645da) to head (d050380).

Additional details and impacted files

@@            Coverage Diff            @@
##             master   #19769   +/-   ##
=========================================
  Coverage     68.40%   68.40%           
  Complexity     1470     1470           
=========================================
  Files          3521     3521           
  Lines        229091   229098    +7     
  Branches      36346    36348    +2     
=========================================
+ Hits         156707   156721   +14     
+ Misses        60104    60092   -12     
- Partials      12280    12285    +5     

Flag	Coverage Δ
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (?)
java-25	68.40% <100.00%> (+<0.01%)	⬆️
lane-a	100.00% <ø> (ø)
lane-b	0.00% <ø> (ø)
temurin	68.40% <100.00%> (+<0.01%)	⬆️
unittests	68.40% <100.00%> (+<0.01%)	⬆️
unittests1	58.20% <100.00%> (+<0.01%)	⬆️
unittests2	40.06% <0.00%> (+<0.01%)	⬆️

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.