diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 360d311c44633..56570927623b0 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -25,12 +25,14 @@ //! implementations — a fair apples-to-apples comparison with the same hashing //! and data layout. Most experiments use `Int32` columns; `bench_fixed_size_binary` //! covers a `(FixedSizeBinary, Int32)` key to exercise the -//! `FixedSizeBinaryGroupValueBuilder`. +//! `FixedSizeBinaryGroupValueBuilder`, and `bench_list` covers a +//! `(List, Int32)` key to exercise the `ListGroupValueBuilder`. use arrow::array::{ ArrayRef, Decimal256Array, DurationMicrosecondArray, Float16Array, Int32Array, - IntervalMonthDayNanoArray, UInt32Array, + IntervalMonthDayNanoArray, ListArray, StringArray, UInt32Array, }; +use arrow::buffer::OffsetBuffer; use arrow::compute::take; use arrow::datatypes::{ DataType, Field, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, TimeUnit, @@ -798,6 +800,217 @@ fn bench_decimal256(c: &mut Criterion) { group.finish(); } +fn make_list_int_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new_list("list", Field::new_list_field(DataType::Int32, true), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(List, Int32)` batches with `num_distinct_groups` distinct +/// keys. +/// +/// Each distinct list is `[g, g]` (a fixed-width sublist keyed on the group +/// id), so the `ListGroupValueBuilder` exercises its child `GroupColumn` +/// (`Int32`) once per element while cardinality stays controlled. The `Int32` +/// column is keyed identically so the combined cardinality equals +/// `num_distinct_groups`. +fn generate_list_int_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + const SUBLIST_LEN: usize = 2; + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + usize::from(remainder > 0); + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids: Vec = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups) + .collect(); + + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n( + SUBLIST_LEN, + current_batch_size, + )); + let values = Int32Array::from_iter_values( + group_ids + .iter() + .flat_map(|&g| std::iter::repeat_n(g as i32, SUBLIST_LEN)), + ); + let list = ListArray::new( + Arc::new(Field::new_list_field(DataType::Int32, true)), + offsets, + Arc::new(values), + None, + ); + let id: Int32Array = group_ids.iter().map(|&g| g as i32).collect(); + + vec![Arc::new(list) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +fn make_list_utf8_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new_list("list", Field::new_list_field(DataType::Utf8, true), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(List, Int32)` batches with `num_distinct_groups` distinct +/// keys. +/// +/// Mirrors [`generate_list_batches`] but with a variable-width child, so the +/// `ListGroupValueBuilder` drives a `ByteGroupValueBuilder` child instead of a +/// fixed-width primitive one. Each distinct list is `["gN", "gN"]`. +fn generate_list_utf8_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + const SUBLIST_LEN: usize = 2; + + // Pool of distinct strings, one per group, so string construction stays out + // of the measured loop. + let pool: Vec = (0..num_distinct_groups).map(|g| format!("g{g}")).collect(); + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + usize::from(remainder > 0); + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids: Vec = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups) + .collect(); + + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n( + SUBLIST_LEN, + current_batch_size, + )); + // `from_iter_values` needs an ExactSizeIterator, which `flat_map` + // is not. + let flat: Vec<&str> = group_ids + .iter() + .flat_map(|&g| std::iter::repeat_n(pool[g].as_str(), SUBLIST_LEN)) + .collect(); + let values = StringArray::from_iter_values(flat); + let list = ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + offsets, + Arc::new(values), + None, + ); + let id: Int32Array = group_ids.iter().map(|&g| g as i32).collect(); + + vec![Arc::new(list) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 12: Group count sweep for a `(List, Int32)` key. +/// +/// Exercises the `ListGroupValueBuilder` on the multi-column path. The +/// `vectorized` baseline before this builder existed was the generic +/// row-encoded `RowsGroupColumn` nested fallback, so this measures the +/// specialized list builder against that. +fn bench_list_int(c: &mut Criterion) { + let mut group = c.benchmark_group("list_int"); + group.sample_size(15); + group.measurement_time(std::time::Duration::from_secs(10)); + + let schema = make_list_int_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_list_int_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 13: Group count sweep for a `(List, Int32)` key. +/// +/// The variable-width counterpart to [`bench_list_int`]. +fn bench_list_utf8(c: &mut Criterion) { + let mut group = c.benchmark_group("list_utf8"); + group.sample_size(15); + group.measurement_time(std::time::Duration::from_secs(10)); + + let schema = make_list_utf8_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_list_utf8_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + criterion_group!( benches, bench_issue_17850_regression, @@ -811,5 +1024,7 @@ criterion_group!( bench_duration, bench_interval, bench_decimal256, + bench_list_int, + bench_list_utf8, ); criterion_main!(benches);