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 51cf38b9ab1f..e401222bc25f 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 @@ -129,6 +129,9 @@ pub(super) struct StreamedBatch { pub num_output_rows: usize, /// Index of currently scanned batch from buffered data pub buffered_batch_idx: Option, + /// Reusable scratch buffer for materializing contiguous buffered index + /// ranges in [`Self::append_output_range`] + buffered_index_scratch: Vec, } impl StreamedBatch { @@ -141,6 +144,7 @@ impl StreamedBatch { output_indices: vec![], num_output_rows: 0, buffered_batch_idx: None, + buffered_index_scratch: vec![], } } @@ -152,6 +156,7 @@ impl StreamedBatch { output_indices: vec![], num_output_rows: 0, buffered_batch_idx: None, + buffered_index_scratch: vec![], } } @@ -199,6 +204,61 @@ impl StreamedBatch { } self.num_output_rows += 1; } + + /// Appends `len` output pairs joining the current streamed index against the + /// contiguous buffered index range `buffered_start..buffered_start + len`. + /// + /// Equivalent to calling [`Self::append_output_pair`] once per index in the + /// range, but appends into the builders in bulk. + fn append_output_range( + &mut self, + buffered_batch_idx: Option, + buffered_start: usize, + len: usize, + batch_size: usize, + ) { + if len == 0 { + return; + } + + if self.output_indices.is_empty() || self.buffered_batch_idx != buffered_batch_idx + { + debug_assert!( + batch_size > self.num_output_rows, + "batch_size ({batch_size}) must be > num_output_rows ({})", + self.num_output_rows + ); + let capacity = batch_size - self.num_output_rows; + self.output_indices.push(StreamedJoinedChunk { + buffered_batch_idx, + streamed_indices: UInt64Builder::with_capacity(capacity), + buffered_indices: UInt64Builder::with_capacity(capacity), + }); + self.buffered_batch_idx = buffered_batch_idx; + }; + + let current_chunk = self.output_indices.last_mut().unwrap(); + if len == 1 { + // Fast path avoids touching the scratch buffer for the common + // one-row-per-key case. + current_chunk.streamed_indices.append_value(self.idx as u64); + current_chunk + .buffered_indices + .append_value(buffered_start as u64); + } else { + self.buffered_index_scratch.clear(); + self.buffered_index_scratch + .extend((buffered_start..buffered_start + len).map(|idx| idx as u64)); + current_chunk + .streamed_indices + .append_value_n(self.idx as u64, len); + current_chunk + .buffered_indices + .append_slice(&self.buffered_index_scratch); + } + + self.num_output_rows += len; + } } /// Per-row filter outcome tracking for full outer joins. @@ -1344,22 +1404,32 @@ impl MaterializingSortMergeJoinStream { while !self.buffered_data.scanning_finished() && self.num_unfrozen_pairs() < self.batch_size { + // The buffered rows remaining in the batch being scanned form a + // contiguous index range, so append them in one go rather than + // one row at a time, capped by the remaining output batch space. + let remaining_capacity = self.batch_size - self.num_unfrozen_pairs(); + let run_len = self + .buffered_data + .scanning_batch_remaining() + .min(remaining_capacity); let scanning_idx = self.buffered_data.scanning_idx(); + if join_streamed { - // Join streamed row and buffered row - self.streamed_batch.append_output_pair( + // Join streamed row and buffered rows + self.streamed_batch.append_output_range( Some(self.buffered_data.scanning_batch_idx), - Some(scanning_idx), + scanning_idx, + run_len, self.batch_size, ); } else { - // Join nulls and buffered row for FULL join + // Join nulls and buffered rows for FULL join self.buffered_data .scanning_batch_mut() .null_joined - .push(scanning_idx); + .extend(scanning_idx..scanning_idx + run_len); } - self.buffered_data.scanning_advance(); + self.buffered_data.scanning_advance_by(run_len); if self.buffered_data.scanning_finished() { self.streamed_joined = join_streamed; @@ -1996,14 +2066,24 @@ impl BufferedData { self.scanning_offset = 0; } - pub fn scanning_advance(&mut self) { - self.scanning_offset += 1; + /// Advances the scanning position by `n` rows. + /// + /// `n` must not exceed [`Self::scanning_batch_remaining`], so that the + /// position never runs past the end of the batch being scanned. + pub fn scanning_advance_by(&mut self, n: usize) { + debug_assert!(n <= self.scanning_batch_remaining()); + self.scanning_offset += n; while !self.scanning_finished() && self.scanning_batch_finished() { self.scanning_batch_idx += 1; self.scanning_offset = 0; } } + /// Number of unscanned rows remaining in the batch currently being scanned + pub fn scanning_batch_remaining(&self) -> usize { + self.scanning_batch().range.len() - self.scanning_offset + } + pub fn scanning_batch(&self) -> &BufferedBatch { &self.batches[self.scanning_batch_idx] }