diff --git a/datafusion/execution/src/memory_pool/merge_memory_pool.rs b/datafusion/execution/src/memory_pool/merge_memory_pool.rs new file mode 100644 index 0000000000000..58459cc1ec0ba --- /dev/null +++ b/datafusion/execution/src/memory_pool/merge_memory_pool.rs @@ -0,0 +1,541 @@ +// 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 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 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 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, +} + +/// 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 { + /// Returns the number of bytes currently held by this loan. + pub fn size(&self) -> usize { + self.size + } + + /// 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) + .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 { + /// 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, + state: Mutex::new(MergeMemoryState { + reservation, + used: 0, + retained: 0, + }), + } + } + + /// 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; + } + + /// 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; + WorkspaceLoan { + pool: Arc::clone(self), + size, + } + } + + /// 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); + } + + 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 crate::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 crate::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/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/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 3ec52cc70c0a9..b5aa5c4d54015 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -28,7 +28,7 @@ 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::sort::get_reserved_bytes_for_record_batch_size; @@ -153,6 +153,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 +193,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 +241,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 +858,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..31f5668a617ac 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -67,7 +67,9 @@ 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, MergeMemoryPool, +}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; @@ -76,6 +78,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 +266,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 +295,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 +320,7 @@ impl ExternalSorter { reservation, spill_manager, merge_reservation, + merge_pool, runtime, batch_size, sort_spill_reservation_bytes, @@ -369,12 +382,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 +485,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( @@ -619,6 +631,10 @@ 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)?; self.in_mem_batches.clear(); @@ -726,8 +742,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 +758,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 +766,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 +821,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..f41ca4507c25a --- /dev/null +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -0,0 +1,1047 @@ +// 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_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(); + 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_single_batch_spill_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); + + // 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); + 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()); + 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..726ae48b1c79c 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -31,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; @@ -95,6 +95,7 @@ pub struct StreamingMergeBuilder<'a> { batch_size: Option, fetch: Option, reservation: Option, + merge_pool: Option>, enable_round_robin_tie_breaker: bool, } @@ -154,6 +155,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 +193,7 @@ impl<'a> StreamingMergeBuilder<'a> { metrics, batch_size, reservation, + merge_pool, fetch, expressions, enable_round_robin_tie_breaker, @@ -226,6 +234,7 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, enable_round_robin_tie_breaker, ) + .with_merge_pool(merge_pool) .create_spillable_merge_stream()); }