tohuya6 opened a new pull request, #23682:
URL: https://github.com/apache/datafusion/pull/23682
## Which issue does this PR close?
Closes #23602.
Part of the epic #23600.
> **Depends on #23628** (Closes #23601). This rule produces a `Struct`-typed
`first_value`
> argument, which only takes the fast `GroupsAccumulator` path once #23628
merges. On current
> `main` that fast path doesn't exist for `Struct`, so enabling the flag is
a regression — the
> config therefore defaults to `false`, and the speedup is not observable
until #23628 lands.
## Rationale for this change
`SELECT ..., first_value(a ORDER BY o), first_value(b ORDER BY o), ... GROUP
BY p` plans one
independent `first_value` / `last_value` aggregate per selected column. When
those peers all
share the same `ORDER BY`, they pick the *same* winning row for every
column, so the work is
duplicated across N expressions even though a single pass already determines
the winner for all.
Two costs on wide "one row per group" payloads:
- **Runtime.** Each `first_value` runs its own argmax pass over the input,
comparing the
`ORDER BY` key row by row. Selecting N winner columns is N independent
scans of the same rows,
so cost scales linearly with the number of peer expressions.
- **State.** Each aggregate holds its own per-group accumulator slot, so the
plan keeps N slots
per group — on wide payloads, N copies of the winner's columns retained
for every group key.
## What changes are included in this PR?
Adds a logical optimizer rule `CoalesceFirstLast`
(`datafusion/optimizer/src/coalesce_first_last.rs`)
that runs bottom-up over `Aggregate` nodes. Within a single `Aggregate`, it
groups peer
`first_value` / `last_value` expressions by `(function, ORDER BY, null
treatment)` and, for any
bucket with two or more members, rewrites them into one struct-valued
aggregate whose argument is
a `named_struct` of the original values. A `Projection` on top unpacks the
struct back into the
original columns with `get_field`, so the rewrite is transparent to
everything above the `Aggregate`:
```text
Aggregate: groupBy=[[p]], aggr=[[first_value(a ORDER BY o), first_value(b
ORDER BY o)]]
```
becomes
```text
Projection: p, get_field(s, 'c0') AS first_value(a ...), get_field(s, 'c1')
AS first_value(b ...)
Aggregate: groupBy=[[p]], aggr=[[first_value(named_struct('c0', a, 'c1',
b) ORDER BY o) AS s]]
```
- Coalesces only within a bucket keyed by `(function, ORDER BY, null
treatment)`; buckets with
fewer than two members are left unchanged.
- Bails on `DISTINCT` / `FILTER` / `IGNORE NULLS`, grouping-set aggregates
(`GROUPING SETS` / `ROLLUP` / `CUBE`), and expressions without an `ORDER
BY`.
- No-ops when `named_struct` / `get_field` are not in the function registry.
- Gated behind a new config `optimizer.enable_coalesce_first_last`, default
`false`.
- Registers the new config/rule in the generated artifacts: `configs.md`,
`information_schema.slt`, `optimizer_rule_reference.md`, and the `EXPLAIN
VERBOSE` list in
`explain.slt`.
## Are these changes tested?
Yes.
- Unit tests in `coalesce_first_last.rs` asserting the rewritten plan shape
(peer `first_value`
collapsed into one `named_struct` aggregate + `get_field` projection; same
for `last_value`; a
non-coalesced aggregate such as `count` preserved alongside), plus no-op
cases: a single
`first_value`, mixed `first_value` + `last_value`, `DISTINCT`, a
grouping-set aggregate, flag
disabled, missing registry.
- An end-to-end sqllogictest (`coalesce_first_last.slt`) that runs the query
with the flag off and
on over multi-group data containing NULLs and asserts identical results,
plus an `EXPLAIN`
pinning the rewritten logical and physical plans.
- `cargo fmt`, `cargo clippy --all-targets --all-features -- -D warnings`,
the full sqllogictest
suite, and the config/rule doc-generation checks all pass.
## Are there any user-facing changes?
Adds one new configuration option,
`datafusion.optimizer.enable_coalesce_first_last`, defaulting
to `false`. No SQL or API surface changes: with the flag off the plan is
unchanged, and with it on
the rewrite is result-preserving, so existing queries and previously-passing
tests behave identically.
## Follow-ups
- Once #23628 (Closes #23601) merges, consider flipping the default after
benchmarking across
payload widths and group cardinalities.
- #23603 (logical rewrite `Filter(row_number() = 1) →
Aggregate(FIRST_VALUE(... ORDER BY))`)
composes with this rule: the peer `FIRST_VALUE` aggregates it emits would
be coalesced here.
--
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]
