perf(sort): fast path for full batch when it sorts before all other streams in SortPreservingMergeStream

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -2839,10 +2839,11 @@ async fn test_sort_with_streaming_table() -> Result<()> {
     let sql = "SELECT a FROM test_table GROUP BY a ORDER BY a";
     let results = ctx.sql(sql).await?.collect().await?;
 
-    assert_eq!(results.len(), 1);
-    assert_eq!(results[0].num_columns(), 1);
+    // the merge may split the output into several batches, so concat first
+    let results = arrow::compute::concat_batches(&results[0].schema(), &results)?;
+    assert_eq!(results.num_columns(), 1);
     let expected = create_array!(Int32, vec![1, 2, 3]) as ArrayRef;
-    assert_eq!(results[0].column(0), &expected);
+    assert_eq!(results.column(0), &expected);
 
     Ok(())
 }

datafusion/ffi/tests/ffi_udaf.rs

@@ -56,13 +56,16 @@ mod tests {
             .sort_by(vec![col("a")])?;
 
         let result = df.collect().await?;
+        // the merge may split the output into several batches, so concat first
+        let result =
+            arrow::compute::concat_batches(&result[0].schema(), &result).unwrap();
 
         let expected = record_batch!(
             ("a", Int32, vec![1, 2, 4]),
             ("sum_b", Float64, vec![1.0, 4.0, 16.0])
         )?;
 
-        assert_eq!(result[0], expected);
+        assert_eq!(result, expected);
 
         Ok(())
     }
@@ -105,7 +108,10 @@ mod tests {
             .sort_by(vec![col("a")])?;
 
         let result = df.collect().await?;
-        let result = result[0].column_by_name("stddev_b").unwrap();
+        // the merge may split the output into several batches, so concat first
+        let result =
+            arrow::compute::concat_batches(&result[0].schema(), &result).unwrap();
+        let result = result.column_by_name("stddev_b").unwrap();
         let result = result
             .as_any()
             .downcast_ref::<Float64Array>()
@@ -162,13 +168,16 @@ mod tests {
         df.clone().show().await?;
 
         let result = df.collect().await?;
+        // the merge may split the output into several batches, so concat first
+        let result =
+            arrow::compute::concat_batches(&result[0].schema(), &result).unwrap();
 
         let expected = record_batch!(
             ("a", Int32, vec![1, 2, 4]),
             ("sum_b", Float64, vec![2.0, 8.0, 32.0])
         )?;
 
-        assert_eq!(result[0], expected);
+        assert_eq!(result, expected);
 
         Ok(())
     }

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<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: Vec<Vec<u64>> = (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<Vec<RecordBatch>>, SchemaRef, LexOrdering) {
+    let partitions = partition_keys(&generate_keys(ordering))
+        .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(|| {
+                            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);

datafusion/physical-plan/src/sorts/builder.rs

@@ -129,6 +129,23 @@ impl BatchBuilder {
         &self.schema
     }
 
+    /// Consume the rows of `stream_idx`'s current batch that have not been
+    /// pushed via [`Self::push_row`] yet, as a zero-copy slice.
+    ///
+    /// Also releases batches that are no longer referenced: skip-heavy merges
+    /// never reach [`Self::build_record_batch`], which otherwise does that.
+    pub fn take_remaining_rows_of(&mut self, stream_idx: usize) -> Option<RecordBatch> {
+        let cursor = &mut self.cursors[stream_idx];
+        let (_, batch) = &self.batches[cursor.batch_idx];
+        let remaining = batch.num_rows() - cursor.row_idx;
+        let slice = (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining));
+        cursor.row_idx = batch.num_rows();
+        if self.indices.is_empty() {
+            self.release_unreferenced_batches();
+        }
+        slice
+    }
+
     /// Try to interleave all columns using the given index slice.
     fn try_interleave_columns(
         &self,
@@ -164,37 +181,48 @@ impl BatchBuilder {
         // hot path. The retention is bounded and short-lived since leftover
         // rows are drained over subsequent polls.
         if self.indices.is_empty() {
-            // New cursors are only created once the previous cursor for the stream
-            // is finished. This means all remaining rows from all but the last batch
-            // for each stream have been yielded to the newly created record batch
-            //
-            // We can therefore drop all but the last batch for each stream
-            let mut batch_idx = 0;
-            let mut retained = 0;
-            self.batches.retain(|(stream_idx, batch)| {
-                let stream_cursor = &mut self.cursors[*stream_idx];
-                let retain = stream_cursor.batch_idx == batch_idx;
-                batch_idx += 1;
-
-                if retain {
-                    stream_cursor.batch_idx = retained;
-                    retained += 1;
-                } else {
-                    self.batches_mem_used -= get_record_batch_memory_size(batch);
-                }
-                retain
-            });
+            self.release_unreferenced_batches();
         }
 
+        RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into)
+    }
+
+    /// Drop buffered batches no stream cursor references anymore and release
+    /// their memory back to the pool.
+    ///
+    /// Must only be called when `indices` is empty, since indices reference
+    /// batches by position.
+    fn release_unreferenced_batches(&mut self) {
+        assert_eq!(self.indices.len(), 0);
+
+        // New cursors are only created once the previous cursor for the stream
+        // is finished. This means all remaining rows from all but the last batch
+        // for each stream have been yielded to the newly created record batch
+        //
+        // We can therefore drop all but the last batch for each stream
+        let mut batch_idx = 0;
+        let mut retained = 0;
+        self.batches.retain(|(stream_idx, batch)| {
+            let stream_cursor = &mut self.cursors[*stream_idx];
+            let retain = stream_cursor.batch_idx == batch_idx;
+            batch_idx += 1;
+
+            if retain {
+                stream_cursor.batch_idx = retained;
+                retained += 1;
+            } else {
+                self.batches_mem_used -= get_record_batch_memory_size(batch);
+            }
+            retain
+        });
+
         // Release excess memory back to the pool, but never shrink below
         // initial_reservation to maintain the anti-starvation guarantee
         // for the merge phase.
         let target = self.batches_mem_used.max(self.initial_reservation);
         if self.reservation.size() > target {
             self.reservation.shrink(self.reservation.size() - target);
         }
-
-        RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into)
     }
 
     /// Drains the in_progress row indexes, and builds a new RecordBatch from them