mkleen opened a new pull request, #25288:
URL: https://github.com/apache/datafusion/pull/25288
## Which issue does this PR close?
- Part of #24929 (the `min`/`max` half).
- Closes #11686.
## Rationale for this change
`SELECT g, min(DISTINCT x) FROM t GROUP BY g` is planned as if `DISTINCT`
mattered. `SingleDistinctToGroupBy` rewrites it into an inner group by at `(g,
x)` grain just to deduplicate `x`, even though `min` returns the same value
with or without duplicates. On a 4M-row table with ~1M distinct `x` values,
that extra stage makes the query **up to 10x slower** and reserves **~190 MB**
where the plain `min(x)` needs less than 1 MB.
The same applies to every aggregate whose result cannot change when
duplicates are removed: `max`, `bool_and`, `bool_or`, `bit_and`, `bit_or`,
`approx_distinct`.
Nothing in DataFusion could say which aggregates those are.
`SingleDistinctToGroupBy` identifies `sum`/`min`/`max` by lowercased function
name, and a user-defined aggregate had no way to declare the property at all.
## What changes are included in this PR?
**A new `AggregateUDFImpl::distinct_handling()` method** returning a new
enum `DistinctHandling`:
| Variant | Meaning | Built-ins |
|---|---|---|
| `Ignored` | Duplicates cannot change the result (the merge is idempotent),
so the planner may drop `DISTINCT`. | `min`, `max`, `bool_and`, `bool_or`,
`bit_and`, `bit_or`, `approx_distinct`, `grouping` |
| `Honored` *(default)* | `DISTINCT` is applied, either by the accumulator
itself (`count`, `sum`, `avg`, `var_samp`, `array_agg`, `bit_xor`, ...) or by
`SingleDistinctToGroupBy` deduplicating the input first (`stddev`,
`approx_median`). | everything not listed |
| `Unsupported` | The result depends on duplicates, but nothing deduplicates
the input, so `f(DISTINCT ...)` errors or silently returns the non-distinct
answer. Declaration only for now; rejecting these at planning time is a
follow-up. | `corr`, `covar_samp`, `covar_pop`, `regr_*`, `nth_value`,
`approx_percentile_cont`, `approx_percentile_cont_with_weight` |
`first_value`, `last_value` and `any_value` stay at `Honored` with a TODO:
whether `DISTINCT` is a no-op for them depends on `ORDER BY`.
**A new logical optimizer rule, `EliminateAggregateDistinct`**, placed
before `SingleDistinctToGroupBy`. It clears the `DISTINCT` flag on `Ignored`
aggregates and keeps the output column name with an alias:
```text
Aggregate: groupBy=[[g]], aggr=[[min(DISTINCT x)]]
-- becomes
Aggregate: groupBy=[[g]], aggr=[[min(x) AS min(DISTINCT x)]]
```
The rule is **all-or-nothing per `Aggregate` node**: it only fires when
every `DISTINCT` on the node is `Ignored`. Stripping only some flags would
change which plans `SingleDistinctToGroupBy` rewrites, because that rule keys
off how many distinct aggregates a node has and whether they share an argument.
So `min(DISTINCT x), sum(DISTINCT x)` is left exactly as on `main`.
## What is the testing strategy for this PR?
- **Unit tests** in
`datafusion/optimizer/src/eliminate_aggregate_distinct.rs`: `min`, `min`+`max`,
a non-distinct aggregate beside the distinct one, `FILTER`, the negative cases
`sum` and `bit_xor`, a mixed node that must be left alone (on its own and
followed by `SingleDistinctToGroupBy`), an `Honored` distinct aggregate in the
group expressions, and a plan with no `Aggregate`.
- **sqllogictest:**
- `aggregates_simplify.slt` gets a new section. For each `Ignored`
aggregate it checks the `EXPLAIN` plan and the query results, over data with
duplicates and `NULL`s. It also covers the negative cases (`bit_xor`, `count`),
a mixed node, and `FILTER`.
- `single_distinct_to_groupby.slt` has an updated expectation for
`min(DISTINCT v)`.
- `explain.slt` lists the new rule.
- **Documentation:** a new "Declaring how an Aggregate UDF treats
`DISTINCT`" section in `adding-udfs.md`, and the rule in
`optimizer_rule_reference.md`.
### Benchmark
Here are some results from a local benchmark. Every query has the shape
`SELECT g, <aggregate> FROM t GROUP BY g`, over 4M rows where `x` has about 1M
distinct values and `b` has two.
Speedup (`main` time / branch time):
| query | 2k groups, 1 part | 2k groups, 10 parts | 500k groups, 1 part |
500k groups, 10 parts |
|---|---:|---:|---:|---:|
| min | 10.10x | 6.56x | 3.44x | 2.49x |
| max | 10.37x | 7.19x | 3.35x | 2.44x |
| bit_and | 10.36x | 7.73x | 3.50x | 2.53x |
| bit_or | 10.25x | 7.10x | 3.33x | 2.52x |
| bool_and | 1.66x | 1.53x | 1.51x | 2.36x |
| bool_or | 1.50x | 1.65x | 1.50x | 1.71x |
| approx_distinct | 1.97x | 1.77x | 1.77x | 1.77x |
Peak memory-pool reservation in MB (`main` → branch):
| query | 2k groups, 1 part | 2k groups, 10 parts | 500k groups, 1 part |
500k groups, 10 parts |
|---|---:|---:|---:|---:|
| min | 191.4 → 0.2 | 167.9 → 1.5 | 191.4 → 24.0 | 167.9 → 42.7 |
| max | 191.4 → 0.2 | 167.9 → 1.5 | 191.4 → 24.0 | 167.8 → 42.9 |
| bit_and | 191.4 → 0.2 | 168.2 → 1.5 | 191.4 → 24.0 | 167.5 → 41.8 |
| bit_or | 191.4 → 0.2 | 167.6 → 1.5 | 191.4 → 24.0 | 167.5 → 42.0 |
| bool_and | 0.3 → 0.1 | 2.7 → 1.3 | 28.2 → 20.0 | 34.1 → 26.0 |
| bool_or | 0.3 → 0.1 | 2.7 → 1.3 | 28.2 → 20.0 | 35.4 → 26.9 |
| approx_distinct | 191.4 → 31.5 | 170.5 → 41.5 | 191.4 → 62.5 | 168.0 →
88.5 |
## Are there any user-facing changes?
- **New public API, not breaking.**
- `AggregateUDFImpl::distinct_handling()` has a default of
`DistinctHandling::Honored`, so existing implementations compile and behave as
before.
- `DistinctHandling` is new and `#[non_exhaustive]`.
## LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated content
has been manually reviewed.
--
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]
