You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/06/16 11:54:29 UTC
(datafusion) branch gh-readonly-queue/main/pr-22230-ae5f3f51736359beb64ef4aaba9345d8cc1fc853 created (now fa271ce8fd)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22230-ae5f3f51736359beb64ef4aaba9345d8cc1fc853
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at fa271ce8fd refactor: Update SortMergeJoin to use async spill abstractions (#22230)
This branch includes the following new commits:
new fa271ce8fd refactor: Update SortMergeJoin to use async spill abstractions (#22230)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: refactor: Update SortMergeJoin to use async spill abstractions (#22230)
Posted by gi...@apache.org.
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-22230-ae5f3f51736359beb64ef4aaba9345d8cc1fc853
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit fa271ce8fdb246321a8606449a09b4c03b405388
Author: pantShrey <12...@users.noreply.github.com>
AuthorDate: Tue Jun 16 17:24:06 2026 +0530
refactor: Update SortMergeJoin to use async spill abstractions (#22230)
~~## Note: This PR depends on #21882 (pluggable SpillFile trait) and
cannot be merged before it. Opening in parallel per @alamb's suggestion
for easier review. The required SpillFile trait used here is defined in
that base PR.To review locally, apply #21882 first and then stack this
branch on top.~~
**Update:** This PR has been rebased to use the existing
`RefCountedTempFile` and is now completely standalone. It can be
reviewed and merged independently
## Which issue does this PR close?
- Contributes to #21215 (and is required by #21882)
## Rationale for this change
`materializing_stream.rs` and `bitwise_stream.rs` were reading spilled
batches via `open_sync_reader` / direct `File::open` calls ~~, bypassing
the `SpillFile` abstraction introduced in #21882~~. This PR migrates
both to use `SpillManager::read_spill_as_stream`. This safely converts
the SMJ to an async I/O path, preparing the ground for custom backends
(Postgres BufFile, object storage) to handle spill reads without
requiring an OS file path.
## What changes are included in this PR?
- `materializing_stream.rs`: Eagerly restores spilled `BufferedBatches`
via async streams before freezing, avoiding new state machine variants.
- `bitwise_stream.rs`: Replaces sync reads with an async
`poll_next_unpin` loop, caching the stream to survive `Poll::Pending`.
~~- `spill_file.rs`: Removes `open_sync_reader` from the `SpillFile`
trait (no longer needed).~~
## Are these changes tested?
Covered by existing SMJ tests. No new tests added, the behavioral change
is internal (sync → async IO path).
## Are there any user-facing changes?
No.
~~Removes `open_sync_reader` from the SpillFile trait, this is a
breaking API change for anyone implementing the trait, but the trait was
introduced in #21882 which has not merged yet so there are no external
implementors.~~
---------
Co-authored-by: Kumar Ujjawal <uj...@gmail.com>
---
.../src/joins/sort_merge_join/bitwise_stream.rs | 228 ++++++++++-------
.../joins/sort_merge_join/materializing_stream.rs | 282 ++++++++++++++-------
.../src/joins/sort_merge_join/tests.rs | 18 ++
3 files changed, 344 insertions(+), 184 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 ad7312426b..99aef6ed82 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
@@ -119,8 +119,6 @@
//! factor than the pair-materialization approach.
use std::cmp::Ordering;
-use std::fs::File;
-use std::io::BufReader;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -134,7 +132,6 @@ use crate::{EmptyRecordBatchStream, RecordBatchStream};
use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch};
use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not};
use arrow::datatypes::SchemaRef;
-use arrow::ipc::reader::StreamReader;
use arrow::util::bit_chunk_iterator::UnalignedBitChunk;
use arrow::util::bit_util::apply_bitwise_binary_op;
use datafusion_common::{
@@ -259,6 +256,13 @@ pub(crate) struct BitwiseSortMergeJoinStream {
inner_key_buffer: Vec<RecordBatch>,
inner_key_spill: Option<RefCountedTempFile>,
+ // Track the active spill_stream
+ spill_stream: Option<SendableRecordBatchStream>,
+ // Whether the active spill stream has produced any batches yet.
+ spill_stream_has_data: bool,
+ // Prevents wiping out the buffer if we yield while evaluating the filter
+ inner_group_buffered: bool,
+
// True when buffer_inner_key_group returned Pending after partially
// filling inner_key_buffer. On re-entry, buffer_inner_key_group
// must skip clear() and resume from poll_next_inner_batch (the
@@ -371,6 +375,9 @@ impl BitwiseSortMergeJoinStream {
matched: BooleanBufferBuilder::new(0),
inner_key_buffer: vec![],
inner_key_spill: None,
+ spill_stream: None,
+ spill_stream_has_data: false,
+ inner_group_buffered: false,
buffering_inner_pending: false,
pending_boundary: None,
on_outer,
@@ -468,6 +475,9 @@ impl BitwiseSortMergeJoinStream {
fn clear_inner_key_group(&mut self) {
self.inner_key_buffer.clear();
self.inner_key_spill = None;
+ self.spill_stream = None;
+ self.spill_stream_has_data = false;
+ self.inner_group_buffered = false;
self.inner_buffer_size = 0;
}
@@ -749,7 +759,10 @@ impl BitwiseSortMergeJoinStream {
/// 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.
- fn process_key_match_with_filter(&mut self) -> Result<()> {
+ fn process_key_match_with_filter(
+ &mut self,
+ cx: &mut Context<'_>,
+ ) -> Poll<Result<()>> {
self.get_outer_self_cmp()?;
let filter = self.filter.as_ref().unwrap();
let outer_batch = self.outer_batch.as_ref().unwrap();
@@ -785,24 +798,47 @@ impl BitwiseSortMergeJoinStream {
)
.count_ones();
- // Process spilled inner batches first (read back from disk).
- if let Some(spill_file) = &self.inner_key_spill {
- let file = BufReader::new(File::open(spill_file.path())?);
- let reader = StreamReader::try_new(file, None)?;
- for batch_result in reader {
- let inner_slice = batch_result?;
- matched_count = eval_filter_for_inner_slice(
- self.outer_is_left,
- filter,
- &outer_slice,
- &inner_slice,
- &mut self.matched,
- self.outer_offset,
- outer_group_len,
- matched_count,
- )?;
- if matched_count == outer_group_len {
- break;
+ // Process spilled inner batches first asynchronously.
+ if matched_count < outer_group_len
+ && (self.inner_key_spill.is_some() || self.spill_stream.is_some())
+ {
+ if self.spill_stream.is_none()
+ && let Some(spill_file) = &self.inner_key_spill
+ {
+ let stream = self
+ .spill_manager
+ .read_spill_as_stream(spill_file.clone(), None)?;
+ self.spill_stream = Some(stream);
+ }
+
+ while matched_count < outer_group_len {
+ let stream = self.spill_stream.as_mut().unwrap();
+ match ready!(stream.poll_next_unpin(cx)) {
+ Some(Ok(inner_slice)) => {
+ self.spill_stream_has_data = true;
+ matched_count = eval_filter_for_inner_slice(
+ self.outer_is_left,
+ filter,
+ &outer_slice,
+ &inner_slice,
+ &mut self.matched,
+ self.outer_offset,
+ outer_group_len,
+ matched_count,
+ )?;
+ }
+ Some(Err(e)) => {
+ self.spill_stream = None;
+ self.spill_stream_has_data = false;
+ return Poll::Ready(Err(e));
+ }
+ None => {
+ self.spill_stream = None;
+ if !self.spill_stream_has_data {
+ return Poll::Ready(internal_err!("Spill file was empty"));
+ }
+ break;
+ }
}
}
}
@@ -830,13 +866,17 @@ impl BitwiseSortMergeJoinStream {
}
self.outer_offset = outer_group_end;
- Ok(())
+
+ self.spill_stream = None;
+ self.spill_stream_has_data = false;
+
+ Poll::Ready(Ok(()))
}
/// Continue processing an outer key group that spans multiple outer
/// batches. Returns `true` if this outer batch was fully consumed
/// by the key group and the caller should load another.
- fn resume_boundary(&mut self) -> Result<bool> {
+ fn resume_boundary(&mut self, cx: &mut Context<'_>) -> Poll<Result<bool>> {
debug_assert!(
self.outer_batch.is_some(),
"caller must load outer_batch first"
@@ -858,7 +898,7 @@ impl BitwiseSortMergeJoinStream {
});
self.emit_outer_batch()?;
self.outer_batch = None;
- return Ok(true);
+ return Poll::Ready(Ok(true));
}
}
}
@@ -874,7 +914,15 @@ impl BitwiseSortMergeJoinStream {
self.null_equality,
)?;
if same_key {
- self.process_key_match_with_filter()?;
+ match self.process_key_match_with_filter(cx) {
+ Poll::Ready(Ok(())) => (),
+ Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
+ Poll::Pending => {
+ self.pending_boundary =
+ Some(PendingBoundary::Filtered { saved_keys });
+ return Poll::Pending;
+ }
+ }
let num_outer = self.outer_batch.as_ref().unwrap().num_rows();
if self.outer_offset >= num_outer {
self.pending_boundary = Some(PendingBoundary::Filtered {
@@ -882,14 +930,63 @@ impl BitwiseSortMergeJoinStream {
});
self.emit_outer_batch()?;
self.outer_batch = None;
- return Ok(true);
+ return Poll::Ready(Ok(true));
}
}
self.clear_inner_key_group();
}
None => {}
}
- Ok(false)
+ Poll::Ready(Ok(false))
+ }
+
+ /// Helper to process an Equal match across potential outer batch boundaries.
+ fn process_filtered_match_loop(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+ loop {
+ ready!(self.process_key_match_with_filter(cx))?;
+
+ let outer_batch = self.outer_batch.as_ref().unwrap();
+ if self.outer_offset >= outer_batch.num_rows() {
+ let saved_keys =
+ slice_keys(&self.outer_key_arrays, outer_batch.num_rows() - 1);
+
+ self.emit_outer_batch()?;
+ self.pending_boundary = Some(PendingBoundary::Filtered { saved_keys });
+
+ // Clear stale batch before polling
+ self.outer_batch = None;
+
+ match ready!(self.poll_next_outer_batch(cx)) {
+ Err(e) => return Poll::Ready(Err(e)),
+ Ok(false) => {
+ self.pending_boundary = None;
+ break;
+ }
+ Ok(true) => {
+ let Some(PendingBoundary::Filtered { saved_keys }) =
+ self.pending_boundary.take()
+ else {
+ unreachable!()
+ };
+ let same = keys_match(
+ &saved_keys,
+ &self.outer_key_arrays,
+ &self.sort_options,
+ self.null_equality,
+ )?;
+ if same {
+ continue;
+ }
+ break;
+ }
+ }
+ } else {
+ break;
+ }
+ }
+
+ self.clear_inner_key_group(); // This resets inner_group_buffered to false
+ Poll::Ready(Ok(()))
}
/// Main loop: drive the merge-scan to produce output batches.
@@ -911,14 +1008,21 @@ impl BitwiseSortMergeJoinStream {
}
return Poll::Ready(Ok(None));
}
- Ok(true) => {
- if self.resume_boundary()? {
- continue;
- }
- }
+ Ok(true) => {} // Loaded batch, move on to checks
}
}
+ // Handles pausing while fetching a NEW outer batch.
+ if self.pending_boundary.is_some() && ready!(self.resume_boundary(cx))? {
+ continue;
+ }
+
+ // Handles pausing while reading the disk stream mid-batch.
+ if self.inner_group_buffered {
+ ready!(self.process_filtered_match_loop(cx))?;
+ continue;
+ }
+
// 2. Ensure we have an inner batch (unless inner is exhausted).
// Skip this when resuming a pending boundary — inner was already
// advanced past the key group before the boundary loop started.
@@ -1043,65 +1147,17 @@ impl BitwiseSortMergeJoinStream {
}
Ordering::Equal => {
if self.filter.is_some() {
+ debug_assert!(!self.inner_group_buffered);
// Buffer inner key group (may span batches)
match ready!(self.buffer_inner_key_group(cx)) {
Err(e) => return Poll::Ready(Err(e)),
- Ok(_inner_exhausted) => {}
+ Ok(_inner_exhausted) => {
+ self.inner_group_buffered = true;
+ }
}
-
// Process outer rows against buffered inner group
// (may need to handle outer batch boundary)
- loop {
- self.process_key_match_with_filter()?;
-
- let outer_batch = self.outer_batch.as_ref().unwrap();
- if self.outer_offset >= outer_batch.num_rows() {
- let saved_keys = slice_keys(
- &self.outer_key_arrays,
- outer_batch.num_rows() - 1,
- );
-
- self.emit_outer_batch()?;
- debug_assert!(
- !self.inner_key_buffer.is_empty()
- || self.inner_key_spill.is_some(),
- "Filtered pending boundary requires inner key data in buffer or spill"
- );
- self.pending_boundary =
- Some(PendingBoundary::Filtered { saved_keys });
-
- match ready!(self.poll_next_outer_batch(cx)) {
- Err(e) => return Poll::Ready(Err(e)),
- Ok(false) => {
- self.pending_boundary = None;
- self.outer_batch = None;
- break;
- }
- Ok(true) => {
- let Some(PendingBoundary::Filtered {
- saved_keys,
- }) = self.pending_boundary.take()
- else {
- unreachable!()
- };
- let same = keys_match(
- &saved_keys,
- &self.outer_key_arrays,
- &self.sort_options,
- self.null_equality,
- )?;
- if same {
- continue;
- }
- break;
- }
- }
- } else {
- break;
- }
- }
-
- self.clear_inner_key_group();
+ ready!(self.process_filtered_match_loop(cx))?;
} else {
// No filter: advance inner past key group, then
// mark all outer rows with this key as matched.
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 9bcc749c23..f1a18aac76 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
@@ -23,8 +23,7 @@
use std::cmp::Ordering;
use std::collections::{HashMap, VecDeque};
-use std::fs::File;
-use std::io::BufReader;
+use std::fmt::Debug;
use std::mem::size_of;
use std::ops::Range;
use std::pin::Pin;
@@ -47,10 +46,9 @@ use crate::{PhysicalExpr, RecordBatchStream, SendableRecordBatchStream};
use arrow::array::{types::UInt64Type, *};
use arrow::compute::{
self, BatchCoalescer, SortOptions, concat_batches, filter_record_batch, interleave,
- take, take_arrays,
+ take_arrays,
};
use arrow::datatypes::SchemaRef;
-use arrow::ipc::reader::StreamReader;
use datafusion_common::cast::as_uint64_array;
use datafusion_common::{JoinType, NullEquality, Result, exec_err, internal_err};
use datafusion_execution::disk_manager::RefCountedTempFile;
@@ -58,7 +56,7 @@ use datafusion_execution::memory_pool::MemoryReservation;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
-use futures::{Stream, StreamExt};
+use futures::{Stream, StreamExt, ready};
/// State of SMJ stream
#[derive(Debug, PartialEq, Eq)]
@@ -382,6 +380,11 @@ pub(super) struct MaterializingSortMergeJoinStream {
/// Manages the process of spilling and reading back intermediate data
pub spill_manager: SpillManager,
+ /// Tracks the active stream when loading spilled buffered batches back in memory
+ pub spill_stream: Option<SendableRecordBatchStream>,
+ /// Tracks the number of batches currently spilled
+ pub spilled_batch_count: usize,
+
// ========================================================================
// CACHED COMPARATORS:
// Pre-built comparators to avoid per-row type dispatch in hot loops.
@@ -599,6 +602,16 @@ impl Stream for MaterializingSortMergeJoinStream {
.filter_mask
.len();
if accumulated >= self.batch_size {
+ // Ensure required spilled batches are restored to memory
+ // before processing, as this path invokes freeze_all().
+ let needed = self.get_required_batch_indices(
+ self.buffered_data.batches.len(),
+ );
+ if let Err(e) = ready!(
+ self.poll_spilled_batches(cx, &needed)
+ ) {
+ return Poll::Ready(Some(Err(e)));
+ }
match self.process_filtered_batches()? {
Poll::Ready(Some(batch)) => {
return Poll::Ready(Some(Ok(batch)));
@@ -684,14 +697,13 @@ impl Stream for MaterializingSortMergeJoinStream {
self.state = SortMergeJoinState::Init;
}
SortMergeJoinState::JoinOutput => {
- self.join_partial()?;
+ // If the batch size limit is reached, restore required spilled batches to memory and freeze.
+ // Guarding at the top of the loop safely handles re-entry from Poll::Pending.
+ if self.num_unfrozen_pairs() >= self.batch_size {
+ let needed = self
+ .get_required_batch_indices(self.buffered_data.batches.len());
+ ready!(self.poll_spilled_batches(cx, &needed))?;
- if self.num_unfrozen_pairs() < self.batch_size {
- if self.buffered_data.scanning_finished() {
- self.buffered_data.scanning_reset();
- self.state = SortMergeJoinState::EmitReadyThenInit;
- }
- } else {
self.freeze_all()?;
// Verify metadata alignment before checking if we have batches to output
@@ -705,7 +717,6 @@ impl Stream for MaterializingSortMergeJoinStream {
}
// For non-filtered joins, only output if we have a completed batch
- // (opportunistic output when target batch size is reached)
if self
.joined_record_batches
.joined_batches
@@ -720,10 +731,26 @@ impl Stream for MaterializingSortMergeJoinStream {
.record_output(&self.join_metrics.baseline_metrics());
return Poll::Ready(Some(Ok(record_batch)));
}
+
// Otherwise keep buffering (don't output yet)
+ continue;
+ }
+
+ self.join_partial()?;
+
+ if self.num_unfrozen_pairs() < self.batch_size
+ && self.buffered_data.scanning_finished()
+ {
+ self.buffered_data.scanning_reset();
+ self.state = SortMergeJoinState::EmitReadyThenInit;
}
+ // Note: If join_partial() reached the batch size, the loop repeats to freeze the data.
}
SortMergeJoinState::Exhausted => {
+ let needed =
+ self.get_required_batch_indices(self.buffered_data.batches.len());
+ ready!(self.poll_spilled_batches(cx, &needed))?;
+
self.freeze_all()?;
// Verify metadata alignment before final output
@@ -843,6 +870,8 @@ impl MaterializingSortMergeJoinStream {
reservation,
runtime_env,
spill_manager,
+ spill_stream: None,
+ spilled_batch_count: 0,
streamed_buffered_cmp: None,
buffered_equality_cmp: None,
streamed_batch_counter: AtomicUsize::new(0),
@@ -917,6 +946,84 @@ impl MaterializingSortMergeJoinStream {
Poll::Pending
}
+ /// Identifies which buffered batches are needed for the upcoming freeze operation
+ 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 {
+ return needed;
+ }
+ // We need all batches that matched with streamed rows
+ for chunk in &self.streamed_batch.output_indices {
+ if let Some(idx) = chunk.buffered_batch_idx {
+ needed.push(idx);
+ }
+ }
+
+ // Full Joins need to emit null-joined rows, so we need batches up to freeze_count
+ if self.join_type == JoinType::Full {
+ needed.extend(0..buffered_freeze_count);
+ }
+
+ needed.sort_unstable();
+ needed.dedup();
+ needed
+ }
+
+ /// Asynchronously reads spilled batches back into memory.
+ /// Only processes the required indices to avoid OOMs.
+ fn poll_spilled_batches(
+ &mut self,
+ cx: &mut Context<'_>,
+ required_indices: &[usize],
+ ) -> Poll<Result<()>> {
+ 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() {
+ continue;
+ }
+
+ let bb = &mut self.buffered_data.batches[idx];
+
+ if let BufferedBatchState::Spilled(spill_file) = &bb.batch {
+ if self.spill_stream.is_none() {
+ let stream = self
+ .spill_manager
+ .read_spill_as_stream(spill_file.clone(), None)?;
+ self.spill_stream = Some(stream);
+ }
+
+ match ready!(self.spill_stream.as_mut().unwrap().poll_next_unpin(cx)) {
+ Some(Ok(batch)) => {
+ // Transition the batch back to InMemory
+ bb.batch = BufferedBatchState::InMemory(batch);
+ self.spilled_batch_count -= 1;
+ // The batch is back in memory, so we must account for its size.
+ let newly_allocated =
+ bb.size_estimation.saturating_sub(bb.reserved_amount);
+ self.reservation.grow(newly_allocated);
+ bb.reserved_amount = bb.size_estimation;
+
+ self.join_metrics
+ .peak_mem_used()
+ .set_max(self.reservation.size());
+
+ self.spill_stream = None;
+ }
+ Some(Err(e)) => {
+ self.spill_stream = None;
+ return Poll::Ready(Err(e));
+ }
+ None => {
+ self.spill_stream = None;
+ return Poll::Ready(internal_err!("Spill file was empty"));
+ }
+ }
+ }
+ }
+ Poll::Ready(Ok(()))
+ }
+
/// Poll next streamed row
fn poll_streamed_row(&mut self, cx: &mut Context) -> Poll<Option<Result<()>>> {
loop {
@@ -931,33 +1038,41 @@ impl MaterializingSortMergeJoinStream {
self.streamed_state = StreamedState::Polling;
}
}
- StreamedState::Polling => match self.streamed.poll_next_unpin(cx)? {
- Poll::Pending => {
- return Poll::Pending;
- }
- Poll::Ready(None) => {
- // Release the streamed input pipeline's resources.
- let streamed_schema = self.streamed.schema();
- self.streamed =
- Box::pin(EmptyRecordBatchStream::new(streamed_schema));
- self.streamed_state = StreamedState::Exhausted;
+ StreamedState::Polling => {
+ let needed =
+ self.get_required_batch_indices(self.buffered_data.batches.len());
+ if let Err(e) = ready!(self.poll_spilled_batches(cx, &needed)) {
+ return Poll::Ready(Some(Err(e)));
}
- Poll::Ready(Some(batch)) => {
- if batch.num_rows() > 0 {
- self.freeze_streamed()?;
- self.join_metrics.input_batches().add(1);
- self.join_metrics.input_rows().add(batch.num_rows());
- self.streamed_batch =
- StreamedBatch::new(batch, &self.on_streamed);
- self.rebuild_streamed_buffered_cmp()?;
- // Every incoming streaming batch should have its unique id
- // Check `JoinedRecordBatches.self.streamed_batch_counter` documentation
- self.streamed_batch_counter
- .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
- self.streamed_state = StreamedState::Ready;
+
+ match self.streamed.poll_next_unpin(cx)? {
+ Poll::Pending => {
+ return Poll::Pending;
+ }
+ Poll::Ready(None) => {
+ // Release the streamed input pipeline's resources.
+ let streamed_schema = self.streamed.schema();
+ self.streamed =
+ Box::pin(EmptyRecordBatchStream::new(streamed_schema));
+ self.streamed_state = StreamedState::Exhausted;
+ }
+ Poll::Ready(Some(batch)) => {
+ if batch.num_rows() > 0 {
+ self.freeze_streamed()?;
+ self.join_metrics.input_batches().add(1);
+ self.join_metrics.input_rows().add(batch.num_rows());
+ self.streamed_batch =
+ StreamedBatch::new(batch, &self.on_streamed);
+ self.rebuild_streamed_buffered_cmp()?;
+ // Every incoming streaming batch should have its unique id
+ // Check `JoinedRecordBatches.self.streamed_batch_counter` documentation
+ self.streamed_batch_counter
+ .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+ self.streamed_state = StreamedState::Ready;
+ }
}
}
- },
+ }
StreamedState::Ready => {
return Poll::Ready(Some(Ok(())));
}
@@ -997,6 +1112,7 @@ impl MaterializingSortMergeJoinStream {
.unwrap(); // Operation only return None if no batches are spilled, here we ensure that at least one batch is spilled
buffered_batch.batch = BufferedBatchState::Spilled(spill_file);
+ self.spilled_batch_count += 1;
// Join key arrays remain in memory after the batch is
// spilled — the comparator needs them for key boundary
@@ -1036,12 +1152,25 @@ impl MaterializingSortMergeJoinStream {
let head_batch = self.buffered_data.head_batch();
// If the head batch is fully processed, dequeue it and produce output of it.
if head_batch.range.end == head_batch.num_rows {
+ // load the spilled head batch before dequeuing
+ let needed = self.get_required_batch_indices(1);
+ if let Err(e) = ready!(self.poll_spilled_batches(cx, &needed))
+ {
+ return Poll::Ready(Some(Err(e)));
+ }
+
self.freeze_dequeuing_buffered()?;
if let Some(mut buffered_batch) =
self.buffered_data.batches.pop_front()
{
self.produce_buffered_not_matched(&mut buffered_batch)?;
self.free_reservation(&buffered_batch);
+ if matches!(
+ buffered_batch.batch,
+ BufferedBatchState::Spilled(_)
+ ) {
+ self.spilled_batch_count -= 1;
+ }
head_changed = true;
}
} else {
@@ -1556,18 +1685,6 @@ impl MaterializingSortMergeJoinStream {
as_uint64_array(&compute::concat(&refs)?)?.clone()
};
- let spill_reservation = self.reservation.new_empty();
- if matches!(
- &self.buffered_data.batches[first_batch_idx].batch,
- BufferedBatchState::Spilled(_)
- ) {
- spill_reservation
- .grow(self.buffered_data.batches[first_batch_idx].size_estimation);
- self.join_metrics
- .peak_mem_used()
- .set_max(self.reservation.size() + spill_reservation.size());
- }
-
return fetch_right_columns_by_idxs(
&self.buffered_data,
first_batch_idx,
@@ -1603,29 +1720,20 @@ impl MaterializingSortMergeJoinStream {
let num_right_cols = self.buffered_schema.fields().len();
// Read each source batch once (spilled batches require disk I/O).
- // Track memory for each spilled batch at the point of deserialization
- // so the pool reflects actual usage as it grows.
- let spill_reservation = self.reservation.new_empty();
- let mut source_data: Vec<Option<RecordBatch>> =
- Vec::with_capacity(source_batches.len());
- for &idx in &source_batches {
- let bb = &self.buffered_data.batches[idx];
- match &bb.batch {
- BufferedBatchState::InMemory(batch) => {
- source_data.push(Some(batch.clone()));
- }
- BufferedBatchState::Spilled(spill_file) => {
- spill_reservation.grow(bb.size_estimation);
- self.join_metrics
- .peak_mem_used()
- .set_max(self.reservation.size() + spill_reservation.size());
-
- let file = BufReader::new(File::open(spill_file.path())?);
- let reader = StreamReader::try_new(file, None)?;
- source_data.push(reader.into_iter().next().transpose()?);
+ let source_data_result: Result<Vec<RecordBatch>> = source_batches
+ .iter()
+ .map(|&idx| {
+ let bb = &self.buffered_data.batches[idx];
+ match &bb.batch {
+ BufferedBatchState::InMemory(batch) => Ok(batch.clone()),
+ BufferedBatchState::Spilled(_) => {
+ internal_err!("Buffered batch should have been unspilled before fetching columns")
+ }
}
- }
- }
+ })
+ .collect();
+
+ let source_data = source_data_result?;
let mut right_columns = Vec::with_capacity(num_right_cols);
for col_idx in 0..num_right_cols {
@@ -1637,14 +1745,7 @@ impl MaterializingSortMergeJoinStream {
source_arrays.push(null_array.as_ref());
for data in &source_data {
- match data {
- Some(batch) => source_arrays.push(batch.column(col_idx).as_ref()),
- None => {
- return internal_err!(
- "Failed to read spilled buffered batch during interleave"
- );
- }
- }
+ source_arrays.push(data.column(col_idx).as_ref());
}
right_columns.push(interleave(&source_arrays, &interleave_indices)?);
}
@@ -1838,32 +1939,17 @@ fn fetch_right_columns_from_batch_by_idxs(
buffered_indices: &UInt64Array,
) -> Result<Vec<ArrayRef>> {
match &buffered_batch.batch {
- // In memory batch
- // In memory batch
BufferedBatchState::InMemory(batch) => {
- // When indices form a contiguous range (common in SMJ since the
- // buffered side is scanned sequentially), use zero-copy slice.
if let Some(range) = is_contiguous_range(buffered_indices) {
Ok(batch.slice(range.start, range.len()).columns().to_vec())
} else {
Ok(take_arrays(batch.columns(), buffered_indices, None)?)
}
}
- // If the batch was spilled to disk, less likely
- BufferedBatchState::Spilled(spill_file) => {
- let mut buffered_cols: Vec<ArrayRef> =
- Vec::with_capacity(buffered_indices.len());
-
- let file = BufReader::new(File::open(spill_file.path())?);
- let reader = StreamReader::try_new(file, None)?;
-
- for batch in reader {
- batch?.columns().iter().for_each(|column| {
- buffered_cols.extend(take(column, &buffered_indices, None))
- });
- }
-
- Ok(buffered_cols)
+ BufferedBatchState::Spilled(_) => {
+ internal_err!(
+ "Buffered batch should have been unspilled before fetching columns"
+ )
}
}
}
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 b1fdf3ddab..0347299dd0 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -2459,6 +2459,24 @@ async fn overallocation_multi_batch_spill() -> Result<()> {
assert!(join.metrics().unwrap().spilled_bytes().unwrap() > 0);
assert!(join.metrics().unwrap().spilled_rows().unwrap() > 0);
+ // For Full joins, get_required_batch_indices extends 0..batches.len(), so
+ // poll_spilled_batches can restore all spilled batches at once via infallible
+ // grow(). Verify accounting tracked the transient spike and cleaned up.
+ let peak_mem = join
+ .metrics()
+ .and_then(|m| m.sum_by_name("peak_mem_used"))
+ .map(|m| m.as_usize())
+ .unwrap_or(0);
+ assert!(
+ peak_mem > 0,
+ "peak_mem_used should be > 0 for {join_type:?} batch_size={batch_size}"
+ );
+ assert_eq!(
+ runtime.memory_pool.reserved(),
+ 0,
+ "memory should be fully released after {join_type:?} completes
+ (batch_size={batch_size}): infallible grow during restore must be balanced"
+ );
// Run the test with no spill configuration as
let task_ctx_no_spill =
TaskContext::default().with_session_config(session_config.clone());
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org