Skip to content

Commit 092743d

Browse files
TinyMurkyTinyMurky
committed
perf(functions-aggregate): coalesce small ordered ARRAY_AGG batches
- Closes [#25199](#25199) However, the accumulator currently stores one ArrayRef per update_batch() call. Each array has a fixed cost from its ArrayData, buffer allocation, and Arc, regardless of how many rows it contains. Ordered `ARRAY_AGG` does not support a native `GroupsAccumulator` because `groups_accumulator_supported` requires `order_bys` to be empty. Grouped execution therefore falls back to `GroupsAccumulatorAdapter`, which calls `update_batch()` once per group per input batch. With a high-cardinality `GROUP BY`, these calls often contain only one or two rows. Without this change, each call retains a separate `ArrayRef`, so the accumulator pays the fixed per-array allocation cost for many small Arrow arrays, increasing the per-row memory footprint. - Coalesce consecutive small ordered ARRAY_AGG payload batches up to 64 rows. - Add tests covering: - coalescing exactly up to the threshold; - creating a new batch after the threshold is reached; - entry indices across coalesced and newly created batches; - ordering across multiple coalesced batches; - coalescing payloads received through partial-state merge_batch(). - Add ordered ARRAY_AGG benchmarks covering: - 1, 8, 64, and 2,048 rows per update_batch(); - random input; - preordered input. Retained memory was measured using `Accumulator::size()` after inserting 2,048 Int64 payloads with an Int64 ordering key and before calling evaluate(). The same measurement code was used for the baseline after #24392 and for this change. | Rows/update | Batches | Total retained (with this change) | Bytes/row (with this change) | Total retained (baseline [#24392](#24392)) | Bytes/row (baseline [#24392](#24392)) | | ----------: | ------: | --------------------------------: | ---------------------------: | ----------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ | | 1 | 32 | 104,701 B | 51.12 B/row | 445,181 B | 217.37 B/row | | 8 | 32 | 90,365 B | 44.12 B/row | 115,453 B | 56.37 B/row | | 64 | 32 | 88,573 B | 43.25 B/row | 88,573 B | 43.25 B/row | | 2048 | 1 | 84,901 B | 41.46 B/row | 84,901 B | 41.46 B/row | The worst-case one-row update footprint decreases from 217.37 B/row to 51.12 B/row. Inputs already at or above the 64-row coalescing threshold retain the existing memory footprint. Added unit tests covering: - coalescing small batches exactly up to the threshold; - creating a new batch when the threshold would be exceeded; - preserving the correct batch_idx and row_idx for entries; - sorting values across multiple coalesced batches; - coalescing small partial-state payloads passed through merge_batch(). Added Criterion benchmarks for ordered ARRAY_AGG using random and preordered input with 1, 8, 64, and 2,048 rows per update_batch(). The following commands have been executed and passed: - `cargo test -p datafusion-functions-aggregate --lib array_agg::tests` - `cargo bench -p datafusion-functions-aggregate --bench array_agg --no-run` - `cargo bench -p datafusion-functions-aggregate --bench array_agg -ordered_array_agg` - `cargo test --profile=ci --test sqllogictests` - `cargo test -p datafusion` - `cargo test -p datafusion-cli
1 parent 7570366 commit 092743d

2 files changed

Lines changed: 285 additions & 9 deletions

File tree

‎datafusion/functions-aggregate/benches/array_agg.rs‎

Lines changed: 164 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,22 +18,27 @@
1818
use std::hint::black_box;
1919
use std::sync::Arc;
2020

21+
use arrow::array::Int64Array;
2122
use arrow::array::{
2223
Array, ArrayRef, ArrowPrimitiveType, AsArray, ListArray, NullBufferBuilder,
2324
StringArray,
2425
};
25-
use arrow::datatypes::{DataType, Field, Int64Type};
26+
use arrow::datatypes::{DataType, Field, FieldRef, Int64Type, Schema};
2627
use criterion::{Criterion, criterion_group, criterion_main};
27-
use datafusion_expr::Accumulator;
28+
use datafusion_expr::function::AccumulatorArgs;
29+
use datafusion_expr::{Accumulator, AggregateUDFImpl};
2830
use datafusion_functions_aggregate::array_agg::{
29-
ArrayAggAccumulator, DistinctArrayAggAccumulator,
31+
ArrayAgg, ArrayAggAccumulator, DistinctArrayAggAccumulator,
3032
};
33+
use datafusion_physical_expr::{PhysicalSortExpr, expressions::col};
34+
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
3135

3236
use arrow::buffer::OffsetBuffer;
3337
use arrow::util::bench_util::create_primitive_array;
3438
use rand::Rng;
3539
use rand::SeedableRng;
3640
use rand::prelude::StdRng;
41+
use rand::seq::SliceRandom;
3742

3843
/// Returns fixed seedable RNG
3944
pub fn seedable_rng() -> StdRng {
@@ -288,5 +293,160 @@ fn distinct_array_agg_benchmark(c: &mut Criterion) {
288293
);
289294
}
290295

291-
criterion_group!(benches, array_agg_benchmark, distinct_array_agg_benchmark);
296+
/// Precomputes the schema, physical expressions, sort expression, and aggregate
297+
/// metadata so each benchmark iteration measures only accumulator creation,
298+
/// `update_batch()`, and `evaluate()`.
299+
struct OrderedArrayAggBenchFixture {
300+
schema: Schema,
301+
value_expr: Arc<dyn PhysicalExpr>,
302+
value_field: FieldRef,
303+
order_by: PhysicalSortExpr,
304+
array_agg: Arc<dyn AggregateUDFImpl>,
305+
return_field: FieldRef,
306+
}
307+
308+
impl OrderedArrayAggBenchFixture {
309+
fn new(input_preordered: bool) -> Self {
310+
let schema = Schema::new(vec![
311+
Field::new("value", DataType::Int64, false),
312+
Field::new("ordering", DataType::Int64, false),
313+
]);
314+
315+
let value_expr = col("value", &schema).unwrap();
316+
let ordering_expr = col("ordering", &schema).unwrap();
317+
318+
let value_field = value_expr.return_field(&schema).unwrap();
319+
320+
let order_by = PhysicalSortExpr::new(
321+
ordering_expr,
322+
arrow::compute::SortOptions {
323+
descending: false,
324+
nulls_first: false,
325+
},
326+
);
327+
328+
let array_agg = Arc::new(ArrayAgg::default())
329+
.with_beneficial_ordering(input_preordered)
330+
.unwrap()
331+
.unwrap();
332+
333+
let return_field = Field::new(
334+
"array_agg",
335+
DataType::List(Field::new_list_field(DataType::Int64, true).into()),
336+
true,
337+
)
338+
.into();
339+
340+
Self {
341+
schema,
342+
value_expr,
343+
value_field,
344+
order_by,
345+
array_agg,
346+
return_field,
347+
}
348+
}
349+
fn create_accumulator(&self) -> Box<dyn Accumulator> {
350+
self.array_agg
351+
.accumulator(AccumulatorArgs {
352+
return_field: Arc::clone(&self.return_field),
353+
schema: &self.schema,
354+
expr_fields: std::slice::from_ref(&self.value_field),
355+
ignore_nulls: false,
356+
order_bys: std::slice::from_ref(&self.order_by),
357+
is_reversed: false,
358+
name: "array_agg(value ORDER BY ordering)",
359+
is_distinct: false,
360+
exprs: std::slice::from_ref(&self.value_expr),
361+
})
362+
.unwrap()
363+
}
364+
}
365+
366+
const ORDERED_ARRAY_AGG_ROWS: usize = 2048;
367+
368+
fn create_ordered_array_agg_batches(
369+
rows_per_batch: usize,
370+
input_preordered: bool,
371+
) -> Vec<[ArrayRef; 2]> {
372+
assert!(
373+
rows_per_batch > 0,
374+
"rows_per_batch must be greater than zero"
375+
);
376+
377+
let upper_value: i64 = ORDERED_ARRAY_AGG_ROWS.try_into().unwrap();
378+
let mut values = (0..upper_value).collect::<Vec<i64>>();
379+
380+
if !input_preordered {
381+
let mut rng = StdRng::seed_from_u64(42);
382+
values.shuffle(&mut rng);
383+
}
384+
385+
values
386+
.chunks(rows_per_batch)
387+
.map(|batch_values| {
388+
let values = Arc::new(Int64Array::from(batch_values.to_vec())) as ArrayRef;
389+
390+
[
391+
Arc::clone(&values),
392+
values, // Reuse the payload values as ordering keys.
393+
]
394+
})
395+
.collect()
396+
}
397+
398+
fn ordered_array_agg_bench(
399+
c: &mut Criterion,
400+
name: &str,
401+
batches: &[[ArrayRef; 2]],
402+
input_preordered: bool,
403+
) {
404+
c.bench_function(name, |b| {
405+
let fixture = OrderedArrayAggBenchFixture::new(input_preordered);
406+
b.iter(|| {
407+
let mut accumulator = fixture.create_accumulator();
408+
409+
for batch in batches {
410+
accumulator
411+
.update_batch(batch)
412+
.expect("update_batch should succeed");
413+
}
414+
415+
let result = accumulator.evaluate().expect("evaluate should succeed");
416+
417+
black_box(result);
418+
})
419+
});
420+
}
421+
422+
fn ordered_array_agg_benchmark(c: &mut Criterion) {
423+
for rows_per_batch in [1, 8, 64, ORDERED_ARRAY_AGG_ROWS] {
424+
let ordered_batches = create_ordered_array_agg_batches(rows_per_batch, true);
425+
ordered_array_agg_bench(
426+
c,
427+
&format!(
428+
"ordered_array_agg i64 ordered input, \
429+
{rows_per_batch} rows per update_batch"
430+
),
431+
&ordered_batches,
432+
true,
433+
);
434+
let shuffled_batches = create_ordered_array_agg_batches(rows_per_batch, false);
435+
ordered_array_agg_bench(
436+
c,
437+
&format!(
438+
"ordered_array_agg i64 random input, \
439+
{rows_per_batch} rows per update_batch"
440+
),
441+
&shuffled_batches,
442+
false,
443+
);
444+
}
445+
}
446+
criterion_group!(
447+
benches,
448+
array_agg_benchmark,
449+
distinct_array_agg_benchmark,
450+
ordered_array_agg_benchmark
451+
);
292452
criterion_main!(benches);

‎datafusion/functions-aggregate/src/array_agg.rs‎

Lines changed: 121 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1246,6 +1246,10 @@ struct OrderedArrayAggEntry {
12461246
row_idx: usize,
12471247
}
12481248

1249+
/// Coalesce consecutive small payload batches up to this row count to reduce
1250+
/// per-array memory overhead while limiting repeated concatenation costs.
1251+
const ORDERED_ARRAY_AGG_COALESCE_ROWS: usize = 64;
1252+
12491253
#[derive(Debug)]
12501254
pub(crate) struct OrderSensitiveArrayAggAccumulator {
12511255
/// Arrow payload arrays. Entries refer to rows in these batches.
@@ -1531,11 +1535,34 @@ impl OrderSensitiveArrayAggAccumulator {
15311535
}
15321536

15331537
let start = self.entries.len();
1534-
let batch_idx = self.batches.len();
1535-
self.batches.push(values);
1536-
self.entries.extend(
1537-
(0..row_count).map(|row_idx| OrderedArrayAggEntry { batch_idx, row_idx }),
1538-
);
1538+
let (batch_idx, row_offset) = match self.batches.last() {
1539+
Some(last_batch)
1540+
if last_batch.len() + row_count <= ORDERED_ARRAY_AGG_COALESCE_ROWS =>
1541+
{
1542+
let merged =
1543+
arrow::compute::concat(&[last_batch.as_ref(), values.as_ref()])?;
1544+
1545+
// Concatenation preserves the existing row offsets, while new rows
1546+
// start at the previous length of the tail batch.
1547+
let batch_idx = self.batches.len() - 1;
1548+
let row_offset = last_batch.len();
1549+
1550+
self.batches[batch_idx] = merged;
1551+
(batch_idx, row_offset)
1552+
}
1553+
_ => {
1554+
let batch_idx = self.batches.len();
1555+
self.batches.push(values);
1556+
(batch_idx, 0)
1557+
}
1558+
};
1559+
1560+
self.entries
1561+
.extend((0..row_count).map(|row_idx| OrderedArrayAggEntry {
1562+
batch_idx,
1563+
row_idx: row_offset + row_idx,
1564+
}));
1565+
15391566
self.sorted_entry_indices = None;
15401567

15411568
debug_assert_eq!(self.entries.len(), self.ordering_rows.num_rows());
@@ -2513,6 +2540,10 @@ mod tests {
25132540
final_acc.merge_batch(&state_a)?;
25142541
final_acc.merge_batch(&state_b)?;
25152542

2543+
// Small partial-state batches are coalesced into one payload batch.
2544+
assert_eq!(final_acc.batches.len(), 1);
2545+
assert_eq!(final_acc.batches[0].len(), 4);
2546+
25162547
let result = print_nulls(str_arr(final_acc.evaluate()?)?);
25172548
assert_eq!(result, vec!["a1", "a2", "b1", "b2"]);
25182549
Ok(())
@@ -2908,6 +2939,91 @@ mod tests {
29082939
Ok(())
29092940
}
29102941

2942+
#[test]
2943+
fn ordered_array_agg_coalesces_small_batches_up_to_threshold() -> Result<()> {
2944+
use arrow::array::Int64Array;
2945+
2946+
let mut acc = ordered_accumulator(
2947+
DataType::Int64,
2948+
DataType::Int64,
2949+
SortOptions::new(false, false),
2950+
false,
2951+
false,
2952+
)?;
2953+
2954+
let threshold = ORDERED_ARRAY_AGG_COALESCE_ROWS;
2955+
let threshold_i64 = i64::try_from(threshold).unwrap();
2956+
2957+
for values in [
2958+
(0..(threshold_i64 - 1)).collect::<Vec<_>>(),
2959+
vec![threshold_i64 - 1],
2960+
vec![threshold_i64],
2961+
((threshold_i64 + 1)..(threshold_i64 * 2)).collect::<Vec<i64>>(),
2962+
vec![threshold_i64 * 2],
2963+
] {
2964+
let values = Arc::new(Int64Array::from(values)) as ArrayRef;
2965+
2966+
acc.update_batch(&[Arc::clone(&values), values])?;
2967+
}
2968+
2969+
assert_eq!(
2970+
acc.batches
2971+
.iter()
2972+
.map(|batch| batch.len())
2973+
.collect::<Vec<usize>>(),
2974+
vec![threshold, threshold, 1],
2975+
);
2976+
2977+
assert_eq!(acc.entries.len(), threshold * 2 + 1);
2978+
assert_eq!(acc.ordering_rows.num_rows(), threshold * 2 + 1);
2979+
2980+
assert_eq!(acc.entries[threshold - 1].batch_idx, 0);
2981+
assert_eq!(acc.entries[threshold - 1].row_idx, threshold - 1);
2982+
2983+
assert_eq!(acc.entries[threshold].batch_idx, 1);
2984+
assert_eq!(acc.entries[threshold].row_idx, 0);
2985+
2986+
assert_eq!(acc.entries[threshold * 2 - 1].batch_idx, 1);
2987+
assert_eq!(acc.entries[threshold * 2 - 1].row_idx, threshold - 1);
2988+
2989+
assert_eq!(acc.entries[threshold * 2].batch_idx, 2);
2990+
assert_eq!(acc.entries[threshold * 2].row_idx, 0);
2991+
Ok(())
2992+
}
2993+
2994+
#[test]
2995+
fn ordered_array_agg_sorts_across_coalesced_batches() -> Result<()> {
2996+
use arrow::array::Int64Array;
2997+
2998+
let mut acc = ordered_accumulator(
2999+
DataType::Int64,
3000+
DataType::Int64,
3001+
SortOptions::new(false, false),
3002+
false,
3003+
false,
3004+
)?;
3005+
3006+
for value in (0_i64..128).rev() {
3007+
let payload = Arc::new(Int64Array::from(vec![value])) as ArrayRef;
3008+
let ordering = Arc::new(Int64Array::from(vec![value + 1000])) as ArrayRef;
3009+
acc.update_batch(&[payload, ordering])?;
3010+
}
3011+
3012+
let ScalarValue::List(result) = acc.evaluate()? else {
3013+
return internal_err!("expect list");
3014+
};
3015+
3016+
assert_eq!(
3017+
result
3018+
.values()
3019+
.as_primitive::<arrow::datatypes::Int64Type>()
3020+
.values(),
3021+
&(0..128).collect::<Vec<i64>>()
3022+
);
3023+
3024+
Ok(())
3025+
}
3026+
29113027
struct ArrayAggAccumulatorBuilder {
29123028
return_field: FieldRef,
29133029
distinct: bool,

0 commit comments

Comments
 (0)