Skip to content

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

Closed
tohuya6 wants to merge 2 commits into
apache:mainfrom
tohuya6:coalesce-first-last-23602
Closed

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

Conversation

@tohuya6

@tohuya6 tohuya6 commented Jul 17, 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 and on 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, 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

@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 17, 2026
tohuya6 added 2 commits July 17, 2026 20:17
… 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.

Closes apache#23602. Part of apache#23600; the struct fast path depends on apache#23628.
…cs and tests

Add the new config and rule to 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.
@tohuya6
tohuya6 force-pushed the coalesce-first-last-23602 branch from e41f217 to d3f9311 Compare July 18, 2026 00:17
@tohuya6 tohuya6 closed this Jul 18, 2026
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