From 017ad87a820170aa0a5a2674dae793a773b665fd Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 1 Mar 2026 21:13:01 +0100 Subject: [PATCH 1/5] try to fix the sort merge oom --- datafusion/physical-plan/src/sorts/builder.rs | 41 ++++++- .../src/sorts/multi_level_merge.rs | 91 ++++++++------ datafusion/physical-plan/src/sorts/sort.rs | 114 +++++++++++++++++- 3 files changed, 205 insertions(+), 41 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 9b2fa968222c4..881b995cf9d51 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -40,9 +40,24 @@ pub struct BatchBuilder { /// Maintain a list of [`RecordBatch`] and their corresponding stream batches: Vec<(usize, RecordBatch)>, - /// Accounts for memory used by buffered batches + /// Accounts for memory used by buffered batches. + /// + /// May include pre-reserved bytes (from `sort_spill_reservation_bytes`) + /// that were transferred via [`MemoryReservation::take()`] to prevent + /// starvation when concurrent sort partitions compete for pool memory. reservation: MemoryReservation, + /// Tracks the actual memory used by buffered batches (not including + /// pre-reserved bytes). This allows [`Self::push_batch`] to skip pool + /// allocation requests when the pre-reserved bytes cover the batch. + batches_mem_used: usize, + + /// The initial reservation size at construction time. When the reservation + /// is pre-loaded with `sort_spill_reservation_bytes` (via `take()`), this + /// records that amount so we never shrink below it, maintaining the + /// anti-starvation guarantee throughout the merge. + initial_reservation: usize, + /// The current [`BatchCursor`] for each stream cursors: Vec, @@ -59,19 +74,29 @@ impl BatchBuilder { batch_size: usize, reservation: MemoryReservation, ) -> Self { + let initial_reservation = reservation.size(); Self { schema, batches: Vec::with_capacity(stream_count * 2), cursors: vec![BatchCursor::default(); stream_count], indices: Vec::with_capacity(batch_size), reservation, + batches_mem_used: 0, + initial_reservation, } } /// Append a new batch in `stream_idx` pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<()> { - self.reservation - .try_grow(get_record_batch_memory_size(&batch))?; + let size = get_record_batch_memory_size(&batch); + self.batches_mem_used += size; + // Only request additional memory from the pool when actual batch + // usage exceeds the current reservation (which may include + // pre-reserved bytes from sort_spill_reservation_bytes). + if self.batches_mem_used > self.reservation.size() { + self.reservation + .try_grow(self.batches_mem_used - self.reservation.size())?; + } let batch_idx = self.batches.len(); self.batches.push((stream_idx, batch)); self.cursors[stream_idx] = BatchCursor { @@ -143,11 +168,19 @@ impl BatchBuilder { stream_cursor.batch_idx = retained; retained += 1; } else { - self.reservation.shrink(get_record_batch_memory_size(batch)); + 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); + } + Ok(Some(RecordBatch::try_new( Arc::clone(&self.schema), columns, diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 2e0d668a29559..d5b0c2ad8e5c1 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -253,7 +253,12 @@ impl MultiLevelMergeBuilder { // Need to merge multiple streams (_, _) => { - let mut memory_reservation = self.reservation.new_empty(); + // Transfer any pre-reserved bytes (from sort_spill_reservation_bytes) + // to the merge memory reservation. This prevents starvation when + // concurrent sort partitions compete for pool memory: the pre-reserved + // bytes cover spill file buffer reservations without additional pool + // allocation. + let mut memory_reservation = self.reservation.take(); // Don't account for existing streams memory // as we are not holding the memory for them @@ -337,8 +342,10 @@ impl MultiLevelMergeBuilder { builder = builder.with_bypass_mempool(); } else { // If we are only merging in-memory streams, we need to use the memory reservation - // because we don't know the maximum size of the batches in the streams - builder = builder.with_reservation(self.reservation.new_empty()); + // because we don't know the maximum size of the batches in the streams. + // Use take() to transfer any pre-reserved bytes so the merge can use them + // as its initial budget without additional pool allocation. + builder = builder.with_reservation(self.reservation.take()); } builder.build() @@ -356,45 +363,57 @@ impl MultiLevelMergeBuilder { ) -> Result<(Vec, usize)> { assert_ne!(buffer_len, 0, "Buffer length must be greater than 0"); let mut number_of_spills_to_read_for_current_phase = 0; + // Track total memory needed for spill file buffers. When the + // reservation has pre-reserved bytes (from sort_spill_reservation_bytes), + // those bytes cover the first N spill files without additional pool + // allocation, preventing starvation under memory pressure. + let mut total_needed: usize = 0; for spill in &self.sorted_spill_files { - // For memory pools that are not shared this is good, for other this is not - // and there should be some upper limit to memory reservation so we won't starve the system - match reservation.try_grow( - get_reserved_bytes_for_record_batch_size( - spill.max_record_batch_memory, - // Size will be the same as the sliced size, bc it is a spilled batch. - spill.max_record_batch_memory, - ) * buffer_len, - ) { - Ok(_) => { - number_of_spills_to_read_for_current_phase += 1; - } - // If we can't grow the reservation, we need to stop - Err(err) => { - // We must have at least 2 streams to merge, so if we don't have enough memory - // fail - if minimum_number_of_required_streams - > number_of_spills_to_read_for_current_phase - { - // Free the memory we reserved for this merge as we either try again or fail - reservation.free(); - if buffer_len > 1 { - // Try again with smaller buffer size, it will be slower but at least we can merge - return self.get_sorted_spill_files_to_merge( - buffer_len - 1, - minimum_number_of_required_streams, - reservation, - ); + let per_spill = get_reserved_bytes_for_record_batch_size( + spill.max_record_batch_memory, + // Size will be the same as the sliced size, bc it is a spilled batch. + spill.max_record_batch_memory, + ) * buffer_len; + total_needed += per_spill; + + // Only request additional memory from the pool when total needed + // exceeds what's already reserved (which may include pre-reserved + // bytes from sort_spill_reservation_bytes). + if total_needed > reservation.size() { + match reservation.try_grow(total_needed - reservation.size()) { + Ok(_) => { + number_of_spills_to_read_for_current_phase += 1; + } + // If we can't grow the reservation, we need to stop + Err(err) => { + // We must have at least 2 streams to merge, so if we don't have enough memory + // fail + if minimum_number_of_required_streams + > number_of_spills_to_read_for_current_phase + { + // Free the memory we reserved for this merge as we either try again or fail + reservation.free(); + if buffer_len > 1 { + // Try again with smaller buffer size, it will be slower but at least we can merge + return self.get_sorted_spill_files_to_merge( + buffer_len - 1, + minimum_number_of_required_streams, + reservation, + ); + } + + return Err(err); } - return Err(err); + // We reached the maximum amount of memory we can use + // for this merge + break; } - - // We reached the maximum amount of memory we can use - // for this merge - break; } + } else { + // Pre-reserved bytes cover this spill file's buffer + number_of_spills_to_read_for_current_phase += 1; } } diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index d02ef48e761bd..57ddc69c602e6 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -355,6 +355,13 @@ impl ExternalSorter { self.sort_and_spill_in_mem_batches().await?; } + // Transfer the pre-reserved merge memory to the streaming merge + // using `take()` instead of `new_empty()`. This ensures the merge + // stream starts with `sort_spill_reservation_bytes` already + // allocated, preventing starvation when concurrent sort partitions + // compete for pool memory. `take()` moves the bytes atomically + // without releasing them back to the pool, so other partitions + // cannot race to consume the freed memory. StreamingMergeBuilder::new() .with_sorted_spill_files(std::mem::take(&mut self.finished_spill_files)) .with_spill_manager(self.spill_manager.clone()) @@ -363,7 +370,7 @@ impl ExternalSorter { .with_metrics(self.metrics.baseline.clone()) .with_batch_size(self.batch_size) .with_fetch(None) - .with_reservation(self.merge_reservation.new_empty()) + .with_reservation(self.merge_reservation.take()) .build() } else { self.in_mem_sort_stream(self.metrics.baseline.clone()) @@ -2716,4 +2723,109 @@ mod tests { Ok(()) } + + /// Test that concurrent sort partitions sharing a tight memory pool + /// don't starve during the merge phase. + /// + /// This reproduces the starvation scenario where: + /// 1. Multiple ExternalSorter instances share a single GreedyMemoryPool + /// 2. Each reserves `sort_spill_reservation_bytes` for its merge phase + /// 3. After spilling, the merge must proceed using the pre-reserved bytes + /// without additional pool allocation + /// + /// Without the fix (using `take()` + smart tracking), the merge's + /// `new_empty()` reservation starts at 0 bytes and the pre-reserved bytes + /// sit unused in ExternalSorter's merge_reservation. When other partitions + /// consume the freed memory, the merge starves. + /// + /// With the fix, the pre-reserved bytes are atomically transferred to the + /// merge stream and used for spill file buffer reservations, preventing + /// starvation. + #[tokio::test] + async fn test_sort_merge_no_starvation_with_concurrent_partitions() -> Result<()> { + use futures::TryStreamExt; + + let sort_spill_reservation_bytes: usize = 10 * 1024; // 10 KB per partition + let num_partitions: usize = 4; + + // Pool: each partition needs sort_spill_reservation_bytes for its merge, + // plus a small amount for data accumulation before spilling. + // Total: 4 * 10KB + 8KB = 48KB -- very tight. + let memory_limit = + sort_spill_reservation_bytes * num_partitions + 8 * 1024; + + let session_config = SessionConfig::new() + .with_batch_size(128) + .with_sort_spill_reservation_bytes(sort_spill_reservation_bytes); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .build_arc()?; + + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + + // Create multiple batches per partition to force spilling. + // Each batch: 100 rows of Int32 ≈ 400 bytes. + // 20 batches per partition ≈ 8KB per partition. + // With only ~2KB of pool headroom per partition, this forces spilling. + let batches_per_partition = 20; + let rows_per_batch: i32 = 100; + + let all_partitions: Vec> = (0..num_partitions) + .map(|_| { + (0..batches_per_partition) + .map(|_| make_partition(rows_per_batch)) + .collect() + }) + .collect(); + + let schema = all_partitions[0][0].schema(); + let input = TestMemoryExec::try_new_exec(&all_partitions, schema.clone(), None)?; + + let sort_exec = Arc::new( + SortExec::new( + [PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions::default(), + }] + .into(), + input, + ) + .with_preserve_partitioning(true), + ); + + // Execute all partitions concurrently -- they share the same pool. + let mut tasks = Vec::new(); + for partition in 0..num_partitions { + let sort = Arc::clone(&sort_exec); + let ctx = Arc::clone(&task_ctx); + tasks.push(tokio::spawn(async move { + let stream = sort.execute(partition, ctx)?; + let batches: Vec = stream.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + Ok::(total_rows) + })); + } + + let mut total_rows = 0; + for task in tasks { + total_rows += task.await.unwrap()?; + } + + let expected_rows = + num_partitions * batches_per_partition * (rows_per_batch as usize); + assert_eq!(total_rows, expected_rows); + + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "All memory should be returned to the pool after sort completes" + ); + + Ok(()) + } } From 76f7d792b9c8a679f5d2e94b25df5b30e7dd17cf Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 2 Mar 2026 10:49:11 +0100 Subject: [PATCH 2/5] remove test --- datafusion/physical-plan/src/sorts/sort.rs | 105 --------------------- 1 file changed, 105 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 57ddc69c602e6..8c3c88a80cbf0 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -2723,109 +2723,4 @@ mod tests { Ok(()) } - - /// Test that concurrent sort partitions sharing a tight memory pool - /// don't starve during the merge phase. - /// - /// This reproduces the starvation scenario where: - /// 1. Multiple ExternalSorter instances share a single GreedyMemoryPool - /// 2. Each reserves `sort_spill_reservation_bytes` for its merge phase - /// 3. After spilling, the merge must proceed using the pre-reserved bytes - /// without additional pool allocation - /// - /// Without the fix (using `take()` + smart tracking), the merge's - /// `new_empty()` reservation starts at 0 bytes and the pre-reserved bytes - /// sit unused in ExternalSorter's merge_reservation. When other partitions - /// consume the freed memory, the merge starves. - /// - /// With the fix, the pre-reserved bytes are atomically transferred to the - /// merge stream and used for spill file buffer reservations, preventing - /// starvation. - #[tokio::test] - async fn test_sort_merge_no_starvation_with_concurrent_partitions() -> Result<()> { - use futures::TryStreamExt; - - let sort_spill_reservation_bytes: usize = 10 * 1024; // 10 KB per partition - let num_partitions: usize = 4; - - // Pool: each partition needs sort_spill_reservation_bytes for its merge, - // plus a small amount for data accumulation before spilling. - // Total: 4 * 10KB + 8KB = 48KB -- very tight. - let memory_limit = - sort_spill_reservation_bytes * num_partitions + 8 * 1024; - - let session_config = SessionConfig::new() - .with_batch_size(128) - .with_sort_spill_reservation_bytes(sort_spill_reservation_bytes); - - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(memory_limit, 1.0) - .build_arc()?; - - let task_ctx = Arc::new( - TaskContext::default() - .with_session_config(session_config) - .with_runtime(runtime), - ); - - // Create multiple batches per partition to force spilling. - // Each batch: 100 rows of Int32 ≈ 400 bytes. - // 20 batches per partition ≈ 8KB per partition. - // With only ~2KB of pool headroom per partition, this forces spilling. - let batches_per_partition = 20; - let rows_per_batch: i32 = 100; - - let all_partitions: Vec> = (0..num_partitions) - .map(|_| { - (0..batches_per_partition) - .map(|_| make_partition(rows_per_batch)) - .collect() - }) - .collect(); - - let schema = all_partitions[0][0].schema(); - let input = TestMemoryExec::try_new_exec(&all_partitions, schema.clone(), None)?; - - let sort_exec = Arc::new( - SortExec::new( - [PhysicalSortExpr { - expr: col("i", &schema)?, - options: SortOptions::default(), - }] - .into(), - input, - ) - .with_preserve_partitioning(true), - ); - - // Execute all partitions concurrently -- they share the same pool. - let mut tasks = Vec::new(); - for partition in 0..num_partitions { - let sort = Arc::clone(&sort_exec); - let ctx = Arc::clone(&task_ctx); - tasks.push(tokio::spawn(async move { - let stream = sort.execute(partition, ctx)?; - let batches: Vec = stream.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - Ok::(total_rows) - })); - } - - let mut total_rows = 0; - for task in tasks { - total_rows += task.await.unwrap()?; - } - - let expected_rows = - num_partitions * batches_per_partition * (rows_per_batch as usize); - assert_eq!(total_rows, expected_rows); - - assert_eq!( - task_ctx.runtime_env().memory_pool.reserved(), - 0, - "All memory should be returned to the pool after sort completes" - ); - - Ok(()) - } } From f4e37376c66df108b2dbf9a98f962a6aed2bda26 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 2 Mar 2026 22:10:16 +0100 Subject: [PATCH 3/5] add an end-to-end test --- datafusion/physical-plan/src/sorts/sort.rs | 137 +++++++++++++++++++++ 1 file changed, 137 insertions(+) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 8c3c88a80cbf0..52e3635d31ef5 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -2723,4 +2723,141 @@ mod tests { Ok(()) } + + /// End-to-end test that verifies `ExternalSorter::sort()` atomically + /// transfers the pre-reserved merge bytes to the merge stream via `take()`. + /// + /// This test directly exercises the `ExternalSorter` code path: + /// 1. Create a sorter with a tight memory pool and insert enough data + /// to force spilling + /// 2. Call `sort()` to get the merge stream + /// 3. Verify that dropping the sorter does NOT free the pre-reserved + /// bytes back to the pool (they should have been transferred to + /// the merge stream) + /// 4. Simulate contention: a task grabs all available pool memory + /// 5. Verify the merge stream still works (it has its own pre-reserved bytes) + /// + /// Before the fix, main (using `new_empty()`), step 3 fails: the sorter drop frees + /// `sort_spill_reservation_bytes` back to the pool, and the task can + /// steal them, causing the merge stream to starve. + /// + /// With the fix (using `take()`), the bytes are atomically transferred + /// to the merge stream. The sorter drop frees 0 bytes, so there's + /// nothing for the task to steal. + #[tokio::test] + async fn test_sort_merge_reservation_transferred_not_freed() -> Result<()> { + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, + }; + use futures::TryStreamExt; + + let sort_spill_reservation_bytes: usize = 10 * 1024; // 10 KB + + // Pool: merge reservation (10KB) + enough room for sort to work. + // The room must accommodate batch data accumulation before spilling. + let sort_working_memory: usize = 40 * 1024; // 40 KB for sort operations + let pool_size = sort_spill_reservation_bytes + sort_working_memory; + let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_size)); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + + let metrics_set = ExecutionPlanMetricsSet::new(); + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); + + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(), + 128, // batch_size + sort_spill_reservation_bytes, + usize::MAX, // sort_in_place_threshold_bytes (high to avoid concat path) + SpillCompression::Uncompressed, + &metrics_set, + Arc::clone(&runtime), + )?; + + // Insert enough data to force spilling. Each batch is ~400 bytes + // (100 rows × 4 bytes). With 40KB of working memory, we'll spill + // after accumulating ~100 batches worth. 200 batches guarantees + // multiple spill cycles. + let num_batches = 200; + for i in 0..num_batches { + let values: Vec = ((i * 100)..((i + 1) * 100)).rev().collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + )?; + sorter.insert_batch(batch).await?; + } + + assert!( + sorter.spilled_before(), + "Test requires spilling to exercise the merge path" + ); + + // Call sort() to get the merge stream. After this: + // - With take() (the fix): merge_reservation = 0, merge stream has R bytes + // - With new_empty() (before fix): merge_reservation = R, merge stream has 0 bytes + let merge_stream = sorter.sort().await?; + + // Record pool state before dropping the sorter + let reserved_before_drop = pool.reserved(); + + // Drop the sorter. This frees merge_reservation: + // - With take() (the fix): frees 0 bytes (already transferred to merge stream) + // - With new_empty() (before fix): frees R bytes back to pool + drop(sorter); + + let reserved_after_drop = pool.reserved(); + + // THE KEY ASSERTION: dropping the sorter should NOT free the + // pre-reserved merge bytes. They must have been transferred to + // the merge stream via take(). + assert_eq!( + reserved_after_drop, + reserved_before_drop, + "Dropping the sorter freed {} bytes back to the pool! \ + The merge reservation bytes should have been transferred \ + to the merge stream (via take()), not freed back to the pool \ + (via new_empty()). Freed bytes can be stolen by concurrent \ + partitions, causing merge starvation.", + reserved_before_drop - reserved_after_drop + ); + + // Simulate contention: a task (representing another partition) + // grabs all available pool memory + let task = MemoryConsumer::new("TaskPartition").register(&pool); + let available = pool_size.saturating_sub(pool.reserved()); + if available > 0 { + task.try_grow(available).unwrap(); + } + + // The merge stream should still work because it holds the + // pre-reserved bytes (transferred via take()) + let batches: Vec = merge_stream.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, + (num_batches * 100) as usize, + "Merge stream should produce all rows even under memory contention" + ); + + // Verify data is sorted + let merged = concat_batches(&schema, &batches)?; + let col = merged.column(0).as_primitive::(); + for i in 1..col.len() { + assert!( + col.value(i - 1) <= col.value(i), + "Output should be sorted, but found {} > {} at index {}", + col.value(i - 1), + col.value(i), + i + ); + } + + drop(task); + Ok(()) + } } From acd7caa3c85b38e9e8c5345e912b2916d4316031 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Mar 2026 20:07:49 +0800 Subject: [PATCH 4/5] update --- datafusion/physical-plan/src/sorts/sort.rs | 115 +++++++++++---------- 1 file changed, 59 insertions(+), 56 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 52e3635d31ef5..0100e54424950 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -342,11 +342,6 @@ impl ExternalSorter { /// 2. A combined streaming merge incorporating both in-memory /// batches and data from spill files on disk. async fn sort(&mut self) -> Result { - // Release the memory reserved for merge back to the pool so - // there is some left when `in_mem_sort_stream` requests an - // allocation. - self.merge_reservation.free(); - if self.spilled_before() { // Sort `in_mem_batches` and spill it first. If there are many // `in_mem_batches` and the memory limit is almost reached, merging @@ -373,6 +368,11 @@ impl ExternalSorter { .with_reservation(self.merge_reservation.take()) .build() } else { + // Release the memory reserved for merge back to the pool so + // there is some left when `in_mem_sort_stream` requests an + // allocation. Only needed for the non-spill path; the spill + // path transfers the reservation to the merge stream instead. + self.merge_reservation.free(); self.in_mem_sort_stream(self.metrics.baseline.clone()) } } @@ -382,6 +382,12 @@ impl ExternalSorter { self.reservation.size() } + /// How much memory is reserved for the merge phase? + #[cfg(test)] + fn merge_reservation_size(&self) -> usize { + self.merge_reservation.size() + } + /// How many bytes have been spilled to disk? fn spilled_bytes(&self) -> usize { self.metrics.spill_metrics.spilled_bytes.value() @@ -2724,26 +2730,22 @@ mod tests { Ok(()) } - /// End-to-end test that verifies `ExternalSorter::sort()` atomically - /// transfers the pre-reserved merge bytes to the merge stream via `take()`. + /// Verifies that `ExternalSorter::sort()` transfers the pre-reserved + /// merge bytes to the merge stream via `take()`, rather than leaving + /// them in the sorter (via `new_empty()`). /// - /// This test directly exercises the `ExternalSorter` code path: /// 1. Create a sorter with a tight memory pool and insert enough data /// to force spilling - /// 2. Call `sort()` to get the merge stream - /// 3. Verify that dropping the sorter does NOT free the pre-reserved - /// bytes back to the pool (they should have been transferred to - /// the merge stream) - /// 4. Simulate contention: a task grabs all available pool memory - /// 5. Verify the merge stream still works (it has its own pre-reserved bytes) - /// - /// Before the fix, main (using `new_empty()`), step 3 fails: the sorter drop frees - /// `sort_spill_reservation_bytes` back to the pool, and the task can - /// steal them, causing the merge stream to starve. + /// 2. Verify `merge_reservation` holds the pre-reserved bytes before sort + /// 3. Call `sort()` to get the merge stream + /// 4. Verify `merge_reservation` is now 0 (bytes transferred to merge stream) + /// 5. Simulate contention: a competing consumer grabs all available pool memory + /// 6. Verify the merge stream still works (it uses its pre-reserved bytes + /// as initial budget, not requesting from pool starting at 0) /// - /// With the fix (using `take()`), the bytes are atomically transferred - /// to the merge stream. The sorter drop frees 0 bytes, so there's - /// nothing for the task to steal. + /// With `new_empty()` (before fix), step 4 fails: `merge_reservation` + /// still holds the bytes, the merge stream starts with 0 budget, and + /// those bytes become unaccounted-for reserved memory that nobody uses. #[tokio::test] async fn test_sort_merge_reservation_transferred_not_freed() -> Result<()> { use datafusion_execution::memory_pool::{ @@ -2778,10 +2780,7 @@ mod tests { Arc::clone(&runtime), )?; - // Insert enough data to force spilling. Each batch is ~400 bytes - // (100 rows × 4 bytes). With 40KB of working memory, we'll spill - // after accumulating ~100 batches worth. 200 batches guarantees - // multiple spill cycles. + // Insert enough data to force spilling. let num_batches = 200; for i in 0..num_batches { let values: Vec = ((i * 100)..((i + 1) * 100)).rev().collect(); @@ -2797,45 +2796,49 @@ mod tests { "Test requires spilling to exercise the merge path" ); - // Call sort() to get the merge stream. After this: - // - With take() (the fix): merge_reservation = 0, merge stream has R bytes - // - With new_empty() (before fix): merge_reservation = R, merge stream has 0 bytes - let merge_stream = sorter.sort().await?; - - // Record pool state before dropping the sorter - let reserved_before_drop = pool.reserved(); - - // Drop the sorter. This frees merge_reservation: - // - With take() (the fix): frees 0 bytes (already transferred to merge stream) - // - With new_empty() (before fix): frees R bytes back to pool - drop(sorter); + // Before sort(), merge_reservation holds sort_spill_reservation_bytes. + assert!( + sorter.merge_reservation_size() >= sort_spill_reservation_bytes, + "merge_reservation should hold the pre-reserved bytes before sort()" + ); - let reserved_after_drop = pool.reserved(); + // Call sort() to get the merge stream. With the fix (take()), + // the pre-reserved merge bytes are transferred to the merge + // stream. Without the fix (free() + new_empty()), the bytes + // are released back to the pool and the merge stream starts + // with 0 bytes. + let merge_stream = sorter.sort().await?; - // THE KEY ASSERTION: dropping the sorter should NOT free the - // pre-reserved merge bytes. They must have been transferred to - // the merge stream via take(). + // THE KEY ASSERTION: after sort(), merge_reservation must be 0. + // This proves take() transferred the bytes to the merge stream, + // rather than them being freed back to the pool where other + // partitions could steal them. assert_eq!( - reserved_after_drop, - reserved_before_drop, - "Dropping the sorter freed {} bytes back to the pool! \ - The merge reservation bytes should have been transferred \ - to the merge stream (via take()), not freed back to the pool \ - (via new_empty()). Freed bytes can be stolen by concurrent \ - partitions, causing merge starvation.", - reserved_before_drop - reserved_after_drop + sorter.merge_reservation_size(), + 0, + "After sort(), merge_reservation should be 0 (bytes transferred \ + to merge stream via take()). If non-zero, the bytes are still \ + held by the sorter and will be freed on drop, allowing other \ + partitions to steal them." ); - // Simulate contention: a task (representing another partition) - // grabs all available pool memory - let task = MemoryConsumer::new("TaskPartition").register(&pool); + // Drop the sorter to free its reservations back to the pool. + drop(sorter); + + // Simulate contention: another partition grabs ALL available + // pool memory. If the merge stream didn't receive the + // pre-reserved bytes via take(), it will fail when it tries + // to allocate memory for reading spill files. + let contender = MemoryConsumer::new("CompetingPartition").register(&pool); let available = pool_size.saturating_sub(pool.reserved()); if available > 0 { - task.try_grow(available).unwrap(); + contender.try_grow(available).unwrap(); } - // The merge stream should still work because it holds the - // pre-reserved bytes (transferred via take()) + // The merge stream must still produce correct results despite + // the pool being fully consumed by the contender. This only + // works if sort() transferred the pre-reserved bytes to the + // merge stream (via take()) rather than freeing them. let batches: Vec = merge_stream.try_collect().await?; let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!( @@ -2857,7 +2860,7 @@ mod tests { ); } - drop(task); + drop(contender); Ok(()) } } From c478d435ba25b2f8b96cf3cb44d8edf8528a354f Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Fri, 13 Mar 2026 19:04:43 +0800 Subject: [PATCH 5/5] address comments --- datafusion/physical-plan/src/sorts/builder.rs | 21 ++++-- .../src/sorts/multi_level_merge.rs | 71 ++++++++++--------- 2 files changed, 55 insertions(+), 37 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 881b995cf9d51..a462b832056bd 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -93,10 +93,7 @@ impl BatchBuilder { // Only request additional memory from the pool when actual batch // usage exceeds the current reservation (which may include // pre-reserved bytes from sort_spill_reservation_bytes). - if self.batches_mem_used > self.reservation.size() { - self.reservation - .try_grow(self.batches_mem_used - self.reservation.size())?; - } + try_grow_reservation_to_at_least(&mut self.reservation, self.batches_mem_used)?; let batch_idx = self.batches.len(); self.batches.push((stream_idx, batch)); self.cursors[stream_idx] = BatchCursor { @@ -187,3 +184,19 @@ impl BatchBuilder { )?)) } } + +/// Try to grow `reservation` so it covers at least `needed` bytes. +/// +/// When a reservation has been pre-loaded with bytes (e.g. via +/// [`MemoryReservation::take()`]), this avoids redundant pool +/// allocations: if the reservation already covers `needed`, this is +/// a no-op; otherwise only the deficit is requested from the pool. +pub(crate) fn try_grow_reservation_to_at_least( + reservation: &mut MemoryReservation, + needed: usize, +) -> Result<()> { + if needed > reservation.size() { + reservation.try_grow(needed - reservation.size())?; + } + Ok(()) +} diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index d5b0c2ad8e5c1..8985e1d8c70ee 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -30,6 +30,7 @@ use arrow::datatypes::SchemaRef; use datafusion_common::Result; use datafusion_execution::memory_pool::MemoryReservation; +use crate::sorts::builder::try_grow_reservation_to_at_least; use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::stream::RecordBatchStreamAdapter; @@ -274,6 +275,15 @@ impl MultiLevelMergeBuilder { let is_only_merging_memory_streams = sorted_spill_files.is_empty(); + // If no spill files were selected (e.g. all too large for + // available memory but enough in-memory streams exist), + // return the pre-reserved bytes to self.reservation so + // create_new_merge_sort can transfer them to the merge + // stream's BatchBuilder. + if is_only_merging_memory_streams { + mem::swap(&mut self.reservation, &mut memory_reservation); + } + for spill in sorted_spill_files { let stream = self .spill_manager @@ -377,43 +387,38 @@ impl MultiLevelMergeBuilder { ) * buffer_len; total_needed += per_spill; - // Only request additional memory from the pool when total needed - // exceeds what's already reserved (which may include pre-reserved - // bytes from sort_spill_reservation_bytes). - if total_needed > reservation.size() { - match reservation.try_grow(total_needed - reservation.size()) { - Ok(_) => { - number_of_spills_to_read_for_current_phase += 1; - } - // If we can't grow the reservation, we need to stop - Err(err) => { - // We must have at least 2 streams to merge, so if we don't have enough memory - // fail - if minimum_number_of_required_streams - > number_of_spills_to_read_for_current_phase - { - // Free the memory we reserved for this merge as we either try again or fail - reservation.free(); - if buffer_len > 1 { - // Try again with smaller buffer size, it will be slower but at least we can merge - return self.get_sorted_spill_files_to_merge( - buffer_len - 1, - minimum_number_of_required_streams, - reservation, - ); - } - - return Err(err); + // For memory pools that are not shared this is good, for other + // this is not and there should be some upper limit to memory + // reservation so we won't starve the system. + match try_grow_reservation_to_at_least(reservation, total_needed) { + Ok(_) => { + number_of_spills_to_read_for_current_phase += 1; + } + // If we can't grow the reservation, we need to stop + Err(err) => { + // We must have at least 2 streams to merge, so if we don't have enough memory + // fail + if minimum_number_of_required_streams + > number_of_spills_to_read_for_current_phase + { + // Free the memory we reserved for this merge as we either try again or fail + reservation.free(); + if buffer_len > 1 { + // Try again with smaller buffer size, it will be slower but at least we can merge + return self.get_sorted_spill_files_to_merge( + buffer_len - 1, + minimum_number_of_required_streams, + reservation, + ); } - // We reached the maximum amount of memory we can use - // for this merge - break; + return Err(err); } + + // We reached the maximum amount of memory we can use + // for this merge + break; } - } else { - // Pre-reserved bytes cover this spill file's buffer - number_of_spills_to_read_for_current_phase += 1; } }