From bc1a503441f9fc69c3afe623defea85767c21d36 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 31 Aug 2026 13:44:29 +0300 Subject: [PATCH 1/4] bench: multi-column GROUP BY with List Add a benchmark to verify optimizations in GroupColumn support for List types, developed separately. --- .../physical-plan/benches/multi_group_by.rs | 108 +++++++++++++++++- 1 file changed, 106 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 360d311c44633..763d3bc637ec6 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, UInt32Array, }; +use arrow::buffer::OffsetBuffer; use arrow::compute::take; use arrow::datatypes::{ DataType, Field, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, TimeUnit, @@ -798,6 +800,107 @@ fn bench_decimal256(c: &mut Criterion) { group.finish(); } +fn make_list_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_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() +} + +/// Experiment 12: Group count sweep for a `(List, Int32)` key. +/// +/// Exercises the `ListGroupValueBuilder` on the multi-column path (previously +/// such a schema fell back to `GroupValuesRows`). +fn bench_list(c: &mut Criterion) { + let mut group = c.benchmark_group("list"); + group.sample_size(15); + + let schema = make_list_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = generate_list_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 +914,6 @@ criterion_group!( bench_duration, bench_interval, bench_decimal256, + bench_list, ); criterion_main!(benches); From ad79220c9bc940e215e3dc47d5b33f765057d2cc Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 31 Aug 2026 16:52:08 +0300 Subject: [PATCH 2/4] bench: GroupValues for list of strings --- .../physical-plan/benches/multi_group_by.rs | 120 +++++++++++++++++- 1 file changed, 114 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 763d3bc637ec6..18f5714e98ad0 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -30,7 +30,7 @@ use arrow::array::{ ArrayRef, Decimal256Array, DurationMicrosecondArray, Float16Array, Int32Array, - IntervalMonthDayNanoArray, ListArray, UInt32Array, + IntervalMonthDayNanoArray, ListArray, StringArray, UInt32Array, }; use arrow::buffer::OffsetBuffer; use arrow::compute::take; @@ -861,12 +861,79 @@ fn generate_list_batches( .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 (previously -/// such a schema fell back to `GroupValuesRows`). -fn bench_list(c: &mut Criterion) { - let mut group = c.benchmark_group("list"); +/// 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); let schema = make_list_schema(); @@ -901,6 +968,46 @@ fn bench_list(c: &mut Criterion) { 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); + + 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, @@ -914,6 +1021,7 @@ criterion_group!( bench_duration, bench_interval, bench_decimal256, - bench_list, + bench_list_int, + bench_list_utf8, ); criterion_main!(benches); From e04cad03c4aefa7a0a29442ce8e727211a7c4706 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 31 Aug 2026 19:40:43 +0300 Subject: [PATCH 3/4] chore: rename for consistency --- datafusion/physical-plan/benches/multi_group_by.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 18f5714e98ad0..a699ae1aa6795 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -800,7 +800,7 @@ fn bench_decimal256(c: &mut Criterion) { group.finish(); } -fn make_list_schema() -> SchemaRef { +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), @@ -815,7 +815,7 @@ fn make_list_schema() -> SchemaRef { /// (`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_batches( +fn generate_list_int_batches( num_distinct_groups: usize, num_rows: usize, batch_size: usize, @@ -936,10 +936,11 @@ fn bench_list_int(c: &mut Criterion) { let mut group = c.benchmark_group("list_int"); group.sample_size(15); - let schema = make_list_schema(); + let schema = make_list_int_schema(); for num_groups in [1_000, 1_000_000] { - let batches = generate_list_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + let batches = + generate_list_int_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); for vectorized in [true, false] { let label = if vectorized { From 7118b24e006e44e974a8d0e7ac17e5413ad6d0b4 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 1 Sep 2026 00:23:43 +0300 Subject: [PATCH 4/4] bench: tune measurement time --- datafusion/physical-plan/benches/multi_group_by.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index a699ae1aa6795..56570927623b0 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -935,6 +935,7 @@ fn generate_list_utf8_batches( 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(); @@ -975,6 +976,7 @@ fn bench_list_int(c: &mut Criterion) { 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();