diff --git a/datafusion/functions-aggregate/benches/array_agg.rs b/datafusion/functions-aggregate/benches/array_agg.rs index ad6d502eb5087..ebfe1f8cdd023 100644 --- a/datafusion/functions-aggregate/benches/array_agg.rs +++ b/datafusion/functions-aggregate/benches/array_agg.rs @@ -19,21 +19,25 @@ use std::hint::black_box; use std::sync::Arc; use arrow::array::{ - Array, ArrayRef, ArrowPrimitiveType, AsArray, ListArray, NullBufferBuilder, - StringArray, + Array, ArrayRef, ArrowPrimitiveType, AsArray, Int64Array, ListArray, + NullBufferBuilder, StringArray, }; -use arrow::datatypes::{DataType, Field, Int64Type}; -use criterion::{Criterion, criterion_group, criterion_main}; -use datafusion_expr::Accumulator; +use arrow::datatypes::{DataType, Field, FieldRef, Int64Type, Schema}; +use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; +use datafusion_expr::function::AccumulatorArgs; +use datafusion_expr::{Accumulator, AggregateUDFImpl}; use datafusion_functions_aggregate::array_agg::{ - ArrayAggAccumulator, DistinctArrayAggAccumulator, + ArrayAgg, ArrayAggAccumulator, DistinctArrayAggAccumulator, }; +use datafusion_physical_expr::{PhysicalSortExpr, expressions::col}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use arrow::buffer::OffsetBuffer; use arrow::util::bench_util::create_primitive_array; use rand::Rng; use rand::SeedableRng; use rand::prelude::StdRng; +use rand::seq::SliceRandom; /// Returns fixed seedable RNG pub fn seedable_rng() -> StdRng { @@ -288,5 +292,229 @@ fn distinct_array_agg_benchmark(c: &mut Criterion) { ); } -criterion_group!(benches, array_agg_benchmark, distinct_array_agg_benchmark); +/// Precomputes the schema, physical expressions, sort expression, and aggregate +/// metadata so each benchmark iteration can focus on accumulator operations. +struct OrderedArrayAggBenchFixture { + schema: Schema, + value_expr: Arc, + value_field: FieldRef, + order_by: PhysicalSortExpr, + array_agg: Arc, + return_field: FieldRef, +} + +impl OrderedArrayAggBenchFixture { + fn new(value_type: &DataType, input_preordered: bool) -> Self { + let schema = Schema::new(vec![ + Field::new("value", value_type.clone(), false), + Field::new("ordering", DataType::Int64, false), + ]); + + let value_expr = col("value", &schema).unwrap(); + let ordering_expr = col("ordering", &schema).unwrap(); + + let value_field = value_expr.return_field(&schema).unwrap(); + + let order_by = PhysicalSortExpr::new( + ordering_expr, + arrow::compute::SortOptions { + descending: false, + nulls_first: false, + }, + ); + + let array_agg = Arc::new(ArrayAgg::default()) + .with_beneficial_ordering(input_preordered) + .unwrap() + .unwrap(); + + let return_field = Field::new( + "array_agg", + DataType::List(Field::new_list_field(value_type.clone(), true).into()), + true, + ) + .into(); + + Self { + schema, + value_expr, + value_field, + order_by, + array_agg, + return_field, + } + } + fn create_accumulator(&self) -> Box { + self.array_agg + .accumulator(AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + expr_fields: std::slice::from_ref(&self.value_field), + ignore_nulls: false, + order_bys: std::slice::from_ref(&self.order_by), + is_reversed: false, + name: "array_agg(value ORDER BY ordering)", + is_distinct: false, + exprs: std::slice::from_ref(&self.value_expr), + }) + .unwrap() + } +} + +const ORDERED_ARRAY_AGG_ROWS: usize = 2048; + +fn create_ordered_array_agg_batches( + rows_per_batch: usize, + input_preordered: bool, +) -> Vec<[ArrayRef; 2]> { + assert!( + rows_per_batch > 0, + "rows_per_batch must be greater than zero" + ); + + let upper_value: i64 = ORDERED_ARRAY_AGG_ROWS.try_into().unwrap(); + let mut values = (0..upper_value).collect::>(); + + if !input_preordered { + let mut rng = StdRng::seed_from_u64(42); + values.shuffle(&mut rng); + } + + values + .chunks(rows_per_batch) + .map(|batch_values| { + let values = Arc::new(Int64Array::from(batch_values.to_vec())) as ArrayRef; + + [ + Arc::clone(&values), + values, // Reuse the payload values as ordering keys. + ] + }) + .collect() +} + +fn ordered_array_agg_bench( + c: &mut Criterion, + name: &str, + value_type: &DataType, + batches: &[[ArrayRef; 2]], + input_preordered: bool, +) { + c.bench_function(name, |b| { + let fixture = OrderedArrayAggBenchFixture::new(value_type, input_preordered); + b.iter(|| { + let mut accumulator = fixture.create_accumulator(); + + for batch in batches { + accumulator + .update_batch(batch) + .expect("update_batch should succeed"); + } + + let result = accumulator.evaluate().expect("evaluate should succeed"); + + black_box(result); + }) + }); +} + +fn ordered_array_agg_benchmark(c: &mut Criterion) { + for rows_per_batch in [1, 8, 64, ORDERED_ARRAY_AGG_ROWS] { + let ordered_batches = create_ordered_array_agg_batches(rows_per_batch, true); + ordered_array_agg_bench( + c, + &format!( + "ordered_array_agg i64 ordered input, \ + {rows_per_batch} rows per update_batch" + ), + &DataType::Int64, + &ordered_batches, + true, + ); + let shuffled_batches = create_ordered_array_agg_batches(rows_per_batch, false); + ordered_array_agg_bench( + c, + &format!( + "ordered_array_agg i64 random input, \ + {rows_per_batch} rows per update_batch" + ), + &DataType::Int64, + &shuffled_batches, + false, + ); + } +} + +/// Creates preordered batches with a 4 KiB Utf8 payload per row. +fn create_wide_utf8_batches(rows_per_batch: usize) -> Vec<[ArrayRef; 2]> { + assert!( + rows_per_batch > 0, + "rows_per_batch must be greater than zero" + ); + + let values = vec!["x".repeat(4 * 1_024); ORDERED_ARRAY_AGG_ROWS]; + + values + .chunks(rows_per_batch) + .enumerate() + .map(|(idx, batch_values)| { + let payload = + Arc::new(StringArray::from_iter_values(batch_values.iter())) as ArrayRef; + let start = idx * rows_per_batch; + let ordering_values = Arc::new(Int64Array::from( + (start..start + batch_values.len()) + .map(|row| i64::try_from(row).unwrap()) + .collect::>(), + )) as ArrayRef; + + [payload, ordering_values] + }) + .collect() +} + +fn ordered_array_agg_update_only_bench( + c: &mut Criterion, + name: &str, + value_type: &DataType, + batches: &[[ArrayRef; 2]], + input_preordered: bool, +) { + let fixture = OrderedArrayAggBenchFixture::new(value_type, input_preordered); + + c.bench_function(name, |b| { + b.iter_batched( + || fixture.create_accumulator(), + |mut accumulator| { + for batch in batches { + accumulator + .update_batch(batch) + .expect("update_batch should succeed"); + } + + // Return the accumulator so Criterion drops it after timing the updates. + accumulator + }, + BatchSize::PerIteration, + ); + }); +} + +fn ordered_array_agg_wide_utf8_benchmark(c: &mut Criterion) { + let ordered_batches = create_wide_utf8_batches(1); + ordered_array_agg_update_only_bench( + c, + "ordered_array_agg utf8 4 KiB, 1 row per update_batch", + &DataType::Utf8, + &ordered_batches, + true, + ); +} + +criterion_group!( + benches, + array_agg_benchmark, + distinct_array_agg_benchmark, + ordered_array_agg_benchmark, + ordered_array_agg_wide_utf8_benchmark +); criterion_main!(benches); diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 546fad97387bb..e40a5a2bec618 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -1246,6 +1246,14 @@ struct OrderedArrayAggEntry { row_idx: usize, } +/// Coalesce consecutive small payload batches up to this row count to reduce +/// per-array memory overhead while limiting repeated concatenation costs. +const ORDERED_ARRAY_AGG_COALESCE_ROWS: usize = 64; + +/// Skip coalescing once the tail batch is large enough that the fixed +/// per-array overhead is negligible relative to the payload. +const ORDERED_ARRAY_AGG_COALESCE_BYTES: usize = 4096; + #[derive(Debug)] pub(crate) struct OrderSensitiveArrayAggAccumulator { /// Arrow payload arrays. Entries refer to rows in these batches. @@ -1531,11 +1539,37 @@ impl OrderSensitiveArrayAggAccumulator { } let start = self.entries.len(); - let batch_idx = self.batches.len(); - self.batches.push(values); - self.entries.extend( - (0..row_count).map(|row_idx| OrderedArrayAggEntry { batch_idx, row_idx }), - ); + let (batch_idx, row_offset) = match self.batches.last() { + Some(last_batch) + if last_batch.len() + row_count <= ORDERED_ARRAY_AGG_COALESCE_ROWS + && last_batch.get_buffer_memory_size() + + values.get_buffer_memory_size() + <= ORDERED_ARRAY_AGG_COALESCE_BYTES => + { + let merged = + arrow::compute::concat(&[last_batch.as_ref(), values.as_ref()])?; + + // Concatenation preserves the existing row offsets, while new rows + // start at the previous length of the tail batch. + let batch_idx = self.batches.len() - 1; + let row_offset = last_batch.len(); + + self.batches[batch_idx] = merged; + (batch_idx, row_offset) + } + _ => { + let batch_idx = self.batches.len(); + self.batches.push(values); + (batch_idx, 0) + } + }; + + self.entries + .extend((0..row_count).map(|row_idx| OrderedArrayAggEntry { + batch_idx, + row_idx: row_offset + row_idx, + })); + self.sorted_entry_indices = None; debug_assert_eq!(self.entries.len(), self.ordering_rows.num_rows()); @@ -2513,6 +2547,10 @@ mod tests { final_acc.merge_batch(&state_a)?; final_acc.merge_batch(&state_b)?; + // Small partial-state batches are coalesced into one payload batch. + assert_eq!(final_acc.batches.len(), 1); + assert_eq!(final_acc.batches[0].len(), 4); + let result = print_nulls(str_arr(final_acc.evaluate()?)?); assert_eq!(result, vec!["a1", "a2", "b1", "b2"]); Ok(()) @@ -2908,6 +2946,91 @@ mod tests { Ok(()) } + #[test] + fn ordered_array_agg_coalesces_small_batches_up_to_threshold() -> Result<()> { + use arrow::array::Int64Array; + + let mut acc = ordered_accumulator( + DataType::Int64, + DataType::Int64, + SortOptions::new(false, false), + false, + false, + )?; + + let threshold = ORDERED_ARRAY_AGG_COALESCE_ROWS; + let threshold_i64 = i64::try_from(threshold).unwrap(); + + for values in [ + (0..(threshold_i64 - 1)).collect::>(), + vec![threshold_i64 - 1], + vec![threshold_i64], + ((threshold_i64 + 1)..(threshold_i64 * 2)).collect::>(), + vec![threshold_i64 * 2], + ] { + let values = Arc::new(Int64Array::from(values)) as ArrayRef; + + acc.update_batch(&[Arc::clone(&values), values])?; + } + + assert_eq!( + acc.batches + .iter() + .map(|batch| batch.len()) + .collect::>(), + vec![threshold, threshold, 1], + ); + + assert_eq!(acc.entries.len(), threshold * 2 + 1); + assert_eq!(acc.ordering_rows.num_rows(), threshold * 2 + 1); + + assert_eq!(acc.entries[threshold - 1].batch_idx, 0); + assert_eq!(acc.entries[threshold - 1].row_idx, threshold - 1); + + assert_eq!(acc.entries[threshold].batch_idx, 1); + assert_eq!(acc.entries[threshold].row_idx, 0); + + assert_eq!(acc.entries[threshold * 2 - 1].batch_idx, 1); + assert_eq!(acc.entries[threshold * 2 - 1].row_idx, threshold - 1); + + assert_eq!(acc.entries[threshold * 2].batch_idx, 2); + assert_eq!(acc.entries[threshold * 2].row_idx, 0); + Ok(()) + } + + #[test] + fn ordered_array_agg_sorts_across_coalesced_batches() -> Result<()> { + use arrow::array::Int64Array; + + let mut acc = ordered_accumulator( + DataType::Int64, + DataType::Int64, + SortOptions::new(false, false), + false, + false, + )?; + + for value in (0_i64..128).rev() { + let payload = Arc::new(Int64Array::from(vec![value])) as ArrayRef; + let ordering = Arc::new(Int64Array::from(vec![value + 1000])) as ArrayRef; + acc.update_batch(&[payload, ordering])?; + } + + let ScalarValue::List(result) = acc.evaluate()? else { + return internal_err!("expect list"); + }; + + assert_eq!( + result + .values() + .as_primitive::() + .values(), + &(0..128).collect::>() + ); + + Ok(()) + } + struct ArrayAggAccumulatorBuilder { return_field: FieldRef, distinct: bool,