diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index 761c492d8dde3..aa9ca0efd06fc 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -26,7 +26,7 @@ use arrow::array::UInt64Array; use arrow::row::{RowConverter, SortField}; use arrow::{array::StringArray, compute::SortOptions, record_batch::RecordBatch}; use arrow_schema::{DataType, Field, Schema}; -use datafusion::common::Result; +use datafusion::common::{DataFusionError, Result}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::expressions::PhysicalSortExpr; @@ -624,7 +624,7 @@ async fn run_sort_test_with_limited_memory( let assert_output_batch_size = args.assert_all_output_batches_roughly_match_batch_size_conf; - let metrics = run_test(args, sort_exec, result).await?; + let metrics = run_test(args, sort_exec, result, false).await?; assert_baseline_metrics_for_non_empty_output( &metrics, @@ -764,15 +764,19 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_ #[tokio::test] async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_sizes_of_record_batch_and_take_all_memory() -> Result<()> { + // Permanent non-spillable pressure must fail within the pool limit instead + // of letting replay overcommit memory after it has produced output. let record_batch_size = 8192; let pool_size = 2 * MB as usize; + let memory_pool = Arc::new(PeakRecordingPool::new(Arc::new(FairSpillPool::new( + pool_size, + )))); let task_ctx = { - let memory_pool = Arc::new(FairSpillPool::new(pool_size)); TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(record_batch_size)) .with_runtime(Arc::new( RuntimeEnvBuilder::new() - .with_memory_pool(memory_pool) + .with_memory_pool(Arc::clone(&memory_pool) as Arc) .build()?, )) }; @@ -793,6 +797,8 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_ }) .await?; + assert!(memory_pool.peak_reserved() <= pool_size); + assert_eq!(memory_pool.reserved(), 0); Ok(()) } @@ -910,18 +916,29 @@ async fn run_test_aggregate_with_high_cardinality( let result = aggregate_final.execute(0, Arc::clone(&args.task_ctx))?; - run_test(args, aggregate_final, result).await + // A non-spilling competitor that permanently consumes every free byte can + // prevent later replay growth. Require a bounded error, not overcommit. + let expect_memory_exhaustion = matches!( + &args.memory_behavior, + MemoryBehavior::TakeAllMemoryAtTheBeginning + ); + run_test(args, aggregate_final, result, expect_memory_exhaustion).await } async fn run_test( args: RunTestWithLimitedMemoryArgs, plan: Arc, result_stream: SendableRecordBatchStream, + expect_memory_exhaustion: bool, ) -> Result { let number_of_record_batches = args.number_of_record_batches; - consume_stream_and_simulate_other_running_memory_consumers(args, result_stream) - .await?; + consume_stream_and_simulate_other_running_memory_consumers( + args, + result_stream, + expect_memory_exhaustion, + ) + .await?; let metrics = plan.metrics().expect("must have metrics"); let spill_count = assert_spill_count_metric(true, plan); @@ -938,6 +955,7 @@ async fn run_test( async fn consume_stream_and_simulate_other_running_memory_consumers( args: RunTestWithLimitedMemoryArgs, mut result_stream: SendableRecordBatchStream, + expect_memory_exhaustion: bool, ) -> Result<()> { let mut number_of_rows = 0; let record_batch_size = args.task_ctx.session_config().batch_size() as u64; @@ -950,6 +968,25 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( let mut memory_took = false; while let Some(batch) = result_stream.next().await { + let batch = match batch { + Ok(batch) => batch, + Err(DataFusionError::ResourcesExhausted(_)) if expect_memory_exhaustion => { + // Do not accept an early replay error before the mock actually + // takes memory away from an operator that has produced rows. + assert!(number_of_rows > 0); + assert!(memory_took && memory_reservation.size() > 0); + assert!(memory_pool.reserved() <= args.pool_size); + drop(result_stream); + drop(memory_reservation); + assert_eq!(memory_pool.reserved(), 0); + let progress = + args.task_ctx.runtime_env().disk_manager.spilling_progress(); + assert_eq!(progress.current_bytes, 0); + assert_eq!(progress.active_files_count, 0); + return Ok(()); + } + Err(error) => return Err(error), + }; match args.memory_behavior { MemoryBehavior::AsIs => { // Do nothing @@ -958,6 +995,11 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( if !memory_took { memory_took = true; grow_memory_as_much_as_possible(10, &mut memory_reservation)?; + if expect_memory_exhaustion { + assert!(memory_reservation.size() > 0); + assert!(memory_pool.reserved() <= args.pool_size); + assert!(args.pool_size - memory_pool.reserved() < 10); + } } } MemoryBehavior::TakeAllMemoryAndReleaseEveryNthBatch(n) => { @@ -974,12 +1016,15 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( } } - let batch = batch?; number_of_rows += batch.num_rows(); index += 1; } + assert!( + !expect_memory_exhaustion, + "expected memory exhaustion after external pressure" + ); assert_eq!( number_of_rows, args.number_of_record_batches * record_batch_size as usize diff --git a/datafusion/execution/src/memory_pool/merge_memory_pool.rs b/datafusion/execution/src/memory_pool/merge_memory_pool.rs index 58459cc1ec0ba..d384257f76872 100644 --- a/datafusion/execution/src/memory_pool/merge_memory_pool.rs +++ b/datafusion/execution/src/memory_pool/merge_memory_pool.rs @@ -251,6 +251,11 @@ impl MemoryPool for MergeMemoryPool { fn memory_limit(&self) -> MemoryLimit { self.parent.memory_limit() } + + fn memory_limit_for(&self, _consumer: &MemoryConsumer) -> MemoryLimit { + let state = self.state.lock(); + self.parent.memory_limit_for(state.reservation.consumer()) + } } #[cfg(test)] diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs index 5df64ba47ba3d..6e5d70faf69c8 100644 --- a/datafusion/execution/src/memory_pool/mod.rs +++ b/datafusion/execution/src/memory_pool/mod.rs @@ -230,6 +230,18 @@ pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display { fn memory_limit(&self) -> MemoryLimit { MemoryLimit::Unknown } + + /// Return the current total allowance for a registered consumer, including + /// all of its sibling reservations, rather than its remaining free memory. + /// + /// This is an advisory snapshot: other consumers may register or reserve + /// memory after the call. Allocations must still use [`Self::try_grow`]. + /// The default is [`MemoryLimit::Unknown`] because the global pool limit + /// need not be available to every consumer. Transparent wrappers should + /// delegate this method to their inner pool. + fn memory_limit_for(&self, _consumer: &MemoryConsumer) -> MemoryLimit { + MemoryLimit::Unknown + } } impl dyn MemoryPool { diff --git a/datafusion/execution/src/memory_pool/peak_recording.rs b/datafusion/execution/src/memory_pool/peak_recording.rs index b407cc0eaf36b..1a2ac0fe342f0 100644 --- a/datafusion/execution/src/memory_pool/peak_recording.rs +++ b/datafusion/execution/src/memory_pool/peak_recording.rs @@ -214,6 +214,10 @@ impl MemoryPool for PeakRecordingPool { fn memory_limit(&self) -> MemoryLimit { self.inner.memory_limit() } + + fn memory_limit_for(&self, consumer: &MemoryConsumer) -> MemoryLimit { + self.inner.memory_limit_for(consumer) + } } #[cfg(test)] diff --git a/datafusion/execution/src/memory_pool/pool.rs b/datafusion/execution/src/memory_pool/pool.rs index d854cbd627cec..f5c9be8b77705 100644 --- a/datafusion/execution/src/memory_pool/pool.rs +++ b/datafusion/execution/src/memory_pool/pool.rs @@ -59,6 +59,10 @@ impl MemoryPool for UnboundedMemoryPool { fn memory_limit(&self) -> MemoryLimit { MemoryLimit::Infinite } + + fn memory_limit_for(&self, _consumer: &MemoryConsumer) -> MemoryLimit { + self.memory_limit() + } } impl Display for UnboundedMemoryPool { @@ -127,6 +131,10 @@ impl MemoryPool for GreedyMemoryPool { fn memory_limit(&self) -> MemoryLimit { MemoryLimit::Finite(self.pool_size) } + + fn memory_limit_for(&self, _consumer: &MemoryConsumer) -> MemoryLimit { + self.memory_limit() + } } impl Display for GreedyMemoryPool { @@ -163,6 +171,11 @@ impl Display for GreedyMemoryPool { /// └───────────────────────z──────────────────────z───────────────┘ /// ``` /// +/// Reservations created with [`MemoryReservation::new_empty`], +/// [`MemoryReservation::split`], or [`MemoryReservation::take`] share their +/// consumer's allowance. Registering a new consumer does not revoke existing +/// reservations, but further fallible growth remains limited by total pool capacity. +/// /// Unspillable memory is allocated in a first-come, first-serve fashion #[derive(Debug)] pub struct FairSpillPool { @@ -180,6 +193,12 @@ struct FairSpillPoolState { /// The total amount of memory reserved that can be spilled spillable: usize, + /// Total reservation across every sibling of each spillable consumer. + /// + /// `MemoryReservation::new_empty`, `split`, and `take` share a consumer + /// registration while maintaining separate reservation-size counters. + spillable_by_consumer: HashMap, + /// The total amount of memory reserved by consumers that cannot spill unspillable: usize, } @@ -193,6 +212,7 @@ impl FairSpillPool { state: Mutex::new(FairSpillPoolState { num_spill: 0, spillable: 0, + spillable_by_consumer: HashMap::default(), unspillable: 0, }), } @@ -206,7 +226,9 @@ impl MemoryPool for FairSpillPool { fn register(&self, consumer: &MemoryConsumer) { if consumer.can_spill { - self.state.lock().num_spill += 1; + let mut state = self.state.lock(); + state.num_spill += 1; + state.spillable_by_consumer.insert(consumer.id(), 0); } } @@ -214,13 +236,22 @@ impl MemoryPool for FairSpillPool { if consumer.can_spill { let mut state = self.state.lock(); state.num_spill = state.num_spill.checked_sub(1).unwrap(); + let released = state.spillable_by_consumer.remove(&consumer.id()); + debug_assert_eq!(released, Some(0)); } } fn grow(&self, reservation: &MemoryReservation, additional: usize) { let mut state = self.state.lock(); match reservation.registration.consumer.can_spill { - true => state.spillable += additional, + true => { + state.spillable += additional; + *state + .spillable_by_consumer + .get_mut(&reservation.consumer().id()) + .expect("spillable memory consumer must remain registered") += + additional; + } false => state.unspillable += additional, } } @@ -228,7 +259,13 @@ impl MemoryPool for FairSpillPool { fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { let mut state = self.state.lock(); match reservation.registration.consumer.can_spill { - true => state.spillable -= shrink, + true => { + state.spillable -= shrink; + *state + .spillable_by_consumer + .get_mut(&reservation.consumer().id()) + .expect("spillable memory consumer must remain registered") -= shrink; + } false => state.unspillable -= shrink, } } @@ -245,8 +282,16 @@ impl MemoryPool for FairSpillPool { let available = spill_available .checked_div(state.num_spill) .unwrap_or(spill_available); - - if reservation.size() + additional > available { + let consumer_used = state + .spillable_by_consumer + .get(&reservation.consumer().id()) + .copied() + .expect("spillable memory consumer must remain registered"); + + if consumer_used + .checked_add(additional) + .is_none_or(|requested| requested > available) + { return Err(insufficient_capacity_err( reservation, additional, @@ -254,12 +299,28 @@ impl MemoryPool for FairSpillPool { self, )); } + let remaining = self + .pool_size + .saturating_sub(state.unspillable.saturating_add(state.spillable)); + if additional > remaining { + return Err(insufficient_capacity_err( + reservation, + additional, + remaining, + self, + )); + } state.spillable += additional; + *state + .spillable_by_consumer + .get_mut(&reservation.consumer().id()) + .expect("spillable memory consumer must remain registered") += + additional; } false => { let available = self .pool_size - .saturating_sub(state.unspillable + state.spillable); + .saturating_sub(state.unspillable.saturating_add(state.spillable)); if available < additional { return Err(insufficient_capacity_err( @@ -283,6 +344,15 @@ impl MemoryPool for FairSpillPool { fn memory_limit(&self) -> MemoryLimit { MemoryLimit::Finite(self.pool_size) } + + fn memory_limit_for(&self, consumer: &MemoryConsumer) -> MemoryLimit { + if !consumer.can_spill() { + return self.memory_limit(); + } + let state = self.state.lock(); + let available = self.pool_size.saturating_sub(state.unspillable); + MemoryLimit::Finite(available.checked_div(state.num_spill).unwrap_or(available)) + } } impl Display for FairSpillPool { @@ -600,6 +670,10 @@ impl MemoryPool for TrackConsumersPool { fn memory_limit(&self) -> MemoryLimit { self.inner.memory_limit() } + + fn memory_limit_for(&self, consumer: &MemoryConsumer) -> MemoryLimit { + self.inner.memory_limit_for(consumer) + } } fn provide_top_memory_consumers_to_error_msg( @@ -694,6 +768,195 @@ mod tests { assert_snapshot!(err, @"Resources exhausted: Failed to allocate additional 30.0 B for s4 with 0.0 B already allocated for this reservation - 20.0 B remain available for the total memory pool: fair(pool_size: 100.0 B)"); } + #[test] + fn test_fair_sibling_reservations_share_one_consumer_limit() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let parent = MemoryConsumer::new("spilling operator") + .with_can_spill(true) + .register(&pool); + let first_partition = parent.new_empty(); + let second_partition = parent.new_empty(); + + first_partition.try_grow(60).unwrap(); + second_partition.try_grow(40).unwrap(); + assert_eq!(pool.reserved(), 100); + assert!(parent.try_grow(1).is_err()); + assert!(second_partition.try_grow(1).is_err()); + assert_eq!(pool.reserved(), 100); + + drop(first_partition); + second_partition.try_grow(60).unwrap(); + assert_eq!(pool.reserved(), 100); + drop(second_partition); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn test_fair_siblings_respect_consumer_shares_and_global_capacity() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let unspillable = MemoryConsumer::new("fixed").register(&pool); + unspillable.try_grow(20).unwrap(); + + let first = MemoryConsumer::new("same name") + .with_can_spill(true) + .register(&pool); + let second = MemoryConsumer::new("same name") + .with_can_spill(true) + .register(&pool); + let sibling = first.new_empty(); + + first.try_grow(25).unwrap(); + sibling.try_grow(15).unwrap(); + assert!(sibling.try_grow(1).is_err()); + second.try_grow(40).unwrap(); + assert_eq!(pool.reserved(), 100); + + let split = first.split(10); + assert_eq!(pool.reserved(), 100); + assert!(split.try_grow(1).is_err()); + drop(split); + second.try_grow(1).unwrap_err(); + sibling.try_grow(10).unwrap(); + assert_eq!(pool.reserved(), 100); + } + + #[test] + fn test_fair_take_retains_usage_until_last_sibling_drops() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let mut parent = MemoryConsumer::new("spilling operator") + .with_can_spill(true) + .register(&pool); + let other = MemoryConsumer::new("other") + .with_can_spill(true) + .register(&pool); + parent.try_grow(50).unwrap(); + let taken = parent.take(); + assert_eq!(parent.size(), 0); + assert_eq!(taken.size(), 50); + assert!(parent.try_grow(1).is_err()); + assert!(taken.try_grow(1).is_err()); + drop(parent); + assert_eq!(pool.reserved(), 50); + taken.shrink(10); + taken.try_grow(10).unwrap(); + assert!(other.try_grow(51).is_err()); + drop(taken); + other.try_grow(100).unwrap(); + assert_eq!(pool.reserved(), 100); + drop(other); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn test_fair_new_consumer_respects_existing_global_usage() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let first = MemoryConsumer::new("first") + .with_can_spill(true) + .register(&pool); + first.try_grow(100).unwrap(); + let second = MemoryConsumer::new("second") + .with_can_spill(true) + .register(&pool); + assert!(second.try_grow(1).is_err()); + assert_eq!(pool.reserved(), 100); + first.shrink(50); + second.try_grow(50).unwrap(); + assert_eq!(pool.reserved(), 100); + } + + #[test] + fn test_fair_infallible_growth_is_charged_to_all_siblings() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let first = MemoryConsumer::new("first") + .with_can_spill(true) + .register(&pool); + let sibling = first.new_empty(); + // Infallible growth remains permitted, including beyond the configured capacity. + first.grow(110); + assert_eq!(pool.reserved(), 110); + assert!(sibling.try_grow(1).is_err()); + first.shrink(20); + sibling.try_grow(10).unwrap(); + assert_eq!(pool.reserved(), 100); + drop(first); + drop(sibling); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn test_fair_oversized_sibling_growth_does_not_overflow() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let parent = MemoryConsumer::new("spilling operator") + .with_can_spill(true) + .register(&pool); + parent.try_grow(1).unwrap(); + let sibling = parent.new_empty(); + assert!(sibling.try_grow(usize::MAX).is_err()); + assert_eq!(pool.reserved(), 1); + } + + #[test] + fn test_fair_consumer_memory_limit() { + use crate::memory_pool::PeakRecordingPool; + + let top = NonZeroUsize::new(2).unwrap(); + let pools: [Arc; 6] = [ + Arc::new(FairSpillPool::new(100)), + Arc::new(TrackConsumersPool::new(FairSpillPool::new(100), top)), + Arc::new(TrackConsumersPool::new( + TrackConsumersPool::new(FairSpillPool::new(100), top), + top, + )), + Arc::new(PeakRecordingPool::new(Arc::new(FairSpillPool::new(100)))), + Arc::new(TrackConsumersPool::new( + PeakRecordingPool::new(Arc::new(FairSpillPool::new(100))), + top, + )), + Arc::new(PeakRecordingPool::new(Arc::new(TrackConsumersPool::new( + FairSpillPool::new(100), + top, + )))), + ]; + for pool in pools { + let check_limit = |reservation: &MemoryReservation, expected| { + let before = (pool.reserved(), reservation.size()); + assert!(matches!( + pool.memory_limit_for(reservation.consumer()), + MemoryLimit::Finite(actual) if actual == expected + )); + assert_eq!((pool.reserved(), reservation.size()), before); + }; + let fixed = MemoryConsumer::new("fixed").register(&pool); + fixed.grow(20); + let first = MemoryConsumer::new("first") + .with_can_spill(true) + .register(&pool); + check_limit(&first, 80); + + let peer = MemoryConsumer::new("peer") + .with_can_spill(true) + .register(&pool); + let sibling = first.new_empty(); + sibling.try_grow(30).unwrap(); + for reservation in [&first, &peer, &sibling] { + check_limit(reservation, 40); + } + check_limit(&fixed, 100); + assert_eq!(pool.reserved(), 50); + assert!(matches!(pool.memory_limit(), MemoryLimit::Finite(100))); + + drop(peer); + check_limit(&first, 80); + fixed.grow(100); + check_limit(&first, 0); + check_limit(&fixed, 100); + assert_eq!(pool.reserved(), 150); + assert!(matches!(pool.memory_limit(), MemoryLimit::Finite(100))); + drop((first, sibling, fixed)); + assert_eq!(pool.reserved(), 0); + } + } + #[test] fn test_tracked_consumers_pool() { let setting = make_settings(); diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 7ebf32c3edfc6..5262e3054cc75 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -34,7 +34,7 @@ use datafusion_common::{ DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, internal_err, }; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryLimit, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; @@ -338,9 +338,15 @@ impl FinalSpillContext { } = self; let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. + // Bound merge buffers while sharing the operator's reservation, so the + // concurrent replay table can use any memory the merge does not need. + let merge_memory_limit = match context + .memory_pool() + .memory_limit_for(reservation.consumer()) + { + MemoryLimit::Finite(limit) => Some(limit / 2), + MemoryLimit::Infinite | MemoryLimit::Unknown => None, + }; let merge_reservation = reservation.new_empty(); let merged = StreamingMergeBuilder::new() .with_schema(spill_schema) @@ -350,6 +356,7 @@ impl FinalSpillContext { .with_metrics(baseline_metrics.intermediate()) .with_batch_size(batch_size) .with_reservation(merge_reservation) + .with_spill_merge_memory_limit(merge_memory_limit) .build()?; let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( &final_agg, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 774c08535c8e5..c98df576ac949 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3461,6 +3461,270 @@ mod tests { )) } + // This high-cardinality memory test would create quadratic collision scratch + // space; small group-value tests cover forced hash-collision correctness. + #[cfg(not(feature = "force_hash_collisions"))] + #[rstest::rstest] + #[case::final_hash(AggregateMode::Final, false)] + #[case::single_hash(AggregateMode::Single, false)] + #[case::ordered_final(AggregateMode::Final, true)] + #[case::ordered_single(AggregateMode::Single, true)] + #[tokio::test] + async fn migrated_aggregate_spill_merge_leaves_memory_for_replay( + #[case] mode: AggregateMode, + #[case] ordered: bool, + ) -> Result<()> { + use arrow::array::{ListArray, StringArray}; + use arrow::buffer::OffsetBuffer; + + const BATCH_SIZE: usize = 8192; + const KEYS_PER_PREFIX: usize = 25 * BATCH_SIZE; + const MEMORY_LIMIT: usize = 2 * 1024 * 1024; + let schema = Arc::new(Schema::new(vec![ + Field::new("prefix", DataType::Int64, false), + Field::new("key", DataType::Int64, false), + Field::new("value", DataType::Utf8, false), + ])); + let group_by = PhysicalGroupBy::new_single(vec![ + (col("prefix", &schema)?, "prefix".to_string()), + (col("key", &schema)?, "key".to_string()), + ]); + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("values") + .build()?, + )]; + let input_schema = if mode == AggregateMode::Single { + Arc::clone(&schema) + } else { + Arc::new(create_schema( + &schema, + &group_by, + &aggregates, + AggregateMode::Partial, + )?) + }; + let mut batches = vec![]; + for prefix in 0..2 { + // Repeat every key across spill runs. Only the prefix is ordered; + // the second pass restarts the descending key sequence. + for value in 1..=2 { + for start in (0..KEYS_PER_PREFIX).step_by(BATCH_SIZE).rev() { + let values: ArrayRef = + Arc::new(StringArray::from(vec![ + format!("{value:08}"); + BATCH_SIZE + ])); + // Final takes singleton ARRAY_AGG states rather than raw strings. + let values = if mode == AggregateMode::Single { + values + } else { + let DataType::List(field) = input_schema.field(2).data_type() + else { + unreachable!("ARRAY_AGG state must be a list") + }; + Arc::new(ListArray::new( + Arc::clone(field), + OffsetBuffer::from_lengths(std::iter::repeat_n( + 1, BATCH_SIZE, + )), + values, + None, + )) as ArrayRef + }; + batches.push(RecordBatch::try_new( + Arc::clone(&input_schema), + vec![ + Arc::new(Int64Array::from(vec![prefix; BATCH_SIZE])), + Arc::new(Int64Array::from_iter_values( + (start..start + BATCH_SIZE).rev().map(|key| key as i64), + )), + values, + ], + )?); + } + } + } + + let mut input = + TestMemoryExec::try_new(&[batches], Arc::clone(&input_schema), None)?; + if ordered { + input = input.try_with_sort_information(vec![ + LexOrdering::new([PhysicalSortExpr::new_default(col( + "prefix", &schema, + )?)]) + .unwrap(), + ])?; + } + let aggregate = AggregateExec::try_new( + mode, + group_by, + aggregates, + vec![None], + Arc::new(input), + Arc::clone(&schema), + )?; + let context = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(BATCH_SIZE)) + .with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(MEMORY_LIMIT))) + .build_arc()?, + ), + ); + let stream = aggregate.execute_typed(0, &context)?; + match (mode, ordered, &stream) { + (AggregateMode::Final, false, StreamType::FinalHash(_)) + | (AggregateMode::Single, false, StreamType::SingleHash(_)) => {} + (AggregateMode::Final, true, StreamType::OrderedFinalAggregate(_)) + | (AggregateMode::Single, true, StreamType::OrderedSingleAggregate(_)) => { + assert_eq!( + aggregate.input_order_mode(), + &InputOrderMode::PartiallySorted(vec![0]) + ); + } + _ => panic!("unexpected stream for {mode:?}, ordered={ordered}"), + } + let result = collect(stream.into()) + .await + .unwrap_or_else(|error| panic!("{mode:?}, ordered={ordered}: {error}")); + let mut seen = HashSet::new(); + for batch in &result { + assert!( + batch + .columns() + .iter() + .all(|column| column.null_count() == 0) + ); + let columns = batch + .columns() + .iter() + .take(2) + .map(|column| column.as_any().downcast_ref::().unwrap()) + .collect::>(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + let prefix = columns[0].value(row); + let key = columns[1].value(row); + assert!((0..2).contains(&prefix)); + assert!((0..KEYS_PER_PREFIX as i64).contains(&key)); + let values = values.value(row); + assert_eq!(values.len(), 2); + assert_eq!(values.null_count(), 0); + let values = values.as_any().downcast_ref::().unwrap(); + let mut values = [values.value(0), values.value(1)]; + values.sort_unstable(); + assert_eq!(values, ["00000001", "00000002"]); + assert!(seen.insert((prefix, key)), "duplicate group"); + } + } + assert_eq!(seen.len(), 2 * KEYS_PER_PREFIX); + let metrics = aggregate.metrics().unwrap(); + assert!(metrics.spill_count().unwrap() > 1); + assert!(metrics.spilled_rows().unwrap() > 0); + assert!(metrics.spilled_bytes().unwrap() > 0); + assert_eq!(context.memory_pool().reserved(), 0); + let progress = context.runtime_env().disk_manager.spilling_progress(); + assert_eq!(progress.current_bytes, 0); + assert_eq!(progress.active_files_count, 0); + Ok(()) + } + + #[tokio::test] + async fn migrated_aggregate_spill_merge_allows_indivisible_rows() -> Result<()> { + use arrow::array::StringArray; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, PeakRecordingPool, + }; + + const KEY_BYTES: usize = 350_000; + const GROUPS: usize = 24; + const MEMORY_LIMIT: usize = 5 * 1024 * 1024 / 2; + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int64, false), + ])); + let mut batches = Vec::new(); + for _ in 0..2 { + for key in (0..GROUPS).rev() { + batches.push(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec![format!( + "{key:02}{}", + "x".repeat(KEY_BYTES - 2) + )])), + Arc::new(Int64Array::from(vec![1])), + ], + )?); + } + } + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)?; + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".into())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + )], + vec![None], + Arc::new(input), + schema, + )?; + let pool = Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new( + MEMORY_LIMIT, + )))); + let context = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(1)) + .with_runtime( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?, + ), + ); + + // Two one-row spill inputs need 1,400,064 bytes for merge buffers, more + // than half the pool. They cannot shrink, but merge plus replay fits. + let result = collect(aggregate.execute(0, Arc::clone(&context))?).await?; + let mut seen = HashSet::new(); + for batch in result { + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let sums = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + let key = keys.value(row); + assert_eq!(key.len(), KEY_BYTES); + let group = key[..2].parse::().unwrap(); + assert!(group < GROUPS && seen.insert(group)); + assert_eq!(sums.value(row), 2); + } + } + assert_eq!(seen.len(), GROUPS); + assert!(aggregate.metrics().unwrap().spill_count().unwrap() > 1); + assert!(pool.peak_reserved() <= MEMORY_LIMIT); + assert_eq!(pool.reserved(), 0); + let progress = context.runtime_env().disk_manager.spilling_progress(); + assert_eq!(progress.current_bytes, 0); + assert_eq!(progress.active_files_count, 0); + Ok(()) + } + async fn check_grouping_sets( input: Arc, spill: bool, @@ -7310,32 +7574,22 @@ mod tests { Field::new("c", DataType::Int64, false), ])); - let batches = vec![vec![ - RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int64Array::from(vec![2])), - Arc::new(Int64Array::from(vec![2])), - Arc::new(Int64Array::from(vec![1])), - ], - )?, - RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int64Array::from(vec![1])), - Arc::new(Int64Array::from(vec![1])), - Arc::new(Int64Array::from(vec![1])), - ], - )?, - RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int64Array::from(vec![0])), - Arc::new(Int64Array::from(vec![0])), - Arc::new(Int64Array::from(vec![1])), - ], - )?, - ]]; + let mut descending_batches = Vec::new(); + for ordered_group in (0_i64..3).rev() { + // Multiple groups sharing the ordered prefix must remain in memory + // until its boundary, which forces an actual aggregation spill. + for unordered_group in 1_i64..=16 { + descending_batches.push(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![ordered_group])), + Arc::new(Int64Array::from(vec![ordered_group])), + Arc::new(Int64Array::from(vec![unordered_group])), + ], + )?); + } + } + let batches = vec![descending_batches]; let scan = TestMemoryExec::try_new(&batches, Arc::clone(&schema), None)?; let scan = scan.try_with_sort_information(vec![ LexOrdering::new([PhysicalSortExpr::new( @@ -7367,8 +7621,11 @@ mod tests { Arc::clone(&schema), )?); - let task_ctx = new_migrated_spill_ctx(1, 600); + // The merge input and replay aggregate now share one allowance. Keep + // enough space for both while the additional groups still force spilling. + let task_ctx = new_migrated_spill_ctx(1, 1024); let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await?; + assert_eq!(task_ctx.memory_pool().reserved(), 0); assert_spill_count_metric(true, Arc::clone(&aggr)); let metrics = aggr.metrics().unwrap(); for phase in ["update", "state", "merge", "evaluate"] { @@ -7382,13 +7639,58 @@ mod tests { allow_duplicates! { assert_snapshot!(batches_to_string(&result), @r" - +---+---+--------+ - | b | c | SUM(c) | - +---+---+--------+ - | 2 | 1 | 1 | - | 1 | 1 | 1 | - | 0 | 1 | 1 | - +---+---+--------+ + +---+----+--------+ + | b | c | SUM(c) | + +---+----+--------+ + | 2 | 1 | 1 | + | 2 | 2 | 2 | + | 2 | 3 | 3 | + | 2 | 4 | 4 | + | 2 | 5 | 5 | + | 2 | 6 | 6 | + | 2 | 7 | 7 | + | 2 | 8 | 8 | + | 2 | 9 | 9 | + | 2 | 10 | 10 | + | 2 | 11 | 11 | + | 2 | 12 | 12 | + | 2 | 13 | 13 | + | 2 | 14 | 14 | + | 2 | 15 | 15 | + | 2 | 16 | 16 | + | 1 | 1 | 1 | + | 1 | 2 | 2 | + | 1 | 3 | 3 | + | 1 | 4 | 4 | + | 1 | 5 | 5 | + | 1 | 6 | 6 | + | 1 | 7 | 7 | + | 1 | 8 | 8 | + | 1 | 9 | 9 | + | 1 | 10 | 10 | + | 1 | 11 | 11 | + | 1 | 12 | 12 | + | 1 | 13 | 13 | + | 1 | 14 | 14 | + | 1 | 15 | 15 | + | 1 | 16 | 16 | + | 0 | 1 | 1 | + | 0 | 2 | 2 | + | 0 | 3 | 3 | + | 0 | 4 | 4 | + | 0 | 5 | 5 | + | 0 | 6 | 6 | + | 0 | 7 | 7 | + | 0 | 8 | 8 | + | 0 | 9 | 9 | + | 0 | 10 | 10 | + | 0 | 11 | 11 | + | 0 | 12 | 12 | + | 0 | 13 | 13 | + | 0 | 14 | 14 | + | 0 | 15 | 15 | + | 0 | 16 | 16 | + +---+----+--------+ "); } Ok(()) diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 8823dcb32f47c..8dc4833f98661 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -25,7 +25,7 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, internal_err}; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryLimit, MemoryReservation}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -231,9 +231,15 @@ impl OrderedFinalSpillContext { } = self; let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. + // Bound merge buffers while sharing the operator's reservation, so the + // concurrent replay table can use any memory the merge does not need. + let merge_memory_limit = match context + .memory_pool() + .memory_limit_for(reservation.consumer()) + { + MemoryLimit::Finite(limit) => Some(limit / 2), + MemoryLimit::Infinite | MemoryLimit::Unknown => None, + }; let merge_reservation = reservation.new_empty(); let merged = StreamingMergeBuilder::new() .with_schema(spill_schema) @@ -243,6 +249,7 @@ impl OrderedFinalSpillContext { .with_metrics(baseline_metrics.intermediate()) .with_batch_size(batch_size) .with_reservation(merge_reservation) + .with_spill_merge_memory_limit(merge_memory_limit) .build()?; let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( &agg, @@ -312,8 +319,8 @@ impl OrderedFinalAggregateStream { reason = "keeps replay metric reuse explicit" )] /// Builds the stream with the reservation of its logical aggregate operator. - /// Replay callers pass a sibling of the reservation used by the merge input, - /// keeping both components under one memory-consumer registration. + /// Replay callers share the memory-consumer registration with their merge + /// input, whose buffer budget leaves room for concurrent aggregation. pub(in crate::aggregates) fn new_with_input_and_metrics( agg: &AggregateExec, context: &Arc, diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 701dee4e8146d..d760166ead6d1 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -25,7 +25,7 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryLimit, MemoryReservation}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -291,9 +291,15 @@ impl OrderedSingleSpillContext { } = self; let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. + // Bound merge buffers while sharing the operator's reservation, so the + // concurrent replay table can use any memory the merge does not need. + let merge_memory_limit = match context + .memory_pool() + .memory_limit_for(reservation.consumer()) + { + MemoryLimit::Finite(limit) => Some(limit / 2), + MemoryLimit::Infinite | MemoryLimit::Unknown => None, + }; let merge_reservation = reservation.new_empty(); let merged = StreamingMergeBuilder::new() .with_schema(spill_schema) @@ -303,6 +309,7 @@ impl OrderedSingleSpillContext { .with_metrics(baseline_metrics.intermediate()) .with_batch_size(batch_size) .with_reservation(merge_reservation) + .with_spill_merge_memory_limit(merge_memory_limit) .build()?; let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( &final_agg, diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 40412385efbc7..59461d7f139e3 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -30,7 +30,7 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryLimit, MemoryReservation}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -290,9 +290,15 @@ impl SingleSpillContext { } = self; let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. + // Bound merge buffers while sharing the operator's reservation, so the + // concurrent replay table can use any memory the merge does not need. + let merge_memory_limit = match context + .memory_pool() + .memory_limit_for(reservation.consumer()) + { + MemoryLimit::Finite(limit) => Some(limit / 2), + MemoryLimit::Infinite | MemoryLimit::Unknown => None, + }; let merge_reservation = reservation.new_empty(); let merged = StreamingMergeBuilder::new() .with_schema(spill_schema) @@ -302,6 +308,7 @@ impl SingleSpillContext { .with_metrics(baseline_metrics.intermediate()) .with_batch_size(batch_size) .with_reservation(merge_reservation) + .with_spill_merge_memory_limit(merge_memory_limit) .build()?; let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( &final_agg, diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index b5aa5c4d54015..1b204efb39112 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -33,6 +33,8 @@ 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; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::gc_view_arrays; +use crate::spill::spill_manager::GetSlicedSize; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -155,6 +157,8 @@ pub(crate) struct MultiLevelMergeBuilder { reservation: MemoryReservation, /// Workspace retained across retries and intermediate spill passes. merge_pool: Option>, + /// Limit concurrent merge buffers while leaving unused quota to replay. + spill_merge_memory_limit: Option, fetch: Option, enable_round_robin_tie_breaker: bool, } @@ -194,6 +198,7 @@ impl MultiLevelMergeBuilder { batch_size, reservation, merge_pool: None, + spill_merge_memory_limit: None, enable_round_robin_tie_breaker, fetch, } @@ -204,6 +209,12 @@ impl MultiLevelMergeBuilder { self } + /// Bound fan-in, not temporary workspace used before the merge starts. + pub(super) fn with_spill_merge_memory_limit(mut self, limit: Option) -> Self { + self.spill_merge_memory_limit = limit; + self + } + pub(crate) fn create_spillable_merge_stream(self) -> SendableRecordBatchStream { Box::pin(RecordBatchStreamAdapter::new( Arc::clone(&self.schema), @@ -212,9 +223,10 @@ impl MultiLevelMergeBuilder { } async fn create_stream(mut self) -> Result { + let mut allow_minimum_over_cap = false; loop { let (mut stream, batch_size_limit) = - match self.merge_sorted_runs_within_mem_limit()? { + match self.merge_sorted_runs_within_mem_limit(allow_minimum_over_cap)? { MergeStep::Stream { stream, batch_size_limit, @@ -225,10 +237,20 @@ impl MultiLevelMergeBuilder { // size so its largest batch shrinks, lowering the per-stream // reservation, then retry. Makes the merge resilient to skewed // (very wide) rows. - self.split_spill_file_in_half(index).await?; + let retry_unsplittable = self.spill_merge_memory_limit.is_some() + && !allow_minimum_over_cap; + if !self + .split_spill_file_in_half(index, retry_unsplittable) + .await? + { + // A single row may exceed the replay fan-in cap while + // the minimum merge still fits the actual shared pool. + allow_minimum_over_cap = true; + } continue; } }; + allow_minimum_over_cap = false; // TODO - add a threshold for number of files to disk even if empty and reading from disk so // we can avoid the memory reservation @@ -278,7 +300,10 @@ impl MultiLevelMergeBuilder { /// This tries to create a stream that merges the most sorted streams and sorted spill files /// as possible within the memory limit. - fn merge_sorted_runs_within_mem_limit(&mut self) -> Result { + fn merge_sorted_runs_within_mem_limit( + &mut self, + allow_minimum_over_cap: bool, + ) -> Result { match (self.sorted_spill_files.len(), self.sorted_streams.len()) { // No data so empty batch (0, 0) => { @@ -352,6 +377,7 @@ impl MultiLevelMergeBuilder { // we must have at least 2 streams to merge minimum_number_of_required_streams, &mut memory_reservation, + allow_minimum_over_cap, )? { SpillFilesToMerge::Ready(sorted_spill_files, buffer_size) => { (sorted_spill_files, buffer_size) @@ -483,6 +509,7 @@ impl MultiLevelMergeBuilder { buffer_len: usize, minimum_number_of_required_streams: usize, reservation: &mut MemoryReservation, + allow_minimum_over_cap: bool, ) -> Result { assert_ne!(buffer_len, 0, "Buffer length must be greater than 0"); let mut number_of_spills_to_read_for_current_phase = 0; @@ -510,10 +537,24 @@ impl MultiLevelMergeBuilder { ) * buffer_len; total_needed += per_spill; - // For memory pools that are not shared this is good, for other - // this is not and there should be some upper limit to memory - // reservation so we won't starve the system. - match try_grow_reservation_to_at_least(reservation, total_needed) { + let exceeds_merge_limit = self + .spill_merge_memory_limit + .is_some_and(|limit| total_needed > limit); + // Only the minimum merge may exceed the cap after splitting fails + // to shrink a run. Keep read-ahead disabled and ask the real pool + // for every byte; later streams must still fit the normal cap. + let allow_over_cap = allow_minimum_over_cap + && buffer_len == 1 + && number_of_spills_to_read_for_current_phase + < minimum_number_of_required_streams; + let admission = if exceeds_merge_limit && !allow_over_cap { + resources_err!( + "Spill merge requires {total_needed} bytes, exceeding its fan-in memory limit" + ) + } else { + try_grow_reservation_to_at_least(reservation, total_needed) + }; + match admission { Ok(_) => { number_of_spills_to_read_for_current_phase += 1; } @@ -532,11 +573,17 @@ impl MultiLevelMergeBuilder { buffer_len - 1, minimum_number_of_required_streams, reservation, + allow_minimum_over_cap, ); } // buffer_len == 1 and we still can't seat the minimum of 2 streams. if number_of_spills_to_read_for_current_phase == 0 { + if exceeds_merge_limit { + // Replay has not started, so re-splitting may use the + // full pool to shrink a batch above the fan-in cap. + return Ok(SpillFilesToMerge::SplitThenRetry(0)); + } // We couldn't even reserve a single stream - one record batch // is larger than the whole merge budget. That's the lone-batch // case, not the 2-stream merge skew we rescue here - surface it. @@ -581,7 +628,14 @@ impl MultiLevelMergeBuilder { /// than one run is re-spilled), the shrunk run records its own smaller batch-size /// limit (tracked alongside the run in `sorted_spill_files`), so only merges that /// actually consume it pay the reduced batch size. - async fn split_spill_file_in_half(&mut self, index: usize) -> Result<()> { + /// + /// Returns whether the largest batch shrank. If `retry_unsplittable` is true, + /// restore an unchanged run for one minimum-merge admission attempt. + async fn split_spill_file_in_half( + &mut self, + index: usize, + retry_unsplittable: bool, + ) -> Result { log::debug!( "2 spilled streams could not be loaded into memory for merge \ (requires 2x of the largest batch from both), re-spilling the larger of the two with half \ @@ -609,6 +663,32 @@ impl MultiLevelMergeBuilder { reservation .try_grow(get_reserved_bytes_for_record_batch_size(old_max, old_max))?; + if self.spill_merge_memory_limit.is_some() { + // A maximum-sized singleton cannot shrink. Find it without writing + // another file: the original runs may already fill the disk quota. + // Use an unbuffered reader so no background read outlives this guard. + let mut source = self.spill_manager.read_spill_as_stream_unbuffered( + Arc::clone(&target.file), + Some(old_max), + )?; + while let Some(batch) = source.next().await { + let batch = batch?; + if batch.num_rows() == 1 + && gc_view_arrays(&batch)?.get_sliced_size()? >= old_max + { + if !retry_unsplittable { + return resources_err!( + "Cannot merge sorted runs: a single record batch of {old_max} bytes \ + exceeds the available merge memory and cannot be split further" + ); + } + self.sorted_spill_files.push((target, 1)); + self.sorted_spill_files.swap(index, last); + return Ok(false); + } + } + } + let source = self .spill_manager .read_spill_as_stream(target.file, Some(old_max))?; @@ -642,10 +722,11 @@ impl MultiLevelMergeBuilder { return internal_err!("re-spilling a skewed spill file produced no data"); }; - // If halving could not reduce the largest batch (e.g. a single row that is - // itself wider than the budget), there is nothing more we can do - surface - // the out-of-memory condition instead of looping forever. - if new_max >= old_max { + // If halving cannot reduce the largest batch, only a requested retry + // against the actual pool can make progress. The caller permits that + // retry once before surfacing this error. + let shrank = new_max < old_max; + if !shrank && !retry_unsplittable { return resources_err!( "Cannot merge sorted runs: a single record batch of {old_max} bytes \ exceeds the available merge memory and cannot be split further" @@ -656,7 +737,13 @@ impl MultiLevelMergeBuilder { // global batch size. Merges that don't touch this run keep the full batch // size. a merge that reads it caps its output at this limit so the merged run // can't rebuild a full-size batch and reintroduce the skew. - let new_batch_size_limit = (old_batch_size / 2).max(1); + let new_batch_size_limit = if shrank { + (old_batch_size / 2).max(1) + } else { + // The cap exception only covers indivisible input rows, not a + // larger output batch formed by concatenating those rows. + 1 + }; // Push the re-spilled (smaller) file and swap it back into `index`, undoing // the swap-to-back above so the order is preserved. @@ -670,7 +757,7 @@ impl MultiLevelMergeBuilder { let last = self.sorted_spill_files.len() - 1; self.sorted_spill_files.swap(index, last); - Ok(()) + Ok(shrank) } fn observe_output( @@ -900,13 +987,13 @@ mod tests { // 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()? + builder.merge_sorted_runs_within_mem_limit(false)? 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?; + builder.split_spill_file_in_half(index, false).await?; assert_eq!(parent.reserved(), capacity); assert!(contender.try_grow(1).is_err()); } @@ -1002,6 +1089,200 @@ mod tests { Ok(()) } + /// A run can exceed the fan-in cap while still fitting the pool used to + /// split it before replay. The resulting merge must honor the smaller cap. + #[tokio::test] + async fn spill_merge_memory_limit_splits_an_oversized_first_run() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + let rows: i64 = 4096; + let first = make_sorted_spill_file(&spill_manager, &schema, (0..rows).collect()); + let second = + make_sorted_spill_file(&spill_manager, &schema, (rows..2 * rows).collect()); + let merge_limit = first.max_record_batch_memory; + // Even one original run needs twice this cap. The full pool can hold + // it during re-splitting, but must not be used to raise merge fan-in. + let pool: Arc = Arc::new(GreedyMemoryPool::new(8 * merge_limit)); + let builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![first, second], + &pool, + rows as usize, + ) + .with_spill_merge_memory_limit(Some(merge_limit)); + let mut stream = builder.create_spillable_merge_stream(); + let mut batches = vec![]; + while let Some(batch) = stream.next().await { + batches.push(batch?); + assert!(pool.reserved() <= merge_limit); + } + let merged = concat_batches(&schema, &batches)?; + let values = merged.column(0).as_primitive::(); + assert_eq!(values.len(), (2 * rows) as usize); + for (expected, value) in values.values().iter().enumerate() { + assert_eq!(*value, expected as i64); + } + drop(stream); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + + #[tokio::test] + async fn spill_merge_memory_limit_allows_only_an_indivisible_minimum() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + let spills = (0..3) + .map(|value| make_sorted_spill_file(&spill_manager, &schema, vec![value])) + .collect::>(); + let batch_memory = spills[0].max_record_batch_memory; + let pool: Arc = Arc::new(GreedyMemoryPool::new(8 * batch_memory)); + let mut builder = + build_merge_builder(spill_manager, Arc::clone(&schema), spills, &pool, 8192) + .with_spill_merge_memory_limit(Some(3 * batch_memory)); + + assert!(!builder.split_spill_file_in_half(0, true).await?); + assert_eq!(builder.sorted_spill_files.len(), 3); + // Actual batches contain one row even though the nominal size is 8192. + assert_eq!(builder.sorted_spill_files[0].1, 1); + let mut reservation = builder.reservation.new_empty(); + let SpillFilesToMerge::Ready(spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(2, 2, &mut reservation, true)? + else { + panic!("minimum merge should fit the pool"); + }; + assert_eq!(buffer_len, 1); + assert_eq!(spills.len(), 2); + assert_eq!(builder.sorted_spill_files.len(), 1); + assert_eq!(reservation.size(), 4 * batch_memory); + builder.sorted_spill_files.splice(0..0, spills); + reservation.free(); + + let mut stream = builder.create_spillable_merge_stream(); + let mut values = Vec::new(); + while let Some(batch) = stream.next().await { + let batch = batch?; + assert_eq!(batch.num_rows(), 1); + values.push(batch.column(0).as_primitive::().value(0)); + } + assert_eq!(values, vec![0, 1, 2]); + drop(stream); + assert_eq!(pool.reserved(), 0); + assert_eq!(env.disk_manager.spilling_progress().current_bytes, 0); + assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0); + Ok(()) + } + + #[tokio::test] + async fn spill_merge_memory_limit_still_enforces_the_actual_pool() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + let first = make_sorted_spill_file(&spill_manager, &schema, vec![1]); + let second = make_sorted_spill_file(&spill_manager, &schema, vec![2]); + let batch_memory = first.max_record_batch_memory; + let pool: Arc = Arc::new(GreedyMemoryPool::new(3 * batch_memory)); + let builder = + build_merge_builder(spill_manager, schema, vec![first, second], &pool, 8192) + .with_spill_merge_memory_limit(Some(batch_memory)); + let mut stream = builder.create_spillable_merge_stream(); + let error = stream.next().await.unwrap().unwrap_err(); + assert!(error.to_string().contains("cannot be split further")); + drop(stream); + assert_eq!(pool.reserved(), 0); + assert_eq!(env.disk_manager.spilling_progress().current_bytes, 0); + assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0); + Ok(()) + } + + #[rstest::rstest] + #[case(1, false, DataType::Utf8)] + #[case(8192, true, DataType::Utf8)] + #[case(8192, true, DataType::Utf8View)] + #[tokio::test] + async fn indivisible_spill_merge_needs_no_extra_disk_space( + #[case] batch_size: usize, + #[case] mixed_batches: bool, + #[case] data_type: DataType, + ) -> Result<()> { + use arrow::array::{ArrayRef, StringArray, StringViewArray}; + + const KEY_BYTES: usize = 300_000; + const POOL_BYTES: usize = 2 * 1024 * 1024; + let schema = Arc::new(Schema::new(vec![Field::new("x", data_type, false)])); + let make_runs = |manager: &SpillManager| -> Result> { + (0..2) + .map(|run| { + let batches = (0..8).map(|batch| { + // A small two-row batch before the large singleton + // prevents treating the first batch as representative. + let keys = if mixed_batches && batch == 0 { + vec![format!("{run:04}"), format!("{:04}", run + 2)] + } else { + let key = run + 2 * (batch + usize::from(mixed_batches)); + vec![format!("{key:04}{}", "x".repeat(KEY_BYTES - 4))] + }; + let values: ArrayRef = match schema.field(0).data_type() { + DataType::Utf8 => Arc::new(StringArray::from(keys)), + DataType::Utf8View => Arc::new(StringViewArray::from(keys)), + _ => unreachable!(), + }; + RecordBatch::try_new(Arc::clone(&schema), vec![values]) + .map_err(Into::into) + }); + let (file, max_record_batch_memory) = manager + .spill_record_batch_iter_and_return_max_batch_memory( + batches, + "indivisible input run", + )? + .unwrap(); + Ok(SortedSpillFile { + file, + max_record_batch_memory, + }) + }) + .collect() + }; + + // Calibrate the quota to exactly the original IPC files, with no room + // for even a replacement header. Keep several batches in each run so + // read-ahead cannot retire the original before the unnecessary write. + let calibration = Arc::new(RuntimeEnv::default()); + let runs = make_runs(&build_spill_manager(&calibration, &schema))?; + let quota = runs.iter().map(|run| run.file.size().unwrap()).sum(); + drop(runs); + let env = RuntimeEnvBuilder::new() + .with_max_temp_directory_size(quota) + .build_arc()?; + let manager = build_spill_manager(&env, &schema); + let runs = make_runs(&manager)?; + assert_eq!(env.disk_manager.spilling_progress().current_bytes, quota); + assert!(4 * runs[0].max_record_batch_memory > POOL_BYTES / 2); + assert!(4 * runs[0].max_record_batch_memory <= POOL_BYTES); + let pool: Arc = Arc::new(GreedyMemoryPool::new(POOL_BYTES)); + let builder = build_merge_builder(manager, schema, runs, &pool, batch_size) + .with_spill_merge_memory_limit(Some(POOL_BYTES / 2)); + let mut stream = builder.create_spillable_merge_stream(); + let mut expected = 0; + while let Some(batch) = stream.next().await { + let batch = batch?; + assert_eq!(batch.num_rows(), 1); + let values = arrow::compute::cast(batch.column(0), &DataType::Utf8)?; + let key = values.as_string::().value(0); + assert_eq!(key[..4].parse::().unwrap(), expected); + expected += 1; + } + assert_eq!(expected, if mixed_batches { 18 } else { 16 }); + drop(stream); + assert_eq!(pool.reserved(), 0); + let progress = env.disk_manager.spilling_progress(); + assert_eq!(progress.current_bytes, 0); + assert_eq!(progress.active_files_count, 0); + Ok(()) + } + /// Tests the `new_max >= old_max` guard: a single-row run cannot be split /// any smaller, so re-spilling it does not shrink the largest batch and the /// rescue surfaces `ResourcesExhausted` rather than looping forever. @@ -1022,7 +1303,7 @@ mod tests { build_merge_builder(spill_manager, schema, vec![f0], &pool, 1024); let err = builder - .split_spill_file_in_half(0) + .split_spill_file_in_half(0, false) .await .expect_err("re-spilling a one-row run cannot shrink it"); assert!( @@ -1195,6 +1476,7 @@ mod tests { 1, 2, &mut merge_reservation, + false, )? { SpillFilesToMerge::Ready(spills, buffer_len) => (spills, buffer_len), SpillFilesToMerge::SplitThenRetry(index) => { diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 726ae48b1c79c..b36150da1b94b 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -96,6 +96,8 @@ pub struct StreamingMergeBuilder<'a> { fetch: Option, reservation: Option, merge_pool: Option>, + /// Optional fan-in budget when spill replay shares a consumer with its parent. + spill_merge_memory_limit: Option, enable_round_robin_tie_breaker: bool, } @@ -161,6 +163,13 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// Limit spill-merge fan-in without reserving unused memory from its parent. + /// Temporary re-splitting workspace still uses the original memory pool. + pub(crate) fn with_spill_merge_memory_limit(mut self, limit: Option) -> Self { + self.spill_merge_memory_limit = limit; + self + } + /// See [SortPreservingMergeExec::with_round_robin_repartition] for more /// information. /// @@ -194,6 +203,7 @@ impl<'a> StreamingMergeBuilder<'a> { batch_size, reservation, merge_pool, + spill_merge_memory_limit, fetch, expressions, enable_round_robin_tie_breaker, @@ -235,6 +245,7 @@ impl<'a> StreamingMergeBuilder<'a> { enable_round_robin_tie_breaker, ) .with_merge_pool(merge_pool) + .with_spill_merge_memory_limit(spill_merge_memory_limit) .create_spillable_merge_stream()); }