Perf: Take the comparison-free fast path in the grouped first/last accumulator when input is pre-ordered - #25050
Merged
comphead merged 5 commits intoSep 10, 2026
Conversation
…tor when input is pre-ordered 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 — but the grouped path never did: create_groups_accumulator dropped it, and FirstLastGroupsAccumulator ran the full per-row lexicographic tournament, per-winner ordering-key boxing, and cross-batch compare_rows on every batch even though the winner was already determined by position. Thread is_input_pre_ordered through create_groups_accumulator into FirstLastGroupsAccumulator, and when set, replace the tournament in update_batch with a single pass over group_indices: the first (FIRST_VALUE) or last (LAST_VALUE) qualifying row of each group in the batch wins, 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 winning row's ordering values are still materialized into the partial state: the final aggregation 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. Closes apache#24771
github-actions
Bot
added
sqllogictest
Contributor
There was a problem hiding this comment.
🟡 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_orderedthrough grouped accumulator creation and add a fast path inFirstLastGroupsAccumulator::update_batchto 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.
zhuqi-lucas
changed the title
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25050 +/- ##
========================================
Coverage 81.74% 81.74%
========================================
Files 1128 1128
Lines 416644 417152 +508
Branches 416644 417152 +508
========================================
+ Hits 340592 341016 +424
- Misses 55995 56003 +8
- Partials 20057 20133 +76 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
zhuqi-lucas
requested review from
2010YOUY01,
adriangb,
alamb and
jayzhan211
September 8, 2026 07:32
jayzhan211
reviewed
Sep 8, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
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.
…ping set With grouping sets 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, not the aggregate's own ORDER BY. Proving the prefixed requirement there marked first/last aggregates as pre-ordered when their groups were not, which the new grouped fast path (and the single-group accumulator through GroupsAccumulatorAdapter, latently on main) turned into wrong results. Suggested by @jayzhan211 in review, with the ROLLUP reproduction now covered in first_last_ordered.slt alongside a grouping-set case that legitimately keeps the fast path.
Contributor
Author
|
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 Added two cases to
Also worth noting for reviewers: as you said, the single-group accumulator consumed this flag through |
comphead
reviewed
Sep 9, 2026
comphead
left a comment
comphead
left a comment
Contributor
There was a problem hiding this comment.
Thanks @zhuqi-lucas I had a quick LLM review first
Blockers
-
LAST_VALUEtie-break. The fast path takes the last row of a tie, the tournament takes the first (verified: tournament10, fast12within a batch,10vs20across batches). At SQL level, addingWITH ORDER (ts ASC)to a table flipsLAST_VALUE(px ORDER BY ts)from10,20to11,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. -
Split out the
is_single()gate, or at least test it where it bites. It fixes a wrong-results bug that exists onmaintoday viaGroupsAccumulatorAdapter(Boolean is not ingroups_accumulator_supported). Add:- a Boolean-valued
ROLLUPSLT case (the currentINT dcase is already correct onmain, so it only covers the new fast path) - a
update_aggr_exprs.rsunit test asserting the flag is not set under grouping sets - ideally a separate PR so it can be backported
- a Boolean-valued
Should fix in this PR
-
Collect winners in O(rows), not O(
total_num_groups). Pushgroup_idxinto a reusableVecfield 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 raisingtotal_num_groupsto 1M takes it from 12.71 to 414.66 us/batch and collapses the speedup to 1.04x. Both regimes set the flag. -
Extend
benches/first_last.rs. Add apre_orderedflag toprepare_typed_groups_accumulatorviawith_beneficial_ordering(true). Sweeptotal_num_groups, rows per group, null density, filter on and off, and one bytes or nested type. -
Give the fuzzer a path that actually reaches the fast path. Today the generated
ORDER BYtakesmin(12, 43)columns with random directions against sort keys of at most 3, so it never triggers. Draw the aggregateORDER BYfrom the dataset'ssort_keysand compare against the non-hinted plan. Item 1 must be settled first, since the baseline uses the tournament tie-break and comparison is exact.
Contributor
Author
|
Thanks @comphead — all five addressed; per-item notes:
|
…zz the fast path - update_batch_pre_ordered records touched groups on the scoreboard's false-to-true transition and visits/clears only those, instead of scanning total_num_groups per batch; a dirty flag handles merge_batch (spill replay) interleaving on the same instance, which leaves winners' bits set. Covered by a sparse-1M-groups test and a merge-interleave test. - Tie handling is documented as a user-visible behavior change in the PR description and pinned by a tied-keys slt case. - The grouping-sets guard gets rule-level unit tests (flag set for a single grouping set, not set under grouping sets) and a Boolean ROLLUP slt case that exercises the GroupsAccumulatorAdapter path, which consumed the flag on main already. - The aggregation fuzzer draws first/last ORDER BYs from a prefix of the dataset's sort keys half the time, so sorted datasets exercise the fast path against the unsorted dataset's comparing path; the argument is the last ORDER BY column so ties carry equal outputs and the documented tie-break difference cannot produce false mismatches. Benchmarks for the pre-ordered path land in a separate bench-only PR.
Contributor
|
run benchmark tpch tpcds |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing first-last-preordered-fastpath (184bb56) to a5c809f (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing first-last-preordered-fastpath (184bb56) to a5c809f (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing first-last-preordered-fastpath (184bb56) to a5c809f (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing first-last-preordered-fastpath (184bb56) to a5c809f (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
Contributor
|
For TPCH it is consistently faster, mixed results for TPCDS which might indicate just noise. |
Contributor
|
run benchmark tpcds |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing first-last-preordered-fastpath (184bb56) to a5c809f (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing first-last-preordered-fastpath (184bb56) to a5c809f (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
Contributor
Author
Thanks @comphead for benchmark, this PR should not affect tpch and tpcds, it may be noise. |
jayzhan211
approved these changes
Sep 10, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
Thanks @zhuqi-lucas, LGTM!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #24771.
Rationale for this change
OptimizeAggregateOrderalready proves, viawith_beneficial_ordering, when the inputordering satisfies the group-by prefix followed by a
first_value/last_valueorderingrequirement, and the single-group accumulators consume that flag (
first_last.rs,get_first_idx/get_last_idx). The grouped path never did:create_groups_accumulatordropped the flag, and
FirstLastGroupsAccumulatorran the full per-row lexicographictournament, per-winner ordering-key materialization, and cross-batch
compare_rowsonevery batch even though the winner was already determined by position.
The motivating workload (a materialized
GROUP BY date, tickerwith threeLAST_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 insideget_filtered_extreme_of_each_group— all of it avoidable comparisons. The sameordering evidence is cashed twice by the planner but only once by the executor.
What changes are included in this PR?
is_input_pre_orderedis threaded throughcreate_groups_accumulatorintoFirstLastGroupsAccumulator.update_batchtakes a comparison-free path: one pass overgroup_indicesrecords 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
LexicographicalComparatoris built andcompare_rowsnever runs on this path.stage merges states from partitions whose relative arrival order is not guaranteed, so
merge_batchkeeps comparing and is intentionally untouched, as areconvert_to_stateand the skip-partialis_sethandling.(
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 keepsthe 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?
incl. null predicates/IGNORE NULLS/RESPECT NULLS with a null winner), explicit
expected values,
EmitTo::First(n)mid-stream draining with index shifting, and apartial→final merge in both arrival orders proving the emitted state carries the
winning ordering keys.
first_last_ordered.slt: requirement proven through declaredorderings, through a projection-alias equivalence, and through a filter-induced
constant; FILTER interaction; and the reverse branch (FIRST over DESC).
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 physicallyfirst for
FIRST_VALUE), matching the single-group pre-ordered accumulatorand
Iterator::max_by. The grouped tournament path keeps the first-seenrow of a tie (its comparisons are strict). Which row of a tie wins is
unspecified, but declaring
WITH ORDERon a table can now change which one aquery returns. Pinned by a dedicated tied-keys case in
first_last_ordered.slt; release notes should carry one line for this.Grouping sets
OptimizeAggregateOrderno longer uses the GROUP BY prefix to prove theaggregate'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: thesingle-group accumulators consumed the flag through
GroupsAccumulatorAdapterfor types outside
groups_accumulator_supported(e.g. Boolean). Covered byROLLUP 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_orderedareO(groups touched by the batch), not O(
total_num_groups): touched groupindices are recorded on the false-to-true scoreboard transition and only those
slots are visited and cleared. A spill replay interleaves
merge_batch(whichleaves scoreboard bits set) with
update_batchon the same accumulatorinstance, 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.