From 1890709cb19455e719bbcbf08a77ec610ef9272a Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sun, 2 Aug 2026 17:14:28 +0300 Subject: [PATCH 1/3] fix: keep every spilled slice of a sort-merge join inner key group BitwiseSortMergeJoinStream kept only a single finished spill file per inner key group. When one group spanned several inner batches and the memory reservation failed more than once, each new spill replaced the previous file, silently dropping its rows: semi joins lost matches, anti joins returned extra rows, and mark joins produced wrong marks under memory pressure. Keep one InProgressSpillFile open per key group instead: every overflow appends the freshly buffered slices to the same file, and the writer is finalized for read-back when evaluation of the group starts (buffering is complete by then). A group therefore produces exactly one spill file that holds all of its spilled rows. --- .../joins/sort_merge_join/bitwise_stream.rs | 42 +++++++---- .../src/joins/sort_merge_join/tests.rs | 71 +++++++++++++++++++ 2 files changed, 101 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index 3716ecde284c5..4f1f0a2ebba04 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/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,13 @@ 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. Spilled to disk when memory reservation fails; one group + // can overflow the pool repeatedly, and every overflow appends to the + // same open spill file. The writer is finalized into `inner_key_spill` + // when evaluation starts (buffering for the group is complete by then). inner_key_buffer: Vec, inner_key_spill: Option>, + inner_key_spill_writer: Option, // Join ON expressions, evaluated against each new batch to produce // the key arrays used for sorted key comparisons. @@ -339,6 +344,7 @@ impl BitwiseSortMergeJoinStream { inner_key_arrays: vec![], matched: BooleanBufferBuilder::new(0), inner_key_buffer: vec![], + inner_key_spill_writer: None, inner_key_spill: None, on_outer, on_inner, @@ -443,18 +449,21 @@ impl BitwiseSortMergeJoinStream { Ok(self.inner_self_cmp.as_ref().unwrap()) } - /// Spill the in-memory inner key buffer to disk and clear it. + /// Spill the in-memory inner key buffer to disk and clear it. One key + /// group can spill repeatedly; every call appends to the group's single + /// open spill file. 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(); + if self.inner_key_spill_writer.is_none() { + self.inner_key_spill_writer = Some( + self.spill_manager + .create_in_progress_file("semi_anti_smj_inner_key_spill")?, + ); + } + let writer = self.inner_key_spill_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,6 +474,7 @@ impl BitwiseSortMergeJoinStream { /// pool interactions (see apache/datafusion#20729). fn clear_inner_key_group(&mut self) { self.inner_key_buffer.clear(); + self.inner_key_spill_writer = None; self.inner_key_spill = None; self.inner_buffer_size = 0; } @@ -703,7 +713,9 @@ impl BitwiseSortMergeJoinStream { // 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() + || self.inner_key_spill_writer.is_some() + || self.inner_key_spill.is_some(), "process_key_match_with_filter called with no inner key data" ); debug_assert!( @@ -732,6 +744,12 @@ impl BitwiseSortMergeJoinStream { ) .count_ones(); + // Evaluation starting means buffering for this group is complete, so + // the group's spill file (if any) can be finalized for reading. + if let Some(mut writer) = self.inner_key_spill_writer.take() { + self.inner_key_spill = writer.finish()?; + } + // Process spilled inner batches first asynchronously. if matched_count < outer_group_len && let Some(spill_file) = &self.inner_key_spill 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 3dbb50eba07d9..175a9c0ea7198 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/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 From 0b0df63641ab1682f595182149f2c8e308f210df Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sun, 2 Aug 2026 19:28:24 +0300 Subject: [PATCH 2/3] add slt tests --- .../test_files/sort_merge_join_spill.slt | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt index 1a3dcafa60f82..69bb718bd8c1f 100644 --- a/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt +++ b/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 +SortMergeJoinExec: join_type=LeftSemi, on=[(k@0, k@0)], filter=v@1 <= x@0 - 300, metrics=[output_rows=1,spill_count=1, spilled_bytes= KB, spilled_rows= K, peak_mem_used= + +# 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 From 626f32506621251f89b9069dc2b315cd68476300 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sun, 2 Aug 2026 22:29:40 +0300 Subject: [PATCH 3/3] delete fields to simplify --- .../joins/sort_merge_join/bitwise_stream.rs | 89 ++++++++++--------- 1 file changed, 46 insertions(+), 43 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index 4f1f0a2ebba04..1b90f24b96acc 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -235,13 +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; one group - // can overflow the pool repeatedly, and every overflow appends to the - // same open spill file. The writer is finalized into `inner_key_spill` - // when evaluation starts (buffering for the group is complete by then). + // 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, - inner_key_spill: Option>, - inner_key_spill_writer: Option, // Join ON expressions, evaluated against each new batch to produce // the key arrays used for sorted key comparisons. @@ -344,8 +340,6 @@ impl BitwiseSortMergeJoinStream { inner_key_arrays: vec![], matched: BooleanBufferBuilder::new(0), inner_key_buffer: vec![], - inner_key_spill_writer: None, - inner_key_spill: None, on_outer, on_inner, filter, @@ -450,16 +444,19 @@ impl BitwiseSortMergeJoinStream { } /// Spill the in-memory inner key buffer to disk and clear it. One key - /// group can spill repeatedly; every call appends to the group's single - /// open spill file. - fn spill_inner_key_buffer(&mut self) -> Result<()> { - if self.inner_key_spill_writer.is_none() { - self.inner_key_spill_writer = Some( + /// 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, + ) -> Result<()> { + if writer.is_none() { + *writer = Some( self.spill_manager .create_in_progress_file("semi_anti_smj_inner_key_spill")?, ); } - let writer = self.inner_key_spill_writer.as_mut().unwrap(); + let writer = writer.as_mut().unwrap(); for batch in self.inner_key_buffer.drain(..) { writer.append_batch(&batch)?; } @@ -474,8 +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_writer = None; - self.inner_key_spill = None; self.inner_buffer_size = 0; } @@ -649,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>> { self.clear_inner_key_group(); + let mut writer: Option = 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 = @@ -670,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| { @@ -683,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 @@ -692,7 +689,7 @@ impl BitwiseSortMergeJoinStream { if !self.next_inner_batch().await? { self.inner_batch = None; - return Ok(()); + break; } if !keys_match( &saved_inner_keys, @@ -700,22 +697,30 @@ impl BitwiseSortMergeJoinStream { &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>, + ) -> 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_writer.is_some() - || 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!( @@ -744,15 +749,9 @@ impl BitwiseSortMergeJoinStream { ) .count_ones(); - // Evaluation starting means buffering for this group is complete, so - // the group's spill file (if any) can be finalized for reading. - if let Some(mut writer) = self.inner_key_spill_writer.take() { - self.inner_key_spill = writer.finish()?; - } - // 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 @@ -817,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>, + ) -> 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() { @@ -890,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.