fix: bound sort-merge join spill readback memory

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.

datafusion/physical-plan/src/joins/sort_merge_join/README.md

@@ -0,0 +1,45 @@
+<!--
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements. See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership. The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing,
+software distributed under the License is distributed on an
+"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied. See the License for the
+specific language governing permissions and limitations
+under the License.
+-->
+
+# 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.

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<dyn SpillFile>)>,
+
     /// 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<usize> {
         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;
             }
         }

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<RecordBatch> = (0..BATCHES)
+            .map(|batch| {
+                let values: Vec<i32> = (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::<usize>();
+        let batch_bytes = buffered_batches[0].get_array_memory_size()
+            + buffered_batches[0].column(1).get_array_memory_size()
+            + BATCH_ROWS * size_of::<usize>()
+            + size_of::<std::ops::Range<usize>>()
+            + size_of::<usize>();
+        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::<usize>(),
+            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

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