diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 424d1faffcf33..39d06121bf111 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,185 @@ 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) -> FanInHeadroom { + 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))); + 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(Some(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(Some(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(Some(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..b1fd5a05e5efe 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -9,6 +9,7 @@ use std::fmt; use std::fmt::{Debug, Formatter}; use std::future::Future; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; use parking_lot::RwLock; @@ -309,6 +310,9 @@ struct ExternalSorter { /// How much memory to reserve for performing in-memory sort/merges /// prior to spilling. sort_spill_reservation_bytes: usize, + /// Partitions of the owning `SortExec`, this one included, still reading + /// their input. Each one will need memory from the pool this sort shares. + partitions_reading_input: Arc, /// 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 +376,25 @@ impl ExternalSorter { batch_size, sort_spill_reservation_bytes, sort_in_place_threshold_bytes, + partitions_reading_input: Arc::new(AtomicUsize::new(1)), sort_concurrency: Arc::new(OnceLock::new()), }) } + /// Share the count of partitions still reading their input with the + /// other partitions of this sort. It must include this partition. + fn with_partitions_reading_input(mut self, reading: Arc) -> Self { + self.partitions_reading_input = reading; + self + } + + /// Records that this partition has read all its input, and returns whether + /// another partition still reads its own. Executing partitions again + /// without `reset_state` can wrap the count, which only keeps headroom. + fn finish_reading_input(&self) -> bool { + self.partitions_reading_input.fetch_sub(1, Ordering::AcqRel) > 1 + } + /// Appends an unsorted [`RecordBatch`] to `in_mem_batches` /// /// Updates memory usage metrics, and possibly triggers spilling to disk @@ -406,6 +425,7 @@ impl ExternalSorter { /// 2. A combined streaming merge incorporating both in-memory /// batches and data from spill files on disk. async fn sort(&mut self) -> Result { + let other_partitions_reading = self.finish_reading_input(); if self.spilled_before() { // Sort `in_mem_batches` and spill it first. If there are many // `in_mem_batches` and the memory limit is almost reached, merging @@ -431,6 +451,21 @@ 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 any of them is + // still reading or anyone else holds memory in the pool; + // partitions already merging hold memory, so the second check + // covers them. A lone partition, or the last one to finish + // reading, 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), + other_partitions_reading, + ) .build() } else { // Final output needs no reserve for future spills. Return unused @@ -1326,6 +1361,10 @@ pub struct SortExec { /// Batch-sort limit shared by partitions and clones. /// `reset_state` returns a new plan with a fresh limit. sort_concurrency: Arc>, + /// Partitions still reading their input, shared with their sorters. Set + /// on the first `execute`, once the input can no longer change. + /// `reset_state` returns a new plan with a fresh count. + partitions_reading_input: Arc>>, } impl SortExec { @@ -1346,6 +1385,7 @@ impl SortExec { cache: Arc::new(cache), filter: None, sort_concurrency: Arc::new(OnceLock::new()), + partitions_reading_input: Arc::new(OnceLock::new()), } } @@ -1423,6 +1463,7 @@ impl SortExec { cache: Arc::clone(&self.cache), filter: self.filter.clone(), sort_concurrency: Arc::clone(&self.sort_concurrency), + partitions_reading_input: Arc::clone(&self.partitions_reading_input), } } @@ -1753,14 +1794,15 @@ impl ExecutionPlan for SortExec { .downcast_ref::() .expect("rebuilt SortExec with new children") .clone(); - // Our dynamic filter, execution metrics and sort concurrency are the state - // we need to reset. + // Our dynamic filter, execution metrics, sort concurrency and count of + // partitions reading input are the state we need to reset. if let Some(filter) = &new_sort.filter { filter.read().expr().reset(); new_sort.filter = Some(new_sort.create_filter()); } new_sort.metrics_set = ExecutionPlanMetricsSet::new(); new_sort.sort_concurrency = Arc::new(OnceLock::new()); + new_sort.partitions_reading_input = Arc::new(OnceLock::new()); Ok(Arc::new(new_sort)) } @@ -1841,7 +1883,14 @@ impl ExecutionPlan for SortExec { context.session_config().spill_compression(), &self.metrics_set, context.runtime_env(), - )?; + )? + .with_partitions_reading_input(Arc::clone( + self.partitions_reading_input.get_or_init(|| { + Arc::new(AtomicUsize::new( + self.input.output_partitioning().partition_count(), + )) + }), + )); sorter.sort_concurrency = Arc::clone(&self.sort_concurrency); Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), @@ -2044,6 +2093,7 @@ impl ExecutionPlan for SortExec { filter, // Execution state, initialized on the runtime polling the decoded plan. sort_concurrency: _, + partitions_reading_input: _, } = self; let input = ctx.encode_child(input)?; let expr = expr diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 625f96be682cb..0b6f97ff4281c 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -3,11 +3,14 @@ // 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, SortExec, get_reserved_bytes_for_record_batch, + get_reserved_bytes_for_record_batch_size, sort_batch, sort_batch_chunked, }; +use crate::ExecutionPlan; use crate::metrics::ExecutionPlanMetricsSet; use crate::spill::get_record_batch_memory_size; use crate::spill::spill_manager::GetSlicedSize; +use crate::test::TestMemoryExec; use arrow::array::{ ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int64Array, StringArray, StringViewArray, @@ -17,6 +20,8 @@ use arrow::datatypes::{DataType, Field, Int8Type, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::config::{ExecutionOptions, SpillCompression}; use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, @@ -26,6 +31,7 @@ use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use futures::TryStreamExt; use std::fmt::{Display, Formatter}; +use std::sync::atomic::AtomicUsize; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -1202,3 +1208,163 @@ 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 still read their input 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_partitions_reading_input(Arc::new(AtomicUsize::new( + 1 + usize::from(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(()) +} + +/// Whether a `SortExec` keeps fan-in headroom follows its partitions still +/// reading input. The pool fits every run's buffers but not the headroom on +/// top, so with headroom the final merge seats fewer runs and rewrites rows in +/// an extra pass. Partition 0 runs alone, so nothing else ever holds pool +/// memory and the outcome is deterministic. In the last case its sibling has +/// already read all its input, so there is nobody to keep headroom for. +#[rstest::rstest] +#[case::lone_partition(1, false, false)] +#[case::idle_sibling(2, false, true)] +#[case::finished_sibling(2, true, false)] +#[tokio::test] +async fn test_sort_exec_keeps_headroom_only_while_siblings_read_input( + #[case] input_partitions: usize, + #[case] finish_sibling_first: bool, + #[case] rewrites_rows: bool, +) -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("payload", DataType::Utf8, false), + ])); + let (batches, rows) = (64, 1024); + let batches = (0..batches) + .map(|batch| { + let keys = (0..rows).map(|row| ((batch * rows + row) * 2654435761) % 1000003); + let payload = (0..rows).map(|row| format!("{row:0120}")); + let columns: Vec = vec![ + Arc::new(Int64Array::from_iter_values(keys)), + Arc::new(StringArray::from_iter_values(payload)), + ]; + Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?) + }) + .collect::>>()?; + let total_rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + let mut partitions = vec![batches]; + partitions.resize(input_partitions, vec![]); + let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "key", 0, + )))] + .into(); + let sort = SortExec::new(ordering, input).with_preserve_partitioning(true); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(3500 * 1024))) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(1024) + .with_sort_spill_reservation_bytes(128 * 1024), + ) + .with_runtime(runtime), + ); + + if finish_sibling_first { + let sibling: Vec = + sort.execute(1, Arc::clone(&context))?.try_collect().await?; + assert!(sibling.is_empty()); + } + let output: Vec = sort.execute(0, context)?.try_collect().await?; + + assert_eq!( + output.iter().map(RecordBatch::num_rows).sum::(), + total_rows + ); + let spilled_rows = sort.metrics().unwrap().spilled_rows().unwrap(); + if rewrites_rows { + assert!(spilled_rows > total_rows, "{spilled_rows} <= {total_rows}"); + } else { + assert_eq!(spilled_rows, total_rows); + } + 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()); }