fix: register the sort-merge join stream as a spillable memory consumer

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(),

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<Vec<(String, bool)>>,
+    }
+    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<dyn MemoryPool>)
+        .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<i32> = (0..STREAMED_ROWS as i32).map(|i| base + i).collect();
+                let b: Vec<i32> = (0..STREAMED_ROWS as i32)
+                    .map(|i| {
+                        if i == 0 {
+                            GROUP_KEY
+                        } else {
+                            STREAMED_ROWS as i32 - i
+                        }
+                    })
+                    .collect();
+                let c: Vec<i32> =
+                    (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<dyn ExecutionPlan>;
+
+    // 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<i32> = (0..BUFFERED_ROWS as i32).map(|i| base + i).collect();
+                let b = vec![GROUP_KEY; BUFFERED_ROWS];
+                let c: Vec<i32> = (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::<usize>(),
+        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<datafusion_execution::runtime_env::RuntimeEnv>,
 ) {
@@ -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<Box<dyn Stream<Item = Result<Bytes>> + Send>> {