From 0124044c1631132403a4e15e8856b43bfec1d36e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 13 Oct 2026 11:38:30 -0600 Subject: [PATCH 1/9] perf: append sort-merge join output pairs in bulk When a streamed row matches a buffered key group, pair it with the rest of the scanning buffered batch at once, up to the room left in the output batch, instead of one buffered row at a time. The pairs and their order are unchanged. Adds a test whose key group spans buffered and output batches, and a benchmark case in which every streamed row matches 100 buffered rows. --- .../physical-plan/benches/sort_merge_join.rs | 16 +++ .../sort_merge_join/materializing_stream.rs | 102 ++++++++++++++---- .../src/joins/sort_merge_join/tests.rs | 67 ++++++++++++ 3 files changed, 167 insertions(+), 18 deletions(-) diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index d5f87a04f0a28..03dec34eaded7 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -168,6 +168,22 @@ fn bench_smj(c: &mut Criterion) { }); } + // 1:100 Inner Join — 20K left, 20K right, 200 distinct keys + // Every streamed row matches 100 buffered rows, so the join emits 2M rows. + { + let n = 20_000; + let key_mod = 200; + let left_batches = build_sorted_batches(n, key_mod, &s); + let right_batches = build_sorted_batches(n, key_mod, &s); + group.bench_function(BenchmarkId::new("inner_1to100", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join(left, right, datafusion_common::JoinType::Inner, &rt) + }) + }); + } + // Left Join — 100K each, ~5% unmatched on left { let n = 100_000; 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 8027d83dbe124..0b797924a71ac 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 @@ -125,8 +125,71 @@ impl StreamedBatch { buffered_idx: Option, batch_size: usize, ) { - // If no current chunk exists or current chunk is not for current buffered batch, - // create a new chunk + let streamed_idx = self.idx as u64; + let current_chunk = self.chunk_for(buffered_batch_idx, batch_size); + + // Append index of streamed batch and index of buffered batch into current chunk + current_chunk.streamed_indices.append_value(streamed_idx); + if let Some(idx) = buffered_idx { + current_chunk.buffered_indices.append_value(idx as u64); + } else { + current_chunk.buffered_indices.append_null(); + } + self.num_output_rows += 1; + } + + /// Appends one pair for each row in `buffered_range` of the buffered batch with + /// `buffered_batch_idx` index, all with the current streamed index. + /// + /// Equivalent to calling [`Self::append_output_pair`] once per buffered row, but + /// appends the indices in bulk, which matters when a streamed row matches many + /// buffered rows. + #[inline] + fn append_output_pairs( + &mut self, + buffered_batch_idx: usize, + buffered_range: Range, + batch_size: usize, + ) { + let len = buffered_range.len(); + if len == 0 { + return; + } + debug_assert!( + self.num_output_rows + len <= batch_size, + "appending {len} pairs to {} would exceed batch_size ({batch_size})", + self.num_output_rows + ); + let streamed_idx = self.idx as u64; + let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); + if len == 1 { + // A single match, as in a one-to-one join, is cheaper to append directly. + current_chunk.streamed_indices.append_value(streamed_idx); + current_chunk + .buffered_indices + .append_value(buffered_range.start as u64); + } else { + current_chunk + .streamed_indices + .append_value_n(streamed_idx, len); + // SAFETY: a `Range` reports its exact length in `size_hint`. + unsafe { + current_chunk.buffered_indices.append_trusted_len_iter( + buffered_range.start as u64..buffered_range.end as u64, + ); + } + } + self.num_output_rows += len; + } + + /// Returns the chunk that pairs for `buffered_batch_idx` go to, starting a new one + /// if there is no chunk yet or the current chunk is for another buffered batch. + #[inline] + fn chunk_for( + &mut self, + buffered_batch_idx: Option, + batch_size: usize, + ) -> &mut StreamedJoinedChunk { if self.output_indices.is_empty() || self.buffered_batch_idx != buffered_batch_idx { // Compute capacity only when creating a new chunk (infrequent operation). @@ -145,16 +208,7 @@ impl StreamedBatch { }); self.buffered_batch_idx = buffered_batch_idx; } - let current_chunk = self.output_indices.last_mut().unwrap(); - - // Append index of streamed batch and index of buffered batch into current chunk - current_chunk.streamed_indices.append_value(self.idx as u64); - if let Some(idx) = buffered_idx { - current_chunk.buffered_indices.append_value(idx as u64); - } else { - current_chunk.buffered_indices.append_null(); - } - self.num_output_rows += 1; + self.output_indices.last_mut().unwrap() } } @@ -725,13 +779,18 @@ impl MaterializingSortMergeJoinStream { 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), + // Pair the streamed row with the rest of the scanning buffered batch at + // once, up to the room left in the output batch. + let range = &self.buffered_data.scanning_batch().range; + let start = range.start + self.buffered_data.scanning_offset; + let len = + (range.end - start).min(self.batch_size - self.num_unfrozen_pairs()); + self.streamed_batch.append_output_pairs( + self.buffered_data.scanning_batch_idx, + start..start + len, self.batch_size, ); - self.buffered_data.scanning_advance(); + self.buffered_data.scanning_advance_by(len); } if self.num_unfrozen_pairs() >= self.batch_size { return false; @@ -2114,7 +2173,14 @@ impl BufferedData { } pub fn scanning_advance(&mut self) { - self.scanning_offset += 1; + self.scanning_advance_by(1); + } + + /// Advances the scan by `n` rows, which must not pass the end of the scanning + /// batch, then skips to the next batch that has rows left. + pub fn scanning_advance_by(&mut self, n: usize) { + self.scanning_offset += n; + debug_assert!(self.scanning_offset <= self.scanning_batch().range.len()); while !self.scanning_finished() && self.scanning_batch_finished() { self.scanning_batch_idx += 1; self.scanning_offset = 0; 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 f921812ec0619..f23234a240d9d 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -1056,6 +1056,73 @@ async fn join_inner_two_two() -> Result<()> { Ok(()) } +/// A streamed row is paired with its buffered key group in runs, each ending where its +/// buffered batch ends or where the output batch fills. Here the group of key 1 spans three +/// buffered batches and the batch size is 4, so runs are cut both ways. Every pair must come +/// out once, in the order of the buffered rows. +#[tokio::test] +async fn join_inner_group_spans_buffered_and_output_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 1, 2]), + ("b1", &vec![0, 0, 0]), + ("c1", &vec![10, 11, 20]), + ); + let right = build_table_from_batches(vec![ + build_table_i32( + ("a2", &vec![1, 1, 1]), + ("b2", &vec![0, 0, 0]), + ("c2", &vec![100, 101, 102]), + ), + build_table_i32( + ("a2", &vec![1, 1]), + ("b2", &vec![0, 0]), + ("c2", &vec![103, 104]), + ), + build_table_i32( + ("a2", &vec![1, 2]), + ("b2", &vec![0, 0]), + ("c2", &vec![105, 200]), + ), + ]); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(4)), + ); + let stream = join(left, right, on, Inner)?.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + assert!(batches.iter().all(|batch| batch.num_rows() <= 4)); + let pairs: Vec<(i32, i32)> = batches + .iter() + .flat_map(|batch| { + let c1 = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let c2 = batch + .column(5) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|i| (c1.value(i), c2.value(i))) + .collect::>() + }) + .collect(); + let expected: Vec<(i32, i32)> = [10, 11] + .into_iter() + .flat_map(|c1| (100..=105).map(move |c2| (c1, c2))) + .chain([(20, 200)]) + .collect(); + assert_eq!(pairs, expected); + Ok(()) +} + #[tokio::test] async fn join_inner_with_nulls() -> Result<()> { let left = build_table_i32_nullable( From f5b8c46747081342ef3daa85aed25ec189411e98 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 13 Oct 2026 12:54:40 -0600 Subject: [PATCH 2/9] test: check the spanning-group join output with a snapshot Match the other sort-merge join tests, which compare the whole output with batches_to_string. --- .../src/joins/sort_merge_join/tests.rs | 44 +++++++++---------- 1 file changed, 20 insertions(+), 24 deletions(-) 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 f23234a240d9d..c9b829f87f723 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -1096,30 +1096,26 @@ async fn join_inner_group_spans_buffered_and_output_batches() -> Result<()> { let batches = common::collect(stream).await?; assert!(batches.iter().all(|batch| batch.num_rows() <= 4)); - let pairs: Vec<(i32, i32)> = batches - .iter() - .flat_map(|batch| { - let c1 = batch - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - let c2 = batch - .column(5) - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()) - .map(|i| (c1.value(i), c2.value(i))) - .collect::>() - }) - .collect(); - let expected: Vec<(i32, i32)> = [10, 11] - .into_iter() - .flat_map(|c1| (100..=105).map(move |c2| (c1, c2))) - .chain([(20, 200)]) - .collect(); - assert_eq!(pairs, expected); + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+-----+ + | 1 | 0 | 10 | 1 | 0 | 100 | + | 1 | 0 | 10 | 1 | 0 | 101 | + | 1 | 0 | 10 | 1 | 0 | 102 | + | 1 | 0 | 10 | 1 | 0 | 103 | + | 1 | 0 | 10 | 1 | 0 | 104 | + | 1 | 0 | 10 | 1 | 0 | 105 | + | 1 | 0 | 11 | 1 | 0 | 100 | + | 1 | 0 | 11 | 1 | 0 | 101 | + | 1 | 0 | 11 | 1 | 0 | 102 | + | 1 | 0 | 11 | 1 | 0 | 103 | + | 1 | 0 | 11 | 1 | 0 | 104 | + | 1 | 0 | 11 | 1 | 0 | 105 | + | 2 | 0 | 20 | 2 | 0 | 200 | + +----+----+----+----+----+-----+ + "); Ok(()) } From 626ef739f7eb6ecac17d5db16fa0d02a29c69337 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 13 Oct 2026 14:00:37 -0600 Subject: [PATCH 3/9] perf: keep the bulk path's setup off one-to-one joins - append_output_pairs hands runs of more than one row to an out-of-line append_output_run, so the one-row case that one-to-one joins take for every row doesn't pay for the bulk path's setup. - chunk_for keeps only its check inline and starts a new chunk in a cold start_chunk. - append_output_pair's only caller passes a null buffered index, so it is now append_null_joined_pair. - BufferedData::scanning_remaining returns the unscanned rows of the scanning batch, and the one-caller scanning_advance is replaced by scanning_advance_by(1). - Move the new test next to the related one-row-batch test. --- .../sort_merge_join/materializing_stream.rs | 151 ++++++++++-------- .../src/joins/sort_merge_join/tests.rs | 127 +++++++-------- 2 files changed, 145 insertions(+), 133 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 0b797924a71ac..106d69d63d221 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 @@ -117,34 +117,24 @@ impl StreamedBatch { self.num_output_rows } - /// Appends new pair consisting of current streamed index and `buffered_idx` - /// index of buffered batch with `buffered_batch_idx` index. - fn append_output_pair( + /// Appends a pair of the current streamed index and a null buffered index, for a + /// streamed row that no buffered row matches. + fn append_null_joined_pair( &mut self, buffered_batch_idx: Option, - buffered_idx: Option, batch_size: usize, ) { let streamed_idx = self.idx as u64; let current_chunk = self.chunk_for(buffered_batch_idx, batch_size); - - // Append index of streamed batch and index of buffered batch into current chunk current_chunk.streamed_indices.append_value(streamed_idx); - if let Some(idx) = buffered_idx { - current_chunk.buffered_indices.append_value(idx as u64); - } else { - current_chunk.buffered_indices.append_null(); - } + current_chunk.buffered_indices.append_null(); self.num_output_rows += 1; } /// Appends one pair for each row in `buffered_range` of the buffered batch with - /// `buffered_batch_idx` index, all with the current streamed index. - /// - /// Equivalent to calling [`Self::append_output_pair`] once per buffered row, but - /// appends the indices in bulk, which matters when a streamed row matches many - /// buffered rows. - #[inline] + /// `buffered_batch_idx` index, all with the current streamed index. The indices + /// are appended in bulk, which matters when a streamed row matches many buffered + /// rows. fn append_output_pairs( &mut self, buffered_batch_idx: usize, @@ -152,39 +142,56 @@ impl StreamedBatch { batch_size: usize, ) { let len = buffered_range.len(); - if len == 0 { - return; - } debug_assert!( - self.num_output_rows + len <= batch_size, - "appending {len} pairs to {} would exceed batch_size ({batch_size})", + len > 0 && self.num_output_rows + len <= batch_size, + "appending {len} pairs to {} with batch_size {batch_size}", self.num_output_rows ); + if len > 1 { + return self.append_output_run( + buffered_batch_idx, + buffered_range, + batch_size, + ); + } + // A single match, as in a one-to-one join, is cheaper to append directly. let streamed_idx = self.idx as u64; let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); - if len == 1 { - // A single match, as in a one-to-one join, is cheaper to append directly. - current_chunk.streamed_indices.append_value(streamed_idx); - current_chunk - .buffered_indices - .append_value(buffered_range.start as u64); - } else { - current_chunk - .streamed_indices - .append_value_n(streamed_idx, len); - // SAFETY: a `Range` reports its exact length in `size_hint`. - unsafe { - current_chunk.buffered_indices.append_trusted_len_iter( - buffered_range.start as u64..buffered_range.end as u64, - ); - } + current_chunk.streamed_indices.append_value(streamed_idx); + current_chunk + .buffered_indices + .append_value(buffered_range.start as u64); + self.num_output_rows += 1; + } + + /// The bulk half of [`Self::append_output_pairs`], kept out of line so that the + /// one-row case, which one-to-one joins take for every row, does not pay for its + /// setup. + #[inline(never)] + fn append_output_run( + &mut self, + buffered_batch_idx: usize, + buffered_range: Range, + batch_size: usize, + ) { + let len = buffered_range.len(); + let streamed_idx = self.idx as u64; + let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); + current_chunk + .streamed_indices + .append_value_n(streamed_idx, len); + // SAFETY: a `Range` reports its exact length in `size_hint`. + unsafe { + current_chunk.buffered_indices.append_trusted_len_iter( + buffered_range.start as u64..buffered_range.end as u64, + ); } self.num_output_rows += len; } /// Returns the chunk that pairs for `buffered_batch_idx` go to, starting a new one /// if there is no chunk yet or the current chunk is for another buffered batch. - #[inline] + #[inline(always)] fn chunk_for( &mut self, buffered_batch_idx: Option, @@ -192,24 +199,31 @@ impl StreamedBatch { ) -> &mut StreamedJoinedChunk { if self.output_indices.is_empty() || self.buffered_batch_idx != buffered_batch_idx { - // Compute capacity only when creating a new chunk (infrequent operation). - // The capacity is the remaining space to reach batch_size. - // This should always be >= 1 since we only call this when num_output_rows < batch_size. - 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; + self.start_chunk(buffered_batch_idx, batch_size); } self.output_indices.last_mut().unwrap() } + + /// Starts the chunk for `buffered_batch_idx`. Kept out of line because a chunk + /// usually takes many pairs, so the check in [`Self::chunk_for`] stays small. + #[cold] + #[inline(never)] + fn start_chunk(&mut self, buffered_batch_idx: Option, batch_size: usize) { + // The capacity is the remaining space to reach batch_size. + // This should always be >= 1 since we only call this when num_output_rows < batch_size. + 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; + } } /// Per-row filter outcome tracking for full outer joins. @@ -781,13 +795,11 @@ impl MaterializingSortMergeJoinStream { { // Pair the streamed row with the rest of the scanning buffered batch at // once, up to the room left in the output batch. - let range = &self.buffered_data.scanning_batch().range; - let start = range.start + self.buffered_data.scanning_offset; - let len = - (range.end - start).min(self.batch_size - self.num_unfrozen_pairs()); + let rows = self.buffered_data.scanning_remaining(); + let len = rows.len().min(self.batch_size - self.num_unfrozen_pairs()); self.streamed_batch.append_output_pairs( self.buffered_data.scanning_batch_idx, - start..start + len, + rows.start..rows.start + len, self.batch_size, ); self.buffered_data.scanning_advance_by(len); @@ -813,11 +825,8 @@ impl MaterializingSortMergeJoinStream { } else { Some(self.buffered_data.scanning_batch_idx) }; - self.streamed_batch.append_output_pair( - scanning_batch_idx, - None, - self.batch_size, - ); + self.streamed_batch + .append_null_joined_pair(scanning_batch_idx, self.batch_size); } self.buffered_data.scanning_reset(); } @@ -833,7 +842,7 @@ impl MaterializingSortMergeJoinStream { .scanning_batch_mut() .null_joined .push(scanning_idx); - self.buffered_data.scanning_advance(); + self.buffered_data.scanning_advance_by(1); } } self.buffered_data.scanning_reset(); @@ -1813,7 +1822,7 @@ impl MaterializingSortMergeJoinStream { // no key comparison. // // The keys a freeze sees are usually a contiguous run, since - // `scanning_advance` walks the deque in order. The exception is a + // `scanning_advance_by` walks the deque in order. The exception is a // freeze that straddles a `scanning_reset`: its window wraps (the // tail of one streamed row's pass, then the head of the next) and // leaves a gap, so the table is sized by the whole group rather than @@ -2172,10 +2181,6 @@ impl BufferedData { self.scanning_offset = 0; } - pub fn scanning_advance(&mut self) { - self.scanning_advance_by(1); - } - /// Advances the scan by `n` rows, which must not pass the end of the scanning /// batch, then skips to the next batch that has rows left. pub fn scanning_advance_by(&mut self, n: usize) { @@ -2199,6 +2204,12 @@ impl BufferedData { self.scanning_batch().range.start + self.scanning_offset } + /// The rows of the scanning batch that the scan has not reached yet. + pub fn scanning_remaining(&self) -> Range { + let range = &self.scanning_batch().range; + range.start + self.scanning_offset..range.end + } + pub fn scanning_batch_finished(&self) -> bool { self.scanning_offset == self.scanning_batch().range.len() } 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 c9b829f87f723..9a77bc08f21ba 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -1056,69 +1056,6 @@ async fn join_inner_two_two() -> Result<()> { Ok(()) } -/// A streamed row is paired with its buffered key group in runs, each ending where its -/// buffered batch ends or where the output batch fills. Here the group of key 1 spans three -/// buffered batches and the batch size is 4, so runs are cut both ways. Every pair must come -/// out once, in the order of the buffered rows. -#[tokio::test] -async fn join_inner_group_spans_buffered_and_output_batches() -> Result<()> { - let left = build_table( - ("a1", &vec![1, 1, 2]), - ("b1", &vec![0, 0, 0]), - ("c1", &vec![10, 11, 20]), - ); - let right = build_table_from_batches(vec![ - build_table_i32( - ("a2", &vec![1, 1, 1]), - ("b2", &vec![0, 0, 0]), - ("c2", &vec![100, 101, 102]), - ), - build_table_i32( - ("a2", &vec![1, 1]), - ("b2", &vec![0, 0]), - ("c2", &vec![103, 104]), - ), - build_table_i32( - ("a2", &vec![1, 2]), - ("b2", &vec![0, 0]), - ("c2", &vec![105, 200]), - ), - ]); - let on = vec![( - Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, - )]; - let task_ctx = Arc::new( - TaskContext::default() - .with_session_config(SessionConfig::new().with_batch_size(4)), - ); - let stream = join(left, right, on, Inner)?.execute(0, task_ctx)?; - let batches = common::collect(stream).await?; - - assert!(batches.iter().all(|batch| batch.num_rows() <= 4)); - // The output order is important as SMJ preserves sortedness - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+----+----+-----+ - | a1 | b1 | c1 | a2 | b2 | c2 | - +----+----+----+----+----+-----+ - | 1 | 0 | 10 | 1 | 0 | 100 | - | 1 | 0 | 10 | 1 | 0 | 101 | - | 1 | 0 | 10 | 1 | 0 | 102 | - | 1 | 0 | 10 | 1 | 0 | 103 | - | 1 | 0 | 10 | 1 | 0 | 104 | - | 1 | 0 | 10 | 1 | 0 | 105 | - | 1 | 0 | 11 | 1 | 0 | 100 | - | 1 | 0 | 11 | 1 | 0 | 101 | - | 1 | 0 | 11 | 1 | 0 | 102 | - | 1 | 0 | 11 | 1 | 0 | 103 | - | 1 | 0 | 11 | 1 | 0 | 104 | - | 1 | 0 | 11 | 1 | 0 | 105 | - | 2 | 0 | 20 | 2 | 0 | 200 | - +----+----+----+----+----+-----+ - "); - Ok(()) -} - #[tokio::test] async fn join_inner_with_nulls() -> Result<()> { let left = build_table_i32_nullable( @@ -5147,6 +5084,70 @@ async fn join_with_group_spanning_batches_rescanned_per_streamed_row() -> Result Ok(()) } +/// A streamed row is paired with its buffered key group in runs, each ending where its +/// buffered batch ends or where the output batch fills. The test above uses one-row buffered +/// batches, so its runs have one row. Here the group of key 1 spans buffered batches of 3, 2 +/// and 1 rows and the batch size is 4, so runs of several rows are cut both ways. Every pair +/// must come out once, in the order of the buffered rows. +#[tokio::test] +async fn join_inner_group_spans_buffered_and_output_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 1, 2]), + ("b1", &vec![0, 0, 0]), + ("c1", &vec![10, 11, 20]), + ); + let right = build_table_from_batches(vec![ + build_table_i32( + ("a2", &vec![1, 1, 1]), + ("b2", &vec![0, 0, 0]), + ("c2", &vec![100, 101, 102]), + ), + build_table_i32( + ("a2", &vec![1, 1]), + ("b2", &vec![0, 0]), + ("c2", &vec![103, 104]), + ), + build_table_i32( + ("a2", &vec![1, 2]), + ("b2", &vec![0, 0]), + ("c2", &vec![105, 200]), + ), + ]); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(4)), + ); + let stream = join(left, right, on, Inner)?.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + assert!(batches.iter().all(|batch| batch.num_rows() <= 4)); + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+-----+ + | 1 | 0 | 10 | 1 | 0 | 100 | + | 1 | 0 | 10 | 1 | 0 | 101 | + | 1 | 0 | 10 | 1 | 0 | 102 | + | 1 | 0 | 10 | 1 | 0 | 103 | + | 1 | 0 | 10 | 1 | 0 | 104 | + | 1 | 0 | 10 | 1 | 0 | 105 | + | 1 | 0 | 11 | 1 | 0 | 100 | + | 1 | 0 | 11 | 1 | 0 | 101 | + | 1 | 0 | 11 | 1 | 0 | 102 | + | 1 | 0 | 11 | 1 | 0 | 103 | + | 1 | 0 | 11 | 1 | 0 | 104 | + | 1 | 0 | 11 | 1 | 0 | 105 | + | 2 | 0 | 20 | 2 | 0 | 200 | + +----+----+----+----+----+-----+ + "); + Ok(()) +} + /// A wrapped multi-source freeze that also carries a null buffered index. /// /// `materialize_right_columns` has two independent offsets in play on the From 2e39009590fb32a28076d1431ac64a71d5bb8291 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 13:07:19 -0600 Subject: [PATCH 4/9] refactor: pass a run's start and length to the bulk append - pair_streamed_row_with_group computes the room left in the output batch once and counts it down, and passes the start and length of each run instead of a Range. - append_output_pairs matches on the length, with the bulk path as the catch-all arm, so a run of no rows would append no pair in a release build. The loop never produces one, and the debug_assert stays. - Rename append_output_run to append_output_range, and drop a comment that repeats the docs. --- .../sort_merge_join/materializing_stream.rs | 67 ++++++++++--------- 1 file changed, 35 insertions(+), 32 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 106d69d63d221..0df3e2aa8afa2 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 @@ -131,60 +131,64 @@ impl StreamedBatch { self.num_output_rows += 1; } - /// Appends one pair for each row in `buffered_range` of the buffered batch with - /// `buffered_batch_idx` index, all with the current streamed index. The indices - /// are appended in bulk, which matters when a streamed row matches many buffered - /// rows. + /// Appends a pair of the current streamed index and each of the `len` rows from + /// `buffered_start` on in the buffered batch with `buffered_batch_idx` index. The + /// indices are appended in bulk, which matters when a streamed row matches many + /// buffered rows. fn append_output_pairs( &mut self, buffered_batch_idx: usize, - buffered_range: Range, + buffered_start: usize, + len: usize, batch_size: usize, ) { - let len = buffered_range.len(); debug_assert!( len > 0 && self.num_output_rows + len <= batch_size, "appending {len} pairs to {} with batch_size {batch_size}", self.num_output_rows ); - if len > 1 { - return self.append_output_run( + match len { + // A single match, as in a one-to-one join, is cheaper to append directly. + 1 => { + let streamed_idx = self.idx as u64; + let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); + current_chunk.streamed_indices.append_value(streamed_idx); + current_chunk + .buffered_indices + .append_value(buffered_start as u64); + self.num_output_rows += 1; + } + _ => self.append_output_range( buffered_batch_idx, - buffered_range, + buffered_start, + len, batch_size, - ); + ), } - // A single match, as in a one-to-one join, is cheaper to append directly. - let streamed_idx = self.idx as u64; - let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); - current_chunk.streamed_indices.append_value(streamed_idx); - current_chunk - .buffered_indices - .append_value(buffered_range.start as u64); - self.num_output_rows += 1; } /// The bulk half of [`Self::append_output_pairs`], kept out of line so that the /// one-row case, which one-to-one joins take for every row, does not pay for its /// setup. #[inline(never)] - fn append_output_run( + fn append_output_range( &mut self, buffered_batch_idx: usize, - buffered_range: Range, + buffered_start: usize, + len: usize, batch_size: usize, ) { - let len = buffered_range.len(); let streamed_idx = self.idx as u64; let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); current_chunk .streamed_indices .append_value_n(streamed_idx, len); + let start = buffered_start as u64; // SAFETY: a `Range` reports its exact length in `size_hint`. unsafe { - current_chunk.buffered_indices.append_trusted_len_iter( - buffered_range.start as u64..buffered_range.end as u64, - ); + current_chunk + .buffered_indices + .append_trusted_len_iter(start..start + len as u64); } self.num_output_rows += len; } @@ -790,21 +794,20 @@ impl MaterializingSortMergeJoinStream { /// 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 - { - // Pair the streamed row with the rest of the scanning buffered batch at - // once, up to the room left in the output batch. + let mut room = self.batch_size - self.num_unfrozen_pairs(); + while !self.buffered_data.scanning_finished() && room > 0 { let rows = self.buffered_data.scanning_remaining(); - let len = rows.len().min(self.batch_size - self.num_unfrozen_pairs()); + let len = rows.len().min(room); self.streamed_batch.append_output_pairs( self.buffered_data.scanning_batch_idx, - rows.start..rows.start + len, + rows.start, + len, self.batch_size, ); self.buffered_data.scanning_advance_by(len); + room -= len; } - if self.num_unfrozen_pairs() >= self.batch_size { + if room == 0 { return false; } From 2f6064f2e9c2a4e0789a49135f82bcda817c89f2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 13:07:19 -0600 Subject: [PATCH 5/9] refactor: append a run's buffered indices without unsafe Replace append_trusted_len_iter with a loop of append_value. The chunk's builders are created with capacity for the room left in the output batch, so the loop needs no reserve. The loop is about 5% slower at 100 matches per row, but leaves the module without unsafe. --- .../src/joins/sort_merge_join/materializing_stream.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 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 0df3e2aa8afa2..9c9fb0724b088 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 @@ -184,11 +184,8 @@ impl StreamedBatch { .streamed_indices .append_value_n(streamed_idx, len); let start = buffered_start as u64; - // SAFETY: a `Range` reports its exact length in `size_hint`. - unsafe { - current_chunk - .buffered_indices - .append_trusted_len_iter(start..start + len as u64); + for idx in start..start + len as u64 { + current_chunk.buffered_indices.append_value(idx); } self.num_output_rows += len; } From 2a112e7d7fd3f411147523701e70712f0548c201 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 13:07:19 -0600 Subject: [PATCH 6/9] bench: add a filtered sort-merge left join with 100 matches per row left_1to100_filtered joins the inputs of inner_1to100 as a left join with the filter left.data < right.data, so that the pairs go through deferred filtering. The filter keeps about half of them, and none of the last streamed row of each key, which is null-joined. --- .../physical-plan/benches/sort_merge_join.rs | 69 +++++++++++++++++-- 1 file changed, 65 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index 03dec34eaded7..9d750976fb909 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -27,11 +27,13 @@ use arrow::array::{Int64Array, RecordBatch, StringArray}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; -use datafusion_common::NullEquality; +use datafusion_common::{JoinSide, NullEquality}; use datafusion_execution::TaskContext; -use datafusion_physical_expr::expressions::col; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, col}; use datafusion_physical_plan::collect; -use datafusion_physical_plan::joins::{SortMergeJoinExec, utils::JoinOn}; +use datafusion_physical_plan::joins::SortMergeJoinExec; +use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; use datafusion_physical_plan::test::TestMemoryExec; use tokio::runtime::Runtime; @@ -104,11 +106,45 @@ fn schema() -> SchemaRef { ])) } +/// The join filter `left.data < right.data`. +fn data_lt_filter() -> JoinFilter { + let column_indices = vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + let intermediate_schema = Schema::new(vec![ + Field::new("left_data", DataType::Int64, false), + Field::new("right_data", DataType::Int64, false), + ]); + let expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("left_data", 0)), + Operator::Lt, + Arc::new(Column::new("right_data", 1)), + )); + JoinFilter::new(expression, column_indices, Arc::new(intermediate_schema)) +} + fn do_join( left: Arc, right: Arc, join_type: datafusion_common::JoinType, rt: &Runtime, +) -> usize { + do_join_with_filter(left, right, join_type, None, rt) +} + +fn do_join_with_filter( + left: Arc, + right: Arc, + join_type: datafusion_common::JoinType, + filter: Option, + rt: &Runtime, ) -> usize { let on: JoinOn = vec![( col("key", &left.schema()).unwrap(), @@ -118,7 +154,7 @@ fn do_join( left, right, on, - None, + filter, join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, @@ -184,6 +220,31 @@ fn bench_smj(c: &mut Criterion) { }); } + // 1:100 Left Join with a filter — 20K left, 20K right, 200 distinct keys + // As in `inner_1to100`, every streamed row matches 100 buffered rows on the key, + // but here the pairs go through deferred filtering. `left.data < right.data` + // keeps about half of them, and none of the last streamed row of each key, which + // is null-joined. + { + let n = 20_000; + let key_mod = 200; + let left_batches = build_sorted_batches(n, key_mod, &s); + let right_batches = build_sorted_batches(n, key_mod, &s); + group.bench_function(BenchmarkId::new("left_1to100_filtered", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join_with_filter( + left, + right, + datafusion_common::JoinType::Left, + Some(data_lt_filter()), + &rt, + ) + }) + }); + } + // Left Join — 100K each, ~5% unmatched on left { let n = 100_000; From 287d0a00b83674d15435aa8a1597f7309b30180d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 16 Oct 2026 13:07:19 -0600 Subject: [PATCH 7/9] test: fuzz sort-merge joins with low-cardinality keys Most key groups in the join fuzz inputs hold one row, so a sort-merge join almost never pairs a streamed row with a run of more than one buffered row there, and never cuts such a run at the output batch. Add Inner, Left, Right and Full cases, unfiltered and filtered, whose keys take 25 values, so that each streamed row matches about 40 buffered rows and runs end both at buffered batch boundaries and where the output batch fills. --- datafusion/core/tests/fuzz_cases/join_fuzz.rs | 133 +++++++++++++++++- 1 file changed, 132 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 296136fd7b3d1..7a30a56d590f7 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -662,6 +662,124 @@ async fn test_right_mark_join_1k_binary_filtered() { } } +// In the tests above, most key groups hold one row: the i32 tests draw the join keys +// `(a, b)` from 10,000 values for 1000 rows, and the binary ones from random bytes. +// The tests below draw them from 25 values, so each streamed row matches about 40 +// buffered rows. Sort-merge join pairs a streamed row with its matches in runs, which +// end where a buffered batch ends or where the output batch fills. + +#[tokio::test] +async fn test_inner_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Inner, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_inner_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Inner, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_left_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Left, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_left_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Left, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_right_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Right, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_right_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Right, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_full_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Full, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_full_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Full, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + type JoinFilterBuilder = Box, Arc) -> JoinFilter>; struct JoinFuzzTestCase { @@ -1392,12 +1510,25 @@ async fn test_nested_loop_join_spill_fuzz() { /// two sorted int32 columns 'a', 'b' ranged from 0..99 as join columns /// two random int32 columns 'x', 'y' as other columns fn make_staggered_batches_i32(len: usize, with_extra_column: bool) -> Vec { + make_staggered_batches_i32_with_key_range(len, 100, with_extra_column) +} + +/// Like [`make_staggered_batches_i32`], but with 'a' and 'b' drawn from +/// `0..key_range`, so that a small `key_range` makes few, large key groups +fn make_staggered_batches_i32_with_key_range( + len: usize, + key_range: i32, + with_extra_column: bool, +) -> Vec { let mut rng = rand::rng(); let mut input12: Vec<(i32, i32)> = vec![(0, 0); len]; let mut input3: Vec = vec![0; len]; let mut input4: Vec = vec![0; len]; for v in &mut input12 { - *v = (rng.random_range(0..100), rng.random_range(0..100)); + *v = ( + rng.random_range(0..key_range), + rng.random_range(0..key_range), + ); } rng.fill(&mut input3[..]); rng.fill(&mut input4[..]); From 413061472573ce769f81573d3f9ca5669831b5c2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 19 Oct 2026 18:31:28 +0100 Subject: [PATCH 8/9] perf: keep sort-merge join chunk indices in Vecs StreamedJoinedChunk keeps its indices in Vecs, with a NullBufferBuilder for the buffered nulls, instead of in UInt64Builders, as hash join does. A run of buffered rows is then one Vec::extend over the range and one append_n_non_nulls, so the per-row null and capacity checks move out of the loop. freeze_streamed builds the arrays with UInt64Array::try_new, which skips the ArrayData that PrimitiveBuilder::finish builds for every chunk. --- .../sort_merge_join/materializing_stream.rs | 48 +++++++++++-------- 1 file changed, 28 insertions(+), 20 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 9c9fb0724b088..2181f8c7e574a 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 @@ -25,7 +25,8 @@ use std::cmp::Ordering; use std::collections::VecDeque; use std::fmt::Debug; -use std::mem::size_of; +use std::iter::repeat_n; +use std::mem::{size_of, take}; use std::ops::Range; use std::sync::Arc; @@ -62,11 +63,12 @@ use futures::StreamExt; pub(super) struct StreamedJoinedChunk { /// Index of batch in buffered_data buffered_batch_idx: Option, - /// Array builder for streamed indices - streamed_indices: UInt64Builder, - /// Array builder for buffered indices - /// This could contain nulls if the join is null-joined - buffered_indices: UInt64Builder, + /// Streamed indices, which are never null + streamed_indices: Vec, + /// Buffered indices. The slot of a null-joined pair holds 0. + buffered_indices: Vec, + /// Validity of `buffered_indices`, null where the pair is null-joined + buffered_nulls: NullBufferBuilder, } /// Represents a record batch from streamed input. @@ -126,8 +128,9 @@ impl StreamedBatch { ) { let streamed_idx = self.idx as u64; let current_chunk = self.chunk_for(buffered_batch_idx, batch_size); - current_chunk.streamed_indices.append_value(streamed_idx); - current_chunk.buffered_indices.append_null(); + current_chunk.streamed_indices.push(streamed_idx); + current_chunk.buffered_indices.push(0); + current_chunk.buffered_nulls.append_null(); self.num_output_rows += 1; } @@ -152,10 +155,9 @@ impl StreamedBatch { 1 => { let streamed_idx = self.idx as u64; let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); - current_chunk.streamed_indices.append_value(streamed_idx); - current_chunk - .buffered_indices - .append_value(buffered_start as u64); + current_chunk.streamed_indices.push(streamed_idx); + current_chunk.buffered_indices.push(buffered_start as u64); + current_chunk.buffered_nulls.append_non_null(); self.num_output_rows += 1; } _ => self.append_output_range( @@ -182,11 +184,12 @@ impl StreamedBatch { let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); current_chunk .streamed_indices - .append_value_n(streamed_idx, len); + .extend(repeat_n(streamed_idx, len)); let start = buffered_start as u64; - for idx in start..start + len as u64 { - current_chunk.buffered_indices.append_value(idx); - } + current_chunk + .buffered_indices + .extend(start..start + len as u64); + current_chunk.buffered_nulls.append_n_non_nulls(len); self.num_output_rows += len; } @@ -220,8 +223,9 @@ impl StreamedBatch { 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), + streamed_indices: Vec::with_capacity(capacity), + buffered_indices: Vec::with_capacity(capacity), + buffered_nulls: NullBufferBuilder::new(capacity), }); self.buffered_batch_idx = buffered_batch_idx; } @@ -1587,11 +1591,15 @@ impl MaterializingSortMergeJoinStream { let mut total_matched_rows: usize = 0; for chunk in self.streamed_batch.output_indices.iter_mut() { - let left_indices = chunk.streamed_indices.finish(); + let left_indices = + UInt64Array::try_new(take(&mut chunk.streamed_indices).into(), None)?; if left_indices.is_empty() { continue; } - let right_indices: UInt64Array = chunk.buffered_indices.finish(); + let right_indices = UInt64Array::try_new( + take(&mut chunk.buffered_indices).into(), + chunk.buffered_nulls.finish(), + )?; if chunk.buffered_batch_idx.is_none() { let left_columns = From 1ce28a4e6aa4443f7fd22976ba277369d10e6c12 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 19 Oct 2026 18:31:37 +0100 Subject: [PATCH 9/9] refactor: saturate the room left in the output batch On main, pair_streamed_row_with_group appends nothing and returns false when it is called with a full output batch, so the caller freezes and calls again. Computing the room with saturating_sub keeps that behavior where the subtraction would underflow. The callers keep the count below batch_size, so this only matters if a later change breaks that. --- .../src/joins/sort_merge_join/materializing_stream.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 2181f8c7e574a..892be4a2f39ec 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 @@ -795,7 +795,7 @@ impl MaterializingSortMergeJoinStream { /// 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 { - let mut room = self.batch_size - self.num_unfrozen_pairs(); + let mut room = self.batch_size.saturating_sub(self.num_unfrozen_pairs()); while !self.buffered_data.scanning_finished() && room > 0 { let rows = self.buffered_data.scanning_remaining(); let len = rows.len().min(room);