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/08/27 12:51:44 UTC
(datafusion) 01/01: fix: spill the NestedLoopJoin build side from the pass that consumed it (#24677)
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-24677-4ea9975291942503045cee0b1a47a8239591d911
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e387d955205c0ad9f0ecf7d9c3fa2545cf984c00
Author: ranflarion <ra...@flarion.io>
AuthorDate: Thu Aug 27 12:51:25 2026 +0000
fix: spill the NestedLoopJoin build side from the pass that consumed it (#24677)
## Which issue does this PR close?
- Closes #24676.
## Rationale for this change
The memory-limited fallback re-executes the build-side child: when the
in-memory load fails with `ResourcesExhausted`, `initiate_fallback`
calls `self.left.execute(0, ctx)` a second time and spills that stream —
but the first execution already consumed part of the input. For a child
that cannot be executed twice, the consumed batches are gone.
Proof on current `main`, wrapping this file's own `build_left_table()`
in an `ExecutionPlan` whose second `execute` returns an empty stream
(the way an exhausted external reader behaves), under the 50-byte limit
the existing memory-limited tests use:
```
executions=2 result=Err(Internal("Left side produced no data to spill"))
```
For a child whose stream resumes instead of restarting, the same shape
loses the consumed prefix silently: embedding DataFusion 54.1.0 over a
JVM-fed source we measured a 200k-row build side returning 134,464 rows
with no error. Replayable inputs are also executed twice for no benefit.
This PR moves the spill into the load: when a reservation fails inside
`collect_left_input`, the batches buffered so far plus the remainder of
the same stream are written to one spill file, and the shared
`OnceAsync` now resolves to a `LeftLoad` enum (`InMemory(JoinLeftData)`
or `Spilled(LeftSpillData)`) that every partition consumes. The build
side is executed exactly once whether or not it fits,
`SpillState::Pending` no longer carries the child plan, and an empty
build side stays on the in-memory path instead of erroring as "no data
to spill".
## What changes are included in this PR?
`collect_left_input` takes an optional `SpillManager` (built up front,
since by the time the limit is hit the stream is already partly
consumed) and returns `LeftLoad`; a new `spill_left_input` writes
buffered batches plus the stream remainder; `initiate_fallback` becomes
`enter_memory_limited_mode(Arc<LeftSpillData>)`; the per-partition chunk
reader opens the shared spill file directly instead of awaiting a spill
future; the load reservation is created `with_can_spill(can_spill)`
since the consumer now genuinely spills.
## Are these changes tested?
New test `memory_limited_left_side_reads_the_child_once`: a one-shot
child under the tight limit must produce the full result with
`spill_count > 0` and exactly one `execute` of the build side — on
unpatched `main` this scenario fails as shown above. All 43
`nested_loop_join` tests pass (the existing memory-limited tests cover
every join type through the new path), 1783 crate tests pass,
`./dev/rust_lint.sh` is clean.
## Are there any user-facing changes?
Memory-limited nested loop joins over non-replayable inputs now return
complete results instead of failing (or silently truncating), and the
build side is no longer executed twice. No API changes;
`LeftLoad`/`LeftSpillData` are `pub(crate)`.
---
.../physical-plan/src/joins/nested_loop_join.rs | 562 +++++++++++++--------
1 file changed, 360 insertions(+), 202 deletions(-)
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index 548a0cf1db..3d966702fb 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -75,7 +75,7 @@ use datafusion_physical_expr::equivalence::{
};
use datafusion_physical_expr::projection::{ProjectionRef, combine_projections};
-use futures::{Stream, StreamExt, TryStreamExt};
+use futures::{Stream, StreamExt};
use log::debug;
use parking_lot::Mutex;
@@ -209,14 +209,7 @@ pub struct NestedLoopJoinExec {
///
/// Each output stream waits on the `OnceAsync` to signal the completion of
/// the build(left) side data, and buffer them all for later joining.
- build_side_data: OnceAsync<JoinLeftData>,
- /// Shared left-side spill data for OOM fallback.
- ///
- /// When `build_side_data` fails with OOM, the first partition to
- /// initiate fallback spills the entire left side to disk. Other
- /// partitions share the same spill file via this `OnceAsync`,
- /// avoiding redundant re-execution of the left child.
- left_spill_data: Arc<OnceAsync<LeftSpillData>>,
+ build_side_data: OnceAsync<LeftLoad>,
/// Information of index and left / right placement of columns
column_indices: Vec<ColumnIndex>,
/// Projection to apply to the output of the join
@@ -300,7 +293,6 @@ impl NestedLoopJoinExecBuilder {
join_type,
join_schema,
build_side_data: Default::default(),
- left_spill_data: Arc::new(OnceAsync::default()),
column_indices,
projection,
metrics: Default::default(),
@@ -454,18 +446,19 @@ impl NestedLoopJoinExec {
pub fn swap_inputs(&self) -> Result<Arc<dyn ExecutionPlan>> {
let left = self.left();
let right = self.right();
- let new_join = NestedLoopJoinExec::try_new(
+ let new_join = NestedLoopJoinExecBuilder::new(
Arc::clone(right),
Arc::clone(left),
- self.filter().map(JoinFilter::swap),
- &self.join_type().swap(),
- swap_join_projection(
- left.schema().fields().len(),
- right.schema().fields().len(),
- self.projection.as_deref(),
- self.join_type(),
- ),
- )?;
+ self.join_type().swap(),
+ )
+ .with_filter(self.filter().map(JoinFilter::swap))
+ .with_projection(swap_join_projection(
+ left.schema().fields().len(),
+ right.schema().fields().len(),
+ self.projection.as_deref(),
+ self.join_type(),
+ ))
+ .build()?;
// For Semi/Anti joins, swap result will produce same output schema,
// no need to wrap them into additional projection
@@ -589,7 +582,6 @@ impl ExecutionPlan for NestedLoopJoinExec {
right,
metrics: ExecutionPlanMetricsSet::new(),
build_side_data: Default::default(),
- left_spill_data: Arc::new(OnceAsync::default()),
cache: Arc::clone(&self.cache),
filter: self.filter.clone(),
join_type: self.join_type,
@@ -657,15 +649,43 @@ impl ExecutionPlan for NestedLoopJoinExec {
let right_partition_count = self.right().output_partitioning().partition_count();
- // Always try to buffer all left data in memory via OnceFut.
- // If that fails with OOM, the stream will fallback to memory-limited
- // mode (if conditions allow).
+ // Always try to buffer all left data in memory via OnceFut. If it does not fit, the load
+ // spills the left side during the same pass and the stream runs in memory-limited mode.
let load_reservation =
MemoryConsumer::new(format!("NestedLoopJoinLoad[{partition}]"))
.register(context.memory_pool());
+ // Determine if memory-limited mode is possible.
+ // Conditions:
+ // 1. Disk manager supports temp files (needed for spilling).
+ // 2. FULL join with multiple right partitions is not yet supported
+ // in the memory-limited path. FULL join needs to track BOTH left-side
+ // matches (for unmatched left rows) AND right-side matches (for
+ // unmatched right rows). That path builds a per-partition
+ // `JoinLeftData` with `probe_threads_counter == 1`, so each
+ // partition emits unmatched left rows based only on its own
+ // right-side matches, producing incorrect duplicate output for
+ // left rows that match in another partition. Other join types
+ // that need only one-sided final emission (LEFT, LEFT SEMI,
+ // LEFT ANTI, LEFT MARK) have a similar latent issue in that path.
+ let full_join_multi_partition =
+ matches!(self.join_type, JoinType::Full) && right_partition_count > 1;
+ let can_spill = context.runtime_env().disk_manager.tmp_files_enabled()
+ && !full_join_multi_partition;
+
let build_side_data = self.build_side_data.try_once(|| {
let stream = self.left.execute(0, Arc::clone(&context))?;
+ // Built here rather than on demand: by the time the load hits the memory limit the
+ // stream is partly consumed, and it has to be spillable at that point without going
+ // back to the left child.
+ let left_spill_manager = can_spill.then(|| {
+ SpillManager::new(
+ context.runtime_env(),
+ metrics.spill_metrics.clone(),
+ stream.schema(),
+ )
+ .with_compression_type(context.session_config().spill_compression())
+ });
Ok(collect_left_input(
stream,
@@ -673,35 +693,15 @@ impl ExecutionPlan for NestedLoopJoinExec {
load_reservation,
need_produce_result_in_final(self.join_type),
right_partition_count,
+ left_spill_manager,
))
})?;
let probe_side_data = self.right.execute(partition, Arc::clone(&context))?;
- // Determine if OOM fallback to memory-limited mode is possible.
- // Conditions:
- // 1. Disk manager supports temp files (needed for spilling).
- // 2. FULL join with multiple right partitions is not yet supported
- // in the fallback path. FULL join needs to track BOTH left-side
- // matches (for unmatched left rows) AND right-side matches (for
- // unmatched right rows). The fallback path builds a per-partition
- // `JoinLeftData` with `probe_threads_counter == 1`, so each
- // partition emits unmatched left rows based only on its own
- // right-side matches, producing incorrect duplicate output for
- // left rows that match in another partition. Other join types
- // that need only one-sided final emission (LEFT, LEFT SEMI,
- // LEFT ANTI, LEFT MARK) have a similar latent issue in the
- // fallback path which predates this change; tracking is out of
- // scope for this PR.
- let full_join_multi_partition =
- matches!(self.join_type, JoinType::Full) && right_partition_count > 1;
- let spill_state = if context.runtime_env().disk_manager.tmp_files_enabled()
- && !full_join_multi_partition
- {
+ let spill_state = if can_spill {
SpillState::Pending {
- left_plan: Arc::clone(&self.left),
task_context: Arc::clone(&context),
- left_spill_data: Arc::clone(&self.left_spill_data),
}
} else {
SpillState::Disabled
@@ -786,14 +786,15 @@ impl ExecutionPlan for NestedLoopJoinExec {
self.filter(),
self.column_indices.as_slice(),
)? {
- Ok(Some(Arc::new(NestedLoopJoinExec::try_new(
- Arc::new(projected_left_child),
- Arc::new(projected_right_child),
- join_filter,
- self.join_type(),
- // Returned early if projection is not None
- None,
- )?)))
+ Ok(Some(Arc::new(
+ NestedLoopJoinExecBuilder::new(
+ Arc::new(projected_left_child),
+ Arc::new(projected_right_child),
+ *self.join_type(),
+ )
+ .with_filter(join_filter)
+ .build()?,
+ )))
} else {
try_embed_projection(projection, self)
}
@@ -817,8 +818,6 @@ impl ExecutionPlan for NestedLoopJoinExec {
join_schema: _,
// runtime build-side state, not part of the plan
build_side_data: _,
- // runtime spill state, not part of the plan
- left_spill_data: _,
// recomputed by `try_new` on decode
column_indices: _,
// runtime metrics, not part of the plan
@@ -1055,34 +1054,57 @@ impl JoinLeftData {
}
}
-/// Asynchronously collect input into a single batch, and creates `JoinLeftData` from it
+/// Asynchronously collect the left input in a single pass over the stream.
+///
+/// The whole side is buffered in memory when it fits the budget. When it does not and
+/// `spill_manager` is available, the batches collected so far and the rest of the same stream are
+/// written to one spill file, which is what the memory-limited mode then reads in chunks. The
+/// stream is consumed exactly once either way, so a left child that cannot be executed twice (or
+/// replayed) stays correct.
async fn collect_left_input(
- stream: SendableRecordBatchStream,
+ mut stream: SendableRecordBatchStream,
join_metrics: BuildProbeJoinMetrics,
reservation: MemoryReservation,
with_visited_left_side: bool,
probe_threads_count: usize,
-) -> Result<JoinLeftData> {
+ spill_manager: Option<SpillManager>,
+) -> Result<LeftLoad> {
let schema = stream.schema();
-
- // Load all batches and count the rows
- let (batches, metrics, reservation) = stream
- .try_fold(
- (Vec::new(), join_metrics, reservation),
- |(mut batches, metrics, reservation), batch| async {
- let batch_size = batch.get_array_memory_size();
- // Reserve memory for incoming batch
- reservation.try_grow(batch_size)?;
- // Update metrics
+ let metrics = join_metrics;
+ let mut batches: Vec<RecordBatch> = Vec::new();
+
+ while let Some(batch) = stream.next().await {
+ let batch = batch?;
+ let batch_size = batch.get_array_memory_size();
+ match reservation.try_grow(batch_size) {
+ Ok(()) => {
metrics.build_mem_used.add(batch_size);
metrics.build_input_batches.add(1);
metrics.build_input_rows.add(batch.num_rows());
- // Push batch to output
batches.push(batch);
- Ok((batches, metrics, reservation))
- },
- )
- .await?;
+ }
+ Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => {
+ let spill_manager = spill_manager.expect("checked by is_spillable_oom");
+ let spilled = spill_left_input(
+ spill_manager,
+ Arc::clone(&schema),
+ batches,
+ Some(batch),
+ stream,
+ metrics,
+ &reservation,
+ )
+ .await?;
+ return Ok(left_load_from_spill(
+ spilled,
+ schema,
+ probe_threads_count,
+ reservation,
+ ));
+ }
+ Err(e) => return Err(e),
+ }
+ }
let merged_batch = concat_batches(&schema, &batches)?;
@@ -1090,7 +1112,30 @@ async fn collect_left_input(
let visited_left_side = if with_visited_left_side {
let n_rows = merged_batch.num_rows();
let buffer_size = n_rows.div_ceil(8);
- reservation.try_grow(buffer_size)?;
+ match reservation.try_grow(buffer_size) {
+ Ok(()) => {}
+ Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => {
+ let spill_manager = spill_manager.expect("checked by is_spillable_oom");
+ drop(batches);
+ let spilled = spill_left_input(
+ spill_manager,
+ Arc::clone(&schema),
+ vec![merged_batch],
+ None,
+ stream,
+ metrics,
+ &reservation,
+ )
+ .await?;
+ return Ok(left_load_from_spill(
+ spilled,
+ schema,
+ probe_threads_count,
+ reservation,
+ ));
+ }
+ Err(e) => return Err(e),
+ }
metrics.build_mem_used.add(buffer_size);
let mut buffer = BooleanBufferBuilder::new(n_rows);
@@ -1100,12 +1145,90 @@ async fn collect_left_input(
BooleanBufferBuilder::new(0)
};
- Ok(JoinLeftData::new(
+ Ok(LeftLoad::InMemory(Arc::new(JoinLeftData::new(
merged_batch,
Mutex::new(visited_left_side),
AtomicUsize::new(probe_threads_count),
reservation,
- ))
+ ))))
+}
+
+/// A left side with no rows needs no spill file, so it stays on the in-memory path.
+fn left_load_from_spill(
+ spilled: Option<LeftSpillData>,
+ schema: SchemaRef,
+ probe_threads_count: usize,
+ reservation: MemoryReservation,
+) -> LeftLoad {
+ match spilled {
+ Some(data) => LeftLoad::Spilled(Arc::new(data)),
+ // No rows means no bitmap either, whatever the join type.
+ None => LeftLoad::InMemory(Arc::new(JoinLeftData::new(
+ RecordBatch::new_empty(schema),
+ Mutex::new(BooleanBufferBuilder::new(0)),
+ AtomicUsize::new(probe_threads_count),
+ reservation,
+ ))),
+ }
+}
+
+/// Whether a failed reservation is an exhausted pool that the caller can spill its way out of.
+fn is_spillable_oom(
+ error: &datafusion_common::DataFusionError,
+ spill_manager: Option<&SpillManager>,
+) -> bool {
+ spill_manager.is_some()
+ && matches!(
+ error.find_root(),
+ datafusion_common::DataFusionError::ResourcesExhausted(_)
+ )
+}
+
+/// Write the already-buffered left batches plus the remainder of the same stream to one spill file.
+/// Returns `None` when the left side carried no rows at all, which needs no spill file.
+async fn spill_left_input(
+ spill_manager: SpillManager,
+ schema: SchemaRef,
+ buffered: Vec<RecordBatch>,
+ pending: Option<RecordBatch>,
+ mut stream: SendableRecordBatchStream,
+ metrics: BuildProbeJoinMetrics,
+ reservation: &MemoryReservation,
+) -> Result<Option<LeftSpillData>> {
+ let mut spill_file =
+ spill_manager.create_in_progress_file("NestedLoopJoin left spill")?;
+
+ for batch in buffered {
+ if batch.num_rows() > 0 {
+ spill_file.append_batch(&batch)?;
+ }
+ }
+ // The in-memory batches are spilled and dropped, so their reservation goes back to the pool
+ // before the rest of the stream is drained.
+ reservation.free();
+
+ for batch in pending.into_iter() {
+ if batch.num_rows() > 0 {
+ metrics.build_input_batches.add(1);
+ metrics.build_input_rows.add(batch.num_rows());
+ spill_file.append_batch(&batch)?;
+ }
+ }
+
+ while let Some(batch) = stream.next().await {
+ let batch = batch?;
+ if batch.num_rows() > 0 {
+ metrics.build_input_batches.add(1);
+ metrics.build_input_rows.add(batch.num_rows());
+ spill_file.append_batch(&batch)?;
+ }
+ }
+
+ Ok(spill_file.finish()?.map(|file| LeftSpillData {
+ spill_manager,
+ spill_file: file,
+ schema,
+ }))
}
/// States for join processing. See `poll_next()` comment for more details about
@@ -1134,11 +1257,16 @@ enum NLJState {
EmitGlobalRightUnmatched,
Done,
}
-/// Shared data for the left-side spill fallback.
-///
-/// When the in-memory `OnceFut` path fails with OOM, the first partition
-/// spills the entire left side to disk. This struct holds the spill file
-/// reference so other partitions can read from the same file.
+/// Outcome of the single pass over the left (build) input.
+pub(crate) enum LeftLoad {
+ /// The left side fit the memory budget and is buffered as one batch.
+ InMemory(Arc<JoinLeftData>),
+ /// The budget ran out, so the left side was spilled during that same pass. Every partition
+ /// shares this handle, and each left chunk pass re-opens the file.
+ Spilled(Arc<LeftSpillData>),
+}
+
+/// The spilled left side, shared by every output partition.
pub(crate) struct LeftSpillData {
/// SpillManager used to read the spill file (has the left schema)
spill_manager: SpillManager,
@@ -1148,31 +1276,25 @@ pub(crate) struct LeftSpillData {
schema: SchemaRef,
}
-/// Tracks the state of the memory-limited spill fallback for NLJ.
+/// Tracks the state of the memory-limited spill mode for NLJ.
///
-/// The NLJ always starts with the standard OnceFut path. If the in-memory
-/// load fails with OOM and conditions allow, the operator falls back to a
-/// multi-pass strategy where left data is loaded in chunks and the right
-/// side is spilled to disk.
+/// The NLJ always tries to buffer the whole left side in memory. If that does not fit and
+/// conditions allow, the load spills the left side (see [`collect_left_input`]) and the operator
+/// switches to a multi-pass strategy where left chunks are read back from that spill file and the
+/// right side is spilled for re-scanning.
pub(crate) enum SpillState {
- /// Fallback is not possible (e.g., join type requires global right bitmap,
+ /// Memory-limited mode is not possible (e.g., join type requires global right bitmap,
/// or disk manager is disabled). OOM errors will propagate as-is.
Disabled,
- /// Fallback is possible but not yet triggered. The operator is still
- /// attempting the standard OnceFut path. Holds the context needed to
- /// initiate fallback if OOM occurs.
+ /// Memory-limited mode is possible but not entered: the left side is still expected to fit.
+ /// Holds the context needed to set the mode up.
Pending {
- /// Left child plan for re-execution
- left_plan: Arc<dyn ExecutionPlan>,
- /// TaskContext for re-execution and SpillManager creation
+ /// TaskContext for reservations and SpillManager creation
task_context: Arc<TaskContext>,
- /// Shared OnceAsync for left-side spill data. The first partition
- /// to initiate fallback spills the left side; others share the file.
- left_spill_data: Arc<OnceAsync<LeftSpillData>>,
},
- /// Fallback has been triggered. Left data is being loaded in chunks
+ /// Memory-limited mode is running. Left data is read back in chunks
/// and the right side is spilled to disk for re-scanning.
Active(Box<SpillStateActive>),
}
@@ -1180,13 +1302,12 @@ pub(crate) enum SpillState {
/// State for active memory-limited spill execution.
/// Boxed inside [`SpillState::Active`] to reduce enum size.
pub(crate) struct SpillStateActive {
- /// Shared future for left-side spill data. All partitions wait on
- /// the same future — the first to poll triggers the actual spill.
- left_spill_fut: OnceFut<LeftSpillData>,
- /// Left input stream for incremental chunk reading (from spill file).
- /// None until `left_spill_fut` resolves.
+ /// The spilled left side, shared by every partition.
+ left_spill: Arc<LeftSpillData>,
+ /// Left input stream for incremental chunk reading (from the spill file).
+ /// None between chunk passes, when the previous stream has been consumed.
left_stream: Option<SendableRecordBatchStream>,
- /// Left-side schema (set once `left_spill_fut` resolves)
+ /// Left-side schema
left_schema: Option<SchemaRef>,
/// Memory reservation for left-side buffering
reservation: MemoryReservation,
@@ -1257,7 +1378,7 @@ pub(crate) struct NestedLoopJoinStream {
/// `Option` is used because memory-limited path requires resetting it.
pub(crate) right_data: Option<SendableRecordBatchStream>,
/// the build-side table data of the nested loop join
- pub(crate) left_data: OnceFut<JoinLeftData>,
+ pub(crate) left_data: OnceFut<LeftLoad>,
/// Projection to construct the output schema from the left and right tables.
/// Example:
/// - output_schema: ['a', 'c']
@@ -1600,7 +1721,7 @@ impl NestedLoopJoinStream {
filter: Option<JoinFilter>,
join_type: JoinType,
right_data: SendableRecordBatchStream,
- left_data: OnceFut<JoinLeftData>,
+ left_data: OnceFut<LeftLoad>,
column_indices: Vec<ColumnIndex>,
metrics: NestedLoopJoinMetrics,
batch_size: usize,
@@ -1636,68 +1757,24 @@ impl NestedLoopJoinStream {
matches!(self.spill_state, SpillState::Active(_))
}
- /// Check if we can fall back to memory-limited mode on this error.
- fn can_fallback_to_spill(&self, error: &datafusion_common::DataFusionError) -> bool {
- matches!(self.spill_state, SpillState::Pending { .. })
- && matches!(
- error.find_root(),
- datafusion_common::DataFusionError::ResourcesExhausted(_)
- )
- }
-
- /// Switch from the standard OnceFut path to memory-limited mode.
+ /// Enter memory-limited mode with the left side already spilled by the load.
///
- /// Uses the shared `left_spill_data` OnceAsync so that only the first
- /// partition to reach this point re-executes the left child and spills
- /// it to disk. Other partitions share the same spill file.
- fn initiate_fallback(&mut self) -> Result<()> {
- // Take ownership of Pending state
+ /// Every partition resolves the same shared `LeftLoad`, so they all read the
+ /// one spill file the load wrote; the left child is never executed a second
+ /// time.
+ fn enter_memory_limited_mode(
+ &mut self,
+ left_spill: Arc<LeftSpillData>,
+ ) -> Result<()> {
let SpillState::Pending {
- left_plan,
task_context: context,
- left_spill_data,
} = std::mem::replace(&mut self.spill_state, SpillState::Disabled)
else {
- return internal_err!("initiate_fallback called in non-Pending spill state");
+ return internal_err!(
+ "enter_memory_limited_mode called in non-Pending spill state"
+ );
};
- // Use OnceAsync to ensure only the first partition spills the left
- // side. Other partitions will get the same OnceFut that resolves
- // to the shared spill file.
- let left_spill_fut = left_spill_data.try_once(|| {
- let plan = Arc::clone(&left_plan);
- let ctx = Arc::clone(&context);
- let spill_metrics = self.metrics.spill_metrics.clone();
- Ok(async move {
- let mut stream = plan.execute(0, Arc::clone(&ctx))?;
- let schema = stream.schema();
- let left_spill_manager = SpillManager::new(
- ctx.runtime_env(),
- spill_metrics,
- Arc::clone(&schema),
- )
- .with_compression_type(ctx.session_config().spill_compression());
-
- let result = left_spill_manager
- .spill_record_batch_stream_and_return_max_batch_memory(
- &mut stream,
- "NestedLoopJoin left spill",
- )
- .await?;
-
- match result {
- Some((file, _max_batch_memory)) => Ok(LeftSpillData {
- spill_manager: left_spill_manager,
- spill_file: file,
- schema,
- }),
- None => {
- internal_err!("Left side produced no data to spill")
- }
- }
- })
- })?;
-
// Create reservation with can_spill for fair memory allocation
let reservation = MemoryConsumer::new("NestedLoopJoinLoad[fallback]".to_string())
.with_can_spill(true)
@@ -1714,12 +1791,12 @@ impl NestedLoopJoinStream {
let right_schema = self
.right_data
.as_ref()
- .expect("right_data must be present before fallback")
+ .expect("right_data must be present before entering memory-limited mode")
.schema();
let right_data = self
.right_data
.take()
- .expect("right_data must be present before fallback");
+ .expect("right_data must be present before entering memory-limited mode");
let right_spill_manager = SpillManager::new(
context.runtime_env(),
self.metrics.spill_metrics.clone(),
@@ -1728,7 +1805,7 @@ impl NestedLoopJoinStream {
.with_compression_type(context.session_config().spill_compression());
self.spill_state = SpillState::Active(Box::new(SpillStateActive {
- left_spill_fut,
+ left_spill,
left_stream: None,
left_schema: None,
reservation,
@@ -1766,28 +1843,25 @@ impl NestedLoopJoinStream {
} else {
// Standard path: use OnceFut
match self.left_data.get_shared(cx) {
- Poll::Ready(Ok(left_data)) => {
- self.buffered_left_data = Some(left_data);
- self.left_exhausted = true;
- self.state = NLJState::FetchingRight;
- ControlFlow::Continue(())
- }
- Poll::Ready(Err(e)) => {
- if self.can_fallback_to_spill(&e) {
+ Poll::Ready(Ok(load)) => match load.as_ref() {
+ LeftLoad::InMemory(left_data) => {
+ self.buffered_left_data = Some(Arc::clone(left_data));
+ self.left_exhausted = true;
+ self.state = NLJState::FetchingRight;
+ ControlFlow::Continue(())
+ }
+ LeftLoad::Spilled(left_spill) => {
debug!(
- "NestedLoopJoin: OnceFut failed with OOM, \
- falling back to memory-limited mode"
+ "NestedLoopJoin: left side exceeded the budget and was spilled, \
+ entering memory-limited mode"
);
- match self.initiate_fallback() {
+ match self.enter_memory_limited_mode(Arc::clone(left_spill)) {
Ok(()) => ControlFlow::Continue(()),
- Err(fallback_err) => {
- ControlFlow::Break(Poll::Ready(Some(Err(fallback_err))))
- }
+ Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
}
- } else {
- ControlFlow::Break(Poll::Ready(Some(Err(e))))
}
- }
+ },
+ Poll::Ready(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
Poll::Pending => ControlFlow::Break(Poll::Pending),
}
}
@@ -1808,31 +1882,21 @@ impl NestedLoopJoinStream {
);
};
- // On first entry (or after re-entry for a new chunk pass when
- // left_stream was consumed), wait for the shared left spill
- // future to resolve and then open a stream from the spill file.
+ // On first entry, and again for every new chunk pass once the previous
+ // stream was consumed, open a fresh stream over the left spill file.
if active.left_stream.is_none() {
- match active.left_spill_fut.get_shared(cx) {
- Poll::Ready(Ok(spill_data)) => {
- match spill_data
- .spill_manager
- .read_spill_as_stream(Arc::clone(&spill_data.spill_file), None)
- {
- Ok(stream) => {
- active.left_schema = Some(Arc::clone(&spill_data.schema));
- active.left_stream = Some(stream);
- }
- Err(e) => {
- return ControlFlow::Break(Poll::Ready(Some(Err(e))));
- }
- }
+ let spill_data = Arc::clone(&active.left_spill);
+ match spill_data
+ .spill_manager
+ .read_spill_as_stream(Arc::clone(&spill_data.spill_file), None)
+ {
+ Ok(stream) => {
+ active.left_schema = Some(Arc::clone(&spill_data.schema));
+ active.left_stream = Some(stream);
}
- Poll::Ready(Err(e)) => {
+ Err(e) => {
return ControlFlow::Break(Poll::Ready(Some(Err(e))));
}
- Poll::Pending => {
- return ControlFlow::Break(Poll::Pending);
- }
}
}
@@ -3359,6 +3423,100 @@ pub(crate) mod tests {
Arc::new(TestMemoryExec::update_cache(&source))
}
+ /// An input that can be executed only once: later executions yield no batches, the way a
+ /// stream backed by an external one-shot iterator behaves.
+ #[derive(Debug)]
+ struct OneShotExec {
+ inner: Arc<dyn ExecutionPlan>,
+ executions: Arc<AtomicUsize>,
+ }
+
+ impl DisplayAs for OneShotExec {
+ fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+ write!(f, "OneShotExec")
+ }
+ }
+
+ impl ExecutionPlan for OneShotExec {
+ fn name(&self) -> &str {
+ "OneShotExec"
+ }
+
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ self.inner.properties()
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.inner]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ if self.executions.fetch_add(1, Ordering::Relaxed) == 0 {
+ self.inner.execute(partition, context)
+ } else {
+ Ok(Box::pin(crate::stream::RecordBatchStreamAdapter::new(
+ self.inner.schema(),
+ futures::stream::empty(),
+ )))
+ }
+ }
+ }
+
+ /// The left side is spilled by the load that consumed it, so nothing asks the left child for
+ /// its batches a second time. Before that, the memory-limited fallback re-executed the child
+ /// and silently dropped every batch the first pass had already consumed.
+ #[tokio::test]
+ async fn memory_limited_left_side_reads_the_child_once() -> Result<()> {
+ let executions = Arc::new(AtomicUsize::new(0));
+ let left = Arc::new(OneShotExec {
+ inner: build_left_table(),
+ executions: Arc::clone(&executions),
+ });
+ let right = build_right_table();
+ let filter = prepare_join_filter();
+ let task_ctx = task_ctx_with_memory_limit(50, 16)?;
+
+ let (_, batches, metrics) =
+ join_collect(left, right, &JoinType::Inner, Some(filter), task_ctx).await?;
+
+ assert!(
+ metrics.spill_count().unwrap_or(0) > 0,
+ "expected the tight memory limit to spill the left side"
+ );
+ assert_eq!(
+ executions.load(Ordering::Relaxed),
+ 1,
+ "the left child must be executed exactly once"
+ );
+ // Same answer as the in-memory path (see test_nlj_memory_limited_inner_join).
+ allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r"
+ +----+----+----+----+----+----+
+ | a1 | b1 | c1 | a2 | b2 | c2 |
+ +----+----+----+----+----+----+
+ | 5 | 5 | 50 | 2 | 2 | 80 |
+ +----+----+----+----+----+----+
+ "));
+ Ok(())
+ }
+
fn build_left_table() -> Arc<dyn ExecutionPlan> {
build_table(
("a1", &vec![5, 9, 11]),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org