gmhelmold opened a new pull request, #23486:
URL: https://github.com/apache/datafusion/pull/23486
## Which issue does this PR close?
- Closes #23401.
## Rationale for this change
`SingleDistinctToGroupBy` rewrites a grouped `AGG(DISTINCT col)` into a
cheaper nested double aggregate. Today the rule only recognizes a **bare**
`Expr::AggregateFunction` in `Aggregate.aggr_expr`, so it fires for plans from
the SQL planner but silently skips the logically-identical plan built through
the DataFrame API — that path runs materially slower for the same query.
The divergence is purely in where each front-end places the output alias:
- **SQL** (`count(DISTINCT v) AS n`): the planner hoists `AS n` into an
outer `Projection`, leaving a bare `Expr::AggregateFunction` in `aggr_expr`.
- **DataFrame API** (`count(col("v")).distinct().alias("n")`): `.alias(..)`
wraps the aggregate directly, so `aggr_expr` holds
`Expr::Alias(AggregateFunction)`.
`is_single_distinct_agg` matched only the bare shape and returned `false`
for `Alias(AggregateFunction)`, so the rule never fired for DataFrame-API
distinct aggregates.
## What changes are included in this PR?
- Unwrap a single `Expr::Alias` layer in the matcher
(`is_single_distinct_agg`) and in the rewrite, mirroring the one-layer
`unalias()` idiom already used in the optimizer (e.g. `decorrelate.rs`), so
both front-end shapes are recognized identically.
- The user-facing alias is preserved without extra work: the rewrite already
rebuilds the final output `Projection` from the original `Aggregate` schema
(`schema.qualified_field(idx)`), so `AS n` survives.
- Bare-aggregate behavior is unchanged (byte-identical): every pre-existing
snapshot passes untouched.
## Are these changes tested?
Yes — new snapshot tests in `single_distinct_to_groupby.rs`:
- `single_distinct_aliased` — the DataFrame-API `Alias(AggregateFunction)`
shape is now rewritten, and the output column stays named `n`.
- `single_distinct_aliased_and_groupby` — same, alongside a real `GROUP BY`
column.
- `non_distinct_aliased` and `two_distinct_aliased_and_groupby` — negative
cases: an aliased non-distinct aggregate, and two aliased distinct aggregates
on different fields, are still (correctly) left untouched.
Reverting the production change makes the two positive tests fail (the rule
no longer fires) while every other test stays green, confirming the tests are
wired to the fix.
## Are there any user-facing changes?
No API changes. DataFrame-API queries using `count(distinct …)` (and other
single-distinct aggregates) now receive the same optimized plan as their SQL
equivalents — a performance improvement, with no change to results or output
column names.
--
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]
