Skip to content

refactor(hash-aggr): share one spill context between the spilling aggregate streams - #25538

Open
jayzhan211 wants to merge 2 commits into
apache:mainfrom
jayzhan211:simplify/agg-spill-context
Open

jayzhan211 wants to merge 2 commits into
apache:mainfrom
jayzhan211:simplify/agg-spill-context

Conversation

@jayzhan211

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

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.

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 20, 2026
@jayzhan211
jayzhan211 marked this pull request as ready for review September 20, 2026 10:26
@jayzhan211
jayzhan211 marked this pull request as draft September 20, 2026 10:28
@codecov-commenter

codecov-commenter commented Sep 20, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 92.99363% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.38%. Comparing base (d20936c) to head (c70ddc2).
⚠️ Report is 7 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-plan/src/aggregates/spill.rs 91.59% 7 Missing and 3 partials ⚠️
...fusion/physical-plan/src/aggregates/hash_stream.rs 91.66% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25538      +/-   ##
==========================================
- Coverage   82.38%   82.38%   -0.01%     
==========================================
  Files        1138     1139       +1     
  Lines      434491   434297     -194     
  Branches   434491   434297     -194     
==========================================
- Hits       357969   357799     -170     
+ Misses      54875    54839      -36     
- Partials    21647    21659      +12     

☔ 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
jayzhan211 marked this pull request as ready for review September 20, 2026 11:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants