feat: add AggregateMode::PartialReduce for tree-reduce aggregation by njsmith · Pull Request #20019 · apache/datafusion

njsmith · 2026-01-27T06:35:49Z

DataFusion's current AggregateMode enum has four variants covering three of the four cells in the input/output matrix:

	Input: raw data	Input: partial state
Output: final values	`Single` / `SinglePartitioned`	`Final` / `FinalPartitioned`
Output: partial state	`Partial`	???

This PR adds AggregateMode::PartialReduce to fill in the missing cell: it takes partially-reduced values as input, and reduces them further, but without finalizing.

This is useful because it's the key component needed to implement distributed tree-reduction (as seen in e.g. the Scuba or Honeycomb papers): a set of worker nodes each perform multithreaded Partial aggregations, feed those into a PartialReduce to reduce all of this node's values into a single row, and then a head node collects the outputs from all nodes' PartialReduce to feed into a Final reduction.

PR can be reviewed commit by commit: first commit is pure refactor/simplification; most places we were matching on AggregateMode we were actually just trying to either check which row of the above table we were in, or else which column. So now we have is_first_stage (tells you which column) and is_last_stage (tells you which row) and we use them everywhere.

Second commit adds PartialReduce, and is pretty small because is_first_stage/is_last_stage do most of the heavy lifting. It also adds a test demonstrating a minimal Partial -> PartialReduce -> Final tree-reduction.

Add is_last_stage() method to AggregateMode (complementing the existing is_first_stage()) and refactor match statements throughout the aggregate execution code to use these helpers instead of explicit variant listing. This makes the code robust to adding new AggregateMode variants -- instead of updating every match arm, new variants only need to be added to is_first_stage() and is_last_stage(). Refactored sites: - create_schema: uses is_last_stage() - aggregate_expressions: uses is_first_stage() - finalize_aggregation: uses is_last_stage() - no_grouping.rs filter expressions: uses is_first_stage() - no_grouping.rs update vs merge dispatch: uses is_first_stage() - row_hash.rs filter expressions: uses is_first_stage() - row_hash.rs update vs merge dispatch: uses is_first_stage() - row_hash.rs create_batch_from_accumulators: uses is_last_stage()

njsmith · 2026-01-27T20:18:22Z

fixed the clippy/cargo fmt failures from CI

adriangb · 2026-01-27T21:17:02Z

Hi @njsmith, nice to see you over here :)

Add a new PartialReduce variant to AggregateMode that takes partially- computed aggregate state as input and produces further-reduced (but still partial) aggregate state as output. This fills the missing cell in the input/output matrix: | Input: raw data | Input: state Output: values | Single/SP | Final/FP Output: state | Partial | PartialReduce <-- new This enables tree-reduce distributed aggregation patterns where multiple rounds of partial reduction happen before the final aggregation step. Changes: - Add PartialReduce variant to AggregateMode enum - Update is_first_stage() (false) and is_last_stage() (false) helpers - Update required_input_distribution() (UnspecifiedDistribution) - Add proto serialization support (PARTIAL_REDUCE = 5) - Add unit tests for mode helpers - Add integration tests for Partial->PartialReduce->Final pipeline - Add integration test for chained tree-reduce pattern

njsmith · 2026-01-27T21:48:44Z

🤦 now hopefully fixed CI (forgot that .clone() can do implicit coercion and Arc::clone can't)

alamb · 2026-01-27T22:31:50Z

FYI @gabotechs -- this may be of interest for the datafusion-distributed

adriangb

This is a very clean way to organize the change, I was immediately able to understand the goal and implementation.

I think it's worth getting another set of eyes (or I can dig a bit deeper into the implications, I'm out of time today) but this looks great from my perspective!

adriangb · 2026-01-28T01:52:49Z

+ /// Checks whether this aggregation step produces final output values
+ /// (as opposed to intermediate accumulator state).
+ pub fn is_last_stage(&self) -> bool {


I think it could be helpful to have your table and the description of "tells us in which row/column" as docstrings on is_last_stage and is_first_stage.
Would it also be easier to understand if this was called output_final_results() or output_mode() -> OutputMode where enum OutputMode { Partial, Final }?

Ping on this before we move to merge

Sure happy to make the change, just didn't want to make noisy changes until people had chance to look things over

gabotechs

👍 This looks pretty good to me. Adding a new building block (even if it's not immediately used in the codebase) while reducing complexity of existing code is a +1 on my side.

Just had one tangential question.

gabotechs · 2026-01-28T10:20:22Z

+ /// ```text
+ /// Final
+ /// / \
+ /// PartialReduce PartialReduce
+ /// / \ / \
+ /// Partial Partial Partial Partial
+ /// ```
+ PartialReduce,


This is looking really good 👍 I don't think a partial reduce implementation can get cleaner than this.
One question: do you think this could be beneficial for single-node DataFusion plans at some point? I can imagine how this is useful for distributed engine, but do you have any thoughts about single-node usage?

Regarding distributed engines, this produces a good alternative for shuffling, and I imagine that for systems that rely on materializing shuffle results, this can be a pretty good alternative. However, I do also imagine that in-memory streaming-based distributed systems like Trino or https://github.com/datafusion-contrib/datafusion-distributed might not benefit as much from this pattern because:
Shuffling is as cheap as it can get, as intermediate results are zero-copied over the network between workers
Shuffling allows running the final aggregation step concurrently in multiple workers, while a tree-reduction approach limits the final aggregation to one worker

I get the impression that https://github.com/apache/datafusion-ballista might be a good candidate for this aggregation pattern, but I'd bet that it might not yield solid performance improvements in https://github.com/datafusion-contrib/datafusion-distributed

If you have heterogenous links between your compute (e.g. a bunch of machines, and each machine contains a bunch of cores, where same-machine cores have much higher bandwidth/lower latency between them than the cross-machine links), then this is still useful in the shuffle world I think?
Each node partitions local data into N threads -> Partial each partition -> within-machine shuffle -> PartialReduce -> cross-machine shuffle -> Final gives a 1/N-reduction in network transfer, right? (compared to shuffling each thread's Partial result directly)
Whether this would be worth the extra complexity depends on a bunch of deployment-specific constants of course.

🤔 Yes, I think you are right there, I can imagine how applying an extra partial reduce right before shuffling over the wire can give solid improvements.

I think this Partial Reduce mode makes sense
I find our other modes somewhat confusing (I am not sure why we have them and suspect we could make this a much cleaner 2 mode switch like "initial", "partial", and "final"

(nothing needed for this PR)

@gabotechs @njsmith the benefits are real in aggregations where the cardinality of the aggregation key is large. That is why all the Aggregation frameworks that I am aware of (spark, hadoop, algebird etc) have a separate Combiner itermediate type/function for merging partial aggregates.
Both hadoop and spark do this by default and it's called map side combine. Big +1 for the work on this PR!
PS: the most elegant definition of generic aggregation that I know of is in the Algebird scala framework (developed originally at twitter ~15y ago): https://github.com/twitter/algebird/blob/develop/algebird-core/src/main/scala/com/twitter/algebird/Aggregator.scala#L367-L390

gabotechs · 2026-01-29T07:10:57Z

IMO this is pretty much ready to do, I'll let @adriangb finish his review. Thanks for this addition!

alamb

Looks good to me too -thanks @njsmith and @gabotechs

It would be nice if we could figure out how to clean up the mess of aggregate modes (Final and FinalPartitioned and Single seem like they should be merged) but that is for some other day perhaps

alamb · 2026-01-29T19:14:25Z

+ /// ```text
+ /// Final
+ /// / \
+ /// PartialReduce PartialReduce
+ /// / \ / \
+ /// Partial Partial Partial Partial
+ /// ```
+ PartialReduce,


(nothing needed for this PR)

…_mode Replace the boolean helpers with typed enums: - is_first_stage() -> input_mode() -> AggregateInputMode { Raw, Partial } - is_last_stage() -> output_mode() -> AggregateOutputMode { Partial, Final } Add a table to AggregateMode docs showing how each variant maps to its input/output modes. This makes the semantics more explicit and self-documenting.

njsmith · 2026-01-30T02:55:16Z

Added the table to the docs and replaced is_first_stage/is_last_stage -> bool with input_mode/output_mode returning easier to interpret enums, as per review comments.

adriangb · 2026-01-30T03:18:34Z

Added the table to the docs and replaced is_first_stage/is_last_stage -> bool with input_mode/output_mode returning easier to interpret enums, as per review comments.

Just to get a second opinion, @gabotechs do you agree this naming is clearer?

gabotechs · 2026-01-30T07:18:14Z

Yes! Not like the previous one was bad, but this one is slightly better

gabotechs · 2026-01-30T09:55:35Z

Merged! thanks @njsmith for the great PR and @adriangb and @alamb for the reviews

…pache#20019) DataFusion's current `AggregateMode` enum has four variants covering three of the four cells in the input/output matrix: | | Input: raw data | Input: partial state | | - | - | - | | Output: final values | `Single` / `SinglePartitioned` | `Final` / `FinalPartitioned` | | Output: partial state | `Partial` | ??? | This PR adds `AggregateMode::PartialReduce` to fill in the missing cell: it takes partially-reduced values as input, and reduces them further, but without finalizing. This is useful because it's the key component needed to implement distributed tree-reduction (as seen in e.g. the Scuba or Honeycomb papers): a set of worker nodes each perform multithreaded `Partial` aggregations, feed those into a `PartialReduce` to reduce all of this node's values into a single row, and then a head node collects the outputs from all nodes' `PartialReduce` to feed into a `Final` reduction. PR can be reviewed commit by commit: first commit is pure refactor/simplification; most places we were matching on `AggregateMode` we were actually just trying to either check which row of the above table we were in, or else which column. So now we have `is_first_stage` (tells you which column) and `is_last_stage` (tells you which row) and we use them everywhere. Second commit adds `PartialReduce`, and is pretty small because `is_first_stage`/`is_last_stage` do most of the heavy lifting. It also adds a test demonstrating a minimal Partial -> PartialReduce -> Final tree-reduction.

github-actionsBot added proto Related to proto crate physical-plan Changes to the physical-plan crate labels Jan 27, 2026

njsmith force-pushed the njsmith/aggregate-partial-reduce-v2 branch 2 times, most recently from 0727a56 to 0ced71aCompare January 27, 2026 20:17

njsmith force-pushed the njsmith/aggregate-partial-reduce-v2 branch from 0ced71a to a854145Compare January 27, 2026 21:31

njsmith force-pushed the njsmith/aggregate-partial-reduce-v2 branch from a854145 to 97b3d15Compare January 27, 2026 21:47

adriangb reviewed Jan 27, 2026
View reviewed changes

adriangb reviewed Jan 28, 2026
View reviewed changes

gabotechs approved these changes Jan 28, 2026
View reviewed changes

gabotechs mentioned this pull request Jan 28, 2026
feat: Support multiple output tasks from NetworkCoalesceExec datafusion-contrib/datafusion-distributed#323
Merged

alamb approved these changes Jan 29, 2026
View reviewed changes

adriangb approved these changes Jan 29, 2026
View reviewed changes

github-actionsBot added the optimizer Optimizer rules label Jan 30, 2026

gabotechs added this pull request to the merge queue Jan 30, 2026

Merged via the queue into apache:main with commit 7c3ea05Jan 30, 2026
32 checks passed

alamb mentioned this pull request Jan 30, 2026
Alamb/consolidate aggreate enu #20085
Closed

fmonjalet mentioned this pull request Feb 26, 2026
Leverage PartialReduce AggregationExec mode to drastically reduce shuffle size datafusion-contrib/datafusion-distributed#360
Closed

2010YOUY01 mentioned this pull request Jul 2, 2026
refactor(hash-aggr): Migrate partial-reduce hash aggregation #23233
Merged

2010YOUY01 mentioned this pull request Aug 11, 2026
refactor(hash-aggr): Let partial-reduce mode aggregation error under memory limit #24262
Draft

Uh oh!

Conversation

njsmith commented Jan 27, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

njsmith commented Jan 27, 2026

Uh oh!

adriangb commented Jan 27, 2026

Uh oh!

njsmith commented Jan 27, 2026

Uh oh!

alamb commented Jan 27, 2026

Uh oh!

adriangb 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!

gabotechs left a comment • 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!

gabotechsJan 28, 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!

Choose a reason for hiding this comment

Uh oh!

njsmithJan 29, 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aditanaseJan 30, 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!

gabotechs commented Jan 29, 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!

njsmith commented Jan 30, 2026

Uh oh!

adriangb commented Jan 30, 2026

Uh oh!

gabotechs commented Jan 30, 2026

Uh oh!

Uh oh!

gabotechs commented Jan 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

njsmith commented Jan 27, 2026•
edited
Loading

gabotechs left a comment •
edited
Loading

gabotechsJan 28, 2026•
edited
Loading

njsmithJan 29, 2026•
edited
Loading

aditanaseJan 30, 2026•
edited
Loading