From b6d5b489287ff454b3daea0a520a5d6a55f26507 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:59:15 +0300 Subject: [PATCH 1/6] bench: add sort benchmarks for various data profile --- .../benches/sort_preserving_merge.rs | 173 +++++++++++++++++- 1 file changed, 171 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 76ebf230a30e0..35b19aa62bfb7 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -21,15 +21,28 @@ use arrow::{ }; use arrow_schema::{SchemaRef, SortOptions}; use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; +use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; use datafusion_physical_plan::test::TestMemoryExec; use datafusion_physical_plan::{ - collect, sorts::sort_preserving_merge::SortPreservingMergeExec, + collect, execute_stream, sorts::sort_preserving_merge::SortPreservingMergeExec, }; +use futures::StreamExt; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use std::hint::black_box; use std::sync::Arc; +/// Consume the stream batch by batch, dropping each batch as it arrives +/// instead of holding the whole result in memory like `collect` would +async fn drain(mut stream: SendableRecordBatchStream) { + while let Some(batch) = stream.next().await { + black_box(batch.unwrap()); + } +} + const BENCH_ROWS: usize = 1_000_000; // 1 million rows fn get_large_string(idx: usize) -> String { @@ -193,5 +206,161 @@ fn bench_merge_sorted_preserving(c: &mut Criterion) { } } -criterion_group!(benches, bench_merge_sorted_preserving); +// --------------------------------------------------------------------------- +// Benchmarks across data orderings (sorted / nearly sorted / reverse / +// unsorted), sort key types (u64 / string / complex) and payload widths +// (5 / 20 / 100 columns). +// --------------------------------------------------------------------------- + +const NUM_PARTITIONS: usize = 4; +const ROWS_PER_PARTITION: usize = 100_000; +const BATCH_SIZE: usize = 8192; + +const ORDERINGS: [&str; 4] = ["sorted", "nearly_sorted", "reverse", "unsorted"]; +const KEY_TYPES: [&str; 3] = ["u64", "string", "complex"]; +const PAYLOAD_WIDTHS: [usize; 3] = [5, 20, 100]; + +/// Generate the keys in their "arrival" order, before partitioning +fn generate_keys(ordering: &str) -> Vec { + let n = (NUM_PARTITIONS * ROWS_PER_PARTITION) as u64; + let mut rng = StdRng::seed_from_u64(42); + match ordering { + "sorted" => (0..n).collect(), + "reverse" => (0..n).rev().collect(), + // Sorted except for ~1% of items misplaced to random positions + "nearly_sorted" => { + let mut keys: Vec = (0..n).collect(); + for _ in 0..(n / 100) { + let a = rng.random_range(0..n as usize); + let b = rng.random_range(0..n as usize); + keys.swap(a, b); + } + keys + } + "unsorted" => (0..n).map(|_| rng.random_range(0..n)).collect(), + _ => unreachable!(), + } +} + +/// Distribute the arrival sequence round-robin (a batch at a time) over the +/// partitions, then sort each partition's keys, as SortExec would before a +/// sort preserving merge +fn partition_keys(keys: &[u64]) -> Vec> { + let mut partitions = (0..NUM_PARTITIONS) + .map(|_| Vec::with_capacity(ROWS_PER_PARTITION)) + .collect::>(); + for (i, chunk) in keys.chunks(BATCH_SIZE).enumerate() { + partitions[i % NUM_PARTITIONS].extend_from_slice(chunk); + } + for partition in &mut partitions { + partition.sort_unstable(); + } + partitions +} + +fn key_columns(key_type: &str, keys: &[u64]) -> Vec<(String, ArrayRef)> { + let as_string = || { + Arc::new(StringArray::from_iter_values( + keys.iter().map(|k| format!("{k:012}")), + )) as ArrayRef + }; + match key_type { + "u64" => vec![( + "key0".to_string(), + Arc::new(UInt64Array::from(keys.to_vec())) as _, + )], + "string" => vec![("key0".to_string(), as_string())], + // Two sort columns force the row-based (normalized key) cursor + "complex" => vec![ + ( + "key0".to_string(), + Arc::new(UInt64Array::from_iter_values(keys.iter().map(|k| k / 8))) as _, + ), + ("key1".to_string(), as_string()), + ], + _ => unreachable!(), + } +} + +fn create_case( + ordering: &str, + key_type: &str, + payload_width: usize, +) -> (Vec>, SchemaRef, LexOrdering) { + let partitions = partition_keys(generate_keys(ordering).as_slice()) + .into_iter() + .map(|keys| { + keys.chunks(BATCH_SIZE) + .map(|chunk| { + let mut columns = key_columns(key_type, chunk); + // All payload columns share the same buffer, so wide + // payloads don't blow up memory + let payload = Arc::new(UInt64Array::from(chunk.to_vec())) as ArrayRef; + for i in 0..payload_width { + columns.push((format!("col{i}"), Arc::clone(&payload))); + } + RecordBatch::try_from_iter(columns).unwrap() + }) + .collect::>() + }) + .collect::>(); + + let schema = partitions[0][0].schema(); + let sort_order = LexOrdering::new( + schema + .fields() + .iter() + .filter(|field| field.name().starts_with("key")) + .map(|field| { + PhysicalSortExpr::new( + col(field.name(), &schema).unwrap(), + SortOptions::default(), + ) + }), + ) + .unwrap(); + + (partitions, schema, sort_order) +} + +fn bench_spm_data_patterns(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let task_ctx = Arc::new(TaskContext::default()); + + for ordering in ORDERINGS { + for key_type in KEY_TYPES { + for payload_width in PAYLOAD_WIDTHS { + let (partitions, schema, sort_order) = + create_case(ordering, key_type, payload_width); + + c.bench_function( + &format!("spm/{ordering}/{key_type}/payload_{payload_width}"), + |b| { + b.iter(|| { + let exec = TestMemoryExec::try_new_exec( + &partitions, + Arc::clone(&schema), + None, + ) + .unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new( + sort_order.clone(), + exec, + )); + rt.block_on(drain( + execute_stream(merge, Arc::clone(&task_ctx)).unwrap(), + )) + }) + }, + ); + } + } + } +} + +criterion_group!( + benches, + bench_merge_sorted_preserving, + bench_spm_data_patterns +); criterion_main!(benches); From ab7b7af5b0570e07d5ca9784dabbb2db54e51f1b Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:13:35 +0300 Subject: [PATCH 2/6] revert sort bench --- .../benches/sort_preserving_merge.rs | 173 +----------------- 1 file changed, 2 insertions(+), 171 deletions(-) diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 35b19aa62bfb7..76ebf230a30e0 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -21,28 +21,15 @@ use arrow::{ }; use arrow_schema::{SchemaRef, SortOptions}; use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; -use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; use datafusion_physical_plan::test::TestMemoryExec; use datafusion_physical_plan::{ - collect, execute_stream, sorts::sort_preserving_merge::SortPreservingMergeExec, + collect, sorts::sort_preserving_merge::SortPreservingMergeExec, }; -use futures::StreamExt; -use rand::rngs::StdRng; -use rand::{Rng, SeedableRng}; -use std::hint::black_box; use std::sync::Arc; -/// Consume the stream batch by batch, dropping each batch as it arrives -/// instead of holding the whole result in memory like `collect` would -async fn drain(mut stream: SendableRecordBatchStream) { - while let Some(batch) = stream.next().await { - black_box(batch.unwrap()); - } -} - const BENCH_ROWS: usize = 1_000_000; // 1 million rows fn get_large_string(idx: usize) -> String { @@ -206,161 +193,5 @@ fn bench_merge_sorted_preserving(c: &mut Criterion) { } } -// --------------------------------------------------------------------------- -// Benchmarks across data orderings (sorted / nearly sorted / reverse / -// unsorted), sort key types (u64 / string / complex) and payload widths -// (5 / 20 / 100 columns). -// --------------------------------------------------------------------------- - -const NUM_PARTITIONS: usize = 4; -const ROWS_PER_PARTITION: usize = 100_000; -const BATCH_SIZE: usize = 8192; - -const ORDERINGS: [&str; 4] = ["sorted", "nearly_sorted", "reverse", "unsorted"]; -const KEY_TYPES: [&str; 3] = ["u64", "string", "complex"]; -const PAYLOAD_WIDTHS: [usize; 3] = [5, 20, 100]; - -/// Generate the keys in their "arrival" order, before partitioning -fn generate_keys(ordering: &str) -> Vec { - let n = (NUM_PARTITIONS * ROWS_PER_PARTITION) as u64; - let mut rng = StdRng::seed_from_u64(42); - match ordering { - "sorted" => (0..n).collect(), - "reverse" => (0..n).rev().collect(), - // Sorted except for ~1% of items misplaced to random positions - "nearly_sorted" => { - let mut keys: Vec = (0..n).collect(); - for _ in 0..(n / 100) { - let a = rng.random_range(0..n as usize); - let b = rng.random_range(0..n as usize); - keys.swap(a, b); - } - keys - } - "unsorted" => (0..n).map(|_| rng.random_range(0..n)).collect(), - _ => unreachable!(), - } -} - -/// Distribute the arrival sequence round-robin (a batch at a time) over the -/// partitions, then sort each partition's keys, as SortExec would before a -/// sort preserving merge -fn partition_keys(keys: &[u64]) -> Vec> { - let mut partitions = (0..NUM_PARTITIONS) - .map(|_| Vec::with_capacity(ROWS_PER_PARTITION)) - .collect::>(); - for (i, chunk) in keys.chunks(BATCH_SIZE).enumerate() { - partitions[i % NUM_PARTITIONS].extend_from_slice(chunk); - } - for partition in &mut partitions { - partition.sort_unstable(); - } - partitions -} - -fn key_columns(key_type: &str, keys: &[u64]) -> Vec<(String, ArrayRef)> { - let as_string = || { - Arc::new(StringArray::from_iter_values( - keys.iter().map(|k| format!("{k:012}")), - )) as ArrayRef - }; - match key_type { - "u64" => vec![( - "key0".to_string(), - Arc::new(UInt64Array::from(keys.to_vec())) as _, - )], - "string" => vec![("key0".to_string(), as_string())], - // Two sort columns force the row-based (normalized key) cursor - "complex" => vec![ - ( - "key0".to_string(), - Arc::new(UInt64Array::from_iter_values(keys.iter().map(|k| k / 8))) as _, - ), - ("key1".to_string(), as_string()), - ], - _ => unreachable!(), - } -} - -fn create_case( - ordering: &str, - key_type: &str, - payload_width: usize, -) -> (Vec>, SchemaRef, LexOrdering) { - let partitions = partition_keys(generate_keys(ordering).as_slice()) - .into_iter() - .map(|keys| { - keys.chunks(BATCH_SIZE) - .map(|chunk| { - let mut columns = key_columns(key_type, chunk); - // All payload columns share the same buffer, so wide - // payloads don't blow up memory - let payload = Arc::new(UInt64Array::from(chunk.to_vec())) as ArrayRef; - for i in 0..payload_width { - columns.push((format!("col{i}"), Arc::clone(&payload))); - } - RecordBatch::try_from_iter(columns).unwrap() - }) - .collect::>() - }) - .collect::>(); - - let schema = partitions[0][0].schema(); - let sort_order = LexOrdering::new( - schema - .fields() - .iter() - .filter(|field| field.name().starts_with("key")) - .map(|field| { - PhysicalSortExpr::new( - col(field.name(), &schema).unwrap(), - SortOptions::default(), - ) - }), - ) - .unwrap(); - - (partitions, schema, sort_order) -} - -fn bench_spm_data_patterns(c: &mut Criterion) { - let rt = tokio::runtime::Runtime::new().unwrap(); - let task_ctx = Arc::new(TaskContext::default()); - - for ordering in ORDERINGS { - for key_type in KEY_TYPES { - for payload_width in PAYLOAD_WIDTHS { - let (partitions, schema, sort_order) = - create_case(ordering, key_type, payload_width); - - c.bench_function( - &format!("spm/{ordering}/{key_type}/payload_{payload_width}"), - |b| { - b.iter(|| { - let exec = TestMemoryExec::try_new_exec( - &partitions, - Arc::clone(&schema), - None, - ) - .unwrap(); - let merge = Arc::new(SortPreservingMergeExec::new( - sort_order.clone(), - exec, - )); - rt.block_on(drain( - execute_stream(merge, Arc::clone(&task_ctx)).unwrap(), - )) - }) - }, - ); - } - } - } -} - -criterion_group!( - benches, - bench_merge_sorted_preserving, - bench_spm_data_patterns -); +criterion_group!(benches, bench_merge_sorted_preserving); criterion_main!(benches); From 843ce8549dcad65a2fc9470ca0a8aac075cb0438 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:26:16 +0300 Subject: [PATCH 3/6] moved to sort bench --- datafusion/core/benches/sort.rs | 301 ++++++++++++++++++++++++++++---- 1 file changed, 267 insertions(+), 34 deletions(-) diff --git a/datafusion/core/benches/sort.rs b/datafusion/core/benches/sort.rs index 7544f7ae26d43..b2b0574f6874f 100644 --- a/datafusion/core/benches/sort.rs +++ b/datafusion/core/benches/sort.rs @@ -71,7 +71,7 @@ use std::sync::Arc; use arrow::array::StringViewArray; use arrow::{ array::{DictionaryArray, Float64Array, Int64Array, StringArray}, - datatypes::{Int32Type, Schema}, + datatypes::{DataType, Field, Int32Type, Schema}, record_batch::RecordBatch, }; use datafusion::physical_plan::sorts::sort::SortExec; @@ -93,6 +93,7 @@ use criterion::{Criterion, criterion_group, criterion_main}; use futures::StreamExt; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use rand::seq::SliceRandom; use tokio::runtime::Runtime; /// Total number of streams to divide each input into @@ -103,10 +104,74 @@ const NUM_STREAMS: usize = 8; const BATCH_SIZE: usize = 1024; /// Input sizes to benchmark. The small size (100K) exercises the -/// in-memory concat-and-sort path; the large size (10M) exercises +/// in-memory concat-and-sort path; the large size (1M) exercises /// the sort-then-merge path with high fan-in. const INPUT_SIZES: &[(u64, &str)] = &[(100_000, "100k"), (1_000_000, "1M")]; +/// Number of extra (non-sort-key) payload columns to carry alongside the sort +/// keys in the axis benchmarks. Measures the cost of reordering wide batches. +const EXTRA_COLUMN_COUNTS: &[usize] = &[0, 5, 20, 100]; + +/// Input ordering profiles for the SortExec axis benchmarks. +#[derive(Clone, Copy, Debug)] +enum DataProfile { + Sorted, + Unsorted, + /// Fully sorted, then 10% of rows swapped to random positions. + NearlySorted, +} + +impl DataProfile { + fn label(self) -> &'static str { + match self { + DataProfile::Sorted => "sorted", + DataProfile::Unsorted => "unsorted", + DataProfile::NearlySorted => "nearly sorted", + } + } + + /// Arrange `v` (whose initial order is irrelevant) into this profile. + fn apply(self, mut v: Vec) -> Vec { + let mut rng = StdRng::seed_from_u64(99); + match self { + DataProfile::Sorted => v.sort_unstable(), + DataProfile::Unsorted => { + v.shuffle(&mut rng) + } + DataProfile::NearlySorted => { + v.sort_unstable(); + let n = v.len(); + + // 10% is globally misplaced + for _ in 0..n / 10 { + v.swap(rng.random_range(0..n), rng.random_range(0..n)); + } + } + } + v + } +} + +/// Sort-key cardinality, i.e. how much the key values overlap across rows and +/// partitions. Only affects the sort keys, not the extra payload columns. +#[derive(Clone, Copy, Debug)] +enum Cardinality { + /// Heavy overlap: i64 in `0..input_size` (~1/3 duplicates), 100 distinct + /// strings repeated across all rows. + Low, + /// Minimal overlap: full-range i64 and random strings (~no duplicates). + High, +} + +impl Cardinality { + fn label(self) -> &'static str { + match self { + Cardinality::Low => "low card", + Cardinality::High => "high card", + } + } +} + type PartitionedBatches = Vec>; type StreamGenerator = Box PartitionedBatches>; @@ -312,11 +377,15 @@ impl BenchCase { } } -/// Make sort exprs for each column in `schema` +const EXTRA_COLUMN_NAME_PREFIX: &str = "extra_"; + +/// Make sort exprs for each column in `schema`, skipping non-sort payload +/// columns added by [`with_extra_columns`]. fn make_sort_exprs(schema: &Schema) -> LexOrdering { let sort_exprs = schema .fields() .iter() + .filter(|f| !f.name().starts_with(EXTRA_COLUMN_NAME_PREFIX)) .map(|f| PhysicalSortExpr::new_default(col(f.name(), schema).unwrap())); LexOrdering::new(sort_exprs).unwrap() } @@ -328,10 +397,19 @@ fn i64_streams(sorted: bool, input_size: u64) -> PartitionedBatches { values.sort_unstable(); } - split_tuples(values, |v| { - let array = Int64Array::from(v); - RecordBatch::try_from_iter(vec![("i64", Arc::new(array) as _)]).unwrap() - }) + split_tuples(values, build_i64_batch) +} + +/// Build a single-column i64 [`RecordBatch`]. +fn build_i64_batch(v: Vec) -> RecordBatch { + let array = Int64Array::from(v); + RecordBatch::try_from_iter(vec![("i64", Arc::new(array) as _)]).unwrap() +} + +/// Build a single-column utf8 view [`RecordBatch`] under the given column name. +fn build_utf8_view_batch(name: &str, v: Vec>>) -> RecordBatch { + let array: StringViewArray = v.into_iter().collect(); + RecordBatch::try_from_iter(vec![(name, Arc::new(array) as _)]).unwrap() } /// Create streams of f64 (where approximately 1/3 values are repeated) @@ -369,10 +447,7 @@ fn utf8_view_low_cardinality_streams( if sorted { values.sort_unstable(); } - split_tuples(values, |v| { - let array: StringViewArray = v.into_iter().collect(); - RecordBatch::try_from_iter(vec![("utf_view_low", Arc::new(array) as _)]).unwrap() - }) + split_tuples(values, |v| build_utf8_view_batch("utf_view_low", v)) } /// Create streams of high cardinality (~ no duplicates) utf8_view values @@ -384,10 +459,7 @@ fn utf8_view_high_cardinality_streams( if sorted { values.sort_unstable(); } - split_tuples(values, |v| { - let array: StringViewArray = v.into_iter().collect(); - RecordBatch::try_from_iter(vec![("utf_view_high", Arc::new(array) as _)]).unwrap() - }) + split_tuples(values, |v| build_utf8_view_batch("utf_view_high", v)) } /// Create streams of high cardinality (~ no duplicates) utf8 values @@ -485,25 +557,32 @@ fn mixed_tuple_streams(sorted: bool, input_size: u64) -> PartitionedBatches { tuples.sort_unstable(); } - split_tuples(tuples, |tuples| { - let (tuples, i64_values): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); - let (tuples, utf8_low2): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); - let (f64_values, utf8_low1): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); - - let f64_values: Float64Array = f64_values.into_iter().map(|v| v as f64).collect(); - - let utf8_low1: StringArray = utf8_low1.into_iter().collect(); - let utf8_low2: StringArray = utf8_low2.into_iter().collect(); - let i64_values: Int64Array = i64_values.into_iter().collect(); + split_tuples(tuples, build_mixed_tuple_batch) +} - RecordBatch::try_from_iter(vec![ - ("f64", Arc::new(f64_values) as _), - ("utf_low1", Arc::new(utf8_low1) as _), - ("utf_low2", Arc::new(utf8_low2) as _), - ("i64", Arc::new(i64_values) as _), - ]) - .unwrap() - }) +/// The tuple shape used by the `mixed tuple` case: (i64, utf8_low, utf8_low, i64) +type MixedTuple = (((i64, Option>), Option>), i64); + +/// Build a (f64, utf8_low, utf8_low, i64) batch from [`MixedTuple`]s +/// (the leading i64 becomes the f64 column). +fn build_mixed_tuple_batch(tuples: Vec) -> RecordBatch { + let (tuples, i64_values): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); + let (tuples, utf8_low2): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); + let (f64_values, utf8_low1): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); + + let f64_values: Float64Array = f64_values.into_iter().map(|v| v as f64).collect(); + + let utf8_low1: StringArray = utf8_low1.into_iter().collect(); + let utf8_low2: StringArray = utf8_low2.into_iter().collect(); + let i64_values: Int64Array = i64_values.into_iter().collect(); + + RecordBatch::try_from_iter(vec![ + ("f64", Arc::new(f64_values) as _), + ("utf_low1", Arc::new(utf8_low1) as _), + ("utf_low2", Arc::new(utf8_low2) as _), + ("i64", Arc::new(i64_values) as _), + ]) + .unwrap() } /// Create a batch of (f64, utf8_view_low, utf8_view_low, i64) @@ -699,6 +778,30 @@ impl DataGenerator { .map(char::from) .collect::() } + + /// i64 values with the given cardinality (initial order is irrelevant since + /// callers reorder via [`DataProfile::apply`]). + fn i64_values_by(&mut self, card: Cardinality) -> Vec { + match card { + Cardinality::Low => self.i64_values(), + // Full i64 range -> effectively unique (minimal overlap) + Cardinality::High => { + (0..self.input_size).map(|_| self.rng.random()).collect() + } + } + } + + /// utf8 values with the given cardinality, unified to `Arc`. + fn utf8_values_by(&mut self, card: Cardinality) -> Vec>> { + match card { + Cardinality::Low => self.utf8_low_cardinality_values(), + Cardinality::High => self + .utf8_high_cardinality_values() + .into_iter() + .map(|o| o.map(Arc::from)) + .collect(), + } + } } /// Splits the `input` tuples randomly into batches of `BATCH_SIZE` distributed across @@ -733,5 +836,135 @@ where .collect() } -criterion_group!(benches, criterion_benchmark); +type AxisGenerator = Box PartitionedBatches>; + +/// Benchmarks `SortExec` (at the 1M input size) across the following axes: +/// 1. Sort columns +/// - single column with a specialized impl (primitive or byte(view)) +/// - multiple columns, which will use fallback impl +/// 2. Number of columns in the record batch - more columns mean more data to +/// copy while reordering and more memory to hold +/// 3. Value cardinality - whether the sort-key values overlap or not +/// 4. Input ordering - already sorted / unsorted / nearly sorted +fn sort_axis_benchmark(c: &mut Criterion) { + let input_size = 1_000_000u64; + let size_label = "1M"; + + let cases: Vec<(&str, AxisGenerator)> = vec![ + ( + "i64", + Box::new(move |p, card, extra| i64_axis(p, card, extra, input_size)), + ), + ( + "utf8 view", + Box::new(move |p, card, extra| utf8_view_axis(p, card, extra, input_size)), + ), + ( + "mixed tuple", + Box::new(move |p, card, extra| mixed_tuple_axis(p, card, extra, input_size)), + ), + ]; + + for (name, f) in &cases { + for card in [Cardinality::Low, Cardinality::High] { + for &extra in EXTRA_COLUMN_COUNTS { + for profile in [ + DataProfile::Sorted, + DataProfile::Unsorted, + DataProfile::NearlySorted, + ] { + c.bench_function( + &format!( + "sort {name} {size_label} {card:?} cardinality {profile:?} +{extra}cols", + ), + |b| { + let data = f(profile, card, extra); + let case = BenchCase::sort(&data); + b.iter(move || case.run()) + }, + ); + } + } + } + } +} + +/// Single-column i64 batches +fn i64_axis( + profile: DataProfile, + card: Cardinality, + extra: usize, + input_size: u64, +) -> PartitionedBatches { + let values = profile.apply(DataGenerator::new(input_size).i64_values_by(card)); + let batches = split_tuples(values, build_i64_batch); + with_extra_columns(batches, extra) +} + +/// Single-column utf8 view batches +fn utf8_view_axis( + profile: DataProfile, + card: Cardinality, + extra: usize, + input_size: u64, +) -> PartitionedBatches { + let values = profile.apply(DataGenerator::new(input_size).utf8_values_by(card)); + let batches = split_tuples(values, |v| build_utf8_view_batch("utf_view", v)); + with_extra_columns(batches, extra) +} + +/// Multi-column (f64, utf8, utf8, i64) batches. +fn mixed_tuple_axis( + profile: DataProfile, + card: Cardinality, + extra: usize, + input_size: u64, +) -> PartitionedBatches { + let mut data_gen = DataGenerator::new(input_size); + let tuples: Vec = data_gen + .i64_values_by(card) + .into_iter() + .zip(data_gen.utf8_values_by(card)) + .zip(data_gen.utf8_values_by(card)) + .zip(data_gen.i64_values_by(card)) + .collect(); + let batches = split_tuples(profile.apply(tuples), build_mixed_tuple_batch); + with_extra_columns(batches, extra) +} + +/// Append `n` extra non-sort-key `Int64` payload columns +fn with_extra_columns(batches: PartitionedBatches, n: usize) -> PartitionedBatches { + if n == 0 { + return batches; + } + let mut rng = StdRng::seed_from_u64(7); + batches + .into_iter() + .map(|stream| { + stream + .into_iter() + .map(|batch| { + let num_rows = batch.num_rows(); + let mut fields = + batch.schema().fields().iter().cloned().collect::>(); + let mut columns = batch.columns().to_vec(); + for i in 0..n { + let arr = Int64Array::from_iter_values( + (0..num_rows).map(|_| rng.random::()), + ); + columns.push(Arc::new(arr) as _); + fields.push(Arc::new(Field::new( + format!("{}{i}", EXTRA_COLUMN_NAME_PREFIX), + DataType::Int64, + false, + ))); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + }) + .collect() + }) + .collect() +} + +criterion_group!(benches, criterion_benchmark, sort_axis_benchmark); criterion_main!(benches); From 5b3ad342702fe66bb8b0f66aa5fa3509c6191ceb Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:11:01 +0300 Subject: [PATCH 4/6] cleanup --- datafusion/core/benches/sort.rs | 108 ++++++++++++++++++++------------ 1 file changed, 67 insertions(+), 41 deletions(-) diff --git a/datafusion/core/benches/sort.rs b/datafusion/core/benches/sort.rs index b2b0574f6874f..52a2fec380be8 100644 --- a/datafusion/core/benches/sort.rs +++ b/datafusion/core/benches/sort.rs @@ -68,10 +68,10 @@ use std::sync::Arc; -use arrow::array::StringViewArray; +use arrow::array::{ArrayRef, StringViewArray, StringViewBuilder}; use arrow::{ - array::{DictionaryArray, Float64Array, Int64Array, StringArray}, - datatypes::{DataType, Field, Int32Type, Schema}, + array::{Array, DictionaryArray, Float64Array, Int64Array, StringArray}, + datatypes::{Field, Int32Type, Schema}, record_batch::RecordBatch, }; use datafusion::physical_plan::sorts::sort::SortExec; @@ -92,8 +92,8 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; use criterion::{Criterion, criterion_group, criterion_main}; use futures::StreamExt; use rand::rngs::StdRng; -use rand::{Rng, SeedableRng}; use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; use tokio::runtime::Runtime; /// Total number of streams to divide each input into @@ -122,22 +122,12 @@ enum DataProfile { } impl DataProfile { - fn label(self) -> &'static str { - match self { - DataProfile::Sorted => "sorted", - DataProfile::Unsorted => "unsorted", - DataProfile::NearlySorted => "nearly sorted", - } - } - /// Arrange `v` (whose initial order is irrelevant) into this profile. fn apply(self, mut v: Vec) -> Vec { let mut rng = StdRng::seed_from_u64(99); match self { DataProfile::Sorted => v.sort_unstable(), - DataProfile::Unsorted => { - v.shuffle(&mut rng) - } + DataProfile::Unsorted => v.shuffle(&mut rng), DataProfile::NearlySorted => { v.sort_unstable(); let n = v.len(); @@ -163,15 +153,6 @@ enum Cardinality { High, } -impl Cardinality { - fn label(self) -> &'static str { - match self { - Cardinality::Low => "low card", - Cardinality::High => "high card", - } - } -} - type PartitionedBatches = Vec>; type StreamGenerator = Box PartitionedBatches>; @@ -760,10 +741,10 @@ impl DataGenerator { } /// Create sorted values of high cardinality (~ no duplicates) utf8 values - fn utf8_high_cardinality_values(&mut self) -> Vec> { + fn utf8_high_cardinality_values(&mut self) -> Vec>> { // make random strings let mut input = (0..self.input_size) - .map(|_| Some(self.random_string())) + .map(|_| Some(self.random_string().into())) .collect::>(); input.sort_unstable(); @@ -791,15 +772,11 @@ impl DataGenerator { } } - /// utf8 values with the given cardinality, unified to `Arc`. + /// utf8 values with the given cardinality. fn utf8_values_by(&mut self, card: Cardinality) -> Vec>> { match card { Cardinality::Low => self.utf8_low_cardinality_values(), - Cardinality::High => self - .utf8_high_cardinality_values() - .into_iter() - .map(|o| o.map(Arc::from)) - .collect(), + Cardinality::High => self.utf8_high_cardinality_values(), } } } @@ -932,12 +909,58 @@ fn mixed_tuple_axis( with_extra_columns(batches, extra) } -/// Append `n` extra non-sort-key `Int64` payload columns +/// Append `n` extra non-sort-key payload columns to every batch, split across i64, string, string view and dictionary fn with_extra_columns(batches: PartitionedBatches, n: usize) -> PartitionedBatches { if n == 0 { return batches; } let mut rng = StdRng::seed_from_u64(7); + + type Generator = Box ArrayRef>; + + let generators: Vec = vec![ + Box::new(|data_gen: &mut DataGenerator| { + let arr = Int64Array::from_iter_values(data_gen.i64_values()); + + Arc::new(arr) + }), + Box::new(|data_gen: &mut DataGenerator| { + let values = data_gen.utf8_low_cardinality_values(); + let arr: StringArray = values.iter().map(|item| item.as_deref()).collect(); + + Arc::new(arr) + }), + Box::new(|data_gen: &mut DataGenerator| { + let values = data_gen.utf8_low_cardinality_values(); + let mut builder = + StringViewBuilder::with_capacity(values.len()).with_deduplicate_strings(); + for v in values { + builder.append_option(v.as_deref()); + } + + let arr = builder.finish(); + + Arc::new(arr) + }), + Box::new(|data_gen: &mut DataGenerator| { + let values = data_gen.utf8_low_cardinality_values(); + + let arr: DictionaryArray = + values.iter().map(|item| item.as_deref()).collect(); + + Arc::new(arr) + }), + ]; + + let generator_index = (0..n) + .map(|_| rng.random_range(0..generators.len())) + .collect::>(); + + let mut generator = DataGenerator { + input_size: 1 as u64, + rng, + }; + batches .into_iter() .map(|stream| { @@ -948,17 +971,20 @@ fn with_extra_columns(batches: PartitionedBatches, n: usize) -> PartitionedBatch let mut fields = batch.schema().fields().iter().cloned().collect::>(); let mut columns = batch.columns().to_vec(); - for i in 0..n { - let arr = Int64Array::from_iter_values( - (0..num_rows).map(|_| rng.random::()), - ); - columns.push(Arc::new(arr) as _); + generator.input_size = num_rows as u64; + + for (col_index, gen_index) in generator_index.iter().enumerate() { + let gen_fn = &generators[*gen_index]; + + let array = gen_fn(&mut generator); fields.push(Arc::new(Field::new( - format!("{}{i}", EXTRA_COLUMN_NAME_PREFIX), - DataType::Int64, - false, + format!("{EXTRA_COLUMN_NAME_PREFIX}{col_index}"), + array.data_type().clone(), + array.logical_null_count() > 0, ))); + columns.push(array); } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() }) .collect() From a16e46562b9d0c4c62753742e23f892b545e95e0 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:00:07 +0300 Subject: [PATCH 5/6] lint --- datafusion/core/benches/sort.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/benches/sort.rs b/datafusion/core/benches/sort.rs index 52a2fec380be8..a640740df33d5 100644 --- a/datafusion/core/benches/sort.rs +++ b/datafusion/core/benches/sort.rs @@ -957,7 +957,7 @@ fn with_extra_columns(batches: PartitionedBatches, n: usize) -> PartitionedBatch .collect::>(); let mut generator = DataGenerator { - input_size: 1 as u64, + input_size: 1, rng, }; From eb67ee3872d14f6df90efbf0394a6ea7084957d0 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:35:49 +0300 Subject: [PATCH 6/6] format --- datafusion/core/benches/sort.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/datafusion/core/benches/sort.rs b/datafusion/core/benches/sort.rs index a640740df33d5..4c4cb2ea1ec92 100644 --- a/datafusion/core/benches/sort.rs +++ b/datafusion/core/benches/sort.rs @@ -956,10 +956,7 @@ fn with_extra_columns(batches: PartitionedBatches, n: usize) -> PartitionedBatch .map(|_| rng.random_range(0..generators.len())) .collect::>(); - let mut generator = DataGenerator { - input_size: 1, - rng, - }; + let mut generator = DataGenerator { input_size: 1, rng }; batches .into_iter()