From 6194e50c1296481326a21558660601f2e9b91bc4 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 6 Nov 2026 13:42:37 +0800 Subject: [PATCH 1/4] Add sort_spill bench for datafusion/physical-plan - Four lone-sort configurations; reports spill counts, rows, bytes. - Validates sorted output and reservation cleanup. - Debug/CI-profile smoke, Clippy, formatting pass. - Full suites timed out. Four-partition case excluded: memory exhaustion. --- datafusion/physical-plan/Cargo.toml | 5 + .../physical-plan/benches/sort_spill.rs | 177 ++++++++++++++++++ 2 files changed, 182 insertions(+) create mode 100644 datafusion/physical-plan/benches/sort_spill.rs 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); From 9136c5795fe08c1e762661c11767e45878663924 Mon Sep 17 00:00:00 2001 From: MisterRaindrop <278811821@qq.com> Date: Tue, 13 Oct 2026 19:39:01 +0800 Subject: [PATCH 2/4] fix: leave memory for others when a spilled sort merges its runs The final merge of a spilled sort seated every run the memory pool would admit, and held those buffers until its output was consumed. Partitions that finished their input early could take the whole pool this way, and a partition still sorting then failed with ResourcesExhausted (#25804, finding 9). Seat runs beyond the minimum two only while the pool can admit what the merge takes beyond its own workspace plus as much again, and return the extra before merging, as #25383 does for aggregate replay. The minimum merge is admitted as before, so a sort under tight memory keeps its batch sizes and does not re-spill runs to leave headroom. --- .../src/sorts/multi_level_merge.rs | 37 +++++++++- datafusion/physical-plan/src/sorts/sort.rs | 5 ++ .../src/sorts/sort/spill_tests.rs | 74 ++++++++++++++++++- .../src/sorts/streaming_merge.rs | 12 +++ 4 files changed, 124 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 424d1faffcf33..8e33eb3920782 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; @@ -146,6 +146,10 @@ pub(crate) struct MultiLevelMergeBuilder { merge_pool: Option>, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, + /// Seat runs beyond the minimum two only while the pool can admit their + /// buffers plus equal headroom, so a merge does not take memory that the + /// operator consuming it and other partitions sharing the pool need. + reserve_fan_in_headroom: bool, /// Avoid merging more runs than needed before the final replay pass. size_intermediate_merges: bool, fetch: Option, @@ -188,6 +192,7 @@ impl MultiLevelMergeBuilder { reservation, merge_pool: None, reserve_replay_headroom: false, + reserve_fan_in_headroom: false, size_intermediate_merges: false, enable_round_robin_tie_breaker, fetch, @@ -206,6 +211,13 @@ impl MultiLevelMergeBuilder { self } + /// Leave headroom when seating runs beyond the minimum two. The minimum + /// merge is admitted against the pool alone, as without headroom. + pub(super) fn with_fan_in_headroom(mut self, reserve: bool) -> Self { + self.reserve_fan_in_headroom = reserve; + self + } + pub(super) fn with_intermediate_merge_sizing( mut self, min_spill_batch_rows: Option, @@ -573,6 +585,12 @@ 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()); for (spill, _) in &self.sorted_spill_files { if number_of_spills_to_read_for_current_phase >= max_spill_files @@ -598,7 +616,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 + && self.reserve_fan_in_headroom + && 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 +697,7 @@ impl MultiLevelMergeBuilder { } } - if self.reserve_replay_headroom { + if self.reserve_replay_headroom || self.reserve_fan_in_headroom { // `total_needed` may include a rejected candidate. Keep only the // buffers that were admitted, releasing temporary replay headroom. reservation.shrink(reservation.size() - accepted_memory); diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 17e66ef339950..c145b49533878 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -431,6 +431,11 @@ impl ExternalSorter { .with_fetch(None) .with_reservation(self.merge_reservation.take()) .with_merge_pool(Arc::clone(&self.merge_pool)) + // The operator consuming the sorted output, and the other + // partitions of this sort still buffering input, share the + // pool with this merge. Seating every run the pool admits + // would leave them nothing. + .with_fan_in_headroom() .build() } else { // Final output needs no reserve for future spills. Return unused diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 625f96be682cb..2fe493bbc106c 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,74 @@ 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 the operator that +/// consumes its output and 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. +#[tokio::test] +async fn test_final_spilled_merge_leaves_headroom() -> 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 (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?; + 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); + let contender = + MemoryConsumer::new("sibling partition").register(&runtime.memory_pool); + contender.try_grow(capacity - workspace - free)?; + + let first = stream + .try_next() + .await? + .expect("the merge must produce output"); + let taken = (pool.reserved() - contender.size()).saturating_sub(workspace); + assert!( + capacity - pool.reserved() >= taken, + "the final merge takes {taken} bytes beyond its workspace and leaves {} free", + capacity - 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..891383bf4da28 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -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 pool when seating more than the + /// minimum two runs. + reserve_fan_in_headroom: bool, min_spill_batch_rows: Option, enable_round_robin_tie_breaker: bool, } @@ -158,6 +161,13 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// Seat spill runs beyond the minimum two only while the pool can admit + /// their buffers plus equal headroom. Release the headroom before merging. + pub(crate) fn with_fan_in_headroom(mut self) -> Self { + self.reserve_fan_in_headroom = true; + 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 +215,7 @@ impl<'a> StreamingMergeBuilder<'a> { reservation, merge_pool, reserve_replay_headroom, + reserve_fan_in_headroom, min_spill_batch_rows, fetch, expressions, @@ -248,6 +259,7 @@ impl<'a> StreamingMergeBuilder<'a> { ) .with_merge_pool(merge_pool) .with_replay_headroom(reserve_replay_headroom) + .with_fan_in_headroom(reserve_fan_in_headroom) .with_intermediate_merge_sizing(min_spill_batch_rows) .create_spillable_merge_stream()); } From 06663ce2f383108ef2c9b6dc8a00416db390bde8 Mon Sep 17 00:00:00 2001 From: MisterRaindrop <278811821@qq.com> Date: Wed, 14 Oct 2026 09:12:34 +0800 Subject: [PATCH 3/4] test: cover fan-in headroom when seating spill runs Check at the builder that runs beyond the minimum two are seated only with equal headroom, and that a headroom request beyond usize::MAX stops seating runs instead of failing the merge. --- .../src/sorts/multi_level_merge.rs | 67 ++++++++++++++++++- 1 file changed, 66 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 8e33eb3920782..6dfd75a41a05d 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -997,7 +997,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}; @@ -1494,6 +1494,71 @@ mod tests { Ok(()) } + #[rstest::rstest] + #[case::without_headroom(false, 4, 8)] + #[case::with_headroom(true, 3, 6)] + fn fan_in_headroom_limits_runs_beyond_the_minimum( + #[case] headroom: 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 mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_fan_in_headroom(headroom); + 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. + assert_eq!(spills.len(), seated); + assert_eq!(reservation.size(), retained_batches * batch_memory); + drop((reservation, spills, builder)); + assert_eq!(pool.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 mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_fan_in_headroom(true); + 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)); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + #[tokio::test] async fn replay_headroom_still_enforces_the_actual_pool() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); From 217bd1f6ae39a530e95a1a2e0f012d9d18234219 Mon Sep 17 00:00:00 2001 From: MisterRaindrop <278811821@qq.com> Date: Thu, 5 Nov 2026 21:56:24 +0800 Subject: [PATCH 4/4] fix: keep spill merge headroom only when the pool is shared The fan-in headroom made a lone sort partition that has the pool to itself rewrite runs in an extra merge pass, for no one: with 60 x 8192 rows of Int64 plus a 120-byte string, 1 MB sort_spill_reservation_bytes and a 16 MB GreedyMemoryPool, it wrote 1,589,248 rows to spill instead of 983,040. Keep the headroom only while the sort has other partitions sharing the pool, or another consumer holds memory in it. Sibling partitions count even when they hold nothing at the moment the merge seats its runs. The operator consuming the output of a lone partition is not counted: it reserves nothing before its first batch, and keeping headroom for it would cost the extra pass whether or not it ever needs the memory. With one partition, the case above writes the same rows as without headroom in all four configurations measured. With four partitions in a 32 MB pool and 1 MB reservations, the sort fails with "Not enough memory to continue external sort" without headroom and succeeds with it. Also cover the minimum two runs under a budget below their headroom, and seating against a retained merge workspace on both sides of 2 * (T - W). --- .../src/sorts/multi_level_merge.rs | 185 ++++++++++++++++-- datafusion/physical-plan/src/sorts/sort.rs | 33 +++- .../src/sorts/sort/spill_tests.rs | 31 ++- .../src/sorts/streaming_merge.rs | 30 ++- 4 files changed, 235 insertions(+), 44 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 6dfd75a41a05d..a2d1b30ee05f7 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -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,10 +164,11 @@ pub(crate) struct MultiLevelMergeBuilder { merge_pool: Option>, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, - /// Seat runs beyond the minimum two only while the pool can admit their - /// buffers plus equal headroom, so a merge does not take memory that the - /// operator consuming it and other partitions sharing the pool need. - reserve_fan_in_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, @@ -192,7 +211,7 @@ impl MultiLevelMergeBuilder { reservation, merge_pool: None, reserve_replay_headroom: false, - reserve_fan_in_headroom: false, + fan_in_headroom: None, size_intermediate_merges: false, enable_round_robin_tie_breaker, fetch, @@ -211,10 +230,14 @@ impl MultiLevelMergeBuilder { self } - /// Leave headroom when seating runs beyond the minimum two. The minimum - /// merge is admitted against the pool alone, as without headroom. - pub(super) fn with_fan_in_headroom(mut self, reserve: bool) -> Self { - self.reserve_fan_in_headroom = reserve; + /// 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 } @@ -591,6 +614,14 @@ impl MultiLevelMergeBuilder { .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 @@ -617,7 +648,7 @@ impl MultiLevelMergeBuilder { < minimum_number_of_required_streams; let check_headroom = self.reserve_replay_headroom && !skip_headroom; let check_fan_in_headroom = !check_headroom - && self.reserve_fan_in_headroom + && fan_in_contended && number_of_spills_to_read_for_current_phase >= minimum_number_of_required_streams; let admission = if check_fan_in_headroom { @@ -697,7 +728,7 @@ impl MultiLevelMergeBuilder { } } - if self.reserve_replay_headroom || self.reserve_fan_in_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); @@ -1494,11 +1525,31 @@ 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, 4, 8)] - #[case::with_headroom(true, 3, 6)] + #[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] headroom: bool, + #[case] with_headroom: bool, + #[case] others: bool, + #[case] siblings: bool, #[case] seated: usize, #[case] retained_batches: usize, ) -> Result<()> { @@ -1511,8 +1562,11 @@ mod tests { 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(headroom); + .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)? @@ -1521,14 +1575,106 @@ mod tests { }; 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. + // 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)); + 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()); @@ -1544,8 +1690,9 @@ mod tests { }) .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(true); + .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)? @@ -1554,7 +1701,7 @@ mod tests { }; assert_eq!(spills.len(), 2); assert_eq!(reservation.size(), 4 * batch_memory); - drop((reservation, spills, builder)); + drop((reservation, spills, builder, other)); assert_eq!(pool.reserved(), 0); Ok(()) } diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index c145b49533878..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,11 +441,19 @@ impl ExternalSorter { .with_fetch(None) .with_reservation(self.merge_reservation.take()) .with_merge_pool(Arc::clone(&self.merge_pool)) - // The operator consuming the sorted output, and the other - // partitions of this sort still buffering input, share the - // pool with this merge. Seating every run the pool admits - // would leave them nothing. - .with_fan_in_headroom() + // 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 @@ -1846,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 2fe493bbc106c..ca996ea3283cd 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -1204,19 +1204,24 @@ async fn test_single_batch_spill_returns_live_workspace_loan_on_drop() -> Result Ok(()) } -/// The final merge of a spilled sort shares the pool with the operator that -/// consumes its output and 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. +/// 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() -> Result<()> { +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 (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?; + 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?; } @@ -1246,9 +1251,17 @@ async fn test_final_spilled_merge_leaves_headroom() -> Result<()> { .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); - contender.try_grow(capacity - workspace - free)?; + let limit = if siblings { + pool.set_limit(workspace + free); + workspace + free + } else { + contender.try_grow(capacity - workspace - free)?; + capacity + }; let first = stream .try_next() @@ -1256,9 +1269,9 @@ async fn test_final_spilled_merge_leaves_headroom() -> Result<()> { .expect("the merge must produce output"); let taken = (pool.reserved() - contender.size()).saturating_sub(workspace); assert!( - capacity - pool.reserved() >= taken, + limit - pool.reserved() >= taken, "the final merge takes {taken} bytes beyond its workspace and leaves {} free", - capacity - pool.reserved() + limit - pool.reserved() ); drop(contender); diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 891383bf4da28..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,9 +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 pool when seating more than the - /// minimum two runs. - reserve_fan_in_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, } @@ -161,10 +161,20 @@ impl<'a> StreamingMergeBuilder<'a> { self } - /// Seat spill runs beyond the minimum two only while the pool can admit - /// their buffers plus equal headroom. Release the headroom before merging. - pub(crate) fn with_fan_in_headroom(mut self) -> Self { - self.reserve_fan_in_headroom = true; + /// 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 } @@ -215,7 +225,7 @@ impl<'a> StreamingMergeBuilder<'a> { reservation, merge_pool, reserve_replay_headroom, - reserve_fan_in_headroom, + fan_in_headroom, min_spill_batch_rows, fetch, expressions, @@ -259,7 +269,7 @@ impl<'a> StreamingMergeBuilder<'a> { ) .with_merge_pool(merge_pool) .with_replay_headroom(reserve_replay_headroom) - .with_fan_in_headroom(reserve_fan_in_headroom) + .with_fan_in_headroom(fan_in_headroom) .with_intermediate_merge_sizing(min_spill_batch_rows) .create_spillable_merge_stream()); }