Optimize GROUP BY GROUPING SETS/ROLLUP/CUBE via base aggregation - #19264
Optimize GROUP BY GROUPING SETS/ROLLUP/CUBE via base aggregation#19264xiangfu0 wants to merge 6 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #19264 +/- ##
=============================================
- Coverage 67.12% 39.08% -28.05%
+ Complexity 1424 1423 -1
=============================================
Files 3462 3463 +1
Lines 220677 220942 +265
Branches 35255 35310 +55
=============================================
- Hits 148136 86351 -61785
- Misses 60708 126719 +66011
+ Partials 11833 7872 -3961
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
xiangfu0
added
aggregation
There was a problem hiding this comment.
Pull request overview
This PR optimizes single-stage GROUP BY GROUPING SETS / ROLLUP / CUBE by aggregating the base (union) grouping once per segment using the standard fast GROUP BY path, then deriving per-grouping-set rows from those base groups (with defensive cloning of object intermediates to keep merges correct). It also improves the legacy expansion key generator for the fallback path and adds test/benchmark coverage for the new behavior.
Changes:
- Add a new query option (
groupingSetsBaseAggregation, enabled by default) and plumb executor/operator logic to use base aggregation + derivation where applicable. - Implement segment-level derivation of grouping-set rows from base groups (including per-set segment trimming parity with the legacy path).
- Improve the legacy expansion generator performance (native dict-id fast path + optional long-packed key storage) and add equivalence tests + a new JMH benchmark.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java | Adds the groupingSetsBaseAggregation query option key. |
| pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java | Adds isGroupingSetsBaseAggregation() decision logic (default-on, with filtered-agg carve-out). |
| pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java | Centralizes base-vs-expansion selection and exposes it via the executor. |
| pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByExecutor.java | Adds an executor flag to indicate base aggregation was used for grouping sets. |
| pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java | Switches grouping-set result building to the base-derivation path when indicated by the executor. |
| pinot-core/src/main/java/org/apache/pinot/core/util/GroupByUtils.java | Implements derivation from base groups with defensive cloning and per-set segment trimming. |
| pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java | Adds per-set bucketing trim for already-built IntermediateRecords (used by the derive path). |
| pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupingSetsGroupKeyGenerator.java | Improves legacy expansion key generation (native dict-id fast path + optional long packing). |
| pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GroupingSetsQueriesTest.java | Adds equivalence tests to ensure base aggregation matches legacy expansion across key scenarios. |
| pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkGroupingSetsQueriesSSE.java | Adds a JMH benchmark covering end-to-end SSE grouping sets and base-vs-expansion comparison. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // use 8 threads to test combine parallelism | ||
| private static final ExecutorService EXECUTOR_SERVICE = Executors.newFixedThreadPool(8); | ||
|
|
Jackie-Jiang
added
enhancement
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Suppressed comments (1)
pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java:267
- This makes a new behavior-changing execution path active whenever the option is absent. Pinot's rollout convention requires new feature flags to default off so production workloads can opt in and validate them before the default changes; the fallback alone does not prevent an immediate blast-radius change. Default this option to false initially and update the related documentation/benchmark description accordingly.
String option = _queryOptions.get(CommonConstants.Broker.Request.QueryOptionKey.GROUPING_SETS_BASE_AGGREGATION);
return option == null || Boolean.parseBoolean(option);
|
|
||
| // Derived group table keyed on (projected union values..., $groupingId). Values layout mirrors the record | ||
| // schema: key columns first, then the aggregation intermediates. | ||
| Map<Key, Record> derived = new HashMap<>(); |
| int packedId = extractPackedId(key, i); | ||
| keys[i] = packedId == _nullPackedIds[i] ? null : _dictionaries[i].getInternal(packedId); | ||
| } | ||
| keys[_numGroupByExpressions] = (int) (key >>> _bitShifts[_numGroupByExpressions]); |
|
Thanks for the review! Addressed all three inline comments in fc597a0: 1. 2. 3. On the suppressed comment about defaulting Note: the failing |
fc597a0 to
c64a993
Compare
|
Pushed d91f07a making base aggregation cardinality-adaptive so it's a strict win in both regimes. Why: base aggregation only pays off when rows collapse into far fewer base groups. For high-cardinality union columns the base grouping barely collapses the rows, so it adds a base pass + a derive pass on top of the same output — benchmarks showed CUBE over near-unique columns ~2x slower than expansion. How: before the scan, estimate the base-group count as the product of the union columns' dictionary cardinalities (saturating to Result (50 segments × 15k rows, base-agg default vs. forced expansion):
Low-cardinality keeps the 1.7–2.5x win; high-cardinality now matches expansion instead of regressing. Also folded in the earlier review fixes from the same round: the |
b137d09 to
5d132e9
Compare
xiangfu0
added
needs-attention
|
Pushed eba1229, which reworks base aggregation to do the grouping-set fan-out once, in parallel, at combine time — removing the need for the cardinality gate entirely (base aggregation is now a strict win in all regimes). Change: segments emit only the BASE groups (union grouping, no Benchmark (50 segments × 15k rows, base aggregation vs.
High-cardinality CUBE went from ~2x slower (the regression the gate protected against) to ~40x faster than expansion. (Small-dataset micro-numbers with wide error bars, but the order-of-magnitude difference is unambiguous.) Correctness fix (found by review): All 75 grouping-sets integration tests + the new unit tests pass; spotless/license/checkstyle clean. |
|
Pushed ad2b4ba fixing a trimming bug in the parallel derive that this design surfaced. Bug: the derive built its output table with Fix: the derive is a bounded transformation of the already-bounded base groups (base is capped at Added a regression test ( |
ad2b4ba to
b359df2
Compare
Grouping-set queries previously expanded every input row into one group per grouping set in each segment, so the per-set fan-out cost scaled with the number of scanned rows (O(rows * numSets)). This makes ROLLUP/CUBE several times slower than a plain GROUP BY over the same columns even though the results are derivable from a single base grouping. This change aggregates only the base grouping (the union of all grouping-set columns) once per segment -- reusing the fast plain-GROUP-BY path -- then derives the individual grouping-set records from those base groups by projecting rolled-up columns to NULL, stamping the $groupingId discriminator, and merging the base groups' aggregation intermediates. The per-set fan-out moves from O(rows) to O(base groups). Because a base group's intermediate result flows into every grouping set and AggregationFunction#merge mutates/returns its argument, each base intermediate is cloned per set (via the function's serialize/deserialize round-trip; scalar intermediates are immutable and skipped) before it can become a merge target, so object-backed accumulators (AVG, DISTINCTCOUNT, percentiles, ...) stay exact. The base path also applies the same per-set bucketed segment trim as the expansion path. Behavior change: this path is enabled by default via the new groupingSetsBaseAggregation query option; set it to false to force the legacy per-row expansion path. It is not used for multi-value group-by columns (an MV column fans a row across its values in the base grouping, which would over-count when rolled up) or filtered aggregations (distinct shared-generator path); both fall back to expansion. Also improves the legacy expansion generator (used for the fallback cases and groupingSetsBaseAggregation=false): resolves dictionary-encoded columns via native dict-ids, packs composite keys into a primitive long when they fit (avoiding per-group FixedIntArray allocation and array hashing), and reuses per-row group-id buffers across blocks. Adds a JMH benchmark (BenchmarkGroupingSetsQueriesSSE) and integration tests that compare the base-aggregation and expansion paths for equivalence across scalar/object accumulators and query shapes, plus coverage for the long-packed generator path.
…executor - GroupByUtils: bound the derived grouping-set map by numGroupsLimit while building it (the per-set fan-out can multiply base groups by up to numSets), instead of only trimming afterward; account the derived fan-out in the numGroups limit/warning flags. - GroupingSetsGroupKeyGenerator: disable long-packing when the union columns consume all 64 bits, so the $groupingId shift stays strictly below 64 (Java masks long shifts mod 64, which would otherwise turn a shift of 64 into 0 and corrupt the packed/unpacked ordinal). - BenchmarkGroupingSetsQueriesSSE: use a per-trial executor created in @setup and shut down in @teardown rather than a static one, so multiple JMH @PARAM trials in the same JVM do not submit to a shut-down pool.
Base aggregation only pays off when input rows collapse into far fewer base groups; for high-cardinality union columns it adds a base pass plus a derive pass on top of the same output and can be slower than per-row expansion (benchmarks showed CUBE over near-unique columns ~2x slower). Gate base aggregation on an estimate of the base-group count -- the product of the union columns' dictionary cardinalities, saturating at Long.MAX_VALUE on overflow or when any union column is non-dictionary-encoded (unknown cardinality). When the estimate exceeds the new groupingSetsBaseAggregationMaxGroups query option (default: numGroupsLimit), fall back to per-row expansion. This keeps the low-cardinality speedup (~1.7-2.5x) while matching expansion on high-cardinality inputs. The threshold option is parsed via QueryOptionsUtils so a malformed value yields a clear error naming the option. Adds a unit test for the estimate (product, non-dictionary, overflow, zero-cardinality branches) and an integration test for the cardinality-gated fallback.
Previously each segment expanded its base groups into all grouping sets before emitting (per-segment derive), so the combine merged numSegments * baseGroups * numSets records into the shared table -- and for high-cardinality union columns (baseGroups ~ rows) this ran slower than plain expansion, which the prior commit worked around with a cardinality gate. Instead, segments now emit only the BASE groups (union grouping, no $groupingId), and GroupByCombineOperator merges those across segments and derives the individual grouping sets ONCE in mergeResults(), parallelized across the combine thread pool by base-group ranges into a shared concurrent grouping-set table. Total work is now minimal (each row grouped once, never per set) and the per-set fan-out is multi-threaded, so base aggregation is a strict win in all cardinality regimes; the cardinality gate and its groupingSetsBaseAggregationMaxGroups option are removed. Object intermediates are cloned per derived record so concurrent cross-set/cross-thread merges stay exact. IndexedTable now derives its key-column count from the schema (columns minus aggregations) instead of always from getNumGroupByKeyColumns(). This fixes a cross-segment merge bug for the base combine table, whose base records omit the $groupingId column: the old fixed count merged aggregations at the wrong offset, dropping the first aggregation's cross-segment contribution and reading out of bounds. Adds DeriveGroupingSetsTest covering the parallel derive across thread counts and the cross-segment base merge (which reproduces the out-of-bounds without the fix).
The parallel combine-phase derive created its output table with resultSize = numGroupsLimit. For a grouping-set query without ORDER BY the concurrent upsert caps at that size and drops brand-new keys once full, so which derived groups survived depended on thread interleaving -- non-deterministic, and able to starve an entire low-magnitude grouping set (e.g. the grand total) when the detail set filled the quota first. The derive is a bounded transformation of the already-bounded base groups (base count is capped at numGroupsLimit per segment, so derived <= numGroupsLimit * numSets), so it must not re-apply that per-segment guardrail. Build the derived table with an unbounded result size and defer the real ORDER BY + LIMIT to the broker. Adds a regression test asserting the grand-total set is never dropped and that parallel derive is deterministic even when derived groups exceed a small numGroupsLimit.
b359df2 to
d7a8a08
Compare
…gation Base aggregation defers all ORDER BY trimming to the broker, which is exact but keeps every derived group (up to numGroupsLimit * numSets) in memory and on the wire per server. Add a groupingSetsServerTrimSize query option that, for a base-aggregation grouping-set query WITH an ORDER BY, keeps at most K groups within each grouping set on the server after the derive -- a per-set top-K bucketed by the $groupingId discriminator, so a global top-K can never starve a low-magnitude set such as the grand total. This bounds each server's output for high-cardinality unions at the cost of an approximate top-K (the broker still applies the exact final ORDER BY + LIMIT). Unset/non-positive (default) keeps the exact defer-to-broker behavior; ignored without ORDER BY. TableResizer.trimTableByGroupingSet applies the bucketed trim over a finished derived table's records. Adds an integration test asserting a large K is exact and a small K bounds the output without starving any set (including the grand total).
|
Made grouping-set trimming configurable per stage so accuracy vs. memory/latency can be tuned (pushed 20fa894). Where trimming happens now (base aggregation):
New knob:
The per-set bucketing (reusing the same anti-starvation logic as the legacy expansion path's segment trim) is the key: it never drops an entire grouping set, only ranks within each. Added an integration test asserting a large K is exact (no group dropped) and a small K (K=1) bounds the output while every set — including the grand total — still survives. All 76 grouping-sets tests pass; checkstyle/license/spotless clean. This gives operators the accuracy/cost dial you asked for; we can add segment-stage or broker-stage variants later on the same pattern if a use case wants them. |
Summary
GROUP BY
GROUPING SETS/ROLLUP/CUBEpreviously expanded every input row into one group per grouping set in each segment, so the per-set fan-out cost scaled with the number of scanned rows (O(rows * numSets)). This makes ROLLUP/CUBE several times slower than a plain GROUP BY over the same columns, even though the results are derivable from a single base grouping.This PR aggregates only the base grouping (the union of all grouping-set columns) once per segment — reusing the fast plain-GROUP-BY path — and then derives the individual grouping sets from the merged base groups, in parallel, at combine time. The per-set fan-out moves from
O(rows)(per segment, single-threaded) toO(base groups)(once, multi-threaded).Approach
Segment phase. A base-aggregation grouping-set query aggregates only the union columns, exactly like a plain GROUP BY, and emits the base groups — no per-segment expansion, no
$groupingIdcolumn.Combine phase.
GroupByCombineOperatormerges base groups across segments into a base-keyedIndexedTable, then inmergeResults()derives the grouping sets once: for each base group and each grouping set, it projects the base key (rolled-up columns →NULL), stamps the$groupingIddiscriminator, and merges the base group's aggregation intermediates into the derived group. The derivation is partitioned across the combine thread pool by base-group ranges into a shared concurrent grouping-set table, so the fan-out is multi-threaded and runs after the row-collapsing base merge rather than repeating per-row.This reuses the existing
AggregationFunction#mergemachinery the combine/reduce phases already rely on, so it is exact for every mergeable aggregation. Because a base group's intermediate flows into every grouping set (and across threads) andmergemutates/returns its argument, each base intermediate is cloned per derived record (via the function's serialize/deserialize round-trip; scalar intermediates are immutable and skipped), keeping object-backed accumulators (AVG, DISTINCTCOUNT, percentiles, ...) correct.Trimming. The derive is a bounded transformation of the already-bounded base groups (base is capped at
numGroupsLimitper segment, so derived ≤numGroupsLimit * numSets), so it does not re-apply that per-segment guardrail — doing so would drop derived groups non-deterministically under the parallel upsert and could starve a low-magnitude set such as the grand total. The real ORDER BY + LIMIT is deferred to the broker.Behavior change
The base-aggregation path is enabled by default via a new
groupingSetsBaseAggregationquery option. SetgroupingSetsBaseAggregation=falseto force the legacy per-row expansion path. Results are identical to the expansion path (verified by equivalence tests).Carve-outs that fall back to the expansion path:
Also included
Improvements to the legacy expansion generator (used for the fallback cases and
groupingSetsBaseAggregation=false):longwhen they fit, using aLong2IntOpenHashMapinstead ofObject2IntOpenHashMap<FixedIntArray>(avoids per-group object allocation and array hashing).A general fix in
IndexedTable: it now derives its key-column count from the schema (columns − aggregations) rather than always fromgetNumGroupByKeyColumns(), so the base combine table (whose records omit the$groupingIdcolumn) merges aggregations at the correct offset.A new JMH benchmark
BenchmarkGroupingSetsQueriesSSEexercises the full server→broker flow.Benchmark
Single-stage engine, 50 segments × 15k rows,
LIMIT 10000, base aggregation vs.groupingSetsBaseAggregation=false(per-row expansion). JMHAverageTime, 2 forks × 5 iterations (10 samples);±is the 99.9% error.Low / moderate cardinality (
EXP(0.5)):ROLLUP(D1, D2)ROLLUP(D1, D2, D3)CUBE(D1, D2, D3)GROUPING SETS ((D1),(D2),(D3),(D1,D2),())High cardinality (
EXP(0.001)), where base groups ≈ rows so there is little to collapse:ROLLUP(D1, D2)ROLLUP(D1, D2, D3)CUBE(D1, D2, D3)GROUPING SETS ((D1),(D2),(D3),(D1,D2),())Takeaways: base aggregation is a 1.8–2.7x win at low/moderate cardinality (the common OLAP case, where many rows collapse into few base groups) and roughly parity at extreme cardinality (base groups ≈ rows, so the base grouping is as expensive as the scan and the derive fan-out is comparable to expansion's per-row work). The one mild high-cardinality regression (
ROLLUP(D1, D2), 0.83x) sits in that near-unique regime.Testing
GroupingSetsQueriesTest(75 tests) passes, including:DeriveGroupingSetsTest(unit) covers the parallel derive:ArrayIndexOutOfBoundsExceptionwithout theIndexedTablekey-column fix).numGroupsLimit.Pre-commit checks (spotless, license, checkstyle) pass on all touched modules.