This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25050-57f861942da849d338b660a029a3d58bcbcf20a0 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 4cee381898d0de0f98c9a74d9489203f8fdc0c52 Author: Qi Zhu <[email protected]> AuthorDate: Thu Sep 10 20:13:51 2026 +0000 Perf: Take the comparison-free fast path in the grouped first/last accumulator when input is pre-ordered (#25050) ## 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. --- .../fuzz_cases/aggregation_fuzzer/query_builder.rs | 45 +- datafusion/core/tests/physical_optimizer/mod.rs | 1 + .../tests/physical_optimizer/update_aggr_exprs.rs | 144 +++++ datafusion/functions-aggregate/src/first_last.rs | 592 ++++++++++++++++++++- .../physical-optimizer/src/update_aggr_exprs.rs | 37 +- .../sqllogictest/test_files/first_last_ordered.slt | 183 +++++++ 6 files changed, 976 insertions(+), 26 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/query_builder.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/query_builder.rs index d32078ec63..2b2d2483d7 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/query_builder.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/query_builder.rs @@ -293,20 +293,51 @@ impl QueryBuilder { while aggregate_functions.len() < num_aggregate_functions { let idx = rng.random_range(0..self.aggregate_functions.len()); let (function_name, is_distinct) = &self.aggregate_functions[idx]; - let argument = self.random_argument(); let alias = format!("col{alias_gen}"); let distinct = if *is_distinct { "DISTINCT " } else { "" }; alias_gen += 1; - let (order_by, null_opt) = if function_name.eq("first_value") + let (argument, order_by, null_opt) = if function_name.eq("first_value") || function_name.eq("last_value") { - ( - self.order_by(&order_by_black_list), /* Among the order by columns, at most one group by column can be included to avoid all order by column values being identical */ - self.null_opt(), - ) + if !self.dataset_sort_keys.is_empty() && rng.random_bool(0.5) { + // Draw the ORDER BY from a prefix of the dataset's own + // sort keys, so on the sorted datasets the plan satisfies + // the aggregate's requirement and the pre-ordered + // first/last fast path is exercised; the unsorted dataset + // runs the same query through the comparing path, giving + // a hinted-vs-unhinted comparison for free. A random + // ORDER BY never does this: it samples up to 12 columns + // with random directions against at most 3 sort keys. + // + // The argument is the last ORDER BY column on purpose: + // rows that tie on the whole ORDER BY then carry equal + // output values, so the documented tie-break difference + // between the fast path (physically last/first) and the + // tournament (first seen) cannot produce false + // mismatches. + let keys = &self.dataset_sort_keys + [rng.random_range(0..self.dataset_sort_keys.len())]; + let prefix_len = rng.random_range(1..=keys.len()); + let prefix = &keys[..prefix_len]; + let order_by = format!( + " order by {}", + prefix + .iter() + .map(|c| format!("{c} ASC")) + .collect::<Vec<_>>() + .join(", ") + ); + (prefix[prefix_len - 1].clone(), order_by, self.null_opt()) + } else { + ( + self.random_argument(), + self.order_by(&order_by_black_list), /* Among the order by columns, at most one group by column can be included to avoid all order by column values being identical */ + self.null_opt(), + ) + } } else { - ("".to_string(), "".to_string()) + (self.random_argument(), "".to_string(), "".to_string()) }; let function = format!( diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs index f3b2884dab..f89d86074b 100644 --- a/datafusion/core/tests/physical_optimizer/mod.rs +++ b/datafusion/core/tests/physical_optimizer/mod.rs @@ -38,6 +38,7 @@ mod replace_with_order_preserving_variants; mod sanity_checker; #[expect(clippy::needless_pass_by_value)] mod test_utils; +mod update_aggr_exprs; mod window_optimize; mod window_topn; diff --git a/datafusion/core/tests/physical_optimizer/update_aggr_exprs.rs b/datafusion/core/tests/physical_optimizer/update_aggr_exprs.rs new file mode 100644 index 0000000000..d2cdbd1761 --- /dev/null +++ b/datafusion/core/tests/physical_optimizer/update_aggr_exprs.rs @@ -0,0 +1,144 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for the [`OptimizeAggregateOrder`] rule. + +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_functions_aggregate::first_last::first_value_udaf; +use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::update_aggr_exprs::OptimizeAggregateOrder; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; + +use crate::physical_optimizer::test_utils::parquet_exec_with_sort; + +fn abc_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + ])) +} + +/// An input whose declared ordering is `(a, b, c)`. +fn ordered_input(schema: &SchemaRef) -> Result<Arc<dyn ExecutionPlan>> { + let ordering: LexOrdering = [ + PhysicalSortExpr::new_default(col("a", schema)?), + PhysicalSortExpr::new_default(col("b", schema)?), + PhysicalSortExpr::new_default(col("c", schema)?), + ] + .into(); + Ok(parquet_exec_with_sort(Arc::clone(schema), vec![ordering]) as _) +} + +/// `FIRST_VALUE(d ORDER BY c)`. +fn first_value_expr(schema: &SchemaRef) -> Result<Arc<AggregateFunctionExpr>> { + let order = PhysicalSortExpr::new_default(col("c", schema)?); + AggregateExprBuilder::new(first_value_udaf(), vec![col("d", schema)?]) + .schema(Arc::clone(schema)) + .alias("first_value(d)") + .order_by(vec![order]) + .build() + .map(Arc::new) +} + +fn optimized_flag(plan: Arc<dyn ExecutionPlan>) -> Result<String> { + let optimized = + OptimizeAggregateOrder::new().optimize(plan, &ConfigOptions::default())?; + let agg = optimized + .downcast_ref::<AggregateExec>() + .expect("optimizer keeps the AggregateExec at the root"); + Ok(format!("{:?}", agg.aggr_expr()[0].fun().inner())) +} + +/// With a plain GROUP BY, the input ordering `(a, b, c)` proves the group-by +/// prefix plus the aggregate's `ORDER BY c`, so the rule marks the aggregate +/// pre-ordered. +#[test] +fn single_group_by_gets_the_group_by_prefix() -> Result<()> { + let schema = abc_schema(); + let group_by = PhysicalGroupBy::new_single(vec![ + (col("a", &schema)?, "a".to_string()), + (col("b", &schema)?, "b".to_string()), + ]); + let agg = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![first_value_expr(&schema)?], + vec![None], + ordered_input(&schema)?, + Arc::clone(&schema), + )?; + let dbg = optimized_flag(Arc::new(agg))?; + assert!( + dbg.contains("is_input_pre_ordered: true"), + "expected the flag set for a single grouping set, got: {dbg}" + ); + Ok(()) +} + +/// With grouping sets the same rows are fed 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. The rule must not use the group-by prefix to +/// prove the requirement there. +#[test] +fn grouping_sets_do_not_get_the_group_by_prefix() -> Result<()> { + let schema = abc_schema(); + // ROLLUP(a, b): grouping sets (a, b), (a), (). + let group_by = PhysicalGroupBy::new( + vec![ + (col("a", &schema)?, "a".to_string()), + (col("b", &schema)?, "b".to_string()), + ], + vec![ + ( + datafusion_physical_expr::expressions::lit(1i32), + "a".to_string(), + ), + ( + datafusion_physical_expr::expressions::lit(1i32), + "b".to_string(), + ), + ], + vec![vec![false, false], vec![false, true], vec![true, true]], + true, + ); + let agg = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![first_value_expr(&schema)?], + vec![None], + ordered_input(&schema)?, + Arc::clone(&schema), + )?; + let dbg = optimized_flag(Arc::new(agg))?; + assert!( + dbg.contains("is_input_pre_ordered: false"), + "the flag must stay unset under grouping sets, got: {dbg}" + ); + Ok(()) +} diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index d36ea5d630..f2d6af5bcf 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -79,6 +79,7 @@ pub fn last_value(expression: Expr, order_by: Vec<SortExpr>) -> Expr { fn create_groups_accumulator_helper<S: ValueState + 'static>( args: &AccumulatorArgs, is_first: bool, + is_input_pre_ordered: bool, state: S, ) -> Result<Box<dyn GroupsAccumulator>> { let Some(ordering) = LexOrdering::new(args.order_bys.to_vec()) else { @@ -96,12 +97,14 @@ fn create_groups_accumulator_helper<S: ValueState + 'static>( args.ignore_nulls, &ordering_dtypes, is_first, + is_input_pre_ordered, )?)) } fn create_groups_accumulator( args: &AccumulatorArgs, is_first: bool, + is_input_pre_ordered: bool, function_name: &str, ) -> Result<Box<dyn GroupsAccumulator>> { let data_type = args.return_field.data_type(); @@ -111,6 +114,7 @@ fn create_groups_accumulator( create_groups_accumulator_helper( args, is_first, + is_input_pre_ordered, PrimitiveValueState::<$t>::new(data_type.clone()), ) }; @@ -168,6 +172,7 @@ fn create_groups_accumulator( | DataType::BinaryView => create_groups_accumulator_helper( args, is_first, + is_input_pre_ordered, BytesValueState::try_new(data_type.clone())?, ), @@ -185,6 +190,7 @@ fn create_groups_accumulator( | DataType::Map(_, _) => create_groups_accumulator_helper( args, is_first, + is_input_pre_ordered, GenericValueState::new(data_type.clone()), ), @@ -346,7 +352,7 @@ impl AggregateUDFImpl for FirstValue { &self, args: AccumulatorArgs, ) -> Result<Box<dyn GroupsAccumulator>> { - create_groups_accumulator(&args, true, self.name()) + create_groups_accumulator(&args, true, self.is_input_pre_ordered, self.name()) } fn with_beneficial_ordering( @@ -393,10 +399,25 @@ struct FirstLastGroupsAccumulator<S: ValueState> { // to avoid calling `ScalarValue::size_of_vec` by Self.size. size_of_orderings: usize, - // buffer for `get_filtered_extreme_of_each_group` - // filter_min_of_each_group_buf.0[group_idx] -> idx_in_val - // only valid if filter_min_of_each_group_buf.1[group_idx] == true + // Per-batch scoreboard shared by `get_filtered_extreme_of_each_group` + // and the pre-ordered fast path: + // extreme_of_each_group_buf.0[group_idx] -> idx_in_val + // only valid if extreme_of_each_group_buf.1[group_idx] == true extreme_of_each_group_buf: (Vec<usize>, BooleanBufferBuilder), + // Set by `get_filtered_extreme_of_each_group` (merge_batch / + // convert_to_state), which clears the scoreboard at its start but leaves + // its winners' bits set on return. `update_batch_pre_ordered` keeps the + // scoreboard all-false between its own batches, but the two can + // interleave on one instance (the aggregation stream calls merge_batch on + // the same accumulators when replaying spilled state), so the fast path + // does one full reset when this is set. + extreme_buf_dirty: bool, + // Batch-local list of groups touched by `update_batch_pre_ordered`, + // reused across batches. It lets winner collection and the scoreboard + // reset run in O(groups touched by the batch) instead of + // O(total_num_groups): with 1M allocated groups and 64 touched per + // batch, a full-scoreboard sweep per batch dominates the runtime. + touched_groups_buf: Vec<usize>, // =========== option ============ @@ -409,6 +430,11 @@ struct FirstLastGroupsAccumulator<S: ValueState> { sort_options: Vec<SortOptions>, // Ignore null values. ignore_nulls: bool, + // When `true` (set through `with_beneficial_ordering` by the + // `OptimizeAggregateOrder` physical-optimizer rule), the optimizer has + // proven that every group's rows already arrive in `ordering_req` order, + // and `update_batch` takes the comparison-free fast path. + is_input_pre_ordered: bool, default_orderings: Vec<ScalarValue>, } @@ -419,6 +445,7 @@ impl<S: ValueState> FirstLastGroupsAccumulator<S> { ignore_nulls: bool, ordering_dtypes: &[DataType], pick_first_in_group: bool, + is_input_pre_ordered: bool, ) -> Result<Self> { let default_orderings = ordering_dtypes .iter() @@ -437,7 +464,10 @@ impl<S: ValueState> FirstLastGroupsAccumulator<S> { is_sets: BooleanBufferBuilder::new(0), size_of_orderings: 0, extreme_of_each_group_buf: (Vec::new(), BooleanBufferBuilder::new(0)), + extreme_buf_dirty: false, + touched_groups_buf: Vec::new(), pick_first_in_group, + is_input_pre_ordered, }) } @@ -553,14 +583,14 @@ impl<S: ValueState> FirstLastGroupsAccumulator<S> { vals: &ArrayRef, is_set_arr: Option<&BooleanArray>, ) -> Result<Vec<(usize, usize)>> { - // Set all values in min_of_each_group_buf.1 to false. + // Set all values in extreme_of_each_group_buf.1 to false. self.extreme_of_each_group_buf.1.truncate(0); self.extreme_of_each_group_buf .1 .append_n(self.is_sets.len(), false); - // No need to call `clear` since `self.min_of_each_group_buf.0[group_idx]` - // is only valid when `self.min_of_each_group_buf.1[group_idx] == true`. + // No need to call `clear` since `self.extreme_of_each_group_buf.0[group_idx]` + // is only valid when `self.extreme_of_each_group_buf.1[group_idx] == true`. let comparator = { assert_eq!(orderings.len(), self.ordering_req.len()); @@ -614,6 +644,11 @@ impl<S: ValueState> FirstLastGroupsAccumulator<S> { } } + // Winners' bits stay set on return; tell the pre-ordered fast path + // that the scoreboard needs a reset before it can trust its all-false + // invariant again. + self.extreme_buf_dirty = true; + Ok(self .extreme_of_each_group_buf .0 @@ -623,6 +658,108 @@ impl<S: ValueState> FirstLastGroupsAccumulator<S> { .map(|(group_idx, idx_in_val)| (group_idx, *idx_in_val)) .collect::<Vec<_>>()) } + + /// Comparison-free `update_batch` for pre-ordered input. + /// + /// `is_input_pre_ordered` is only set (through `with_beneficial_ordering`, + /// by the `OptimizeAggregateOrder` physical-optimizer rule) after the + /// optimizer has proven that the input ordering satisfies the ordered + /// group-by prefix followed by this aggregate's own ordering requirement. + /// Under that guarantee each group's rows arrive in `ordering_req` order, + /// so: + /// - within a batch, the extreme row of a group is simply its first + /// (FIRST_VALUE) or last (LAST_VALUE) qualifying row, and + /// - across batches, a later batch's winner always beats an earlier one + /// for LAST_VALUE, and never does for FIRST_VALUE. + /// + /// No lexicographic comparator is built and `compare_rows` never runs. + /// The winner's ordering values are still materialized into + /// `self.orderings`: partial state must carry them, because the final + /// aggregation stage merges states coming from different partitions whose + /// relative order is not guaranteed, so `merge_batch` keeps comparing. + /// + /// Tie handling: among rows whose ordering keys compare equal, this path + /// picks the physically last qualifying row for LAST_VALUE (and the 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 instead (its comparisons are strict). Both are valid answers — + /// which row of a tie wins is unspecified — but results can differ on + /// tied keys. + fn update_batch_pre_ordered( + &mut self, + values_and_order_cols: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + ) -> Result<()> { + let vals = &values_and_order_cols[0]; + + // `extreme_of_each_group_buf.1` is sized by `resize_states` and kept + // all-false between batches: each batch records the groups it touched + // in `touched_groups_buf` and clears exactly those bits before + // returning. Neither the reset nor the winner collection may scan + // `total_num_groups` -- with 1M allocated groups and 64 touched, a + // full sweep turns a ~13us batch into ~400us and erases the win. + debug_assert!(self.touched_groups_buf.is_empty()); + if self.extreme_buf_dirty { + // A merge (spill replay) ran on this instance and left its + // winners' bits set; restore the all-false invariant once. + self.extreme_of_each_group_buf.1.truncate(0); + self.extreme_of_each_group_buf + .1 + .append_n(self.is_sets.len(), false); + self.extreme_buf_dirty = false; + } + + for (idx_in_val, &group_idx) in group_indices.iter().enumerate() { + // A row passes the FILTER clause only when the predicate is + // `true`; rows whose predicate evaluates to `null` are excluded. + let passed_filter = + opt_filter.is_none_or(|x| x.is_valid(idx_in_val) && x.value(idx_in_val)); + if !passed_filter { + continue; + } + if self.ignore_nulls && vals.is_null(idx_in_val) { + continue; + } + + let touched_this_batch = self.extreme_of_each_group_buf.1.get_bit(group_idx); + if self.pick_first_in_group + && (self.is_sets.get_bit(group_idx) || touched_this_batch) + { + // The first qualifying row wins; groups decided by an earlier + // batch (or earlier in this batch) never change again. + continue; + } + if !touched_this_batch { + self.extreme_of_each_group_buf.1.set_bit(group_idx, true); + self.touched_groups_buf.push(group_idx); + } + // For LAST_VALUE, later qualifying rows unconditionally overwrite. + self.extreme_of_each_group_buf.0[group_idx] = idx_in_val; + } + + let mut ordering_buf = Vec::with_capacity(self.ordering_req.len()); + // Take the buffer to appease the borrow checker; `update_state` + // needs `&mut self`. + let touched = std::mem::take(&mut self.touched_groups_buf); + for &group_idx in &touched { + let idx = self.extreme_of_each_group_buf.0[group_idx]; + extract_row_at_idx_to_buf( + &values_and_order_cols[1..], + idx, + &mut ordering_buf, + )?; + self.update_state(group_idx, &ordering_buf, vals, idx)?; + } + // Restore the all-false invariant by clearing only the touched bits. + for &group_idx in &touched { + self.extreme_of_each_group_buf.1.set_bit(group_idx, false); + } + self.touched_groups_buf = touched; + self.touched_groups_buf.clear(); + + Ok(()) + } } impl<S: ValueState + 'static> GroupsAccumulator for FirstLastGroupsAccumulator<S> { @@ -636,6 +773,14 @@ impl<S: ValueState + 'static> GroupsAccumulator for FirstLastGroupsAccumulator<S ) -> Result<()> { self.resize_states(total_num_groups); + if self.is_input_pre_ordered { + return self.update_batch_pre_ordered( + values_and_order_cols, + group_indices, + opt_filter, + ); + } + let vals = &values_and_order_cols[0]; let mut ordering_buf = Vec::with_capacity(self.ordering_req.len()); @@ -1147,7 +1292,7 @@ impl AggregateUDFImpl for LastValue { &self, args: AccumulatorArgs, ) -> Result<Box<dyn GroupsAccumulator>> { - create_groups_accumulator(&args, false, self.name()) + create_groups_accumulator(&args, false, self.is_input_pre_ordered, self.name()) } } @@ -1587,6 +1732,7 @@ mod tests { true, &[DataType::Int64], true, + false, )?; let mut val_with_orderings = { @@ -1679,6 +1825,7 @@ mod tests { true, &[DataType::Int64], true, + false, )?; let val_with_orderings = { @@ -1760,6 +1907,7 @@ mod tests { true, &[DataType::Int64], false, + false, )?; let mut val_with_orderings = { @@ -1834,6 +1982,7 @@ mod tests { true, &[DataType::Int64], true, + false, )?; let val_with_orderings: Vec<ArrayRef> = vec![ @@ -1888,6 +2037,7 @@ mod tests { true, &[DataType::Int64], true, + false, )?; let val_with_orderings: Vec<ArrayRef> = vec![ @@ -1910,6 +2060,7 @@ mod tests { true, &[DataType::Int64], true, + false, )?; merging_acc.merge_batch(&state, &[0, 0], 1)?; @@ -2104,6 +2255,7 @@ mod tests { false, &[DataType::Int64], /* pick_first = */ true, + false, )?; // Batch 1: four rows across two groups. @@ -2187,6 +2339,7 @@ mod tests { false, &[DataType::Int64], true, + false, )?; // 10 groups × 10_000 candidate rows per group (100_000 total). Each @@ -2277,6 +2430,7 @@ mod tests { false, &[DataType::Int64], true, + false, )?; const GROUPS: usize = 4; @@ -2385,4 +2539,426 @@ mod tests { } Ok(()) } + + // ==================== pre-ordered fast path (#24771) ==================== + + use arrow::datatypes::{Field, Int64Type}; + + /// Builds a grouped first/last accumulator over Int64 values with a + /// two-column Int64 ordering requirement `(o1, o2)`. + #[expect(clippy::fn_params_excessive_bools)] + fn grouped_acc( + pick_first: bool, + ignore_nulls: bool, + pre_ordered: bool, + descending: bool, + ) -> Result<FirstLastGroupsAccumulator<PrimitiveValueState<Int64Type>>> { + let schema = Schema::new(vec![ + Field::new("v", DataType::Int64, true), + Field::new("o1", DataType::Int64, true), + Field::new("o2", DataType::Int64, true), + ]); + let options = SortOptions { + descending, + nulls_first: false, + }; + let ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new(col("o1", &schema)?, options), + PhysicalSortExpr::new(col("o2", &schema)?, options), + ]) + .unwrap(); + FirstLastGroupsAccumulator::try_new( + PrimitiveValueState::<Int64Type>::new(DataType::Int64), + ordering, + ignore_nulls, + &[DataType::Int64, DataType::Int64], + pick_first, + pre_ordered, + ) + } + + /// `[value, o1, o2]` columns for one batch. + fn vo_batch(vals: &[Option<i64>], o1: &[i64], o2: &[i64]) -> Vec<ArrayRef> { + assert_eq!(vals.len(), o1.len()); + assert_eq!(vals.len(), o2.len()); + vec![ + Arc::new(Int64Array::from(vals.to_vec())) as ArrayRef, + Arc::new(Int64Array::from(o1.to_vec())) as ArrayRef, + Arc::new(Int64Array::from(o2.to_vec())) as ArrayRef, + ] + } + + type PreOrderedBatch = (Vec<ArrayRef>, Vec<usize>, Option<BooleanArray>, usize); + + /// Feeds identical batches through the tournament path and the pre-ordered + /// fast path, then asserts that every emitted state column (value, + /// ordering columns, and the is_set flags) is identical. + fn assert_matches_tournament( + pick_first: bool, + ignore_nulls: bool, + descending: bool, + batches: &[PreOrderedBatch], + ) -> Result<()> { + let mut slow = grouped_acc(pick_first, ignore_nulls, false, descending)?; + let mut fast = grouped_acc(pick_first, ignore_nulls, true, descending)?; + for (cols, group_indices, filter, total_num_groups) in batches { + slow.update_batch(cols, group_indices, filter.as_ref(), *total_num_groups)?; + fast.update_batch(cols, group_indices, filter.as_ref(), *total_num_groups)?; + } + let slow_state = slow.state(EmitTo::All)?; + let fast_state = fast.state(EmitTo::All)?; + assert_eq!(slow_state.len(), fast_state.len()); + for (col_idx, (s, f)) in slow_state.iter().zip(fast_state.iter()).enumerate() { + assert_eq!( + s.to_data(), + f.to_data(), + "state column {col_idx} differs between tournament and \ + pre-ordered fast path" + ); + } + Ok(()) + } + + fn int64_values(arr: &ArrayRef) -> Vec<Option<i64>> { + let arr = arr.as_primitive::<Int64Type>(); + (0..arr.len()) + .map(|i| arr.is_valid(i).then(|| arr.value(i))) + .collect() + } + + /// LAST_VALUE, ascending, groups interleaved (PartiallySorted-style shape: + /// each group's own rows are ordered even though groups mix), duplicate + /// `o1` broken by `o2`, one group appearing only in the second batch. + #[test] + fn pre_ordered_last_value_matches_tournament() -> Result<()> { + let batches: Vec<PreOrderedBatch> = vec![ + ( + vo_batch( + &[Some(10), Some(20), Some(11), Some(12), Some(21)], + &[1, 1, 2, 2, 3], + &[1, 1, 1, 2, 1], + ), + vec![0, 1, 0, 0, 1], + None, + 2, + ), + ( + vo_batch(&[Some(13), Some(30), Some(22)], &[4, 1, 5], &[1, 1, 1]), + vec![0, 2, 1], + None, + 3, + ), + ]; + assert_matches_tournament(false, false, false, &batches)?; + + // Explicit expected values, so both paths being wrong together is + // caught too: group 0 last row is (o=4,1)->13, group 1 is (5,1)->22, + // group 2 only saw (1,1)->30. + let mut fast = grouped_acc(false, false, true, false)?; + for (cols, gids, filter, total) in &batches { + fast.update_batch(cols, gids, filter.as_ref(), *total)?; + } + let out = fast.evaluate(EmitTo::All)?; + assert_eq!(int64_values(&out), vec![Some(13), Some(22), Some(30)]); + Ok(()) + } + + /// FIRST_VALUE, ascending: the first batch decides every group it saw; + /// later batches must not overwrite. + #[test] + fn pre_ordered_first_value_locks_after_first_qualifying_row() -> Result<()> { + let batches: Vec<PreOrderedBatch> = vec![ + ( + vo_batch(&[Some(10), Some(20)], &[1, 1], &[1, 2]), + vec![0, 1], + None, + 2, + ), + ( + vo_batch(&[Some(99), Some(98), Some(30)], &[2, 3, 1], &[1, 1, 1]), + vec![0, 1, 2], + None, + 3, + ), + ]; + assert_matches_tournament(true, false, false, &batches)?; + + let mut fast = grouped_acc(true, false, true, false)?; + for (cols, gids, filter, total) in &batches { + fast.update_batch(cols, gids, filter.as_ref(), *total)?; + } + let out = fast.evaluate(EmitTo::All)?; + assert_eq!(int64_values(&out), vec![Some(10), Some(20), Some(30)]); + Ok(()) + } + + /// Descending ordering requirement with input laid out descending: the + /// physically-last row is still the requirement's extreme. + #[test] + fn pre_ordered_descending_matches_tournament() -> Result<()> { + let batches: Vec<PreOrderedBatch> = vec![ + ( + vo_batch( + &[Some(1), Some(2), Some(3), Some(4)], + &[9, 9, 7, 5], + &[5, 3, 1, 1], + ), + vec![0, 0, 1, 0], + None, + 2, + ), + ( + vo_batch(&[Some(5), Some(6)], &[4, 2], &[9, 9]), + vec![0, 1], + None, + 2, + ), + ]; + assert_matches_tournament(false, false, true, &batches)?; + assert_matches_tournament(true, false, true, &batches)?; + Ok(()) + } + + /// FILTER interaction: a row with a `null` predicate is excluded; a group + /// whose rows are all filtered in the last batch keeps its earlier winner; + /// a group filtered everywhere stays unset (emits null + is_set=false). + #[test] + fn pre_ordered_respects_filter() -> Result<()> { + let batches: Vec<PreOrderedBatch> = vec![ + ( + vo_batch( + &[Some(10), Some(11), Some(20), Some(30)], + &[1, 2, 1, 1], + &[1, 1, 1, 1], + ), + vec![0, 0, 1, 2], + Some(BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + Some(false), + ])), + 3, + ), + ( + vo_batch(&[Some(12), Some(21), Some(31)], &[3, 2, 2], &[1, 1, 1]), + vec![0, 1, 2], + Some(BooleanArray::from(vec![Some(false), None, Some(false)])), + 3, + ), + ]; + assert_matches_tournament(false, false, false, &batches)?; + + let mut fast = grouped_acc(false, false, true, false)?; + for (cols, gids, filter, total) in &batches { + fast.update_batch(cols, gids, filter.as_ref(), *total)?; + } + let state = fast.state(EmitTo::All)?; + // value column: group 0 keeps batch-1 winner 11 (batch 2 filtered), + // group 1 keeps 20 (null predicate excluded), group 2 never set. + assert_eq!(int64_values(&state[0]), vec![Some(11), Some(20), None]); + let is_sets = state.last().unwrap().as_boolean(); + assert_eq!( + (0..3).map(|i| is_sets.value(i)).collect::<Vec<_>>(), + vec![true, true, false] + ); + Ok(()) + } + + /// IGNORE NULLS: null values are skipped, so the winner is the last + /// non-null row; an all-null group stays unset. + #[test] + fn pre_ordered_respects_ignore_nulls() -> Result<()> { + let batches: Vec<PreOrderedBatch> = vec![ + ( + vo_batch( + &[Some(10), None, None, Some(20)], + &[1, 2, 1, 1], + &[1, 1, 1, 2], + ), + vec![0, 0, 1, 1], + None, + 2, + ), + ( + vo_batch(&[None, None], &[3, 2], &[1, 1]), + vec![0, 2], + None, + 3, + ), + ]; + assert_matches_tournament(false, true, false, &batches)?; + assert_matches_tournament(true, true, false, &batches)?; + + let mut fast = grouped_acc(false, true, true, false)?; + for (cols, gids, filter, total) in &batches { + fast.update_batch(cols, gids, filter.as_ref(), *total)?; + } + let state = fast.state(EmitTo::All)?; + assert_eq!(int64_values(&state[0]), vec![Some(10), Some(20), None]); + Ok(()) + } + + /// Ties: among rows with equal ordering keys the fast path picks the + /// physically last row for LAST_VALUE / first row for FIRST_VALUE, + /// matching the single-group pre-ordered accumulator. (The tournament + /// path keeps the first-seen row of a tie, so the two paths may pick + /// different — equally valid — rows; this test pins the fast path's + /// choice rather than asserting equivalence.) + #[test] + fn pre_ordered_tie_picks_positional_extreme() -> Result<()> { + let cols = vo_batch(&[Some(10), Some(11), Some(12)], &[1, 1, 1], &[1, 1, 1]); + let mut last = grouped_acc(false, false, true, false)?; + last.update_batch(&cols, &[0, 0, 0], None, 1)?; + assert_eq!(int64_values(&last.evaluate(EmitTo::All)?), vec![Some(12)]); + + let mut first = grouped_acc(true, false, true, false)?; + first.update_batch(&cols, &[0, 0, 0], None, 1)?; + assert_eq!(int64_values(&first.evaluate(EmitTo::All)?), vec![Some(10)]); + Ok(()) + } + + /// RESPECT NULLS (the default): a null value can itself be the winner; + /// the fast path must not treat value-nulls specially. + #[test] + fn pre_ordered_respect_nulls_null_can_win() -> Result<()> { + let batches: Vec<PreOrderedBatch> = vec![ + ( + vo_batch(&[Some(10), None], &[1, 2], &[1, 1]), + vec![0, 0], + None, + 1, + ), + ( + vo_batch(&[None, Some(7)], &[1, 2], &[1, 1]), + vec![1, 1], + None, + 2, + ), + ]; + assert_matches_tournament(false, false, false, &batches)?; + assert_matches_tournament(true, false, false, &batches)?; + + let mut fast = grouped_acc(false, false, true, false)?; + for (cols, gids, filter, total) in &batches { + fast.update_batch(cols, gids, filter.as_ref(), *total)?; + } + // Group 0's last row is the null; it wins under RESPECT NULLS. + let out = fast.evaluate(EmitTo::All)?; + assert_eq!(int64_values(&out), vec![None, Some(7)]); + Ok(()) + } + + /// Draining with `EmitTo::First(n)` mid-stream shifts group indices; the + /// fast path must stay in lockstep with the tournament path across the + /// shift. + #[test] + fn pre_ordered_survives_partial_emit() -> Result<()> { + let mut slow = grouped_acc(false, false, false, false)?; + let mut fast = grouped_acc(false, false, true, false)?; + + let b1 = vo_batch(&[Some(10), Some(20), Some(30)], &[1, 1, 1], &[1, 1, 1]); + for acc in [&mut slow, &mut fast] { + acc.update_batch(&b1, &[0, 1, 2], None, 3)?; + } + + // Emit the first two groups; group 2 shifts down to index 0. + let s1 = slow.state(EmitTo::First(2))?; + let f1 = fast.state(EmitTo::First(2))?; + for (s, f) in s1.iter().zip(f1.iter()) { + assert_eq!(s.to_data(), f.to_data()); + } + assert_eq!(int64_values(&f1[0]), vec![Some(10), Some(20)]); + + // Keep feeding the surviving group (now index 0) plus a new group. + let b2 = vo_batch(&[Some(31), Some(40)], &[2, 1], &[1, 1]); + for acc in [&mut slow, &mut fast] { + acc.update_batch(&b2, &[0, 1], None, 2)?; + } + let s2 = slow.state(EmitTo::All)?; + let f2 = fast.state(EmitTo::All)?; + for (s, f) in s2.iter().zip(f2.iter()) { + assert_eq!(s.to_data(), f.to_data()); + } + assert_eq!(int64_values(&f2[0]), vec![Some(31), Some(40)]); + Ok(()) + } + + /// End-to-end partial→final: states produced by the fast path carry the + /// winning ordering values, so a (never pre-ordered) final-stage + /// accumulator merging two partitions in either arrival order picks the + /// true global extreme. + #[test] + fn pre_ordered_partial_states_merge_correctly() -> Result<()> { + let make_partition_state = |o1: i64, val: i64| -> Result<Vec<ArrayRef>> { + let mut partial = grouped_acc(false, false, true, false)?; + partial.update_batch(&vo_batch(&[Some(val)], &[o1], &[1]), &[0], None, 1)?; + partial.state(EmitTo::All) + }; + // Partition A saw the later row (o1=9), partition B the earlier one. + let a = make_partition_state(9, 900)?; + let b = make_partition_state(3, 300)?; + + for order in [[&a, &b], [&b, &a]] { + let mut final_acc = grouped_acc(false, false, false, false)?; + for state in order { + final_acc.merge_batch(state, &[0], 1)?; + } + let out = final_acc.evaluate(EmitTo::All)?; + assert_eq!(int64_values(&out), vec![Some(900)]); + } + Ok(()) + } + + /// Winner collection and scoreboard reset must scale with the groups a + /// batch touches, not with `total_num_groups`. Correctness side of that: + /// sparse high group indices against a large total still resolve. + #[test] + fn pre_ordered_sparse_groups_large_total() -> Result<()> { + let total = 100_000; + let mut fast = grouped_acc(false, false, true, false)?; + fast.update_batch( + &vo_batch(&[Some(10), Some(20)], &[1, 1], &[1, 1]), + &[7, 42_000], + None, + total, + )?; + fast.update_batch( + &vo_batch(&[Some(11), Some(30)], &[2, 1], &[1, 1]), + &[7, 99_999], + None, + total, + )?; + let out = fast.evaluate(EmitTo::All)?; + let vals = int64_values(&out); + assert_eq!(vals.len(), total); + assert_eq!(vals[7], Some(11)); + assert_eq!(vals[42_000], Some(20)); + assert_eq!(vals[99_999], Some(30)); + assert_eq!(vals[0], None); + Ok(()) + } + + /// The aggregation stream calls `merge_batch` on the same accumulators + /// when replaying spilled state, and the tournament helper leaves its + /// winners' scoreboard bits set. A later pre-ordered `update_batch` must + /// not mistake those for "touched in this batch", or a merged group's + /// newer row is dropped from the winner list. + #[test] + fn pre_ordered_update_after_merge_interleave() -> Result<()> { + let merged_state = { + let mut partial = grouped_acc(false, false, true, false)?; + partial.update_batch(&vo_batch(&[Some(100)], &[1], &[1]), &[0], None, 1)?; + partial.state(EmitTo::All)? + }; + + let mut acc = grouped_acc(false, false, true, false)?; + acc.merge_batch(&merged_state, &[0], 1)?; + // A newer row (higher ordering key) for the merged group arrives + // through the fast path afterwards. + acc.update_batch(&vo_batch(&[Some(500)], &[5], &[1]), &[0], None, 1)?; + let out = acc.evaluate(EmitTo::All)?; + assert_eq!(int64_values(&out), vec![Some(500)]); + Ok(()) + } } diff --git a/datafusion/physical-optimizer/src/update_aggr_exprs.rs b/datafusion/physical-optimizer/src/update_aggr_exprs.rs index 2430918e2c..a047ab11f4 100644 --- a/datafusion/physical-optimizer/src/update_aggr_exprs.rs +++ b/datafusion/physical-optimizer/src/update_aggr_exprs.rs @@ -89,20 +89,35 @@ impl PhysicalOptimizerRule for OptimizeAggregateOrder { let input = aggr_exec.input(); let mut aggr_exprs = aggr_exec.aggr_expr().to_vec(); - 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<_>>(); + // + // Only do this for a single grouping set. With grouping sets + // (e.g. ROLLUP), the stream feeds the same rows once per + // grouping set, and within a coarser set's group the rows are + // ordered by the *full* group-by prefix, not by the + // aggregate's own ORDER BY. Proving the prefixed requirement + // there would mark first/last aggregates as pre-ordered when + // their groups are not, producing wrong results in positional + // paths. Without the prefix, a globally satisfied ORDER BY + // still holds within every group of every 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![] + }; aggr_exprs = try_convert_aggregate_if_better( aggr_exprs, diff --git a/datafusion/sqllogictest/test_files/first_last_ordered.slt b/datafusion/sqllogictest/test_files/first_last_ordered.slt new file mode 100644 index 0000000000..63a7cc1cf3 --- /dev/null +++ b/datafusion/sqllogictest/test_files/first_last_ordered.slt @@ -0,0 +1,183 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# End-to-end coverage for the grouped first_value/last_value pre-ordered fast +# path (issue #24771): when `OptimizeAggregateOrder` proves the input ordering +# satisfies the group-by prefix followed by the aggregate's ordering +# requirement, the grouped accumulator skips its per-row tournament. These +# queries exercise the whole chain — the requirement being proven through +# declared orderings, through projection-alias equivalences, and through +# filter-induced constants — and pin the results. + +statement ok +CREATE EXTERNAL TABLE multiple_ordered_table ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER +) +STORED AS CSV +WITH ORDER (a ASC, b ASC) +WITH ORDER (c ASC) +LOCATION '../core/tests/data/window_2.csv' +OPTIONS ('format.has_header' 'true'); + +# Direct match: requirement (a) ++ (c ASC) is satisfied by the declared +# orderings, so the fast path runs. +query III +SELECT a, FIRST_VALUE(c ORDER BY c), LAST_VALUE(c ORDER BY c) +FROM multiple_ordered_table GROUP BY a ORDER BY a; +---- +0 0 49 +1 50 99 + +# The winner's non-key column rides along (argmax semantics). +query III +SELECT a, FIRST_VALUE(d ORDER BY c), LAST_VALUE(d ORDER BY c) +FROM multiple_ordered_table GROUP BY a ORDER BY a; +---- +0 0 3 +1 0 3 + +# Requirement proven through a projection-alias equivalence class: the input +# is ordered by (a, b) / (c), the query groups by an alias of `a`, and the +# aggregate orders by an alias of `c`. +query II +SELECT a2, LAST_VALUE(d2 ORDER BY c2) +FROM (SELECT a AS a2, c AS c2, d AS d2 FROM multiple_ordered_table) +GROUP BY a2 ORDER BY a2; +---- +0 3 +1 3 + +# Requirement proven with the help of a filter-induced constant: `a` is +# constant under `a = 1`, so the (b) prefix comes from the (a, b) ordering and +# `c` stays ordered inside every b-segment. +query III +SELECT b, FIRST_VALUE(c ORDER BY c), LAST_VALUE(c ORDER BY c) +FROM multiple_ordered_table WHERE a = 1 GROUP BY b ORDER BY b; +---- +2 50 74 +3 75 99 + +# FILTER clause interacts with the fast path. +query II +SELECT a, LAST_VALUE(c ORDER BY c) FILTER (WHERE c % 2 = 0) +FROM multiple_ordered_table GROUP BY a ORDER BY a; +---- +0 48 +1 98 + +# A DESC requirement over an ASC input takes the reverse branch: the rule +# swaps FIRST_VALUE for LAST_VALUE before setting the flag, and the result is +# unchanged. +query II +SELECT a, FIRST_VALUE(c ORDER BY c DESC) +FROM multiple_ordered_table GROUP BY a ORDER BY a; +---- +0 49 +1 99 + +statement ok +DROP TABLE multiple_ordered_table; + +# With grouping sets the same rows are fed once per grouping set, and within a +# coarser set's group the rows are ordered by the full group-by prefix rather +# than the aggregate's ORDER BY. The rule must not mark first/last aggregates +# pre-ordered there (github.com/apache/datafusion/pull/25050 review): input +# below is ordered on (a, b, c) but within the ROLLUP's (a) set, group a=1 +# arrives ordered by (b, c), so FIRST/LAST over `ORDER BY c` need comparisons. + +statement ok +COPY (VALUES + (1, 1, 5, 100, false), (1, 1, 6, 101, true), + (1, 2, 1, 200, true), (1, 2, 2, 201, false)) +TO 'test_files/scratch/first_last_ordered/rollup.csv' +OPTIONS ('format.has_header' 'false'); + +statement ok +CREATE EXTERNAL TABLE rollup_source (a INT, b INT, c INT, d INT, e BOOLEAN) +STORED AS CSV +LOCATION 'test_files/scratch/first_last_ordered/rollup.csv' +WITH ORDER (a ASC, b ASC, c ASC) +OPTIONS ('format.has_header' 'false'); + +query IIII +SELECT a, b, FIRST_VALUE(d ORDER BY c), LAST_VALUE(d ORDER BY c) +FROM rollup_source GROUP BY ROLLUP(a, b) ORDER BY a NULLS FIRST, b NULLS FIRST; +---- +NULL NULL 200 101 +1 NULL 200 101 +1 1 100 101 +1 2 200 201 + +# An aggregate ORDER BY that is satisfied without any group-by prefix stays on +# the fast path even under ROLLUP: `a` leads the input ordering, so it is +# ordered inside every group of every grouping set. +query III +SELECT a, b, LAST_VALUE(d ORDER BY a) +FROM rollup_source GROUP BY ROLLUP(a, b) ORDER BY a NULLS FIRST, b NULLS FIRST; +---- +NULL NULL 201 +1 NULL 201 +1 1 101 +1 2 201 + +# Boolean is outside groups_accumulator_supported, so these run through +# GroupsAccumulatorAdapter and the single-group accumulators, which consumed +# the pre-ordered flag on main already: this covers the latent wrong-results +# path that existed before this PR, not just the new grouped fast path. +query IIBB +SELECT a, b, FIRST_VALUE(e ORDER BY c), LAST_VALUE(e ORDER BY c) +FROM rollup_source GROUP BY ROLLUP(a, b) ORDER BY a NULLS FIRST, b NULLS FIRST; +---- +NULL NULL true true +1 NULL true true +1 1 false true +1 2 true false + +statement ok +DROP TABLE rollup_source; + +# Tie handling is a user-visible behavior change: among rows whose ordering +# keys compare equal, the pre-ordered fast path picks the physically last +# qualifying row for LAST_VALUE (first for FIRST_VALUE), while the tournament +# path keeps the first-seen row of a tie. Both are valid answers -- which row +# of a tie wins is unspecified -- but declaring WITH ORDER on a table can +# change which one a query returns. This pins the fast path's choice. +statement ok +COPY (VALUES (1, 1, 10), (1, 1, 11), (2, 1, 20), (2, 1, 21)) +TO 'test_files/scratch/first_last_ordered/ties.csv' +OPTIONS ('format.has_header' 'false'); + +statement ok +CREATE EXTERNAL TABLE tie_source (g INT, ts INT, px INT) +STORED AS CSV +LOCATION 'test_files/scratch/first_last_ordered/ties.csv' +WITH ORDER (g ASC, ts ASC) +OPTIONS ('format.has_header' 'false'); + +query III +SELECT g, FIRST_VALUE(px ORDER BY ts), LAST_VALUE(px ORDER BY ts) +FROM tie_source GROUP BY g ORDER BY g; +---- +1 10 11 +2 20 21 + +statement ok +DROP TABLE tie_source; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
