refactor(hash-aggr): Support spilling for partial and final mode aggregation - #24061
Open
2010YOUY01 wants to merge 2 commits into
Open
refactor(hash-aggr): Support spilling for partial and final mode aggregation#240612010YOUY01 wants to merge 2 commits into
partial and final mode aggregation#240612010YOUY01 wants to merge 2 commits into
Conversation
2010YOUY01
changed the title
partial and final mode aggregation
14 tasks
2010YOUY01
commented
Aug 3, 2026
| /// into an ordered streaming aggregation, which ensures bounded memory usage and | ||
| /// evaluates the final result. | ||
| /// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. | ||
| pub(crate) struct PartialHashAggregateStream { |
Contributor
Author
There was a problem hiding this comment.
We could put partial and final stream to two different files, I plan to this after this PR.
| | 3 | 1 | 2.0 | | ||
| | 3 | 2 | 5.0 | | ||
| | 4 | 3 | 11.0 | | ||
| | 4 | 1 | 4.0 | |
Contributor
Author
There was a problem hiding this comment.
they're both valid partial aggregation output, due to early emitting under memory limit.
| // enlarge memory limit to let the final aggregation finish | ||
| new_spill_ctx(2, 2600) | ||
| // Enlarge the memory limit enough to replay spilled states. | ||
| new_spill_ctx(2, 4640) |
Contributor
Author
There was a problem hiding this comment.
The legacy implementation underestimate the memory usage somehow, so here it relaxes the memory budget
|
|
||
| assert_eq!(3, output_rows); | ||
| if spill { | ||
| // When spilling, the output rows metrics become partial output size + final output size |
Contributor
Author
There was a problem hiding this comment.
This should be a bug in the legacy implementation. Final aggregation should return the same output_rows regardless of the spilling condition.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24061 +/- ##
==========================================
+ Coverage 80.88% 80.89% +0.01%
==========================================
Files 1101 1102 +1
Lines 375720 376376 +656
Branches 375720 376376 +656
==========================================
+ Hits 303895 304487 +592
- Misses 53729 53777 +48
- Partials 18096 18112 +16 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2010YOUY01
marked this pull request as draft
August 3, 2026 11:57
github-actions
Bot
added
core
2010YOUY01
marked this pull request as ready for review
August 4, 2026 08:14
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
2 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?
Part of #22710
Rationale for this change
This PR adds existing spilling feature into the new 2-staged (partial and final) aggregation. The high-level implementation idea is the same as the legacy implementation.
For the algorithm description for this feature, see top comment change at
datafusion/physical-plan/src/aggregates/hash_stream.rsWhat changes are included in this PR?
The key changes to the operator state machine are:
In file
datafusion/physical-plan/src/aggregates/hash_stream.rsPartialHashAggregateStream::poll_next()FinalHashAggregateStream::poll_next()Use this as the starting point, you can navigate to all the related changes, for example adding new states to implement larger-than-memory execution.
This PR also includes small fixes to memory reservation in
OrderedFinalAggregateStream. The bugs are caught by existing tests on aggregation spilling that is enabled in this PR.Are these changes tested?
Existing tests
codecovAre there any user-facing changes?
No