Perf: Take the comparison-free fast path in the grouped first/last accumulator when input is pre-ordered

[zhuqi-lucas]

Which issue does this PR close?

Closes #24771.

Rationale for this change

OptimizeAggregateOrder already proves, via with_beneficial_ordering, when the input
ordering satisfies the group-by prefix followed by a first_value/last_value ordering
requirement, and the single-group accumulators consume that flag (first_last.rs,
get_first_idx/get_last_idx). The grouped path never did: create_groups_accumulator
dropped the flag, and FirstLastGroupsAccumulator ran the full per-row lexicographic
tournament, per-winner ordering-key materialization, and cross-batch compare_rows on
every batch even though the winner was already determined by position.

The motivating workload (a materialized GROUP BY date, ticker with three
LAST_VALUE(... ORDER BY ts, seq) over an NBBO table whose file sort order is exactly
(ticker, ts, seq), ~1.5B rows/day, ~1.8M groups) spends ~20% of on-CPU time inside
get_filtered_extreme_of_each_group — all of it avoidable comparisons. The same
ordering evidence is cashed twice by the planner but only once by the executor.

What changes are included in this PR?

• is_input_pre_ordered is threaded through create_groups_accumulator into
  FirstLastGroupsAccumulator.
• When set, update_batch takes a comparison-free path: one pass over group_indices
  records each group's first (FIRST_VALUE) or last (LAST_VALUE) qualifying row; later
  batches unconditionally overwrite for LAST_VALUE and never overwrite for FIRST_VALUE.
  No LexicographicalComparator is built and compare_rows never runs on this path.
• The winner's ordering values are still materialized into the partial state: the final
  stage merges states from partitions whose relative arrival order is not guaranteed, so
  merge_batch keeps comparing and is intentionally untouched, as are
  convert_to_state and the skip-partial is_set handling.
• Drive-by: three stale comments still referring to the field's old name
  (min_of_each_group_buf) are updated.

Tie semantics note: among rows whose ordering keys compare equal, this path picks
the physically last qualifying row for LAST_VALUE (first for FIRST_VALUE), matching the
single-group pre-ordered accumulator and Iterator::max_by. The tournament path keeps
the first-seen row of a tie (its comparisons are strict), so the two paths may pick
different — equally valid — rows on tied keys. Documented on the method and pinned by a
dedicated test.

Are these changes tested?

• 9 new unit tests: fast-vs-tournament full-state equivalence (LAST/FIRST/DESC/FILTER
  incl. null predicates/IGNORE NULLS/RESPECT NULLS with a null winner), explicit
  expected values, EmitTo::First(n) mid-stream draining with index shifting, and a
  partial→final merge in both arrival orders proving the emitted state carries the
  winning ordering keys.
• A new end-to-end first_last_ordered.slt: requirement proven through declared
  orderings, through a projection-alias equivalence, and through a filter-induced
  constant; FILTER interaction; and the reverse branch (FIRST over DESC).
• Existing suites pass: 215 crate tests, aggregate.slt, first_last_nested.slt,
  group_by.slt, distinct_on.slt, array_agg.slt, subquery_sort.slt.

Are there any user-facing changes?

No API changes. Queries whose input ordering already satisfies a grouped
first_value/last_value requirement get faster; on the motivating workload the
aggregation stage improved by ~20% end-to-end wall time.

User-visible behavior change: tie handling

Among rows whose ordering keys compare equal, the pre-ordered fast path picks
the physically last qualifying row for LAST_VALUE (and the physically
first for FIRST_VALUE), matching the single-group pre-ordered accumulator
and Iterator::max_by. The grouped tournament path keeps the first-seen
row of a tie (its comparisons are strict). Which row of a tie wins is
unspecified, but declaring WITH ORDER on a table can now change which one a
query returns. Pinned by a dedicated tied-keys case in
first_last_ordered.slt; release notes should carry one line for this.

Grouping sets

OptimizeAggregateOrder no longer uses the GROUP BY prefix to prove the
aggregate's requirement when grouping sets are present: the stream feeds the
same rows once per grouping set, and within a coarser set's group the rows
follow the full group-by prefix rather than the aggregate's own ORDER BY.
This also fixes a latent wrong-results path that existed on main: the
single-group accumulators consumed the flag through GroupsAccumulatorAdapter
for types outside groups_accumulator_supported (e.g. Boolean). Covered by
ROLLUP slt cases (including a Boolean one that exercises the adapter path) and
two rule-level unit tests.

Fast-path complexity

Winner collection and the scoreboard reset in update_batch_pre_ordered are
O(groups touched by the batch), not O(total_num_groups): touched group
indices are recorded on the false-to-true scoreboard transition and only those
slots are visited and cleared. A spill replay interleaves merge_batch (which
leaves scoreboard bits set) with update_batch on the same accumulator
instance, so a dirty flag triggers one full reset in that case; covered by a
merge-interleave unit test and a sparse-1M-groups test.

Fuzzing

The aggregation fuzzer now draws first/last ORDER BYs from a prefix of the
dataset's sort keys half of the time, so sorted datasets exercise the fast
path while the unsorted dataset runs the comparing path on the same query.
The aggregate argument is the last ORDER BY column so ties carry equal output
values and the tie-break difference above cannot cause false mismatches.

Benchmarks for the pre-ordered path (group-count / rows-per-group / null /
filter sweeps) land in a separate bench-only PR so before/after numbers can be
taken on main.

[Copilot]

🟡 Changes recommended

The newly added sqllogictest file appears to reference window_2.csv using an incorrect relative LOCATION path, likely causing test failures.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Improves performance of grouped first_value/last_value aggregates by honoring optimizer-proven beneficial input ordering, enabling a comparison-free update path while preserving correct partial-state merge behavior.

Changes:

• Thread is_input_pre_ordered through grouped accumulator creation and add a fast path in FirstLastGroupsAccumulator::update_batch to avoid per-row comparisons when input is pre-ordered.
• Add extensive unit tests covering fast-path correctness (including FILTER/NULLS/ties/partial emits) and partial→final merge correctness.
• Add an end-to-end sqllogictest (first_last_ordered.slt) to validate optimizer→executor propagation of the pre-ordered optimization.

File summaries

File	Description
datafusion/sqllogictest/test_files/first_last_ordered.slt	Adds end-to-end coverage for grouped ordered first/last fast path via declared orderings and optimizer-driven proofs.
datafusion/functions-aggregate/src/first_last.rs	Threads the pre-ordered flag into grouped accumulators and implements a comparison-free update_batch path with added unit tests.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 83.28358% with 56 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.74%. Comparing base (a5c809f) to head (768fee4).

Files with missing lines	Patch %	Lines
datafusion/functions-aggregate/src/first_last.rs	82.91%	1 Missing and 54 partials ⚠️
...fusion/physical-optimizer/src/update_aggr_exprs.rs	92.30%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25050      +/-   ##
==========================================
- Coverage   81.74%   81.74%   -0.01%     
==========================================
  Files        1128     1128              
  Lines      416644   416967     +323     
  Branches   416644   416967     +323     
==========================================
+ Hits       340592   340849     +257     
- Misses      55995    56003       +8     
- Partials    20057    20115      +58     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[jayzhan211]

Thanks @zhuqi-lucas , here is a suggestion:

The fast path is sound for a single grouping set, but OptimizeAggregateOrder sets the flag for grouping sets too, and there the guarantee does not hold. group_expr().input_exprs() for ROLLUP(a, b) is [a, b], so an input ordered (a, b, c) satisfies [a, b, c ASC] and with_beneficial_ordering(true) is applied. GroupedHashAggregateStream then calls update_batch once per grouping set, and for the (a) set the rows of group a=1 arrive in (b, c) order, not c order. The tournament path tolerated that; the positional path does not.

Repro on this branch (target_partitions=1, CSV declared WITH ORDER (a ASC, b ASC, c ASC), rows (1,1,5,100) (1,1,6,101) (1,2,1,200) (1,2,2,201)):

SELECT a, b, FIRST_VALUE(d ORDER BY c), LAST_VALUE(d ORDER BY c)
FROM t GROUP BY ROLLUP(a, b) ORDER BY a NULLS FIRST, b NULLS FIRST;

expected            actual
NULL NULL 200 101   NULL NULL 100 201
1    NULL 200 101   1    NULL 100 201
1    1    100 101   1    1    100 101
1    2    200 201   1    2    200 201

main is fine because the grouped accumulator ignored the flag. The single-group accumulator (used through GroupsAccumulatorAdapter for types outside groups_accumulator_supported) has had the same latent issue, so I'd fix it in the rule rather than in the accumulator. With grouping sets, drop the group-by prefix and require the aggregate's own ORDER BY to be satisfied on its own; a globally sorted input stays sorted in every subsequence, so the fast path is still taken in that case:

--- a/datafusion/physical-optimizer/src/update_aggr_exprs.rs
+++ b/datafusion/physical-optimizer/src/update_aggr_exprs.rs
@@
-                let groupby_exprs = aggr_exec.group_expr().input_exprs();
-                // If the existing ordering satisfies a prefix of the GROUP BY
-                // expressions, prefix requirements with this section. In this
-                // case, aggregation will work more efficiently.
-                let indices = get_ordered_partition_by_indices(&groupby_exprs, input)?;
-                let requirement = indices
-                    .iter()
-                    .map(|&idx| {
-                        PhysicalSortRequirement::new(
-                            Arc::clone(&groupby_exprs[idx]),
-                            None,
-                        )
-                    })
-                    .collect::<Vec<_>>();
+                // With grouping sets the same rows are fed once per grouping
+                // set, and a coarser set's group is only ordered by the
+                // aggregate's ORDER BY if the input is ordered by it without
+                // any group-by prefix. So only use the prefix for a single
+                // grouping set.
+                let requirement = if aggr_exec.group_expr().is_single() {
+                    let groupby_exprs = aggr_exec.group_expr().input_exprs();
+                    let indices =
+                        get_ordered_partition_by_indices(&groupby_exprs, input)?;
+                    indices
+                        .iter()
+                        .map(|&idx| {
+                            PhysicalSortRequirement::new(
+                                Arc::clone(&groupby_exprs[idx]),
+                                None,
+                            )
+                        })
+                        .collect::<Vec<_>>()
+                } else {
+                    vec![]
+                };

Please also add the ROLLUP query above to first_last_ordered.slt so this stays covered.

[zhuqi-lucas]

Thank you @jayzhan211 — confirmed, and great catch. I reproduced your ROLLUP example on this branch before touching anything (wrong results exactly as you showed), then applied the fix in OptimizeAggregateOrder as you suggested: the group-by prefix is only used for a single grouping set, so with grouping sets the aggregate's own ORDER BY must be satisfied on its own. Your repro now returns the expected rows.

Added two cases to first_last_ordered.slt:

• your ROLLUP repro (comparisons required, correct results), and
• a ROLLUP where the aggregate ORDER BY leads the input ordering, which stays on the fast path since a globally sorted input is sorted within every group of every grouping set.

Also worth noting for reviewers: as you said, the single-group accumulator consumed this flag through GroupsAccumulatorAdapter on main already, so the rule-level guard fixes that latent path too, not just the new grouped fast path.

[comphead]

Thanks @zhuqi-lucas I had a quick LLM review first

Blockers

1. LAST_VALUE tie-break. The fast path takes the last row of a tie, the tournament takes the first (verified: tournament 10, fast 12 within a batch, 10 vs 20 across batches). At SQL level, adding WITH ORDER (ts ASC) to a table flips LAST_VALUE(px ORDER BY ts) from 10,20 to 11,21. Not fixable without comparisons, so: document it as a user-visible behavior change in the PR description and release notes, not just a code comment. Add an SLT case with a tied ordering key.

2. Split out the is_single() gate, or at least test it where it bites. It fixes a wrong-results bug that exists on main today via GroupsAccumulatorAdapter (Boolean is not in groups_accumulator_supported). Add:

   • a Boolean-valued ROLLUP SLT case (the current INT d case is already correct on main, so it only covers the new fast path)
   • a update_aggr_exprs.rs unit test asserting the flag is not set under grouping sets
   • ideally a separate PR so it can be backported

Should fix in this PR

3. Collect winners in O(rows), not O(total_num_groups). Push group_idx into a reusable Vec field on the false-to-true transition of the scoreboard bit, and clear only those slots on reset. Measured: at 8192 rows and 64 touched groups the fast path is 1.99x, but raising total_num_groups to 1M takes it from 12.71 to 414.66 us/batch and collapses the speedup to 1.04x. Both regimes set the flag.

4. Extend benches/first_last.rs. Add a pre_ordered flag to prepare_typed_groups_accumulator via with_beneficial_ordering(true). Sweep total_num_groups, rows per group, null density, filter on and off, and one bytes or nested type.

5. Give the fuzzer a path that actually reaches the fast path. Today the generated ORDER BY takes min(12, 43) columns with random directions against sort keys of at most 3, so it never triggers. Draw the aggregate ORDER BY from the dataset's sort_keys and compare against the non-hinted plan. Item 1 must be settled first, since the baseline uses the tournament tie-break and comparison is exact.