Skip to content

feat: stream aggregates for partition-disjoint input - #24497

Draft
xavlee wants to merge 1 commit into
apache:mainfrom
xavlee:feat/issue-24438-partition-disjoint-aggregates
Draft

feat: stream aggregates for partition-disjoint input#24497
xavlee wants to merge 1 commit into
apache:mainfrom
xavlee:feat/issue-24438-partition-disjoint-aggregates

Conversation

@xavlee

@xavlee xavlee commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

Data can be range-partitioned into logical runs that are each locally grouped by (key, time) and non-overlapping on (key, date_bin(time)). Concatenating several runs into one DataFusion execution partition resets the global sort order, so AggregateExec currently falls back to blocking hash aggregation even though every group remains one contiguous range.

This change lets a data source communicate that stronger contiguity guarantee independently of physical sort order, allowing the first aggregate after an optional projection to emit completed groups incrementally.

What changes are included in this PR?

  • Add ExecutionPlan::partition_disjoint_exprs and the corresponding DataSource hook. The returned expressions form one composite tuple whose values must occur in a single contiguous range within each output partition.
  • Propagate the complete tuple through ProjectionExec; if any component cannot be projected, clear the property.
  • Let AggregateExec derive a private group-completion mode from the property while retaining InputOrderMode::Linear, so execution can reuse the existing ordered aggregate paths without advertising a false output ordering.
  • Support exact group-key matches and strict subsets:
    • exact match uses full group completion;
    • a disjoint-key subset uses partial group completion for additional grouping keys.
  • Consume the property at the first aggregate and keep grouping sets on their conservative non-streaming path.
  • Prevent round-robin repartitioning from destroying the source guarantee when it is enabling streaming.
  • Propagate the property across FFI execution plans and add source/test helpers plus optimizer and cross-library coverage.
  • Add a Criterion benchmark for concatenated, descending logical runs whose global key order resets.

A source may declare a derived expression such as (key, date_bin(time)) after verifying that logical-run boundaries do not split a bin. The projection path maps that expression but deliberately does not infer it from (key, time).

Are these changes tested?

Yes. Coverage includes:

  • early output before end-of-input;
  • order resets between logical runs;
  • exact and partial composite-key implication;
  • date_bin projection;
  • all-or-nothing projection propagation;
  • aggregate combination and distribution optimizer behavior;
  • spill replay and both aggregate implementations;
  • local and cross-library FFI round trips.

Validation run locally:

  • cargo test -p datafusion-physical-plan --lib (1,765 tests)
  • cargo test -p datafusion-datasource --lib (178 tests)
  • cargo test -p datafusion --test core_integration physical_optimizer (553 tests)
  • FFI unit/integration tests, including the cross-library execution-plan test
  • proto-feature focused tests
  • ./dev/rust_lint.sh
  • cargo clippy --all-targets --all-features -- -D warnings

The added benchmark measured approximately 885 us for unordered hash aggregation versus 663 us for partition-disjoint streaming on the development machine (about 25% faster for this workload).

Are there any user-facing changes?

Yes. Data source and execution-plan implementors can opt into the new property. It is a correctness contract: an invalid declaration can cause an aggregate group to be emitted before all of its rows have been observed.

The new FFI callback adds a field to the versioned FFI_ExecutionPlan struct. cargo-semver-checks therefore reports the expected major-version FFI layout change; datafusion-physical-plan and datafusion-datasource otherwise pass all semver checks.

@github-actions github-actions Bot added core Core DataFusion crate datasource Changes to the datasource crate ffi Changes to the ffi crate physical-plan Changes to the physical-plan crate labels Aug 19, 2026
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 77.17231% with 155 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.23%. Comparing base (c429919) to head (e3b1c0a).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-plan/src/aggregates/mod.rs 75.52% 66 Missing and 50 partials ⚠️
datafusion/physical-plan/src/test.rs 70.21% 13 Missing and 1 partial ⚠️
datafusion/physical-plan/src/projection.rs 77.77% 4 Missing and 8 partials ⚠️
datafusion/datasource/src/source.rs 55.55% 7 Missing and 1 partial ⚠️
datafusion/ffi/src/execution_plan.rs 95.83% 0 Missing and 2 partials ⚠️
.../aggregates/aggregate_hash_table/common_ordered.rs 50.00% 0 Missing and 1 partial ⚠️
...hysical-plan/src/aggregates/grouped_hash_stream.rs 0.00% 0 Missing and 1 partial ⚠️
...ysical-plan/src/aggregates/ordered_final_stream.rs 92.30% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #24497      +/-   ##
==========================================
- Coverage   81.24%   81.23%   -0.01%     
==========================================
  Files        1113     1113              
  Lines      392744   393384     +640     
  Branches   392744   393384     +640     
==========================================
+ Hits       319090   319582     +492     
- Misses      54900    54981      +81     
- Partials    18754    18821      +67     

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate datasource Changes to the datasource crate ffi Changes to the ffi crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support streaming aggregates when partitions are unsorted but non‑overlapping

2 participants