Skip to content

Commit 8775702

Browse files
authored
fix: bound sort-merge join spill readback memory (#94)
Merge the standalone SMJ spill-readback memory fix from ff68b05. All 35 other CI checks passed on the reviewed head. Two known non-SMJ CI failures were explicitly accepted for this merge and are deferred: datafusion-cli MinIO image pull access denied, and the pre-existing missing generated array_reduce documentation. No test or gate is disabled or weakened. Autoresearch integration remains deferred.
1 parent 31a96ad commit 8775702

5 files changed

Lines changed: 218 additions & 7 deletions

File tree

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Sort-merge join maintenance
2+
3+
Keep spill-backed payload readback limited to the current output working set.
4+
Preserve backing files until their buffered batches are removed: repeated keys
5+
on the streamed side can revisit an already-read group. Keep memory reservation
6+
growth/shrink balanced, including eviction, dequeuing, errors and cancellation.
7+
8+
When changing restoration or freeze logic, test full-join null output and
9+
deferred filter failures as well as ordinary matches. Do not remove required
10+
restores merely because a batch has no current matched output indices.
11+
12+
Run the SMJ unit tests and the core SMJ memory-limit validation before accepting
13+
spill changes. Preserve the RSS allowance; a lower accounting metric alone is
14+
not evidence that resident memory decreased. Document changes to the readback
15+
lifecycle and its remaining memory limits in README.md.
Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
<!--
2+
Licensed to the Apache Software Foundation (ASF) under one
3+
or more contributor license agreements. See the NOTICE file
4+
distributed with this work for additional information
5+
regarding copyright ownership. The ASF licenses this file
6+
to you under the Apache License, Version 2.0 (the
7+
"License"); you may not use this file except in compliance
8+
with the License. You may obtain a copy of the License at
9+
10+
http://www.apache.org/licenses/LICENSE-2.0
11+
12+
Unless required by applicable law or agreed to in writing,
13+
software distributed under the License is distributed on an
14+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
KIND, either express or implied. See the License for the
16+
specific language governing permissions and limitations
17+
under the License.
18+
-->
19+
20+
# Sort-merge join spill readback
21+
22+
The materializing stream buffers the current group of equal join keys. Payloads
23+
that cannot be admitted to the memory pool spill to temporary files; join keys
24+
and per-row match state remain resident.
25+
26+
Readback retains the original spill file and caches decoded payloads only while
27+
the next output operation needs them. Before restoring another working set, it
28+
evicts decoded payloads no longer needed and releases their reservations. A
29+
later streamed row can revisit the same spilled group without rewriting files
30+
or retaining the whole decoded group. Removing the buffered head also removes
31+
its cached file handle and adjusts the remaining working-set indices.
32+
33+
Full joins restore batches with pending null output, plus the dequeued head for
34+
deferred filter failures; they do not eagerly reload every matched batch.
35+
This bounds accumulation of **visited spill payloads**, not total process RSS:
36+
join keys, match state, output batches and the current readback window still
37+
consume memory, and readback may temporarily exceed the configured pool limit.
38+
Repeated streamed keys can require reading a spill file again after its decoded
39+
payload is evicted. This trades possible extra read I/O for bounded accumulation
40+
of restored payloads; it is not a claim of faster execution for every spilled join.
41+
42+
`spill_restored_batches_do_not_accumulate` checks bounded reservations, exact
43+
spill/non-spill output for inner/left/right/full joins, revisiting the group and
44+
early cancellation. The core SMJ RSS validation tests additionally exercise a
45+
large single-key group through SQL and real spill files.

‎datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs‎

Lines changed: 52 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,11 @@ pub(super) struct MaterializingSortMergeJoinStream {
348348
/// Tracks the number of batches currently spilled
349349
pub spilled_batch_count: usize,
350350

351+
/// Spill-backed batches loaded for the current output window. Keep the
352+
/// original files so a later streamed row can revisit the group without
353+
/// retaining every previously read payload in memory or writing it again.
354+
restored_batches: Vec<(usize, Arc<dyn SpillFile>)>,
355+
351356
/// Time spent doing the join's own work (including spill write and
352357
/// read-back). The clock is stopped while awaiting the child inputs or
353358
/// the consumer taking an emitted batch — see [`Self::stop_join_time`].
@@ -584,6 +589,7 @@ impl MaterializingSortMergeJoinStream {
584589
runtime_env,
585590
spill_manager,
586591
spilled_batch_count: 0,
592+
restored_batches: vec![],
587593
join_time,
588594
join_time_start: None,
589595
streamed_buffered_cmp: None,
@@ -949,7 +955,7 @@ impl MaterializingSortMergeJoinStream {
949955
fn get_required_batch_indices(&self, buffered_freeze_count: usize) -> Vec<usize> {
950956
let mut needed = vec![];
951957
// Avoid scanning if no spilled batches exist
952-
if self.spilled_batch_count == 0 {
958+
if self.spilled_batch_count == 0 && self.restored_batches.is_empty() {
953959
return needed;
954960
}
955961
// We need all batches that matched with streamed rows
@@ -959,9 +965,20 @@ impl MaterializingSortMergeJoinStream {
959965
}
960966
}
961967

962-
// Full Joins need to emit null-joined rows, so we need batches up to freeze_count
968+
// Only batches with pending null output need to be restored here.
969+
// Restoring every batch in a Full join would load the entire key group
970+
// even when all its rows matched and none needs null materialization.
963971
if self.join_type == JoinType::Full {
964-
needed.extend(0..buffered_freeze_count);
972+
needed.extend(
973+
self.buffered_data
974+
.batches
975+
.iter()
976+
.take(buffered_freeze_count)
977+
.enumerate()
978+
.filter_map(|(idx, batch)| {
979+
(!batch.null_joined.is_empty()).then_some(idx)
980+
}),
981+
);
965982
}
966983

967984
needed.sort_unstable();
@@ -975,6 +992,23 @@ impl MaterializingSortMergeJoinStream {
975992
&mut self,
976993
required_indices: &[usize],
977994
) -> Result<()> {
995+
// Drop decoded payloads outside this output's working set before
996+
// reading more. Join arrays and per-row match state remain resident.
997+
// Only visit the previous working set, not the entire buffered group.
998+
let mut retained = Vec::new();
999+
for (idx, file) in std::mem::take(&mut self.restored_batches) {
1000+
if required_indices.contains(&idx) {
1001+
retained.push((idx, file));
1002+
} else {
1003+
let batch = &mut self.buffered_data.batches[idx];
1004+
batch.batch = BufferedBatchState::Spilled(file);
1005+
self.reservation
1006+
.shrink(batch.reserved_amount - batch.join_arrays_mem);
1007+
batch.reserved_amount = batch.join_arrays_mem;
1008+
self.spilled_batch_count += 1;
1009+
}
1010+
}
1011+
self.restored_batches = retained;
9781012
for &idx in required_indices {
9791013
// Guard against indices that might be out of bounds if the queue was cleared
9801014
if idx >= self.buffered_data.batches.len() {
@@ -990,6 +1024,7 @@ impl MaterializingSortMergeJoinStream {
9901024

9911025
match spill_stream.next().await.transpose()? {
9921026
Some(batch) => {
1027+
self.restored_batches.push((idx, Arc::clone(spill_file)));
9931028
// Transition the batch back to InMemory
9941029
bb.batch = BufferedBatchState::InMemory(batch);
9951030
self.spilled_batch_count -= 1;
@@ -1200,7 +1235,12 @@ impl MaterializingSortMergeJoinStream {
12001235
break;
12011236
}
12021237
// load the spilled head batch before dequeuing
1203-
let needed = self.get_required_batch_indices(1);
1238+
let mut needed = self.get_required_batch_indices(1);
1239+
if self.join_type == JoinType::Full && !needed.contains(&0) {
1240+
// Deferred filter failures can create null output as the head
1241+
// is removed, even if it had no earlier null-joined indices.
1242+
needed.push(0);
1243+
}
12041244
self.restore_spilled_batches(&needed).await?;
12051245

12061246
self.freeze_dequeuing_buffered()?;
@@ -1210,6 +1250,14 @@ impl MaterializingSortMergeJoinStream {
12101250
if matches!(buffered_batch.batch, BufferedBatchState::Spilled(_)) {
12111251
self.spilled_batch_count -= 1;
12121252
}
1253+
self.restored_batches.retain_mut(|(idx, _)| {
1254+
if *idx == 0 {
1255+
false
1256+
} else {
1257+
*idx -= 1;
1258+
true
1259+
}
1260+
});
12131261
head_changed = true;
12141262
}
12151263
}

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

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2512,6 +2512,105 @@ async fn overallocation_multi_batch_spill() -> Result<()> {
25122512
Ok(())
25132513
}
25142514

2515+
/// Reading a spilled key group must not progressively restore the entire group
2516+
/// into memory. Two streamed rows revisit the same group, so evicting a restored
2517+
/// batch must preserve its spill file and exact output on the second visit.
2518+
#[tokio::test]
2519+
async fn spill_restored_batches_do_not_accumulate() -> Result<()> {
2520+
const BATCH_ROWS: usize = 1024;
2521+
const BATCHES: usize = 64;
2522+
const MEMORY_LIMIT: usize = 32 * 1024;
2523+
2524+
for join_type in [Inner, Left, Right, Full] {
2525+
let buffered_batches: Vec<RecordBatch> = (0..BATCHES)
2526+
.map(|batch| {
2527+
let values: Vec<i32> = (0..BATCH_ROWS)
2528+
.map(|row| (batch * BATCH_ROWS + row) as i32)
2529+
.collect();
2530+
build_table_i32(
2531+
("a", &values),
2532+
("k", &vec![1; BATCH_ROWS]),
2533+
("v", &values),
2534+
)
2535+
})
2536+
.collect();
2537+
let key_bytes = buffered_batches
2538+
.iter()
2539+
.map(|batch| batch.column(1).get_array_memory_size())
2540+
.sum::<usize>();
2541+
let batch_bytes = buffered_batches[0].get_array_memory_size()
2542+
+ buffered_batches[0].column(1).get_array_memory_size()
2543+
+ BATCH_ROWS * size_of::<usize>()
2544+
+ size_of::<std::ops::Range<usize>>()
2545+
+ size_of::<usize>();
2546+
let buffered = build_table_from_batches(buffered_batches);
2547+
let streamed = build_table_from_batches(vec![build_table_i32(
2548+
("a", &vec![0, 1]),
2549+
("k", &vec![1, 1]),
2550+
("v", &vec![-1, -2]),
2551+
)]);
2552+
let (left, right) = if join_type == Right {
2553+
(buffered, streamed)
2554+
} else {
2555+
(streamed, buffered)
2556+
};
2557+
let on = vec![(
2558+
Arc::new(Column::new_with_schema("k", &left.schema())?) as _,
2559+
Arc::new(Column::new_with_schema("k", &right.schema())?) as _,
2560+
)];
2561+
let make_join = || {
2562+
join_with_options(
2563+
Arc::clone(&left),
2564+
Arc::clone(&right),
2565+
on.clone(),
2566+
join_type,
2567+
vec![SortOptions::default()],
2568+
NullEquality::NullEqualsNothing,
2569+
)
2570+
};
2571+
let config = SessionConfig::default().with_batch_size(BATCH_ROWS);
2572+
let unlimited =
2573+
Arc::new(TaskContext::default().with_session_config(config.clone()));
2574+
let expected = common::collect(make_join()?.execute(0, unlimited)?).await?;
2575+
assert_eq!(
2576+
expected.iter().map(RecordBatch::num_rows).sum::<usize>(),
2577+
2 * BATCHES * BATCH_ROWS,
2578+
);
2579+
let runtime = RuntimeEnvBuilder::new()
2580+
.with_memory_limit(MEMORY_LIMIT, 1.0)
2581+
.with_disk_manager_builder(
2582+
DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
2583+
)
2584+
.build_arc()?;
2585+
let context = Arc::new(
2586+
TaskContext::default()
2587+
.with_session_config(config)
2588+
.with_runtime(Arc::clone(&runtime)),
2589+
);
2590+
let join = make_join()?;
2591+
let actual = common::collect(join.execute(0, Arc::clone(&context))?).await?;
2592+
assert_eq!(actual, expected, "{join_type:?} spilled output differs");
2593+
let metrics = join.metrics().unwrap();
2594+
assert!(metrics.spill_count().unwrap() > 0);
2595+
let peak = metrics.sum_by_name("peak_mem_used").unwrap().as_usize();
2596+
// Keys remain resident by design. Only the current output's source
2597+
// batches, not all visited payloads, may be restored above the pool.
2598+
let bound = MEMORY_LIMIT + key_bytes + 3 * batch_bytes;
2599+
assert!(
2600+
peak <= bound,
2601+
"{join_type:?}: peak {peak} > bounded restore {bound}"
2602+
);
2603+
assert_eq!(runtime.memory_pool.reserved(), 0);
2604+
2605+
// Cancellation after the first output also releases restored buffers.
2606+
let mut partial = join.execute(0, context)?;
2607+
assert!(partial.next().await.transpose()?.is_some());
2608+
drop(partial);
2609+
assert_eq!(runtime.memory_pool.reserved(), 0);
2610+
}
2611+
Ok(())
2612+
}
2613+
25152614
/// Verifies that `peak_mem_used` reflects join_arrays memory on the spill path.
25162615
///
25172616
/// Uses a memory limit smaller than a single batch's `size_estimation` so that

‎datafusion/sqllogictest/test_files/aggregate.slt‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8070,11 +8070,15 @@ select max(v1), * exclude (v1, v2) from having_test having max(v1) = 3
80708070
----
80718071
3
80728072

8073-
# because v1, v2 is not in the group by clause, the sql is invalid
8074-
query III
8073+
# REPLACE targets must exist in the wildcard input, even with grouping/HAVING.
8074+
query error Column 'v3' specified in REPLACE does not exist
80758075
select max(v1), * replace ('v1' as v3) from having_test group by v1, v2 having max(v1) = 3
8076+
8077+
# Replacing an existing grouped column remains valid.
8078+
query III
8079+
select max(v1), * replace (v2 + 1 as v2) from having_test group by v1, v2 having max(v1) = 3
80768080
----
8077-
3 3 4
8081+
3 3 5
80788082

80798083
query III
80808084
select max(v1), t.* from having_test t group by v1, v2 having max(v1) = 3

0 commit comments

Comments
 (0)