fangchenli opened a new pull request, #27080:
URL: https://github.com/apache/datafusion/pull/27080
## Which issue does this PR close?
- Closes #23401.
This picks up #23403, which the stale bot closed. It has been reworked on
current `main`; see the note on the constant-argument guard below.
## Rationale for this change
The same query gets a different plan depending on how it is written:
```rust
// Not rewritten: the alias sits directly on the aggregate
df.aggregate(vec![col("g")], vec![count_distinct(col("x")).alias("n")])?
```
```sql
-- Rewritten into a two-phase GROUP BY
SELECT g, count(DISTINCT x) AS n FROM t GROUP BY g
```
`SingleDistinctToGroupBy` only matches a bare `Expr::AggregateFunction`, so
an alias on the aggregate turns it off. DataFrame users therefore miss the
rewrite that SQL users get, including the parallelism it provides.
Measured on 4M rows and 2,000 groups, using the DataFrame API (median of 3
rounds, each the median of 9 runs). Peak RSS includes the in-memory table:
| Aggregate | Before | After | SQL, for reference | Peak RSS before → after |
| --- | --- | --- | --- | --- |
| `count(DISTINCT int)` | 32 ms | 34 ms | 33 ms | 0.80 → 0.82 GB |
| `sum(DISTINCT int)` | 54 ms | 35 ms | 32 ms | 0.65 → 0.84 GB |
| `count(DISTINCT varchar)` | 71 ms | 34 ms | 37 ms | 1.28 → 0.98 GB |
The main point is that both APIs now reach the same plan. The result is not
always an improvement: `count(DISTINCT int)` has a specialized accumulator and
gains nothing, and `sum(DISTINCT int)` trades memory for speed, which is the
tradeoff described in #25932. Any policy on when the rewrite should apply, such
as #26000, is orthogonal and would now apply to both APIs equally.
## What changes are included in this PR?
- `is_single_distinct_agg` looks through aliases on the aggregate, and the
rewrite strips them. Output names are unaffected because the outer projection
restores them from the schema.
- The doc comment on `can_strip_every_distinct` said this rule does not
rewrite an aliased aggregate. That is no longer true, so the comment now gives
the reason that still applies. Its behavior does not change.
An earlier version of this PR also skipped the rewrite for constant
`DISTINCT` arguments, because the rewrite exposed a wrong result on empty input
in `EliminateGroupByConstant`. #22132 has since fixed that, so the guard is
gone.
## What is the testing strategy for this PR?
- A DataFrame API test, `aliased_single_distinct_aggregate_is_rewritten`,
checks the optimized plan and the results end to end.
- Unit tests in `single_distinct_to_groupby.rs` cover aliased aggregates
with and without `GROUP BY`, nested aliases, partly aliased aggregates, the
case that must stay unrewritten (two distinct fields), and the `count` rollup
path.
## Are there any user-facing changes?
DataFrame queries with an aliased single `DISTINCT` aggregate now get the
same plan as the equivalent SQL, so their performance and memory use can change
as shown above. Results and output names do not change. There are no public API
changes.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
