adriangb opened a new pull request, #25376:
URL: https://github.com/apache/datafusion/pull/25376
## Which issue does this PR close?
- No issue. This is an alternative take on apache/datafusion#25257
("Deprecate public APIs for internal usage in AggregateExec"), which marks the
`AggregateExec` limit 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.
Every one of these plans can be built today, and each one only misbehaves
once it is executed:
| plan | what happens when it runs |
| --- | --- |
| limit, no MIN/MAX aggregate, no ordering direction |
`internal_err!("Ordering direction required for DISTINCT with limit")` |
| limit with more than one group by expression |
`aggr.group_expr().expr()[0]` panics |
| limit of `0` on a top-k aggregate | `worst_val().expect("Missing root")`
panics: a queue of capacity 0 reports itself full with an empty root |
| limit on an aggregate with a `FILTER` | `GroupedTopKAggregateStream` never
reads `filter_expr`, so the filter is silently dropped and the query returns
**wrong results** |
| limit on a `COUNT`/`AVG`/… aggregate | the TopK stream ignores the
aggregate, wrong results |
| limit with an unsupported key/value type | `debug_assert!` in debug builds
only |
On top of that, `AggregateExec` had two nearly identical
clone-with-one-change methods with different semantics
(`with_new_limit_options` resets the metrics, `with_limit_options` keeps them),
each hand-copying twelve fields, and a `try_new` with six positional arguments,
two of which are schemas that are easy to transpose (`input` vs `input_schema`).
Validating the node when it is built makes these unrepresentable in a built
plan, which means the API no longer needs to be dangerous — only internal.
## 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:
```rust
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"
- `build()` validates the node once: a limit is checked against the shape of
the aggregate (the rows in the table above), and a replacement of the aggregate
expressions is checked against both the output schema it inherits and the
dynamic filter it inherits
- 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 rewrites behave exactly as they did before.
A dynamic filter records which aggregate expressions are `MIN`/`MAX`, at
which index, and over which column — a `MIN` pushes down `col < bound`, a `MAX`
pushes `col > bound`. Carrying it verbatim across a rewrite that turns a `MIN`
into a `MAX` would push down the wrong predicate and prune rows the aggregate
needs, and the schema check cannot see it because both produce the same output
field. `with_aggr_exprs` therefore re-derives the state and keeps the inherited
filter only while it still describes the new expressions, so a rewrite that
merely reorders or reverses them keeps the original filter (and with it the
link to whichever child accepted it during pushdown).
**Migrated to the builder:** `TopKAggregation`,
`LimitedDistinctAggregation`, `CombinePartialFinalAggregate`,
`OptimizeAggregateOrder`, the protobuf decoder, and every test. The optimizer
rules use `.build().ok()?`, so an invalid limit means "skip this optimization"
instead of a plan that fails at runtime. 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.
Two existing tests were building plans that cannot execute and had to be
adjusted: a statistics test that put a limit on a `COUNT(a)` aggregate, and the
protobuf roundtrip test that put an ordered limit on `AVG(b)`. Both now use
shapes the optimizer actually produces.
## What is the testing strategy for this PR?
Eight new unit tests in `aggregates/builder.rs` cover the builder itself:
defaulted filter expressions, mismatched filter arity, derived state being
reused on a rewrite (asserted with `Arc::ptr_eq` on the plan properties) and
metrics being reset, the schema recompute when the mode changes, rejection of
incompatible aggregate expressions, each limit validation rule (including the
zero limit and the ordered-input case that used to produce an internal error at
execution time), and all three dynamic-filter outcomes when the aggregate
expressions are replaced — kept, rebuilt, dropped. The drop case uses `SUM`,
whose output field matches `MIN`'s, so it gets past the schema check and
actually reaches the dynamic filter. Two doctests on the builder and one on
`to_builder` cover the documented usage.
Both panics in the table above were reproduced before being fixed, not
reasoned about: the zero limit panics at `heap.rs:133`.
Everything else is covered by the existing suites, which all pass:
- `datafusion-physical-plan` lib (1874) and doctests
- `datafusion-physical-optimizer` (33)
- `datafusion` `core_integration physical_optimizer` (560)
- `datafusion-proto` aggregate roundtrips
- the full 505-file `sqllogictest` suite
- `cargo clippy --all-targets --all-features --workspace -- -D warnings`,
`cargo fmt --all`, and `cargo doc -p datafusion-physical-plan` under `-D
warnings`
Not run: the `sql_planner` planning benchmarks — the environment ran out of
disk during the release build. The change is plan-time only (execution paths
are untouched) and adds one `create_schema` call per aggregate rewritten by
`OptimizeAggregateOrder`, 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)]`.
- A limit that the aggregate cannot execute is now an error when the plan is
built, including when decoding from protobuf, instead of an internal error, a
panic, or silently ignored `FILTER` expressions at execution time. No plan
DataFusion's own optimizer produces is affected — the rules already check these
conditions before pushing a limit down.
- `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.
## Follow-up notes
Validation catches the bad limit shapes, but `LimitOptions` still lets you
construct them. Splitting it into `SoftLimit { limit }` / `TopK { limit,
descending }` would make two of the failure modes unrepresentable rather than
merely rejected. That's a wider rename, so it was left out of this PR — happy
to do it here if a reviewer wants to go further.
The deprecated `with_new_aggr_exprs` keeps the dynamic-filter hazard
described above: it copies the state verbatim, and this PR does not change its
behaviour. Only the builder path re-derives it.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_01PwTc51ca2XHDCyVB7MbJoz
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
