fix: keep every spilled slice of a sort-merge join inner key group

datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs

@@ -127,6 +127,7 @@ use crate::joins::utils::{JoinFilter, JoinKeyComparator, compare_join_arrays};
 use crate::metrics::{
     BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, Time,
 };
+use crate::spill::in_progress_spill_file::InProgressSpillFile;
 use crate::spill::spill_manager::SpillManager;
 use crate::stream::{ObservedStream, RecordBatchStreamAdapter};
 use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
@@ -234,9 +235,9 @@ pub(crate) struct BitwiseSortMergeJoinStream {
     // Inner key group buffer: all inner rows sharing the current join key.
     // Only populated when a filter is present. Unbounded — a single key
     // with many inner rows will buffer them all. See "Degenerate cases"
-    // in exec.rs. Spilled to disk when memory reservation fails.
+    // in exec.rs. On memory pool overflow the buffered slices move to a
+    // per-group spill file (see [`Self::buffer_inner_key_group`]).
     inner_key_buffer: Vec<RecordBatch>,
-    inner_key_spill: Option<Arc<dyn SpillFile>>,
 
     // Join ON expressions, evaluated against each new batch to produce
     // the key arrays used for sorted key comparisons.
@@ -339,7 +340,6 @@ impl BitwiseSortMergeJoinStream {
             inner_key_arrays: vec![],
             matched: BooleanBufferBuilder::new(0),
             inner_key_buffer: vec![],
-            inner_key_spill: None,
             on_outer,
             on_inner,
             filter,
@@ -443,18 +443,24 @@ impl BitwiseSortMergeJoinStream {
         Ok(self.inner_self_cmp.as_ref().unwrap())
     }
 
-    /// Spill the in-memory inner key buffer to disk and clear it.
-    fn spill_inner_key_buffer(&mut self) -> Result<()> {
-        let spill_file = self
-            .spill_manager
-            .spill_record_batch_and_finish(
-                &self.inner_key_buffer,
-                "semi_anti_smj_inner_key_spill",
-            )?
-            .expect("inner_key_buffer is non-empty when spilling");
-        self.inner_key_buffer.clear();
+    /// Spill the in-memory inner key buffer to disk and clear it. One key
+    /// group can spill repeatedly; every call appends to `writer` — the
+    /// group's single open spill file — creating it on first use.
+    fn spill_inner_key_buffer(
+        &mut self,
+        writer: &mut Option<InProgressSpillFile>,
+    ) -> Result<()> {
+        if writer.is_none() {
+            *writer = Some(
+                self.spill_manager
+                    .create_in_progress_file("semi_anti_smj_inner_key_spill")?,
+            );
+        }
+        let writer = writer.as_mut().unwrap();
+        for batch in self.inner_key_buffer.drain(..) {
+            writer.append_batch(&batch)?;
+        }
         self.inner_buffer_size = 0;
-        self.inner_key_spill = Some(spill_file);
         // Should succeed now — inner buffer has been spilled.
         self.try_resize_reservation()
     }
@@ -465,7 +471,6 @@ impl BitwiseSortMergeJoinStream {
     /// pool interactions (see apache/datafusion#20729).
     fn clear_inner_key_group(&mut self) {
         self.inner_key_buffer.clear();
-        self.inner_key_spill = None;
         self.inner_buffer_size = 0;
     }
 
@@ -639,13 +644,15 @@ impl BitwiseSortMergeJoinStream {
     /// cursor past the group. Collects all inner rows with the current key
     /// across batch boundaries. Sets `inner_batch` to `None` if inner is
     /// exhausted.
-    async fn buffer_inner_key_group(&mut self) -> Result<()> {
+    ///
+    /// Slices that overflow the memory pool are appended to a single spill
+    /// file, returned finished — ready for reading — once the whole group
+    /// has been buffered. `None` means the group fit in memory.
+    async fn buffer_inner_key_group(&mut self) -> Result<Option<Arc<dyn SpillFile>>> {
         self.clear_inner_key_group();
+        let mut writer: Option<InProgressSpillFile> = None;
 
-        loop {
-            let Some(inner_batch) = &self.inner_batch else {
-                return Ok(());
-            };
+        while let Some(inner_batch) = &self.inner_batch {
             let num_inner = inner_batch.num_rows();
             let from = self.inner_offset;
             let group_end =
@@ -660,7 +667,7 @@ impl BitwiseSortMergeJoinStream {
             // is exhausted, spill the entire buffer to disk.
             if self.try_resize_reservation().is_err() {
                 if self.runtime_env.disk_manager.tmp_files_enabled() {
-                    self.spill_inner_key_buffer()?;
+                    self.spill_inner_key_buffer(&mut writer)?;
                 } else {
                     // Re-attempt to get the error message
                     self.try_resize_reservation().map_err(|e| {
@@ -673,7 +680,7 @@ impl BitwiseSortMergeJoinStream {
 
             if group_end < num_inner {
                 self.inner_offset = group_end;
-                return Ok(());
+                break;
             }
 
             // Key group extends to the end of the batch — it may continue
@@ -682,28 +689,38 @@ impl BitwiseSortMergeJoinStream {
 
             if !self.next_inner_batch().await? {
                 self.inner_batch = None;
-                return Ok(());
+                break;
             }
             if !keys_match(
                 &saved_inner_keys,
                 &self.inner_key_arrays,
                 &self.sort_options,
                 self.null_equality,
             )? {
-                return Ok(());
+                break;
             }
         }
+
+        match writer {
+            Some(mut writer) => writer.finish(),
+            None => Ok(None),
+        }
     }
 
     /// Process a key match with a filter. For each inner row in the buffered
-    /// key group, evaluates the filter against the outer key group and ORs
-    /// the results into the matched bitset using u64-chunked bitwise ops.
-    async fn process_key_match_with_filter(&mut self) -> Result<()> {
+    /// key group — the spilled slices in `spill` plus the in-memory
+    /// `inner_key_buffer` — evaluates the filter against the outer key group
+    /// and ORs the results into the matched bitset using u64-chunked bitwise
+    /// ops.
+    async fn process_key_match_with_filter(
+        &mut self,
+        spill: Option<&Arc<dyn SpillFile>>,
+    ) -> Result<()> {
         let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
 
         // buffer_inner_key_group must be called before this function
         debug_assert!(
-            !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(),
+            !self.inner_key_buffer.is_empty() || spill.is_some(),
             "process_key_match_with_filter called with no inner key data"
         );
         debug_assert!(
@@ -734,7 +751,7 @@ impl BitwiseSortMergeJoinStream {
 
         // Process spilled inner batches first asynchronously.
         if matched_count < outer_group_len
-            && let Some(spill_file) = &self.inner_key_spill
+            && let Some(spill_file) = spill
         {
             let mut spill_stream = self
                 .spill_manager
@@ -799,10 +816,14 @@ impl BitwiseSortMergeJoinStream {
 
     /// Evaluate the filter for the buffered inner key group against the
     /// outer key group. If the outer key group continues into subsequent
-    /// outer batches, keep evaluating there too.
-    async fn process_filtered_match_loop(&mut self) -> Result<()> {
+    /// outer batches, keep evaluating there too. Dropping `spill` on return
+    /// deletes the group's temp file.
+    async fn process_filtered_match_loop(
+        &mut self,
+        spill: Option<Arc<dyn SpillFile>>,
+    ) -> Result<()> {
         loop {
-            self.process_key_match_with_filter().await?;
+            self.process_key_match_with_filter(spill.as_ref()).await?;
 
             let outer_batch = self.outer_batch.as_ref().unwrap();
             if self.outer_offset < outer_batch.num_rows() {
@@ -872,8 +893,8 @@ impl BitwiseSortMergeJoinStream {
             // Buffer the inner key group so each inner row can be evaluated
             // against the outer key group, OR-ing filter results into the
             // matched bitset.
-            self.buffer_inner_key_group().await?;
-            self.process_filtered_match_loop().await
+            let spill = self.buffer_inner_key_group().await?;
+            self.process_filtered_match_loop(spill).await
         } else {
             // Without a filter, key equality alone means every outer row in
             // the group matches; the inner rows themselves are not needed.

datafusion/physical-plan/src/joins/sort_merge_join/tests.rs

@@ -5033,6 +5033,77 @@ async fn bitwise_spill_with_filter() -> Result<()> {
     Ok(())
 }
 
+/// A single inner key group spanning several inner batches can spill more
+/// than once under memory pressure. Every spilled slice must still be
+/// evaluated against the outer rows — an earlier spill file must not be
+/// dropped when a later slice of the same group spills.
+#[tokio::test]
+async fn bitwise_multi_spill_inner_key_group() -> Result<()> {
+    // Outer: one row with key 1, c1 = 5.
+    let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![5]));
+
+    // Inner: one key group (b2 = 1) spanning two batches. Only the first
+    // batch satisfies the filter c1 < c2 (5 < 10); the second (5 < 0) does
+    // not, so dropping the first spilled slice flips the semi-join result.
+    let right_batches = vec![
+        build_table_i32(("a2", &vec![10]), ("b2", &vec![1]), ("c2", &vec![10])),
+        build_table_i32(("a2", &vec![20]), ("b2", &vec![1]), ("c2", &vec![0])),
+    ];
+    let right = build_table_from_batches(right_batches);
+
+    let on = vec![(
+        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
+    )];
+    let sort_options = vec![SortOptions::default(); on.len()];
+    let filter = build_c1_lt_c2_filter(left.schema().as_ref(), right.schema().as_ref());
+
+    // 100-byte pool: every buffered slice fails its reservation, so each
+    // inner batch of the key group spills separately.
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_limit(100, 1.0)
+        .with_disk_manager_builder(
+            DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
+        )
+        .build_arc()?;
+    let task_ctx = Arc::new(
+        TaskContext::default()
+            .with_session_config(SessionConfig::default().with_batch_size(1))
+            .with_runtime(runtime),
+    );
+
+    let join = SortMergeJoinExec::try_new(
+        left,
+        right,
+        on,
+        Some(filter),
+        LeftSemi,
+        sort_options,
+        NullEquality::NullEqualsNothing,
+    )?;
+    let stream = join.execute(0, task_ctx)?;
+    let batches = common::collect(stream).await?;
+
+    let output_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(
+        output_rows, 1,
+        "left row must match the group's first (spilled) inner slice",
+    );
+
+    let metrics = join.metrics().expect("must have metrics");
+    assert_eq!(
+        metrics.spill_count(),
+        Some(1),
+        "all overflows of one key group must share a single spill file",
+    );
+    assert_eq!(
+        metrics.spilled_rows(),
+        Some(2),
+        "both inner slices of the group must be spilled",
+    );
+    Ok(())
+}
+
 /// Once the inner key group has spilled, an outer key group spanning a batch
 /// boundary must still be evaluated against the spilled inner rows — the
 /// second outer batch's rows must not be treated as having no inner group to

datafusion/sqllogictest/test_files/sort_merge_join_spill.slt

@@ -49,6 +49,13 @@ SELECT 2 AS k,
        lpad(cast(value AS varchar), 512, 'x') AS p
 FROM generate_series(1, 2000);
 
+# One narrow 20,000-row key group for the bitwise semi-join regression. Its
+# 200-row input batches fit in 64 KB, while the complete group does not.
+statement ok
+CREATE VIEW bitwise_wide AS
+SELECT 2 AS k, value AS v
+FROM generate_series(1, 20000);
+
 # Keep output narrow while retaining the payload in the buffered input.
 
 query TT
@@ -106,6 +113,19 @@ SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k
 ----
 6009 values hashing to 126cd87356bc636448b65c5fb5f4bd2b
 
+# Only the first input batch satisfies this filtered semi join.
+query I
+SELECT pr.k
+FROM probe pr
+WHERE EXISTS (
+  SELECT 1
+  FROM bitwise_wide wi
+  WHERE pr.k = wi.k
+    AND wi.v <= pr.x - 300
+)
+----
+2
+
 # A 64 KB pool spills all 10 buffered batches; each result must match its
 # unlimited-memory hash-join reference.
 
@@ -236,6 +256,42 @@ SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k
 ----
 6009 values hashing to 126cd87356bc636448b65c5fb5f4bd2b
 
+# Let the required narrow input sorts merge within the constrained pool.
+statement ok
+SET datafusion.execution.sort_spill_reservation_bytes = 0
+
+# Prove the correlated EXISTS uses the filtered bitwise LeftSemi stream and
+# spills the complete multi-batch key group.
+query TT
+EXPLAIN ANALYZE
+SELECT pr.k
+FROM probe pr
+WHERE EXISTS (
+  SELECT 1
+  FROM bitwise_wide wi
+  WHERE pr.k = wi.k
+    AND wi.v <= pr.x - 300
+)
+----
+Plan with Metrics
+<slt:ignore>SortMergeJoinExec: join_type=LeftSemi, on=[(k@0, k@0)], filter=v@1 <= x@0 - 300, metrics=[output_rows=1,<slt:ignore>spill_count=1, spilled_bytes=<slt:ignore> KB, spilled_rows=<slt:ignore> K, peak_mem_used=<slt:ignore>
+
+# The same query must retain the matching first slice after later overflows.
+query I
+SELECT pr.k
+FROM probe pr
+WHERE EXISTS (
+  SELECT 1
+  FROM bitwise_wide wi
+  WHERE pr.k = wi.k
+    AND wi.v <= pr.x - 300
+)
+----
+2
+
+statement ok
+RESET datafusion.execution.sort_spill_reservation_bytes
+
 statement ok
 RESET datafusion.runtime.memory_limit