namanjain24-sudo opened a new pull request, #25461:
URL: https://github.com/apache/datafusion/pull/25461
## Which issue does this PR close?
- Closes #25417.
## Rationale for this change
`is_single_distinct_agg` requires that every distinct argument in the
aggregation is the same expression. It checks that by collecting the arguments
into a set and requiring the set to hold exactly one element:
```rust
if *distinct {
for e in args {
fields_set.insert(e);
}
```
A call that passes that expression more than once satisfies the check,
because its arguments collapse to a single entry. `corr(DISTINCT x, x)` is such
a call, and the rewrite below then assumed one argument per distinct aggregate
and returned an internal error.
This is also why the neighbouring cases in the issue are unaffected:
`corr(DISTINCT x, y)` and `corr(DISTINCT x, x + 0.0)` put two entries in the
set, so the rule bails out early and leaves the plan alone.
## What changes are included in this PR?
The distinct branch keeps the call's arity, repeating the alias that the
inner group by produces, so `corr(DISTINCT x, x)` becomes a group by `x` with
`corr(alias1, alias1)` above it.
That is sound because of the check above. Every distinct argument is the
same expression, so the distinct argument tuples the call aggregates are
exactly the distinct values of that expression, which is what grouping by it
produces.
The assertion stays, now stating the invariant the rewrite actually relies
on rather than an arity the rule does not enforce. It also covers an empty
argument list, which would otherwise panic in `args.swap_remove(0)` instead of
reporting an internal error.
## What is the testing strategy for this PR?
Rule tests in `single_distinct_to_groupby.rs`:
- `single_distinct_repeated_arg_and_groupby`: `corr(DISTINCT b, b)` is
rewritten, and the plan keeps both arguments and the original output column
name.
- `single_distinct_two_args_and_groupby`: `corr(DISTINCT b, c)` is still
left alone.
Two queries in `aggregate.slt` cover it end to end. One is `covar_samp`,
chosen because its value depends on the de-duplication rather than only on the
absence of the error: over the same input it returns `0.5` with `DISTINCT` and
`0.333333333333` without, so the expected value pins the semantics.
Measured, rather than assumed:
| check | result |
| --- | --- |
| both halves of the fix reverted | the two `aggregate.slt` queries and
`single_distinct_repeated_arg_and_groupby` fail with an `Assertion failed:
args.len() == 1` internal error, which is the failure in the issue |
| the rewrite against a reference the rule does not touch | the same
aggregates over a manually de-duplicated input return the same values |
`cargo test -p datafusion-optimizer` (915 tests), the sqllogictest suite
(521 files), and `./ci/scripts/rust_clippy.sh` pass.
## Are there any user-facing changes?
Queries that failed with an internal error now run. No plan that the rule
rewrote before is rewritten differently: a single argument call produces the
same outer aggregate as it did, and the only calls that reach this branch with
more than one argument are the ones that previously returned the error.
--
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]
