From 94c1c16f8c07582fc1da065bae4cd2e02b99caf8 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 27 Aug 2026 15:57:26 -0700 Subject: [PATCH 1/2] fix: preserve external sort workspace across spilling Keep acquired spill workspace available to the sorter's cursor, row, and batch reservations instead of returning it to the execution pool before the spill merge. Reuse idle workspace for chunked sorted output while keeping additional output growth under the original sort consumer. Release the idle reserve after the final output path has selected its budget, preserving ownership and accounting for live reservations. Add regressions for reduced memory availability, intermediate and skewed merge retries, parent pool limits, and cleanup on completion or cancellation. Closes #24739. --- .../src/sorts/merge_memory_pool.rs | 459 +++++++++ datafusion/physical-plan/src/sorts/mod.rs | 1 + .../src/sorts/multi_level_merge.rs | 112 ++ datafusion/physical-plan/src/sorts/sort.rs | 84 +- .../src/sorts/sort/spill_tests.rs | 958 ++++++++++++++++++ .../src/sorts/streaming_merge.rs | 10 + 6 files changed, 1598 insertions(+), 26 deletions(-) create mode 100644 datafusion/physical-plan/src/sorts/merge_memory_pool.rs create mode 100644 datafusion/physical-plan/src/sorts/sort/spill_tests.rs diff --git a/datafusion/physical-plan/src/sorts/merge_memory_pool.rs b/datafusion/physical-plan/src/sorts/merge_memory_pool.rs new file mode 100644 index 0000000000000..4e8ae269ccb79 --- /dev/null +++ b/datafusion/physical-plan/src/sorts/merge_memory_pool.rs @@ -0,0 +1,459 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shares the sorter's reserved merge workspace without returning it to the +//! execution pool between sorting batches and the merge's cursor reservations. + +use std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion_common::{Result, resources_err}; +use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use parking_lot::Mutex; + +#[derive(Debug)] +pub(super) struct MergeMemoryPool { + parent: Arc, + state: Mutex, +} + +#[derive(Debug)] +struct MergeMemoryState { + /// The only reservation charged to the execution pool. + reservation: MemoryReservation, + /// Total usage of all reservations and workspace loans in this pool, + /// including siblings created with `new_empty` or `split`. + used: usize, + /// Workspace already acquired for the next merge. It is reusable by any + /// child reservation, but unavailable to other execution-pool consumers. + retained: usize, +} + +/// Temporarily lends already-reserved workspace to per-batch sorting. Any +/// additional capacity must still be acquired by the ordinary sort consumer. +pub(super) struct WorkspaceLoan { + pool: Arc, + size: usize, +} + +impl WorkspaceLoan { + pub(super) fn size(&self) -> usize { + self.size + } + + pub(super) fn shrink(&mut self, size: usize) { + self.size = self + .size + .checked_sub(size) + .expect("workspace loan underflow"); + if size != 0 { + self.pool.release(size); + } + } +} + +impl Drop for WorkspaceLoan { + fn drop(&mut self) { + self.shrink(self.size); + } +} + +impl MergeMemoryPool { + pub(super) fn new(parent: Arc, consumer: MemoryConsumer) -> Self { + let reservation = consumer.register(&parent); + Self { + parent, + state: Mutex::new(MergeMemoryState { + reservation, + used: 0, + retained: 0, + }), + } + } + + /// Retain workspace only after it has been successfully acquired. + pub(super) fn retain(&self, size: usize) { + let mut state = self.state.lock(); + assert!(size <= state.reservation.size()); + state.retained = size; + } + + /// Lend up to `size` unused bytes without acquiring more parent memory. + /// Count the loan alongside merge reservations so they cannot spend the + /// same credit. Dropping the loan returns any outstanding bytes. + pub(super) fn borrow(self: &Arc, size: usize) -> WorkspaceLoan { + let mut state = self.state.lock(); + let size = size.min(state.reservation.size() - state.used); + state.used += size; + WorkspaceLoan { + pool: Arc::clone(self), + size, + } + } + + /// Stop retaining idle workspace when no further spill or intermediate merge + /// needs it. Live child reservations and loans remain charged. + pub(super) fn release_unused(&self) { + let mut state = self.state.lock(); + state.retained = 0; + state.reservation.resize(state.used); + } + + fn release(&self, size: usize) { + let mut state = self.state.lock(); + state.used = state.used.checked_sub(size).expect("memory underflow"); + state.reservation.resize(state.used.max(state.retained)); + } +} + +impl Display for MergeMemoryPool { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "MergeMemoryPool") + } +} + +impl MemoryPool for MergeMemoryPool { + fn name(&self) -> &str { + "MergeMemoryPool" + } + + fn grow(&self, _reservation: &MemoryReservation, additional: usize) { + let mut state = self.state.lock(); + let used = state.used.checked_add(additional).expect("memory overflow"); + if used > state.reservation.size() { + state.reservation.resize(used); + } + state.used = used; + } + + fn try_grow( + &self, + _reservation: &MemoryReservation, + additional: usize, + ) -> Result<()> { + let mut state = self.state.lock(); + let Some(used) = state.used.checked_add(additional) else { + return resources_err!("Sort merge memory reservation overflow"); + }; + if used > state.reservation.size() { + state.reservation.try_resize(used)?; + } + state.used = used; + Ok(()) + } + + fn shrink(&self, _reservation: &MemoryReservation, subtractive: usize) { + self.release(subtractive); + } + + fn reserved(&self) -> usize { + self.state.lock().reservation.size() + } + + fn memory_limit(&self) -> MemoryLimit { + self.parent.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_execution::memory_pool::GreedyMemoryPool; + + fn reservation( + parent: &Arc, + ) -> (Arc, MemoryReservation) { + let pool = Arc::new(MergeMemoryPool::new( + Arc::clone(parent), + MemoryConsumer::new("merge workspace"), + )); + let reservation = MemoryConsumer::new("merge children") + .register(&(Arc::clone(&pool) as Arc)); + (pool, reservation) + } + + #[test] + fn reuses_workspace_across_child_reservations() -> Result<()> { + let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); + let (pool, reservation) = reservation(&parent); + reservation.try_grow(60)?; + pool.retain(60); + reservation.free(); + + let contender = MemoryConsumer::new("contender").register(&parent); + contender.try_grow(40)?; + assert_eq!(parent.reserved(), 100); + + let cursor = reservation.new_empty(); + let rows = reservation.new_empty(); + cursor.try_grow(20)?; + rows.try_grow(40)?; + assert!(rows.try_grow(1).is_err()); + assert_eq!(rows.size(), 40); + assert_eq!(parent.reserved(), 100); + + drop(cursor); + drop(rows); + assert_eq!(pool.reserved(), 60); + reservation.try_grow(60)?; // Workspace for the next spill needs no new grant. + assert_eq!(parent.reserved(), 100); + + drop(reservation); + drop(pool); + assert_eq!(parent.reserved(), 40); + drop(contender); + assert_eq!(parent.reserved(), 0); + Ok(()) + } + + #[test] + fn child_keeps_workspace_alive_after_sorter_drops() -> Result<()> { + let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); + let (pool, mut reservation) = reservation(&parent); + reservation.try_grow(60)?; + pool.retain(60); + let child = reservation.take(); + let sibling = child.split(20); + drop(reservation); + drop(pool); + assert_eq!(parent.reserved(), 60); + + child.free(); + assert_eq!(parent.reserved(), 60); + child.try_grow(40)?; + drop(child); + assert_eq!(parent.reserved(), 60); + drop(sibling); + assert_eq!(parent.reserved(), 0); + Ok(()) + } + + #[test] + fn releases_excess_capacity_but_retains_workspace() -> Result<()> { + let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); + let (pool, reservation) = reservation(&parent); + reservation.try_grow(40)?; + pool.retain(40); + reservation.try_grow(50)?; + assert_eq!(parent.reserved(), 90); + reservation.shrink(60); + assert_eq!(parent.reserved(), 40); + let loan = pool.borrow(10); + assert_eq!(loan.size(), 10); + pool.release_unused(); + assert_eq!(parent.reserved(), 40); // Live loans remain charged. + drop(loan); + assert_eq!(parent.reserved(), 30); + drop(reservation); + assert_eq!(parent.reserved(), 0); + Ok(()) + } + + #[test] + fn does_not_reserve_workspace_implicitly() -> Result<()> { + let parent: Arc = Arc::new(GreedyMemoryPool::new(8)); + let (pool, reservation) = reservation(&parent); + reservation.grow(0); + reservation.try_grow(0)?; + assert_eq!(parent.reserved(), 0); + reservation.try_grow(8)?; + assert!(reservation.try_grow(1).is_err()); + assert!(reservation.try_grow(usize::MAX).is_err()); + assert_eq!(reservation.size(), 8); + assert_eq!(pool.reserved(), 8); + reservation.free(); + assert_eq!(parent.reserved(), 0); + Ok(()) + } + + #[test] + fn delegates_infallible_growth_to_parent() { + let parent: Arc = Arc::new(GreedyMemoryPool::new(8)); + let (pool, reservation) = reservation(&parent); + reservation.grow(16); + assert_eq!(parent.reserved(), 16); + assert_eq!(pool.reserved(), 16); + drop(reservation); + assert_eq!(parent.reserved(), 0); + } + + #[test] + fn workspace_loans_share_credit_with_cursors() -> Result<()> { + let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); + let (pool, reservation) = reservation(&parent); + reservation.try_grow(60)?; + pool.retain(60); + reservation.free(); + let contender = MemoryConsumer::new("contender").register(&parent); + contender.try_grow(40)?; + let cursor = reservation.new_empty(); + cursor.try_grow(20)?; + + // Both loans stay alive until after both requests complete. Together they + // may claim only the 40 bytes not already assigned to the cursor. + let gate = std::sync::Barrier::new(2); + let (mut first, second) = std::thread::scope(|scope| { + let first = scope.spawn(|| { + gate.wait(); + pool.borrow(30) + }); + let second = scope.spawn(|| { + gate.wait(); + pool.borrow(30) + }); + (first.join().unwrap(), second.join().unwrap()) + }); + assert_eq!(first.size() + second.size(), 40); + assert!(matches!((first.size(), second.size()), (30, 10) | (10, 30))); + assert_eq!(pool.borrow(1).size(), 0); + assert_eq!(parent.reserved(), 100); + + first.shrink(5); + cursor.try_grow(5)?; + assert_eq!(first.size() + second.size() + cursor.size(), 60); + assert_eq!(pool.borrow(1).size(), 0); + assert!(cursor.try_grow(1).is_err()); + assert_eq!(cursor.size(), 25); + assert_eq!(parent.reserved(), 100); + + let returned = first.size(); + drop(first); + let replacement = pool.borrow(usize::MAX); + assert_eq!(replacement.size(), returned); + assert_eq!(parent.reserved(), 100); + + // The final loan owns the workspace even after the sorter and its child + // reservations are gone; dropping it releases the parent allocation once. + drop(reservation); + drop(pool); + drop(cursor); + drop(second); + assert_eq!(parent.reserved(), 100); + drop(replacement); + assert_eq!(parent.reserved(), 40); + drop(contender); + assert_eq!(parent.reserved(), 0); + Ok(()) + } + + #[test] + fn workspace_loan_preserves_ordinary_consumer_fair_share() -> Result<()> { + use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool}; + use std::num::NonZeroUsize; + + let tracked = Arc::new(TrackConsumersPool::new( + FairSpillPool::new(100), + NonZeroUsize::new(3).unwrap(), + )); + let parent: Arc = Arc::clone(&tracked) as Arc; + let (pool, reservation) = reservation(&parent); + reservation.try_grow(20)?; + pool.retain(20); + reservation.free(); + let ordinary = MemoryConsumer::new("ordinary sort") + .with_can_spill(true) + .register(&parent); + let contender = MemoryConsumer::new("other spillable sort") + .with_can_spill(true) + .register(&parent); + ordinary.try_grow(40)?; + assert_eq!(parent.reserved(), 60); + + let snapshot = || { + let mut metrics = tracked + .metrics() + .into_iter() + .map(|m| (m.name, m.can_spill, m.reserved, m.peak)) + .collect::>(); + metrics.sort_unstable(); + metrics + }; + let before = snapshot(); + assert_eq!(before.len(), 3); + + // Although 40 bytes remain globally free, the ordinary consumer's share + // is (100 - 20) / 2 = 40. Only the acquired 20 bytes may be borrowed; + // the remaining 5 must be denied under the original consumer's policy. + let result: Result<()> = (|| { + let sorted_size = 65; + let loan = pool.borrow(sorted_size - ordinary.size()); + assert_eq!(loan.size(), 20); + ordinary.try_resize(sorted_size - loan.size())?; + Ok(()) + })(); + assert!(result.is_err()); + assert_eq!(ordinary.size(), 40); + assert_eq!(parent.reserved(), 60); + assert_eq!(snapshot(), before); + + // Failed resize returned its loan, so a smaller expansion can now use + // existing credit without a new grant or a different parent consumer. + let loan = pool.borrow(55 - ordinary.size()); + assert_eq!(loan.size(), 15); + ordinary.try_resize(55 - loan.size())?; + assert_eq!(ordinary.size(), 40); + assert_eq!(parent.reserved(), 60); + assert_eq!(snapshot(), before); + + drop(loan); + drop(ordinary); + drop(contender); + drop(reservation); + drop(pool); + assert_eq!(parent.reserved(), 0); + assert!(tracked.metrics().is_empty()); + Ok(()) + } + + #[test] + fn failed_sorted_resize_returns_workspace_loan() -> Result<()> { + let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); + let (pool, reservation) = reservation(&parent); + assert_eq!(pool.borrow(usize::MAX).size(), 0); + assert_eq!(parent.reserved(), 0); + reservation.try_grow(60)?; + pool.retain(60); + reservation.free(); + let ordinary = MemoryConsumer::new("ordinary sort") + .with_can_spill(true) + .register(&parent); + ordinary.try_grow(40)?; + + let result: Result<()> = (|| { + let sorted_size = 120; + let loan = pool.borrow(sorted_size - ordinary.size()); + assert_eq!(loan.size(), 60); + ordinary.try_resize(sorted_size - loan.size())?; + Ok(()) + })(); + assert!(result.is_err()); + assert_eq!(ordinary.size(), 40); + assert_eq!(parent.reserved(), 100); + let restored = pool.borrow(usize::MAX); + assert_eq!(restored.size(), 60); + assert_eq!(parent.reserved(), 100); + + drop(restored); + drop(ordinary); + drop(reservation); + drop(pool); + assert_eq!(parent.reserved(), 0); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs index 4af2a3629fa9c..125fca1abbf51 100644 --- a/datafusion/physical-plan/src/sorts/mod.rs +++ b/datafusion/physical-plan/src/sorts/mod.rs @@ -20,6 +20,7 @@ mod builder; mod cursor; mod merge; +mod merge_memory_pool; #[cfg(test)] mod merge_trace; mod multi_level_merge; diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 3ec52cc70c0a9..8ff845ef92652 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -31,6 +31,7 @@ use datafusion_common::{Result, internal_err, resources_err}; use datafusion_execution::memory_pool::MemoryReservation; use crate::sorts::builder::try_grow_reservation_to_at_least; +use crate::sorts::merge_memory_pool::MergeMemoryPool; use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; @@ -153,6 +154,8 @@ pub(crate) struct MultiLevelMergeBuilder { metrics: BaselineMetrics, batch_size: usize, reservation: MemoryReservation, + /// Workspace retained across retries and intermediate spill passes. + merge_pool: Option>, fetch: Option, enable_round_robin_tie_breaker: bool, } @@ -191,11 +194,17 @@ impl MultiLevelMergeBuilder { metrics, batch_size, reservation, + merge_pool: None, enable_round_robin_tie_breaker, fetch, } } + pub(super) fn with_merge_pool(mut self, pool: Option>) -> Self { + self.merge_pool = pool; + self + } + pub(crate) fn create_spillable_merge_stream(self) -> SendableRecordBatchStream { Box::pin(RecordBatchStreamAdapter::new( Arc::clone(&self.schema), @@ -233,6 +242,12 @@ impl MultiLevelMergeBuilder { "We should not have any sorted streams left" ); + // The final pass has its buffer budget and needs no future spill + // workspace. Keep live reservations, but release the idle floor. + if let Some(pool) = &self.merge_pool { + pool.release_unused(); + } + return Ok(stream); } @@ -844,6 +859,103 @@ mod tests { /// Two sorted runs whose largest batches are too big to both /// be seated in the merge budget at once are re-spilled (halved) until they /// fit, and the merge then completes with fully sorted, complete output. + #[tokio::test] + async fn skewed_runs_reuse_retained_workspace_across_retries() -> Result<()> { + // These budgets require one and two re-spills, respectively. + for (budget_halves, expected_splits) in [(7, 1), (5, 2)] { + let capacity = 1024 * 1024; + let parent: Arc = Arc::new(GreedyMemoryPool::new(capacity)); + let env = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&parent)) + .build_arc()?; + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + let spill_count = spill_manager.metrics.spill_file_count.clone(); + let n: i64 = 16384; + let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); + let workspace = m * budget_halves / 2; + assert!(workspace < capacity); + let merge_pool = Arc::new(MergeMemoryPool::new( + Arc::clone(&parent), + MemoryConsumer::new("merge workspace"), + )); + let pool: Arc = + Arc::clone(&merge_pool) as Arc; + let batch_size = 8192; + let mut builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![f0, f1], + &pool, + batch_size, + ) + .with_merge_pool(Some(Arc::clone(&merge_pool))); + builder.reservation.try_grow(workspace)?; + merge_pool.retain(workspace); + let contender = MemoryConsumer::new("contender").register(&parent); + contender.try_grow(capacity - workspace)?; + + // Drive each retry so its freed child reservation cannot silently + // release the workspace and reacquire it from the parent pool. + for _ in 0..expected_splits { + let MergeStep::SplitThenRetry(index) = + builder.merge_sorted_runs_within_mem_limit()? + else { + panic!("the merge must re-spill a skewed run"); + }; + assert_eq!(parent.reserved(), capacity); + assert!(contender.try_grow(1).is_err()); + builder.split_spill_file_in_half(index).await?; + assert_eq!(parent.reserved(), capacity); + assert!(contender.try_grow(1).is_err()); + } + + let final_bytes: usize = builder + .sorted_spill_files + .iter() + .map(|(file, _)| { + get_reserved_bytes_for_record_batch_size( + file.max_record_batch_memory, + file.max_record_batch_memory, + ) + }) + .sum(); + assert!(final_bytes < workspace && workspace < 2 * final_bytes); + + let mut stream = builder.create_spillable_merge_stream(); + let first = stream.try_next().await?.expect("nonempty merge"); + assert_eq!(spill_count.value(), 2 + expected_splits); + assert_eq!(merge_pool.reserved(), final_bytes); + assert_eq!(parent.reserved(), contender.size() + final_bytes); + + let mut batches = vec![first]; + while let Some(batch) = stream.try_next().await? { + batches.push(batch); + } + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).max(), + Some(batch_size / 2), + "splitting both runs must halve, not quarter, the output batch size" + ); + let merged = concat_batches(&schema, &batches)?; + let expected = + Int64Array::from_iter_values((0..n).flat_map(|value| [value, value])); + assert_eq!(merged.column(0).as_primitive::(), &expected); + + // The pool stays alive: EOF must release live bytes and the idle floor. + assert_eq!(merge_pool.reserved(), 0); + assert_eq!(parent.reserved(), contender.size()); + drop(stream); + assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0); + assert_eq!(env.disk_manager.used_disk_space(), 0); + drop(contender); + assert_eq!(parent.reserved(), 0); + } + Ok(()) + } + #[tokio::test] async fn skewed_runs_are_respilled_so_the_merge_fits() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 09655648ab62b..7726fb00c193b 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -42,6 +42,7 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::IncrementalSortIterator; +use crate::sorts::merge_memory_pool::MergeMemoryPool; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::get_record_batch_memory_size; use crate::spill::in_progress_spill_file::InProgressSpillFile; @@ -67,7 +68,7 @@ use datafusion_common::{ unwrap_or_internal_err, }; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; @@ -76,6 +77,9 @@ use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; use futures::{StreamExt, TryStreamExt}; use log::{debug, trace}; +#[cfg(test)] +mod spill_tests; + struct ExternalSorterMetrics { /// metrics baseline: BaselineMetrics, @@ -261,6 +265,9 @@ struct ExternalSorter { /// might spill, `sort_spill_reservation_bytes` will be /// pre-reserved to ensure there is some space for this sort/merge. merge_reservation: MemoryReservation, + /// Keeps that workspace available to the merge's cursor, row, and batch + /// reservations even when the execution pool cannot grant more memory. + merge_pool: Arc, /// How much memory to reserve for performing in-memory sort/merges /// prior to spilling. sort_spill_reservation_bytes: usize, @@ -287,9 +294,13 @@ impl ExternalSorter { .with_can_spill(true) .register(&runtime.memory_pool); - let merge_reservation = - MemoryConsumer::new(format!("ExternalSorterMerge[{partition_id}]")) - .register(&runtime.memory_pool); + let merge_name = format!("ExternalSorterMerge[{partition_id}]"); + let merge_pool = Arc::new(MergeMemoryPool::new( + Arc::clone(&runtime.memory_pool), + MemoryConsumer::new(&merge_name), + )); + let merge_reservation = MemoryConsumer::new(merge_name) + .register(&(Arc::clone(&merge_pool) as Arc)); let spill_manager = SpillManager::new( Arc::clone(&runtime), @@ -308,6 +319,7 @@ impl ExternalSorter { reservation, spill_manager, merge_reservation, + merge_pool, runtime, batch_size, sort_spill_reservation_bytes, @@ -369,12 +381,12 @@ impl ExternalSorter { .with_batch_size(self.batch_size) .with_fetch(None) .with_reservation(self.merge_reservation.take()) + .with_merge_pool(Arc::clone(&self.merge_pool)) .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. + // Final output needs no reserve for future spills. Return unused + // workspace so another sorter can start while this stream is alive. + self.merge_pool.release_unused(); self.merge_reservation.free(); self.in_mem_sort_stream(true, true) } @@ -472,10 +484,9 @@ impl ExternalSorter { "in_mem_batches must not be empty when attempting to sort and spill" ); - // Release the memory reserved for merge back to the pool so - // there is some left when `in_mem_sort_stream` requests an - // allocation. At the end of this function, memory will be - // reserved again for the next spill. + // Reuse the pre-reserved workspace across cursor, encoded-row, and + // batch reservations. Returning it to the execution pool here can + // make spilling fail if another task consumes it or our share shrinks. self.merge_reservation.free(); let mut sorted_stream = self.in_mem_sort_stream( @@ -594,6 +605,7 @@ impl ExternalSorter { coalesce_runs: bool, ) -> Result { if self.in_mem_batches.is_empty() { + self.merge_pool.release_unused(); let empty_stream = Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); return Ok(self.observe_if_output(empty_stream, is_output_stream)); @@ -611,6 +623,7 @@ impl ExternalSorter { // consumer of the stream. if self.in_mem_batches.len() == 1 { + self.merge_pool.release_unused(); let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); let sorted_stream = self.sort_batch_stream(batch, reservation)?; @@ -619,6 +632,7 @@ impl ExternalSorter { // If less than sort_in_place_threshold_bytes, concatenate and sort in place if self.reservation.size() < self.sort_in_place_threshold_bytes { + self.merge_pool.release_unused(); // Concatenate memory batches together and sort let batch = concat_batches(&self.schema, &self.in_mem_batches)?; self.in_mem_batches.clear(); @@ -726,8 +740,8 @@ impl ExternalSorter { /// sorted data and the target batch size. /// For single-batch output cases, `reservation` will be freed immediately after sorting, /// as the batch will be output and is expected to be reserved by the consumer of the stream. - /// For multi-batch output cases, `reservation` will be grown to match the actual - /// size of sorted output, and as each batch is output, its memory will be freed from the reservation. + /// For multi-batch output cases, `reservation` and any borrowed spill workspace + /// cover the sorted output, releasing its memory as each batch is output. /// (This leads to the same behaviour, as futures are only evaluated when polled by the consumer.) fn sort_batch_stream( &self, @@ -742,6 +756,7 @@ impl ExternalSorter { let schema = batch.schema(); let expressions = self.expr.clone(); let batch_size = self.batch_size; + let merge_pool = Arc::clone(&self.merge_pool); let stream = futures::stream::once(async move { let schema = batch.schema(); @@ -749,26 +764,42 @@ impl ExternalSorter { // Sort the batch immediately and get all output batches let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?; - // Resize the reservation to match the actual sorted output size. - // Using try_resize avoids a release-then-reacquire cycle, which - // matters for MemoryPool implementations where grow/shrink have - // non-trivial cost (e.g. JNI calls in Comet). + // Chunked output can retain shared buffers in every batch and + // exceed the input estimate. Borrow only already-reserved spill + // workspace; any remainder still uses the original sort consumer. let total_sorted_size: usize = sorted_batches .iter() .map(get_record_batch_memory_size) .sum(); + let mut workspace = + merge_pool.borrow(total_sorted_size.saturating_sub(reservation.size())); reservation - .try_resize(total_sorted_size) + .try_resize(total_sorted_size - workspace.size()) .map_err(Self::err_with_oom_context)?; - // Wrap in ReservationStream to hold the reservation - Result::<_, DataFusionError>::Ok(Box::pin(ReservationStream::new( - Arc::clone(&schema), - Box::pin(RecordBatchStreamAdapter::new( + if workspace.size() == 0 { + return Ok(Box::pin(ReservationStream::new( Arc::clone(&schema), - futures::stream::iter(sorted_batches.into_iter().map(Ok)), - )), - reservation, + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(sorted_batches.into_iter().map(Ok)), + )), + reservation, + )) as SendableRecordBatchStream); + } + + // Return borrowed workspace first so the merge's cursors can reuse + // it immediately. Both reservations also release on stream drop. + let batches = sorted_batches.into_iter().map(move |batch| { + let size = get_record_batch_memory_size(&batch); + let borrowed = size.min(workspace.size()); + workspace.shrink(borrowed); + reservation.shrink(size - borrowed); + Ok(batch) + }); + Result::<_, DataFusionError>::Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(batches), )) as SendableRecordBatchStream) }) .try_flatten(); @@ -788,6 +819,7 @@ impl ExternalSorter { .try_resize(size) .map_err(Self::err_with_oom_context)?; } + self.merge_pool.retain(size); } Ok(()) diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs new file mode 100644 index 0000000000000..57a41984cbfcd --- /dev/null +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -0,0 +1,958 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::{ + ExternalSorter, get_reserved_bytes_for_record_batch, sort_batch, sort_batch_chunked, +}; +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, +}; +use arrow::compute::{SortOptions, concat_batches}; +use arrow::datatypes::{DataType, Field, Int8Type, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::config::{ExecutionOptions, SpillCompression}; +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +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 futures::TryStreamExt; +use std::fmt::{Display, Formatter}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +const MERGE_BYTES: usize = 64 * 1024; +const BUFFERED_BATCHES: usize = 9; + +#[derive(Debug, Default)] +struct AllocationState { + used: usize, + limit: usize, + peak: usize, + denied: usize, + unchecked_over_limit: usize, +} + +/// Existing allocations survive a lower limit, but fresh allocations must fit. +/// The test changes the limit before insertion, independently of spill internals. +#[derive(Debug)] +struct AdjustablePool { + capacity: usize, + state: Mutex, +} + +impl AdjustablePool { + fn new(capacity: usize) -> Arc { + Arc::new(Self { + capacity, + state: Mutex::new(AllocationState { + limit: capacity, + ..Default::default() + }), + }) + } + + fn set_limit(&self, limit: usize) { + assert!(limit <= self.capacity); + self.state.lock().unwrap().limit = limit; + } +} + +impl Display for AdjustablePool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "AdjustablePool") + } +} + +impl MemoryPool for AdjustablePool { + fn name(&self) -> &str { + "AdjustablePool" + } + + fn grow(&self, _: &MemoryReservation, additional: usize) { + // Honor the infallible API, but detect attempts to evade try_grow. + let mut state = self.state.lock().unwrap(); + state.used = state.used.checked_add(additional).unwrap(); + state.peak = state.peak.max(state.used); + if additional > 0 && state.used > state.limit { + state.unchecked_over_limit += 1; + } + } + + fn shrink(&self, _: &MemoryReservation, shrink: usize) { + let mut state = self.state.lock().unwrap(); + state.used = state.used.checked_sub(shrink).unwrap(); + } + + fn try_grow(&self, _: &MemoryReservation, additional: usize) -> Result<()> { + let mut state = self.state.lock().unwrap(); + if additional == 0 { + return Ok(()); + } + let requested = state.used.checked_add(additional).unwrap(); + if requested > state.limit { + state.denied += 1; + return Err(DataFusionError::ResourcesExhausted( + "allocation limit reached".into(), + )); + } + state.used = requested; + state.peak = state.peak.max(requested); + Ok(()) + } + + fn reserved(&self) -> usize { + self.state.lock().unwrap().used + } + fn memory_limit(&self) -> MemoryLimit { + MemoryLimit::Finite(self.capacity) + } +} + +struct Fixture { + parent: RecordBatch, + batches: Vec, + ordering: LexOrdering, + capacity: usize, +} + +fn fixture() -> Result { + let rows = 1024; + let count = 30; + let total = rows * count; + let schema = Arc::new(Schema::new(vec![ + Field::new("category", DataType::Utf8, true), + Field::new("sales", DataType::Decimal128(28, 2), true), + Field::new("payload", DataType::Utf8, false), + ])); + let categories = StringArray::from_iter( + (0..total).map(|i| (i % 11 != 0).then(|| format!("category-{:04}", i % 100))), + ); + // Every non-null sales value is distinct, so sorting has no ambiguous ties. + let sales = Decimal128Array::from_iter( + (0..total).map(|i| (i != 1).then_some(((total - i) * 100) as i128)), + ) + .with_precision_and_scale(28, 2)?; + let payload = StringArray::from_iter_values( + (0..total).map(|i| format!("row-{i:08}-{}", "x".repeat(64))), + ); + let columns: Vec = + vec![Arc::new(categories), Arc::new(sales), Arc::new(payload)]; + let parent = RecordBatch::try_new(schema, columns)?; + let batches: Vec<_> = (0..count).map(|i| parent.slice(i * rows, rows)).collect(); + let ordering = [ + PhysicalSortExpr::new( + Arc::new(Column::new("category", 0)), + SortOptions { + descending: false, + nulls_first: false, + }, + ), + PhysicalSortExpr::new( + Arc::new(Column::new("sales", 1)), + SortOptions { + descending: true, + nulls_first: true, + }, + ), + ] + .into(); + let buffered = batches[..BUFFERED_BATCHES] + .iter() + .map(get_reserved_bytes_for_record_batch) + .collect::>>()? + .into_iter() + .sum::(); + let capacity = MERGE_BYTES + + buffered + + get_reserved_bytes_for_record_batch(&batches[BUFFERED_BATCHES])? / 2; + Ok(Fixture { + parent, + batches, + ordering, + capacity, + }) +} + +fn sorter( + fixture: &Fixture, + pool: Arc, +) -> Result<(ExternalSorter, Arc)> { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc()?; + let sorter = ExternalSorter::new( + 0, + fixture.parent.schema(), + fixture.ordering.clone(), + 128, + MERGE_BYTES, + 0, // Force per-batch sorting and merging, rather than concatenation. + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + )?; + Ok((sorter, runtime)) +} + +async fn assert_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 + || runtime.disk_manager.spilling_progress().active_files_count != 0 + || runtime.disk_manager.used_disk_space() != 0 + { + tokio::task::yield_now().await; + } + }) + .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); +} + +#[tokio::test] +async fn test_spill_preserves_merge_workspace_after_limit_decreases() -> Result<()> { + let fixture = fixture()?; + let pool = AdjustablePool::new(fixture.capacity); + let (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?; + for batch in &fixture.batches[..BUFFERED_BATCHES] { + sorter.insert_batch(batch.clone()).await?; + } + assert!(!sorter.spilled_before()); + let reduced_limit = fixture.capacity * 3 / 4; + assert!(pool.reserved() > reduced_limit); + pool.set_limit(reduced_limit); + // Unlike a hook on free(), this pressure transition still occurs when the + // implementation preserves its merge reservation instead of releasing it. + for batch in &fixture.batches[BUFFERED_BATCHES..] { + sorter.insert_batch(batch.clone()).await?; + assert!( + runtime.disk_manager.spilling_progress().active_files_count + <= fixture.batches.len() + ); + } + assert!(pool.state.lock().unwrap().denied > 0); + assert!(sorter.spilled_before()); + let stream = sorter.sort().await?; + drop(sorter); // The output stream must retain any transferred workspace. + let batches: Vec = stream.try_collect().await?; + assert!(batches.len() > 1); + let actual = concat_batches(&fixture.parent.schema(), &batches)?; + let expected = sort_batch(&fixture.parent, &fixture.ordering, None)?; + assert_eq!(actual, expected); + assert_released(&pool, &runtime).await; + Ok(()) +} + +#[tokio::test] +async fn test_spill_workspace_does_not_hide_insufficient_memory() -> Result<()> { + let fixture = fixture()?; + let capacity = + MERGE_BYTES + get_reserved_bytes_for_record_batch(&fixture.batches[0])? - 1; + let pool = AdjustablePool::new(capacity); + let (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?; + let error = sorter + .insert_batch(fixture.batches[0].clone()) + .await + .unwrap_err(); + assert!(matches!( + error.find_root(), + DataFusionError::ResourcesExhausted(_) + )); + drop(sorter); + assert_released(&pool, &runtime).await; + Ok(()) +} + +#[tokio::test] +async fn test_spill_workspace_cleanup_after_drop_or_error() -> Result<()> { + for inject_error in [false, true] { + let fixture = fixture()?; + let pool = AdjustablePool::new(fixture.capacity); + let (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?; + for batch in &fixture.batches[..BUFFERED_BATCHES + 2] { + sorter.insert_batch(batch.clone()).await?; + } + assert!(runtime.disk_manager.spilling_progress().active_files_count > 0); + if inject_error { + pool.set_limit(0); + let error = sorter + .insert_batch(fixture.batches[BUFFERED_BATCHES + 2].clone()) + .await + .unwrap_err(); + assert!(matches!( + error.find_root(), + DataFusionError::ResourcesExhausted(_) + )); + } + drop(sorter); + assert_released(&pool, &runtime).await; + } + Ok(()) +} + +#[tokio::test] +async fn test_spill_workspace_cleanup_after_output_is_dropped() -> Result<()> { + let fixture = fixture()?; + let pool = AdjustablePool::new(fixture.capacity); + let (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?; + for batch in &fixture.batches[..BUFFERED_BATCHES + 2] { + sorter.insert_batch(batch.clone()).await?; + } + assert!(runtime.disk_manager.spilling_progress().active_files_count > 0); + let mut stream = sorter.sort().await?; + drop(sorter); + assert!(stream.try_next().await?.is_some()); + drop(stream); + assert_released(&pool, &runtime).await; + Ok(()) +} + +fn aliased_batches() -> Result<(Vec, LexOrdering)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("alias_of_a", DataType::Int64, false), + ])); + let batches = [0, 128] + .into_iter() + .map(|start| { + let values: ArrayRef = + Arc::new(Int64Array::from_iter_values((start..start + 128).rev())); + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&values), values]) + .map_err(Into::into) + }) + .collect::>>()?; + let ordering = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + Ok((batches, ordering)) +} + +#[tokio::test] +async fn test_concat_sort_releases_unused_merge_workspace() -> Result<()> { + let (batches, ordering) = aliased_batches()?; + let schema = batches[0].schema(); + let input_bytes = batches + .iter() + .map(get_reserved_bytes_for_record_batch) + .collect::>>()? + .into_iter() + .sum::(); + let concatenated = concat_batches(&schema, &batches)?; + let concat_bytes = get_reserved_bytes_for_record_batch(&concatenated)?; + let expected = sort_batch(&concatenated, &ordering, None)?; + let headroom = 4096; + let pool_capacity = input_bytes + headroom; + // Concatenation breaks the aliases and needs more space than its inputs. + assert!(concat_bytes > input_bytes); + assert!(concat_bytes <= pool_capacity); + + let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_capacity)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + ordering, + 64, + headroom, + usize::MAX, // Force concatenation instead of a streaming merge. + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + runtime, + )?; + for batch in batches { + sorter.insert_batch(batch).await?; + } + assert_eq!(pool.reserved(), pool_capacity); + assert!(!sorter.spilled_before()); + let stream = sorter.sort().await?; + assert_eq!(pool.reserved(), concat_bytes); + drop(sorter); + let output: Vec = stream.try_collect().await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn test_disabled_spilling_does_not_reserve_merge_workspace() -> Result<()> { + let (batches, ordering) = aliased_batches()?; + let schema = batches[0].schema(); + let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; + let pool: Arc = Arc::new(GreedyMemoryPool::new(16 * 1024)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + ordering, + 64, + 10 * 1024 * 1024, // Must not be acquired when spilling is disabled. + 0, + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + runtime, + )?; + assert_eq!(pool.reserved(), 0); + for batch in batches { + sorter.insert_batch(batch).await?; + } + assert_eq!(sorter.merge_reservation_size(), 0); + let stream = sorter.sort().await?; + drop(sorter); + let output: Vec = stream.try_collect().await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +async fn check_chunked_string_view_workspace(during_spill: bool) -> Result<()> { + let options = ExecutionOptions::default(); + let rows = 4096; + let batch_size = 1024; + let batch_count = if during_spill { 3 } else { 2 }; + let schema = Arc::new(Schema::new(vec![Field::new( + "key", + DataType::Utf8View, + false, + )])); + let batches = (0..batch_count) + .map(|batch_id| { + let values: ArrayRef = + Arc::new(StringViewArray::from_iter_values((0..rows).rev().map( + |i| format!("row-{:08}-{}", batch_id * rows + i, "x".repeat(87)), + ))); + RecordBatch::try_new(Arc::clone(&schema), vec![values]) + }) + .collect::, _>>()?; + let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "key", 0, + )))] + .into(); + let input_bytes = batches[..2] + .iter() + .map(get_reserved_bytes_for_record_batch) + .collect::>>()? + .into_iter() + .sum::(); + assert!(input_bytes > options.sort_in_place_threshold_bytes); + + // Every chunk retains the long-string buffers. The ordinary sort + // reservation must grow even though all input was already reserved. + let sorted_bytes = sort_batch_chunked(&batches[0], &ordering, batch_size)? + .iter() + .map(get_record_batch_memory_size) + .sum::(); + assert!(sorted_bytes > get_reserved_bytes_for_record_batch(&batches[0])?); + + // Leave no unreserved capacity: the sort must be able to use its workspace. + let capacity = options.sort_spill_reservation_bytes + input_bytes; + 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(), + batch_size, + options.sort_spill_reservation_bytes, + options.sort_in_place_threshold_bytes, + options.spill_compression, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + )?; + for batch in &batches[..2] { + sorter.insert_batch(batch.clone()).await?; + } + assert_eq!(pool.reserved(), capacity); + assert!(!sorter.spilled_before()); + if during_spill { + sorter.insert_batch(batches[2].clone()).await?; + assert!(pool.state.lock().unwrap().denied > 0); + assert!(sorter.spilled_before()); + } + + let stream = sorter.sort().await?; + drop(sorter); + let output: Vec = stream.try_collect().await?; + let actual = concat_batches(&schema, &output)?; + let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; + assert_eq!(actual, expected); + assert_eq!(actual.num_rows(), batch_count * rows); + assert_released(&pool, &runtime).await; + Ok(()) +} + +#[tokio::test] +async fn test_chunked_string_view_final_sort_can_use_workspace() -> Result<()> { + check_chunked_string_view_workspace(false).await +} + +#[tokio::test] +async fn test_chunked_string_view_spill_can_use_workspace() -> Result<()> { + check_chunked_string_view_workspace(true).await +} + +#[tokio::test] +async fn test_chunked_dictionary_sort_can_use_workspace_with_defaults() -> Result<()> { + let options = ExecutionOptions::default(); + let rows = 32_768; + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new( + "payload", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + false, + ), + ])); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..64).map(|i| format!("{i:02}{}", "x".repeat(8190))), + )); + let mut batches = Vec::new(); + for start in [0, rows] { + let key: ArrayRef = Arc::new(Int64Array::from_iter_values( + (start..start + rows).rev().map(|i| i as i64), + )); + let payload: ArrayRef = Arc::new(DictionaryArray::::try_new( + Int8Array::from_iter_values((0..rows).map(|i| (i % 64) as i8)), + Arc::clone(&values), + )?); + batches.push(RecordBatch::try_new( + Arc::clone(&schema), + vec![key, payload], + )?); + } + let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "key", 0, + )))] + .into(); + let input_bytes = batches + .iter() + .map(get_reserved_bytes_for_record_batch) + .collect::>>()?; + let buffered_bytes = input_bytes.iter().sum::(); + assert!(buffered_bytes > options.sort_in_place_threshold_bytes); + + // Dictionary values remain shared, but each output batch must account for + // its backing buffers. Splitting the input therefore increases its charge. + let sorted_bytes = + sort_batch_chunked(&batches[0], &ordering, options.batch_size.get())? + .iter() + .map(get_record_batch_memory_size) + .sum::(); + assert!(sorted_bytes > input_bytes[0]); + + let capacity = options.sort_spill_reservation_bytes + buffered_bytes; + 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(), + options.batch_size.get(), + options.sort_spill_reservation_bytes, + options.sort_in_place_threshold_bytes, + options.spill_compression, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + )?; + for batch in &batches { + sorter.insert_batch(batch.clone()).await?; + } + assert_eq!(pool.reserved(), capacity); + assert!(!sorter.spilled_before()); + + let stream = sorter.sort().await?; + drop(sorter); + let output: Vec = stream.try_collect().await?; + let actual = concat_batches(&schema, &output)?; + let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; + // RecordBatch equality compares dictionary values, so this checks payloads + // as well as keys without depending on the dictionary's physical encoding. + assert_eq!(actual, expected); + assert_eq!(actual.num_rows(), rows * 2); + assert_released(&pool, &runtime).await; + Ok(()) +} + +#[tokio::test] +async fn test_overlapping_final_outputs_release_unused_workspace() -> Result<()> { + let options = ExecutionOptions::default(); + let rows = 80_000; + let capacity = 16 * 1024 * 1024; + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false)])); + let batches = (0..2) + .map(|batch_id| { + let values: ArrayRef = Arc::new(Int64Array::from_iter_values( + (batch_id * rows..(batch_id + 1) * rows) + .rev() + .map(|value| value as i64), + )); + RecordBatch::try_new(Arc::clone(&schema), vec![values]) + }) + .collect::, _>>()?; + let input_bytes = batches + .iter() + .map(get_reserved_bytes_for_record_batch) + .collect::>>()? + .into_iter() + .sum::(); + assert!(input_bytes > options.sort_in_place_threshold_bytes); + assert!(options.sort_spill_reservation_bytes + input_bytes < capacity); + assert!(capacity < 2 * options.sort_spill_reservation_bytes); + let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "key", 0, + )))] + .into(); + let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; + let pool = AdjustablePool::new(capacity); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?; + let make_sorter = |partition| { + ExternalSorter::new( + partition, + Arc::clone(&schema), + ordering.clone(), + options.batch_size.get(), + options.sort_spill_reservation_bytes, + options.sort_in_place_threshold_bytes, + options.spill_compression, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + ) + }; + + let mut first_sorter = make_sorter(0)?; + for batch in &batches { + first_sorter.insert_batch(batch.clone()).await?; + } + assert!(!first_sorter.spilled_before()); + let mut first_stream = first_sorter.sort().await?; + drop(first_sorter); + let first_batch = first_stream + .try_next() + .await? + .expect("first sort must produce output"); + assert_eq!(first_batch.num_rows(), options.batch_size.get()); + + // Keep the first merge alive and partially consumed. Its remaining data + // fits comfortably, but an idle spill-workspace floor would block another + // sorter from acquiring its own workspace and input reservations. + assert!( + pool.reserved() + options.sort_spill_reservation_bytes + input_bytes <= capacity, + "final output must release unused spill workspace for the next sorter" + ); + let mut second_sorter = make_sorter(1)?; + for batch in &batches { + second_sorter.insert_batch(batch.clone()).await?; + } + assert!(!second_sorter.spilled_before()); + let second_stream = second_sorter.sort().await?; + drop(second_sorter); + let second_output: Vec = second_stream.try_collect().await?; + + let mut first_output = vec![first_batch]; + first_output.extend(first_stream.try_collect::>().await?); + for output in [first_output, second_output] { + let actual = concat_batches(&schema, &output)?; + assert_eq!(actual, expected); + assert_eq!(actual.num_rows(), rows * 2); + } + assert_released(&pool, &runtime).await; + Ok(()) +} + +async fn check_final_spilled_merge_releases_unused_workspace( + intermediate: bool, +) -> Result<()> { + let options = ExecutionOptions::default(); + let workspace = options.sort_spill_reservation_bytes; + let rows = options.batch_size.get(); + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("payload", DataType::Utf8, false), + ])); + let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "key", 0, + )))] + .into(); + let make_batch = |rows: usize, start: usize, payload_bytes: usize| { + let key: ArrayRef = Arc::new(Int64Array::from_iter_values( + (start..start + rows).rev().map(|value| value as i64), + )); + let value = "x".repeat(payload_bytes); + let payload: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..rows).map(|_| value.as_str()), + )); + RecordBatch::try_new(Arc::clone(&schema), vec![key, payload]) + }; + let batch_count = if intermediate { 16 } else { 8 }; + let big = make_batch(rows, 0, 200)?; + let tail = make_batch(1, rows * batch_count, 200)?; + let big_bytes = get_reserved_bytes_for_record_batch(&big)?; + let tail_bytes = get_reserved_bytes_for_record_batch(&tail)?; + let capacity = workspace + 4 * big_bytes + tail_bytes / 2; + assert!(4 * big_bytes > options.sort_in_place_threshold_bytes); + assert!(tail_bytes > 1); + + // Generate the initial spill files with a fixed capacity and normal pressure. + let pool = AdjustablePool::new(capacity); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?; + let make_sorter = |partition| { + ExternalSorter::new( + partition, + Arc::clone(&schema), + ordering.clone(), + options.batch_size.get(), + options.sort_spill_reservation_bytes, + options.sort_in_place_threshold_bytes, + options.spill_compression, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + ) + }; + + let mut first_sorter = make_sorter(0)?; + let mut original = Vec::new(); + for index in 0..batch_count { + let batch = make_batch(rows, index * rows, 200)?; + assert_eq!(get_reserved_bytes_for_record_batch(&batch)?, big_bytes); + first_sorter.insert_batch(batch.clone()).await?; + original.push(batch); + } + assert_eq!(first_sorter.finished_spill_files.len(), batch_count / 4 - 1); + // Four batches fit, but this tiny insertion exceeds the remaining room + // by half its size and triggers another spill through normal pressure. + first_sorter.insert_batch(tail.clone()).await?; + original.push(tail.clone()); + assert_eq!(first_sorter.finished_spill_files.len(), batch_count / 4); + assert_eq!(first_sorter.in_mem_batches.len(), 1); + + let large_spill_bytes = first_sorter.finished_spill_files[0].max_record_batch_memory; + assert!( + first_sorter + .finished_spill_files + .iter() + .all(|spill| spill.max_record_batch_memory == large_spill_bytes) + ); + let tail_spill_bytes = sort_batch(&tail, &ordering, None)?.get_sliced_size()?; + let single_buffer_bytes = if intermediate { + 4 * large_spill_bytes // Two intermediate files remain for the final merge. + } else { + 4 * large_spill_bytes + 2 * tail_spill_bytes + }; + let available_during_merge = workspace + 128 * 1024; + assert!(tail_bytes < 128 * 1024); + assert!(single_buffer_bytes < workspace); + assert!(8 * large_spill_bytes > available_during_merge); + assert!(capacity > available_during_merge); + + // Leave room for the final input spill, but not two buffers for the first + // two files. With five files, two intermediate spill passes are also needed. + let transient = + MemoryConsumer::new("final merge contender").register(&runtime.memory_pool); + transient.try_grow(capacity - available_during_merge)?; + let spill_file_count = first_sorter.metrics.spill_metrics.spill_file_count.clone(); + let mut first_stream = first_sorter.sort().await?; + let spills_before = spill_file_count.value(); + drop(first_sorter); + if intermediate { + // Deny all fresh parent grants: intermediate passes must reuse the + // workspace acquired before the sort's share of memory decreased. + pool.set_limit(transient.size()); + } + let denied_before = pool.state.lock().unwrap().denied; + let first_batch = first_stream + .try_next() + .await? + .expect("the final spilled merge must produce output"); + assert_eq!(first_batch.num_rows(), rows); + assert!( + pool.state.lock().unwrap().denied > denied_before, + "the initial two-buffer reservation must fail before retrying" + ); + if intermediate { + assert!(spill_file_count.value() >= spills_before + 2); + } else { + assert_eq!(spill_file_count.value(), spills_before); + } + assert_eq!( + pool.reserved() - transient.size(), + single_buffer_bytes, + "the final disk merge must release its unused workspace" + ); + pool.set_limit(capacity); + drop(transient); + + // Leave the first final merge alive with most of its output still unread. + // The next sort's input fits alongside the single-buffer reservation, but + // not alongside the old workspace floor. Use a default-sized input again. + let next_batch = make_batch(rows, rows * 16, 400)?; + let next_bytes = get_reserved_bytes_for_record_batch(&next_batch)?; + assert!(single_buffer_bytes + workspace + next_bytes <= capacity); + assert!(2 * workspace + next_bytes > capacity); + let mut second_sorter = make_sorter(1)?; + second_sorter.insert_batch(next_batch.clone()).await?; + let second_stream = second_sorter.sort().await?; + drop(second_sorter); + let second_output: Vec = second_stream.try_collect().await?; + assert_eq!( + concat_batches(&schema, &second_output)?, + sort_batch(&next_batch, &ordering, None)? + ); + + let mut first_output = vec![first_batch]; + while let Some(batch) = first_stream.try_next().await? { + first_output.push(batch); + } + assert_eq!( + concat_batches(&schema, &first_output)?, + sort_batch(&concat_batches(&schema, &original)?, &ordering, None)? + ); + // EOF must release resources even while the exhausted outer stream lives. + assert_released(&pool, &runtime).await; + drop(first_stream); + Ok(()) +} + +#[tokio::test] +async fn test_final_spilled_merge_releases_unused_workspace() -> Result<()> { + check_final_spilled_merge_releases_unused_workspace(false).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_final_spilled_merge_releases_unused_workspace_multithreaded() -> Result<()> +{ + check_final_spilled_merge_releases_unused_workspace(false).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_final_spilled_merge_retains_workspace_between_passes() -> Result<()> { + check_final_spilled_merge_releases_unused_workspace(true).await +} + +#[tokio::test] +async fn test_chunked_sort_returns_live_workspace_loan_on_drop() -> Result<()> { + let options = ExecutionOptions::default(); + let rows = 4096; + let batch_size = 1024; + let schema = Arc::new(Schema::new(vec![Field::new( + "key", + DataType::Utf8View, + false, + )])); + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..rows) + .rev() + .map(|i| format!("row-{i:08}-{}", "x".repeat(87))), + )); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![values])?; + let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "key", 0, + )))] + .into(); + let input_bytes = get_reserved_bytes_for_record_batch(&batch)?; + let expected = sort_batch_chunked(&batch, &ordering, batch_size)?; + let sorted_bytes = expected + .iter() + .map(get_record_batch_memory_size) + .sum::(); + let first_bytes = get_record_batch_memory_size(&expected[0]); + assert!(sorted_bytes > input_bytes); + let initial_loan = sorted_bytes - input_bytes; + assert!(initial_loan <= options.sort_spill_reservation_bytes); + assert!( + initial_loan > first_bytes, + "cancellation must leave a positive workspace loan" + ); + let remaining_loan = initial_loan - first_bytes; + + let capacity = options.sort_spill_reservation_bytes + input_bytes; + 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, + batch_size, + options.sort_spill_reservation_bytes, + options.sort_in_place_threshold_bytes, + options.spill_compression, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + )?; + sorter.insert_batch(batch).await?; + assert_eq!(pool.reserved(), capacity); + + // Match the beginning of spill preparation, retaining its workspace but + // assigning it no merge cursors. Call sort_batch_stream directly: the usual + // one-input path would release unused workspace before this stream starts. + sorter.merge_reservation.free(); + let merge_pool = Arc::clone(&sorter.merge_pool); + let batch = sorter.in_mem_batches.pop().expect("one buffered input"); + assert!(sorter.in_mem_batches.is_empty()); + let reservation = sorter.reservation.take(); + let mut stream = sorter.sort_batch_stream(batch, reservation)?; + drop(sorter); + + let first_batch = stream + .try_next() + .await? + .expect("the sorted stream must produce a batch"); + assert_eq!(first_batch, expected[0]); + assert_eq!(first_batch.num_rows(), batch_size); + assert_eq!(pool.reserved(), capacity); + + // Check actual loan accounting rather than inferring a loan from the input + // type. Only the workspace not held by the remaining output is available. + let unused_workspace = merge_pool.borrow(usize::MAX); + assert_eq!( + unused_workspace.size(), + options.sort_spill_reservation_bytes - remaining_loan + ); + drop(unused_workspace); + drop(merge_pool); + assert_eq!(pool.reserved(), capacity); + + // The stream is the remaining owner of the positive loan. Drop it without + // consuming the remaining batches and require both reservations to retire. + drop(stream); + assert_released(&pool, &runtime).await; + Ok(()) +} diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 81adad8e9ec84..af7097282d1aa 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -19,6 +19,7 @@ //! This is an order-preserving merge. use crate::metrics::BaselineMetrics; +use crate::sorts::merge_memory_pool::MergeMemoryPool; use crate::sorts::multi_level_merge::MultiLevelMergeBuilder; use crate::sorts::{ merge::SortPreservingMergeStream, @@ -95,6 +96,7 @@ pub struct StreamingMergeBuilder<'a> { batch_size: Option, fetch: Option, reservation: Option, + merge_pool: Option>, enable_round_robin_tie_breaker: bool, } @@ -154,6 +156,12 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// Keep spill workspace until the final merge pass selects its buffer budget. + pub(super) fn with_merge_pool(mut self, pool: Arc) -> Self { + self.merge_pool = Some(pool); + self + } + /// See [SortPreservingMergeExec::with_round_robin_repartition] for more /// information. /// @@ -186,6 +194,7 @@ impl<'a> StreamingMergeBuilder<'a> { metrics, batch_size, reservation, + merge_pool, fetch, expressions, enable_round_robin_tie_breaker, @@ -226,6 +235,7 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, enable_round_robin_tie_breaker, ) + .with_merge_pool(merge_pool) .create_spillable_merge_stream()); } From 804fbca9c566795e5a7dfc76882bbc1a696c4c2d Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 3 Sep 2026 05:24:22 +0000 Subject: [PATCH 2/2] fix: retain single-batch spill workspace and expose merge pool Keep reserved workspace through the single-batch spill path and explain why concatenation still releases idle capacity before its reservation grows. Add reduced-limit and live-loan regressions through the real spill dispatch. Move MergeMemoryPool and WorkspaceLoan into datafusion-execution with public API documentation, a usage example, and the existing pool tests. --- .../src/memory_pool}/merge_memory_pool.rs | 128 ++++++++++++++---- datafusion/execution/src/memory_pool/mod.rs | 5 + datafusion/physical-plan/src/sorts/mod.rs | 1 - .../src/sorts/multi_level_merge.rs | 3 +- datafusion/physical-plan/src/sorts/sort.rs | 10 +- .../src/sorts/sort/spill_tests.rs | 103 +++++++++++++- .../src/sorts/streaming_merge.rs | 3 +- 7 files changed, 214 insertions(+), 39 deletions(-) rename datafusion/{physical-plan/src/sorts => execution/src/memory_pool}/merge_memory_pool.rs (73%) diff --git a/datafusion/physical-plan/src/sorts/merge_memory_pool.rs b/datafusion/execution/src/memory_pool/merge_memory_pool.rs similarity index 73% rename from datafusion/physical-plan/src/sorts/merge_memory_pool.rs rename to datafusion/execution/src/memory_pool/merge_memory_pool.rs index 4e8ae269ccb79..58459cc1ec0ba 100644 --- a/datafusion/physical-plan/src/sorts/merge_memory_pool.rs +++ b/datafusion/execution/src/memory_pool/merge_memory_pool.rs @@ -15,20 +15,73 @@ // specific language governing permissions and limitations // under the License. -//! Shares the sorter's reserved merge workspace without returning it to the -//! execution pool between sorting batches and the merge's cursor reservations. +//! Shares reserved workspace across merge reservations and temporary loans. use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use datafusion_common::{Result, resources_err}; -use datafusion_execution::memory_pool::{ - MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, -}; use parking_lot::Mutex; +use super::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation}; + +/// Shares reserved merge workspace across child reservations and temporary loans. +/// +/// This pool charges a single reservation to its parent [`MemoryPool`]. After +/// acquiring workspace through a child [`MemoryReservation`], call [`Self::retain`] +/// to keep that capacity available even when children release their reservations. +/// Children and [`WorkspaceLoan`]s share this capacity; only usage above the +/// existing parent reservation requires additional parent memory. Call +/// [`Self::release_unused`] when no further merge needs the idle workspace. +/// Live children and loans keep the pool and their reserved memory alive. +/// +/// All parent accounting and allocation policy uses the [`MemoryConsumer`] +/// supplied to [`Self::new`]. Child consumers are not registered with the parent, +/// so their names and [`MemoryConsumer::can_spill`] flags do not affect the +/// parent's accounting or fair shares. Consumers that need a separate parent +/// allocation policy should register directly with the parent. They can borrow +/// already-acquired workspace with [`Self::borrow`] and request any additional +/// capacity through their own reservation. +/// +/// # Example +/// +/// ``` +/// # use std::sync::Arc; +/// # use datafusion_common::Result; +/// # use datafusion_execution::memory_pool::{ +/// # GreedyMemoryPool, MemoryConsumer, MemoryPool, MergeMemoryPool, +/// # }; +/// # fn main() -> Result<()> { +/// let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); +/// let workspace = Arc::new(MergeMemoryPool::new( +/// Arc::clone(&parent), +/// MemoryConsumer::new("merge workspace"), +/// )); +/// let pool: Arc = Arc::clone(&workspace) as _; +/// let reservation = MemoryConsumer::new("merge children").register(&pool); +/// reservation.try_grow(60)?; +/// workspace.retain(60); +/// reservation.free(); +/// +/// // Another operator can use only the space outside the retained workspace. +/// let contender = MemoryConsumer::new("contender").register(&parent); +/// contender.try_grow(40)?; +/// assert_eq!(parent.reserved(), 100); +/// +/// // Children and loans reuse that workspace without another parent grant. +/// let cursor = reservation.new_empty(); +/// cursor.try_grow(20)?; +/// let loan = workspace.borrow(60); +/// assert_eq!(loan.size(), 40); +/// workspace.release_unused(); +/// assert_eq!(parent.reserved(), 100); // Both the cursor and loan remain charged. +/// drop(loan); +/// assert_eq!(parent.reserved(), 60); // Only the cursor and contender remain. +/// # Ok(()) +/// # } +/// ``` #[derive(Debug)] -pub(super) struct MergeMemoryPool { +pub struct MergeMemoryPool { parent: Arc, state: Mutex, } @@ -45,19 +98,30 @@ struct MergeMemoryState { retained: usize, } -/// Temporarily lends already-reserved workspace to per-batch sorting. Any -/// additional capacity must still be acquired by the ordinary sort consumer. -pub(super) struct WorkspaceLoan { +/// A temporary claim on workspace already reserved by a [`MergeMemoryPool`]. +/// +/// Created by [`MergeMemoryPool::borrow`]. The loan shares capacity with child +/// reservations and keeps the pool alive until it is dropped. Shrinking or +/// dropping the loan returns its bytes to the pool; the parent reservation is +/// reduced only when those bytes are no longer retained as workspace. +#[derive(Debug)] +pub struct WorkspaceLoan { pool: Arc, size: usize, } impl WorkspaceLoan { - pub(super) fn size(&self) -> usize { + /// Returns the number of bytes currently held by this loan. + pub fn size(&self) -> usize { self.size } - pub(super) fn shrink(&mut self, size: usize) { + /// Returns `size` bytes to the pool. + /// + /// # Panics + /// + /// Panics if `size` exceeds [`Self::size`]. + pub fn shrink(&mut self, size: usize) { self.size = self .size .checked_sub(size) @@ -75,7 +139,11 @@ impl Drop for WorkspaceLoan { } impl MergeMemoryPool { - pub(super) fn new(parent: Arc, consumer: MemoryConsumer) -> Self { + /// Creates a pool whose memory is charged to `consumer` in `parent`. + /// + /// Registers the consumer without reserving any memory. The consumer's + /// allocation policy applies to the combined usage of all children and loans. + pub fn new(parent: Arc, consumer: MemoryConsumer) -> Self { let reservation = consumer.register(&parent); Self { parent, @@ -87,17 +155,29 @@ impl MergeMemoryPool { } } - /// Retain workspace only after it has been successfully acquired. - pub(super) fn retain(&self, size: usize) { + /// Sets the amount of acquired workspace to retain when children release it. + /// + /// The workspace must first be acquired through a child reservation. This + /// method sets the retained floor without changing the parent reservation. + /// Use [`Self::release_unused`] to stop retaining idle workspace immediately. + /// + /// # Panics + /// + /// Panics if `size` exceeds this pool's current reservation in the parent, + /// as reported by its [`MemoryPool::reserved`] method. + pub fn retain(&self, size: usize) { let mut state = self.state.lock(); assert!(size <= state.reservation.size()); state.retained = size; } - /// Lend up to `size` unused bytes without acquiring more parent memory. - /// Count the loan alongside merge reservations so they cannot spend the - /// same credit. Dropping the loan returns any outstanding bytes. - pub(super) fn borrow(self: &Arc, size: usize) -> WorkspaceLoan { + /// Lends up to `size` unused bytes without acquiring more parent memory. + /// + /// Returns a smaller loan if less workspace is available, including an empty + /// loan if all capacity is in use. The loan is counted alongside child + /// reservations so they cannot spend the same credit. Dropping it returns + /// any outstanding bytes. + pub fn borrow(self: &Arc, size: usize) -> WorkspaceLoan { let mut state = self.state.lock(); let size = size.min(state.reservation.size() - state.used); state.used += size; @@ -107,9 +187,11 @@ impl MergeMemoryPool { } } - /// Stop retaining idle workspace when no further spill or intermediate merge - /// needs it. Live child reservations and loans remain charged. - pub(super) fn release_unused(&self) { + /// Stops retaining idle workspace and returns unused capacity to the parent. + /// + /// Live child reservations and loans remain charged. Their future releases + /// also return memory to the parent, unless [`Self::retain`] is called again. + pub fn release_unused(&self) { let mut state = self.state.lock(); state.retained = 0; state.reservation.resize(state.used); @@ -174,7 +256,7 @@ impl MemoryPool for MergeMemoryPool { #[cfg(test)] mod tests { use super::*; - use datafusion_execution::memory_pool::GreedyMemoryPool; + use crate::memory_pool::GreedyMemoryPool; fn reservation( parent: &Arc, @@ -354,7 +436,7 @@ mod tests { #[test] fn workspace_loan_preserves_ordinary_consumer_fair_share() -> Result<()> { - use datafusion_execution::memory_pool::{FairSpillPool, TrackConsumersPool}; + use crate::memory_pool::{FairSpillPool, TrackConsumersPool}; use std::num::NonZeroUsize; let tracked = Arc::new(TrackConsumersPool::new( diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs index 40a79d136b84e..7f8efe9cf8878 100644 --- a/datafusion/execution/src/memory_pool/mod.rs +++ b/datafusion/execution/src/memory_pool/mod.rs @@ -24,6 +24,7 @@ use std::fmt::Display; use std::hash::{Hash, Hasher}; use std::{cmp::Ordering, sync::Arc, sync::atomic}; +mod merge_memory_pool; mod peak_recording; mod pool; @@ -37,6 +38,7 @@ pub mod proxy { pub use datafusion_common::{ human_readable_count, human_readable_duration, human_readable_size, units, }; +pub use merge_memory_pool::{MergeMemoryPool, WorkspaceLoan}; pub use peak_recording::*; pub use pool::*; @@ -185,6 +187,9 @@ pub use pool::*; /// /// * [`TrackConsumersPool`]: Wraps another [`MemoryPool`] and tracks consumers, /// providing better error messages on the largest memory users. +/// +/// * [`MergeMemoryPool`]: Shares retained workspace across child reservations and +/// temporary loans, charging a single consumer in its parent pool. pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display { /// Return pool name fn name(&self) -> &str; diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs index 125fca1abbf51..4af2a3629fa9c 100644 --- a/datafusion/physical-plan/src/sorts/mod.rs +++ b/datafusion/physical-plan/src/sorts/mod.rs @@ -20,7 +20,6 @@ mod builder; mod cursor; mod merge; -mod merge_memory_pool; #[cfg(test)] mod merge_trace; mod multi_level_merge; diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 8ff845ef92652..b5aa5c4d54015 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -28,10 +28,9 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use datafusion_common::{Result, internal_err, resources_err}; -use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::memory_pool::{MemoryReservation, MergeMemoryPool}; use crate::sorts::builder::try_grow_reservation_to_at_least; -use crate::sorts::merge_memory_pool::MergeMemoryPool; use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 7726fb00c193b..31f5668a617ac 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -42,7 +42,6 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::IncrementalSortIterator; -use crate::sorts::merge_memory_pool::MergeMemoryPool; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::get_record_batch_memory_size; use crate::spill::in_progress_spill_file::InProgressSpillFile; @@ -68,7 +67,9 @@ use datafusion_common::{ unwrap_or_internal_err, }; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; +use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryPool, MemoryReservation, MergeMemoryPool, +}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; @@ -605,7 +606,6 @@ impl ExternalSorter { coalesce_runs: bool, ) -> Result { if self.in_mem_batches.is_empty() { - self.merge_pool.release_unused(); let empty_stream = Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); return Ok(self.observe_if_output(empty_stream, is_output_stream)); @@ -623,7 +623,6 @@ impl ExternalSorter { // consumer of the stream. if self.in_mem_batches.len() == 1 { - self.merge_pool.release_unused(); let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); let sorted_stream = self.sort_batch_stream(batch, reservation)?; @@ -632,6 +631,9 @@ impl ExternalSorter { // If less than sort_in_place_threshold_bytes, concatenate and sort in place if self.reservation.size() < self.sort_in_place_threshold_bytes { + // Concatenation can grow the ordinary sort reservation, which cannot + // borrow merge workspace. Return idle workspace to the execution pool + // so that growth can use it. self.merge_pool.release_unused(); // Concatenate memory batches together and sort let batch = concat_batches(&self.schema, &self.in_mem_batches)?; diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 57a41984cbfcd..f41ca4507c25a 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -522,6 +522,97 @@ async fn test_chunked_string_view_spill_can_use_workspace() -> Result<()> { check_chunked_string_view_workspace(true).await } +#[tokio::test] +async fn test_single_batch_spill_preserves_workspace_after_limit_decreases() -> Result<()> +{ + let rows = 4096; + let batch_size = 1024; + let schema = Arc::new(Schema::new(vec![Field::new( + "key", + DataType::Utf8View, + false, + )])); + let batches = [rows, rows - batch_size] + .into_iter() + .enumerate() + .map(|(batch_id, batch_rows)| { + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..batch_rows).rev().map(|i| { + format!("row-{:08}-{}", batch_id * rows + i, "x".repeat(87)) + }), + )); + RecordBatch::try_new(Arc::clone(&schema), vec![values]) + }) + .collect::, _>>()?; + let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "key", 0, + )))] + .into(); + let input_bytes = batches + .iter() + .map(get_reserved_bytes_for_record_batch) + .collect::>>()?; + let sorted_bytes = batches + .iter() + .map(|batch| { + Ok(sort_batch_chunked(batch, &ordering, batch_size)? + .iter() + .map(get_record_batch_memory_size) + .sum::()) + }) + .collect::>>()?; + assert!(sorted_bytes[0] > input_bytes[0]); + assert!(sorted_bytes[1] > input_bytes[1]); + let workspace = sorted_bytes[0] - input_bytes[0]; + assert!(sorted_bytes[1] - input_bytes[1] <= workspace); + let capacity = workspace + input_bytes[0]; + let reduced_limit = workspace + input_bytes[1]; + // Both spills need a workspace loan. The first sorted batch cannot be + // charged to the execution pool if its workspace is released and regranted. + assert!(reduced_limit < sorted_bytes[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(), + batch_size, + workspace, + 0, + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + )?; + sorter.insert_batch(batches[0].clone()).await?; + assert_eq!(pool.reserved(), capacity); + assert_eq!(sorter.in_mem_batches.len(), 1); + assert!(!sorter.spilled_before()); + + pool.set_limit(reduced_limit); + sorter.insert_batch(batches[1].clone()).await?; + assert_eq!(sorter.finished_spill_files.len(), 1); + assert_eq!(sorter.in_mem_batches.len(), 1); + assert_eq!(pool.reserved(), reduced_limit); + + // Finalizing the sort must spill its one remaining input using the same + // workspace, then merge both files within the reduced memory limit. + let spill_count = sorter.metrics.spill_metrics.spill_file_count.clone(); + let stream = sorter.sort().await?; + drop(sorter); + let output: Vec = stream.try_collect().await?; + assert_eq!(spill_count.value(), 2); + assert!(output.iter().all(|batch| batch.num_rows() <= batch_size)); + let actual = concat_batches(&schema, &output)?; + let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; + assert_eq!(actual, expected); + assert_eq!(actual.num_rows(), 2 * rows - batch_size); + assert_released(&pool, &runtime).await; + Ok(()) +} + #[tokio::test] async fn test_chunked_dictionary_sort_can_use_workspace_with_defaults() -> Result<()> { let options = ExecutionOptions::default(); @@ -866,7 +957,7 @@ async fn test_final_spilled_merge_retains_workspace_between_passes() -> Result<( } #[tokio::test] -async fn test_chunked_sort_returns_live_workspace_loan_on_drop() -> Result<()> { +async fn test_single_batch_spill_returns_live_workspace_loan_on_drop() -> Result<()> { let options = ExecutionOptions::default(); let rows = 4096; let batch_size = 1024; @@ -920,15 +1011,13 @@ async fn test_chunked_sort_returns_live_workspace_loan_on_drop() -> Result<()> { sorter.insert_batch(batch).await?; assert_eq!(pool.reserved(), capacity); - // Match the beginning of spill preparation, retaining its workspace but - // assigning it no merge cursors. Call sort_batch_stream directly: the usual - // one-input path would release unused workspace before this stream starts. + // Follow spill preparation through the one-input branch, retaining the + // workspace but assigning it no merge cursors. sorter.merge_reservation.free(); let merge_pool = Arc::clone(&sorter.merge_pool); - let batch = sorter.in_mem_batches.pop().expect("one buffered input"); + assert_eq!(sorter.in_mem_batches.len(), 1); + let mut stream = sorter.in_mem_sort_stream(false, false)?; assert!(sorter.in_mem_batches.is_empty()); - let reservation = sorter.reservation.take(); - let mut stream = sorter.sort_batch_stream(batch, reservation)?; drop(sorter); let first_batch = stream diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index af7097282d1aa..726ae48b1c79c 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -19,7 +19,6 @@ //! This is an order-preserving merge. use crate::metrics::BaselineMetrics; -use crate::sorts::merge_memory_pool::MergeMemoryPool; use crate::sorts::multi_level_merge::MultiLevelMergeBuilder; use crate::sorts::{ merge::SortPreservingMergeStream, @@ -32,7 +31,7 @@ use datafusion_common::human_readable_size; use datafusion_common::{Result, assert_or_internal_err, internal_err}; use datafusion_execution::SpillFile; use datafusion_execution::memory_pool::{ - MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, + MemoryConsumer, MemoryPool, MemoryReservation, MergeMemoryPool, UnboundedMemoryPool, }; use datafusion_physical_expr_common::sort_expr::LexOrdering; use std::sync::Arc;