Repository navigation
refactor(hash-aggr): share one spill context between the spilling aggregate streams - #25538
Merged
Merged
Conversation
jayzhan211
marked this pull request as ready for review
September 20, 2026 10:26
jayzhan211
marked this pull request as draft
September 20, 2026 10:28
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #25538 +/- ##
==========================================
+ Coverage 82.38% 82.41% +0.02%
==========================================
Files 1138 1139 +1
Lines 434491 435209 +718
Branches 434491 435209 +718
==========================================
+ Hits 357969 358680 +711
+ Misses 54875 54830 -45
- Partials 21647 21699 +52 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
jayzhan211
marked this pull request as ready for review
September 20, 2026 11:38
2010YOUY01
approved these changes
Sep 21, 2026
2010YOUY01
left a comment
2010YOUY01
left a comment
Contributor
There was a problem hiding this comment.
Thank you, love it (and all entropy-reducing PRs in general!)
| /// Sorts `state_batch`, the intermediate state of all currently buffered | ||
| /// groups (`None` if there are no groups), and writes it as one spill file. | ||
| /// Memory reservation should be updated by the caller. | ||
| pub(super) fn spill(&mut self, state_batch: Option<RecordBatch>) -> Result<()> { |
Contributor
There was a problem hiding this comment.
perhaps rename it to sort_and_spill?
Co-authored-by: Yongting You <2010youy01@gmail.com>
Contributor
|
run benchmark clickbench_partitioned |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing simplify/agg-spill-context (4a6cfd6) to d20936c (merge-base) diff Run configurationrun benchmark clickbench_partitionedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing simplify/agg-spill-context (4a6cfd6) to d20936c (merge-base) diff Run configurationrun benchmark clickbench_partitionedCPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
Omega359
pushed a commit
to Omega359/arrow-datafusion
that referenced
this pull request
Oct 11, 2026
…regate streams (apache#25538) ## Which issue does this PR close? - Part of apache#25537 (unified spill-replay driver for grouped aggregation). Follow-up to apache#22710. ## Rationale for this change The four grouped aggregation streams that can spill — `FinalHashAggregateStream`, `SingleHashAggregateStream`, `OrderedSingleAggregateStream` and `OrderedFinalAggregateStream` — each define their own spill context (`FinalSpillContext`, `SingleSpillContext`, `OrderedSingleSpillContext`, `OrderedFinalSpillContext`). The four types have the same seven fields and the same three operations: - build the spill sort key and the aggregate configuration used for replay, - sort the table's state batch and write it as one spill file, - merge all runs with `StreamingMergeBuilder` and replay them through `OrderedFinalAggregateStream::new_with_input_and_metrics`. The bodies are copies of each other; the only real differences are constructor-time: | | sort key | replay aggregate | |---|---|---| | Final hash | group columns in natural order | same aggregate | | Single hash | group columns in natural order | `Single → Final`, `group_by.as_final()` | | Ordered final | already-ordered columns first, then the rest | same aggregate | | Ordered single | already-ordered columns first, then the rest | `Single → Final`, `group_by.as_final()` | Both axes are functions of values the constructor already receives (`AggregateMode` and `InputOrderMode`), and the natural order is just the ordered-columns-first formula with no ordered columns. ## What changes are included in this PR? No behaviour change. - Add `aggregates/spill.rs` with a single non-generic `AggregateSpill` (`try_new`, `has_spills`, `spill`, `into_replay_stream`). It is non-generic because the only thing a spill needs from a table is the batch returned by `take_state_batch()`, so the caller passes that batch in. - Delete the four per-stream spill context types and use `AggregateSpill` instead. - Spill request descriptions (`"FinalHashAggregateSpill"` etc.) and memory consumer names are unchanged. The text of four `internal_err!` messages that cannot be reached by users is now shared. Net: 4 stream files −665 lines, +1 new file of ~235 lines. This is the first step towards a single spill-replay driver for these streams (see the linked issue), but it stands on its own. ## What is the testing strategy for this PR? Existing tests: the aggregate unit tests in `aggregates/mod.rs` and `ordered_final_stream.rs` (spill + replay, OOM, drop/cancel, memory accounting), `aggregate_memory_spill.slt`, `ordered_aggregate_spill.slt`, the `memory_limit` integration tests and the aggregate fuzz tests. No new tests since there is no new behaviour. ## Are there any user-facing changes? No. --------- Co-authored-by: Yongting You <2010youy01@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
The four grouped aggregation streams that can spill —
FinalHashAggregateStream,SingleHashAggregateStream,OrderedSingleAggregateStreamandOrderedFinalAggregateStream— each define their own spill context (FinalSpillContext,SingleSpillContext,OrderedSingleSpillContext,OrderedFinalSpillContext). The four types have the same seven fields and the same three operations:StreamingMergeBuilderand replay them throughOrderedFinalAggregateStream::new_with_input_and_metrics.The bodies are copies of each other; the only real differences are constructor-time:
Single → Final,group_by.as_final()Single → Final,group_by.as_final()Both axes are functions of values the constructor already receives (
AggregateModeandInputOrderMode), and the natural order is just the ordered-columns-first formula with no ordered columns.What changes are included in this PR?
No behaviour change.
aggregates/spill.rswith a single non-genericAggregateSpill(try_new,has_spills,spill,into_replay_stream). It is non-generic because the only thing a spill needs from a table is the batch returned bytake_state_batch(), so the caller passes that batch in.AggregateSpillinstead."FinalHashAggregateSpill"etc.) and memory consumer names are unchanged. The text of fourinternal_err!messages that cannot be reached by users is now shared.Net: 4 stream files −665 lines, +1 new file of ~235 lines.
This is the first step towards a single spill-replay driver for these streams (see the linked issue), but it stands on its own.
What is the testing strategy for this PR?
Existing tests: the aggregate unit tests in
aggregates/mod.rsandordered_final_stream.rs(spill + replay, OOM, drop/cancel, memory accounting),aggregate_memory_spill.slt,ordered_aggregate_spill.slt, thememory_limitintegration tests and the aggregate fuzz tests. No new tests since there is no new behaviour.Are there any user-facing changes?
No.