Skip to content

refactor: build and rewrite AggregateExec through an AggregateExecBuilder - #25376

Open
adriangb wants to merge 3 commits into
apache:mainfrom
pydantic:claude/datafusion-25257-builder-api-5pd3bs
Open

adriangb wants to merge 3 commits into
apache:mainfrom
pydantic:claude/datafusion-25257-builder-api-5pd3bs

Conversation

@adriangb

@adriangb adriangb commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

The methods #25257 hides are hidden because they are dangerous. They are just as dangerous for the code inside DataFusion that keeps using them, and hiding them does nothing about that. This PR gives those callers something to move to.

AggregateExec::try_new takes six positional arguments, two of which are schemas that are easy to transpose (input vs input_schema). On top of it sit three clone-with-one-change methods — with_limit_options, with_new_limit_options and with_new_aggr_exprs — two of which are near-identical with different semantics (one resets the metrics, the other keeps them), and each of which hand-copies twelve fields. A caller has to know which of those twelve fields the one it is replacing feeds.

What changes are included in this PR?

AggregateExecBuilder (datafusion/physical-plan/src/aggregates/builder.rs), reached through AggregateExec::builder(mode, input) for a new node and AggregateExec::to_builder() for a rewrite:

let exec = AggregateExec::builder(AggregateMode::Single, input)
    .with_group_by(group_by)
    .with_aggr_exprs(aggr_exprs)
    .with_limit_options(LimitOptions::new(10))
    .build()?;

let with_limit = exec
    .to_builder()
    .with_limit_options(LimitOptions::new_with_order(10, true))
    .build()?;
  • every argument is named, and filter_expr defaults to "no filter per aggregate"
  • the body of try_new_with_schema moves into the builder, and try_new and try_new_with_schema delegate to it, so there is one place an AggregateExec is constructed
  • a rewrite carries over the derived state of the node it came from — output schema, plan properties, ordering requirements, dynamic filter — instead of asking each caller to copy the fields. Structural setters (with_mode, with_group_by, with_input, with_filter_exprs) drop that state and recompute; the others keep it, so a rewrite costs what the method it replaces cost.
  • build returns a Result so that the node can be checked in the one place it is built rather than at execution time. It currently checks only what try_new already checked, that the aggregate and FILTER expressions have the same length.

Migrated to the builder: TopKAggregation, LimitedDistinctAggregation, CombinePartialFinalAggregate, OptimizeAggregateOrder, the protobuf decoder, and the tests that build an AggregateExec by hand. No #[expect(deprecated)] anywhere.

Deprecated and #[doc(hidden)]: with_limit_options, with_new_limit_options, with_new_aggr_exprs.

#[doc(hidden)] but not deprecated: AggregateExec::builder, AggregateExec::to_builder, AggregateExecBuilder, and the limit_options() getter. Building and rewriting an aggregate is how DataFusion's own optimizer rules work, not a public API, so the whole surface is hidden — as #25257 does. The getter is not deprecated because reading a limit is safe and has no replacement; deprecating it would only push #[expect(deprecated)] back into CombinePartialFinalAggregate, which lives in another crate and cannot reach the field directly.

This is a refactor. The node the builder produces is the node the method it replaces produced, on every path. Nothing new is rejected, and no plan changes shape.

Follow-ups

An earlier revision of this PR also validated the node in build. That is where the interesting work is, but it changes behaviour and deserves its own review, so it is carved out:

Both optimizer rules that push a limit down already read build().ok()?, so #25393 can land without touching them.

What is the testing strategy for this PR?

Five unit tests in aggregates/builder.rs cover the builder: defaulted filter expressions, mismatched filter arity, the derived state being reused on a rewrite (asserted with Arc::ptr_eq on the plan properties) and the metrics being reset, that setting a field to the value it already has does not invalidate that state, and the schema recompute when the mode changes. One doctest covers the documented usage.

The rest is covered by the existing suites, unchanged, which is the point of a refactor:

  • datafusion-physical-plan lib (2247) and doctests
  • datafusion-physical-optimizer (37)
  • datafusion core_integration physical_optimizer (601)
  • datafusion-proto (282)
  • the full 520-file sqllogictest suite
  • cargo clippy --all-targets --all-features --workspace -- -D warnings and cargo fmt --all, on each of the three commits

Not run: the sql_planner planning benchmarks. Plan-time only, and the builder performs the same work the methods it replaces performed, so it should be in the noise, but it has not been measured.

Are there any user-facing changes?

Yes, and docs/source/library-user-guide/upgrading/56.0.0.md has a section for them.

  • The three methods above are deprecated with a replacement, and this whole API is now #[doc(hidden)].
  • No behavioural change: no plan that builds today stops building, and no plan changes shape.
  • cargo-semver-checks classifies this as requiring a major version: #[doc(hidden)] on the pre-existing AggregateExec::limit_options removes it from the public API (major), and the three deprecations are a minor change. That is the intended consequence of hiding an internal API and is expected for the 56.0.0 release; the Check semver job reports it as a note and passes. Nothing else in the four checked crates moved.

🤖 Generated with Claude Code

https://claude.ai/code/session_01D7arPq4Frxu8byr17mqVKA

@github-actions github-actions Bot added documentation Improvements or additions to documentation optimizer Optimizer rules core Core DataFusion crate proto Related to proto crate physical-plan Changes to the physical-plan crate labels Sep 16, 2026

@adriangb adriangb left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some nits

SELECT map_extract(MAP {'a': 1}, 'missing'); -- []
```rust
# /* comment to avoid running as a doctest
let exec = AggregateExec::builder(AggregateMode::Single, input)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use the rust,ignore directive instead of a comment


aggr_exprs = try_convert_aggregate_if_better(
aggr_exprs,
&requirement,
input.equivalence_properties(),
)?;

let aggr_exec = aggr_exec.with_new_aggr_exprs(aggr_exprs);
let aggr_exec =
aggr_exec.to_builder().with_aggr_exprs(aggr_exprs).build()?;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please lets check when to_builder() is used (and discriminate tests vs. production code) and if it should be into_builder() and take ownership instead (or maybe we need both). It'd be nice to avoid clones if we can.

Comment on lines +35 to +61
/// Builder for [`AggregateExec`].
///
/// This is the recommended way to create an [`AggregateExec`], and the only
/// supported way to derive a new [`AggregateExec`] from an existing one (see
/// [`AggregateExec::to_builder`]).
///
/// Like the methods it replaces, this is public for internal use only and is
/// not part of the public API: it is how DataFusion's own physical optimizer
/// rules build and rewrite aggregates, and it may change without notice. It is
/// `#[doc(hidden)]` for that reason, not because it is unfinished.
///
/// Compared to calling [`AggregateExec::try_new`] and then mutating individual
/// fields, the builder:
///
/// 1. Names every argument, so `input` / `input_schema` and `aggr_expr` /
/// `filter_expr` can't be transposed by accident.
/// 2. Defaults `filter_expr` to "no filter for each aggregate", which is what
/// the vast majority of callers want and removes a common source of
/// length-mismatch panics.
/// 3. Validates the plan once, at the end, so combinations that would panic or
/// return an internal error during execution (for example a limit pushed
/// into an aggregate that cannot execute it) are rejected up front.
/// 4. Keeps derived state (output schema, plan properties, ordering
/// requirements, dynamic filter) consistent when rewriting an existing
/// node, instead of asking every caller to copy the fields by hand.
///
/// # Example: creating a new aggregate

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is way too much docstring

///
/// ```
/// # use std::sync::Arc;
/// # use arrow::datatypes::{DataType, Field, Schema};

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also way too much docstring for an internal / non public method.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if we could minimize code duplication by having try_new_with_schema / other existing constructors delegate to the builder. That way there is only one code block that builds and verifies.

existing: Option<Arc<AggrDynFilter>>,
aggr_expr: &[Arc<AggregateFunctionExpr>],
) -> Option<Arc<AggrDynFilter>> {
let existing = existing?;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems a bit superfluous. Maybe the caller should do this and we accept &Arc<AggrDynFilter>?

Comment on lines +387 to +394
let compatible =
computed.fields().len() == schema.fields().len()
&& computed.fields().iter().zip(schema.fields()).all(
|(computed, existing)| {
computed.data_type() == existing.data_type()
&& computed.is_nullable() == existing.is_nullable()
},
);

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comparison is impossible for a human to reason through. Can we break it down into steps w/ named variables? E.g.

let field_count_matches: bool = computed.fields().len() == schema.fields().len();
if !field_count_matches { ... }
...

Comment on lines +198 to +199
self.mode = mode;
self.invalidate_derived()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should these methods check if mode = self.mode { return self } or something to avoid invalidating the derived state if not needed?

@adriangb adriangb changed the title Claude/datafusion 25257 builder api 5pd3bs feat: add a validated AggregateExecBuilder for building and rewriting AggregateExec Sep 16, 2026
@adriangb
adriangb force-pushed the claude/datafusion-25257-builder-api-5pd3bs branch from 674921f to 2ba6d84 Compare September 16, 2026 16:43
@github-actions

github-actions Bot commented Sep 16, 2026

Copy link
Copy Markdown

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details
     Cloning apache/main
    Building datafusion v55.1.0 (current)
       Built [  68.036s] (current)
     Parsing datafusion v55.1.0 (current)
      Parsed [   0.037s] (current)
    Building datafusion v55.1.0 (baseline)
       Built [  66.533s] (baseline)
     Parsing datafusion v55.1.0 (baseline)
      Parsed [   0.042s] (baseline)
    Checking datafusion v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   0.612s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [ 137.031s] datafusion
    Building datafusion-physical-optimizer v55.1.0 (current)
       Built [  46.310s] (current)
     Parsing datafusion-physical-optimizer v55.1.0 (current)
      Parsed [   0.025s] (current)
    Building datafusion-physical-optimizer v55.1.0 (baseline)
       Built [  44.426s] (baseline)
     Parsing datafusion-physical-optimizer v55.1.0 (baseline)
      Parsed [   0.023s] (baseline)
    Checking datafusion-physical-optimizer v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   0.110s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  92.140s] datafusion-physical-optimizer
    Building datafusion-physical-plan v55.1.0 (current)
       Built [  42.201s] (current)
     Parsing datafusion-physical-plan v55.1.0 (current)
      Parsed [   0.171s] (current)
    Building datafusion-physical-plan v55.1.0 (baseline)
       Built [  44.263s] (baseline)
     Parsing datafusion-physical-plan v55.1.0 (baseline)
      Parsed [   0.176s] (baseline)
    Checking datafusion-physical-plan v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   0.684s] 223 checks: 221 pass, 2 fail, 0 warn, 31 skip

--- failure inherent_method_now_doc_hidden: inherent method #[doc(hidden)] added ---

Description:
A method or associated fn is now #[doc(hidden)], removing it from the crate's public API.
        ref: https://doc.rust-lang.org/rustdoc/write-documentation/the-doc-attribute.html#hidden
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.50.0/src/lints/inherent_method_now_doc_hidden.ron

Failed in:
  AggregateExec::limit_options in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/98eb685c29a9751590a0d617e1db8df8bb9f819d/datafusion/physical-plan/src/aggregates/mod.rs:1125

--- failure type_method_marked_deprecated: type method #[deprecated] added ---

Description:
A type method is now #[deprecated]. Downstream crates will get a compiler warning when using this method.
        ref: https://doc.rust-lang.org/reference/attributes/diagnostics.html#the-deprecated-attribute
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.50.0/src/lints/type_method_marked_deprecated.ron

Failed in:
  method datafusion_physical_plan::aggregates::AggregateExec::with_new_aggr_exprs in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/aggregates/mod.rs:945
  method datafusion_physical_plan::aggregates::AggregateExec::with_new_limit_options in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/aggregates/mod.rs:973
  method datafusion_physical_plan::aggregates::AggregateExec::with_limit_options in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/aggregates/mod.rs:1054

     Summary semver requires new major version: 1 major and 1 minor checks failed
    Finished [  88.692s] datafusion-physical-plan
    Building datafusion-proto v55.1.0 (current)
       Built [  59.899s] (current)
     Parsing datafusion-proto v55.1.0 (current)
      Parsed [   0.018s] (current)
    Building datafusion-proto v55.1.0 (baseline)
       Built [  59.879s] (baseline)
     Parsing datafusion-proto v55.1.0 (baseline)
      Parsed [   0.019s] (baseline)
    Checking datafusion-proto v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   0.111s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [ 121.148s] datafusion-proto

@github-actions github-actions Bot added the auto detected api change Auto detected API change label Sep 16, 2026
@codecov-commenter

codecov-commenter commented Sep 16, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 89.70588% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.32%. Comparing base (bf67e97) to head (cfba7d3).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-plan/src/aggregates/builder.rs 90.43% 4 Missing and 27 partials ⚠️
datafusion/physical-plan/src/aggregates/mod.rs 88.13% 1 Missing and 6 partials ⚠️
...ical-optimizer/src/limited_distinct_aggregation.rs 80.00% 0 Missing and 1 partial ⚠️
...afusion/physical-optimizer/src/topk_aggregation.rs 80.00% 0 Missing and 1 partial ⚠️
...fusion/physical-optimizer/src/update_aggr_exprs.rs 50.00% 0 Missing and 1 partial ⚠️
...hysical-plan/src/aggregates/grouped_topk_stream.rs 83.33% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25376      +/-   ##
==========================================
- Coverage   82.33%   82.32%   -0.02%     
==========================================
  Files        1137     1138       +1     
  Lines      432049   432289     +240     
  Branches   432049   432289     +240     
==========================================
+ Hits       355734   355880     +146     
- Misses      54815    54875      +60     
- Partials    21500    21534      +34     

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

@adriangb
adriangb requested review from 2010YOUY01 and a balanced review from Copilot September 16, 2026 23:12

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Soft limits on non-MIN/MAX aggregates can stop input early and return incomplete aggregate values.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds a validated builder for safely constructing and rewriting AggregateExec nodes.

Changes:

  • Introduces AggregateExecBuilder with schema, limit, and dynamic-filter validation.
  • Migrates optimizer rules, protobuf decoding, and tests to the builder.
  • Deprecates unsafe mutation APIs and documents migration.
File summaries
File Description
docs/source/library-user-guide/upgrading/56.0.0.md Documents the new builder and migration.
datafusion/proto/tests/cases/plans/aggregates.rs Updates aggregate roundtrip coverage.
datafusion/physical-plan/src/aggregates/mod.rs Exposes the builder and delegates construction.
datafusion/physical-plan/src/aggregates/builder.rs Implements construction and validation.
datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs Migrates Top-K tests to the builder.
datafusion/physical-optimizer/src/update_aggr_exprs.rs Uses the builder for expression rewrites.
datafusion/physical-optimizer/src/topk_aggregation.rs Validates Top-K limit pushdown.
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs Validates distinct-limit pushdown.
datafusion/physical-optimizer/src/combine_partial_final_agg.rs Rebuilds combined aggregates through the builder.
datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs Migrates optimizer tests.
datafusion/core/tests/execution/coop.rs Migrates cooperative execution tests.
Review details

Suppressed comments (1)

datafusion/physical-plan/src/aggregates/builder.rs:839

  • This assertion codifies an unsafe case: the regular hash stream honors this soft limit by stopping input consumption after enough groups, leaving COUNT values incomplete when later batches contain the same groups. This case should expect rejection; direction-less limits are only safe here when there are no aggregate expressions.
        // a direction-less limit is a soft limit for any aggregate
        AggregateExec::builder(AggregateMode::Single, test_input(&schema))
            .with_group_by(group_by_a(&schema)?)
            .with_aggr_exprs(vec![count_b(&schema)?])
            .with_limit_options(LimitOptions::new(10))
  • Files reviewed: 11/11 changed files
  • Comments generated: 2
  • Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread datafusion/physical-plan/src/aggregates/builder.rs Outdated
Comment thread datafusion/proto/tests/cases/plans/aggregates.rs Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Limit validation permits conflicting MIN/MAX and explicit ordering directions, which can retain the wrong TopK groups.

Get a fresh assessment by requesting another Copilot review.

Review details
  • Files reviewed: 11/11 changed files
  • Comments generated: 1
  • Review effort level: Balanced

Comment thread datafusion/physical-plan/src/aggregates/builder.rs Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Aggregate-expression rewrites can be silently ignored or retain stale ordering requirements, risking incorrect execution.

Get a fresh assessment by requesting another Copilot review.

Review details
  • Files reviewed: 11/11 changed files
  • Comments generated: 3
  • Review effort level: Balanced

Comment thread datafusion/physical-plan/src/aggregates/builder.rs Outdated
Comment thread datafusion/physical-plan/src/aggregates/builder.rs Outdated
Comment thread docs/source/library-user-guide/upgrading/56.0.0.md Outdated
@adriangb
adriangb force-pushed the claude/datafusion-25257-builder-api-5pd3bs branch from fb79b5d to aee2afb Compare September 17, 2026 02:10
@adriangb adriangb changed the title feat: add a validated AggregateExecBuilder for building and rewriting AggregateExec refactor: build and rewrite AggregateExec through an AggregateExecBuilder Sep 17, 2026
`AggregateExec::try_new` takes six positional arguments, two of which are
schemas that are easy to transpose (`input` vs `input_schema`), and the
fields optimizer rules change afterwards were set with `with_*` methods
that copy the remaining twelve fields by hand.

Move the body of `try_new_with_schema` into a new
`AggregateExecBuilder`, reached through `AggregateExec::builder` for a
new node and `AggregateExec::to_builder` for a rewrite of an existing
one. Every argument is named, `filter_expr` defaults to "no filter per
aggregate", and a rewrite carries over the derived state of the node it
came from (output schema, plan properties, ordering requirements,
dynamic filter) unless a field that state is computed from changes, so
it costs no more than the clone-with-one-change methods it replaces.

`try_new` and `try_new_with_schema` now delegate to the builder, so
there is one place an `AggregateExec` is constructed. `build` returns a
`Result` so the node can be checked there rather than at execution time;
for now it checks only what `try_new` already checked, that the
aggregate and `FILTER` expressions have the same length.

No behaviour change: every path produces the node it produced before.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D7arPq4Frxu8byr17mqVKA
Move every caller that builds or rewrites an `AggregateExec` onto
`AggregateExec::builder` and `AggregateExec::to_builder`:
`TopKAggregation`, `LimitedDistinctAggregation`,
`CombinePartialFinalAggregate`, `OptimizeAggregateOrder`, the protobuf
decoder, and the tests that construct one by hand.

`CombinePartialFinalAggregate` replanned the partial aggregate with
`try_new` and then copied the limit onto it; it now derives the combined
node from the partial one, which is the same thing said once.

The two optimizer rules that push a limit down take `build().ok()?`, so
a node the builder cannot produce means "skip this optimization" rather
than a failed plan. Nothing rejects a limit today, so this cannot
trigger; it is how the rules should read once something does.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D7arPq4Frxu8byr17mqVKA
`with_limit_options`, `with_new_limit_options` and
`with_new_aggr_exprs` each clone the node with one field replaced and
leave the caller to know which of the other twelve fields that field
feeds. Deprecate them in favour of the builder.

Building and rewriting an `AggregateExec` is how DataFusion's own
physical optimizer rules work, not a public API, so the whole surface is
`#[doc(hidden)]`: the builder, `AggregateExec::builder`,
`AggregateExec::to_builder`, the three deprecated methods, and the
`limit_options` getter. The getter is not deprecated, because reading a
limit is safe and has no replacement.

This is what apache#25257 asked for, with a replacement to
point callers at and no `#[expect(deprecated)]` left inside DataFusion.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D7arPq4Frxu8byr17mqVKA
@adriangb
adriangb force-pushed the claude/datafusion-25257-builder-api-5pd3bs branch from aee2afb to cfba7d3 Compare September 17, 2026 18:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

auto detected api change Auto detected API change core Core DataFusion crate documentation Improvements or additions to documentation optimizer Optimizer rules physical-plan Changes to the physical-plan crate proto Related to proto crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants