From ff68b05d554c47e205e541c8287b67577f48302d Mon Sep 17 00:00:00 2001 From: Denys Tsomenko Date: Tue, 15 Sep 2026 13:37:42 +0300 Subject: [PATCH] fix: bound sort-merge join spill readback memory Retain spill backing files and evict decoded payloads outside the next output working set. Preserve Full-join null output and deferred filter failures. Add exact spill/no-spill, memory-bound, revisit and cancellation coverage. Include the independent stale aggregate.slt REPLACE fixture repair needed by the full workspace gate; SQL planner behavior and the original SMJ RSS allowance remain unchanged. Standalone extended workspace, all-feature Clippy, full rust_lint.sh, and three original RSS guard repeats pass. The separate autoresearch runtime optimizations are excluded. --- .../src/joins/sort_merge_join/AGENTS.md | 15 +++ .../src/joins/sort_merge_join/README.md | 45 +++++++++ .../sort_merge_join/materializing_stream.rs | 56 ++++++++++- .../src/joins/sort_merge_join/tests.rs | 99 +++++++++++++++++++ .../sqllogictest/test_files/aggregate.slt | 10 +- 5 files changed, 218 insertions(+), 7 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/sort_merge_join/AGENTS.md create mode 100644 datafusion/physical-plan/src/joins/sort_merge_join/README.md diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/AGENTS.md b/datafusion/physical-plan/src/joins/sort_merge_join/AGENTS.md new file mode 100644 index 0000000000000..46da23f7982c9 --- /dev/null +++ b/datafusion/physical-plan/src/joins/sort_merge_join/AGENTS.md @@ -0,0 +1,15 @@ +# Sort-merge join maintenance + +Keep spill-backed payload readback limited to the current output working set. +Preserve backing files until their buffered batches are removed: repeated keys +on the streamed side can revisit an already-read group. Keep memory reservation +growth/shrink balanced, including eviction, dequeuing, errors and cancellation. + +When changing restoration or freeze logic, test full-join null output and +deferred filter failures as well as ordinary matches. Do not remove required +restores merely because a batch has no current matched output indices. + +Run the SMJ unit tests and the core SMJ memory-limit validation before accepting +spill changes. Preserve the RSS allowance; a lower accounting metric alone is +not evidence that resident memory decreased. Document changes to the readback +lifecycle and its remaining memory limits in README.md. diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/README.md b/datafusion/physical-plan/src/joins/sort_merge_join/README.md new file mode 100644 index 0000000000000..0021e4024db8a --- /dev/null +++ b/datafusion/physical-plan/src/joins/sort_merge_join/README.md @@ -0,0 +1,45 @@ + + +# Sort-merge join spill readback + +The materializing stream buffers the current group of equal join keys. Payloads +that cannot be admitted to the memory pool spill to temporary files; join keys +and per-row match state remain resident. + +Readback retains the original spill file and caches decoded payloads only while +the next output operation needs them. Before restoring another working set, it +evicts decoded payloads no longer needed and releases their reservations. A +later streamed row can revisit the same spilled group without rewriting files +or retaining the whole decoded group. Removing the buffered head also removes +its cached file handle and adjusts the remaining working-set indices. + +Full joins restore batches with pending null output, plus the dequeued head for +deferred filter failures; they do not eagerly reload every matched batch. +This bounds accumulation of **visited spill payloads**, not total process RSS: +join keys, match state, output batches and the current readback window still +consume memory, and readback may temporarily exceed the configured pool limit. +Repeated streamed keys can require reading a spill file again after its decoded +payload is evicted. This trades possible extra read I/O for bounded accumulation +of restored payloads; it is not a claim of faster execution for every spilled join. + +`spill_restored_batches_do_not_accumulate` checks bounded reservations, exact +spill/non-spill output for inner/left/right/full joins, revisiting the group and +early cancellation. The core SMJ RSS validation tests additionally exercise a +large single-key group through SQL and real spill files. diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs index 3baa0c4a3e792..37cedf99fab4b 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs @@ -348,6 +348,11 @@ pub(super) struct MaterializingSortMergeJoinStream { /// Tracks the number of batches currently spilled pub spilled_batch_count: usize, + /// Spill-backed batches loaded for the current output window. Keep the + /// original files so a later streamed row can revisit the group without + /// retaining every previously read payload in memory or writing it again. + restored_batches: Vec<(usize, Arc)>, + /// Time spent doing the join's own work (including spill write and /// read-back). The clock is stopped while awaiting the child inputs or /// the consumer taking an emitted batch — see [`Self::stop_join_time`]. @@ -584,6 +589,7 @@ impl MaterializingSortMergeJoinStream { runtime_env, spill_manager, spilled_batch_count: 0, + restored_batches: vec![], join_time, join_time_start: None, streamed_buffered_cmp: None, @@ -949,7 +955,7 @@ impl MaterializingSortMergeJoinStream { fn get_required_batch_indices(&self, buffered_freeze_count: usize) -> Vec { let mut needed = vec![]; // Avoid scanning if no spilled batches exist - if self.spilled_batch_count == 0 { + if self.spilled_batch_count == 0 && self.restored_batches.is_empty() { return needed; } // We need all batches that matched with streamed rows @@ -959,9 +965,20 @@ impl MaterializingSortMergeJoinStream { } } - // Full Joins need to emit null-joined rows, so we need batches up to freeze_count + // Only batches with pending null output need to be restored here. + // Restoring every batch in a Full join would load the entire key group + // even when all its rows matched and none needs null materialization. if self.join_type == JoinType::Full { - needed.extend(0..buffered_freeze_count); + needed.extend( + self.buffered_data + .batches + .iter() + .take(buffered_freeze_count) + .enumerate() + .filter_map(|(idx, batch)| { + (!batch.null_joined.is_empty()).then_some(idx) + }), + ); } needed.sort_unstable(); @@ -975,6 +992,23 @@ impl MaterializingSortMergeJoinStream { &mut self, required_indices: &[usize], ) -> Result<()> { + // Drop decoded payloads outside this output's working set before + // reading more. Join arrays and per-row match state remain resident. + // Only visit the previous working set, not the entire buffered group. + let mut retained = Vec::new(); + for (idx, file) in std::mem::take(&mut self.restored_batches) { + if required_indices.contains(&idx) { + retained.push((idx, file)); + } else { + let batch = &mut self.buffered_data.batches[idx]; + batch.batch = BufferedBatchState::Spilled(file); + self.reservation + .shrink(batch.reserved_amount - batch.join_arrays_mem); + batch.reserved_amount = batch.join_arrays_mem; + self.spilled_batch_count += 1; + } + } + self.restored_batches = retained; for &idx in required_indices { // Guard against indices that might be out of bounds if the queue was cleared if idx >= self.buffered_data.batches.len() { @@ -990,6 +1024,7 @@ impl MaterializingSortMergeJoinStream { match spill_stream.next().await.transpose()? { Some(batch) => { + self.restored_batches.push((idx, Arc::clone(spill_file))); // Transition the batch back to InMemory bb.batch = BufferedBatchState::InMemory(batch); self.spilled_batch_count -= 1; @@ -1200,7 +1235,12 @@ impl MaterializingSortMergeJoinStream { break; } // load the spilled head batch before dequeuing - let needed = self.get_required_batch_indices(1); + let mut needed = self.get_required_batch_indices(1); + if self.join_type == JoinType::Full && !needed.contains(&0) { + // Deferred filter failures can create null output as the head + // is removed, even if it had no earlier null-joined indices. + needed.push(0); + } self.restore_spilled_batches(&needed).await?; self.freeze_dequeuing_buffered()?; @@ -1210,6 +1250,14 @@ impl MaterializingSortMergeJoinStream { if matches!(buffered_batch.batch, BufferedBatchState::Spilled(_)) { self.spilled_batch_count -= 1; } + self.restored_batches.retain_mut(|(idx, _)| { + if *idx == 0 { + false + } else { + *idx -= 1; + true + } + }); head_changed = true; } } 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 175a9c0ea7198..6fb25e79b53e3 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -2512,6 +2512,105 @@ async fn overallocation_multi_batch_spill() -> Result<()> { Ok(()) } +/// Reading a spilled key group must not progressively restore the entire group +/// into memory. Two streamed rows revisit the same group, so evicting a restored +/// batch must preserve its spill file and exact output on the second visit. +#[tokio::test] +async fn spill_restored_batches_do_not_accumulate() -> Result<()> { + const BATCH_ROWS: usize = 1024; + const BATCHES: usize = 64; + const MEMORY_LIMIT: usize = 32 * 1024; + + for join_type in [Inner, Left, Right, Full] { + let buffered_batches: Vec = (0..BATCHES) + .map(|batch| { + let values: Vec = (0..BATCH_ROWS) + .map(|row| (batch * BATCH_ROWS + row) as i32) + .collect(); + build_table_i32( + ("a", &values), + ("k", &vec![1; BATCH_ROWS]), + ("v", &values), + ) + }) + .collect(); + let key_bytes = buffered_batches + .iter() + .map(|batch| batch.column(1).get_array_memory_size()) + .sum::(); + let batch_bytes = buffered_batches[0].get_array_memory_size() + + buffered_batches[0].column(1).get_array_memory_size() + + BATCH_ROWS * size_of::() + + size_of::>() + + size_of::(); + let buffered = build_table_from_batches(buffered_batches); + let streamed = build_table_from_batches(vec![build_table_i32( + ("a", &vec![0, 1]), + ("k", &vec![1, 1]), + ("v", &vec![-1, -2]), + )]); + let (left, right) = if join_type == Right { + (buffered, streamed) + } else { + (streamed, buffered) + }; + let on = vec![( + Arc::new(Column::new_with_schema("k", &left.schema())?) as _, + Arc::new(Column::new_with_schema("k", &right.schema())?) as _, + )]; + let make_join = || { + join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + ) + }; + let config = SessionConfig::default().with_batch_size(BATCH_ROWS); + let unlimited = + Arc::new(TaskContext::default().with_session_config(config.clone())); + let expected = common::collect(make_join()?.execute(0, unlimited)?).await?; + assert_eq!( + expected.iter().map(RecordBatch::num_rows).sum::(), + 2 * BATCHES * BATCH_ROWS, + ); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(MEMORY_LIMIT, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime(Arc::clone(&runtime)), + ); + let join = make_join()?; + let actual = common::collect(join.execute(0, Arc::clone(&context))?).await?; + assert_eq!(actual, expected, "{join_type:?} spilled output differs"); + let metrics = join.metrics().unwrap(); + assert!(metrics.spill_count().unwrap() > 0); + let peak = metrics.sum_by_name("peak_mem_used").unwrap().as_usize(); + // Keys remain resident by design. Only the current output's source + // batches, not all visited payloads, may be restored above the pool. + let bound = MEMORY_LIMIT + key_bytes + 3 * batch_bytes; + assert!( + peak <= bound, + "{join_type:?}: peak {peak} > bounded restore {bound}" + ); + assert_eq!(runtime.memory_pool.reserved(), 0); + + // Cancellation after the first output also releases restored buffers. + let mut partial = join.execute(0, context)?; + assert!(partial.next().await.transpose()?.is_some()); + drop(partial); + assert_eq!(runtime.memory_pool.reserved(), 0); + } + Ok(()) +} + /// Verifies that `peak_mem_used` reflects join_arrays memory on the spill path. /// /// Uses a memory limit smaller than a single batch's `size_estimation` so that diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 460d4cd2ffda3..119f3c649239a 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -8070,11 +8070,15 @@ select max(v1), * exclude (v1, v2) from having_test having max(v1) = 3 ---- 3 -# because v1, v2 is not in the group by clause, the sql is invalid -query III +# REPLACE targets must exist in the wildcard input, even with grouping/HAVING. +query error Column 'v3' specified in REPLACE does not exist select max(v1), * replace ('v1' as v3) from having_test group by v1, v2 having max(v1) = 3 + +# Replacing an existing grouped column remains valid. +query III +select max(v1), * replace (v2 + 1 as v2) from having_test group by v1, v2 having max(v1) = 3 ---- -3 3 4 +3 3 5 query III select max(v1), t.* from having_test t group by v1, v2 having max(v1) = 3