fix: split multi-column sort runs into smaller batches under a memory limit

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

@@ -30,10 +30,12 @@ use crate::metrics::{
 };
 use crate::projection::{ProjectionExec, make_with_child, update_ordering};
 use crate::sorts::IncrementalSortIterator;
-use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
-use crate::spill::get_record_batch_memory_size;
+use crate::sorts::streaming_merge::{
+    SortedSpillFile, StreamingMergeBuilder, merges_in_row_format,
+};
 use crate::spill::in_progress_spill_file::InProgressSpillFile;
 use crate::spill::spill_manager::{GetSlicedSize, SpillManager};
+use crate::spill::{gc_view_arrays, get_record_batch_memory_size};
 use crate::statistics::{ChildStats, StatisticsArgs, with_per_partition_fetch};
 use crate::stream::ReservationStream;
 use crate::stream::{ObservedStream, RecordBatchStreamAdapter};
@@ -60,7 +62,7 @@ use datafusion_common::{
 use datafusion_common_runtime::SpawnedTask;
 use datafusion_execution::TaskContext;
 use datafusion_execution::memory_pool::{
-    MemoryConsumer, MemoryPool, MemoryReservation, MergeMemoryPool,
+    MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, MergeMemoryPool,
 };
 use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_physical_expr::LexOrdering;
@@ -708,7 +710,8 @@ impl ExternalSorter {
         if self.in_mem_batches.len() == 1 {
             let batch = self.in_mem_batches.swap_remove(0);
             let reservation = self.reservation.take();
-            let sorted_stream = self.sort_batch_stream(batch, reservation, None)?;
+            let sorted_stream =
+                self.sort_batch_stream(batch, reservation, None, self.batch_size)?;
             return Ok(self.observe_if_output(sorted_stream, is_output_stream));
         }
 
@@ -725,7 +728,8 @@ impl ExternalSorter {
                 .try_resize(get_reserved_bytes_for_record_batch(&batch)?)
                 .map_err(Self::err_with_oom_context)?;
             let reservation = self.reservation.take();
-            let sorted_stream = self.sort_batch_stream(batch, reservation, None)?;
+            let sorted_stream =
+                self.sort_batch_stream(batch, reservation, None, self.batch_size)?;
             return Ok(self.observe_if_output(sorted_stream, is_output_stream));
         }
 
@@ -781,13 +785,19 @@ impl ExternalSorter {
         let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?;
         let shared_input =
             shared_input_size.map(|size| Arc::new(self.reservation.split(size)));
+        let batch_size = self.merge_run_batch_size();
         let jobs =
             runs.into_iter()
                 .zip(sizes)
                 .enumerate()
                 .map(move |(index, (batch, size))| {
                     let reservation = self.reservation.split(size);
-                    let sort = self.sort_run(batch, reservation, shared_input.clone())?;
+                    let sort = self.sort_run(
+                        batch,
+                        reservation,
+                        shared_input.clone(),
+                        batch_size,
+                    )?;
                     let permits = Arc::clone(&concurrency.permits);
                     Ok::<_, DataFusionError>(async move {
                         let stream = Self::prepare_sort_run(sort, permits).await?;
@@ -969,11 +979,12 @@ impl ExternalSorter {
         let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?;
         let shared_input =
             shared_input_size.map(|size| Arc::new(self.reservation.split(size)));
+        let batch_size = self.merge_run_batch_size();
         runs.into_iter()
             .zip(sizes)
             .map(|(run, size)| {
                 let reservation = self.reservation.split(size);
-                self.sort_batch_stream(run, reservation, shared_input.clone())
+                self.sort_batch_stream(run, reservation, shared_input.clone(), batch_size)
             })
             .collect()
     }
@@ -996,11 +1007,11 @@ impl ExternalSorter {
         batch: RecordBatch,
         reservation: MemoryReservation,
         shared_input: Option<Arc<MemoryReservation>>,
+        batch_size: usize,
     ) -> Result<SendableRecordBatchStream> {
         let schema = batch.schema();
-        let stream =
-            futures::stream::once(self.sort_run(batch, reservation, shared_input)?)
-                .try_flatten();
+        let sort = self.sort_run(batch, reservation, shared_input, batch_size)?;
+        let stream = futures::stream::once(sort).try_flatten();
         Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
     }
 
@@ -1013,6 +1024,7 @@ impl ExternalSorter {
         batch: RecordBatch,
         reservation: MemoryReservation,
         shared_input: Option<Arc<MemoryReservation>>,
+        batch_size: usize,
     ) -> Result<impl Future<Output = Result<SendableRecordBatchStream>> + Send + use<>>
     {
         let reserved_bytes = get_reserved_bytes_for_record_batch(&batch)?;
@@ -1023,9 +1035,10 @@ impl ExternalSorter {
         }
 
         let expressions = self.expr.clone();
-        let batch_size = self.batch_size;
         let merge_pool = Arc::clone(&self.merge_pool);
         let elapsed_compute = self.metrics.baseline.elapsed_compute().clone();
+        // Set when `merge_run_batch_size` split this run
+        let split = batch_size < self.batch_size;
 
         Ok(async move {
             // Held until `batch` is sorted, when this block ends. The last run to
@@ -1036,7 +1049,16 @@ impl ExternalSorter {
             let schema = batch.schema();
 
             // Sort the batch immediately and get all output batches
-            let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?;
+            let mut sorted_batches =
+                sort_batch_chunked(&batch, &expressions, batch_size)?;
+            // `take` keeps the data buffers of view arrays, so splitting would
+            // charge every batch for all of them.
+            if split && sorted_batches.len() > 1 {
+                sorted_batches = sorted_batches
+                    .iter()
+                    .map(gc_view_arrays)
+                    .collect::<Result<_>>()?;
+            }
 
             // Chunked output can retain shared buffers in every batch and
             // exceed the input estimate. Borrow only already-reserved spill
@@ -1078,6 +1100,31 @@ impl ExternalSorter {
         })
     }
 
+    /// Batch size of the sorted runs that feed a merge. A merge in the row format
+    /// holds the keys of every run's current batch at once, and these keys can be
+    /// larger than the batches. With a memory limit, runs therefore use an eighth
+    /// of `batch_size`, but at least 1024 rows to keep the per-batch merge cost
+    /// small. The merge output still uses `batch_size`. Runs keep `batch_size`
+    /// when their batches would share buffers other than the data of top-level
+    /// view columns, such as dictionary values, since every batch is charged for
+    /// them.
+    fn merge_run_batch_size(&self) -> usize {
+        let limited = !matches!(
+            self.runtime.memory_pool.memory_limit(),
+            MemoryLimit::Infinite
+        );
+        let shares_buffers = self.schema.fields().iter().any(|field| {
+            let data_type = field.data_type();
+            sorted_copy_keeps_buffers(data_type)
+                && !matches!(data_type, DataType::Utf8View | DataType::BinaryView)
+        });
+        if limited && !shares_buffers && merges_in_row_format(&self.expr, &self.schema) {
+            (self.batch_size / 8).max(1024).min(self.batch_size)
+        } else {
+            self.batch_size
+        }
+    }
+
     /// If this sort may spill, pre-allocates
     /// `sort_spill_reservation_bytes` of memory to guarantee memory
     /// left for the in memory sort/merge.
@@ -2557,7 +2604,8 @@ mod tests {
             let reservation = sorter.reservation.new_empty();
             reservation.try_grow(get_reserved_bytes_for_record_batch(&batch)?)?;
             let before = elapsed_compute.value();
-            let sort = sorter.sort_run(batch.clone(), reservation, None)?;
+            let sort =
+                sorter.sort_run(batch.clone(), reservation, None, sorter.batch_size)?;
             assert_eq!(elapsed_compute.value(), before);
 
             let stream = if spawned {

datafusion/physical-plan/src/sorts/sort/spill_tests.rs

@@ -9,10 +9,10 @@ use crate::metrics::ExecutionPlanMetricsSet;
 use crate::spill::get_record_batch_memory_size;
 use crate::spill::spill_manager::GetSlicedSize;
 use arrow::array::{
-    ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int64Array, StringArray,
-    StringViewArray,
+    ArrayRef, Decimal128Array, DictionaryArray, FixedSizeBinaryArray, Int8Array,
+    Int32Array, Int64Array, StringArray, StringViewArray,
 };
-use arrow::compute::{SortOptions, concat_batches};
+use arrow::compute::{SortOptions, cast, concat_batches};
 use arrow::datatypes::{DataType, Field, Int8Type, Schema};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::config::{ExecutionOptions, SpillCompression};
@@ -679,6 +679,107 @@ async fn test_sort_preparation_keeps_output_reserved() -> Result<()> {
     Ok(())
 }
 
+/// Merging sorted runs in the row format holds the keys of every run's current
+/// batch. Two `Int32` keys encode to 18 bytes per row against 8 bytes in the
+/// batches, and a single `FixedSizeBinary(16)` key also merges in the row
+/// format, so merging full-size runs needs more memory than the sort reserves
+/// while buffering them. The `Utf8View` and dictionary payloads keep buffers
+/// shared after sorting, which must not be charged once per split batch.
+async fn check_merge_keys_larger_than_batches() -> Result<()> {
+    let rows = 8192;
+    let int_keys = |batch: i32| -> Vec<ArrayRef> {
+        vec![
+            Arc::new(Int32Array::from_iter_values(
+                (0..rows).map(|i| (i * 31 + batch) % 1000),
+            )),
+            Arc::new(Int32Array::from_iter_values((0..rows).rev())),
+        ]
+    };
+    let int_fields = [("a", DataType::Int32), ("b", DataType::Int32)];
+    check_merge_keys(&int_fields, 2, |batch| Ok(int_keys(batch))).await?;
+    // Longer than 12 bytes, so views keep them in data buffers.
+    let strings = |batch: i32| -> StringArray {
+        (0..rows)
+            .map(|i| Some(format!("{batch}-{i:0>64}")))
+            .collect()
+    };
+    let dictionary =
+        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
+    for payload in [DataType::Utf8View, dictionary] {
+        let fields = [&int_fields[..], &[("c", payload.clone())]].concat();
+        check_merge_keys(&fields, 2, |batch| {
+            let payload = cast(&strings(batch), &payload)?;
+            Ok([int_keys(batch), vec![payload]].concat())
+        })
+        .await?;
+    }
+    check_merge_keys(&[("k", DataType::FixedSizeBinary(16))], 1, |batch| {
+        let keys = (0..rows).map(|i| u128::from(((i * 31 + batch) % 1000) as u32));
+        let keys = FixedSizeBinaryArray::try_from_iter(keys.map(u128::to_be_bytes))?;
+        Ok(vec![Arc::new(keys)])
+    })
+    .await
+}
+
+/// Sorts 64 batches of `columns` by the first `keys` fields with room to buffer
+/// 16 of them, then checks the spilled output and that all memory is released.
+async fn check_merge_keys(
+    fields: &[(&str, DataType)],
+    keys: usize,
+    columns: impl Fn(i32) -> Result<Vec<ArrayRef>>,
+) -> Result<()> {
+    let schema = Arc::new(Schema::new(
+        fields
+            .iter()
+            .map(|(name, data_type)| Field::new(*name, data_type.clone(), false))
+            .collect::<Vec<_>>(),
+    ));
+    let batches = (0..64)
+        .map(|batch| Ok(RecordBatch::try_new(Arc::clone(&schema), columns(batch)?)?))
+        .collect::<Result<Vec<_>>>()?;
+    let ordering =
+        LexOrdering::new((0..keys).map(|i| {
+            PhysicalSortExpr::new_default(Arc::new(Column::new(fields[i].0, i)))
+        }))
+        .expect("at least one sort key");
+    let capacity = MERGE_BYTES + 16 * get_reserved_bytes_for_record_batch(&batches[0])?;
+    let pool = AdjustablePool::new(capacity);
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::clone(&pool) as Arc<dyn MemoryPool>)
+        .build_arc()?;
+    let mut sorter = ExternalSorter::new(
+        0,
+        Arc::clone(&schema),
+        ordering.clone(),
+        batches[0].num_rows(),
+        MERGE_BYTES,
+        0,
+        SpillCompression::Uncompressed,
+        &ExecutionPlanMetricsSet::new(),
+        Arc::clone(&runtime),
+    )?;
+    for batch in &batches {
+        sorter.insert_batch(batch.clone()).await?;
+    }
+    assert!(sorter.spilled_before());
+    let output: Vec<RecordBatch> = sorter.sort().await?.try_collect().await?;
+    let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?;
+    assert_eq!(concat_batches(&schema, &output)?, expected);
+    drop(sorter);
+    assert_released(&pool, &runtime).await;
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_merge_keys_larger_than_batches() -> Result<()> {
+    check_merge_keys_larger_than_batches().await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_merge_keys_larger_than_batches_multithreaded() -> Result<()> {
+    check_merge_keys_larger_than_batches().await
+}
+
 #[tokio::test]
 async fn test_single_batch_spill_preserves_workspace_after_limit_decreases() -> Result<()>
 {

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

@@ -13,7 +13,7 @@ use crate::sorts::{
 };
 use crate::{EmptyRecordBatchStream, SendableRecordBatchStream, SpillManager};
 use arrow::array::*;
-use arrow::datatypes::{DataType, SchemaRef};
+use arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion_common::human_readable_size;
 use datafusion_common::{Result, assert_or_internal_err, internal_err};
 use datafusion_execution::SpillFile;
@@ -263,7 +263,8 @@ impl<'a> StreamingMergeBuilder<'a> {
         let reservation =
             reservation.expect("Reservation cannot be empty for streaming merge");
 
-        // Special case single column comparisons with optimized cursor implementations
+        // Special case single column comparisons with optimized cursor implementations.
+        // `merges_in_row_format` must match this choice.
         if expressions.len() == 1 {
             let sort = expressions[0].clone();
             let data_type = sort.expr.data_type(schema.as_ref())?;
@@ -297,6 +298,26 @@ impl<'a> StreamingMergeBuilder<'a> {
     }
 }
 
+/// Returns whether [`StreamingMergeBuilder::build`] merges on `expressions` in
+/// the row format rather than with a cursor for a single column.
+pub(crate) fn merges_in_row_format(expressions: &LexOrdering, schema: &Schema) -> bool {
+    if expressions.len() > 1 {
+        return true;
+    }
+    let Ok(data_type) = expressions[0].expr.data_type(schema) else {
+        return true;
+    };
+    !(data_type.is_primitive()
+        || matches!(
+            data_type,
+            DataType::Utf8
+                | DataType::Utf8View
+                | DataType::LargeUtf8
+                | DataType::Binary
+                | DataType::LargeBinary
+        ))
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{common::collect, stream::RecordBatchStreamAdapter};