Skip to content

feat(optimizer): coalesce peer first_value / last_value into a single struct aggregate#23682

Open
tohuya6 wants to merge 3 commits into
apache:mainfrom
tohuya6:coalesce-first-last
Open

feat(optimizer): coalesce peer first_value / last_value into a single struct aggregate#23682
tohuya6 wants to merge 3 commits into
apache:mainfrom
tohuya6:coalesce-first-last

Conversation

@tohuya6

@tohuya6 tohuya6 commented Jul 18, 2026

Copy link
Copy Markdown

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:

Aggregate: groupBy=[[p]], aggr=[[first_value(a ORDER BY o), first_value(b ORDER BY o)]]

becomes

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

… struct aggregate

Add the CoalesceFirstLast logical optimizer rule, gated behind
optimizer.enable_coalesce_first_last (default off). Peer first_value /
last_value aggregates that share an ORDER BY key are rewritten into one
struct-valued aggregate plus a get_field projection, so the input is
scanned once instead of once per expression and one per-group state slot
is kept instead of N.

Skips DISTINCT / FILTER / IGNORE NULLS, expressions without an ORDER BY,
and grouping-set aggregates, and no-ops when named_struct / get_field are
not registered.

Registers the config and rule in the generated artifacts that enumerate
every setting and optimizer rule: configs.md, information_schema.slt,
optimizer_rule_reference.md, and the EXPLAIN VERBOSE rule list in
explain.slt.

Closes apache#23602. Part of apache#23600; the struct fast path depends on apache#23628.
@github-actions github-actions Bot added documentation Improvements or additions to documentation optimizer Optimizer rules core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) common Related to common crate labels Jul 18, 2026
@tohuya6
tohuya6 marked this pull request as ready for review July 18, 2026 03:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

common Related to common crate core Core DataFusion crate documentation Improvements or additions to documentation optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt)

Projects

None yet

1 participant