diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 25f8e437407a1..3e70e3078551e 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -134,6 +134,11 @@ name = "union_schema" harness = false name = "spill_io" +[[bench]] +harness = false +name = "sort_spill" +required-features = ["test_utils"] + [[bench]] harness = false name = "sort_preserving_merge" diff --git a/datafusion/physical-plan/benches/sort_spill.rs b/datafusion/physical-plan/benches/sort_spill.rs new file mode 100644 index 0000000000000..cc3532c11653d --- /dev/null +++ b/datafusion/physical-plan/benches/sort_spill.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Single-partition external sort with no competing pool consumers. +//! +//! Measures spill-merge overhead for 60 batches of 8192 rows, each containing +//! an Int64 key and a 120-byte Utf8 payload. Pool and workspace sizes use decimal MB. +//! Input generation, plan setup, output validation, and metric reporting are untimed. +//! Spill counts include intermediate rewrites; they are not merge-pass counts. +//! +//! Run the same benchmark source on both revisions, on the same machine: +//! `cargo bench -p datafusion-physical-plan --bench sort_spill --features test_utils -- --save-baseline before` +//! `cargo bench -p datafusion-physical-plan --bench sort_spill --features test_utils -- --baseline before` +//! Preserve Criterion's target directory between runs. For a headroom-on/off A/B, +//! use otherwise identical revisions differing only in the sort's headroom opt-in. +//! There is no public policy toggle; workspace=0 does not turn headroom off. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{Array, Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::Result; +use datafusion_common::config::SpillCompression; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::test::TestMemoryExec; +use futures::TryStreamExt; +use rand::{SeedableRng, seq::SliceRandom}; + +const BATCH_SIZE: usize = 8192; +const BATCHES: usize = 60; +const ROWS: usize = BATCHES * BATCH_SIZE; +const PAYLOAD: &str = "012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789"; + +fn input() -> Vec { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("payload", DataType::Utf8, false), + ])); + let mut keys: Vec = (0..ROWS as i64).collect(); + keys.shuffle(&mut rand::rngs::StdRng::seed_from_u64(42)); + keys.chunks(BATCH_SIZE) + .map(|chunk| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(chunk.to_vec())), + Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + PAYLOAD, + chunk.len(), + ))), + ], + ) + .unwrap() + }) + .collect() +} + +fn setup( + input: &[RecordBatch], + workspace_mb: usize, + pool_mb: usize, +) -> ( + Arc, + Arc, + Arc, +) { + let schema = input[0].schema(); + let source = + TestMemoryExec::try_new_exec(&[input.to_vec()], Arc::clone(&schema), None) + .unwrap(); + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(col("key", &schema).unwrap())]) + .unwrap(); + let plan: Arc = Arc::new(SortExec::new(ordering, source)); + let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_mb * 1_000_000)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc() + .unwrap(); + let config = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(BATCH_SIZE) + .with_sort_spill_reservation_bytes(workspace_mb * 1_000_000) + .with_spill_compression(SpillCompression::Uncompressed); + let context = Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime(runtime), + ); + (plan, context, pool) +} + +async fn run( + plan: Arc, + context: Arc, + validate: bool, +) -> Result { + let mut stream = plan.execute(0, context)?; + let mut count = 0; + while let Some(batch) = stream.try_next().await? { + if validate { + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let payload = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(keys.null_count(), 0); + assert_eq!(payload.null_count(), 0); + for row in 0..batch.num_rows() { + // The shuffled keys are exactly the integers 0..ROWS. + assert_eq!(keys.value(row), (count + row) as i64); + assert_eq!(payload.value(row), PAYLOAD); + } + } + count += batch.num_rows(); + } + Ok(count) +} + +fn benchmark(c: &mut Criterion) { + assert_eq!(PAYLOAD.len(), 120); + let input = input(); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let mut group = c.benchmark_group("sort_spill"); + group.sample_size(10); + for (workspace_mb, pool_mb) in [(1, 16), (1, 32), (10, 32), (10, 24)] { + let id = format!("workspace{workspace_mb}_pool{pool_mb}"); + group.bench_function(BenchmarkId::new("mb", &id), |b| { + let (plan, context, pool) = setup(&input, workspace_mb, pool_mb); + let count = runtime + .block_on(run(Arc::clone(&plan), context, true)) + .unwrap(); + assert_eq!(count, ROWS); + let metrics = plan.metrics().unwrap(); + let spills = metrics.spill_count().unwrap(); + assert!(spills > 0, "{id}: expected spills"); + eprintln!( + "{id}: spill_count={spills}, spilled_rows={}, spilled_bytes={}", + metrics.spilled_rows().unwrap(), + metrics.spilled_bytes().unwrap() + ); + drop(plan); + assert_eq!(pool.reserved(), 0, "{id}: memory reservation leaked"); + b.iter_batched( + || setup(&input, workspace_mb, pool_mb), + |(plan, context, _pool)| { + let count = runtime.block_on(run(plan, context, false)).unwrap(); + assert_eq!(count, ROWS); + black_box(count); + }, + BatchSize::PerIteration, + ); + }); + } + group.finish(); +} + +criterion_group!(benches, benchmark); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 424d1faffcf33..a2d1b30ee05f7 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -15,7 +15,7 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use datafusion_common::{Result, internal_err, resources_err}; -use datafusion_execution::memory_pool::{MemoryReservation, MergeMemoryPool}; +use datafusion_execution::memory_pool::{MemoryPool, MemoryReservation, MergeMemoryPool}; use crate::sorts::builder::try_grow_reservation_to_at_least; use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; @@ -28,6 +28,24 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::TryStreamExt; use futures::{Stream, StreamExt}; +/// Leave memory for others sharing the execution pool when a merge seats +/// more than the minimum two runs. +#[derive(Debug, Clone)] +pub(crate) struct FanInHeadroom { + /// The execution pool: the parent of the merge pool when there is one. + pub(crate) pool: Arc, + /// Other partitions share `pool`, even if they hold nothing right now. + pub(crate) siblings: bool, +} + +impl FanInHeadroom { + /// Whether another consumer may need the pool: sibling partitions exist, + /// or the pool holds more than the `owned` bytes this merge holds in it. + fn is_contended(&self, owned: usize) -> bool { + self.siblings || self.pool.reserved() > owned + } +} + /// Merges a stream of sorted cursors and record batches into a single sorted stream /// /// This is a wrapper around [`SortPreservingMergeStream`](crate::sorts::merge::SortPreservingMergeStream) @@ -146,6 +164,11 @@ pub(crate) struct MultiLevelMergeBuilder { merge_pool: Option>, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, + /// When set and the pool is shared, seat runs beyond the minimum two only + /// if the pool can admit their buffers plus as much again as they exceed + /// the retained workspace, so a merge does not take the memory that other + /// partitions sharing the pool need. + fan_in_headroom: Option, /// Avoid merging more runs than needed before the final replay pass. size_intermediate_merges: bool, fetch: Option, @@ -188,6 +211,7 @@ impl MultiLevelMergeBuilder { reservation, merge_pool: None, reserve_replay_headroom: false, + fan_in_headroom: None, size_intermediate_merges: false, enable_round_robin_tie_breaker, fetch, @@ -206,6 +230,17 @@ impl MultiLevelMergeBuilder { self } + /// Leave headroom when seating runs beyond the minimum two while the pool + /// is shared. The minimum merge is admitted against the pool alone, as + /// without headroom. + pub(super) fn with_fan_in_headroom( + mut self, + headroom: Option, + ) -> Self { + self.fan_in_headroom = headroom; + self + } + pub(super) fn with_intermediate_merge_sizing( mut self, min_spill_batch_rows: Option, @@ -573,6 +608,20 @@ impl MultiLevelMergeBuilder { // allocation, preventing starvation under memory pressure. let mut total_needed: usize = 0; let mut accepted_memory: usize = 0; + // What this merge already holds from the pool: its retained workspace + // when it has one. Seating runs within it takes nothing from others. + let owned = self + .merge_pool + .as_ref() + .map_or_else(|| reservation.size(), |pool| pool.reserved()); + // Headroom is for others sharing the pool. With no sibling partitions + // and nobody else holding memory in it, seat every run that fits, as + // without headroom, rather than rewrite runs in an extra merge pass + // for no one. + let fan_in_contended = self + .fan_in_headroom + .as_ref() + .is_some_and(|headroom| headroom.is_contended(owned)); for (spill, _) in &self.sorted_spill_files { if number_of_spills_to_read_for_current_phase >= max_spill_files @@ -598,7 +647,20 @@ impl MultiLevelMergeBuilder { && number_of_spills_to_read_for_current_phase < minimum_number_of_required_streams; let check_headroom = self.reserve_replay_headroom && !skip_headroom; - let admission = if check_headroom { + let check_fan_in_headroom = !check_headroom + && fan_in_contended + && number_of_spills_to_read_for_current_phase + >= minimum_number_of_required_streams; + let admission = if check_fan_in_headroom { + // Ask the pool for what this merge takes beyond its own + // workspace, plus as much again, then return the spare bytes. + match total_needed.checked_add(total_needed.saturating_sub(owned)) { + Some(with_headroom) => { + try_grow_reservation_to_at_least(reservation, with_headroom) + } + None => resources_err!("Spill merge headroom exceeds usize::MAX"), + } + } else if check_headroom { // Ask the pool for merge buffers plus equal replay space, then // return the spare bytes before exposing the merge stream. match total_needed.checked_mul(2) { @@ -666,7 +728,7 @@ impl MultiLevelMergeBuilder { } } - if self.reserve_replay_headroom { + if self.reserve_replay_headroom || fan_in_contended { // `total_needed` may include a rejected candidate. Keep only the // buffers that were admitted, releasing temporary replay headroom. reservation.shrink(reservation.size() - accepted_memory); @@ -966,7 +1028,7 @@ mod tests { use arrow::compute::concat_batches; use arrow::datatypes::{DataType, Field, Int64Type, Schema}; use datafusion_execution::memory_pool::{ - GreedyMemoryPool, MemoryConsumer, MemoryPool, + GreedyMemoryPool, MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion_physical_expr::expressions::{Column, col}; @@ -1463,6 +1525,187 @@ mod tests { Ok(()) } + /// Registers a consumer holding `size` bytes of `pool`, so a merge sharing + /// it sees another user and leaves headroom. + fn hold(pool: &Arc, size: usize) -> Result { + let other = MemoryConsumer::new("other consumer").register(pool); + other.try_grow(size)?; + Ok(other) + } + + /// Headroom for others sharing `pool`, with or without sibling partitions. + fn headroom(pool: &Arc, siblings: bool) -> Option { + Some(FanInHeadroom { + pool: Arc::clone(pool), + siblings, + }) + } + + #[rstest::rstest] + #[case::without_headroom(false, true, false, 4, 8)] + #[case::headroom_without_others(true, false, false, 4, 8)] + #[case::headroom_with_others(true, true, false, 3, 6)] + #[case::headroom_with_idle_siblings(true, false, true, 3, 6)] + fn fan_in_headroom_limits_runs_beyond_the_minimum( + #[case] with_headroom: bool, + #[case] others: bool, + #[case] siblings: bool, + #[case] seated: usize, + #[case] retained_batches: usize, + ) -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + let spills = (0..4) + .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(13 * batch_memory)); + let other = others.then(|| hold(&pool, 1)).transpose()?; + let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_fan_in_headroom( + with_headroom.then(|| headroom(&pool, siblings).unwrap()), + ); + let mut reservation = builder.reservation.new_empty(); + let SpillFilesToMerge::Ready(spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)? + else { + panic!("the minimum merge should fit the pool"); + }; + assert_eq!(buffer_len, 1); + // Each run takes two batches. With headroom, the third run needs 12 + // batches and is seated; the fourth needs 16 and is not. When nobody + // else holds memory and there are no sibling partitions, all four are + // seated, as without headroom. + assert_eq!(spills.len(), seated); + assert_eq!(reservation.size(), retained_batches * batch_memory); + drop((reservation, spills, builder, other)); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + + /// The minimum two runs are admitted against the pool alone, even when + /// their buffers plus headroom would not fit. + #[test] + fn fan_in_headroom_admits_the_minimum_merge_without_headroom() -> 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(5 * batch_memory)); + let other = hold(&pool, 1)?; + let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_fan_in_headroom(headroom(&pool, false)); + let mut reservation = builder.reservation.new_empty(); + let SpillFilesToMerge::Ready(spills, _) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)? + else { + panic!("the minimum merge should fit the pool"); + }; + // Two runs take four batches; with headroom they would take eight. The + // third would take twelve. + assert_eq!(spills.len(), 2); + assert_eq!(reservation.size(), 4 * batch_memory); + drop((reservation, spills, builder, other)); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + + /// With a retained workspace `W`, seating the third run, for buffers `T` + /// in all, asks the parent for `2 * (T - W)` more, and nothing when + /// `T <= W`. Afterwards the child holds `T` and the parent is charged + /// `max(W, T)`. + #[rstest::rstest] + #[case::runs_within_workspace(6, 0, 0, 3)] + #[case::runs_below_workspace(8, 0, 0, 3)] + #[case::exactly_the_headroom(2, 8, 0, 3)] + #[case::just_below_the_headroom(2, 8, 1, 2)] + fn fan_in_headroom_counts_the_retained_workspace( + #[case] workspace_batches: usize, + #[case] free_batches: usize, + #[case] short_by_bytes: usize, + #[case] seated: usize, + ) -> 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 per_run = + get_reserved_bytes_for_record_batch_size(batch_memory, batch_memory); + assert_eq!(per_run, 2 * batch_memory); + let workspace = workspace_batches * batch_memory; + let free = free_batches * batch_memory - short_by_bytes; + let capacity = 32 * batch_memory; + let parent: Arc = Arc::new(GreedyMemoryPool::new(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 mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_merge_pool(Some(Arc::clone(&merge_pool))) + .with_fan_in_headroom(headroom(&parent, false)); + builder.reservation.try_grow(workspace)?; + merge_pool.retain(workspace); + builder.reservation.free(); + assert_eq!(merge_pool.reserved(), workspace); + let other = hold(&parent, capacity - workspace - free)?; + + let mut reservation = builder.reservation.new_empty(); + let SpillFilesToMerge::Ready(spills, _) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)? + else { + panic!("the minimum merge should fit the pool"); + }; + assert_eq!(spills.len(), seated); + let total = seated * per_run; + assert_eq!(reservation.size(), total); + assert_eq!(merge_pool.reserved(), workspace.max(total)); + assert_eq!(parent.reserved(), other.size() + workspace.max(total)); + drop((reservation, spills, builder, other)); + merge_pool.release_unused(); + assert_eq!(parent.reserved(), 0); + Ok(()) + } + + #[test] + fn fan_in_headroom_beyond_usize_stops_seating_runs() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + // Three runs whose buffers fit in a usize, but whose headroom for the + // third does not. + let batch_memory = usize::MAX / 10; + let spills = (0..3) + .map(|value| SortedSpillFile { + max_record_batch_memory: batch_memory, + ..make_sorted_spill_file(&spill_manager, &schema, vec![value]) + }) + .collect::>(); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let other = hold(&pool, 1)?; + let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_fan_in_headroom(headroom(&pool, false)); + let mut reservation = builder.reservation.new_empty(); + let SpillFilesToMerge::Ready(spills, _) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)? + else { + panic!("the minimum merge should fit an unbounded pool"); + }; + assert_eq!(spills.len(), 2); + assert_eq!(reservation.size(), 4 * batch_memory); + drop((reservation, spills, builder, other)); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + #[tokio::test] async fn replay_headroom_still_enforces_the_actual_pool() -> 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 17e66ef339950..7ab2224865f30 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -309,6 +309,9 @@ struct ExternalSorter { /// How much memory to reserve for performing in-memory sort/merges /// prior to spilling. sort_spill_reservation_bytes: usize, + /// Other partitions of this sort run alongside this one and share its + /// memory pool. + has_sibling_partitions: bool, /// Shared by the owning SortExec's partitions. Initialized on the runtime /// that polls the sort, which may differ from the one calling `execute`. sort_concurrency: Arc>, @@ -372,10 +375,17 @@ impl ExternalSorter { batch_size, sort_spill_reservation_bytes, sort_in_place_threshold_bytes, + has_sibling_partitions: false, sort_concurrency: Arc::new(OnceLock::new()), }) } + /// Record whether other partitions of this sort share its memory pool. + fn with_sibling_partitions(mut self, has_sibling_partitions: bool) -> Self { + self.has_sibling_partitions = has_sibling_partitions; + self + } + /// Appends an unsorted [`RecordBatch`] to `in_mem_batches` /// /// Updates memory usage metrics, and possibly triggers spilling to disk @@ -431,6 +441,19 @@ impl ExternalSorter { .with_fetch(None) .with_reservation(self.merge_reservation.take()) .with_merge_pool(Arc::clone(&self.merge_pool)) + // Other partitions of this sort, still buffering input, share + // the pool with this merge, and may hold nothing at the moment + // it seats its runs. Seating every run the pool admits would + // leave them nothing, so keep headroom while they exist or + // anyone else holds memory in the pool. A lone partition with + // the pool to itself seats every run that fits: the operator + // consuming its output reserves nothing before its first + // batch, and keeping headroom for it would cost an extra merge + // pass whether or not it ever needs the memory. + .with_fan_in_headroom( + Arc::clone(&self.runtime.memory_pool), + self.has_sibling_partitions, + ) .build() } else { // Final output needs no reserve for future spills. Return unused @@ -1841,7 +1864,10 @@ impl ExecutionPlan for SortExec { context.session_config().spill_compression(), &self.metrics_set, context.runtime_env(), - )?; + )? + .with_sibling_partitions( + self.input.output_partitioning().partition_count() > 1, + ); sorter.sort_concurrency = Arc::clone(&self.sort_concurrency); Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 625f96be682cb..ca996ea3283cd 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -3,7 +3,8 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements use super::{ - ExternalSorter, get_reserved_bytes_for_record_batch, sort_batch, sort_batch_chunked, + ExternalSorter, get_reserved_bytes_for_record_batch, + get_reserved_bytes_for_record_batch_size, sort_batch, sort_batch_chunked, }; use crate::metrics::ExecutionPlanMetricsSet; use crate::spill::get_record_batch_memory_size; @@ -1202,3 +1203,87 @@ async fn test_single_batch_spill_returns_live_workspace_loan_on_drop() -> Result assert_released(&pool, &runtime).await; Ok(()) } + +/// The final merge of a spilled sort shares the pool with sibling partitions +/// still sorting their input. Seating every run the pool admits leaves them +/// nothing, so whatever the merge takes beyond its own workspace, it must leave +/// as much free: while another consumer holds memory, and while sibling +/// partitions exist even if they hold nothing at the moment. +#[rstest::rstest] +#[case::other_consumer_holds_memory(false)] +#[case::idle_sibling_partitions(true)] +#[tokio::test] +async fn test_final_spilled_merge_leaves_headroom(#[case] siblings: bool) -> Result<()> { + let fixture = fixture()?; + let schema = fixture.parent.schema(); + let capacity = 16 * fixture.capacity; + let pool = AdjustablePool::new(capacity); + // Spill under the fixture's limit, then merge with room for every run. + pool.set_limit(fixture.capacity); + let (sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?; + let mut sorter = sorter.with_sibling_partitions(siblings); + for batch in &fixture.batches { + sorter.insert_batch(batch.clone()).await?; + } + if !sorter.in_mem_batches.is_empty() { + sorter.sort_and_spill_in_mem_batches().await?; + } + // What the merge reserves to read each run with two buffers. + let needs: Vec = sorter + .finished_spill_files + .iter() + .map(|spill| { + let size = spill.max_record_batch_memory; + 2 * get_reserved_bytes_for_record_batch_size(size, size) + }) + .collect(); + assert!(needs.len() >= 3, "need three runs, got {}", needs.len()); + let all: usize = needs.iter().sum(); + let smallest = *needs.iter().min().unwrap(); + pool.set_limit(capacity); + + let mut stream = sorter.sort().await?; + drop(sorter); + // The merge already holds its workspace. Leave it room to seat every + // run, but not to seat them all and leave as much free. + let workspace = pool.reserved(); + let free = (2 * all) + .checked_sub(2 * workspace + smallest) + .expect("the runs must need more than the workspace"); + assert!(all <= free + workspace); + // Either another consumer holds the rest of the pool, or idle siblings + // hold nothing and the pool is only as large as the merge may use. + let contender = + MemoryConsumer::new("sibling partition").register(&runtime.memory_pool); + let limit = if siblings { + pool.set_limit(workspace + free); + workspace + free + } else { + contender.try_grow(capacity - workspace - free)?; + capacity + }; + + let first = stream + .try_next() + .await? + .expect("the merge must produce output"); + let taken = (pool.reserved() - contender.size()).saturating_sub(workspace); + assert!( + limit - pool.reserved() >= taken, + "the final merge takes {taken} bytes beyond its workspace and leaves {} free", + limit - pool.reserved() + ); + drop(contender); + + let mut output = vec![first]; + while let Some(batch) = stream.try_next().await? { + output.push(batch); + } + assert_eq!( + concat_batches(&schema, &output)?, + sort_batch(&fixture.parent, &fixture.ordering, None)? + ); + 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 da341cf0cf047..197cd43b1414f 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -6,7 +6,7 @@ //! This is an order-preserving merge. use crate::metrics::BaselineMetrics; -use crate::sorts::multi_level_merge::MultiLevelMergeBuilder; +use crate::sorts::multi_level_merge::{FanInHeadroom, MultiLevelMergeBuilder}; use crate::sorts::{ merge::SortPreservingMergeStream, stream::{FieldCursorStream, RowCursorStream}, @@ -85,6 +85,9 @@ pub struct StreamingMergeBuilder<'a> { merge_pool: Option>, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, + /// Leave memory for others sharing the execution pool when seating more + /// than the minimum two runs. + fan_in_headroom: Option, min_spill_batch_rows: Option, enable_round_robin_tie_breaker: bool, } @@ -158,6 +161,23 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// While `siblings` is true or another consumer holds memory in `pool`, + /// seat spill runs beyond the minimum two only if the pool can admit their + /// buffers plus as much again as they exceed the retained workspace. + /// Release the headroom before merging. + /// + /// `pool` is the execution pool: the parent of the merge pool when one is + /// set, not the merge pool itself. `siblings` says other partitions share + /// that pool even when they hold nothing at the moment. + pub(crate) fn with_fan_in_headroom( + mut self, + pool: Arc, + siblings: bool, + ) -> Self { + self.fan_in_headroom = Some(FanInHeadroom { pool, siblings }); + self + } + /// Limit intermediate rewrites when the retained buffer reservation can /// also admit the final replay pass with its headroom. /// `min_spill_batch_rows` is the minimum largest-batch row count across input @@ -205,6 +225,7 @@ impl<'a> StreamingMergeBuilder<'a> { reservation, merge_pool, reserve_replay_headroom, + fan_in_headroom, min_spill_batch_rows, fetch, expressions, @@ -248,6 +269,7 @@ impl<'a> StreamingMergeBuilder<'a> { ) .with_merge_pool(merge_pool) .with_replay_headroom(reserve_replay_headroom) + .with_fan_in_headroom(fan_in_headroom) .with_intermediate_merge_sizing(min_spill_batch_rows) .create_spillable_merge_stream()); }