Skip to content

bench: first_last remove noisy benchmarks, add update_batch #21487

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alamb merged 7 commits into from
Apr 14, 2026
+172 −75
Merged

bench: first_last remove noisy benchmarks, add update_batch #21487

Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
97beaa1
Remove noisy benchmarks for first_last
theirix Apr 8, 2026
fbef405
Measure using iter_custom
theirix Apr 8, 2026
f6216a3
Add benches for update_batch, merge_batch
theirix Apr 8, 2026
f1e327e
Prepopulate with worst-case ordering
theirix Apr 8, 2026
5d98dc0
Fixup bench name
theirix Apr 9, 2026
4f29c61
Add bench for Accumulator
theirix Apr 9, 2026
1a58f05
Run update/merge multiple times for each accumulator
theirix Apr 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
292 changes: 198 additions & 94 deletions datafusion/functions-aggregate/benches/first_last.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,18 +15,20 @@
// specific language governing permissions and limitations
// under the License.

use std::hint::black_box;
use std::sync::Arc;

use arrow::array::{ArrayRef, BooleanArray};
use arrow::array::{ArrayRef, BooleanArray, Int64Array};
use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Int64Type, Schema};
use arrow::util::bench_util::{create_boolean_array, create_primitive_array};
use datafusion_common::instant::Instant;
use std::hint::black_box;
use std::sync::Arc;

use datafusion_expr::{
Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, function::AccumulatorArgs,
};
use datafusion_functions_aggregate::first_last::{FirstValue, LastValue};
use datafusion_functions_aggregate::first_last::{
FirstValue, LastValue, TrivialFirstValueAccumulator, TrivialLastValueAccumulator,
};
use datafusion_physical_expr::PhysicalSortExpr;
use datafusion_physical_expr::expressions::col;

Expand Down Expand Up @@ -72,92 +74,129 @@ fn prepare_groups_accumulator(is_first: bool) -> Box<dyn GroupsAccumulator> {
}
}

fn prepare_accumulator(is_first: bool) -> Box<dyn Accumulator> {
let schema = Arc::new(Schema::new(vec![
Field::new("value", DataType::Int64, true),
Field::new("ord", DataType::Int64, true),
]));

let order_expr = col("ord", &schema).unwrap();
let sort_expr = PhysicalSortExpr {
expr: order_expr,
options: SortOptions::default(),
};

let value_field: Arc<Field> = Field::new("value", DataType::Int64, true).into();
let accumulator_args = AccumulatorArgs {
return_field: Arc::clone(&value_field),
schema: &schema,
expr_fields: &[value_field],
ignore_nulls: false,
order_bys: std::slice::from_ref(&sort_expr),
is_reversed: false,
name: if is_first {
"FIRST_VALUE(value ORDER BY ord)"
} else {
"LAST_VALUE(value ORDER BY ord)"
},
is_distinct: false,
exprs: &[col("value", &schema).unwrap()],
};

fn create_trivial_accumulator(
is_first: bool,
ignore_nulls: bool,
) -> Box<dyn Accumulator> {
if is_first {
FirstValue::new().accumulator(accumulator_args).unwrap()
Box::new(
TrivialFirstValueAccumulator::try_new(&DataType::Int64, ignore_nulls)
.unwrap(),
)
} else {
LastValue::new().accumulator(accumulator_args).unwrap()
Box::new(
TrivialLastValueAccumulator::try_new(&DataType::Int64, ignore_nulls).unwrap(),
)
}
}

#[expect(clippy::needless_pass_by_value)]
fn convert_to_state_bench(
#[expect(clippy::too_many_arguments)]
fn evaluate_bench(
c: &mut Criterion,
is_first: bool,
emit_to: EmitTo,
name: &str,
values: ArrayRef,
ord: ArrayRef,
opt_filter: Option<&BooleanArray>,
num_groups: usize,
) {
let n = values.len();
let group_indices: Vec<usize> = (0..n).map(|i| i % num_groups).collect();

c.bench_function(name, |b| {
b.iter(|| {
let accumulator = prepare_groups_accumulator(is_first);
black_box(
accumulator
.convert_to_state(std::slice::from_ref(&values), opt_filter)
.unwrap(),
)
b.iter_custom(|iters| {
// Every `evaluate` call mutates the accumulator, so prebuild `iters` accumulators
let mut accumulators: Vec<Box<dyn GroupsAccumulator>> = (0..iters)
Copy link
Copy Markdown
Contributor

@alamb alamb Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are we creating more than one accumulator?

I think the important benchmark is a single accumuatlor and then merge multiple batches per iteration. This seems like it merges one batch per iteration across multiple accumulators?

Copy link
Copy Markdown
Contributor Author

@theirix theirix Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's for scaling runs. One run of evaluate / update_bench / merge_bench is diminishingly fast, so multiple accumulators are created just to run a benchmark function on them. We cannot just run 1000 iterations for evaluate/... since it mutates the state.

However, for update_bench / merge_bench we can merge multiple batches for one bench run - I'll try this

Copy link
Copy Markdown
Contributor Author

@theirix theirix Apr 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are we creating more than one accumulator?

I think the important benchmark is a single accumuatlor and then merge multiple batches per iteration. This seems like it merges one batch per iteration across multiple accumulators?

I've implemented these ideas, but the new update/merge benchmarks are still quite noisy, fluctuating by up to 10%. Still, the effect of UDF optimisation, like in the original PR, is clearly visible.

.map(|_| {
let mut accumulator = prepare_groups_accumulator(is_first);
accumulator
.update_batch(
&[Arc::clone(&values), Arc::clone(&ord)],
&group_indices,
opt_filter,
num_groups,
)
.unwrap();
accumulator
})
.collect();

let start = Instant::now();
for accumulator in &mut accumulators {
black_box(accumulator.evaluate(emit_to).unwrap());
}
start.elapsed()
})
});
}

#[expect(clippy::needless_pass_by_value)]
fn evaluate_accumulator_bench(
fn update_bench(
c: &mut Criterion,
is_first: bool,
name: &str,
values: ArrayRef,
ord: ArrayRef,
opt_filter: Option<&BooleanArray>,
num_groups: usize,
) {
let n = values.len();
let group_indices: Vec<usize> = (0..n).map(|i| i % num_groups).collect();

// Initialize with worst-case ordering so update_batch forces rows comparison for all groups.
let worst_ord: ArrayRef = Arc::new(Int64Array::from(vec![
if is_first {
i64::MAX
} else {
i64::MIN
};
n
]));

c.bench_function(name, |b| {
b.iter_batched(
|| {
// setup, not timed
let mut accumulator = prepare_accumulator(is_first);
accumulator
.update_batch(&[Arc::clone(&values), Arc::clone(&ord)])
.unwrap();
accumulator
},
|mut accumulator| black_box(accumulator.evaluate().unwrap()),
criterion::BatchSize::SmallInput,
)
b.iter_custom(|iters| {
// Every `update_batch` call mutates the accumulator, so prebuild `iters` accumulators.
// Each is pre-populated with worst_ord so all is_sets=true before the timed call.
let mut accumulators: Vec<Box<dyn GroupsAccumulator>> = (0..iters)
.map(|_| {
let mut accumulator = prepare_groups_accumulator(is_first);
accumulator
.update_batch(
&[Arc::clone(&values), Arc::clone(&worst_ord)],
&group_indices,
None, // no filter: ensure all groups are initialised
num_groups,
)
.unwrap();
accumulator
})
.collect();

let start = Instant::now();
for accumulator in &mut accumulators {
#[expect(clippy::unit_arg)]
black_box(
accumulator
.update_batch(
&[Arc::clone(&values), Arc::clone(&ord)],
&group_indices,
opt_filter,
num_groups,
)
.unwrap(),
);
}
start.elapsed()
})
});
}

#[expect(clippy::needless_pass_by_value)]
#[expect(clippy::too_many_arguments)]
fn evaluate_bench(
fn merge_bench(
c: &mut Criterion,
is_first: bool,
emit_to: EmitTo,
name: &str,
values: ArrayRef,
ord: ArrayRef,
Expand All @@ -166,25 +205,76 @@ fn evaluate_bench(
) {
let n = values.len();
let group_indices: Vec<usize> = (0..n).map(|i| i % num_groups).collect();
let is_set: ArrayRef = Arc::new(BooleanArray::from(vec![true; n]));

// Initialize with worst-case ordering so update_batch forces rows comparison for all groups.
let worst_ord: ArrayRef = Arc::new(Int64Array::from(vec![
if is_first {
i64::MAX
} else {
i64::MIN
};
n
]));

c.bench_function(name, |b| {
b.iter_batched(
|| {
// setup, not timed
let mut accumulator = prepare_groups_accumulator(is_first);
accumulator
.update_batch(
&[Arc::clone(&values), Arc::clone(&ord)],
&group_indices,
opt_filter,
num_groups,
)
.unwrap();
accumulator
},
|mut accumulator| black_box(accumulator.evaluate(emit_to).unwrap()),
criterion::BatchSize::SmallInput,
)
b.iter_custom(|iters| {
// Every `merge_batch` call mutates the accumulator, so prebuild `iters` accumulators
let mut accumulators: Vec<Box<dyn GroupsAccumulator>> = (0..iters)
.map(|_| {
let mut accumulator = prepare_groups_accumulator(is_first);
accumulator
.update_batch(
&[Arc::clone(&values), Arc::clone(&worst_ord)],
&group_indices,
opt_filter,
num_groups,
)
.unwrap();
accumulator
})
.collect();

let start = Instant::now();
for accumulator in &mut accumulators {
#[expect(clippy::unit_arg)]
black_box(
accumulator
.merge_batch(
&[Arc::clone(&values), Arc::clone(&ord), Arc::clone(&is_set)],
&group_indices,
opt_filter,
num_groups,
)
.unwrap(),
);
}
start.elapsed()
})
});
}

#[expect(clippy::needless_pass_by_value)]
fn trivial_update_bench(
c: &mut Criterion,
is_first: bool,
ignore_nulls: bool,
name: &str,
values: ArrayRef,
) {
c.bench_function(name, |b| {
b.iter_custom(|iters| {
// The bench is way too fast, so apply 10x factor
let mut accumulators: Vec<Box<dyn Accumulator>> = (0..iters * 10)
.map(|_| create_trivial_accumulator(is_first, ignore_nulls))
.collect();
let start = Instant::now();
for acc in &mut accumulators {
#[expect(clippy::unit_arg)]
black_box(acc.update_batch(&[Arc::clone(&values)]).unwrap());
}
start.elapsed()
})
});
}

Expand All @@ -208,27 +298,10 @@ fn first_last_benchmark(c: &mut Criterion) {
let ord = Arc::new(create_primitive_array::<Int64Type>(N, null_density))
as ArrayRef;

evaluate_accumulator_bench(
c,
is_first,
&format!("{fn_name} evaluate_accumulator_bench nulls={pct}%"),
values.clone(),
ord.clone(),
);

for with_filter in [false, true] {
let filter = create_boolean_array(N, 0.0, 0.5);
let opt_filter = if with_filter { Some(&filter) } else { None };

convert_to_state_bench(
c,
is_first,
&format!(
"{fn_name} convert_to_state nulls={pct}%, filter={with_filter}"
),
values.clone(),
opt_filter,
);
evaluate_bench(
c,
is_first,
Expand All @@ -253,6 +326,37 @@ fn first_last_benchmark(c: &mut Criterion) {
opt_filter,
NUM_GROUPS,
);

update_bench(
c,
is_first,
&format!("{fn_name} update_bench nulls={pct}%, filter={with_filter}"),
values.clone(),
ord.clone(),
opt_filter,
NUM_GROUPS,
);
merge_bench(
c,
is_first,
&format!("{fn_name} merge_bench nulls={pct}%, filter={with_filter}"),
values.clone(),
ord.clone(),
opt_filter,
NUM_GROUPS,
);
}

for ignore_nulls in [false, true] {
trivial_update_bench(
c,
is_first,
ignore_nulls,
&format!(
"{fn_name} trivial_update_bench nulls={pct}%, ignore_nulls={ignore_nulls}"
),
values.clone(),
);
}
}
}
Expand Down
Loading