kosiew commented on code in PR #23346:
URL: https://github.com/apache/datafusion/pull/23346#discussion_r3549172415
##########
datafusion/physical-plan/benches/sort_preserving_merge.rs:
##########
@@ -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<u64> {
+ 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<u64> = (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<Vec<u64>> {
+ let mut partitions = (0..NUM_PARTITIONS)
+ .map(|_| Vec::with_capacity(ROWS_PER_PARTITION))
+ .collect::<Vec<_>>();
+ 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<Vec<RecordBatch>>, 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::<Vec<_>>()
+ })
+ .collect::<Vec<_>>();
+
+ 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(|| {
Review Comment:
Nice addition. One small thought: this benchmark builds `TestMemoryExec` and
`SortPreservingMergeExec` inside `b.iter`. The setup cost is probably small
compared with draining 400k rows, but moving it out of the timed section would
make optimizer and runtime comparisons a bit cleaner.
Could we use `iter_batched`, like the benchmark above, so the measured part
focuses more directly on stream execution and drain?
##########
datafusion/physical-plan/benches/sort_preserving_merge.rs:
##########
@@ -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<u64> {
+ 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
Review Comment:
This makes sense, but I think a short comment would help future readers
interpret the benchmark correctly. Since the `nearly_sorted` disorder is
introduced before round-robin batch partitioning and per-partition sorting, SPM
is seeing sorted input partitions whose membership came from a near-sorted
arrival stream, rather than directly nearly-sorted partitions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
