refactor(hash-aggr): Let partial-reduce mode aggregation error under memory limit

[2010YOUY01]

Which issue does this PR close?

Part of #22710

Rationale for this change

This PR is part of the aggregation refactor/migration; see the epic for details. Specifically, it implements larger-than-memory behavior for the PartialReduce aggregation mode. Note the OOM behavior differs from the legacy implementation , due to the following reasons:

Behavior difference

For PartialReduce aggregation, this PR handles the memory limit by returning an error immediately. The existing behavior when PartialReduce runs out of memory is to spill and then perform final streaming aggregation. This behavior is implemented in the legacy code:

• feat: add AggregateMode::PartialReduce for tree-reduce aggregation #20019

Why change the behavior

1. The original PR does not discuss the intended OOM behavior, and there are no tests specifically covering it.
2. For the major use case of PartialReduce aggregation—reducing data locally before shuffle (Leverage PartialReduce AggregationExec mode to drastically reduce shuffle size datafusion-contrib/datafusion-distributed#360) think early emission may actually be the intended behavior.

So I tend to believe the existing behavior was an auto-completion from AI, instead of an intentional design decision. Therefore, this PR takes the more conservative approach of returning an error on OOM. We can revisit and improve this behavior in follow-up PRs.

What changes are included in this PR?

• partial_reduce_stream.rs: key changes, let aggregation update memory usage during execution, and return error when OOM

Are these changes tested?

UTs

Are there any user-facing changes?

[codecov-commenter]

Codecov Report

❌ Patch coverage is 62.79070% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.99%. Comparing base (c429919) to head (961ee8e).
⚠️ Report is 35 commits behind head on main.

Files with missing lines	Patch %	Lines
...sical-plan/src/aggregates/partial_reduce_stream.rs	57.14%	13 Missing and 2 partials ⚠️
datafusion/physical-plan/src/aggregates/mod.rs	87.50%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24486      +/-   ##
==========================================
+ Coverage   81.24%   81.99%   +0.75%     
==========================================
  Files        1113     1117       +4     
  Lines      392744   409120   +16376     
  Branches   392744   409120   +16376     
==========================================
+ Hits       319090   335467   +16377     
- Misses      54900    55165     +265     
+ Partials    18754    18488     -266     

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

[jayzhan211]

So I tend to believe the existing behavior was an auto-completion from AI, instead of an intentional design decision. Therefore, this PR takes the more conservative approach of returning an error on OOM. We can revisit and improve this behavior in follow-up PRs.

How about we come out the idea design for example early emission, instead of revert it back to error state 🤔

[2010YOUY01]

So I tend to believe the existing behavior was an auto-completion from AI, instead of an intentional design decision. Therefore, this PR takes the more conservative approach of returning an error on OOM. We can revisit and improve this behavior in follow-up PRs.

How about we come out the idea design for example early emission, instead of revert it back to error state 🤔

That idea was just my intuition, maybe we can leave the decision to real users (cc @gabotechs) for this feature, and design+implement that in a separate PR.

[alamb]

For the major use case of PartialReduce aggregation—reducing data locally before shuffle (datafusion-contrib/datafusion-distributed#360) think early emission may actually be the intended behavior.

I agree this is the intended behavior -- and I think it would be better behavior to avoid errors entirely, and instead dump the hash table we have accumulated so far, and start aggregating again.

[alamb]

As I mentioned elsewhere, I think that since PartialReduce mode is just converting one partial state to another partial intermediate state as an optimization before sending over the network, it would actually be better here to emit any gathered state on OOM pressure and start re-aggregating (or just start copying the input directly to the output) rather than error here

[alamb]

I made a PR to validate this expectation

• docs: Document expectations of PartialReduce aggregate mode #24640

[gabotechs]

👍 Agree with this. Partial reduction is not needed for correctness, it's purely for performance, so if it accumulates too much memory worth of a hash table, it might be better to flush it to the output, and even just fallback to a bypass that does not aggregate anything.

Whether we should be re-aggregating or just bypassing the input to the output directly, my guess is that whatever is simpler and introduces less code is probably the best initial approach.

[2010YOUY01]

Based on the discussion, let’s switch to the early-emission approach instead, since it is (a) simpler and (b) good enough for the target workload. I’ll put together a new PR.

Thank you all for the feedback.