From db2fd59bcbc69f866d536bc5914380d87fa81077 Mon Sep 17 00:00:00 2001 From: hushen <190065939+918154429@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:42:25 +0800 Subject: [PATCH] perf(aggregate): specialize fully ordered group keys without hashing --- .../file_stream_provider.rs | 7 + datafusion/common/src/rounding.rs | 2 +- datafusion/physical-plan/Cargo.toml | 4 + .../benches/ordered_group_values.rs | 265 ++++++++++ .../src/aggregates/group_values/mod.rs | 10 +- .../group_values/multi_group_by/mod.rs | 2 + .../group_values/multi_group_by/ordered.rs | 471 ++++++++++++++++++ 7 files changed, 759 insertions(+), 2 deletions(-) create mode 100644 datafusion/physical-plan/benches/ordered_group_values.rs create mode 100644 datafusion/physical-plan/src/aggregates/group_values/multi_group_by/ordered.rs diff --git a/datafusion-examples/examples/custom_data_source/file_stream_provider.rs b/datafusion-examples/examples/custom_data_source/file_stream_provider.rs index 5b43072d43f80..49be4469c3947 100644 --- a/datafusion-examples/examples/custom_data_source/file_stream_provider.rs +++ b/datafusion-examples/examples/custom_data_source/file_stream_provider.rs @@ -28,6 +28,13 @@ /// with DataFusion without needing to reload the entire dataset each time. /// /// This example does not work on Windows. +#[cfg_attr( + target_os = "windows", + expect( + clippy::unused_async, + reason = "keep the same async entry point on all platforms" + ) +)] pub async fn file_stream_provider() -> datafusion::error::Result<()> { #[cfg(target_os = "windows")] { diff --git a/datafusion/common/src/rounding.rs b/datafusion/common/src/rounding.rs index 1796143d7cf1a..b3b514e8631c8 100644 --- a/datafusion/common/src/rounding.rs +++ b/datafusion/common/src/rounding.rs @@ -254,7 +254,7 @@ where } } _ => {} - }; + } Ok(result) } diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 534cea8ea9cbb..a828e61cf8164 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -169,3 +169,7 @@ name = "window_filter" [[bench]] harness = false name = "range_repartition" + +[[bench]] +harness = false +name = "ordered_group_values" diff --git a/datafusion/physical-plan/benches/ordered_group_values.rs b/datafusion/physical-plan/benches/ordered_group_values.rs new file mode 100644 index 0000000000000..ae0e44a859d5e --- /dev/null +++ b/datafusion/physical-plan/benches/ordered_group_values.rs @@ -0,0 +1,265 @@ +// 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. + +//! Fully ordered grouping: compare the selected implementation with the +//! existing streaming hash table on identical batches. Include emission and +//! cross-batch continuation, not just lookup. Input preparation is untimed. +//! The physical-plan benchmark includes aggregation but excludes SQL planning, +//! sorting and I/O; the data already has a proven ordering. + +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ArrayRef, Int32Array, Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_execution::TaskContext; +use datafusion_expr::EmitTo; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn; +use datafusion_physical_plan::aggregates::group_values::{GroupValues, new_group_values}; +use datafusion_physical_plan::aggregates::order::GroupOrdering; +use datafusion_physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, collect}; +use tokio::runtime::Runtime; + +const ROWS: usize = 131_072; + +fn inputs( + run_length: usize, + batch_size: usize, + strings: bool, +) -> (SchemaRef, Vec>) { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new( + "b", + if strings { + DataType::Utf8 + } else { + DataType::Int32 + }, + false, + ), + ])); + let batches = (0..ROWS) + .step_by(batch_size) + .map(|start| { + let end = (start + batch_size).min(ROWS); + let first: ArrayRef = Arc::new(Int32Array::from_iter_values( + (start..end).map(|row| ((row / run_length) / 4) as i32), + )); + let second: ArrayRef = if strings { + Arc::new(StringArray::from_iter_values((start..end).map(|row| { + format!("key-{}-a-longer-than-inline-string", (row / run_length) % 4) + }))) + } else { + Arc::new(Int32Array::from_iter_values( + (start..end).map(|row| ((row / run_length) % 4) as i32), + )) + }; + vec![first, second] + }) + .collect(); + (schema, batches) +} + +fn grouping(c: &mut Criterion) { + let mut group = c.benchmark_group("fully_ordered_grouping"); + group.sample_size(10); + group.warm_up_time(Duration::from_millis(250)); + group.measurement_time(Duration::from_secs(1)); + for strings in [false, true] { + for batch_size in [127, 8192] { + for run_length in [1, 8, 128, 8192] { + let (schema, batches) = inputs(run_length, batch_size, strings); + let case = format!( + "{}_batch{batch_size}_run{run_length}", + if strings { "string" } else { "int" } + ); + let (hashed_bytes, selected_bytes) = check_case(&schema, &batches); + eprintln!( + "group_state_bytes {case} hashed={hashed_bytes} selected={selected_bytes}" + ); + for selected in [false, true] { + let name = if selected { "selected" } else { "hashed" }; + group.bench_function(BenchmarkId::new(name, &case), |b| { + b.iter_batched_ref( + || { + let values: Box = if selected { + new_group_values( + Arc::clone(&schema), + &GroupOrdering::try_new(&InputOrderMode::Sorted) + .unwrap(), + ) + .unwrap() + } else { + Box::new( + GroupValuesColumn::::try_new(Arc::clone( + &schema, + )) + .unwrap(), + ) + }; + (values, Vec::new()) + }, + |(values, ids)| { + for batch in &batches { + values.intern(batch, ids).unwrap(); + black_box(&*ids); + let completed = values.len().saturating_sub(1); + if completed > 0 { + black_box( + values + .emit(EmitTo::First(completed)) + .unwrap(), + ); + } + } + black_box(values.emit(EmitTo::All).unwrap()); + }, + criterion::BatchSize::LargeInput, + ); + }); + } + } + } + } + group.finish(); +} + +// Validate the benchmark input and report retained grouping-state memory. +// These counters exclude the shared input and transient kernel allocations; +// they are not process RSS measurements. Neither validation nor reporting is timed. +fn check_case(schema: &SchemaRef, batches: &[Vec]) -> (usize, usize) { + let mut hashed = GroupValuesColumn::::try_new(Arc::clone(schema)).unwrap(); + let mut selected = new_group_values( + Arc::clone(schema), + &GroupOrdering::try_new(&InputOrderMode::Sorted).unwrap(), + ) + .unwrap(); + let mut expected = Vec::new(); + let mut actual = Vec::new(); + let mut hashed_peak = 0; + let mut selected_peak = 0; + for batch in batches { + hashed.intern(batch, &mut expected).unwrap(); + selected.intern(batch, &mut actual).unwrap(); + assert_eq!(actual, expected); + hashed_peak = hashed_peak.max(hashed.size()); + selected_peak = selected_peak.max(selected.size()); + let completed = hashed.len().saturating_sub(1); + assert_eq!( + selected.emit(EmitTo::First(completed)).unwrap(), + hashed.emit(EmitTo::First(completed)).unwrap() + ); + } + assert_eq!( + selected.emit(EmitTo::All).unwrap(), + hashed.emit(EmitTo::All).unwrap() + ); + (hashed_peak, selected_peak) +} + +fn aggregate_plan( + schema: &SchemaRef, + keys: Vec>, +) -> Arc { + let mut fields = schema.fields().to_vec(); + fields.push(Arc::new(Field::new("v", DataType::Int64, false))); + let schema = Arc::new(Schema::new(fields)); + let batches = keys + .into_iter() + .map(|mut cols| { + cols.push(Arc::new(Int64Array::from(vec![1; cols[0].len()]))); + RecordBatch::try_new(Arc::clone(&schema), cols).unwrap() + }) + .collect::>(); + let ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new_default(col("a", &schema).unwrap()), + PhysicalSortExpr::new_default(col("b", &schema).unwrap()), + ]) + .unwrap(); + let input = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None).unwrap(); + let input = Arc::new( + input + .as_ref() + .clone() + .try_with_sort_information(vec![ordering]) + .unwrap(), + ); + let expr = AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema).unwrap()]) + .schema(Arc::clone(&schema)) + .alias("sum_v") + .build() + .unwrap(); + let plan = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + (col("a", &schema).unwrap(), "a".into()), + (col("b", &schema).unwrap(), "b".into()), + ]), + vec![Arc::new(expr)], + vec![None], + input, + schema, + ) + .unwrap(); + assert_eq!(plan.input_order_mode(), &InputOrderMode::Sorted); + Arc::new(plan) +} + +fn aggregation(c: &mut Criterion) { + let runtime = Runtime::new().unwrap(); + let mut group = c.benchmark_group("fully_ordered_aggregate_exec"); + group.sample_size(10); + group.warm_up_time(Duration::from_millis(250)); + group.measurement_time(Duration::from_secs(1)); + for strings in [false, true] { + for run_length in [1, 8, 128, 8192] { + let (schema, keys) = inputs(run_length, 8192, strings); + let plan = aggregate_plan(&schema, keys); + let name = + format!("{}_run{run_length}", if strings { "string" } else { "int" }); + group.bench_function(name, |b| { + b.iter(|| { + black_box( + runtime + .block_on(collect( + Arc::clone(&plan), + Arc::new(TaskContext::default()), + )) + .unwrap(), + ) + }); + }); + } + } + group.finish(); +} + +criterion_group!(benches, grouping, aggregation); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index b764fd7792b39..e5a818c204ecd 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -34,7 +34,7 @@ mod row; pub use row::GroupValuesRows; mod single_group_by; use datafusion_physical_expr::binary_map::OutputType; -use multi_group_by::GroupValuesColumn; +use multi_group_by::{GroupValuesColumn, GroupValuesOrdered}; pub(crate) use single_group_by::primitive::HashValue; @@ -144,6 +144,9 @@ pub trait GroupValues: Send { /// /// [`GroupValues`] implementations choosing logic: /// +/// - Fully ordered multi-column keys with supported scalar types use adjacent +/// comparisons and column builders, without a hash table. +/// /// - If group by single column, and type of this column has /// the specific [`GroupValues`] implementation, such implementation /// will be chosen. @@ -160,6 +163,11 @@ pub fn new_group_values( schema: SchemaRef, group_ordering: &GroupOrdering, ) -> Result> { + if matches!(group_ordering, GroupOrdering::Full(_)) + && GroupValuesOrdered::supports_schema(&schema) + { + return Ok(Box::new(GroupValuesOrdered::try_new(schema)?)); + } if schema.fields.len() == 1 { let d = schema.fields[0].data_type(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 6c1926b402cee..02799e6175311 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -23,6 +23,8 @@ pub mod bytes_view; mod dictionary; mod fixed_size_binary; mod list; +mod ordered; +pub(super) use ordered::GroupValuesOrdered; pub mod primitive; pub mod row_backed; diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/ordered.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/ordered.rs new file mode 100644 index 0000000000000..30978b1f3e4e2 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/ordered.rs @@ -0,0 +1,471 @@ +// 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. + +use std::mem; + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Schema, SchemaRef}; +use arrow_ord::partition::partition; +use datafusion_common::{Result, not_impl_err}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::{EmitTo, GroupSelection}; + +use super::{GroupColumn, GroupValuesColumn}; +use crate::aggregates::group_values::GroupValues; + +/// Columnar group keys for input fully ordered by all grouping expressions. +/// +/// Equal keys must be contiguous across input batches. Within a batch Arrow's +/// partition kernel finds those runs. Only the first run can continue the last +/// buffered group; compare it with the stored columns, without retaining the +/// input batch. New group representatives use the existing column builders. +/// +/// This removes hashing and hash-table storage, but does not change the dense +/// group-id or emission contracts. In particular, `First(n)` shifts the remaining +/// keys and their ids together. Partially ordered inputs must use a hash table. +pub(crate) struct GroupValuesOrdered { + schema: SchemaRef, + columns: Vec>, + new_groups: Vec, +} + +impl GroupValuesOrdered { + /// Types for which adjacent Arrow equality agrees with GROUP BY equality. + /// Keep single-column specializations and floats/nested/encoded keys on the + /// established path until their semantics and performance are validated. + pub(crate) fn supports_schema(schema: &Schema) -> bool { + schema.fields().len() > 1 + && schema.fields().iter().all(|field| { + matches!( + field.data_type(), + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::LargeBinary + | DataType::BinaryView + ) + }) + } + + pub(crate) fn try_new(schema: SchemaRef) -> Result { + if !Self::supports_schema(&schema) { + return not_impl_err!( + "Unsupported schema for fully ordered group values: {schema}" + ); + } + let columns = GroupValuesColumn::::build_group_columns(&schema)?; + Ok(Self { + schema, + columns, + new_groups: Vec::new(), + }) + } +} + +impl GroupValues for GroupValuesOrdered { + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + groups.clear(); + let ranges = partition(cols)?.ranges(); + if ranges.is_empty() { + return Ok(()); + } + + let old_len = self.len(); + let continues = old_len > 0 + && self + .columns + .iter() + .zip(cols) + .all(|(stored, input)| stored.equal_to(old_len - 1, input, 0)); + + self.new_groups.clear(); + self.new_groups.extend( + ranges + .iter() + .skip(usize::from(continues)) + .map(|range| range.start), + ); + for (stored, input) in self.columns.iter_mut().zip(cols) { + stored.vectorized_append(input, &self.new_groups)?; + } + + let first_group = old_len - usize::from(continues); + for (index, range) in ranges.iter().enumerate() { + groups.resize(range.end, first_group + index); + } + Ok(()) + } + + fn size(&self) -> usize { + self.columns + .iter() + .map(|column| column.size()) + .sum::() + + self.new_groups.allocated_size() + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn len(&self) -> usize { + self.columns[0].len() + } + + fn emit(&mut self, emit_to: EmitTo) -> Result> { + Ok(match emit_to { + EmitTo::All => { + let fresh = GroupValuesColumn::::build_group_columns(&self.schema)?; + mem::replace(&mut self.columns, fresh) + .into_iter() + .map(|column| column.build()) + .collect() + } + EmitTo::First(n) => self + .columns + .iter_mut() + .map(|column| column.take_n(n)) + .collect(), + }) + } + + fn values_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + selection.validate_num_groups(self.len())?; + self.columns + .iter() + .map(|column| column.values_preserving(selection)) + .collect() + } + + fn supports_values_preserving(&self) -> bool { + true + } + + fn clear_shrink(&mut self, num_rows: usize) { + self.columns = GroupValuesColumn::::build_group_columns(&self.schema) + .expect("schema validated by GroupValuesOrdered::try_new"); + self.new_groups.clear(); + self.new_groups.shrink_to(num_rows); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::UInt32Array; + use arrow::array::{Array, BooleanArray, Int32Array, Int64Array, StringArray}; + use arrow::compute::{cast, take}; + use arrow::datatypes::Field; + use datafusion_expr::GroupSelection; + + use super::*; + + // Compare with the existing hash implementation, changing actual input + // boundaries and removing completed groups between batches. This checks the + // semantic contract rather than duplicating the adjacent-run algorithm. + #[test] + fn ordered_keys_match_hash_grouping_across_batches_and_emits() -> Result<()> { + for key_type in [ + DataType::Boolean, + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + DataType::UInt8, + DataType::UInt16, + DataType::UInt32, + DataType::UInt64, + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + ] { + let rows = 1025; + let first: ArrayRef = Arc::new(Int32Array::from_iter((0..rows).map(|row| { + let group = row / 7; + (group >= 4).then_some(group / 4) + }))); + let second: ArrayRef = match key_type { + DataType::Boolean => { + Arc::new(BooleanArray::from_iter((0..rows).map(|row| { + let key = (row / 7) % 4; + (key != 0).then_some(key >= 2) + }))) + } + DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::LargeBinary + | DataType::BinaryView => { + let strings = StringArray::from_iter((0..rows).map(|row| { + let key = (row / 7) % 4; + (key != 0).then(|| { + format!( + "key-{key}-{}", + "x".repeat(if key == 3 { 128 } else { 0 }) + ) + }) + })); + cast(&strings, &key_type)? + } + _ => { + let values = Int32Array::from_iter((0..rows).map(|row| { + let key = (row / 7) % 4; + (key != 0).then_some(key) + })); + cast(&values, &key_type)? + } + }; + for descending in [false, true] { + let mut input = vec![Arc::clone(&first), Arc::clone(&second)]; + if descending { + let reverse = UInt32Array::from_iter_values((0..rows as u32).rev()); + input = input + .iter() + .map(|array| take(array, &reverse, None)) + .collect::>()?; + } + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", key_type.clone(), true), + ])); + for batch_size in [1, 2, 7, 8, 63, 1024] { + for emit_limit in [0, 1, 17, usize::MAX] { + let mut ordered = + GroupValuesOrdered::try_new(Arc::clone(&schema))?; + let mut hashed = + GroupValuesColumn::::try_new(Arc::clone(&schema))?; + let mut actual = Vec::new(); + let mut expected = Vec::new(); + for offset in (0..rows as usize).step_by(batch_size) { + let length = batch_size.min(rows as usize - offset); + let batch = input + .iter() + .map(|array| array.slice(offset, length)) + .collect::>(); + ordered.intern(&batch, &mut actual)?; + hashed.intern(&batch, &mut expected)?; + assert_eq!( + actual, expected, + "{key_type:?}, batch={batch_size}, offset={offset}, descending={descending}" + ); + assert_eq!(ordered.len(), hashed.len()); + let selection = GroupSelection::all(ordered.len()); + assert_eq!( + ordered.values_preserving(selection)?, + hashed.values_preserving(selection)? + ); + // A zero-row batch must neither forget the boundary + // key nor leave stale group ids in the output buffer. + let empty = batch + .iter() + .map(|array| array.slice(0, 0)) + .collect::>(); + ordered.intern(&empty, &mut actual)?; + assert!(actual.is_empty()); + let emit = emit_limit.min(ordered.len().saturating_sub(1)); + assert_eq!( + ordered.emit(EmitTo::First(emit))?, + hashed.emit(EmitTo::First(emit))? + ); + } + assert_eq!(ordered.emit(EmitTo::All)?, hashed.emit(EmitTo::All)?); + assert!(ordered.is_empty()); + // All and clear_shrink reset the cross-batch state. + ordered.clear_shrink(0); + hashed.clear_shrink(0); + ordered.intern(&input, &mut actual)?; + hashed.intern(&input, &mut expected)?; + assert_eq!(actual, expected); + assert_eq!(ordered.emit(EmitTo::All)?, hashed.emit(EmitTo::All)?); + } + } + } + } + Ok(()) + } + + #[test] + fn ordered_schema_gate_keeps_unvalidated_types_on_existing_paths() { + for data_type in [ + DataType::Float32, + DataType::Float64, + DataType::Null, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + ] { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", data_type, true), + ]); + assert!(!GroupValuesOrdered::supports_schema(&schema)); + } + assert!(!GroupValuesOrdered::supports_schema(&Schema::new(vec![ + Field::new("a", DataType::Int32, false) + ]))); + } + + #[tokio::test] + async fn ordered_single_and_partial_final_match_unordered_execution() -> Result<()> { + use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + use crate::test::TestMemoryExec; + use crate::{ExecutionPlan, InputOrderMode, collect}; + use arrow::record_batch::RecordBatch; + use arrow::row::{RowConverter, SortField}; + use datafusion_execution::TaskContext; + use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, false), + Field::new("v", DataType::Int64, false), + Field::new("include", DataType::Boolean, false), + ])); + let source = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter( + (0..257).map(|row| (row >= 28).then_some(row / 28)), + )), + Arc::new(StringArray::from_iter_values( + (0..257).map(|row| format!("key-{}", (row / 7) % 4)), + )), + Arc::new(Int64Array::from_iter_values(0..257)), + Arc::new(BooleanArray::from_iter( + (0..257).map(|row| Some(row % 3 != 0)), + )), + ], + )?; + let aggregate = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum_v") + .build()?, + ); + + for input_batch_size in [1, 7, 31, 128] { + let batches = (0..257) + .step_by(input_batch_size) + .map(|offset| source.slice(offset, input_batch_size.min(257 - offset))) + .collect::>(); + for two_stage in [false, true] { + let mut results = Vec::new(); + for sorted in [false, true] { + let input = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?; + let input: Arc = if sorted { + let ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new_default(col("a", &schema)?), + PhysicalSortExpr::new_default(col("b", &schema)?), + ]) + .unwrap(); + Arc::new( + input + .as_ref() + .clone() + .try_with_sort_information(vec![ordering])?, + ) + } else { + input + }; + let group_by = PhysicalGroupBy::new_single(vec![ + (col("a", &schema)?, "a".into()), + (col("b", &schema)?, "b".into()), + ]); + let plan = AggregateExec::try_new( + if two_stage { + AggregateMode::Partial + } else { + AggregateMode::Single + }, + group_by.clone(), + vec![Arc::clone(&aggregate)], + vec![Some(col("include", &schema)?)], + input, + Arc::clone(&schema), + )?; + if sorted { + assert_eq!(plan.input_order_mode(), &InputOrderMode::Sorted); + } + let plan: Arc = if two_stage { + let plan = AggregateExec::try_new( + AggregateMode::Final, + group_by, + vec![Arc::clone(&aggregate)], + vec![None], + Arc::new(plan), + Arc::clone(&schema), + )?; + if sorted { + assert_eq!(plan.input_order_mode(), &InputOrderMode::Sorted); + } + Arc::new(plan) + } else { + Arc::new(plan) + }; + let converter = RowConverter::new( + plan.schema() + .fields() + .iter() + .map(|field| SortField::new(field.data_type().clone())) + .collect(), + )?; + let output = collect(plan, Arc::new(TaskContext::default())).await?; + let mut rows = Vec::new(); + for batch in output { + rows.extend( + converter + .convert_columns(batch.columns())? + .iter() + .map(|row| row.as_ref().to_vec()), + ); + } + rows.sort(); + results.push(rows); + } + assert_eq!( + results[0], results[1], + "batch_size={input_batch_size}, two_stage={two_stage}" + ); + } + } + Ok(()) + } +}