Skip to content

Support GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage query engine#18817

Open
xiangfu0 wants to merge 2 commits into
apache:masterfrom
xiangfu0:mse-grouping-sets
Open

Support GROUP BY GROUPING SETS / ROLLUP / CUBE in the multi-stage query engine#18817
xiangfu0 wants to merge 2 commits into
apache:masterfrom
xiangfu0:mse-grouping-sets

Conversation

@xiangfu0

@xiangfu0 xiangfu0 commented Jun 20, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds native execution of GROUP BY GROUPING SETS (...) / ROLLUP(...) / CUBE(...) and the
GROUPING() / GROUPING_ID() functions to the multi-stage query engine (MSE), building on the
single-stage support in #18662.

Two execution paths, both producing identical results:

  • Aggregate directly over a table scan — the whole aggregate is pushed down to the single-stage (leaf)
    engine
    , which expands each row across the grouping sets and appends the synthetic $groupingId discriminator.
    Reuses the single-stage engine wholesale (the Doris/StarRocks backend-expansion model).
  • Aggregate over any other input (e.g. after a JOIN) — a new RepeatOperator in the multi-stage runtime
    performs the same per-set row expansion (NULLing rolled-up columns, appending $groupingId), then an ordinary
    GROUP BY over [union keys…, $groupingId] runs — no grouping-set-specific aggregation logic.

Plan split (PinotAggregateExchangeNodeInsertRule): LEAF → EXCHANGE(hash union+$groupingId) → FINAL(group by union+$groupingId) → PROJECT. The PROJECT computes GROUPING() / GROUPING_ID() from
$groupingId (bit extraction), mirroring the single-stage post-aggregation handler, and drops $groupingId.
GROUPING / GROUPING_ID are registered in PinotOperatorTable, and explicit
GROUP BY GROUPING SETS ((a, b), …) tuple syntax is accepted by the validator.

Tests

GroupingSetsQueriesTest — grouping-set aggregation, GROUPING() / GROUPING_ID() (SELECT and HAVING),
filtered aggregation, composite/nested sets, explicit GROUPING SETS syntax, and grouping sets over a JOIN
run against both engines (the order-independent cases via the useBothQueryEngines data provider), plus the
single-stage genuine-vs-rolled-up-NULL discrimination cases. All green.

Known limitations (rejected with a clear error → never wrong results)

  • The v2 physical planner (usePhysicalOptimizer=true, opt-in, default off) still rejects grouping sets; the
    default planner is fully supported. (Follow-up.)
  • More than 31 distinct grouping columns (the $groupingId bitmask is a 32-bit int) — same cap as the
    single-stage engine.
  • Leaf group-trim (ORDER BY … LIMIT pushdown) is not pushed for grouping-set queries; results are still correct
    (the broker applies the final ORDER BY + LIMIT).

@xiangfu0 xiangfu0 marked this pull request as ready for review June 20, 2026 10:08
@codecov-commenter

codecov-commenter commented Jun 20, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 66.83805% with 129 lines in your changes missing coverage. Please review.
✅ Project coverage is 64.82%. Comparing base (9ceca0f) to head (ffaebbf).
⚠️ Report is 4 commits behind head on master.

Files with missing lines Patch % Lines
...r/physical/v2/opt/rules/AggregatePushdownRule.java 9.37% 56 Missing and 2 partials ⚠️
...he/pinot/query/runtime/plan/PlanNodeToOpChain.java 0.00% 28 Missing and 1 partial ⚠️
...org/apache/pinot/core/data/table/TableResizer.java 12.50% 7 Missing ⚠️
...inot/core/query/reduce/PostAggregationHandler.java 12.50% 7 Missing ⚠️
...e/pinot/query/runtime/operator/RepeatOperator.java 89.47% 6 Missing ⚠️
...egation/groupby/GroupingSetsGroupKeyGenerator.java 0.00% 4 Missing ⚠️
...el/rules/PinotAggregateExchangeNodeInsertRule.java 93.18% 0 Missing and 3 partials ⚠️
...y/planner/physical/v2/nodes/PhysicalAggregate.java 40.00% 1 Missing and 2 partials ⚠️
...not/query/runtime/operator/MultiStageOperator.java 25.00% 3 Missing ⚠️
...org/apache/pinot/sql/parsers/CalciteSqlParser.java 93.54% 0 Missing and 2 partials ⚠️
... and 4 more
Additional details and impacted files
@@            Coverage Diff             @@
##             master   #18817    +/-   ##
==========================================
  Coverage     64.81%   64.82%            
  Complexity     1347     1347            
==========================================
  Files          3396     3398     +2     
  Lines        212503   212803   +300     
  Branches      33484    33534    +50     
==========================================
+ Hits         137732   137943   +211     
- Misses        63610    63679    +69     
- Partials      11161    11181    +20     
Flag Coverage Δ
custom-integration1 100.00% <ø> (ø)
integration 100.00% <ø> (ø)
integration1 100.00% <ø> (ø)
integration2 0.00% <ø> (ø)
java-21 64.82% <66.83%> (+<0.01%) ⬆️
temurin 64.82% <66.83%> (+<0.01%) ⬆️
unittests 64.81% <66.83%> (+<0.01%) ⬆️
unittests1 56.81% <66.83%> (+0.02%) ⬆️
unittests2 37.17% <5.91%> (-0.06%) ⬇️

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.

@xiangfu0 xiangfu0 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Found one high-signal compatibility issue; see inline comment.

@xiangfu0 xiangfu0 force-pushed the mse-grouping-sets branch from 513b62b to 616c781 Compare June 21, 2026 02:39

@xiangfu0 xiangfu0 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Found one high-signal issue in the new v2 grouping-sets path; see inline comment.

@gortiz

gortiz commented Jun 22, 2026

Copy link
Copy Markdown
Contributor

Nice contribution!

@gortiz gortiz requested review from Jackie-Jiang, gortiz and yashmayya and removed request for Jackie-Jiang June 22, 2026 14:00
@yashmayya yashmayya added multi-stage Related to the multi-stage query engine feature New functionality labels Jun 22, 2026
@xiangfu0 xiangfu0 force-pushed the mse-grouping-sets branch 10 times, most recently from b957437 to 600145c Compare June 30, 2026 12:10
@xiangfu0 xiangfu0 force-pushed the mse-grouping-sets branch from e4529c1 to 81145e9 Compare July 3, 2026 09:33
@xiangfu0 xiangfu0 requested review from Jackie-Jiang July 5, 2026 19:02
…lti-stage query engine

Adds native execution of GROUP BY GROUPING SETS (...) / ROLLUP(...) / CUBE(...) and the
GROUPING() / GROUPING_ID() functions to the multi-stage query engine (MSE), in both the
default planner and the v2 physical planner (usePhysicalOptimizer).

A grouping-set query is represented as the union of all grouping columns plus, per set, a
synthetic $groupingId bitmask (bit i set iff union column i is rolled up). The bitmask is
carried as an internal key column that keeps the sets distinct and powers GROUPING() /
GROUPING_ID(); the shared conventions live in GroupingSets (pinot-common).

Execution paths (identical results on both):
- Aggregate directly over a table scan: the whole aggregate is pushed to the single-stage
  (leaf) engine, which expands each row across the sets and appends $groupingId.
- Aggregate over any other input (e.g. after a JOIN): a new RepeatOperator in the MSE
  runtime performs the same per-set row expansion, then an ordinary GROUP BY over
  [union keys..., $groupingId] runs -- no grouping-set-specific aggregation logic.

Planner split (both planners): LEAF -> EXCHANGE(hash union + $groupingId) ->
FINAL(group by union + $groupingId) -> PROJECT. The PROJECT computes GROUPING() /
GROUPING_ID() from $groupingId (shared GroupingSetsRexUtils) and drops $groupingId.

Guards (explicit failure, never wrong results): at most 31 distinct grouping columns (the
$groupingId bitmask is a 32-bit int); WITHIN GROUP ordered aggregates under grouping sets
are rejected in the v2 planner (run on single-stage).

Tests: GroupingSetsQueriesTest exercises grouping-set aggregation, GROUPING() / GROUPING_ID()
(SELECT and HAVING), filtered aggregation, composite/nested sets, explicit GROUPING SETS
syntax, grouping sets over a JOIN, and the v2 physical planner -- against both engines via
data providers. PlanNodeSerDeTest and RepeatOperatorTest cover the wire format and the
runtime row expansion.

Rolling upgrade: the feature ships within a single release, so brokers and servers are
upgraded together; AggregateNode.groupingSets is additive (proto field 8).
@xiangfu0 xiangfu0 force-pushed the mse-grouping-sets branch from 81145e9 to d1c5838 Compare July 7, 2026 00:56
…ving the 31-column limit

Encode grouping sets the way Calcite does: each set is the list of participating grouping-column
indexes (an unbounded per-set membership list) rather than a 32-bit membership bitmask, so the number
of grouping columns is no longer capped at 31. A set's position in the plan's set list is its ordinal,
carried as the synthetic $groupingId discriminator through the leaf/exchange/final split and the wire.

Also addresses review feedback:
- RepeatOperator now appends NULLable per-set group-key copies and leaves the original input columns
  untouched, so aggregation arguments that reference a grouping column (e.g. SUM(b) under ROLLUP(a, b))
  aggregate the real values; expansion is streamed one set per block with a termination/usage check.
- The per-set expansion is wired in PlanNodeToOpChain (not inside a specific AggregateOperator), so every
  AggregateOperatorFactory receives already-expanded input and needs no grouping-set awareness.
- The GROUPING()/GROUPING_ID() split, projection, LEAF row type, and set encoding are shared between the
  default and v2 physical planners via GroupingSetsPlanUtils (replaces GroupingSetsRexUtils).
- Both planners reject combinations that would otherwise produce broken plans or silently wrong results:
  WITHIN GROUP / skip-leaf / partitioned-by hints with grouping sets, aggregation-free grouping sets,
  GROUPING() over a plain GROUP BY, and (v2) grouping sets over an input with no repartitioning exchange.
  The v2 grouping-set LEAF no longer forwards group-trim collations/limit indexed against the pre-$groupingId
  layout.
- The single-stage engine treats GROUP BY GROUPING SETS (()) (grand-total-only, no grouping columns) as a
  plain aggregation instead of crashing the reducer.
- Wire: plan.proto grouping sets field renumbered to 8; query.thrift/plan.proto rolling-upgrade notes reworded.
@xiangfu0 xiangfu0 force-pushed the mse-grouping-sets branch from d1c5838 to ffaebbf Compare July 7, 2026 06:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature New functionality multi-stage Related to the multi-stage query engine

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants