From f2c9d5ba1f30e4c15092fb6f9d3e7d17d2e75ea1 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Fri, 21 Aug 2026 14:54:09 +0800 Subject: [PATCH 1/2] bench(functions-aggregate): add coalesce-peers first_value comparison Adds coalesce_comparison_bench to the first_last benchmark: N independent primitive first_value accumulators (the pre-rewrite plan) vs one struct-valued accumulator carrying the same N columns (the post-rewrite plan produced by the CoalesceFirstLast optimizer rule). The struct path uses the native nested GroupsAccumulator from #23628. Runs as 'first_value coalesce_peers(i64,utf8,f64) separate x3' vs '... coalesced struct' so the head-to-head win is visible in the same run. Part of #23600. --- .../functions-aggregate/benches/first_last.rs | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/datafusion/functions-aggregate/benches/first_last.rs b/datafusion/functions-aggregate/benches/first_last.rs index 235f11ff30f63..9e21d4a89da97 100644 --- a/datafusion/functions-aggregate/benches/first_last.rs +++ b/datafusion/functions-aggregate/benches/first_last.rs @@ -452,6 +452,103 @@ fn create_list_of_struct_array(n: usize, list_len: usize, null_density: f32) -> )) } +/// Head-to-head for the coalesce-peers rewrite: N independent primitive +/// `first_value` accumulators (the pre-rewrite plan) vs one struct-valued +/// accumulator carrying the same N columns (the post-rewrite plan). Uses the +/// same worst-case ordering as [`update_bench`] so every row forces an +/// ordering comparison in every accumulator. +#[expect(clippy::needless_pass_by_value)] +fn coalesce_comparison_bench( + c: &mut Criterion, + name: &str, + column_values: Vec, + struct_values: ArrayRef, + ord: ArrayRef, + num_groups: usize, +) { + let n = ord.len(); + let group_indices: Vec = (0..n).map(|i| i % num_groups).collect(); + let worst_ord: ArrayRef = Arc::new(Int64Array::from(vec![i64::MAX; n])); + + // Pre-rewrite: one accumulator per column. + c.bench_function(&format!("{name} separate x{}", column_values.len()), |b| { + b.iter_batched( + || { + column_values + .iter() + .map(|values| { + let mut acc = prepare_typed_groups_accumulator( + true, + values.data_type().clone(), + ); + acc.update_batch( + &[Arc::clone(values), Arc::clone(&worst_ord)], + &group_indices, + None, + num_groups, + ) + .unwrap(); + acc + }) + .collect::>() + }, + |mut accumulators| { + for _ in 0..100 { + for (acc, values) in accumulators.iter_mut().zip(&column_values) { + #[expect(clippy::unit_arg)] + black_box( + acc.update_batch( + &[Arc::clone(values), Arc::clone(&ord)], + &group_indices, + None, + num_groups, + ) + .unwrap(), + ); + } + } + }, + BatchSize::SmallInput, + ) + }); + + // Post-rewrite: a single struct-valued accumulator. + c.bench_function(&format!("{name} coalesced struct"), |b| { + b.iter_batched( + || { + let mut acc = prepare_typed_groups_accumulator( + true, + struct_values.data_type().clone(), + ); + acc.update_batch( + &[Arc::clone(&struct_values), Arc::clone(&worst_ord)], + &group_indices, + None, + num_groups, + ) + .unwrap(); + acc + }, + |mut accumulator| { + for _ in 0..100 { + #[expect(clippy::unit_arg)] + black_box( + accumulator + .update_batch( + &[Arc::clone(&struct_values), Arc::clone(&ord)], + &group_indices, + None, + num_groups, + ) + .unwrap(), + ); + } + }, + BatchSize::SmallInput, + ) + }); +} + fn first_last_nested_benchmark(c: &mut Criterion) { const N: usize = 65536; const NUM_GROUPS: usize = 1024; @@ -510,6 +607,19 @@ fn first_last_nested_benchmark(c: &mut Criterion) { ); } } + // Coalesce-peers head-to-head on null-free columns. + let a = Arc::new(create_primitive_array::(N, 0.0)) as ArrayRef; + let b = Arc::new(create_string_array_with_len::(N, 0.0, 16)) as ArrayRef; + let d = Arc::new(create_primitive_array::(N, 0.0)) as ArrayRef; + let struct_values = create_struct_array(N, 0.0); + coalesce_comparison_bench( + c, + "first_value coalesce_peers(i64,utf8,f64)", + vec![a, b, d], + struct_values, + ord, + NUM_GROUPS, + ); } fn first_last_benchmark(c: &mut Criterion) { From 7b438cbc941338a301dbbc1ffd69f1ba16a3a4fd Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Fri, 21 Aug 2026 22:54:17 +0800 Subject: [PATCH 2/2] bench: add winner-changing variant to coalesce-peers comparison The existing coalesce-peers head-to-head reused one `ord` array every iteration, so after the first iteration it only exercised compare-and-reject (the pure ordering-comparison path). Add a `(winner changes)` variant that feeds a strictly-decreasing `ord` per iteration, forcing the running value to be replaced+copied on every row, and label the original `(winner stable)`. This covers both the compare-reject and compare-replace paths, so the struct plan's wider running-value copy is measured too, per review feedback. --- .../functions-aggregate/benches/first_last.rs | 179 ++++++++++++------ 1 file changed, 124 insertions(+), 55 deletions(-) diff --git a/datafusion/functions-aggregate/benches/first_last.rs b/datafusion/functions-aggregate/benches/first_last.rs index 9e21d4a89da97..2b79ab1c4632a 100644 --- a/datafusion/functions-aggregate/benches/first_last.rs +++ b/datafusion/functions-aggregate/benches/first_last.rs @@ -452,68 +452,70 @@ fn create_list_of_struct_array(n: usize, list_len: usize, null_density: f32) -> )) } -/// Head-to-head for the coalesce-peers rewrite: N independent primitive -/// `first_value` accumulators (the pre-rewrite plan) vs one struct-valued -/// accumulator carrying the same N columns (the post-rewrite plan). Uses the -/// same worst-case ordering as [`update_bench`] so every row forces an -/// ordering comparison in every accumulator. -#[expect(clippy::needless_pass_by_value)] -fn coalesce_comparison_bench( +/// Drive one coalesce-peers variant: `separate xN` (pre-rewrite, one primitive +/// accumulator per column) vs `coalesced struct` (post-rewrite, one struct +/// accumulator). Each accumulator is primed with `worst_ord` (the always-losing +/// ordering value), then fed one `ord` per iteration from `iter_ords`. Varying +/// `iter_ords` selects which path dominates — see [`coalesce_comparison_bench`]. +#[expect(clippy::too_many_arguments)] +fn run_coalesce_variant( c: &mut Criterion, name: &str, - column_values: Vec, - struct_values: ArrayRef, - ord: ArrayRef, + variant: &str, + column_values: &[ArrayRef], + struct_values: &ArrayRef, + worst_ord: &ArrayRef, + iter_ords: &[ArrayRef], + group_indices: &[usize], num_groups: usize, ) { - let n = ord.len(); - let group_indices: Vec = (0..n).map(|i| i % num_groups).collect(); - let worst_ord: ArrayRef = Arc::new(Int64Array::from(vec![i64::MAX; n])); - // Pre-rewrite: one accumulator per column. - c.bench_function(&format!("{name} separate x{}", column_values.len()), |b| { - b.iter_batched( - || { - column_values - .iter() - .map(|values| { - let mut acc = prepare_typed_groups_accumulator( - true, - values.data_type().clone(), - ); - acc.update_batch( - &[Arc::clone(values), Arc::clone(&worst_ord)], - &group_indices, - None, - num_groups, - ) - .unwrap(); - acc - }) - .collect::>() - }, - |mut accumulators| { - for _ in 0..100 { - for (acc, values) in accumulators.iter_mut().zip(&column_values) { - #[expect(clippy::unit_arg)] - black_box( + c.bench_function( + &format!("{name} separate x{} {variant}", column_values.len()), + |b| { + b.iter_batched( + || { + column_values + .iter() + .map(|values| { + let mut acc = prepare_typed_groups_accumulator( + true, + values.data_type().clone(), + ); acc.update_batch( - &[Arc::clone(values), Arc::clone(&ord)], - &group_indices, + &[Arc::clone(values), Arc::clone(worst_ord)], + group_indices, None, num_groups, ) - .unwrap(), - ); + .unwrap(); + acc + }) + .collect::>() + }, + |mut accumulators| { + for ord in iter_ords { + for (acc, values) in accumulators.iter_mut().zip(column_values) { + #[expect(clippy::unit_arg)] + black_box( + acc.update_batch( + &[Arc::clone(values), Arc::clone(ord)], + group_indices, + None, + num_groups, + ) + .unwrap(), + ); + } } - } - }, - BatchSize::SmallInput, - ) - }); + }, + BatchSize::SmallInput, + ) + }, + ); // Post-rewrite: a single struct-valued accumulator. - c.bench_function(&format!("{name} coalesced struct"), |b| { + c.bench_function(&format!("{name} coalesced struct {variant}"), |b| { b.iter_batched( || { let mut acc = prepare_typed_groups_accumulator( @@ -521,8 +523,8 @@ fn coalesce_comparison_bench( struct_values.data_type().clone(), ); acc.update_batch( - &[Arc::clone(&struct_values), Arc::clone(&worst_ord)], - &group_indices, + &[Arc::clone(struct_values), Arc::clone(worst_ord)], + group_indices, None, num_groups, ) @@ -530,13 +532,13 @@ fn coalesce_comparison_bench( acc }, |mut accumulator| { - for _ in 0..100 { + for ord in iter_ords { #[expect(clippy::unit_arg)] black_box( accumulator .update_batch( - &[Arc::clone(&struct_values), Arc::clone(&ord)], - &group_indices, + &[Arc::clone(struct_values), Arc::clone(ord)], + group_indices, None, num_groups, ) @@ -549,6 +551,73 @@ fn coalesce_comparison_bench( }); } +/// Head-to-head for the coalesce-peers rewrite: N independent primitive +/// `first_value` accumulators (the pre-rewrite plan) vs one struct-valued +/// accumulator carrying the same N columns (the post-rewrite plan). Runs two +/// variants, since the two plans differ on two separable costs: +/// +/// - `(winner stable)` reuses one `ord` every iteration. After the first +/// iteration the running winner already holds the best value in `ord`, so +/// this is dominated by compare-and-reject — the per-row ordering comparison +/// the rewrite collapses (N compares -> 1). +/// - `(winner changes)` feeds a distinct, strictly-decreasing `ord` per +/// iteration so every row becomes a new winner (smaller ord wins; `worst_ord` +/// = i64::MAX primes the loser). This forces the running value to be +/// replaced+copied on every row, exercising the update path the winner-stable +/// case skips — where the struct plan copies one wider row vs N narrow ones. +#[expect(clippy::needless_pass_by_value)] +fn coalesce_comparison_bench( + c: &mut Criterion, + name: &str, + column_values: Vec, + struct_values: ArrayRef, + ord: ArrayRef, + num_groups: usize, +) { + const ITERS: usize = 100; + let n = ord.len(); + let group_indices: Vec = (0..n).map(|i| i % num_groups).collect(); + let worst_ord: ArrayRef = Arc::new(Int64Array::from(vec![i64::MAX; n])); + + // Winner-stable: the same `ord` every iteration. + let stable_ords: Vec = (0..ITERS).map(|_| Arc::clone(&ord)).collect(); + + // Winner-changing: strictly decreasing across (iteration, row), and below + // `worst_ord`, so every row wins and updates the running value. Built once, + // outside the timed loop. + let changing_ords: Vec = (0..ITERS) + .map(|k| { + let base = i64::MAX - 1 - (k as i64) * (n as i64); + Arc::new(Int64Array::from( + (0..n as i64).map(|i| base - i).collect::>(), + )) as ArrayRef + }) + .collect(); + + run_coalesce_variant( + c, + name, + "(winner stable)", + &column_values, + &struct_values, + &worst_ord, + &stable_ords, + &group_indices, + num_groups, + ); + run_coalesce_variant( + c, + name, + "(winner changes)", + &column_values, + &struct_values, + &worst_ord, + &changing_ords, + &group_indices, + num_groups, + ); +} + fn first_last_nested_benchmark(c: &mut Criterion) { const N: usize = 65536; const NUM_GROUPS: usize = 1024;