This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24056-62650eff81847928fadaf1d5e8ba2303cb13b819 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit b068aa649e18fd136deaec4e89b33dd9ad4eca67 Author: Burak Şen <[email protected]> AuthorDate: Mon Aug 3 17:27:08 2026 +0300 fix: keep every spilled slice of a sort-merge join inner key group (#24056) ## Which issue does this PR close? - No issue. Found while adding test coverage in scope of #13431. ## Rationale for this change Filtered semi, anti, and mark sort-merge joins kept only one finished spill file per inner key group. If a group spilled more than once, each spill replaced the previous file, so read-back could lose earlier rows and produce incorrect results. ## What changes are included in this PR? Keep one `InProgressSpillFile` open per inner key group, append every overflow to it, and finalize it before filter evaluation. ## Are these changes tested? Yes: ``` cargo test -p datafusion-physical-plan bitwise_multi_spill_inner_key_group cargo test -p datafusion-physical-plan cargo test --test sqllogictests -- sort_merge_join_spill ``` ## Are there any user-facing changes? no --- .../src/joins/sort_merge_join/bitwise_stream.rs | 89 +++++++++++++--------- .../src/joins/sort_merge_join/tests.rs | 71 +++++++++++++++++ .../test_files/sort_merge_join_spill.slt | 56 ++++++++++++++ 3 files changed, 182 insertions(+), 34 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 3716ecde28..1b90f24b96 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,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,7 +689,7 @@ impl BitwiseSortMergeJoinStream { if !self.next_inner_batch().await? { self.inner_batch = None; - return Ok(()); + break; } if !keys_match( &saved_inner_keys, @@ -690,20 +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<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. 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 3dbb50eba0..175a9c0ea7 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 diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt index 1a3dcafa60..69bb718bd8 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 +<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 --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
