From b02c75ebbc7575797eec69b9304fd74b8b6e1b04 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 30 Oct 2026 13:16:51 -0700 Subject: [PATCH 1/2] fix: split multi-column sort runs into smaller batches under a memory limit A multi-column merge holds the row-format keys of every sorted run's current batch. When the sort keys are most of each row, those keys are larger than the batches, and the merge needs more memory than the sort reserved while buffering them, so the sort fails with ResourcesExhausted instead of spilling. With a memory limit, sort the runs of a multi-column merge into batches of an eighth of `batch_size`, but at least 1024 rows. The merge output keeps `batch_size`, and sorts without a memory limit are unchanged. Closes #26940 --- datafusion/physical-plan/src/sorts/sort.rs | 47 +++++++++++--- .../src/sorts/sort/spill_tests.rs | 65 ++++++++++++++++++- 2 files changed, 100 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 8663515c07a28..3b76348854ed4 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -60,7 +60,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 +708,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 +726,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 +783,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 +977,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 +1005,11 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + batch_size: usize, ) -> Result { 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 +1022,7 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + batch_size: usize, ) -> Result> + Send + use<>> { let reserved_bytes = get_reserved_bytes_for_record_batch(&batch)?; @@ -1023,7 +1033,6 @@ 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(); @@ -1078,6 +1087,23 @@ impl ExternalSorter { }) } + /// Batch size of the sorted runs that feed a merge. A multi-column merge holds + /// the row-format 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`. + fn merge_run_batch_size(&self) -> usize { + let limited = !matches!( + self.runtime.memory_pool.memory_limit(), + MemoryLimit::Infinite + ); + if limited && self.expr.len() > 1 { + (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. @@ -2552,7 +2578,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 { diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 625f96be682cb..34284d0e0c109 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -9,8 +9,8 @@ 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, Int8Array, Int32Array, Int64Array, + StringArray, StringViewArray, }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::{DataType, Field, Int8Type, Schema}; @@ -679,6 +679,67 @@ async fn test_sort_preparation_keeps_output_reserved() -> Result<()> { Ok(()) } +/// Two `Int32` keys encode to 18 bytes per row for the merge, more than twice +/// the 8 bytes per row of the batches. Merging full-size runs would need more +/// memory than the sort reserves while buffering them. +async fn check_merge_keys_larger_than_batches() -> Result<()> { + let rows = 8192; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batches = (0..64) + .map(|batch| { + let a = + Int32Array::from_iter_values((0..rows).map(|i| (i * 31 + batch) % 1000)); + let b = Int32Array::from_iter_values((0..rows).rev()); + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(a), Arc::new(b)]) + }) + .collect::, _>>()?; + let ordering: LexOrdering = [ + PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))), + PhysicalSortExpr::new_default(Arc::new(Column::new("b", 1))), + ] + .into(); + // Spill every 16 batches. + 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) + .build_arc()?; + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + ordering.clone(), + rows as usize, + 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 = 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<()> { From 571f7f986c56af262c91965af2a4599d0ad80cd2 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 3 Nov 2026 08:54:02 -0800 Subject: [PATCH 2/2] fix: split runs for every row-format merge without raising their charge - Split runs whenever the merge uses the row format, including a single key without a specialized cursor, such as `FixedSizeBinary`. - Compact the view arrays of split runs, so each batch is charged for its own strings rather than for the run's whole data buffers. - Keep `batch_size` when split batches would still share buffers, such as dictionary values. --- datafusion/physical-plan/src/sorts/sort.rs | 37 ++++++-- .../src/sorts/sort/spill_tests.rs | 88 ++++++++++++++----- .../src/sorts/streaming_merge.rs | 25 +++++- 3 files changed, 116 insertions(+), 34 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 3b76348854ed4..0ebe6962a7cfb 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/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}; @@ -1035,6 +1037,8 @@ impl ExternalSorter { let expressions = self.expr.clone(); 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 @@ -1045,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::>()?; + } // Chunked output can retain shared buffers in every batch and // exceed the input estimate. Borrow only already-reserved spill @@ -1087,17 +1100,25 @@ impl ExternalSorter { }) } - /// Batch size of the sorted runs that feed a merge. A multi-column merge holds - /// the row-format 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 + /// 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`. + /// 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 ); - if limited && self.expr.len() > 1 { + 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 diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 34284d0e0c109..090d9972805f5 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/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, Int32Array, 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,29 +679,69 @@ async fn test_sort_preparation_keeps_output_reserved() -> Result<()> { Ok(()) } -/// Two `Int32` keys encode to 18 bytes per row for the merge, more than twice -/// the 8 bytes per row of the batches. Merging full-size runs would need more -/// memory than the sort reserves while buffering them. +/// 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 schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let batches = (0..64) - .map(|batch| { - let a = - Int32Array::from_iter_values((0..rows).map(|i| (i * 31 + batch) % 1000)); - let b = Int32Array::from_iter_values((0..rows).rev()); - RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(a), Arc::new(b)]) + let int_keys = |batch: i32| -> Vec { + 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()) }) - .collect::, _>>()?; - let ordering: LexOrdering = [ - PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))), - PhysicalSortExpr::new_default(Arc::new(Column::new("b", 1))), - ] - .into(); - // Spill every 16 batches. + .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>, +) -> Result<()> { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, data_type)| Field::new(*name, data_type.clone(), false)) + .collect::>(), + )); + let batches = (0..64) + .map(|batch| Ok(RecordBatch::try_new(Arc::clone(&schema), columns(batch)?)?)) + .collect::>>()?; + 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() @@ -711,7 +751,7 @@ async fn check_merge_keys_larger_than_batches() -> Result<()> { 0, Arc::clone(&schema), ordering.clone(), - rows as usize, + batches[0].num_rows(), MERGE_BYTES, 0, SpillCompression::Uncompressed, diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index da341cf0cf047..d08f387702cd0 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/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};