From a5ba530f49e7f66ef5a0cbfd2a7b1b422668268c Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 30 Oct 2026 13:16:51 -0700 Subject: [PATCH 1/6] 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 96e673e4043bc..8c20d86df3ec8 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -59,7 +59,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; @@ -707,7 +707,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)); } @@ -724,7 +725,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)); } @@ -780,13 +782,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?; @@ -968,11 +976,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() } @@ -995,11 +1004,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))) } @@ -1012,6 +1021,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)?; @@ -1024,7 +1034,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(); @@ -1068,6 +1077,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. @@ -2601,7 +2627,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 f21eacce9c4be4238e64dce340500168206e6a52 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 3 Nov 2026 08:54:02 -0800 Subject: [PATCH 2/6] 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 8c20d86df3ec8..87e8dfd9e6ef3 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::{ObservedStream, RecordBatchStreamAdapter}; use crate::topk::TopK; @@ -1036,6 +1038,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 @@ -1046,7 +1050,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 @@ -1077,17 +1090,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 2e5e0790accec..1668a55ac923d 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; @@ -273,7 +273,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())?; @@ -308,6 +309,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}; From 6eef7cb891a47ced996bfb6ae2f3ec81c0f4b4f3 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 6 Nov 2026 10:09:09 -0800 Subject: [PATCH 3/6] fix: keep shared view bytes and spilled view buffers charged once - Keep `batch_size` for a run whose views share bytes, as dictionary-encoded Parquet strings do. Compacting such views copies the shared bytes once per row, so every split batch would keep and be charged for the whole buffer. - Compact view arrays with several data buffers before spilling unless their views share bytes. Read back, each buffer is a slice of one IPC message and is charged for all of it. --- datafusion/physical-plan/src/sorts/sort.rs | 10 ++- .../src/sorts/sort/spill_tests.rs | 74 ++++++++++++++++--- datafusion/physical-plan/src/spill/mod.rs | 45 ++++++++++- 3 files changed, 111 insertions(+), 18 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 87e8dfd9e6ef3..2ffb3e889410a 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -35,7 +35,7 @@ use crate::sorts::streaming_merge::{ }; 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::spill::{gc_view_arrays, get_record_batch_memory_size, views_share_bytes}; use crate::statistics::{ChildStats, StatisticsArgs, with_per_partition_fetch}; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use crate::topk::TopK; @@ -1038,6 +1038,12 @@ impl ExternalSorter { let expressions = self.expr.clone(); let merge_pool = Arc::clone(&self.merge_pool); let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); + // Split batches of views sharing bytes would each be charged for them + let batch_size = if batch_size < self.batch_size && views_share_bytes(&batch) { + self.batch_size + } else { + batch_size + }; // Set when `merge_run_batch_size` split this run let split = batch_size < self.batch_size; @@ -1097,7 +1103,7 @@ impl ExternalSorter { /// 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. + /// them. `sort_run` also keeps it for a run whose views share bytes. fn merge_run_batch_size(&self) -> usize { let limited = !matches!( self.runtime.memory_pool.memory_limit(), diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 090d9972805f5..c4bb94fc94a71 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -215,6 +215,13 @@ fn sorter( } async fn assert_released(pool: &AdjustablePool, runtime: &RuntimeEnv) { + wait_until_released(pool, runtime).await; + let state = pool.state.lock().unwrap(); + assert!(state.peak <= pool.capacity); + assert_eq!(state.unchecked_over_limit, 0); +} + +async fn wait_until_released(pool: &AdjustablePool, runtime: &RuntimeEnv) { // Buffered readers may finish dropping their state after the owning stream. tokio::time::timeout(Duration::from_secs(5), async { while pool.reserved() != 0 @@ -227,9 +234,6 @@ async fn assert_released(pool: &AdjustablePool, runtime: &RuntimeEnv) { .await .expect("sort reservations and spill files must be released"); assert_eq!(runtime.disk_manager.used_disk_space(), 0); - let state = pool.state.lock().unwrap(); - assert!(state.peak <= pool.capacity); - assert_eq!(state.unchecked_over_limit, 0); } async fn check_spill_preserves_merge_workspace_after_limit_decreases() -> Result<()> { @@ -696,7 +700,7 @@ async fn check_merge_keys_larger_than_batches() -> Result<()> { ] }; let int_fields = [("a", DataType::Int32), ("b", DataType::Int32)]; - check_merge_keys(&int_fields, 2, |batch| Ok(int_keys(batch))).await?; + check_merge_keys(&int_fields, 2, 16, true, |batch| Ok(int_keys(batch))).await?; // Longer than 12 bytes, so views keep them in data buffers. let strings = |batch: i32| -> StringArray { (0..rows) @@ -707,25 +711,66 @@ async fn check_merge_keys_larger_than_batches() -> Result<()> { 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| { + check_merge_keys(&fields, 2, 16, true, |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)]) + // Views into 1024 dictionary values share their bytes, as a + // dictionary-encoded Parquet column produces. Sorting the buffered batches + // before a spill grows past the pool, on `main` too. + let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat(); + check_merge_keys(&fields, 2, 16, false, |batch| { + let values: StringArray = (0..1024) + .map(|v| Some(format!("{batch:>4}-{v:0>59}"))) + .collect(); + let keys = + Int32Array::from_iter_values((0..rows).map(|i| (i * 7 + batch) % 1024)); + let views = cast( + &DictionaryArray::try_new(keys, Arc::new(values))?, + &DataType::Utf8View, + )?; + Ok([int_keys(batch), vec![views]].concat()) + }) + .await?; + // Batches that do not overlap, so a merged batch takes the strings of + // several split batches of one run, in as many buffers. + let view_keys = [("a", DataType::Utf8View), ("b", DataType::Utf8View)]; + check_merge_keys(&view_keys, 2, 8, true, |batch| { + let keys = |key: &str| -> Result { + let strings: StringArray = (0..rows) + .rev() + .map(|i| Some(format!("{batch:0>4}{key}{i:0>61}"))) + .collect(); + Ok(cast(&strings, &DataType::Utf8View)?) + }; + Ok(vec![keys("a")?, keys("b")?]) }) + .await?; + check_merge_keys( + &[("k", DataType::FixedSizeBinary(16))], + 1, + 16, + true, + |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. +/// `buffered` of them, then checks the spilled output and that all memory is +/// released. Unless `within_pool`, the pool may have been exceeded by an +/// unchecked `grow`. async fn check_merge_keys( fields: &[(&str, DataType)], keys: usize, + buffered: usize, + within_pool: bool, columns: impl Fn(i32) -> Result>, ) -> Result<()> { let schema = Arc::new(Schema::new( @@ -742,7 +787,8 @@ async fn check_merge_keys( 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 capacity = + MERGE_BYTES + buffered * 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) @@ -766,7 +812,11 @@ async fn check_merge_keys( let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; assert_eq!(concat_batches(&schema, &output)?, expected); drop(sorter); - assert_released(&pool, &runtime).await; + if within_pool { + assert_released(&pool, &runtime).await; + } else { + wait_until_released(&pool, &runtime).await; + } Ok(()) } diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index f2ec70c6f3544..07dcda1f599dc 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -962,9 +962,11 @@ fn gc_array(array: &ArrayRef, dictionaries: bool) -> Result<(ArrayRef, bool)> { /// batch, and `gc()` copies it once per row /// (). /// -/// Repeated buffer entries are merged first. `gc()` then runs only if it -/// writes less than the merged buffers hold, and runs on the original array -/// so that the views are rewritten at most once. +/// Repeated buffer entries are merged first. `gc()` then runs if it writes +/// less than the merged buffers hold, or if several buffers remain and the +/// views share no bytes, since each buffer read back from the spill is a slice +/// of one IPC message and is charged for all of it. It runs on the original +/// array so that the views are rewritten at most once. fn compact_view_array( array: &GenericByteViewArray, ) -> Option> { @@ -975,7 +977,9 @@ fn compact_view_array( let (remap, unique) = dedup_buffer_entries(buffers); let unique_size: usize = unique.iter().map(|b| b.len()).sum(); - if should_gc_view_array(array, unique_size) { + if should_gc_view_array(array, unique_size) + || (unique.len() > 1 && array.total_buffer_bytes_used() <= unique_size) + { return Some(array.gc()); } if unique.len() == buffers.len() { @@ -1034,6 +1038,21 @@ fn dedup_buffer_entries(buffers: &[Buffer]) -> (Vec, Vec) { (remap, unique) } +/// Returns whether a view column of `batch` references more bytes than its +/// data buffers hold, as the views of a dictionary-encoded Parquet column do. +/// Compacting such views would copy the shared bytes once per row. +pub(crate) fn views_share_bytes(batch: &RecordBatch) -> bool { + fn shares_bytes(array: &GenericByteViewArray) -> bool { + let (_, unique) = dedup_buffer_entries(array.data_buffers()); + array.total_buffer_bytes_used() > unique.iter().map(|b| b.len()).sum::() + } + batch.columns().iter().any(|array| match array.data_type() { + DataType::Utf8View => shares_bytes(array.as_string_view()), + DataType::BinaryView => shares_bytes(array.as_binary_view()), + _ => false, + }) +} + fn gc_array_children(array: &ArrayRef, dictionaries: bool) -> Result<(ArrayRef, bool)> { let data = array.to_data(); if data.child_data().is_empty() { @@ -2669,6 +2688,24 @@ mod tests { Ok((dictionary, merged.as_string_view().clone())) } + #[test] + fn test_views_share_bytes() -> Result<()> { + use arrow::array::UInt32Array; + use arrow::compute::take; + + // 8192 views over 1024 values reference 8x the dictionary's bytes + let dictionary = dictionary_array(1024)?; + let indices = UInt32Array::from_iter_values((0..8192).map(|i| i % 1024)); + let batch = + RecordBatch::try_from_iter([("c", take(&dictionary, &indices, None)?)])?; + assert!(views_share_bytes(&batch)); + // 8192 values used once each reference exactly the buffer's bytes + let views: ArrayRef = Arc::new(dictionary_array(8192)?); + let batch = RecordBatch::try_from_iter([("c", views)])?; + assert!(!views_share_bytes(&batch)); + Ok(()) + } + #[test] fn test_gc_keeps_shared_dictionary_buffer() -> Result<()> { let (dictionary, merged) = shared_dictionary_views(1000, 8, 1024)?; From 8eb93eb643e370e4c448cb897efbe36d14708316 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 9 Nov 2026 09:13:38 -0800 Subject: [PATCH 4/6] perf: split only runs whose encoded keys outgrow them The sort reserves about twice each run, so a merge in the row format fits when the keys are at most as large as the run. Splitting smaller keys only added work: with integer keys and a `Utf8View` payload, the compaction of split batches copied every string. `merge_run_batch_size` now checks each run, from the bytes of its key columns, and also covers views that share bytes. Also test that spill compaction leaves views sharing bytes across several buffers alone, and update the `gc_view_arrays` docs. --- datafusion/physical-plan/src/sorts/sort.rs | 61 ++++++++++++------- .../src/sorts/sort/spill_tests.rs | 59 ++++++++++++++++-- datafusion/physical-plan/src/spill/mod.rs | 30 +++++++-- 3 files changed, 119 insertions(+), 31 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 2ffb3e889410a..b39f1207a6611 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -66,7 +66,7 @@ use datafusion_execution::memory_pool::{ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; +use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -784,13 +784,13 @@ 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 batch_size = self.merge_run_batch_size(&batch)?; let sort = self.sort_run( batch, reservation, @@ -978,11 +978,11 @@ 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); + let batch_size = self.merge_run_batch_size(&run)?; self.sort_batch_stream(run, reservation, shared_input.clone(), batch_size) }) .collect() @@ -1038,12 +1038,6 @@ impl ExternalSorter { let expressions = self.expr.clone(); let merge_pool = Arc::clone(&self.merge_pool); let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); - // Split batches of views sharing bytes would each be charged for them - let batch_size = if batch_size < self.batch_size && views_share_bytes(&batch) { - self.batch_size - } else { - batch_size - }; // Set when `merge_run_batch_size` split this run let split = batch_size < self.batch_size; @@ -1096,15 +1090,17 @@ 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. `sort_run` also keeps it for a run whose views share bytes. - fn merge_run_batch_size(&self) -> usize { + /// Batch size of the sorted run `run` when it feeds a merge. A merge in the + /// row format holds the keys of every run's current batch at once. The sort + /// reserves about twice the run, so the keys fit next to the sorted copy if + /// they are at most as large as the run. Otherwise, with a memory limit, the + /// run uses 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, or when the + /// views of `run` share bytes, since every batch is charged for them. + fn merge_run_batch_size(&self, run: &RecordBatch) -> Result { let limited = !matches!( self.runtime.memory_pool.memory_limit(), MemoryLimit::Infinite @@ -1114,11 +1110,34 @@ impl ExternalSorter { 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) + if limited + && !shares_buffers + && merges_in_row_format(&self.expr, &self.schema) + && self.encoded_keys_may_exceed(run)? + && !views_share_bytes(run) + { + Ok((self.batch_size / 8).max(1024).min(self.batch_size)) } else { - self.batch_size + Ok(self.batch_size) + } + } + + /// Returns whether the row-format keys of `run` may be larger than `run`. A + /// computed key counts as larger, since sizing it would evaluate it again. + fn encoded_keys_may_exceed(&self, run: &RecordBatch) -> Result { + // An 8-byte offset per row, then a null byte and the bytes of each key + let mut encoded = run.num_rows() * size_of::(); + for sort_expr in self.expr.iter() { + let Some(column) = sort_expr.expr.downcast_ref::() else { + return Ok(true); + }; + encoded += run.num_rows() + + run + .column(column.index()) + .to_data() + .get_slice_memory_size()?; } + Ok(encoded > run.get_sliced_size()?) } /// If this sort may spill, pre-allocates diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index c4bb94fc94a71..2c4b5adef4b5b 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -685,10 +685,11 @@ async fn test_sort_preparation_keeps_output_reserved() -> Result<()> { /// 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. +/// batches, and a single `FixedSizeBinary(16)` key and two long `Utf8View` keys +/// also encode to more than their runs, so merging full-size runs needs more +/// memory than the sort reserves while buffering them. The payloads make runs +/// larger than their keys or keep buffers shared after sorting, so these runs +/// keep `batch_size` and must not be charged once per batch either. async fn check_merge_keys_larger_than_batches() -> Result<()> { let rows = 8192; let int_keys = |batch: i32| -> Vec { @@ -711,7 +712,9 @@ async fn check_merge_keys_larger_than_batches() -> Result<()> { 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, 16, true, |batch| { + // Runs of this `Utf8View` payload are not split, as on `main` + let within_pool = payload != DataType::Utf8View; + check_merge_keys(&fields, 2, 16, within_pool, |batch| { let payload = cast(&strings(batch), &payload)?; Ok([int_keys(batch), vec![payload]].concat()) }) @@ -830,6 +833,52 @@ async fn test_merge_keys_larger_than_batches_multithreaded() -> Result<()> { check_merge_keys_larger_than_batches().await } +/// Only runs whose encoded keys are larger than the run are split. +#[test] +fn test_merge_run_batch_size_follows_key_size() -> Result<()> { + let batch_size = 8192; + let rows = batch_size as i32; + let check = |payload: Option, expected: usize| -> Result<()> { + let mut columns: Vec<(&str, ArrayRef)> = vec![ + ("a", Arc::new(Int64Array::from_iter_values(0..rows as i64))), + ("b", Arc::new(Int64Array::from_iter_values(0..rows as i64))), + ]; + columns.extend(payload.map(|payload| ("c", payload))); + let batch = RecordBatch::try_from_iter(columns)?; + let ordering = LexOrdering::new((0..2).map(|i| { + PhysicalSortExpr::new_default(Arc::new(Column::new( + batch.schema().field(i).name(), + i, + ))) + })) + .expect("two sort keys"); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(1024 * 1024 * 1024, 1.0) + .build_arc()?; + let sorter = ExternalSorter::new( + 0, + batch.schema(), + ordering, + batch_size, + 0, + 0, + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + runtime, + )?; + assert_eq!(sorter.merge_run_batch_size(&batch)?, expected); + Ok(()) + }; + // Encoded keys: 26 bytes a row, against 16 bytes a row + check(None, batch_size / 8)?; + // Against 16 + 16 bytes a row: inlined views add more than the keys need + let short: StringViewArray = (0..rows).map(|i| Some(format!("{i}"))).collect(); + check(Some(Arc::new(short)), batch_size)?; + // Against 16 + 16 + 64 bytes a row + let long: StringViewArray = (0..rows).map(|i| Some(format!("{i:0>64}"))).collect(); + check(Some(Arc::new(long)), batch_size) +} + #[tokio::test] async fn test_single_batch_spill_preserves_workspace_after_limit_decreases() -> Result<()> { diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 07dcda1f599dc..3bfd0c37eeebb 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -833,15 +833,17 @@ const VIEW_SIZE_BYTES: usize = 16; /// /// # When GC is triggered /// -/// View GC is only performed when data buffers exceed a threshold (currently 10KB). -/// This balances memory savings against the CPU overhead of garbage collection. -/// Small view buffers are passed through unchanged since the GC overhead would exceed -/// any memory savings. +/// View GC runs when data buffers above a threshold (currently 10KB) hold more +/// bytes than the views use. This balances memory savings against the CPU +/// overhead of garbage collection. It also runs at any size when several data +/// buffers remain and the views share no bytes, since each buffer read back +/// from the spill is charged for the whole IPC message. /// /// # Performance considerations /// /// - If no arrays need compaction, the original batch is cloned cheaply -/// - View GC is skipped for small buffers to avoid unnecessary CPU overhead +/// - View GC is skipped for a single small data buffer to avoid unnecessary CPU +/// overhead /// - Nested container types are traversed recursively so view arrays inside /// `List`, `Map`, `Union`, `Dictionary`, and other child-bearing arrays are compacted too /// - Traversal uses physical children, not logical reachability. For example, @@ -2706,6 +2708,24 @@ mod tests { Ok(()) } + #[test] + fn test_gc_keeps_shared_views_in_two_buffers() -> Result<()> { + use arrow::array::AsArray; + use arrow::compute::concat; + + let (_, first) = shared_dictionary_views(1000, 8, 1024)?; + let (_, second) = shared_dictionary_views(1000, 8, 1024)?; + let merged = concat(&[&first as &dyn Array, &second])?; + let merged = merged.as_string_view(); + + // `gc()` would copy 16384 values, 8x what the two dictionaries hold + let compacted = compact_view_array(merged).expect("buffer entries repeat"); + assert_eq!(compacted.data_buffers().len(), 2); + assert_eq!(&compacted, merged); + + Ok(()) + } + #[test] fn test_gc_keeps_shared_dictionary_buffer() -> Result<()> { let (dictionary, merged) = shared_dictionary_views(1000, 8, 1024)?; From b29d5b02a12a52f00e9085fcd803c4281b9c80e1 Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 11 Nov 2026 08:24:56 -0800 Subject: [PATCH 5/6] test: reach the shared-views gate with a computed sort key Since runs whose keys fit keep `batch_size`, the shared-views case sorted by column keys never reached `views_share_bytes`. Sort it by a computed key, which counts as larger than the run, so the case fails without the gate. --- .../src/sorts/sort/spill_tests.rs | 36 ++++++++++++++----- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 2c4b5adef4b5b..99f97f1562c26 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -22,8 +22,8 @@ use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, }; use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; -use datafusion_physical_expr::expressions::Column; -use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_expr::expressions::{CastExpr, Column}; +use datafusion_physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr}; use futures::TryStreamExt; use std::fmt::{Display, Formatter}; use std::sync::{Arc, Mutex}; @@ -721,10 +721,17 @@ async fn check_merge_keys_larger_than_batches() -> Result<()> { .await?; } // Views into 1024 dictionary values share their bytes, as a - // dictionary-encoded Parquet column produces. Sorting the buffered batches - // before a spill grows past the pool, on `main` too. + // dictionary-encoded Parquet column produces. A computed key counts as + // larger than the run, so only the shared views keep these runs at + // `batch_size`. Sorting the buffered batches before a spill grows past the + // pool, on `main` too. let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat(); - check_merge_keys(&fields, 2, 16, false, |batch| { + let a: Arc = Arc::new(Column::new("a", 0)); + let keys = vec![ + Arc::new(CastExpr::new(a, DataType::Int64, None)) as Arc, + Arc::new(Column::new("b", 1)), + ]; + check_merge_exprs(&fields, keys, 16, false, |batch| { let values: StringArray = (0..1024) .map(|v| Some(format!("{batch:>4}-{v:0>59}"))) .collect(); @@ -775,6 +782,20 @@ async fn check_merge_keys( buffered: usize, within_pool: bool, columns: impl Fn(i32) -> Result>, +) -> Result<()> { + let keys = (0..keys) + .map(|i| Arc::new(Column::new(fields[i].0, i)) as Arc) + .collect(); + check_merge_exprs(fields, keys, buffered, within_pool, columns).await +} + +/// Like [`check_merge_keys`], sorting by `keys`. +async fn check_merge_exprs( + fields: &[(&str, DataType)], + keys: Vec>, + buffered: usize, + within_pool: bool, + columns: impl Fn(i32) -> Result>, ) -> Result<()> { let schema = Arc::new(Schema::new( fields @@ -785,10 +806,7 @@ async fn check_merge_keys( 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))) - })) + let ordering = LexOrdering::new(keys.into_iter().map(PhysicalSortExpr::new_default)) .expect("at least one sort key"); let capacity = MERGE_BYTES + buffered * get_reserved_bytes_for_record_batch(&batches[0])?; From 0e71291f0e46643b9902007951e9c316bff41cdb Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 11 Nov 2026 08:59:43 -0800 Subject: [PATCH 6/6] perf: size fixed-width computed sort keys from their type A computed key counted as larger than the run, so its runs were always split under a memory limit. With a computed integer key and a `Utf8View` payload, the compaction of split batches copied every string. Size such keys from their type when it has a fixed width. The shared-views case now sorts by a string cast, which still counts as larger. --- datafusion/physical-plan/src/sorts/sort.rs | 25 ++++++++---- .../src/sorts/sort/spill_tests.rs | 40 +++++++++++-------- 2 files changed, 42 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index b39f1207a6611..c045ccfe3833c 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1123,19 +1123,30 @@ impl ExternalSorter { } /// Returns whether the row-format keys of `run` may be larger than `run`. A - /// computed key counts as larger, since sizing it would evaluate it again. + /// computed key is sized from its type if that has a fixed width, and + /// otherwise counts as larger, since sizing it would evaluate it again. fn encoded_keys_may_exceed(&self, run: &RecordBatch) -> Result { // An 8-byte offset per row, then a null byte and the bytes of each key let mut encoded = run.num_rows() * size_of::(); for sort_expr in self.expr.iter() { - let Some(column) = sort_expr.expr.downcast_ref::() else { - return Ok(true); - }; - encoded += run.num_rows() - + run + let key_bytes = match sort_expr.expr.downcast_ref::() { + Some(column) => run .column(column.index()) .to_data() - .get_slice_memory_size()?; + .get_slice_memory_size()?, + None => { + let width = sort_expr + .expr + .data_type(self.schema.as_ref()) + .ok() + .and_then(|data_type| data_type.primitive_width()); + let Some(width) = width else { + return Ok(true); + }; + run.num_rows() * width + } + }; + encoded += run.num_rows() + key_bytes; } Ok(encoded > run.get_sliced_size()?) } diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 99f97f1562c26..fa132fb91b410 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -721,14 +721,14 @@ async fn check_merge_keys_larger_than_batches() -> Result<()> { .await?; } // Views into 1024 dictionary values share their bytes, as a - // dictionary-encoded Parquet column produces. A computed key counts as - // larger than the run, so only the shared views keep these runs at - // `batch_size`. Sorting the buffered batches before a spill grows past the - // pool, on `main` too. + // dictionary-encoded Parquet column produces. A computed key without a + // fixed width counts as larger than the run, so only the shared views keep + // these runs at `batch_size`. Sorting the buffered batches before a spill + // grows past the pool, on `main` too. let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat(); let a: Arc = Arc::new(Column::new("a", 0)); let keys = vec![ - Arc::new(CastExpr::new(a, DataType::Int64, None)) as Arc, + Arc::new(CastExpr::new(a, DataType::Utf8, None)) as Arc, Arc::new(Column::new("b", 1)), ]; check_merge_exprs(&fields, keys, 16, false, |batch| { @@ -856,20 +856,21 @@ async fn test_merge_keys_larger_than_batches_multithreaded() -> Result<()> { fn test_merge_run_batch_size_follows_key_size() -> Result<()> { let batch_size = 8192; let rows = batch_size as i32; - let check = |payload: Option, expected: usize| -> Result<()> { + // Sorts by `first`, then `b` + let check = |first: Arc, + payload: Option, + expected: usize| + -> Result<()> { let mut columns: Vec<(&str, ArrayRef)> = vec![ ("a", Arc::new(Int64Array::from_iter_values(0..rows as i64))), ("b", Arc::new(Int64Array::from_iter_values(0..rows as i64))), ]; columns.extend(payload.map(|payload| ("c", payload))); let batch = RecordBatch::try_from_iter(columns)?; - let ordering = LexOrdering::new((0..2).map(|i| { - PhysicalSortExpr::new_default(Arc::new(Column::new( - batch.schema().field(i).name(), - i, - ))) - })) - .expect("two sort keys"); + let b: Arc = Arc::new(Column::new("b", 1)); + let ordering = + LexOrdering::new([first, b].into_iter().map(PhysicalSortExpr::new_default)) + .expect("two sort keys"); let runtime = RuntimeEnvBuilder::new() .with_memory_limit(1024 * 1024 * 1024, 1.0) .build_arc()?; @@ -887,14 +888,21 @@ fn test_merge_run_batch_size_follows_key_size() -> Result<()> { assert_eq!(sorter.merge_run_batch_size(&batch)?, expected); Ok(()) }; + let a = || Arc::new(Column::new("a", 0)) as Arc; // Encoded keys: 26 bytes a row, against 16 bytes a row - check(None, batch_size / 8)?; + check(a(), None, batch_size / 8)?; // Against 16 + 16 bytes a row: inlined views add more than the keys need let short: StringViewArray = (0..rows).map(|i| Some(format!("{i}"))).collect(); - check(Some(Arc::new(short)), batch_size)?; + check(a(), Some(Arc::new(short)), batch_size)?; // Against 16 + 16 + 64 bytes a row let long: StringViewArray = (0..rows).map(|i| Some(format!("{i:0>64}"))).collect(); - check(Some(Arc::new(long)), batch_size) + let long: ArrayRef = Arc::new(long); + check(a(), Some(Arc::clone(&long)), batch_size)?; + // A computed key is sized from its type if that has a fixed width, and + // otherwise counts as larger than the run + let cast = |to| Arc::new(CastExpr::new(a(), to, None)) as Arc; + check(cast(DataType::Int64), Some(Arc::clone(&long)), batch_size)?; + check(cast(DataType::Utf8), Some(long), batch_size / 8) } #[tokio::test]