Conversation
| SELECT map_extract(MAP {'a': 1}, 'missing'); -- [] | ||
| ```rust | ||
| # /* comment to avoid running as a doctest | ||
| let exec = AggregateExec::builder(AggregateMode::Single, input) |
There was a problem hiding this comment.
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()?; |
There was a problem hiding this comment.
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.
| /// 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 |
There was a problem hiding this comment.
This is way too much docstring
| /// | ||
| /// ``` | ||
| /// # use std::sync::Arc; | ||
| /// # use arrow::datatypes::{DataType, Field, Schema}; |
There was a problem hiding this comment.
Also way too much docstring for an internal / non public method.
There was a problem hiding this comment.
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?; |
There was a problem hiding this comment.
This seems a bit superfluous. Maybe the caller should do this and we accept &Arc<AggrDynFilter>?
| 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() | ||
| }, | ||
| ); |
There was a problem hiding this comment.
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 { ... }
...| self.mode = mode; | ||
| self.invalidate_derived() |
There was a problem hiding this comment.
Should these methods check if mode = self.mode { return self } or something to avoid invalidating the derived state if not needed?
AggregateExecBuilder for building and rewriting AggregateExec
674921f to
2ba6d84
Compare
|
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 |
Codecov Report❌ Patch coverage is 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. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
🟡 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
AggregateExecBuilderwith 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
COUNTvalues 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.
There was a problem hiding this comment.
🟡 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
There was a problem hiding this comment.
🟡 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
fb79b5d to
aee2afb
Compare
AggregateExecBuilder for building and rewriting AggregateExecAggregateExec through an AggregateExecBuilder
`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
aee2afb to
cfba7d3
Compare
Which issue does this PR close?
AggregateExec"), which marks theAggregateExeclimit setters#[doc(hidden)] #[deprecated]and adds#[expect(deprecated)]at the call sites.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_newtakes six positional arguments, two of which are schemas that are easy to transpose (inputvsinput_schema). On top of it sit three clone-with-one-change methods —with_limit_options,with_new_limit_optionsandwith_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 throughAggregateExec::builder(mode, input)for a new node andAggregateExec::to_builder()for a rewrite:filter_exprdefaults to "no filter per aggregate"try_new_with_schemamoves into the builder, andtry_newandtry_new_with_schemadelegate to it, so there is one place anAggregateExecis constructedwith_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.buildreturns aResultso that the node can be checked in the one place it is built rather than at execution time. It currently checks only whattry_newalready checked, that the aggregate andFILTERexpressions have the same length.Migrated to the builder:
TopKAggregation,LimitedDistinctAggregation,CombinePartialFinalAggregate,OptimizeAggregateOrder, the protobuf decoder, and the tests that build anAggregateExecby 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 thelimit_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 intoCombinePartialFinalAggregate, 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:execute_typed's dispatch rather than enumerated.with_new_aggr_exprsbehaves the same way onmain.LimitOptionsintoSoftLimit { limit }/TopK { limit, descending }would make two of Reject a limit an AggregateExec cannot execute when the plan is built #25393's failure modes unrepresentable rather than merely rejected. Wider rename, noted on that issue.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.rscover the builder: defaulted filter expressions, mismatched filter arity, the derived state being reused on a rewrite (asserted withArc::ptr_eqon 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-planlib (2247) and doctestsdatafusion-physical-optimizer(37)datafusioncore_integration physical_optimizer(601)datafusion-proto(282)sqllogictestsuitecargo clippy --all-targets --all-features --workspace -- -D warningsandcargo fmt --all, on each of the three commitsNot run: the
sql_plannerplanning 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.mdhas a section for them.#[doc(hidden)].cargo-semver-checksclassifies this as requiring a major version:#[doc(hidden)]on the pre-existingAggregateExec::limit_optionsremoves 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; theCheck semverjob 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