jayzhan211 opened a new pull request, #25567: URL: https://github.com/apache/datafusion/pull/25567
## Which issue does this PR close? - Part of #20773 and #24704. Does not close them. Draft, opened to share measurements and get feedback on the direction. The option is off by default, so nothing changes unless it is set. ## Rationale for this change A hash aggregation with millions of groups per partition spends most of its time finding group ids: 69-97% of aggregate compute on ClickBench Q15-18 / Q31-35. The hash table and the group values no longer fit the CPU caches, and once several such tables of hundreds of MB compete (one per partition) the cost per row falls off a cliff (Q18: 42-51 ns per row up to 1M groups per partition, 530 ns at 4-8M). This PR lets a final aggregation stop growing one table: past a threshold it moves its state and all further input into 64 hash buckets and aggregates the buckets one after another with a small, reused table. Buckets are emitted, released, compacted, spilled and split again independently. Measured on an Apple M4 Pro, 12 partitions, `hash_aggregate_bucket_threshold = 262144`, time relative to `main` (interleaved runs, medians): | query | time | peak RSS | |---|---|---| | ClickBench Q32 | 0.34x | 0.61x | | ClickBench Q18 | 0.54x | 0.73x | | TPC-H SF10 `GROUP BY l_orderkey, l_linenumber` | 0.62x | 0.40x | | ClickBench Q15 | 0.77x | 0.53x | | nested aggregation over lineitem (`SinglePartitioned`) | 0.80x | 0.47x | | ClickBench Q35 / Q31 | 0.85x / 0.86x | 0.70x / 0.58x | | ClickBench, all 43 queries | ~0.75x | | | TPC-H SF10, all 22 queries | 1.01x | | Known costs, which is why this is a draft and off by default: | case | time | |---|---| | ClickBench Q13 / Q14 / Q12 (0.5-1M string-keyed groups per partition) | 1.07x / 1.06x / 1.03x | | `GROUP BY l_partkey, l_suppkey` (input repeats its groups 7x) | 0.86x, but 1.48x peak RSS | Where the remaining loss comes from (per final input row): Q14 costs 79-85 ns with one table and 91-99 ns bucketed, of which 42-46 ns is moving the row into its bucket (hash, gather, copy) and 49-54 ns the aggregation itself; Q31 goes from 85 ns to 24.5 + 19.5 ns. Moving rows is 35-55% of the bucketed aggregation everywhere, and `RepartitionExec` has already hashed, gathered and copied every one of these rows once. A trigger cannot fix that (a larger threshold, a self-timing trigger and holding input back were all measured and only move the loss); letting the repartition produce the buckets would remove it. I am prototyping that separately and will report on #20773. ## What changes are included in this PR? One commit per step: 1. `datafusion.execution.hash_aggregate_bucket_threshold` (groups, default 0 = off) and the bucketed final aggregation: `aggregates/final_buckets.rs` (routing by hash with one gather per batch, per-bucket compaction, per-bucket spill through `SpillManager`) and the bucket output in `FinalHashAggregateStream`. State with nested types is excluded; a soft group limit disables it. 2. A partial aggregation table that reaches the threshold is flushed downstream while its groups do not recur (about 1024 sampled group hashes per flush, kept for the last 64 flushes; flushing stops once more than 20% of a sample recurs). 3. One hash table is reused across the buckets, and buckets of unique groups are not compacted again. 4. `aggregates/bucketed_aggregation.rs` shares the bucket logic with the single-stage stream, which buckets when it runs on every partition (`SinglePartitioned`). A lone `Single` stream keeps its table: measured 6-25% slower with buckets, because every row is aggregated twice. 5. A final table whose input holds about one row per group moves into buckets at a quarter of the threshold, so that little work is repeated; buckets are still split again at the full threshold. ## What is the testing strategy for this PR? - Unit tests in `aggregates/final_buckets.rs` and `aggregates/hash_stream.rs`: bucketed results equal the single table for the final, partial and single-stage streams; buckets split again; buckets spill under a memory limit; repeated groups are compacted; the partial flush stops when groups recur. - `aggregate_bucketed.slt`: 14 queries (FILTER, ROLLUP, ordered `array_agg`, DISTINCT, nested aggregation) at 4 and 1 partitions, with the option off and at 100 groups, all sections identical, with `bucket_splits` / `table_flush_count` checked in `EXPLAIN ANALYZE`. - With the default temporarily set to 100, the aggregate unit tests, `memory_limit` tests and the aggregate / group by / distinct sqllogictests pass. One tight-memory file (`ordered_aggregate_spill.slt`, limits of 500K-2M) failed once in a loaded parallel run and passed in 6 isolated runs; memory behaviour under such limits differs with a 100-group threshold. `aggregate_fuzz::streaming_aggregate_test` does not, as expected: it compares partial state rows one by one, and a flushing partial aggregation legitimately emits a group more than once. It would have to merge partial rows first if the default ever became non-zero. - `cargo fmt`, workspace clippy and the extended test suite pass. ## Are there any user-facing changes? A new experimental configuration option, `datafusion.execution.hash_aggregate_bucket_threshold`, off by default, documented in `configs.md` together with its known costs. New metrics on `AggregateExec`: `bucket_splits`, `bucket_compactions`, `table_flush_count`. -- 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]
