diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 911eca0a97928..b0433250c0c51 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -606,7 +606,11 @@ impl ExecutionPlan for SortMergeJoinExec { let buffered = buffered.execute(partition, Arc::clone(&context))?; let batch_size = context.session_config().batch_size(); + // The stream spills its buffered batches when it cannot grow, so a + // pool that budgets spillable and unspillable consumers differently + // (`FairSpillPool`) has to know it can. let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]")) + .with_can_spill(true) .register(context.memory_pool()); let spill_manager = SpillManager::new( context.runtime_env(), diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index 4d817df738e9c..ea9299abdd985 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -34,6 +34,7 @@ use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; use crate::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; use crate::projection::{ProjectionExec, ProjectionExpr}; +use crate::sorts::sort::SortExec; use crate::spill::spill_manager::SpillManager; use crate::test::TestMemoryExec; use crate::test::exec::BarrierExec; @@ -55,24 +56,28 @@ use bytes::Bytes; use datafusion_common::JoinType::*; use datafusion_common::instant::Instant; use datafusion_common::{ - JoinSide, internal_err, - test_util::{batches_to_sort_string, batches_to_string}, + DataFusionError, JoinType, NullEquality, Result, ScalarValue, assert_batches_eq, + assert_contains, }; use datafusion_common::{ - JoinType, NullEquality, Result, ScalarValue, assert_batches_eq, assert_contains, + JoinSide, internal_err, + test_util::{batches_to_sort_string, batches_to_string}, }; use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; use datafusion_execution::disk_manager::{ DiskManager, DiskManagerBuilder, DiskManagerMode, }; -use datafusion_execution::memory_pool::MemoryConsumer; +use datafusion_execution::memory_pool::{ + FairSpillPool, MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_execution::spill_file::{SpillFile, SpillWriter, TempFileFactory}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::BinaryExpr; use datafusion_physical_expr::expressions::Literal; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use futures::{Stream, StreamExt}; use insta::assert_snapshot; @@ -2465,6 +2470,203 @@ async fn overallocation_multi_batch_no_spill() -> Result<()> { Ok(()) } +/// The stream spills its buffered side when it cannot grow, so it has to be +/// registered as a consumer that can spill. A `FairSpillPool` otherwise treats +/// it as unspillable: it is left out of the fair share the spillable consumers +/// split, and may take everything they have not yet claimed, starving the +/// sorts the join usually runs on top of. +#[tokio::test] +async fn stream_registers_as_a_spillable_consumer() -> Result<()> { + /// Records how each consumer registered, and otherwise never limits anything. + #[derive(Debug, Default)] + struct RecordingPool { + inner: UnboundedMemoryPool, + registered: std::sync::Mutex>, + } + impl std::fmt::Display for RecordingPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "RecordingPool") + } + } + impl MemoryPool for RecordingPool { + fn name(&self) -> &str { + "RecordingPool" + } + fn register(&self, consumer: &MemoryConsumer) { + self.registered + .lock() + .unwrap() + .push((consumer.name().to_string(), consumer.can_spill())); + } + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional) + } + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink) + } + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> Result<()> { + self.inner.try_grow(reservation, additional) + } + fn reserved(&self) -> usize { + self.inner.reserved() + } + } + + let pool = Arc::new(RecordingPool::default()); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?; + let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime)); + + let left = build_table( + ("a1", &vec![1, 2]), + ("b1", &vec![1, 2]), + ("c1", &vec![7, 8]), + ); + let right = build_table( + ("a2", &vec![1, 2]), + ("b2", &vec![1, 2]), + ("c2", &vec![9, 10]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let join = join(left, right, on, Inner)?; + common::collect(join.execute(0, task_ctx)?).await?; + + let registered = pool.registered.lock().unwrap(); + let (_, can_spill) = registered + .iter() + .find(|(name, _)| name == "SMJStream[0]") + .expect("the stream registers a reservation under its own name"); + assert!( + can_spill, + "the sort-merge join stream must register as able to spill" + ); + Ok(()) +} + +/// Runs the join under a `FairSpillPool` too small for its buffered side, on +/// top of a sort that competes with it for that pool. +/// +/// The buffered side is a single equal-key group larger than the pool, so the +/// join spills part of it and holds the rest while it drains the streamed +/// side. The streamed side is an in-memory sort of several batches whose merge +/// asks the pool for every chunk it pulls in. Registered as unspillable, the +/// join would take every byte the sort has not claimed before it starts +/// spilling, and that merge would fail to allocate. As a spillable consumer +/// it is held to its fair share and both finish. +#[tokio::test] +async fn fair_spill_pool_leaves_room_for_the_streamed_sort() -> Result<()> { + // Measured window: the join passes as a spillable consumer from 240 KB + // up, and still starves the sort as an unspillable one up to 704 KB. + const POOL_SIZE: usize = 384 * 1024; + const GROUP_KEY: i32 = i32::MAX; + const STREAMED_BATCHES: usize = 2; + const STREAMED_ROWS: usize = 512; + const BUFFERED_BATCHES: usize = 32; + const BUFFERED_ROWS: usize = 1024; + + // Unsorted: every batch carries one row of the buffered group's key, which + // the sort moves to the end, so the join holds the whole buffered group + // while it pulls the streamed side through the sort's merge. + let streamed = build_table_from_batches( + (0..STREAMED_BATCHES) + .map(|batch| { + let base = (batch * STREAMED_ROWS) as i32; + let a: Vec = (0..STREAMED_ROWS as i32).map(|i| base + i).collect(); + let b: Vec = (0..STREAMED_ROWS as i32) + .map(|i| { + if i == 0 { + GROUP_KEY + } else { + STREAMED_ROWS as i32 - i + } + }) + .collect(); + let c: Vec = + (0..STREAMED_ROWS as i32).map(|i| base + i + 1).collect(); + build_table_i32(("a1", &a), ("b1", &b), ("c1", &c)) + }) + .collect(), + ); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new_with_schema("b1", &streamed.schema())?, + ))]) + .unwrap(); + let streamed = Arc::new(SortExec::new(ordering, streamed)) as Arc; + + // Already sorted: one group, larger than the pool. + let buffered = build_table_from_batches( + (0..BUFFERED_BATCHES) + .map(|batch| { + let base = (batch * BUFFERED_ROWS) as i32; + let a: Vec = (0..BUFFERED_ROWS as i32).map(|i| base + i).collect(); + let b = vec![GROUP_KEY; BUFFERED_ROWS]; + let c: Vec = (0..BUFFERED_ROWS as i32).map(|i| base - i).collect(); + build_table_i32(("a2", &a), ("b2", &b), ("c2", &c)) + }) + .collect(), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &streamed.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &buffered.schema())?) as _, + )]; + let session_config = SessionConfig::default() + .with_batch_size(128) + // Keep the sort in memory and merging chunk by chunk rather than + // sorting one concatenated batch in place. + .with_sort_in_place_threshold_bytes(0) + .with_sort_spill_reservation_bytes(0); + + let run = |runtime| async { + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(runtime), + ); + let join = join( + Arc::clone(&streamed), + Arc::clone(&buffered), + on.clone(), + Inner, + )?; + let batches = common::collect(join.execute(0, task_ctx)?).await?; + let metrics = join.metrics().unwrap(); + Ok::<_, DataFusionError>((batches, metrics.spill_count().unwrap())) + }; + + let fair = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(POOL_SIZE))) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + let (spilled, spill_count) = run(fair).await?; + assert!( + spill_count > 0, + "the join must have spilled its buffered side" + ); + assert_eq!( + spilled.iter().map(|b| b.num_rows()).sum::(), + STREAMED_BATCHES * BUFFERED_BATCHES * BUFFERED_ROWS + ); + + let (unbounded, spill_count) = run(RuntimeEnvBuilder::new().build_arc()?).await?; + assert_eq!(spill_count, 0); + assert_eq!( + batches_to_sort_string(&spilled), + batches_to_sort_string(&unbounded) + ); + Ok(()) +} + #[tokio::test] async fn overallocation_single_batch_spill() -> Result<()> { let left = build_table( @@ -4477,7 +4679,7 @@ fn test_stream_resources( inner_schema: SchemaRef, metrics: &ExecutionPlanMetricsSet, ) -> ( - datafusion_execution::memory_pool::MemoryReservation, + MemoryReservation, SpillManager, Arc, ) { @@ -6305,7 +6507,7 @@ impl SpillFile for PendingSpillFile { tokio::fs::read(&path) .await .map(Bytes::from) - .map_err(datafusion_common::DataFusionError::IoError) + .map_err(DataFusionError::IoError) }) .flat_map( |read_result| -> Pin> + Send>> {