refactor(hash-aggr): Support spilling for `partial` and `final` mode aggregation by 2010YOUY01 · Pull Request #24061 · apache/datafusion

2010YOUY01 · 2026-08-03T09:51:14Z

Which issue does this PR close?

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.rs

What changes are included in this PR?

The key changes to the operator state machine are:
In file datafusion/physical-plan/src/aggregates/hash_stream.rs

PartialHashAggregateStream::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

TODO for myself: double check the codecov

Are there any user-facing changes?

No

2010YOUY01 · 2026-08-03T09:55:20Z

+///    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 {


We could put partial and final stream to two different files, I plan to this after this PR.

2010YOUY01 · 2026-08-03T09:56:23Z

            | 3 | 1             | 2.0         |
            | 3 | 2             | 5.0         |
-            | 4 | 3             | 11.0        |
+            | 4 | 1             | 4.0         |


they're both valid partial aggregation output, due to early emitting under memory limit.

2010YOUY01 · 2026-08-03T09:59:04Z

-            // 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)


The legacy implementation underestimate the memory usage somehow, so here it relaxes the memory budget

2010YOUY01 · 2026-08-03T10:00:49Z


+        assert_eq!(3, output_rows);
        if spill {
-            // When spilling, the output rows metrics become partial output size + final output size


This should be a bug in the legacy implementation. Final aggregation should return the same output_rows regardless of the spilling condition.

codecov-commenter · 2026-08-03T11:54:33Z

Codecov Report

❌ Patch coverage is 84.45808% with 76 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.05%. Comparing base (21ad189) to head (084f633).
⚠️ Report is 40 commits behind head on main.

Files with missing lines	Patch %	Lines
...fusion/physical-plan/src/aggregates/hash_stream.rs	84.11%	64 Missing and 11 partials ⚠️
datafusion/physical-plan/src/aggregates/mod.rs	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24061      +/-   ##
==========================================
+ Coverage   80.88%   81.05%   +0.17%     
==========================================
  Files        1101     1105       +4     
  Lines      375720   380386    +4666     
  Branches   375720   380386    +4666     
==========================================
+ Hits       303895   308318    +4423     
- Misses      53729    53848     +119     
- Partials    18096    18220     +124

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

alamb · 2026-08-05T18:45:37Z

👀

alamb

Thank you @2010YOUY01

I went through this PR carefully and I (again) found it clear and well written and easy to follow

It is somewhat unfortunate, perhaps, that the state machines for the different spillable hash streams are so similar (we now have I think 3 state machines that have the various spilling / merging states) -- where I think the difference in the streams is largely related to what triggers a spill (vs an emit) and the type of hash table they have

Maybe we can find some way to extract out the state machine into a common structure (with inline functions or something) to capture the differences

alamb · 2026-08-05T18:52:17Z

        let spill_schema = Arc::clone(spill_manager.schema());
+        // The merge and replay table are two components of the same aggregate
+        // operator. Keep them under one consumer registration so a fair memory
+        // pool does not divide this operator's quota between its own phases.


is this a behavior change? Or does it mirror what the old operator does?

(It seems reasonable to me, but I wanted to check)

Yes, it mirrors the existing behavior. It's only get caught by the existing test after this PR, likely because partial/final 2-staged aggregation has better test coverage.

alamb · 2026-08-05T18:57:18Z

                }

-                if self.hit_soft_group_limit(original_state.hash_table()) {
+                if self.hit_soft_group_limit(&hash_table) {


Claude code flagged this as a potential bug -- if the limit is set, this code path appears to simply drop the spill_context and any spilled content thus far

Maybe update to

let spilled = spill_context .as_ref() .is_some_and(|ctx| ctx.has_spills()); if self.hit_soft_group_limit(&hash_table) && !spilled { // existing early-output path }

🤔

But somehow you have to switch to emitting 🤔

Updated in 084f633

I applied the same logic: if spilled then don't trigger soft limit optimization. Those limits are usually small constants, so they're unlikely to be co-exist with spilling, so I think this extra check would be enough.

BTW I found only good AI models are able to find such tricky bugs effectively, they're really hard to construct tests for. It's quite valuable to let coding agents scan the codebase to find logic inconsistencies.

alamb · 2026-08-05T19:02:37Z

-                        original_state,
-                    ));
+                // Check memory reservation, and potentially spill.
+                let timer = elapsed_compute.timer();


a nit here is that you could potentially use the same timer as above rather than making a new one here

It seems not possible, timer.done() would take self.

As in I was thinking that you could avopid timer.done() -- as the destructor also stops the timer I think

alamb · 2026-08-05T19:05:33Z

-            &original_state,
-            PartialHashAggregateState::SkippingAggregation { .. }
-        ));
+        let PartialHashAggregateState::SkippingAggregation { mut hash_table } =


you could potentially avoid these internal errors by simply passing in hash_table rather than an original_state

The call site already matched on the type of PartialHashAggregateState

I think the benefit of this pattern is, all the logic are implemented within this state function, and the main event loop poll_next() can be kept minimal, the entry-condition check like this one won't get scattered to 2 places.

The downside, as you pointed out, is that it is more verbose. I don’t have a strong preference at this point.

Me neither - this way is fine too

alamb · 2026-08-05T19:06:12Z

+        self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
+    }
+
+    fn break_with_err(error: DataFusionError) -> FinalHashAggregateStateTransition {


this structure is quite elegant and easy to follow 👍

2010YOUY01 · 2026-08-06T04:18:37Z

@alamb thank you for the review

alamb · 2026-08-06T13:01:34Z

Thaks again @2010YOUY01

alamb · 2026-08-06T13:01:56Z

I am keen to get this refactor wrapped up before DF 55 so merging it in

feat: implement spilling for migrated partial/final hash aggregation

3066f51

2010YOUY01 changed the title ~~feat: implement spilling for migrated partial/final hash aggregation~~ refactor(hash-aggr): Support spilling for partial and final mode aggregation Aug 3, 2026

2010YOUY01 mentioned this pull request Aug 3, 2026

[EPIC] Split Aggregation Logic into Dedicated Streams #22710

Open

14 tasks

2010YOUY01 commented Aug 3, 2026

View reviewed changes

2010YOUY01 marked this pull request as draft August 3, 2026 11:57

sesteves mentioned this pull request Aug 3, 2026

Grouped approx_distinct can overflow Arrow BinaryArray state offsets #24072

Open

fix memory accounting bugs

c11f63f

github-actions Bot added core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) physical-plan Changes to the physical-plan crate labels Aug 4, 2026

2010YOUY01 marked this pull request as ready for review August 4, 2026 08:14

2010YOUY01 requested a review from alamb August 4, 2026 08:14

alamb approved these changes Aug 5, 2026

View reviewed changes

review: fix potential bug

084f633

alamb added this pull request to the merge queue Aug 6, 2026

Merged via the queue into apache:main with commit e6b4221 Aug 6, 2026
40 checks passed

Uh oh!

Conversation

2010YOUY01 commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

alamb commented Aug 5, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 commented Aug 6, 2026

Uh oh!

alamb commented Aug 6, 2026

Uh oh!

alamb commented Aug 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

2010YOUY01 commented Aug 3, 2026 •

edited

Loading

2010YOUY01 Aug 3, 2026 •

edited

Loading

codecov-commenter commented Aug 3, 2026 •

edited

Loading