diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 548a0cf1db9b..3d966702fb10 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, - /// 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>, + build_side_data: OnceAsync, /// Information of index and left / right placement of columns column_indices: Vec, /// 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> { 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 { + spill_manager: Option, +) -> Result { 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 = 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, + 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, + pending: Option, + mut stream: SendableRecordBatchStream, + metrics: BuildProbeJoinMetrics, + reservation: &MemoryReservation, +) -> Result> { + 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), + /// 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), +} + +/// 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, - /// TaskContext for re-execution and SpillManager creation + /// TaskContext for reservations and SpillManager creation task_context: Arc, - /// 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>, }, - /// 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), } @@ -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, - /// 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, + /// 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, - /// Left-side schema (set once `left_spill_fut` resolves) + /// Left-side schema left_schema: Option, /// 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, /// the build-side table data of the nested loop join - pub(crate) left_data: OnceFut, + pub(crate) left_data: OnceFut, /// 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, join_type: JoinType, right_data: SendableRecordBatchStream, - left_data: OnceFut, + left_data: OnceFut, column_indices: Vec, 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, + ) -> 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, + executions: Arc, + } + + 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) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + 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 { build_table( ("a1", &vec![5, 9, 11]),