malinjawi opened a new pull request, #12554:
URL: https://github.com/apache/gluten/pull/12554
## What changes were proposed in this pull request?
Adds `LazyAggregateExpandRule` to the Velox backend: a post-transform rule
that, for grouping analytics (ROLLUP / CUBE / GROUPING SETS), moves the partial
aggregation below the Expand operator so the input is aggregated once at the
finest grain, then expands only the intermediate aggregation buffers and merges
them before shuffle.
Before:
```
ColumnarExchange
HashAggregate(Partial) <- hashes (input rows x grouping sets)
Expand <- input rows x grouping sets
child
```
After:
```
ColumnarExchange
HashAggregate(PartialMerge, flushable) <- collapses coarse-grain
duplicates pre-shuffle
Expand <- fine-grain groups x grouping
sets (buffers only)
HashAggregate(Partial, flushable) <- input rows, finest grain
child
```
Gated by
`spark.gluten.sql.columnar.backend.velox.lazyAggregateExpand.enabled` (default
**false**). No native changes; the rewrite reuses existing companion-function,
extract-struct/row-construct, and flush/abandon machinery.
## Why are the changes needed?
Expand multiplies every input row by the number of grouping sets (TPC-DS
q67: 9x on an 8-key rollup), and the partial aggregate pays hashing and shuffle
for every copy. When the finest grain reduces the row count, aggregating first
is substantially cheaper, and the coarser grouping sets collapse from
already-aggregated states.
Safety on high-cardinality keys (where the finest grain does not reduce):
both inserted aggregates are flushable, so Velox's abandon-partial-aggregation
turns the bottom aggregate into a streaming to-intermediate conversion and the
merge stage into an identity pass-through - the worst case degrades to roughly
the original plan plus one cheap pass. This is the failure mode that regressed
q67 for the CH backend's lazy expand (#7986, fixed there by a native
re-aggregation step in #7995); here the protection comes from existing Velox
machinery.
Prior art: the CH backend has this optimization since #7649 (GLUTEN-7647,
default-on there). #12052 explored the same direction for the Velox backend;
this PR adds the pre-shuffle merge stage, exprId-based attribute binding, the
adaptive worst-case protection, and a dedicated test suite.
## Eligibility (v1, conservative)
- `sum` / `count` / `min` / `max` / `avg`; no DISTINCT, no FILTER, no
`try_sum`, no `bloom_filter_agg`
- all modes Partial; grouping and shuffle keys are attributes/literals; at
least one attribute-backed grouping key of atomic type (guards the degenerate
`GROUPING SETS ((),())` empty-input case)
- aggregate inputs must be pass-through columns of the Expand (this also
auto-rejects the `RewriteDistinctAggregates` Expand); deterministic
pre-projections and filters only
- float `sum`/`avg` follow the same `floatingPointMode` policy as flushable
aggregation
- every new node passes native validation, otherwise the rule leaves the
plan untouched
## Benchmarks (algorithm validation)
SQL-level simulation of the rewrite on vanilla Spark 4.0.1 (local[10],
12-core machine; result sets checksum-verified identical; the simulation pays
one extra shuffle the actual rule avoids, so these numbers understate the rule):
| dataset | fine-grain profile | baseline | rewrite | delta |
|---|---|---|---|---|
| 50M rows, 4 keys, 100k fine groups | strong reduction | 7.5s | 2.3s | 3.2x
faster |
| 30M rows, 8 keys, 9 sets, 3.1 rows/group (q67-shaped) | weak fine
reduction, strong coarse collapse | 134.6s | 43.3s | 3.1x faster |
| 20M rows, near-unique keys | no reduction | 25.0s | 47.6s | unprotected
worst case; motivates the flushable/abandon requirement above |
Gluten A/B numbers on TPC-DS (q67, q18, q22, q36, q77, q80) with the flag
on/off will follow - keeping this as a draft until then.
## How was this patch tested?
New `LazyAggregateExpandSuite` (20 cases): plan-shape assertions (flushable
aggregates on both sides of the Expand) plus result comparison against vanilla
Spark, covering rollup/cube with sum/count/min/max/avg, genuine-NULL keys vs
rolled-up NULL with `grouping_id`, decimal sum buffers, the pre-projected q67
shape, aggregates over grouping keys, empty input, degenerate grouping sets,
duplicate grouping sets, single and multi count-distinct, FILTER / strict-float
/ non-whitelisted-function / non-atomic-key rejections, early-abandon
interplay, AQE on/off, and flushable-disabled behavior.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
