From d53f78a3d0427d09e68f302b7d5ab733e8a302c8 Mon Sep 17 00:00:00 2001 From: HarshilMital Date: Sat, 24 Oct 2026 10:27:37 +0000 Subject: [PATCH] bench: add nested-column gather cases (#7080) Part of #7080. What's added: shuffle_gather_nested, the BenchRepartitioner place and place+gather harness under CometPartitioning::Hash (50 partitions, 8 x 8192 rows), over four nested shapes: List, List>, Map and Struct. Each schema is an Int64 key plus four columns of the shape. List lengths vary 0-8 and map lengths 1-7 by row; every value is per-row and every batch is built separately. Code path: the row-level gather (interleave_record_batch) that a hash scatter forces on flush, on list and map offsets and their child arrays. The hash reads only the key, and the discarding writer keeps IPC and disk out, so place+gather minus place is the nested gather. 8 bench IDs. Reproduce on the VM (from native/): cargo bench -p datafusion-comet-shuffle --bench shuffle_writer -- 'shuffle_gather_nested' --save-baseline main Results: TODO (VM) --- native/shuffle/benches/shuffle_writer.rs | 178 ++++++++++++++++++++++- 1 file changed, 177 insertions(+), 1 deletion(-) diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index c099fc0413f..28b2ae45e61 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -376,6 +376,182 @@ fn partitioning_benchmark(c: &mut Criterion) { group.finish(); } +/// The row-level gather on nested columns, through the same place+gather harness as +/// [`partitioning_benchmark`] but under hash partitioning, which is what most shuffles use. +/// +/// Hash placement scatters rows, so the flush has to gather every output partition back out a +/// row at a time with `interleave_record_batch`. For a primitive or string column that is a copy +/// of values; for a nested column it also has to rebuild list and map offsets and recurse into +/// every child array, and the fixtures above have no list or map columns at all. The shapes +/// measured here are `List`, `List>`, `Map` and a flat +/// `Struct`. Each schema puts four columns of its shape behind an Int64 +/// key, and the hash reads only the key, so the gap between `place` and `place+gather` is the +/// gather of the nested columns and not the hash of them. IPC encoding and the file write stay +/// out, as above. List and map lengths vary from row to row (0 to 8 and 1 to 7), so the gather +/// cannot fall into a fixed-stride pattern. Every batch is built separately, for the reasons +/// given in [`partitioning_benchmark`]. +fn nested_gather_benchmark(c: &mut Criterion) { + const BATCH_SIZE: usize = 8192; + const NUM_BATCHES: usize = 8; + const NUM_PARTITIONS: usize = 50; + + let partitioning = + CometPartitioning::Hash(vec![Arc::new(Column::new("key", 0))], NUM_PARTITIONS); + let mut group = c.benchmark_group("shuffle_gather_nested"); + for shape in [ + NestedShape::ListInt64, + NestedShape::ListStruct, + NestedShape::MapUtf8Int64, + NestedShape::Struct, + ] { + let batches = (0..NUM_BATCHES) + .map(|b| nested_shape_batch(shape, b * BATCH_SIZE, BATCH_SIZE, 4)) + .collect::>(); + let name = shape.name(); + group.bench_function(format!("place ({name})"), |b| { + b.iter_batched( + || BenchRepartitioner::try_new(partitioning.clone(), BATCH_SIZE).unwrap(), + |mut repartitioner| { + repartitioner.place(&batches).unwrap(); + repartitioner + }, + BatchSize::LargeInput, + ); + }); + group.bench_function(format!("place+gather ({name})"), |b| { + b.iter_batched( + || BenchRepartitioner::try_new(partitioning.clone(), BATCH_SIZE).unwrap(), + |mut repartitioner| { + repartitioner.place(&batches).unwrap(); + repartitioner.gather().unwrap(); + repartitioner + }, + BatchSize::LargeInput, + ); + }); + } + group.finish(); +} + +/// The column shapes of [`nested_gather_benchmark`]. +#[derive(Clone, Copy)] +enum NestedShape { + ListInt64, + ListStruct, + MapUtf8Int64, + Struct, +} + +impl NestedShape { + fn name(self) -> &'static str { + match self { + NestedShape::ListInt64 => "list_int64", + NestedShape::ListStruct => "list_struct", + NestedShape::MapUtf8Int64 => "map_utf8_int64", + NestedShape::Struct => "struct", + } + } +} + +/// An Int64 `key` column, distinct on every row, followed by `num_cols` columns of `shape`. All +/// values run from `first_row`, so successive batches differ. +fn nested_shape_batch( + shape: NestedShape, + first_row: usize, + num_rows: usize, + num_cols: usize, +) -> RecordBatch { + use arrow::array::builder::{Int64Builder, ListBuilder, MapBuilder, StructBuilder}; + use arrow::array::{Float64Array, Int64Array, StringArray, StructArray}; + use arrow::datatypes::Fields; + + let rows = first_row..first_row + num_rows; + let mut fields = vec![Field::new("key", DataType::Int64, false)]; + let mut columns: Vec> = vec![Arc::new(Int64Array::from_iter_values( + rows.clone().map(|row| row as i64), + ))]; + let entry_fields = || { + Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + ]) + }; + for col in 0..num_cols { + let column: Arc = match shape { + NestedShape::ListInt64 => { + let mut builder = ListBuilder::new(Int64Builder::new()); + for row in rows.clone() { + let values = builder.values(); + for i in 0..row % 9 { + values.append_value((row * 16 + i + col) as i64); + } + builder.append(true); + } + Arc::new(builder.finish()) + } + NestedShape::ListStruct => { + let mut builder = ListBuilder::new(StructBuilder::from_fields(entry_fields(), 0)); + for row in rows.clone() { + let entries = builder.values(); + for i in 0..row % 9 { + entries + .field_builder::(0) + .unwrap() + .append_value((row * 16 + i + col) as i64); + entries + .field_builder::(1) + .unwrap() + .append_value(format!("entry {row} {i}")); + entries.append(true); + } + builder.append(true); + } + Arc::new(builder.finish()) + } + NestedShape::MapUtf8Int64 => { + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int64Builder::new()); + for row in rows.clone() { + for i in 0..1 + row % 7 { + builder.keys().append_value(format!("key {i} {}", row % 97)); + builder.values().append_value((row * 16 + i + col) as i64); + } + builder.append(true).unwrap(); + } + Arc::new(builder.finish()) + } + NestedShape::Struct => Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Int64, false)), + Arc::new(Int64Array::from_iter_values( + rows.clone().map(|row| (row + col) as i64), + )) as Arc, + ), + ( + Arc::new(Field::new("b", DataType::Utf8, false)), + Arc::new( + rows.clone() + .map(|row| Some(format!("value {row} {col}"))) + .collect::(), + ) as Arc, + ), + ( + Arc::new(Field::new("c", DataType::Float64, false)), + Arc::new(Float64Array::from_iter_values( + rows.clone().map(|row| row as f64 * 1.5 + col as f64), + )) as Arc, + ), + ])), + }; + fields.push(Field::new( + format!("col{col}"), + column.data_type().clone(), + true, + )); + columns.push(column); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() +} + /// Benchmarks the per-block IPC encoding cost (schema + record batch) in isolation, using the /// `None` codec so that compression does not obscure the schema-encoding cost. Covers a wide flat /// schema and a deeply nested schema, where the schema flatbuffer is largest. @@ -593,6 +769,6 @@ fn config() -> Criterion { criterion_group! { name = benches; config = config(); - targets = criterion_benchmark, partitioning_benchmark, schema_encoding_benchmark, ipc_context_reuse_benchmark + targets = criterion_benchmark, partitioning_benchmark, nested_gather_benchmark, schema_encoding_benchmark, ipc_context_reuse_benchmark } criterion_main!(benches);