From 000b8a7f5486aa335f4a0eb94e80fef14a362307 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:45:33 +0300 Subject: [PATCH 1/8] chore: refactor `MaterializingSortMergeJoinStream` into generators --- .../src/joins/sort_merge_join/exec.rs | 4 +- .../sort_merge_join/materializing_stream.rs | 625 ++++++++++++------ 2 files changed, 443 insertions(+), 186 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 3b597323b2e7b..85e80d0d2cec0 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -546,7 +546,7 @@ impl ExecutionPlan for SortMergeJoinExec { context.runtime_env(), ) } else { - Ok(Box::pin(MaterializingSortMergeJoinStream::try_new( + MaterializingSortMergeJoinStream::try_new( Arc::clone(&self.schema), self.sort_options.clone(), self.null_equality, @@ -561,7 +561,7 @@ impl ExecutionPlan for SortMergeJoinExec { reservation, spill_manager, context.runtime_env(), - )?)) + ) } } 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 51cf38b9ab1f7..f149b54dc6c5e 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 @@ -40,7 +40,7 @@ use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; use crate::joins::utils::{JoinFilter, JoinKeyComparator}; use crate::metrics::RecordOutput; use crate::spill::spill_manager::SpillManager; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{PhysicalExpr, RecordBatchStream, SendableRecordBatchStream}; use arrow::array::{types::UInt64Type, *}; @@ -50,8 +50,8 @@ use arrow::compute::{ }; use arrow::datatypes::SchemaRef; use datafusion_common::cast::as_uint64_array; -use datafusion_common::{JoinType, NullEquality, Result, exec_err, internal_err}; -use datafusion_execution::SpillFile; +use datafusion_common::{exec_err, internal_err, DataFusionError, JoinType, NullEquality, Result}; +use datafusion_execution::{async_try_stream, SpillFile, TryEmitter}; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; @@ -560,28 +560,105 @@ impl JoinedRecordBatches { self.debug_assert_empty_consistency(); } } -impl RecordBatchStream for MaterializingSortMergeJoinStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) + +impl MaterializingSortMergeJoinStream { + #[expect(clippy::too_many_arguments)] + pub fn try_new( + schema: SchemaRef, + sort_options: Vec, + null_equality: NullEquality, + streamed: SendableRecordBatchStream, + buffered: SendableRecordBatchStream, + on_streamed: Vec>, + on_buffered: Vec>, + filter: Option, + join_type: JoinType, + batch_size: usize, + join_metrics: SortMergeJoinMetrics, + reservation: MemoryReservation, + spill_manager: SpillManager, + runtime_env: Arc, + ) -> Result { + let streamed_schema = streamed.schema(); + let buffered_schema = buffered.schema(); + debug_assert!( + matches!( + join_type, + JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full + ), + "MaterializingSortMergeJoinStream does not handle {join_type:?}; \ + semi/anti/mark joins use BitwiseSortMergeJoinStream" + ); + let mut this = Self { + state: SortMergeJoinState::Init, + sort_options, + null_equality, + schema: Arc::clone(&schema), + streamed_schema: Arc::clone(&streamed_schema), + buffered_schema, + streamed, + buffered, + streamed_batch: StreamedBatch::new_empty(streamed_schema), + buffered_data: BufferedData::default(), + streamed_joined: false, + buffered_joined: false, + streamed_state: StreamedState::Init, + buffered_state: BufferedState::Init, + current_ordering: Ordering::Equal, + on_streamed, + on_buffered, + filter, + joined_record_batches: JoinedRecordBatches { + joined_batches: BatchCoalescer::new(Arc::clone(&schema), batch_size) + .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), + filter_metadata: FilterMetadata::new(), + }, + output: BatchCoalescer::new(schema, batch_size) + .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), + batch_size, + join_type, + join_metrics, + 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), + }; + + let schema = Arc::clone(&this.schema); + let baseline_metrics = this.join_metrics.baseline_metrics(); + + let stream = async_try_stream(|mut emitter| async move { + let result = this.join(&mut emitter).await; + result + }); + // ObservedStream records the baseline metrics (output rows/batches, + // end time) exactly as the former hand-written poll_next did. + Ok(Box::pin(ObservedStream::new( + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + baseline_metrics, + None, + ))) } -} -impl Stream for MaterializingSortMergeJoinStream { - type Item = Result; - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - let join_time = self.join_metrics.join_time().clone(); - let _timer = join_time.timer(); + + /// Main loop + async fn join( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + // TODO - add join time metric loop { match &self.state { SortMergeJoinState::Init => { let streamed_exhausted = - self.streamed_state == StreamedState::Exhausted; + self.streamed_state == StreamedState::Exhausted; let buffered_exhausted = - self.buffered_state == BufferedState::Exhausted; + self.buffered_state == BufferedState::Exhausted; self.state = if streamed_exhausted && buffered_exhausted { SortMergeJoinState::Exhausted } else { @@ -605,27 +682,21 @@ impl Stream for MaterializingSortMergeJoinStream { self.join_type, ) { let accumulated = self.num_unfrozen_pairs() - + self - .joined_record_batches - .filter_metadata - .filter_mask - .len(); + + self + .joined_record_batches + .filter_metadata + .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))); - } - Poll::Ready(None) | Poll::Pending => {} + self.poll_spilled_batches_async(&needed).await?; + if let Some(batch) = self.process_filtered_batches2()? { + emitter.emit(batch).await; + continue; } } } @@ -646,26 +717,20 @@ impl Stream for MaterializingSortMergeJoinStream { } SortMergeJoinState::Polling => { if ![StreamedState::Exhausted, StreamedState::Ready] - .contains(&self.streamed_state) + .contains(&self.streamed_state) { - match self.poll_streamed_row(cx)? { - Poll::Ready(_) => {} - Poll::Pending => return Poll::Pending, - } + self.poll_streamed_row_async().await?; } if ![BufferedState::Exhausted, BufferedState::Ready] - .contains(&self.buffered_state) + .contains(&self.buffered_state) { - match self.poll_buffered_batches(cx)? { - Poll::Ready(_) => {} - Poll::Pending => return Poll::Pending, - } + self.poll_buffered_batches_async().await?; } let streamed_exhausted = - self.streamed_state == StreamedState::Exhausted; + self.streamed_state == StreamedState::Exhausted; let buffered_exhausted = - self.buffered_state == BufferedState::Exhausted; + self.buffered_state == BufferedState::Exhausted; if streamed_exhausted && buffered_exhausted { self.state = SortMergeJoinState::Exhausted; continue; @@ -678,8 +743,8 @@ impl Stream for MaterializingSortMergeJoinStream { // Verify metadata alignment before checking if we have batches to output self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + .filter_metadata + .debug_assert_metadata_aligned(); // For filtered joins, skip output and let Init state handle it if needs_deferred_filtering(&self.filter, self.join_type) { @@ -690,18 +755,19 @@ 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 - .has_completed_batch() + .joined_record_batches + .joined_batches + .has_completed_batch() { let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); + .joined_record_batches + .joined_batches + .next_completed_batch() + .expect("has_completed_batch was true"); (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); + .record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; } self.state = SortMergeJoinState::Init; } @@ -710,15 +776,15 @@ impl Stream for MaterializingSortMergeJoinStream { // 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))?; + .get_required_batch_indices(self.buffered_data.batches.len()); + self.poll_spilled_batches_async(&needed).await?; self.freeze_all()?; // Verify metadata alignment before checking if we have batches to output self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + .filter_metadata + .debug_assert_metadata_aligned(); // For filtered joins, skip output and let Init state handle it if needs_deferred_filtering(&self.filter, self.join_type) { @@ -727,18 +793,19 @@ impl Stream for MaterializingSortMergeJoinStream { // For non-filtered joins, only output if we have a completed batch if self - .joined_record_batches - .joined_batches - .has_completed_batch() + .joined_record_batches + .joined_batches + .has_completed_batch() { let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); + .joined_record_batches + .joined_batches + .next_completed_batch() + .expect("has_completed_batch was true"); (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); + .record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; } // Otherwise keep buffering (don't output yet) @@ -748,7 +815,7 @@ impl Stream for MaterializingSortMergeJoinStream { self.join_partial()?; if self.num_unfrozen_pairs() < self.batch_size - && self.buffered_data.scanning_finished() + && self.buffered_data.scanning_finished() { self.buffered_data.scanning_reset(); self.state = SortMergeJoinState::EmitReadyThenInit; @@ -757,135 +824,69 @@ impl Stream for MaterializingSortMergeJoinStream { } SortMergeJoinState::Exhausted => { let needed = - self.get_required_batch_indices(self.buffered_data.batches.len()); - ready!(self.poll_spilled_batches(cx, &needed))?; + self.get_required_batch_indices(self.buffered_data.batches.len()); + self.poll_spilled_batches_async(&needed).await?; self.freeze_all()?; // Verify metadata alignment before final output self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + .filter_metadata + .debug_assert_metadata_aligned(); // For filtered joins, must concat and filter ALL data at once if needs_deferred_filtering(&self.filter, self.join_type) - && !self.joined_record_batches.joined_batches.is_empty() + && !self.joined_record_batches.joined_batches.is_empty() { let record_batch = self.filter_joined_batch()?; (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); + .record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; } // For non-filtered joins, finish buffered data first if !self.joined_record_batches.joined_batches.is_empty() { self.joined_record_batches - .joined_batches - .finish_buffered_batch()?; + .joined_batches + .finish_buffered_batch()?; } // Output one completed batch at a time (stay in Exhausted until empty) if self - .joined_record_batches - .joined_batches - .has_completed_batch() + .joined_record_batches + .joined_batches + .has_completed_batch() { let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); + .joined_record_batches + .joined_batches + .next_completed_batch() + .expect("has_completed_batch was true"); (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); + .record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; } // Finally check self.output BatchCoalescer (used by filtered joins) - return if !self.output.is_empty() { + if !self.output.is_empty() { self.output.finish_buffered_batch()?; let record_batch = self - .output - .next_completed_batch() - .expect("Failed to get last batch"); + .output + .next_completed_batch() + .expect("Failed to get last batch"); (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - Poll::Ready(Some(Ok(record_batch))) + .record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; } else { - Poll::Ready(None) + return Ok(()); }; } } } } -} - -impl MaterializingSortMergeJoinStream { - #[expect(clippy::too_many_arguments)] - pub fn try_new( - schema: SchemaRef, - sort_options: Vec, - null_equality: NullEquality, - streamed: SendableRecordBatchStream, - buffered: SendableRecordBatchStream, - on_streamed: Vec>, - on_buffered: Vec>, - filter: Option, - join_type: JoinType, - batch_size: usize, - join_metrics: SortMergeJoinMetrics, - reservation: MemoryReservation, - spill_manager: SpillManager, - runtime_env: Arc, - ) -> Result { - let streamed_schema = streamed.schema(); - let buffered_schema = buffered.schema(); - debug_assert!( - matches!( - join_type, - JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full - ), - "MaterializingSortMergeJoinStream does not handle {join_type:?}; \ - semi/anti/mark joins use BitwiseSortMergeJoinStream" - ); - Ok(Self { - state: SortMergeJoinState::Init, - sort_options, - null_equality, - schema: Arc::clone(&schema), - streamed_schema: Arc::clone(&streamed_schema), - buffered_schema, - streamed, - buffered, - streamed_batch: StreamedBatch::new_empty(streamed_schema), - buffered_data: BufferedData::default(), - streamed_joined: false, - buffered_joined: false, - streamed_state: StreamedState::Init, - buffered_state: BufferedState::Init, - current_ordering: Ordering::Equal, - on_streamed, - on_buffered, - filter, - joined_record_batches: JoinedRecordBatches { - joined_batches: BatchCoalescer::new(Arc::clone(&schema), batch_size) - .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), - filter_metadata: FilterMetadata::new(), - }, - output: BatchCoalescer::new(schema, batch_size) - .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), - batch_size, - join_type, - join_metrics, - 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), - }) - } /// Build a comparator for streamed vs buffered head batch keys. fn rebuild_streamed_buffered_cmp(&mut self) -> Result<()> { @@ -933,20 +934,20 @@ impl MaterializingSortMergeJoinStream { self.freeze_all()?; self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + .filter_metadata + .debug_assert_metadata_aligned(); if !self.joined_record_batches.joined_batches.is_empty() { let out_filtered_batch = self.filter_joined_batch()?; self.output - .push_batch(out_filtered_batch) - .expect("Failed to push output batch"); + .push_batch(out_filtered_batch) + .expect("Failed to push output batch"); if self.output.has_completed_batch() { let record_batch = self - .output - .next_completed_batch() - .expect("Failed to get output batch"); + .output + .next_completed_batch() + .expect("Failed to get output batch"); (&record_batch).record_output(&self.join_metrics.baseline_metrics()); return Poll::Ready(Some(Ok(record_batch))); } @@ -955,6 +956,36 @@ impl MaterializingSortMergeJoinStream { Poll::Pending } + /// Process accumulated batches for filtered joins + /// + /// Freezes unfrozen pairs, applies deferred filtering, and outputs if ready. + /// Returns Poll::Ready with a batch if one is available, otherwise Poll::Pending. + fn process_filtered_batches2(&mut self) -> Result> { + self.freeze_all()?; + + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + if !self.joined_record_batches.joined_batches.is_empty() { + let out_filtered_batch = self.filter_joined_batch()?; + self.output + .push_batch(out_filtered_batch) + .expect("Failed to push output batch"); + + if self.output.has_completed_batch() { + let record_batch = self + .output + .next_completed_batch() + .expect("Failed to get output batch"); + (&record_batch).record_output(&self.join_metrics.baseline_metrics()); + return Ok(Some(record_batch)); + } + } + + Ok(None) + } + /// Identifies which buffered batches are needed for the upcoming freeze operation fn get_required_batch_indices(&self, buffered_freeze_count: usize) -> Vec { let mut needed = vec![]; @@ -997,8 +1028,8 @@ impl MaterializingSortMergeJoinStream { if let BufferedBatchState::Spilled(spill_file) = &bb.batch { if self.spill_stream.is_none() { let stream = self - .spill_manager - .read_spill_as_stream(Arc::clone(spill_file), None)?; + .spill_manager + .read_spill_as_stream(Arc::clone(spill_file), None)?; self.spill_stream = Some(stream); } @@ -1009,13 +1040,13 @@ impl MaterializingSortMergeJoinStream { 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); + 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()); + .peak_mem_used() + .set_max(self.reservation.size()); self.spill_stream = None; } @@ -1033,6 +1064,56 @@ impl MaterializingSortMergeJoinStream { Poll::Ready(Ok(())) } + /// Asynchronously reads spilled batches back into memory. + /// Only processes the required indices to avoid OOMs. + async fn poll_spilled_batches_async( + &mut self, + required_indices: &[usize], + ) -> 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(Arc::clone(spill_file), None)?; + self.spill_stream = Some(stream); + } + + match self.spill_stream.as_mut().unwrap().next().await.transpose()? { + Some(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; + } + None => { + self.spill_stream = None; + return internal_err!("Spill file was empty"); + } + } + } + } + + Ok(()) + } + /// Poll next streamed row fn poll_streamed_row(&mut self, cx: &mut Context) -> Poll>> { loop { @@ -1049,7 +1130,7 @@ impl MaterializingSortMergeJoinStream { } StreamedState::Polling => { let needed = - self.get_required_batch_indices(self.buffered_data.batches.len()); + 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))); } @@ -1062,7 +1143,7 @@ impl MaterializingSortMergeJoinStream { // Release the streamed input pipeline's resources. let streamed_schema = self.streamed.schema(); self.streamed = - Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + Box::pin(EmptyRecordBatchStream::new(streamed_schema)); self.streamed_state = StreamedState::Exhausted; } Poll::Ready(Some(batch)) => { @@ -1071,12 +1152,12 @@ impl MaterializingSortMergeJoinStream { 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); + 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); + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); self.streamed_state = StreamedState::Ready; } } @@ -1092,6 +1173,57 @@ impl MaterializingSortMergeJoinStream { } } + /// Poll next streamed row + async fn poll_streamed_row_async(&mut self) -> Result<()> { + loop { + match &self.streamed_state { + StreamedState::Init => { + if self.streamed_batch.idx + 1 < self.streamed_batch.batch.num_rows() + { + self.streamed_batch.idx += 1; + self.streamed_state = StreamedState::Ready; + return Ok(()); + } else { + self.streamed_state = StreamedState::Polling; + } + } + StreamedState::Polling => { + let needed = + self.get_required_batch_indices(self.buffered_data.batches.len()); + self.poll_spilled_batches_async(&needed).await?; + + match self.streamed.next().await.transpose()? { + 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; + } + 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 | StreamedState::Exhausted => { + return Ok(()); + } + } + } + } + fn free_reservation(&mut self, buffered_batch: &BufferedBatch) { if buffered_batch.reserved_amount > 0 { self.reservation.shrink(buffered_batch.reserved_amount); @@ -1170,7 +1302,7 @@ impl MaterializingSortMergeJoinStream { self.freeze_dequeuing_buffered()?; if let Some(mut buffered_batch) = - self.buffered_data.batches.pop_front() + self.buffered_data.batches.pop_front() { self.produce_buffered_not_matched(&mut buffered_batch)?; self.free_reservation(&buffered_batch); @@ -1209,7 +1341,7 @@ impl MaterializingSortMergeJoinStream { // Release the buffered input pipeline's resources. let buffered_schema = self.buffered.schema(); self.buffered = - Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); self.buffered_state = BufferedState::Exhausted; return Poll::Ready(None); } @@ -1219,7 +1351,7 @@ impl MaterializingSortMergeJoinStream { if batch.num_rows() > 0 { let buffered_batch = - BufferedBatch::new(batch, 0..1, &self.on_buffered); + BufferedBatch::new(batch, 0..1, &self.on_buffered); self.allocate_reservation(buffered_batch)?; self.streamed_buffered_cmp = None; @@ -1229,13 +1361,13 @@ impl MaterializingSortMergeJoinStream { }, BufferedState::PollingRest => { if self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows + < self.buffered_data.tail_batch().num_rows { if self.buffered_equality_cmp.is_none() { self.rebuild_buffered_equality_cmp()?; } while self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows + < self.buffered_data.tail_batch().num_rows { if self.buffered_equality_cmp.as_ref().unwrap().is_equal( self.buffered_data.head_batch().range.start, @@ -1287,6 +1419,131 @@ impl MaterializingSortMergeJoinStream { } } + /// Poll next buffered batches + async fn poll_buffered_batches_async(&mut self) -> Result<()> { + loop { + match &self.buffered_state { + BufferedState::Init => { + // pop previous buffered batches + let mut head_changed = false; + while !self.buffered_data.batches.is_empty() { + 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); + self.poll_spilled_batches_async(&needed).await?; + + 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 { + // If the head batch is not fully processed, break the loop. + // Streamed batch will be joined with the head batch in the next step. + break; + } + } + if head_changed { + self.streamed_buffered_cmp = None; + self.buffered_equality_cmp = None; + } + if self.buffered_data.batches.is_empty() { + self.buffered_state = BufferedState::PollingFirst; + } else { + let tail_batch = self.buffered_data.tail_batch_mut(); + tail_batch.range.start = tail_batch.range.end; + tail_batch.range.end += 1; + self.buffered_state = BufferedState::PollingRest; + } + } + BufferedState::PollingFirst => match self.buffered.next().await.transpose()? { + None => { + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + self.buffered_state = BufferedState::Exhausted; + return Ok(()); + } + Some(batch) => { + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + + if batch.num_rows() > 0 { + let buffered_batch = + BufferedBatch::new(batch, 0..1, &self.on_buffered); + + self.allocate_reservation(buffered_batch)?; + self.streamed_buffered_cmp = None; + self.buffered_state = BufferedState::PollingRest; + } + } + }, + BufferedState::PollingRest => { + if self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.is_none() { + self.rebuild_buffered_equality_cmp()?; + } + while self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.as_ref().unwrap().is_equal( + self.buffered_data.head_batch().range.start, + self.buffered_data.tail_batch().range.end, + ) { + self.buffered_data.tail_batch_mut().range.end += 1; + } else { + self.buffered_state = BufferedState::Ready; + return Ok(()); + } + } + } else { + match self.buffered.next().await.transpose()? { + None => { + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = Box::pin(EmptyRecordBatchStream::new( + buffered_schema, + )); + self.buffered_state = BufferedState::Ready; + } + Some(batch) => { + // Polling batches coming concurrently as multiple partitions + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + if batch.num_rows() > 0 { + let buffered_batch = BufferedBatch::new( + batch, + 0..0, + &self.on_buffered, + ); + self.allocate_reservation(buffered_batch)?; + self.buffered_equality_cmp = None; + } + } + } + } + } + BufferedState::Ready | BufferedState::Exhausted => { + return Ok(()); + } + } + } + } + /// Get comparison result of streamed row and buffered batches fn compare_streamed_buffered(&mut self) -> Result { if self.streamed_state == StreamedState::Exhausted { From 88e24715372d20cb7d9c0ab720488a30f627e8b3 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:46:44 +0300 Subject: [PATCH 2/8] remove duplicate functions --- .../sort_merge_join/materializing_stream.rs | 282 ------------------ 1 file changed, 282 deletions(-) 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 f149b54dc6c5e..1fa0ee69c354d 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 @@ -644,8 +644,6 @@ impl MaterializingSortMergeJoinStream { ))) } - - /// Main loop async fn join( &mut self, @@ -926,36 +924,6 @@ impl MaterializingSortMergeJoinStream { self.streamed_batch.num_output_rows() } - /// Process accumulated batches for filtered joins - /// - /// Freezes unfrozen pairs, applies deferred filtering, and outputs if ready. - /// Returns Poll::Ready with a batch if one is available, otherwise Poll::Pending. - fn process_filtered_batches(&mut self) -> Poll>> { - self.freeze_all()?; - - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); - - if !self.joined_record_batches.joined_batches.is_empty() { - let out_filtered_batch = self.filter_joined_batch()?; - self.output - .push_batch(out_filtered_batch) - .expect("Failed to push output batch"); - - if self.output.has_completed_batch() { - let record_batch = self - .output - .next_completed_batch() - .expect("Failed to get output batch"); - (&record_batch).record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); - } - } - - Poll::Pending - } - /// Process accumulated batches for filtered joins /// /// Freezes unfrozen pairs, applies deferred filtering, and outputs if ready. @@ -1010,60 +978,6 @@ impl MaterializingSortMergeJoinStream { 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> { - 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(Arc::clone(spill_file), 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(())) - } - /// Asynchronously reads spilled batches back into memory. /// Only processes the required indices to avoid OOMs. async fn poll_spilled_batches_async( @@ -1114,65 +1028,6 @@ impl MaterializingSortMergeJoinStream { Ok(()) } - /// Poll next streamed row - fn poll_streamed_row(&mut self, cx: &mut Context) -> Poll>> { - loop { - match &self.streamed_state { - StreamedState::Init => { - if self.streamed_batch.idx + 1 < self.streamed_batch.batch.num_rows() - { - self.streamed_batch.idx += 1; - self.streamed_state = StreamedState::Ready; - return Poll::Ready(Some(Ok(()))); - } else { - self.streamed_state = StreamedState::Polling; - } - } - 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))); - } - - 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(()))); - } - StreamedState::Exhausted => { - return Poll::Ready(None); - } - } - } - } - /// Poll next streamed row async fn poll_streamed_row_async(&mut self) -> Result<()> { loop { @@ -1282,143 +1137,6 @@ impl MaterializingSortMergeJoinStream { Ok(()) } - /// Poll next buffered batches - fn poll_buffered_batches(&mut self, cx: &mut Context) -> Poll>> { - loop { - match &self.buffered_state { - BufferedState::Init => { - // pop previous buffered batches - let mut head_changed = false; - while !self.buffered_data.batches.is_empty() { - 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 { - // If the head batch is not fully processed, break the loop. - // Streamed batch will be joined with the head batch in the next step. - break; - } - } - if head_changed { - self.streamed_buffered_cmp = None; - self.buffered_equality_cmp = None; - } - if self.buffered_data.batches.is_empty() { - self.buffered_state = BufferedState::PollingFirst; - } else { - let tail_batch = self.buffered_data.tail_batch_mut(); - tail_batch.range.start = tail_batch.range.end; - tail_batch.range.end += 1; - self.buffered_state = BufferedState::PollingRest; - } - } - BufferedState::PollingFirst => match self.buffered.poll_next_unpin(cx)? { - Poll::Pending => { - return Poll::Pending; - } - Poll::Ready(None) => { - // Release the buffered input pipeline's resources. - let buffered_schema = self.buffered.schema(); - self.buffered = - Box::pin(EmptyRecordBatchStream::new(buffered_schema)); - self.buffered_state = BufferedState::Exhausted; - return Poll::Ready(None); - } - Poll::Ready(Some(batch)) => { - self.join_metrics.input_batches().add(1); - self.join_metrics.input_rows().add(batch.num_rows()); - - if batch.num_rows() > 0 { - let buffered_batch = - BufferedBatch::new(batch, 0..1, &self.on_buffered); - - self.allocate_reservation(buffered_batch)?; - self.streamed_buffered_cmp = None; - self.buffered_state = BufferedState::PollingRest; - } - } - }, - BufferedState::PollingRest => { - if self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows - { - if self.buffered_equality_cmp.is_none() { - self.rebuild_buffered_equality_cmp()?; - } - while self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows - { - if self.buffered_equality_cmp.as_ref().unwrap().is_equal( - self.buffered_data.head_batch().range.start, - self.buffered_data.tail_batch().range.end, - ) { - self.buffered_data.tail_batch_mut().range.end += 1; - } else { - self.buffered_state = BufferedState::Ready; - return Poll::Ready(Some(Ok(()))); - } - } - } else { - match self.buffered.poll_next_unpin(cx)? { - Poll::Pending => { - return Poll::Pending; - } - Poll::Ready(None) => { - // Release the buffered input pipeline's resources. - let buffered_schema = self.buffered.schema(); - self.buffered = Box::pin(EmptyRecordBatchStream::new( - buffered_schema, - )); - self.buffered_state = BufferedState::Ready; - } - Poll::Ready(Some(batch)) => { - // Polling batches coming concurrently as multiple partitions - self.join_metrics.input_batches().add(1); - self.join_metrics.input_rows().add(batch.num_rows()); - if batch.num_rows() > 0 { - let buffered_batch = BufferedBatch::new( - batch, - 0..0, - &self.on_buffered, - ); - self.allocate_reservation(buffered_batch)?; - self.buffered_equality_cmp = None; - } - } - } - } - } - BufferedState::Ready => { - return Poll::Ready(Some(Ok(()))); - } - BufferedState::Exhausted => { - return Poll::Ready(None); - } - } - } - } - /// Poll next buffered batches async fn poll_buffered_batches_async(&mut self) -> Result<()> { loop { From de03b851d99a3f9cd91024d2096c97c2f1349b4d Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:00:57 +0300 Subject: [PATCH 3/8] extract on exhausted --- .../sort_merge_join/materializing_stream.rs | 377 +++++++++--------- 1 file changed, 198 insertions(+), 179 deletions(-) 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 1fa0ee69c354d..07ed00498dcc8 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 @@ -50,10 +50,12 @@ use arrow::compute::{ }; use arrow::datatypes::SchemaRef; use datafusion_common::cast::as_uint64_array; -use datafusion_common::{exec_err, internal_err, DataFusionError, JoinType, NullEquality, Result}; -use datafusion_execution::{async_try_stream, SpillFile, TryEmitter}; +use datafusion_common::{ + DataFusionError, JoinType, NullEquality, Result, exec_err, internal_err, +}; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_execution::{SpillFile, TryEmitter, async_try_stream}; use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use futures::{Stream, StreamExt, ready}; @@ -637,6 +639,8 @@ impl MaterializingSortMergeJoinStream { }); // ObservedStream records the baseline metrics (output rows/batches, // end time) exactly as the former hand-written poll_next did. + // TODO - some have (&record_batch).record_output(&self.join_metrics.baseline_metrics()); + // REMOVE THOSE OR DECIDE SOMETHING TO AVOID DOUBLE COUNTING Ok(Box::pin(ObservedStream::new( Box::pin(RecordBatchStreamAdapter::new(schema, stream)), baseline_metrics, @@ -646,7 +650,7 @@ impl MaterializingSortMergeJoinStream { /// Main loop async fn join( - &mut self, + mut self, emitter: &mut TryEmitter, ) -> Result<()> { // TODO - add join time metric @@ -654,84 +658,80 @@ impl MaterializingSortMergeJoinStream { match &self.state { SortMergeJoinState::Init => { let streamed_exhausted = - self.streamed_state == StreamedState::Exhausted; + self.streamed_state == StreamedState::Exhausted; let buffered_exhausted = - self.buffered_state == BufferedState::Exhausted; - self.state = if streamed_exhausted && buffered_exhausted { - SortMergeJoinState::Exhausted - } else { - match self.current_ordering { - Ordering::Less | Ordering::Equal => { - if !streamed_exhausted { - // Batch deferred filtering: process_filtered_batches() - // only when >= batch_size rows have accumulated. - // Without this gate, unique keys cause per-row pipeline - // execution (concat + correct_mask + filter_by_type), - // which dominates runtime. - // - // Accumulated rows are bounded to ~2*batch_size: - // one batch_size worth from freeze_dequeuing_buffered() - // (when an input batch is fully consumed), plus up to - // batch_size pairs accumulating toward the next freeze. - // This does not reintroduce the unbounded buffering - // fixed by PR #20482. Exhausted state flushes remainder. - if needs_deferred_filtering( - &self.filter, - self.join_type, - ) { - let accumulated = self.num_unfrozen_pairs() - + self - .joined_record_batches - .filter_metadata - .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(), - ); - self.poll_spilled_batches_async(&needed).await?; - if let Some(batch) = self.process_filtered_batches2()? { - emitter.emit(batch).await; - continue; - } + self.buffered_state == BufferedState::Exhausted; + if streamed_exhausted && buffered_exhausted { + break; + } + match self.current_ordering { + Ordering::Less | Ordering::Equal => { + if !streamed_exhausted { + // Batch deferred filtering: process_filtered_batches() + // only when >= batch_size rows have accumulated. + // Without this gate, unique keys cause per-row pipeline + // execution (concat + correct_mask + filter_by_type), + // which dominates runtime. + // + // Accumulated rows are bounded to ~2*batch_size: + // one batch_size worth from freeze_dequeuing_buffered() + // (when an input batch is fully consumed), plus up to + // batch_size pairs accumulating toward the next freeze. + // This does not reintroduce the unbounded buffering + // fixed by PR #20482. Exhausted state flushes remainder. + if needs_deferred_filtering(&self.filter, self.join_type) + { + let accumulated = self.num_unfrozen_pairs() + + self + .joined_record_batches + .filter_metadata + .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(), + ); + self.poll_spilled_batches_async(&needed).await?; + if let Some(batch) = + self.process_filtered_batches2()? + { + emitter.emit(batch).await; + continue; } } - - self.streamed_joined = false; - self.streamed_state = StreamedState::Init; } + + self.streamed_joined = false; + self.streamed_state = StreamedState::Init; } - Ordering::Greater => { - if !buffered_exhausted { - self.buffered_joined = false; - self.buffered_state = BufferedState::Init; - } + } + Ordering::Greater => { + if !buffered_exhausted { + self.buffered_joined = false; + self.buffered_state = BufferedState::Init; } } - SortMergeJoinState::Polling - }; + } + self.state = SortMergeJoinState::Polling; } SortMergeJoinState::Polling => { if ![StreamedState::Exhausted, StreamedState::Ready] - .contains(&self.streamed_state) + .contains(&self.streamed_state) { self.poll_streamed_row_async().await?; } if ![BufferedState::Exhausted, BufferedState::Ready] - .contains(&self.buffered_state) + .contains(&self.buffered_state) { self.poll_buffered_batches_async().await?; } - let streamed_exhausted = - self.streamed_state == StreamedState::Exhausted; - let buffered_exhausted = - self.buffered_state == BufferedState::Exhausted; - if streamed_exhausted && buffered_exhausted { - self.state = SortMergeJoinState::Exhausted; - continue; + if self.streamed_state == StreamedState::Exhausted + && self.buffered_state == BufferedState::Exhausted + { + break; } self.current_ordering = self.compare_streamed_buffered()?; self.state = SortMergeJoinState::JoinOutput; @@ -741,8 +741,8 @@ impl MaterializingSortMergeJoinStream { // Verify metadata alignment before checking if we have batches to output self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + .filter_metadata + .debug_assert_metadata_aligned(); // For filtered joins, skip output and let Init state handle it if needs_deferred_filtering(&self.filter, self.join_type) { @@ -753,17 +753,17 @@ impl 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 - .has_completed_batch() + .joined_record_batches + .joined_batches + .has_completed_batch() { let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); + .joined_record_batches + .joined_batches + .next_completed_batch() + .expect("has_completed_batch was true"); (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); + .record_output(&self.join_metrics.baseline_metrics()); emitter.emit(record_batch).await; continue; } @@ -774,15 +774,15 @@ impl MaterializingSortMergeJoinStream { // 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()); + .get_required_batch_indices(self.buffered_data.batches.len()); self.poll_spilled_batches_async(&needed).await?; self.freeze_all()?; // Verify metadata alignment before checking if we have batches to output self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + .filter_metadata + .debug_assert_metadata_aligned(); // For filtered joins, skip output and let Init state handle it if needs_deferred_filtering(&self.filter, self.join_type) { @@ -791,17 +791,17 @@ impl MaterializingSortMergeJoinStream { // For non-filtered joins, only output if we have a completed batch if self - .joined_record_batches - .joined_batches - .has_completed_batch() + .joined_record_batches + .joined_batches + .has_completed_batch() { let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); + .joined_record_batches + .joined_batches + .next_completed_batch() + .expect("has_completed_batch was true"); (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); + .record_output(&self.join_metrics.baseline_metrics()); emitter.emit(record_batch).await; continue; } @@ -813,7 +813,7 @@ impl MaterializingSortMergeJoinStream { self.join_partial()?; if self.num_unfrozen_pairs() < self.batch_size - && self.buffered_data.scanning_finished() + && self.buffered_data.scanning_finished() { self.buffered_data.scanning_reset(); self.state = SortMergeJoinState::EmitReadyThenInit; @@ -821,69 +821,79 @@ impl MaterializingSortMergeJoinStream { // 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()); - self.poll_spilled_batches_async(&needed).await?; + unreachable!("should be out of the loop"); + } + } + } - self.freeze_all()?; + self.on_children_exhausted(emitter).await + } - // Verify metadata alignment before final output - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + async fn on_children_exhausted( + mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + loop { + let needed = + self.get_required_batch_indices(self.buffered_data.batches.len()); + self.poll_spilled_batches_async(&needed).await?; - // For filtered joins, must concat and filter ALL data at once - if needs_deferred_filtering(&self.filter, self.join_type) - && !self.joined_record_batches.joined_batches.is_empty() - { - let record_batch = self.filter_joined_batch()?; - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - emitter.emit(record_batch).await; - continue; - } + self.freeze_all()?; - // For non-filtered joins, finish buffered data first - if !self.joined_record_batches.joined_batches.is_empty() { - self.joined_record_batches - .joined_batches - .finish_buffered_batch()?; - } + // Verify metadata alignment before final output + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); - // Output one completed batch at a time (stay in Exhausted until empty) - if self - .joined_record_batches - .joined_batches - .has_completed_batch() - { - let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - emitter.emit(record_batch).await; - continue; - } + // For filtered joins, must concat and filter ALL data at once + if needs_deferred_filtering(&self.filter, self.join_type) + && !self.joined_record_batches.joined_batches.is_empty() + { + let record_batch = self.filter_joined_batch()?; + (&record_batch).record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; + } - // Finally check self.output BatchCoalescer (used by filtered joins) - if !self.output.is_empty() { - self.output.finish_buffered_batch()?; - let record_batch = self - .output - .next_completed_batch() - .expect("Failed to get last batch"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - emitter.emit(record_batch).await; - continue; - } else { - return Ok(()); - }; - } + // For non-filtered joins, finish buffered data first + if !self.joined_record_batches.joined_batches.is_empty() { + self.joined_record_batches + .joined_batches + .finish_buffered_batch()?; + } + + // Output one completed batch at a time (stay in Exhausted until empty) + if self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + let record_batch = self + .joined_record_batches + .joined_batches + .next_completed_batch() + .expect("has_completed_batch was true"); + (&record_batch).record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; } + + // Finally check self.output BatchCoalescer (used by filtered joins) + if !self.output.is_empty() { + self.output.finish_buffered_batch()?; + let record_batch = self + .output + .next_completed_batch() + .expect("Failed to get last batch"); + (&record_batch).record_output(&self.join_metrics.baseline_metrics()); + emitter.emit(record_batch).await; + continue; + } else { + break; + }; } + + Ok(()) } /// Build a comparator for streamed vs buffered head batch keys. @@ -932,20 +942,20 @@ impl MaterializingSortMergeJoinStream { self.freeze_all()?; self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + .filter_metadata + .debug_assert_metadata_aligned(); if !self.joined_record_batches.joined_batches.is_empty() { let out_filtered_batch = self.filter_joined_batch()?; self.output - .push_batch(out_filtered_batch) - .expect("Failed to push output batch"); + .push_batch(out_filtered_batch) + .expect("Failed to push output batch"); if self.output.has_completed_batch() { let record_batch = self - .output - .next_completed_batch() - .expect("Failed to get output batch"); + .output + .next_completed_batch() + .expect("Failed to get output batch"); (&record_batch).record_output(&self.join_metrics.baseline_metrics()); return Ok(Some(record_batch)); } @@ -995,25 +1005,32 @@ impl MaterializingSortMergeJoinStream { if let BufferedBatchState::Spilled(spill_file) = &bb.batch { if self.spill_stream.is_none() { let stream = self - .spill_manager - .read_spill_as_stream(Arc::clone(spill_file), None)?; + .spill_manager + .read_spill_as_stream(Arc::clone(spill_file), None)?; self.spill_stream = Some(stream); } - match self.spill_stream.as_mut().unwrap().next().await.transpose()? { + match self + .spill_stream + .as_mut() + .unwrap() + .next() + .await + .transpose()? + { Some(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); + 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()); + .peak_mem_used() + .set_max(self.reservation.size()); self.spill_stream = None; } @@ -1044,7 +1061,7 @@ impl MaterializingSortMergeJoinStream { } StreamedState::Polling => { let needed = - self.get_required_batch_indices(self.buffered_data.batches.len()); + self.get_required_batch_indices(self.buffered_data.batches.len()); self.poll_spilled_batches_async(&needed).await?; match self.streamed.next().await.transpose()? { @@ -1052,7 +1069,7 @@ impl MaterializingSortMergeJoinStream { // Release the streamed input pipeline's resources. let streamed_schema = self.streamed.schema(); self.streamed = - Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + Box::pin(EmptyRecordBatchStream::new(streamed_schema)); self.streamed_state = StreamedState::Exhausted; } Some(batch) => { @@ -1061,12 +1078,12 @@ impl MaterializingSortMergeJoinStream { 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); + 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); + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); self.streamed_state = StreamedState::Ready; } } @@ -1154,7 +1171,7 @@ impl MaterializingSortMergeJoinStream { self.freeze_dequeuing_buffered()?; if let Some(mut buffered_batch) = - self.buffered_data.batches.pop_front() + self.buffered_data.batches.pop_front() { self.produce_buffered_not_matched(&mut buffered_batch)?; self.free_reservation(&buffered_batch); @@ -1185,38 +1202,40 @@ impl MaterializingSortMergeJoinStream { self.buffered_state = BufferedState::PollingRest; } } - BufferedState::PollingFirst => match self.buffered.next().await.transpose()? { - None => { - // Release the buffered input pipeline's resources. - let buffered_schema = self.buffered.schema(); - self.buffered = - Box::pin(EmptyRecordBatchStream::new(buffered_schema)); - self.buffered_state = BufferedState::Exhausted; - return Ok(()); - } - Some(batch) => { - self.join_metrics.input_batches().add(1); - self.join_metrics.input_rows().add(batch.num_rows()); + BufferedState::PollingFirst => { + match self.buffered.next().await.transpose()? { + None => { + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + self.buffered_state = BufferedState::Exhausted; + return Ok(()); + } + Some(batch) => { + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); - if batch.num_rows() > 0 { - let buffered_batch = - BufferedBatch::new(batch, 0..1, &self.on_buffered); + if batch.num_rows() > 0 { + let buffered_batch = + BufferedBatch::new(batch, 0..1, &self.on_buffered); - self.allocate_reservation(buffered_batch)?; - self.streamed_buffered_cmp = None; - self.buffered_state = BufferedState::PollingRest; + self.allocate_reservation(buffered_batch)?; + self.streamed_buffered_cmp = None; + self.buffered_state = BufferedState::PollingRest; + } } } - }, + } BufferedState::PollingRest => { if self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows + < self.buffered_data.tail_batch().num_rows { if self.buffered_equality_cmp.is_none() { self.rebuild_buffered_equality_cmp()?; } while self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows + < self.buffered_data.tail_batch().num_rows { if self.buffered_equality_cmp.as_ref().unwrap().is_equal( self.buffered_data.head_batch().range.start, From ccdd36340b13836fbae8289748495bdad236d042 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:53:22 +0300 Subject: [PATCH 4/8] simplify --- .../src/joins/sort_merge_join/exec.rs | 6 +- .../sort_merge_join/materializing_stream.rs | 812 ++++++++---------- .../src/joins/sort_merge_join/metrics.rs | 3 + 3 files changed, 387 insertions(+), 434 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 85e80d0d2cec0..1abcd9d6c7ce4 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -81,8 +81,7 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequiremen /// on the output batch size of the execution plan. There is no spilling support for streamed input. /// The comparisons are performed from values of join keys in streamed input with the values of /// join keys in buffered input. One row in streamed record batch could be matched with multiple rows in -/// buffered input batches. The streamed input is managed through the states in `StreamedState` -/// and streamed input batches are represented by `StreamedBatch`. +/// buffered input batches. Streamed input batches are represented by `StreamedBatch`. /// /// Buffered input is buffered for all record batches having the same value of join key. /// If the memory limit increases beyond the specified value and spilling is enabled, @@ -92,8 +91,7 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequiremen /// memory/disk depends on the number of rows of buffered input having the same value /// of join key as that of streamed input rows currently present in memory. Due to pre-sorted inputs, /// the algorithm understands when it is not needed anymore, and releases the buffered batches -/// from memory/disk. The buffered input is managed through the states in `BufferedState` -/// and buffered input batches are represented by `BufferedBatch`. +/// from memory/disk. Buffered input batches are represented by `BufferedBatch`. /// /// Depending on the type of join, left or right input may be selected as streamed or buffered /// respectively. For example, in a left-outer join, the left execution plan will be selected as 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 07ed00498dcc8..b7962717d8b0e 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 @@ -26,11 +26,9 @@ use std::collections::{HashMap, VecDeque}; use std::fmt::Debug; use std::mem::size_of; use std::ops::Range; -use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::Relaxed; -use std::task::{Context, Poll}; use crate::joins::sort_merge_join::filter::{ FilterMetadata, filter_record_batch_by_join_type, get_corrected_filter_mask, @@ -38,10 +36,9 @@ use crate::joins::sort_merge_join::filter::{ }; use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; use crate::joins::utils::{JoinFilter, JoinKeyComparator}; -use crate::metrics::RecordOutput; use crate::spill::spill_manager::SpillManager; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; -use crate::{PhysicalExpr, RecordBatchStream, SendableRecordBatchStream}; +use crate::{PhysicalExpr, SendableRecordBatchStream}; use arrow::array::{types::UInt64Type, *}; use arrow::compute::{ @@ -58,50 +55,7 @@ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_execution::{SpillFile, TryEmitter, async_try_stream}; use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; -use futures::{Stream, StreamExt, ready}; - -/// State of SMJ stream -#[derive(Debug, PartialEq, Eq)] -pub(super) enum SortMergeJoinState { - /// Init joining with a new streamed row or a new buffered batches - Init, - /// Polling one streamed row or one buffered batch, or both - Polling, - /// Joining polled data and making output - JoinOutput, - /// Emit ready data if have any and then go back to [`Self::Init`] state - EmitReadyThenInit, - /// No more output - Exhausted, -} - -/// State of streamed data stream -#[derive(Debug, PartialEq, Eq)] -pub(super) enum StreamedState { - /// Init polling - Init, - /// Polling one streamed row - Polling, - /// Ready to produce one streamed row - Ready, - /// No more streamed row - Exhausted, -} - -/// State of buffered data stream -#[derive(Debug, PartialEq, Eq)] -pub(super) enum BufferedState { - /// Init polling - Init, - /// Polling first row in the next batch - PollingFirst, - /// Polling rest rows in the next batch - PollingRest, - /// Ready to produce one batch - Ready, - /// No more buffered batches - Exhausted, -} +use futures::StreamExt; /// Represents a chunk of joined data from streamed and buffered side pub(super) struct StreamedJoinedChunk { @@ -337,6 +291,9 @@ pub(super) struct MaterializingSortMergeJoinStream { pub filter: Option, /// How the join is performed pub join_type: JoinType, + /// Cached `needs_deferred_filtering(filter, join_type)` — both inputs + /// are fixed at construction time. + pub deferred_filtering: bool, /// Target output batch size pub batch_size: usize, @@ -352,8 +309,8 @@ pub(super) struct MaterializingSortMergeJoinStream { pub streamed_batch: StreamedBatch, /// (used in outer join) Is current streamed row joined at least once? pub streamed_joined: bool, - /// State of streamed - pub streamed_state: StreamedState, + /// True once the streamed input has no more rows + pub streamed_exhausted: bool, /// Join key columns of streamed pub on_streamed: Vec, @@ -369,8 +326,8 @@ pub(super) struct MaterializingSortMergeJoinStream { pub buffered_data: BufferedData, /// (used in outer join) Is current buffered batches joined at least once? pub buffered_joined: bool, - /// State of buffered - pub buffered_state: BufferedState, + /// True once the buffered input has no more rows and no group remains + pub buffered_exhausted: bool, /// Join key columns of buffered pub on_buffered: Vec, @@ -379,8 +336,6 @@ pub(super) struct MaterializingSortMergeJoinStream { // These fields track the execution state of merge join and are updated // during the execution. // ======================================================================== - /// Current state of the stream - pub state: SortMergeJoinState, /// Staging output array builders pub joined_record_batches: JoinedRecordBatches, /// Output buffer. Currently used by filtering as it requires double buffering @@ -592,7 +547,6 @@ impl MaterializingSortMergeJoinStream { semi/anti/mark joins use BitwiseSortMergeJoinStream" ); let mut this = Self { - state: SortMergeJoinState::Init, sort_options, null_equality, schema: Arc::clone(&schema), @@ -604,11 +558,12 @@ impl MaterializingSortMergeJoinStream { buffered_data: BufferedData::default(), streamed_joined: false, buffered_joined: false, - streamed_state: StreamedState::Init, - buffered_state: BufferedState::Init, + streamed_exhausted: false, + buffered_exhausted: false, current_ordering: Ordering::Equal, on_streamed, on_buffered, + deferred_filtering: needs_deferred_filtering(&filter, join_type), filter, joined_record_batches: JoinedRecordBatches { joined_batches: BatchCoalescer::new(Arc::clone(&schema), batch_size) @@ -633,14 +588,10 @@ impl MaterializingSortMergeJoinStream { let schema = Arc::clone(&this.schema); let baseline_metrics = this.join_metrics.baseline_metrics(); - let stream = async_try_stream(|mut emitter| async move { - let result = this.join(&mut emitter).await; - result - }); + let stream = + async_try_stream(|mut emitter| async move { this.join(&mut emitter).await }); // ObservedStream records the baseline metrics (output rows/batches, // end time) exactly as the former hand-written poll_next did. - // TODO - some have (&record_batch).record_output(&self.join_metrics.baseline_metrics()); - // REMOVE THOSE OR DECIDE SOMETHING TO AVOID DOUBLE COUNTING Ok(Box::pin(ObservedStream::new( Box::pin(RecordBatchStreamAdapter::new(schema, stream)), baseline_metrics, @@ -648,249 +599,205 @@ impl MaterializingSortMergeJoinStream { ))) } - /// Main loop + /// Main loop: a merge scan over the two sorted inputs. + /// + /// Compares the current streamed row with the current buffered key + /// group and advances the smaller side; on equality the scan phase + /// pairs the streamed row with every row of the buffered group. Pairs + /// are "frozen" (materialized into output columns) once a full batch + /// of them has accumulated. async fn join( mut self, emitter: &mut TryEmitter, ) -> Result<()> { // TODO - add join time metric + + // Load the first buffered key group before the merge loop; the + // streamed side is advanced by the first loop iteration (the + // initial `current_ordering` is Equal). + self.advance_buffered_group().await?; + loop { - match &self.state { - SortMergeJoinState::Init => { - let streamed_exhausted = - self.streamed_state == StreamedState::Exhausted; - let buffered_exhausted = - self.buffered_state == BufferedState::Exhausted; - if streamed_exhausted && buffered_exhausted { - break; + // 1. Based on the previous comparison, advance the consumed + // side(s). + match self.current_ordering { + Ordering::Less | Ordering::Equal if !self.streamed_exhausted => { + // Deferred-filtering gate: run the filter pipeline once at + // least a full batch of rows has accumulated. Without this + // gate, unique keys cause per-row pipeline execution + // (concat + correct_mask + filter_by_type), which + // dominates runtime. See `deferred_rows_accumulated` for + // why the accumulation stays bounded. + if self.deferred_filtering + && self.deferred_rows_accumulated() >= self.batch_size + { + self.emit_deferred_output(emitter).await?; } - match self.current_ordering { - Ordering::Less | Ordering::Equal => { - if !streamed_exhausted { - // Batch deferred filtering: process_filtered_batches() - // only when >= batch_size rows have accumulated. - // Without this gate, unique keys cause per-row pipeline - // execution (concat + correct_mask + filter_by_type), - // which dominates runtime. - // - // Accumulated rows are bounded to ~2*batch_size: - // one batch_size worth from freeze_dequeuing_buffered() - // (when an input batch is fully consumed), plus up to - // batch_size pairs accumulating toward the next freeze. - // This does not reintroduce the unbounded buffering - // fixed by PR #20482. Exhausted state flushes remainder. - if needs_deferred_filtering(&self.filter, self.join_type) - { - let accumulated = self.num_unfrozen_pairs() - + self - .joined_record_batches - .filter_metadata - .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(), - ); - self.poll_spilled_batches_async(&needed).await?; - if let Some(batch) = - self.process_filtered_batches2()? - { - emitter.emit(batch).await; - continue; - } - } - } - self.streamed_joined = false; - self.streamed_state = StreamedState::Init; - } - } - Ordering::Greater => { - if !buffered_exhausted { - self.buffered_joined = false; - self.buffered_state = BufferedState::Init; - } - } + self.streamed_joined = false; + // Sync fast path: next row within the current batch. + // The async helper is only entered at batch boundaries. + if !self.try_advance_streamed_row() { + self.load_next_streamed_batch().await?; } - self.state = SortMergeJoinState::Polling; } - SortMergeJoinState::Polling => { - if ![StreamedState::Exhausted, StreamedState::Ready] - .contains(&self.streamed_state) - { - self.poll_streamed_row_async().await?; + Ordering::Greater if !self.buffered_exhausted => { + self.buffered_joined = false; + // Sync fast path: next group within the current buffered + // batch. The async helper is only entered at batch + // boundaries. + if !self.try_advance_buffered_group()? { + self.advance_buffered_group().await?; } + } + _ => {} + } - if ![BufferedState::Exhausted, BufferedState::Ready] - .contains(&self.buffered_state) - { - self.poll_buffered_batches_async().await?; - } - if self.streamed_state == StreamedState::Exhausted - && self.buffered_state == BufferedState::Exhausted - { - break; - } - self.current_ordering = self.compare_streamed_buffered()?; - self.state = SortMergeJoinState::JoinOutput; + if self.streamed_exhausted && self.buffered_exhausted { + break; + } + + // 2. Compare the join keys at both cursors. + self.current_ordering = self.compare_streamed_buffered()?; + + // 3. Scan the buffered group, collecting output pairs for the + // current streamed row; freeze (materialize) them whenever a + // full batch has accumulated. + debug_assert!(self.num_unfrozen_pairs() < self.batch_size); + loop { + self.join_partial()?; + + if self.num_unfrozen_pairs() < self.batch_size { + // join_partial only stops early when a full batch of + // pairs accumulated, so the scan is complete here. + debug_assert!(self.buffered_data.scanning_finished()); + break; } - SortMergeJoinState::EmitReadyThenInit => { - // If have data to emit, emit it and if no more, change to next - // Verify metadata alignment before checking if we have batches to output - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + // A full batch of pairs is ready — materialize it, + // restoring any spilled batches it needs first. + self.restore_spilled_batches_for_freeze().await?; + self.freeze_all()?; - // For filtered joins, skip output and let Init state handle it - if needs_deferred_filtering(&self.filter, self.join_type) { - self.state = SortMergeJoinState::Init; - continue; - } + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); - // For non-filtered joins, only output if we have a completed batch - // (opportunistic output when target batch size is reached) - if self + // For filtered joins, output happens through step 1's + // deferred-filtering gate instead. + if !self.deferred_filtering + && self .joined_record_batches .joined_batches .has_completed_batch() - { - let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - emitter.emit(record_batch).await; - continue; - } - self.state = SortMergeJoinState::Init; + { + self.emit_completed_joined_batches(emitter).await; } - SortMergeJoinState::JoinOutput => { - // 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()); - self.poll_spilled_batches_async(&needed).await?; - - self.freeze_all()?; - - // Verify metadata alignment before checking if we have batches to output - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); - - // For filtered joins, skip output and let Init state handle it - if needs_deferred_filtering(&self.filter, self.join_type) { - continue; - } + } + self.buffered_data.scanning_reset(); - // For non-filtered joins, only output if we have a completed batch - if self - .joined_record_batches - .joined_batches - .has_completed_batch() - { - let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - emitter.emit(record_batch).await; - continue; - } + // 4. Emit any completed output — opportunistic output when the + // target batch size is reached (non-filtered joins only; + // filtered joins defer output to step 1's gate). + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + if !self.deferred_filtering + && self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + self.emit_completed_joined_batches(emitter).await; + } + } - // Otherwise keep buffering (don't output yet) - continue; - } + self.on_children_exhausted(emitter).await + } - self.join_partial()?; + /// Number of rows currently waiting in the deferred-filtering pipeline. + /// + /// Bounded to ~2*batch_size: one batch_size worth from + /// freeze_dequeuing_buffered() (when an input batch is fully consumed), + /// plus up to batch_size pairs accumulating toward the next freeze. + /// This does not reintroduce the unbounded buffering fixed by + /// PR #20482; `on_children_exhausted` flushes the remainder. + fn deferred_rows_accumulated(&self) -> usize { + self.num_unfrozen_pairs() + + self.joined_record_batches.filter_metadata.filter_mask.len() + } - 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 => { - unreachable!("should be out of the loop"); - } - } + /// Run the deferred-filtering pipeline over everything accumulated so + /// far and emit its completed output, if any. Clears the accumulation + /// it processed. + async fn emit_deferred_output( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + // Ensure required spilled batches are restored to memory before + // processing, as this path invokes freeze_all(). + self.restore_spilled_batches_for_freeze().await?; + if let Some(batch) = self.process_filtered_batches()? { + emitter.emit(batch).await; } + Ok(()) + } - self.on_children_exhausted(emitter).await + /// Restore every spilled buffered batch that the next freeze needs. + async fn restore_spilled_batches_for_freeze(&mut self) -> Result<()> { + let needed = self.get_required_batch_indices(self.buffered_data.batches.len()); + self.restore_spilled_batches(&needed).await } + /// Emit all completed joined batches to the stream consumer. + async fn emit_completed_joined_batches( + &mut self, + emitter: &mut TryEmitter, + ) { + while let Some(record_batch) = self + .joined_record_batches + .joined_batches + .next_completed_batch() + { + emitter.emit(record_batch).await; + } + } + + /// Flush everything that remains once both inputs are exhausted. async fn on_children_exhausted( mut self, emitter: &mut TryEmitter, ) -> Result<()> { - loop { - let needed = - self.get_required_batch_indices(self.buffered_data.batches.len()); - self.poll_spilled_batches_async(&needed).await?; - - self.freeze_all()?; + // Freeze the remaining pairs, restoring any spilled batches needed. + self.restore_spilled_batches_for_freeze().await?; + self.freeze_all()?; - // Verify metadata alignment before final output - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + // Verify metadata alignment before final output + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); - // For filtered joins, must concat and filter ALL data at once - if needs_deferred_filtering(&self.filter, self.join_type) - && !self.joined_record_batches.joined_batches.is_empty() - { + if self.deferred_filtering { + // Filtered joins must concat and filter ALL remaining data at once + if !self.joined_record_batches.joined_batches.is_empty() { let record_batch = self.filter_joined_batch()?; - (&record_batch).record_output(&self.join_metrics.baseline_metrics()); emitter.emit(record_batch).await; - continue; - } - - // For non-filtered joins, finish buffered data first - if !self.joined_record_batches.joined_batches.is_empty() { - self.joined_record_batches - .joined_batches - .finish_buffered_batch()?; } - - // Output one completed batch at a time (stay in Exhausted until empty) - if self - .joined_record_batches + } else if !self.joined_record_batches.joined_batches.is_empty() { + // For non-filtered joins, finish buffered data first, then emit + // every completed batch. + self.joined_record_batches .joined_batches - .has_completed_batch() - { - let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); - (&record_batch).record_output(&self.join_metrics.baseline_metrics()); - emitter.emit(record_batch).await; - continue; - } + .finish_buffered_batch()?; + self.emit_completed_joined_batches(emitter).await; + } - // Finally check self.output BatchCoalescer (used by filtered joins) - if !self.output.is_empty() { - self.output.finish_buffered_batch()?; - let record_batch = self - .output - .next_completed_batch() - .expect("Failed to get last batch"); - (&record_batch).record_output(&self.join_metrics.baseline_metrics()); + // Drain the double-buffering coalescer used by filtered joins. + if !self.output.is_empty() { + self.output.finish_buffered_batch()?; + while let Some(record_batch) = self.output.next_completed_batch() { emitter.emit(record_batch).await; - continue; - } else { - break; - }; + } } Ok(()) @@ -936,9 +843,9 @@ impl MaterializingSortMergeJoinStream { /// Process accumulated batches for filtered joins /// - /// Freezes unfrozen pairs, applies deferred filtering, and outputs if ready. - /// Returns Poll::Ready with a batch if one is available, otherwise Poll::Pending. - fn process_filtered_batches2(&mut self) -> Result> { + /// Freezes unfrozen pairs, applies deferred filtering, and returns a + /// completed output batch if one is ready. + fn process_filtered_batches(&mut self) -> Result> { self.freeze_all()?; self.joined_record_batches @@ -956,7 +863,6 @@ impl MaterializingSortMergeJoinStream { .output .next_completed_batch() .expect("Failed to get output batch"); - (&record_batch).record_output(&self.join_metrics.baseline_metrics()); return Ok(Some(record_batch)); } } @@ -990,7 +896,7 @@ impl MaterializingSortMergeJoinStream { /// Asynchronously reads spilled batches back into memory. /// Only processes the required indices to avoid OOMs. - async fn poll_spilled_batches_async( + async fn restore_spilled_batches( &mut self, required_indices: &[usize], ) -> Result<()> { @@ -1045,53 +951,52 @@ impl MaterializingSortMergeJoinStream { Ok(()) } - /// Poll next streamed row - async fn poll_streamed_row_async(&mut self) -> Result<()> { + /// Sync fast path of advancing the streamed cursor: move to the next row + /// of the current batch. Returns false at the batch boundary, where the + /// caller must load the next batch via + /// [`Self::load_next_streamed_batch`]. + fn try_advance_streamed_row(&mut self) -> bool { + if self.streamed_batch.idx + 1 < self.streamed_batch.batch.num_rows() { + self.streamed_batch.idx += 1; + return true; + } + false + } + + /// Load the next streamed batch (freezing the finished one) and point + /// the streamed cursor at its first row. Sets `streamed_exhausted` when + /// the streamed input has no more rows. + async fn load_next_streamed_batch(&mut self) -> Result<()> { loop { - match &self.streamed_state { - StreamedState::Init => { - if self.streamed_batch.idx + 1 < self.streamed_batch.batch.num_rows() - { - self.streamed_batch.idx += 1; - self.streamed_state = StreamedState::Ready; - return Ok(()); - } else { - self.streamed_state = StreamedState::Polling; - } + // Loading a new streamed batch freezes the current one, which + // materializes buffered columns — restore any spilled buffered + // batches it needs first. + self.restore_spilled_batches_for_freeze().await?; + + match self.streamed.next().await.transpose()? { + None => { + // Release the streamed input pipeline's resources. + let streamed_schema = self.streamed.schema(); + self.streamed = + Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + self.streamed_exhausted = true; + return Ok(()); } - StreamedState::Polling => { - let needed = - self.get_required_batch_indices(self.buffered_data.batches.len()); - self.poll_spilled_batches_async(&needed).await?; - - match self.streamed.next().await.transpose()? { - 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; - } - 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; - } - } + 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); + return Ok(()); } } - StreamedState::Ready | StreamedState::Exhausted => { - return Ok(()); - } } } } @@ -1154,128 +1059,175 @@ impl MaterializingSortMergeJoinStream { Ok(()) } - /// Poll next buffered batches - async fn poll_buffered_batches_async(&mut self) -> Result<()> { + /// Advance the buffered side to the next key group: dequeue batches + /// fully consumed by the previous group, then collect all contiguous + /// rows sharing the next join key (the group may span multiple buffered + /// batches). Sets `buffered_exhausted` when no group remains. + /// Sync fast path of [`Self::advance_buffered_group`]: when the next + /// group starts in the single remaining buffered batch and provably ends + /// within it (the common case — a group only reaches a batch boundary + /// once per batch), advance entirely synchronously. Returns false — + /// leaving all state unchanged — when the async path must run instead. + fn try_advance_buffered_group(&mut self) -> Result { + if self.buffered_data.batches.len() != 1 { + return Ok(false); + } + let head_batch = self.buffered_data.head_batch(); + if head_batch.range.end == head_batch.num_rows { + // Fully consumed — needs dequeuing (and loading the next batch). + return Ok(false); + } + + if self.buffered_equality_cmp.is_none() { + self.rebuild_buffered_equality_cmp()?; + } + let cmp = self.buffered_equality_cmp.as_ref().unwrap(); + + // Scan the next group's extent before committing any state, so a + // bail-out (the group may span into the next batch) leaves + // everything untouched for the async path. + let batch = self.buffered_data.head_batch(); + let group_start = batch.range.end; + let mut group_end = group_start + 1; + while group_end < batch.num_rows && cmp.is_equal(group_start, group_end) { + group_end += 1; + } + if group_end == batch.num_rows { + return Ok(false); + } + + let batch = self.buffered_data.tail_batch_mut(); + batch.range.start = group_start; + batch.range.end = group_end; + Ok(true) + } + + async fn advance_buffered_group(&mut self) -> Result<()> { + self.dequeue_consumed_buffered_batches().await?; + + if self.buffered_data.batches.is_empty() { + // Load the batch holding the first row of the next group. + if !self.load_next_buffered_batch().await? { + self.buffered_exhausted = true; + return Ok(()); + } + } else { + // Seed the next group at the first unconsumed row of the + // remaining batch. + let tail_batch = self.buffered_data.tail_batch_mut(); + tail_batch.range.start = tail_batch.range.end; + tail_batch.range.end += 1; + } + + self.extend_buffered_group().await + } + + /// Dequeue buffered batches fully consumed by the previous group, + /// producing their pending output (e.g. Full-join null-joined rows). + async fn dequeue_consumed_buffered_batches(&mut self) -> Result<()> { + let mut head_changed = false; + while !self.buffered_data.batches.is_empty() { + let head_batch = self.buffered_data.head_batch(); + if head_batch.range.end != head_batch.num_rows { + // The next group starts within the head batch: streamed rows + // will be joined with the head batch in the next step. + break; + } + // load the spilled head batch before dequeuing + let needed = self.get_required_batch_indices(1); + self.restore_spilled_batches(&needed).await?; + + 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; + } + } + if head_changed { + self.streamed_buffered_cmp = None; + self.buffered_equality_cmp = None; + } + Ok(()) + } + + /// Load the next non-empty buffered batch and seed a new group with its + /// first row. Returns false when the buffered input is exhausted. + async fn load_next_buffered_batch(&mut self) -> Result { loop { - match &self.buffered_state { - BufferedState::Init => { - // pop previous buffered batches - let mut head_changed = false; - while !self.buffered_data.batches.is_empty() { - 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); - self.poll_spilled_batches_async(&needed).await?; - - 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 { - // If the head batch is not fully processed, break the loop. - // Streamed batch will be joined with the head batch in the next step. - break; - } - } - if head_changed { + match self.buffered.next().await.transpose()? { + None => { + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + return Ok(false); + } + Some(batch) => { + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + + if batch.num_rows() > 0 { + let buffered_batch = + BufferedBatch::new(batch, 0..1, &self.on_buffered); + self.allocate_reservation(buffered_batch)?; self.streamed_buffered_cmp = None; - self.buffered_equality_cmp = None; - } - if self.buffered_data.batches.is_empty() { - self.buffered_state = BufferedState::PollingFirst; - } else { - let tail_batch = self.buffered_data.tail_batch_mut(); - tail_batch.range.start = tail_batch.range.end; - tail_batch.range.end += 1; - self.buffered_state = BufferedState::PollingRest; + return Ok(true); } } - BufferedState::PollingFirst => { - match self.buffered.next().await.transpose()? { - None => { - // Release the buffered input pipeline's resources. - let buffered_schema = self.buffered.schema(); - self.buffered = - Box::pin(EmptyRecordBatchStream::new(buffered_schema)); - self.buffered_state = BufferedState::Exhausted; - return Ok(()); - } - Some(batch) => { - self.join_metrics.input_batches().add(1); - self.join_metrics.input_rows().add(batch.num_rows()); - - if batch.num_rows() > 0 { - let buffered_batch = - BufferedBatch::new(batch, 0..1, &self.on_buffered); + } + } + } - self.allocate_reservation(buffered_batch)?; - self.streamed_buffered_cmp = None; - self.buffered_state = BufferedState::PollingRest; - } - } - } + /// Extend the current group with every following row that shares its + /// key, loading more buffered batches as needed. + async fn extend_buffered_group(&mut self) -> Result<()> { + loop { + if self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.is_none() { + self.rebuild_buffered_equality_cmp()?; } - BufferedState::PollingRest => { - if self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows - { - if self.buffered_equality_cmp.is_none() { - self.rebuild_buffered_equality_cmp()?; - } - while self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows - { - if self.buffered_equality_cmp.as_ref().unwrap().is_equal( - self.buffered_data.head_batch().range.start, - self.buffered_data.tail_batch().range.end, - ) { - self.buffered_data.tail_batch_mut().range.end += 1; - } else { - self.buffered_state = BufferedState::Ready; - return Ok(()); - } - } + while self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.as_ref().unwrap().is_equal( + self.buffered_data.head_batch().range.start, + self.buffered_data.tail_batch().range.end, + ) { + self.buffered_data.tail_batch_mut().range.end += 1; } else { - match self.buffered.next().await.transpose()? { - None => { - // Release the buffered input pipeline's resources. - let buffered_schema = self.buffered.schema(); - self.buffered = Box::pin(EmptyRecordBatchStream::new( - buffered_schema, - )); - self.buffered_state = BufferedState::Ready; - } - Some(batch) => { - // Polling batches coming concurrently as multiple partitions - self.join_metrics.input_batches().add(1); - self.join_metrics.input_rows().add(batch.num_rows()); - if batch.num_rows() > 0 { - let buffered_batch = BufferedBatch::new( - batch, - 0..0, - &self.on_buffered, - ); - self.allocate_reservation(buffered_batch)?; - self.buffered_equality_cmp = None; - } - } - } + // Group complete within the current batch. + return Ok(()); } } - BufferedState::Ready | BufferedState::Exhausted => { - return Ok(()); + } else { + match self.buffered.next().await.transpose()? { + None => { + // Group complete; the input is done but the group is + // still valid — `buffered_exhausted` is only set once + // it has been fully consumed and dequeued. + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + return Ok(()); + } + Some(batch) => { + // Polling batches coming concurrently as multiple partitions + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + if batch.num_rows() > 0 { + let buffered_batch = + BufferedBatch::new(batch, 0..0, &self.on_buffered); + self.allocate_reservation(buffered_batch)?; + self.buffered_equality_cmp = None; + } + } } } } @@ -1283,7 +1235,7 @@ impl MaterializingSortMergeJoinStream { /// Get comparison result of streamed row and buffered batches fn compare_streamed_buffered(&mut self) -> Result { - if self.streamed_state == StreamedState::Exhausted { + if self.streamed_exhausted { return Ok(Ordering::Greater); } if !self.buffered_data.has_buffered_rows() { @@ -1507,7 +1459,7 @@ impl MaterializingSortMergeJoinStream { // but must flow through the same pipeline as matched rows to // preserve output ordering. Use null metadata as a sentinel so // get_corrected_filter_mask() passes them through unchanged. - if needs_deferred_filtering(&self.filter, self.join_type) { + if self.deferred_filtering { self.joined_record_batches .push_batch_with_null_metadata(batch, self.join_type); } else { @@ -1610,7 +1562,7 @@ impl MaterializingSortMergeJoinStream { filter_result_mask.clone() }; - if needs_deferred_filtering(&self.filter, self.join_type) { + if self.deferred_filtering { self.joined_record_batches.push_batch_with_filter_metadata( output_batch, &combined_left_indices, diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs b/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs index 6f52a2234b3dc..c50bb5fa04458 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs @@ -60,6 +60,9 @@ impl SortMergeJoinMetrics { } } + // TODO: unused while the generator refactor settles; join_time will be + // re-wired once the timing accounting is added back. + #[expect(dead_code)] pub fn join_time(&self) -> Time { self.join_time.clone() } From fc208f5df0f279b06c5f98c8f1e53243fa1f2a61 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 17:51:37 +0300 Subject: [PATCH 5/8] simplify to be text book like --- .../sort_merge_join/materializing_stream.rs | 319 +++++++++--------- 1 file changed, 152 insertions(+), 167 deletions(-) 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 b7962717d8b0e..9d7bd0a5b3399 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 @@ -307,8 +307,6 @@ pub(super) struct MaterializingSortMergeJoinStream { pub streamed: SendableRecordBatchStream, /// Current processing record batch of streamed pub streamed_batch: StreamedBatch, - /// (used in outer join) Is current streamed row joined at least once? - pub streamed_joined: bool, /// True once the streamed input has no more rows pub streamed_exhausted: bool, /// Join key columns of streamed @@ -324,8 +322,9 @@ pub(super) struct MaterializingSortMergeJoinStream { pub buffered: SendableRecordBatchStream, /// Current buffered data pub buffered_data: BufferedData, - /// (used in outer join) Is current buffered batches joined at least once? - pub buffered_joined: bool, + /// Has any streamed row matched the current buffered key group? + /// (FULL join: an unmatched group is emitted null-joined when passed.) + pub buffered_group_matched: bool, /// True once the buffered input has no more rows and no group remains pub buffered_exhausted: bool, /// Join key columns of buffered @@ -341,8 +340,6 @@ pub(super) struct MaterializingSortMergeJoinStream { /// Output buffer. Currently used by filtering as it requires double buffering /// to avoid small/empty batches. Non-filtered join outputs directly from `staging_output_record_batches.batches` pub output: BatchCoalescer, - /// The comparison result of current streamed row and buffered batches - pub current_ordering: Ordering, /// Manages the process of spilling and reading back intermediate data pub spill_manager: SpillManager, @@ -556,11 +553,9 @@ impl MaterializingSortMergeJoinStream { buffered, streamed_batch: StreamedBatch::new_empty(streamed_schema), buffered_data: BufferedData::default(), - streamed_joined: false, - buffered_joined: false, + buffered_group_matched: false, streamed_exhausted: false, buffered_exhausted: false, - current_ordering: Ordering::Equal, on_streamed, on_buffered, deferred_filtering: needs_deferred_filtering(&filter, join_type), @@ -599,109 +594,85 @@ impl MaterializingSortMergeJoinStream { ))) } - /// Main loop: a merge scan over the two sorted inputs. + /// Main loop: the textbook sort-merge join. /// - /// Compares the current streamed row with the current buffered key - /// group and advances the smaller side; on equality the scan phase - /// pairs the streamed row with every row of the buffered group. Pairs - /// are "frozen" (materialized into output columns) once a full batch - /// of them has accumulated. + /// Both inputs arrive sorted on the join keys. The streamed side is + /// consumed one row at a time; the buffered side one key *group* (all + /// contiguous rows sharing a key) at a time: + /// + /// ```text + /// load the first streamed row and the first buffered key group + /// while either input still has rows: + /// compare the join keys at both cursors + /// Less → the streamed row can never match: + /// null-join it (outer joins), advance streamed + /// Greater → the buffered group can never match again: + /// null-join it if nothing matched it (FULL join), + /// advance to the next buffered group + /// Equal → pair the streamed row with the whole group, advance + /// streamed (the group stays — the next streamed row + /// may share its key) + /// materialize ("freeze") pairs once batch_size of them accumulate + /// emit completed output batches + /// flush everything that remains + /// ``` async fn join( mut self, emitter: &mut TryEmitter, ) -> Result<()> { // TODO - add join time metric - // Load the first buffered key group before the merge loop; the - // streamed side is advanced by the first loop iteration (the - // initial `current_ordering` is Equal). + self.load_next_streamed_batch().await?; self.advance_buffered_group().await?; - loop { - // 1. Based on the previous comparison, advance the consumed - // side(s). - match self.current_ordering { - Ordering::Less | Ordering::Equal if !self.streamed_exhausted => { - // Deferred-filtering gate: run the filter pipeline once at - // least a full batch of rows has accumulated. Without this - // gate, unique keys cause per-row pipeline execution - // (concat + correct_mask + filter_by_type), which - // dominates runtime. See `deferred_rows_accumulated` for - // why the accumulation stays bounded. - if self.deferred_filtering - && self.deferred_rows_accumulated() >= self.batch_size - { - self.emit_deferred_output(emitter).await?; - } + while !(self.streamed_exhausted && self.buffered_exhausted) { + // Flush the deferred-filtering pipeline once a full batch of + // rows accumulated (filtered outer joins output through it). + if self.deferred_filtering + && self.deferred_rows_accumulated() >= self.batch_size + { + self.emit_deferred_output(emitter).await?; + } - self.streamed_joined = false; - // Sync fast path: next row within the current batch. - // The async helper is only entered at batch boundaries. + // An exhausted side compares as the larger one, so the other + // side keeps draining through its own arm. + match self.compare_streamed_buffered()? { + // The streamed row can never match: null-join it (outer + // joins emit it; inner joins drop it), then advance. + Ordering::Less => { + self.null_join_streamed_row(); + if self.num_unfrozen_pairs() >= self.batch_size { + self.freeze_full_batch(emitter).await?; + } if !self.try_advance_streamed_row() { self.load_next_streamed_batch().await?; } } - Ordering::Greater if !self.buffered_exhausted => { - self.buffered_joined = false; - // Sync fast path: next group within the current buffered - // batch. The async helper is only entered at batch - // boundaries. + // The buffered group can never match again: null-join it + // if nothing matched it (FULL join), then advance to the + // next key group. + Ordering::Greater => { + self.null_join_buffered_group(); if !self.try_advance_buffered_group()? { self.advance_buffered_group().await?; } } - _ => {} - } - - if self.streamed_exhausted && self.buffered_exhausted { - break; - } - - // 2. Compare the join keys at both cursors. - self.current_ordering = self.compare_streamed_buffered()?; - - // 3. Scan the buffered group, collecting output pairs for the - // current streamed row; freeze (materialize) them whenever a - // full batch has accumulated. - debug_assert!(self.num_unfrozen_pairs() < self.batch_size); - loop { - self.join_partial()?; - - if self.num_unfrozen_pairs() < self.batch_size { - // join_partial only stops early when a full batch of - // pairs accumulated, so the scan is complete here. - debug_assert!(self.buffered_data.scanning_finished()); - break; - } - - // A full batch of pairs is ready — materialize it, - // restoring any spilled batches it needs first. - self.restore_spilled_batches_for_freeze().await?; - self.freeze_all()?; - - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); - - // For filtered joins, output happens through step 1's - // deferred-filtering gate instead. - if !self.deferred_filtering - && self - .joined_record_batches - .joined_batches - .has_completed_batch() - { - self.emit_completed_joined_batches(emitter).await; + // Match: pair the streamed row with the whole group — + // materializing ("freezing") mid-scan whenever a full + // batch of pairs accumulates — then advance streamed. + // The group stays for the next streamed row. + Ordering::Equal => { + while !self.pair_streamed_row_with_group() { + self.freeze_full_batch(emitter).await?; + } + if !self.try_advance_streamed_row() { + self.load_next_streamed_batch().await?; + } } } - self.buffered_data.scanning_reset(); - // 4. Emit any completed output — opportunistic output when the - // target batch size is reached (non-filtered joins only; - // filtered joins defer output to step 1's gate). - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); + // Emit completed output batches (filtered joins emit through + // the deferred-filtering pipeline above instead). if !self.deferred_filtering && self .joined_record_batches @@ -715,6 +686,72 @@ impl MaterializingSortMergeJoinStream { self.on_children_exhausted(emitter).await } + /// `Equal`: pair the current streamed row with every row of the + /// buffered key group, and mark the group as matched. + /// + /// Returns false when a full batch of pairs accumulated mid-scan: the + /// caller must materialize (`freeze_full_batch`) and call again to + /// resume the scan where it paused. Returns true when the group scan + /// is complete. + fn pair_streamed_row_with_group(&mut self) -> bool { + while !self.buffered_data.scanning_finished() + && self.num_unfrozen_pairs() < self.batch_size + { + let scanning_idx = self.buffered_data.scanning_idx(); + self.streamed_batch.append_output_pair( + Some(self.buffered_data.scanning_batch_idx), + Some(scanning_idx), + self.batch_size, + ); + self.buffered_data.scanning_advance(); + } + if self.num_unfrozen_pairs() >= self.batch_size { + return false; + } + + self.buffered_group_matched = true; + self.buffered_data.scanning_reset(); + true + } + + /// `Less` (outer joins): no buffered row matches the current streamed + /// row — emit it joined to NULLs. Inner joins emit nothing. + fn null_join_streamed_row(&mut self) { + if matches!( + self.join_type, + JoinType::Left | JoinType::Right | JoinType::Full + ) { + let scanning_batch_idx = if self.buffered_data.scanning_finished() { + None + } else { + Some(self.buffered_data.scanning_batch_idx) + }; + self.streamed_batch.append_output_pair( + scanning_batch_idx, + None, + self.batch_size, + ); + } + self.buffered_data.scanning_reset(); + } + + /// `Greater` (FULL join): the buffered group can never match a streamed + /// row anymore — if nothing matched it, mark all its rows for + /// null-joined output (produced when the group's batches are dequeued). + fn null_join_buffered_group(&mut self) { + if self.join_type == JoinType::Full && !self.buffered_group_matched { + while !self.buffered_data.scanning_finished() { + let scanning_idx = self.buffered_data.scanning_idx(); + self.buffered_data + .scanning_batch_mut() + .null_joined + .push(scanning_idx); + self.buffered_data.scanning_advance(); + } + } + self.buffered_data.scanning_reset(); + } + /// Number of rows currently waiting in the deferred-filtering pipeline. /// /// Bounded to ~2*batch_size: one batch_size worth from @@ -730,6 +767,10 @@ impl MaterializingSortMergeJoinStream { /// Run the deferred-filtering pipeline over everything accumulated so /// far and emit its completed output, if any. Clears the accumulation /// it processed. + /// + /// The caller gates this on `deferred_rows_accumulated() >= batch_size`: + /// running the pipeline per row instead (concat + correct_mask + + /// filter_by_type) would dominate runtime for unique keys. async fn emit_deferred_output( &mut self, emitter: &mut TryEmitter, @@ -1099,10 +1140,12 @@ impl MaterializingSortMergeJoinStream { let batch = self.buffered_data.tail_batch_mut(); batch.range.start = group_start; batch.range.end = group_end; + self.buffered_group_matched = false; Ok(true) } async fn advance_buffered_group(&mut self) -> Result<()> { + self.buffered_group_matched = false; self.dequeue_consumed_buffered_batches().await?; if self.buffered_data.batches.is_empty() { @@ -1251,81 +1294,23 @@ impl MaterializingSortMergeJoinStream { )) } - /// Produce join and fill output buffer until reaching target batch size - /// or the join is finished - fn join_partial(&mut self) -> Result<()> { - // Whether to join streamed rows - let mut join_streamed = false; - // Whether to join buffered rows - let mut join_buffered = false; - - // determine whether we need to join streamed/buffered rows - match self.current_ordering { - Ordering::Less => { - if matches!( - self.join_type, - JoinType::Left | JoinType::Right | JoinType::Full - ) { - join_streamed = !self.streamed_joined; - } - } - Ordering::Equal => { - join_streamed = true; - join_buffered = true; - } - Ordering::Greater => { - if self.join_type == JoinType::Full { - join_buffered = !self.buffered_joined; - }; - } - } - if !join_streamed && !join_buffered { - // no joined data - self.buffered_data.scanning_finish(); - return Ok(()); - } - - if join_buffered { - // joining streamed/nulls and buffered - while !self.buffered_data.scanning_finished() - && self.num_unfrozen_pairs() < self.batch_size - { - let scanning_idx = self.buffered_data.scanning_idx(); - if join_streamed { - // Join streamed row and buffered row - self.streamed_batch.append_output_pair( - Some(self.buffered_data.scanning_batch_idx), - Some(scanning_idx), - self.batch_size, - ); - } else { - // Join nulls and buffered row for FULL join - self.buffered_data - .scanning_batch_mut() - .null_joined - .push(scanning_idx); - } - self.buffered_data.scanning_advance(); + /// Materialize ("freeze") the accumulated pairs — restoring any spilled + /// batches they reference first — and emit completed output batches + /// (filtered joins emit through the deferred-filtering gate instead). + async fn freeze_full_batch( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + self.restore_spilled_batches_for_freeze().await?; + self.freeze_all()?; - if self.buffered_data.scanning_finished() { - self.streamed_joined = join_streamed; - self.buffered_joined = true; - } - } - } else { - // joining streamed and nulls - let scanning_batch_idx = if self.buffered_data.scanning_finished() { - None - } else { - Some(self.buffered_data.scanning_batch_idx) - }; - self.streamed_batch.append_output_pair( - scanning_batch_idx, - None, - self.batch_size, - ); - self.buffered_data.scanning_finish(); - self.streamed_joined = true; + if !self.deferred_filtering + && self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + self.emit_completed_joined_batches(emitter).await; } Ok(()) } @@ -1914,9 +1899,9 @@ fn fetch_right_columns_from_batch_by_idxs( pub(super) struct BufferedData { /// Buffered batches with the same key pub batches: VecDeque, - /// current scanning batch index used in join_partial() + /// current scanning batch index used by the group-scan phase pub scanning_batch_idx: usize, - /// current scanning offset used in join_partial() + /// current scanning offset used by the group-scan phase pub scanning_offset: usize, } From 5bab8f5cacb00de1e8c507a809093de0ada1939a Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:09:25 +0300 Subject: [PATCH 6/8] comments --- .../sort_merge_join/materializing_stream.rs | 58 ++++++++++--------- 1 file changed, 31 insertions(+), 27 deletions(-) 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 9d7bd0a5b3399..f1aafd164cec2 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 @@ -601,20 +601,20 @@ impl MaterializingSortMergeJoinStream { /// contiguous rows sharing a key) at a time: /// /// ```text - /// load the first streamed row and the first buffered key group - /// while either input still has rows: - /// compare the join keys at both cursors - /// Less → the streamed row can never match: - /// null-join it (outer joins), advance streamed - /// Greater → the buffered group can never match again: - /// null-join it if nothing matched it (FULL join), - /// advance to the next buffered group - /// Equal → pair the streamed row with the whole group, advance - /// streamed (the group stays — the next streamed row - /// may share its key) - /// materialize ("freeze") pairs once batch_size of them accumulate - /// emit completed output batches - /// flush everything that remains + /// 1. load the first streamed row and the first buffered key group + /// 2. while either input still has rows: + /// 3. compare the join keys at both cursors + /// 3a. Less → the streamed row can never match: + /// null-join it (outer joins), advance streamed + /// 3b. Greater → the buffered group can never match again: + /// null-join it if nothing matched it (FULL join), + /// advance to the next buffered group + /// 3c. Equal → pair the streamed row with the whole group, + /// advance streamed (the group stays — the next + /// streamed row may share its key) + /// (materialize ("freeze") pairs once batch_size accumulate) + /// 4. emit completed output batches + /// 5. flush everything that remains /// ``` async fn join( mut self, @@ -622,9 +622,11 @@ impl MaterializingSortMergeJoinStream { ) -> Result<()> { // TODO - add join time metric + // 1. Load the first streamed row and the first buffered key group. self.load_next_streamed_batch().await?; self.advance_buffered_group().await?; + // 2. Merge-scan while either input still has rows. while !(self.streamed_exhausted && self.buffered_exhausted) { // Flush the deferred-filtering pipeline once a full batch of // rows accumulated (filtered outer joins output through it). @@ -634,11 +636,12 @@ impl MaterializingSortMergeJoinStream { self.emit_deferred_output(emitter).await?; } - // An exhausted side compares as the larger one, so the other - // side keeps draining through its own arm. + // 3. Compare the join keys at both cursors. An exhausted side + // compares as the larger one, so the other side keeps + // draining through its own arm. match self.compare_streamed_buffered()? { - // The streamed row can never match: null-join it (outer - // joins emit it; inner joins drop it), then advance. + // 3a. The streamed row can never match: null-join it (outer + // joins emit it; inner joins drop it), then advance. Ordering::Less => { self.null_join_streamed_row(); if self.num_unfrozen_pairs() >= self.batch_size { @@ -648,19 +651,19 @@ impl MaterializingSortMergeJoinStream { self.load_next_streamed_batch().await?; } } - // The buffered group can never match again: null-join it - // if nothing matched it (FULL join), then advance to the - // next key group. + // 3b. The buffered group can never match again: null-join + // it if nothing matched it (FULL join), then advance to + // the next key group. Ordering::Greater => { self.null_join_buffered_group(); if !self.try_advance_buffered_group()? { self.advance_buffered_group().await?; } } - // Match: pair the streamed row with the whole group — - // materializing ("freezing") mid-scan whenever a full - // batch of pairs accumulates — then advance streamed. - // The group stays for the next streamed row. + // 3c. Match: pair the streamed row with the whole group — + // materializing ("freezing") mid-scan whenever a full + // batch of pairs accumulates — then advance streamed. + // The group stays for the next streamed row. Ordering::Equal => { while !self.pair_streamed_row_with_group() { self.freeze_full_batch(emitter).await?; @@ -671,8 +674,8 @@ impl MaterializingSortMergeJoinStream { } } - // Emit completed output batches (filtered joins emit through - // the deferred-filtering pipeline above instead). + // 4. Emit completed output batches (filtered joins emit + // through the deferred-filtering pipeline above instead). if !self.deferred_filtering && self .joined_record_batches @@ -683,6 +686,7 @@ impl MaterializingSortMergeJoinStream { } } + // 5. Flush everything that remains. self.on_children_exhausted(emitter).await } From 2f132c43b53cf0d72b70296914536869f6d56930 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:24:57 +0300 Subject: [PATCH 7/8] update comments and add timing --- .../sort_merge_join/materializing_stream.rs | 167 +++++++++++------- .../src/joins/sort_merge_join/metrics.rs | 3 - .../src/joins/sort_merge_join/tests.rs | 141 +++++++++++++++ 3 files changed, 244 insertions(+), 67 deletions(-) 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 f1aafd164cec2..cf11c15b67aef 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 @@ -17,9 +17,10 @@ //! Sort-Merge Join execution //! -//! This module implements the runtime state machine for the Sort-Merge Join -//! operator. It drives two sorted input streams (the *streamed* side and the -//! *buffered* side), compares join keys, and produces joined `RecordBatch`es. +//! This module implements the Sort-Merge Join operator as an async +//! generator running a merge scan: it drives two sorted input streams (the +//! *streamed* side and the *buffered* side), compares join keys, and +//! produces joined `RecordBatch`es. use std::cmp::Ordering; use std::collections::{HashMap, VecDeque}; @@ -27,8 +28,6 @@ use std::fmt::Debug; use std::mem::size_of; use std::ops::Range; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering::Relaxed; use crate::joins::sort_merge_join::filter::{ FilterMetadata, filter_record_batch_by_join_type, get_corrected_filter_mask, @@ -36,6 +35,7 @@ use crate::joins::sort_merge_join::filter::{ }; use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; use crate::joins::utils::{JoinFilter, JoinKeyComparator}; +use crate::metrics::Time; use crate::spill::spill_manager::SpillManager; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{PhysicalExpr, SendableRecordBatchStream}; @@ -47,6 +47,7 @@ use arrow::compute::{ }; use arrow::datatypes::SchemaRef; use datafusion_common::cast::as_uint64_array; +use datafusion_common::instant::Instant; use datafusion_common::{ DataFusionError, JoinType, NullEquality, Result, exec_err, internal_err, }; @@ -338,16 +339,23 @@ pub(super) struct MaterializingSortMergeJoinStream { /// Staging output array builders pub joined_record_batches: JoinedRecordBatches, /// Output buffer. Currently used by filtering as it requires double buffering - /// to avoid small/empty batches. Non-filtered join outputs directly from `staging_output_record_batches.batches` + /// to avoid small/empty batches. Non-filtered joins output directly from + /// `joined_record_batches.joined_batches` pub output: BatchCoalescer, /// 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, /// Tracks the number of batches currently spilled pub spilled_batch_count: usize, + /// Time spent doing the join's own work (including spill write and + /// read-back). The clock is stopped while awaiting the child inputs or + /// the consumer taking an emitted batch — see [`Self::stop_join_time`]. + pub join_time: Time, + /// Start of the currently running `join_time` span; `None` while the + /// clock is stopped. + pub join_time_start: Option, + // ======================================================================== // CACHED COMPARATORS: // Pre-built comparators to avoid per-row type dispatch in hot loops. @@ -367,8 +375,9 @@ pub(super) struct MaterializingSortMergeJoinStream { pub reservation: MemoryReservation, /// Runtime env pub runtime_env: Arc, - /// A unique number for each batch - pub streamed_batch_counter: AtomicUsize, + /// A unique id per streamed batch, tagging deferred-filter metadata so + /// `get_corrected_filter_mask` can group output rows by input batch. + pub streamed_batch_counter: usize, } /// Staging area for joined data before output @@ -543,6 +552,7 @@ impl MaterializingSortMergeJoinStream { "MaterializingSortMergeJoinStream does not handle {join_type:?}; \ semi/anti/mark joins use BitwiseSortMergeJoinStream" ); + let join_time = join_metrics.join_time(); let mut this = Self { sort_options, null_equality, @@ -573,20 +583,25 @@ impl MaterializingSortMergeJoinStream { reservation, runtime_env, spill_manager, - spill_stream: None, spilled_batch_count: 0, + join_time, + join_time_start: None, streamed_buffered_cmp: None, buffered_equality_cmp: None, - streamed_batch_counter: AtomicUsize::new(0), + streamed_batch_counter: 0, }; let schema = Arc::clone(&this.schema); let baseline_metrics = this.join_metrics.baseline_metrics(); - let stream = - async_try_stream(|mut emitter| async move { this.join(&mut emitter).await }); + let stream = async_try_stream(|mut emitter| async move { + this.start_join_time(); + let result = this.join(&mut emitter).await; + this.stop_join_time(); + result + }); // ObservedStream records the baseline metrics (output rows/batches, - // end time) exactly as the former hand-written poll_next did. + // end time). Ok(Box::pin(ObservedStream::new( Box::pin(RecordBatchStreamAdapter::new(schema, stream)), baseline_metrics, @@ -617,11 +632,9 @@ impl MaterializingSortMergeJoinStream { /// 5. flush everything that remains /// ``` async fn join( - mut self, + &mut self, emitter: &mut TryEmitter, ) -> Result<()> { - // TODO - add join time metric - // 1. Load the first streamed row and the first buffered key group. self.load_next_streamed_batch().await?; self.advance_buffered_group().await?; @@ -645,7 +658,7 @@ impl MaterializingSortMergeJoinStream { Ordering::Less => { self.null_join_streamed_row(); if self.num_unfrozen_pairs() >= self.batch_size { - self.freeze_full_batch(emitter).await?; + self.freeze_and_emit(emitter).await?; } if !self.try_advance_streamed_row() { self.load_next_streamed_batch().await?; @@ -666,7 +679,7 @@ impl MaterializingSortMergeJoinStream { // The group stays for the next streamed row. Ordering::Equal => { while !self.pair_streamed_row_with_group() { - self.freeze_full_batch(emitter).await?; + self.freeze_and_emit(emitter).await?; } if !self.try_advance_streamed_row() { self.load_next_streamed_batch().await?; @@ -693,10 +706,11 @@ impl MaterializingSortMergeJoinStream { /// `Equal`: pair the current streamed row with every row of the /// buffered key group, and mark the group as matched. /// - /// Returns false when a full batch of pairs accumulated mid-scan: the - /// caller must materialize (`freeze_full_batch`) and call again to - /// resume the scan where it paused. Returns true when the group scan - /// is complete. + /// Returns false when a full batch of pairs has accumulated (the scan + /// may or may not be complete): the caller must materialize + /// (`freeze_and_emit`) and call again, which resumes the scan where it + /// paused. Returns true when the group scan is complete and there is + /// room for more pairs. fn pair_streamed_row_with_group(&mut self) -> bool { while !self.buffered_data.scanning_finished() && self.num_unfrozen_pairs() < self.batch_size @@ -756,13 +770,34 @@ impl MaterializingSortMergeJoinStream { self.buffered_data.scanning_reset(); } + /// Start (resume) the `join_time` clock. + fn start_join_time(&mut self) { + debug_assert!(self.join_time_start.is_none(), "join_time already running"); + self.join_time_start = Some(Instant::now()); + } + + /// Stop (pause) the `join_time` clock, accumulating the elapsed span. + /// + /// Called around awaits whose duration is not the join's own work: the + /// child input streams' `next()` and `emitter.emit()` (where the + /// consumer processes the batch). The join's own spill write and + /// read-back are NOT excluded — that time is join work. + fn stop_join_time(&mut self) { + if let Some(start) = self.join_time_start.take() { + self.join_time.add_elapsed(start); + } + } + /// Number of rows currently waiting in the deferred-filtering pipeline. /// - /// Bounded to ~2*batch_size: one batch_size worth from + /// Typically bounded to ~2*batch_size: one batch_size worth from /// freeze_dequeuing_buffered() (when an input batch is fully consumed), - /// plus up to batch_size pairs accumulating toward the next freeze. - /// This does not reintroduce the unbounded buffering fixed by - /// PR #20482; `on_children_exhausted` flushes the remainder. + /// plus up to batch_size pairs accumulating toward the next freeze. A + /// single streamed row matching a very large key group can exceed that + /// (its pairs freeze into the pipeline before the gate runs again — same + /// as the pre-generator design). This does not reintroduce the unbounded + /// buffering fixed by PR #20482; `on_children_exhausted` flushes the + /// remainder. fn deferred_rows_accumulated(&self) -> usize { self.num_unfrozen_pairs() + self.joined_record_batches.filter_metadata.filter_mask.len() @@ -783,7 +818,11 @@ impl MaterializingSortMergeJoinStream { // processing, as this path invokes freeze_all(). self.restore_spilled_batches_for_freeze().await?; if let Some(batch) = self.process_filtered_batches()? { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); emitter.emit(batch).await; + self.start_join_time(); } Ok(()) } @@ -804,13 +843,17 @@ impl MaterializingSortMergeJoinStream { .joined_batches .next_completed_batch() { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); emitter.emit(record_batch).await; + self.start_join_time(); } } /// Flush everything that remains once both inputs are exhausted. async fn on_children_exhausted( - mut self, + &mut self, emitter: &mut TryEmitter, ) -> Result<()> { // Freeze the remaining pairs, restoring any spilled batches needed. @@ -826,7 +869,9 @@ impl MaterializingSortMergeJoinStream { // Filtered joins must concat and filter ALL remaining data at once if !self.joined_record_batches.joined_batches.is_empty() { let record_batch = self.filter_joined_batch()?; + self.stop_join_time(); emitter.emit(record_batch).await; + self.start_join_time(); } } else if !self.joined_record_batches.joined_batches.is_empty() { // For non-filtered joins, finish buffered data first, then emit @@ -841,7 +886,9 @@ impl MaterializingSortMergeJoinStream { if !self.output.is_empty() { self.output.finish_buffered_batch()?; while let Some(record_batch) = self.output.next_completed_batch() { + self.stop_join_time(); emitter.emit(record_batch).await; + self.start_join_time(); } } @@ -954,21 +1001,11 @@ impl MaterializingSortMergeJoinStream { 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(Arc::clone(spill_file), None)?; - self.spill_stream = Some(stream); - } + let mut spill_stream = self + .spill_manager + .read_spill_as_stream(Arc::clone(spill_file), None)?; - match self - .spill_stream - .as_mut() - .unwrap() - .next() - .await - .transpose()? - { + match spill_stream.next().await.transpose()? { Some(batch) => { // Transition the batch back to InMemory bb.batch = BufferedBatchState::InMemory(batch); @@ -982,11 +1019,8 @@ impl MaterializingSortMergeJoinStream { self.join_metrics .peak_mem_used() .set_max(self.reservation.size()); - - self.spill_stream = None; } None => { - self.spill_stream = None; return internal_err!("Spill file was empty"); } } @@ -1018,7 +1052,11 @@ impl MaterializingSortMergeJoinStream { // batches it needs first. self.restore_spilled_batches_for_freeze().await?; - match self.streamed.next().await.transpose()? { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.streamed.next().await.transpose(); + self.start_join_time(); + match item? { None => { // Release the streamed input pipeline's resources. let streamed_schema = self.streamed.schema(); @@ -1035,10 +1073,8 @@ impl MaterializingSortMergeJoinStream { 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); + // Every incoming streamed batch gets a unique id. + self.streamed_batch_counter += 1; return Ok(()); } } @@ -1104,10 +1140,6 @@ impl MaterializingSortMergeJoinStream { Ok(()) } - /// Advance the buffered side to the next key group: dequeue batches - /// fully consumed by the previous group, then collect all contiguous - /// rows sharing the next join key (the group may span multiple buffered - /// batches). Sets `buffered_exhausted` when no group remains. /// Sync fast path of [`Self::advance_buffered_group`]: when the next /// group starts in the single remaining buffered batch and provably ends /// within it (the common case — a group only reaches a batch boundary @@ -1148,6 +1180,10 @@ impl MaterializingSortMergeJoinStream { Ok(true) } + /// Advance the buffered side to the next key group: dequeue batches + /// fully consumed by the previous group, then collect all contiguous + /// rows sharing the next join key (the group may span multiple buffered + /// batches). Sets `buffered_exhausted` when no group remains. async fn advance_buffered_group(&mut self) -> Result<()> { self.buffered_group_matched = false; self.dequeue_consumed_buffered_batches().await?; @@ -1205,7 +1241,11 @@ impl MaterializingSortMergeJoinStream { /// first row. Returns false when the buffered input is exhausted. async fn load_next_buffered_batch(&mut self) -> Result { loop { - match self.buffered.next().await.transpose()? { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.buffered.next().await.transpose(); + self.start_join_time(); + match item? { None => { // Release the buffered input pipeline's resources. let buffered_schema = self.buffered.schema(); @@ -1253,7 +1293,11 @@ impl MaterializingSortMergeJoinStream { } } } else { - match self.buffered.next().await.transpose()? { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.buffered.next().await.transpose(); + self.start_join_time(); + match item? { None => { // Group complete; the input is done but the group is // still valid — `buffered_exhausted` is only set once @@ -1301,7 +1345,7 @@ impl MaterializingSortMergeJoinStream { /// Materialize ("freeze") the accumulated pairs — restoring any spilled /// batches they reference first — and emit completed output batches /// (filtered joins emit through the deferred-filtering gate instead). - async fn freeze_full_batch( + async fn freeze_and_emit( &mut self, emitter: &mut TryEmitter, ) -> Result<()> { @@ -1556,7 +1600,7 @@ impl MaterializingSortMergeJoinStream { output_batch, &combined_left_indices, &mask, - self.streamed_batch_counter.load(Relaxed), + self.streamed_batch_counter, self.join_type, ); } else { @@ -1958,11 +2002,6 @@ impl BufferedData { pub fn scanning_finished(&self) -> bool { self.scanning_batch_idx == self.batches.len() } - - pub fn scanning_finish(&mut self) { - self.scanning_batch_idx = self.batches.len(); - self.scanning_offset = 0; - } } /// Get join array refs of given batch and join columns diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs b/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs index c50bb5fa04458..6f52a2234b3dc 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs @@ -60,9 +60,6 @@ impl SortMergeJoinMetrics { } } - // TODO: unused while the generator refactor settles; join_time will be - // re-wired once the timing accounting is added back. - #[expect(dead_code)] pub fn join_time(&self) -> Time { self.join_time.clone() } 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 313818cf6c6b5..3dbb50eba07d9 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -4315,6 +4315,147 @@ async fn join_time_excludes_consumer_wait() -> Result<()> { .await } +/// Three 2-row batches with unique matching keys, right-side column names. +fn join_time_batches_right() -> Vec { + vec![ + build_table_i32( + ("a2", &vec![0, 1]), + ("b2", &vec![1, 2]), + ("c2", &vec![7, 8]), + ), + build_table_i32( + ("a2", &vec![2, 3]), + ("b2", &vec![3, 4]), + ("c2", &vec![7, 8]), + ), + build_table_i32( + ("a2", &vec![4, 5]), + ("b2", &vec![5, 6]), + ("c2", &vec![7, 8]), + ), + ] +} + +/// Build a no-filter Inner materializing join over the given input streams. +/// The small batch size makes the output surface as multiple batches, so a +/// slow consumer test sees multiple emits. +fn materializing_join_time_test_join( + streamed: SendableRecordBatchStream, + buffered: SendableRecordBatchStream, +) -> (SendableRecordBatchStream, ExecutionPlanMetricsSet) { + use crate::joins::sort_merge_join::materializing_stream::MaterializingSortMergeJoinStream; + use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; + + let metrics = ExecutionPlanMetricsSet::new(); + let out_schema = Arc::new(Schema::new( + streamed + .schema() + .fields() + .iter() + .chain(buffered.schema().fields().iter()) + .map(|f| f.as_ref().clone()) + .collect::>(), + )); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(buffered.schema(), &metrics); + let stream = MaterializingSortMergeJoinStream::try_new( + out_schema, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + streamed, + buffered, + vec![Arc::new(Column::new("b1", 1)) as _], + vec![Arc::new(Column::new("b2", 1)) as _], + None, + Inner, + 2, + SortMergeJoinMetrics::new(0, &metrics), + reservation, + spill_manager, + runtime_env, + ) + .unwrap(); + (stream, metrics) +} + +/// join_time must not include time spent waiting for the streamed input. +#[tokio::test] +async fn materializing_join_time_excludes_streamed_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), delay); + let buffered = delayed_stream(join_time_batches_right(), Duration::ZERO); + let (stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all rows should match"); + assert!( + wall >= delay * 3, + "streamed delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time spent waiting for the buffered input. +#[tokio::test] +async fn materializing_join_time_excludes_buffered_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), Duration::ZERO); + let buffered = delayed_stream(join_time_batches_right(), delay); + let (stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all rows should match"); + assert!( + wall >= delay * 3, + "buffered delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time the consumer spends holding an emitted +/// batch (the generator is suspended inside `emitter.emit` meanwhile). +#[tokio::test] +async fn materializing_join_time_excludes_consumer_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), Duration::ZERO); + let buffered = delayed_stream(join_time_batches_right(), Duration::ZERO); + let (mut stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let mut output_batches = 0u32; + while let Some(batch) = stream.next().await { + batch?; + output_batches += 1; + // Simulate a slow consumer between emitted batches. + tokio::time::sleep(delay).await; + } + let wall = start.elapsed(); + + assert!( + output_batches >= 3, + "expected multiple emitted batches, got {output_batches}" + ); + assert!( + wall >= delay * output_batches, + "consumer delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + /// An inner key group spanning multiple inner batches must survive the inner /// input returning Pending mid-way: inner rows delivered before the Pending /// still take part in the filter evaluation. From cb881c764a1ca7d87fb3fb872ee26b944ad387e4 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:27:02 +0300 Subject: [PATCH 8/8] update comments --- .../sort_merge_join/materializing_stream.rs | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) 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 cf11c15b67aef..3baa0c4a3e792 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 @@ -613,24 +613,7 @@ impl MaterializingSortMergeJoinStream { /// /// Both inputs arrive sorted on the join keys. The streamed side is /// consumed one row at a time; the buffered side one key *group* (all - /// contiguous rows sharing a key) at a time: - /// - /// ```text - /// 1. load the first streamed row and the first buffered key group - /// 2. while either input still has rows: - /// 3. compare the join keys at both cursors - /// 3a. Less → the streamed row can never match: - /// null-join it (outer joins), advance streamed - /// 3b. Greater → the buffered group can never match again: - /// null-join it if nothing matched it (FULL join), - /// advance to the next buffered group - /// 3c. Equal → pair the streamed row with the whole group, - /// advance streamed (the group stays — the next - /// streamed row may share its key) - /// (materialize ("freeze") pairs once batch_size accumulate) - /// 4. emit completed output batches - /// 5. flush everything that remains - /// ``` + /// contiguous rows sharing a key) at a time async fn join( &mut self, emitter: &mut TryEmitter,