From b3d430c4ad8a3987f41fe3e5296816ec2cc15984 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Sat, 22 Aug 2026 15:59:02 +0800
Subject: [PATCH 1/7] fix: ensure deferred-filtered outer joins preserve
 streamed output order

---
 .../sort_merge_join/materializing_stream.rs   | 136 ++++++++++--------
 .../src/joins/sort_merge_join/tests.rs        |  94 ++++++++++++
 2 files changed, 167 insertions(+), 63 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 3baa0c4a3e792..48839e44ca082 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
@@ -23,7 +23,7 @@
 //! produces joined `RecordBatch`es.
 
 use std::cmp::Ordering;
-use std::collections::{HashMap, VecDeque};
+use std::collections::VecDeque;
 use std::fmt::Debug;
 use std::mem::size_of;
 use std::ops::Range;
@@ -89,16 +89,16 @@ pub(super) struct StreamedBatch {
 }
 
 impl StreamedBatch {
-    fn new(batch: RecordBatch, on_column: &[Arc<dyn PhysicalExpr>]) -> Self {
-        let join_arrays = join_arrays(&batch, on_column);
-        StreamedBatch {
+    fn try_new(batch: RecordBatch, on_column: &[Arc<dyn PhysicalExpr>]) -> Result<Self> {
+        let join_arrays = join_arrays(&batch, on_column)?;
+        Ok(StreamedBatch {
             batch,
             idx: 0,
             join_arrays,
             output_indices: vec![],
             num_output_rows: 0,
             buffered_batch_idx: None,
-        }
+        })
     }
 
     fn new_empty(schema: SchemaRef) -> Self {
@@ -213,12 +213,12 @@ pub(super) struct BufferedBatch {
 }
 
 impl BufferedBatch {
-    fn new(
+    fn try_new(
         batch: RecordBatch,
         range: Range<usize>,
         on_column: &[PhysicalExprRef],
-    ) -> Self {
-        let join_arrays = join_arrays(&batch, on_column);
+    ) -> Result<Self> {
+        let join_arrays = join_arrays(&batch, on_column)?;
 
         // Estimation is calculated as
         //   inner batch size
@@ -238,7 +238,7 @@ impl BufferedBatch {
             + size_of::<usize>();
 
         let num_rows = batch.num_rows();
-        BufferedBatch {
+        Ok(BufferedBatch {
             batch: BufferedBatchState::InMemory(batch),
             range,
             join_arrays,
@@ -248,7 +248,7 @@ impl BufferedBatch {
             reserved_amount: 0,
             join_filter_status: vec![FilterState::Unvisited; num_rows],
             num_rows,
-        }
+        })
     }
 }
 
@@ -414,8 +414,7 @@ impl JoinedRecordBatches {
 
     /// Clears batches without touching metadata (for early return when no filtering needed)
     fn clear_batches(&mut self, schema: &SchemaRef, batch_size: usize) {
-        self.joined_batches = BatchCoalescer::new(Arc::clone(schema), batch_size)
-            .with_biggest_coalesce_batch_size(Option::from(batch_size / 2));
+        self.joined_batches = new_output_coalescer(Arc::clone(schema), batch_size);
     }
 
     /// Asserts that if batches is empty, metadata is also empty
@@ -517,8 +516,7 @@ impl JoinedRecordBatches {
     }
 
     fn clear(&mut self, schema: &SchemaRef, batch_size: usize) {
-        self.joined_batches = BatchCoalescer::new(Arc::clone(schema), batch_size)
-            .with_biggest_coalesce_batch_size(Option::from(batch_size / 2));
+        self.joined_batches = new_output_coalescer(Arc::clone(schema), batch_size);
         self.filter_metadata = FilterMetadata::new();
         self.debug_assert_empty_consistency();
     }
@@ -571,12 +569,10 @@ impl MaterializingSortMergeJoinStream {
             deferred_filtering: needs_deferred_filtering(&filter, join_type),
             filter,
             joined_record_batches: JoinedRecordBatches {
-                joined_batches: BatchCoalescer::new(Arc::clone(&schema), batch_size)
-                    .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)),
+                joined_batches: new_output_coalescer(Arc::clone(&schema), batch_size),
                 filter_metadata: FilterMetadata::new(),
             },
-            output: BatchCoalescer::new(schema, batch_size)
-                .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)),
+            output: new_output_coalescer(schema, batch_size),
             batch_size,
             join_type,
             join_metrics,
@@ -800,14 +796,28 @@ impl MaterializingSortMergeJoinStream {
         // 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()? {
+        self.stage_filtered_output()?;
+        self.emit_completed_output(emitter).await;
+        Ok(())
+    }
+
+    /// Emit every completed batch of the deferred-filtering output buffer.
+    ///
+    /// All deferred-filtered output must leave through this single buffer:
+    /// emitting a batch around it would reorder it ahead of rows still
+    /// buffered here, breaking the streamed-side ordering the operator
+    /// advertises via `maintains_input_order`.
+    async fn emit_completed_output(
+        &mut self,
+        emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+    ) {
+        while let Some(record_batch) = self.output.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(batch).await;
+            emitter.emit(record_batch).await;
             self.start_join_time();
         }
-        Ok(())
     }
 
     /// Restore every spilled buffered batch that the next freeze needs.
@@ -849,12 +859,15 @@ impl MaterializingSortMergeJoinStream {
             .debug_assert_metadata_aligned();
 
         if self.deferred_filtering {
-            // Filtered joins must concat and filter ALL remaining data at once
+            // Filtered joins must concat and filter ALL remaining data at
+            // once. The result is staged in `output` rather than emitted
+            // directly: `output` may still hold rows from earlier flushes,
+            // and those precede these on the streamed side.
             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();
+                self.output
+                    .push_batch(record_batch)
+                    .expect("Failed to push output batch");
             }
         } else if !self.joined_record_batches.joined_batches.is_empty() {
             // For non-filtered joins, finish buffered data first, then emit
@@ -868,11 +881,7 @@ impl MaterializingSortMergeJoinStream {
         // 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() {
-                self.stop_join_time();
-                emitter.emit(record_batch).await;
-                self.start_join_time();
-            }
+            self.emit_completed_output(emitter).await;
         }
 
         Ok(())
@@ -916,11 +925,12 @@ impl MaterializingSortMergeJoinStream {
         self.streamed_batch.num_output_rows()
     }
 
-    /// Process accumulated batches for filtered joins
+    /// Process accumulated batches for filtered joins.
     ///
-    /// Freezes unfrozen pairs, applies deferred filtering, and returns a
-    /// completed output batch if one is ready.
-    fn process_filtered_batches(&mut self) -> Result<Option<RecordBatch>> {
+    /// Freezes unfrozen pairs, applies deferred filtering and stages the
+    /// result in [`Self::output`]. Completed batches are emitted separately
+    /// by [`Self::emit_completed_output`].
+    fn stage_filtered_output(&mut self) -> Result<()> {
         self.freeze_all()?;
 
         self.joined_record_batches
@@ -932,17 +942,9 @@ impl MaterializingSortMergeJoinStream {
             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");
-                return Ok(Some(record_batch));
-            }
         }
 
-        Ok(None)
+        Ok(())
     }
 
     /// Identifies which buffered batches are needed for the upcoming freeze operation
@@ -1054,7 +1056,7 @@ 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::try_new(batch, &self.on_streamed)?;
                         self.rebuild_streamed_buffered_cmp()?;
                         // Every incoming streamed batch gets a unique id.
                         self.streamed_batch_counter += 1;
@@ -1242,7 +1244,7 @@ impl MaterializingSortMergeJoinStream {
 
                     if batch.num_rows() > 0 {
                         let buffered_batch =
-                            BufferedBatch::new(batch, 0..1, &self.on_buffered);
+                            BufferedBatch::try_new(batch, 0..1, &self.on_buffered)?;
                         self.allocate_reservation(buffered_batch)?;
                         self.streamed_buffered_cmp = None;
                         return Ok(true);
@@ -1297,7 +1299,7 @@ impl MaterializingSortMergeJoinStream {
                         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);
+                                BufferedBatch::try_new(batch, 0..0, &self.on_buffered)?;
                             self.allocate_reservation(buffered_batch)?;
                             self.buffered_equality_cmp = None;
                         }
@@ -1665,20 +1667,19 @@ impl MaterializingSortMergeJoinStream {
 
         // Multiple source batches: map each buffered_batch_idx to a
         // contiguous source index, reserving source 0 for a null sentinel.
-        let mut batch_idx_to_source: HashMap<usize, usize> = HashMap::new();
+        // A group spans only a handful of buffered batches, so a linear
+        // scan beats hashing here.
         let mut source_batches: Vec<usize> = Vec::new();
-        for (batch_idx, _, _) in matched_chunks {
-            batch_idx_to_source.entry(*batch_idx).or_insert_with(|| {
-                let idx = source_batches.len() + 1;
-                source_batches.push(*batch_idx);
-                idx
-            });
-        }
-
         let mut interleave_indices: Vec<(usize, usize)> =
             Vec::with_capacity(total_matched_rows);
         for (batch_idx, _, right) in matched_chunks {
-            let source = batch_idx_to_source[batch_idx];
+            let source = match source_batches.iter().position(|b| b == batch_idx) {
+                Some(pos) => pos + 1,
+                None => {
+                    source_batches.push(*batch_idx);
+                    source_batches.len()
+                }
+            };
             for i in 0..right.len() {
                 if right.is_null(i) {
                     interleave_indices.push((0, 0));
@@ -1987,14 +1988,23 @@ impl BufferedData {
     }
 }
 
-/// Get join array refs of given batch and join columns
-fn join_arrays(batch: &RecordBatch, on_column: &[PhysicalExprRef]) -> Vec<ArrayRef> {
+/// Build the `BatchCoalescer` used for staging join output.
+///
+/// `biggest_coalesce_batch_size` lets batches larger than half the target
+/// pass through without being copied into the coalescer's buffer.
+fn new_output_coalescer(schema: SchemaRef, batch_size: usize) -> BatchCoalescer {
+    BatchCoalescer::new(schema, batch_size)
+        .with_biggest_coalesce_batch_size(Some(batch_size / 2))
+}
+
+/// Evaluate the join key expressions against `batch`.
+fn join_arrays(
+    batch: &RecordBatch,
+    on_column: &[PhysicalExprRef],
+) -> Result<Vec<ArrayRef>> {
+    let num_rows = batch.num_rows();
     on_column
         .iter()
-        .map(|c| {
-            let num_rows = batch.num_rows();
-            let c = c.evaluate(batch).unwrap();
-            c.into_array(num_rows).unwrap()
-        })
+        .map(|c| c.evaluate(batch)?.into_array(num_rows))
         .collect()
 }
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 91d1b893f1b29..4cf862b4ca8b0 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -5985,3 +5985,97 @@ async fn bitwise_spill_pending_stream() -> Result<()> {
 
     Ok(())
 }
+
+/// Regression test: deferred-filtered outer joins must not reorder their
+/// output.
+///
+/// `LEFT JOIN` advertises `maintains_input_order = [true, false]`, so the
+/// output must stay ordered on the streamed side. The final flush used to
+/// emit its batch directly instead of through the `output` coalescer, so any
+/// rows still buffered there from an earlier flush were emitted *after* it.
+///
+/// The shape below reproduces that: the first five keys each match a large
+/// buffered group, so the deferred-filter gate fires once per key and pushes
+/// a single-row batch into `output` (too small to complete a batch), while
+/// the last two keys match a single row each and so never trip the gate —
+/// leaving their rows for the final flush.
+#[tokio::test]
+async fn left_join_with_filter_preserves_streamed_order() -> Result<()> {
+    let num_keys = 7i32;
+
+    let keys: Vec<i32> = (0..num_keys).collect();
+    let left = build_table_i32(("a1", &keys), ("b1", &keys), ("c1", &keys));
+
+    let mut r_a = vec![];
+    let mut r_b = vec![];
+    let mut r_c = vec![];
+    for k in 0..num_keys {
+        let dup = if k < 5 { 20 } else { 1 };
+        for j in 0..dup {
+            r_a.push(k * 100 + j);
+            r_b.push(k);
+            r_c.push(j);
+        }
+    }
+    let right = build_table_i32(("a2", &r_a), ("b2", &r_b), ("c2", &r_c));
+
+    let left_schema = left.schema();
+    let right_schema = right.schema();
+    let left = TestMemoryExec::try_new_exec(&[vec![left]], left_schema, None)?;
+    let right = TestMemoryExec::try_new_exec(&[vec![right]], right_schema, None)?;
+
+    let on: JoinOn = vec![(
+        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
+    )];
+
+    // A filter that never passes, so every streamed row is emitted
+    // null-joined by the deferred-filtering pipeline.
+    let intermediate_schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
+    let filter = JoinFilter::new(
+        Arc::new(BinaryExpr::new(
+            Arc::new(Column::new("x", 0)),
+            Operator::Lt,
+            Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
+        )) as PhysicalExprRef,
+        vec![ColumnIndex {
+            index: 0,
+            side: JoinSide::Left,
+        }],
+        Arc::new(intermediate_schema),
+    );
+
+    let join = SortMergeJoinExec::try_new(
+        left,
+        right,
+        on,
+        Some(filter),
+        Left,
+        vec![SortOptions::default()],
+        NullEquality::NullEqualsNothing,
+    )?;
+
+    let task_ctx = Arc::new(
+        TaskContext::default()
+            .with_session_config(SessionConfig::default().with_batch_size(8)),
+    );
+    let batches = common::collect(join.execute(0, task_ctx)?).await?;
+
+    let streamed_keys: Vec<i32> = batches
+        .iter()
+        .flat_map(|b| {
+            b.column(0)
+                .as_any()
+                .downcast_ref::<Int32Array>()
+                .unwrap()
+                .values()
+                .to_vec()
+        })
+        .collect();
+
+    assert_eq!(
+        streamed_keys, keys,
+        "LEFT JOIN output must stay ordered on the streamed side"
+    );
+    Ok(())
+}

From ee003d586f3976894d81117abb398aa0fe61eb60 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Sat, 22 Aug 2026 21:20:15 +0800
Subject: [PATCH 2/7] test: add right join and partial filter tests to preserve
 streamed order

---
 .../sort_merge_join/materializing_stream.rs   | 130 ++++++++++---
 .../src/joins/sort_merge_join/tests.rs        | 184 ++++++++++++++++++
 2 files changed, 284 insertions(+), 30 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 48839e44ca082..c3e04f09bb203 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
@@ -1642,7 +1642,7 @@ impl MaterializingSortMergeJoinStream {
     /// gathers columns across sources. A null-row sentinel at source index 0
     /// handles null right indices (unmatched streamed rows).
     fn materialize_right_columns(
-        &mut self,
+        &self,
         matched_chunks: &[(usize, UInt64Array, UInt64Array)],
         total_matched_rows: usize,
     ) -> Result<Vec<ArrayRef>> {
@@ -1666,25 +1666,92 @@ impl MaterializingSortMergeJoinStream {
         }
 
         // Multiple source batches: map each buffered_batch_idx to a
-        // contiguous source index, reserving source 0 for a null sentinel.
+        // contiguous source index. A null sentinel array is prepended as
+        // source 0 only when some right index is actually null (an
+        // unmatched streamed row inside an otherwise matched chunk);
+        // `interleave` walks a null buffer for *every* output row as soon as
+        // any input is nullable, so an always-present sentinel would tax the
+        // common all-matched case.
+        let needs_null_sentinel = matched_chunks
+            .iter()
+            .any(|(_, _, right)| right.null_count() > 0);
+        let source_offset = usize::from(needs_null_sentinel);
+
         // A group spans only a handful of buffered batches, so a linear
-        // scan beats hashing here.
+        // scan beats hashing here. Measured over 8192 rows in 2048 chunks,
+        // against a `HashMap<usize, usize>` built in one pass and read back
+        // in a second (what this used to do):
+        //
+        //   distinct sources |  hashmap  |  linear scan
+        //   -----------------+-----------+-------------
+        //                  4 |  21.5 us  |   5.0 us
+        //                 16 |  22.0 us  |   9.4 us
+        //                 32 |  22.4 us  |  13.6 us
+        //                 64 |  22.9 us  |  23.5 us
+        //                128 |  24.1 us  |  44.8 us
+        //
+        // `std::collections::HashMap` hashes with SipHash-1-3, so a single
+        // `usize` lookup costs several ns of serial latency before the probe
+        // begins, while a scan over a handful of `usize` is one L1-resident
+        // cache line with a perfectly predicted trip count. The map is also
+        // purely additive state: `source_batches` has to be built regardless
+        // (`source_data` is gathered from it), so hashing means maintaining
+        // two containers holding the same keys.
+        //
+        // The crossover is ~32 distinct sources. That bound follows from how
+        // pairs accumulate, not from any assumption about key skew:
+        //
+        //   1. `pair_streamed_row_with_group` appends exactly one pair per
+        //      buffered row and re-checks `num_unfrozen_pairs() < batch_size`
+        //      before each append, so at most `batch_size` pairs accumulate
+        //      between two `freeze_streamed()` calls.
+        //   2. `BufferedData::scanning_advance` walks the group's rows in
+        //      order, so those pairs cover a *contiguous run* of buffered
+        //      rows.
+        //   3. So the distinct `buffered_batch_idx` values seen here are the
+        //      batches spanned by at most `batch_size` consecutive buffered
+        //      rows: `len(source_batches) <= batch_size / R + 1`, where `R`
+        //      is the smallest buffered batch in that run.
+        //
+        // The assumption is therefore not "key groups are narrow" — a group
+        // of any width still only contributes `batch_size` rows per freeze —
+        // but "buffered batches are not tiny relative to `batch_size`".
+        // Exceeding 32 sources needs `R < batch_size / 31`, i.e. under ~264
+        // rows per batch at the default `batch_size` of 8192. The buffered
+        // side of a merge join is sorted input, and every operator that
+        // normally feeds it emits ~`batch_size` batches: `SortExec` chunks
+        // its output with `sort_batch_chunked(.., batch_size)`, and
+        // `FilterExec` and `RepartitionExec` each embed a
+        // `LimitedBatchCoalescer` targeting `batch_size`.
+        //
+        // If something does feed tiny batches, this degrades gradually rather
+        // than falling off a cliff, and never affects correctness: at 4
+        // sources this loop is ~13% of the cost of the `interleave` calls it
+        // feeds (3 columns, 8192 rows), so even the 128-source case above
+        // leaves `interleave` the dominant term.
         let mut source_batches: Vec<usize> = Vec::new();
         let mut interleave_indices: Vec<(usize, usize)> =
             Vec::with_capacity(total_matched_rows);
         for (batch_idx, _, right) in matched_chunks {
             let source = match source_batches.iter().position(|b| b == batch_idx) {
-                Some(pos) => pos + 1,
+                Some(pos) => pos + source_offset,
                 None => {
                     source_batches.push(*batch_idx);
-                    source_batches.len()
+                    source_batches.len() - 1 + source_offset
                 }
             };
-            for i in 0..right.len() {
-                if right.is_null(i) {
-                    interleave_indices.push((0, 0));
-                } else {
-                    interleave_indices.push((source, right.value(i) as usize));
+            if right.null_count() == 0 {
+                // Hot path: no per-row null check, and `values()` avoids
+                // the bounds check `value(i)` would repeat.
+                interleave_indices
+                    .extend(right.values().iter().map(|&idx| (source, idx as usize)));
+            } else {
+                for i in 0..right.len() {
+                    if right.is_null(i) {
+                        interleave_indices.push((0, 0));
+                    } else {
+                        interleave_indices.push((source, right.value(i) as usize));
+                    }
                 }
             }
         }
@@ -1692,33 +1759,36 @@ impl MaterializingSortMergeJoinStream {
         let num_right_cols = self.buffered_schema.fields().len();
 
         // Read each source batch once (spilled batches require disk I/O).
-        let source_data_result: Result<Vec<RecordBatch>> = source_batches
+        let source_data: Vec<&RecordBatch> = source_batches
             .iter()
-            .map(|&idx| {
-                let bb = &self.buffered_data.batches[idx];
-                match &bb.batch {
-                    BufferedBatchState::InMemory(batch) => Ok(batch.clone()),
-                    BufferedBatchState::Spilled(_) => {
-                        internal_err!("Buffered batch should have been unspilled before fetching columns")
-                    }
-                }
+            .map(|&idx| match &self.buffered_data.batches[idx].batch {
+                BufferedBatchState::InMemory(batch) => Ok(batch),
+                BufferedBatchState::Spilled(_) => internal_err!(
+                    "Buffered batch should have been unspilled before fetching columns"
+                ),
             })
-            .collect();
+            .collect::<Result<_>>()?;
 
-        let source_data = source_data_result?;
+        // One single-row null array per column, built up front so the
+        // per-column `source_arrays` can borrow them.
+        let null_arrays: Vec<ArrayRef> = if needs_null_sentinel {
+            self.buffered_schema
+                .fields()
+                .iter()
+                .map(|f| new_null_array(f.data_type(), 1))
+                .collect()
+        } else {
+            vec![]
+        };
 
+        let mut source_arrays: Vec<&dyn Array> =
+            Vec::with_capacity(source_data.len() + source_offset);
         let mut right_columns = Vec::with_capacity(num_right_cols);
         for col_idx in 0..num_right_cols {
-            let dtype = self.buffered_schema.field(col_idx).data_type();
-            let null_array = new_null_array(dtype, 1);
-
-            let mut source_arrays: Vec<&dyn Array> =
-                Vec::with_capacity(source_batches.len() + 1);
-            source_arrays.push(null_array.as_ref());
+            source_arrays.clear();
+            source_arrays.extend(null_arrays.get(col_idx).map(|a| a.as_ref()));
+            source_arrays.extend(source_data.iter().map(|d| d.column(col_idx).as_ref()));
 
-            for data in &source_data {
-                source_arrays.push(data.column(col_idx).as_ref());
-            }
             right_columns.push(interleave(&source_arrays, &interleave_indices)?);
         }
 
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 4cf862b4ca8b0..21ee46b831a0d 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -6079,3 +6079,187 @@ async fn left_join_with_filter_preserves_streamed_order() -> Result<()> {
     );
     Ok(())
 }
+
+/// Mirror of [`left_join_with_filter_preserves_streamed_order`] for
+/// `RIGHT JOIN`, which advertises `maintains_input_order = [false, true]`
+/// and therefore streams its *right* input. The buffered (left) side carries
+/// the duplicate groups here, and the streamed key lands in the output after
+/// the buffered columns.
+#[tokio::test]
+async fn right_join_with_filter_preserves_streamed_order() -> Result<()> {
+    let num_keys = 7i32;
+
+    // Buffered (left) side: large groups for the first five keys, so the
+    // deferred-filter gate fires once per key; single rows for the last two,
+    // whose output only leaves through the final flush.
+    let mut l_a = vec![];
+    let mut l_b = vec![];
+    let mut l_c = vec![];
+    for k in 0..num_keys {
+        let dup = if k < 5 { 20 } else { 1 };
+        for j in 0..dup {
+            l_a.push(k * 100 + j);
+            l_b.push(k);
+            l_c.push(j);
+        }
+    }
+    let left = build_table_i32(("a1", &l_a), ("b1", &l_b), ("c1", &l_c));
+
+    // Streamed (right) side: one row per key, in key order.
+    let keys: Vec<i32> = (0..num_keys).collect();
+    let right = build_table_i32(("a2", &keys), ("b2", &keys), ("c2", &keys));
+
+    let left_schema = left.schema();
+    let right_schema = right.schema();
+    let left = TestMemoryExec::try_new_exec(&[vec![left]], left_schema, None)?;
+    let right = TestMemoryExec::try_new_exec(&[vec![right]], right_schema, None)?;
+
+    let on: JoinOn = vec![(
+        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
+    )];
+
+    // A filter that never passes, so every streamed row is emitted
+    // null-joined by the deferred-filtering pipeline.
+    let intermediate_schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
+    let filter = JoinFilter::new(
+        Arc::new(BinaryExpr::new(
+            Arc::new(Column::new("x", 0)),
+            Operator::Lt,
+            Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
+        )) as PhysicalExprRef,
+        vec![ColumnIndex {
+            index: 0,
+            side: JoinSide::Left,
+        }],
+        Arc::new(intermediate_schema),
+    );
+
+    let join = SortMergeJoinExec::try_new(
+        left,
+        right,
+        on,
+        Some(filter),
+        Right,
+        vec![SortOptions::default()],
+        NullEquality::NullEqualsNothing,
+    )?;
+
+    let task_ctx = Arc::new(
+        TaskContext::default()
+            .with_session_config(SessionConfig::default().with_batch_size(8)),
+    );
+    let batches = common::collect(join.execute(0, task_ctx)?).await?;
+
+    // Output layout for RIGHT JOIN is [left cols.., right cols..], so the
+    // streamed key `a2` sits at index 3.
+    let streamed_keys: Vec<i32> = batches
+        .iter()
+        .flat_map(|b| {
+            b.column(3)
+                .as_any()
+                .downcast_ref::<Int32Array>()
+                .unwrap()
+                .values()
+                .to_vec()
+        })
+        .collect();
+
+    assert_eq!(
+        streamed_keys, keys,
+        "RIGHT JOIN output must stay ordered on the streamed side"
+    );
+    Ok(())
+}
+
+/// Same shape as [`left_join_with_filter_preserves_streamed_order`], but with
+/// a filter that passes for *some* rows. The all-fail case only exercises the
+/// null-joined path; here matched rows survive the filter too, so the output
+/// mixes filter-passing and null-joined rows and must still be non-decreasing
+/// on the streamed key.
+#[tokio::test]
+async fn left_join_with_partial_filter_preserves_streamed_order() -> Result<()> {
+    let num_keys = 7i32;
+
+    let keys: Vec<i32> = (0..num_keys).collect();
+    let left = build_table_i32(("a1", &keys), ("b1", &keys), ("c1", &keys));
+
+    let mut r_a = vec![];
+    let mut r_b = vec![];
+    let mut r_c = vec![];
+    for k in 0..num_keys {
+        let dup = if k < 5 { 20 } else { 1 };
+        for j in 0..dup {
+            r_a.push(k * 100 + j);
+            r_b.push(k);
+            r_c.push(j);
+        }
+    }
+    let right = build_table_i32(("a2", &r_a), ("b2", &r_b), ("c2", &r_c));
+
+    let left_schema = left.schema();
+    let right_schema = right.schema();
+    let left = TestMemoryExec::try_new_exec(&[vec![left]], left_schema, None)?;
+    let right = TestMemoryExec::try_new_exec(&[vec![right]], right_schema, None)?;
+
+    let on: JoinOn = vec![(
+        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
+    )];
+
+    // `c2 < 3`: keys 0..5 keep three of their twenty buffered rows, keys 5
+    // and 6 keep their single row.
+    let intermediate_schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
+    let filter = JoinFilter::new(
+        Arc::new(BinaryExpr::new(
+            Arc::new(Column::new("x", 0)),
+            Operator::Lt,
+            Arc::new(Literal::new(ScalarValue::Int32(Some(3)))),
+        )) as PhysicalExprRef,
+        vec![ColumnIndex {
+            index: 2,
+            side: JoinSide::Right,
+        }],
+        Arc::new(intermediate_schema),
+    );
+
+    let join = SortMergeJoinExec::try_new(
+        left,
+        right,
+        on,
+        Some(filter),
+        Left,
+        vec![SortOptions::default()],
+        NullEquality::NullEqualsNothing,
+    )?;
+
+    let task_ctx = Arc::new(
+        TaskContext::default()
+            .with_session_config(SessionConfig::default().with_batch_size(8)),
+    );
+    let batches = common::collect(join.execute(0, task_ctx)?).await?;
+
+    let streamed_keys: Vec<i32> = batches
+        .iter()
+        .flat_map(|b| {
+            b.column(0)
+                .as_any()
+                .downcast_ref::<Int32Array>()
+                .unwrap()
+                .values()
+                .to_vec()
+        })
+        .collect();
+
+    assert!(
+        streamed_keys.windows(2).all(|w| w[0] <= w[1]),
+        "LEFT JOIN output must stay ordered on the streamed side, got {streamed_keys:?}"
+    );
+    // Every streamed key must still be represented exactly once per
+    // surviving match: 3 per key for keys 0..5, 1 each for keys 5 and 6.
+    let expected: Vec<i32> = (0..num_keys)
+        .flat_map(|k| std::iter::repeat_n(k, if k < 5 { 3 } else { 1 }))
+        .collect();
+    assert_eq!(streamed_keys, expected);
+    Ok(())
+}

From 25dae23f6d3c4509a9ae8646aecee7da8fdfcc41 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Sat, 22 Aug 2026 21:26:08 +0800
Subject: [PATCH 3/7] test: enhance deferred-filtered outer join tests to
 ensure streamed order preservation

---
 .../src/joins/sort_merge_join/tests.rs        | 316 ++++++------------
 1 file changed, 107 insertions(+), 209 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 21ee46b831a0d..c7b3c574e316e 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -5986,41 +5986,58 @@ async fn bitwise_spill_pending_stream() -> Result<()> {
     Ok(())
 }
 
-/// Regression test: deferred-filtered outer joins must not reorder their
-/// output.
-///
-/// `LEFT JOIN` advertises `maintains_input_order = [true, false]`, so the
-/// output must stay ordered on the streamed side. The final flush used to
-/// emit its batch directly instead of through the `output` coalescer, so any
-/// rows still buffered there from an earlier flush were emitted *after* it.
+/// Number of distinct join keys used by the streamed-order regression tests.
+const ORDER_KEYS: i32 = 7;
+
+/// Streamed side of the streamed-order tests: one row per key, ascending.
+fn order_unique_side(names: [&str; 3]) -> RecordBatch {
+    let keys: Vec<i32> = (0..ORDER_KEYS).collect();
+    build_table_i32((names[0], &keys), (names[1], &keys), (names[2], &keys))
+}
+
+/// Buffered side of the streamed-order tests.
 ///
-/// The shape below reproduces that: the first five keys each match a large
-/// buffered group, so the deferred-filter gate fires once per key and pushes
-/// a single-row batch into `output` (too small to complete a batch), while
-/// the last two keys match a single row each and so never trip the gate —
-/// leaving their rows for the final flush.
-#[tokio::test]
-async fn left_join_with_filter_preserves_streamed_order() -> Result<()> {
-    let num_keys = 7i32;
-
-    let keys: Vec<i32> = (0..num_keys).collect();
-    let left = build_table_i32(("a1", &keys), ("b1", &keys), ("c1", &keys));
-
-    let mut r_a = vec![];
-    let mut r_b = vec![];
-    let mut r_c = vec![];
-    for k in 0..num_keys {
-        let dup = if k < 5 { 20 } else { 1 };
-        for j in 0..dup {
-            r_a.push(k * 100 + j);
-            r_b.push(k);
-            r_c.push(j);
+/// Keys 0..5 carry 20 rows each — wide enough that the deferred-filter gate
+/// fires once per key and leaves a partial batch sitting in `output` — while
+/// keys 5 and 6 carry a single row each, so their output only ever leaves
+/// through the final flush. Mixing the two paths is what exposes reordering
+/// between them.
+fn order_skewed_side(names: [&str; 3]) -> RecordBatch {
+    let (mut a, mut b, mut c) = (vec![], vec![], vec![]);
+    for k in 0..ORDER_KEYS {
+        for j in 0..if k < 5 { 20 } else { 1 } {
+            a.push(k * 100 + j);
+            b.push(k);
+            c.push(j);
         }
     }
-    let right = build_table_i32(("a2", &r_a), ("b2", &r_b), ("c2", &r_c));
+    build_table_i32((names[0], &a), (names[1], &b), (names[2], &c))
+}
+
+/// Run a deferred-filtered outer join over the skew shape above and return
+/// the streamed key column of the output, concatenated across batches.
+///
+/// The filter is `<filter_column> < filter_lt` over the intermediate schema.
+async fn collect_streamed_keys(
+    join_type: JoinType,
+    filter_column: ColumnIndex,
+    filter_lt: i32,
+) -> Result<Vec<i32>> {
+    // RIGHT streams its *right* input (`maintains_input_order = [false, true]`),
+    // so the duplicate groups always belong on whichever side is buffered.
+    let (left, right) = if join_type == Right {
+        (
+            order_skewed_side(["a1", "b1", "c1"]),
+            order_unique_side(["a2", "b2", "c2"]),
+        )
+    } else {
+        (
+            order_unique_side(["a1", "b1", "c1"]),
+            order_skewed_side(["a2", "b2", "c2"]),
+        )
+    };
 
-    let left_schema = left.schema();
-    let right_schema = right.schema();
+    let (left_schema, right_schema) = (left.schema(), right.schema());
     let left = TestMemoryExec::try_new_exec(&[vec![left]], left_schema, None)?;
     let right = TestMemoryExec::try_new_exec(&[vec![right]], right_schema, None)?;
 
@@ -6029,20 +6046,14 @@ async fn left_join_with_filter_preserves_streamed_order() -> Result<()> {
         Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
     )];
 
-    // A filter that never passes, so every streamed row is emitted
-    // null-joined by the deferred-filtering pipeline.
-    let intermediate_schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
     let filter = JoinFilter::new(
         Arc::new(BinaryExpr::new(
             Arc::new(Column::new("x", 0)),
             Operator::Lt,
-            Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
+            Arc::new(Literal::new(ScalarValue::Int32(Some(filter_lt)))),
         )) as PhysicalExprRef,
-        vec![ColumnIndex {
-            index: 0,
-            side: JoinSide::Left,
-        }],
-        Arc::new(intermediate_schema),
+        vec![filter_column],
+        Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)])),
     );
 
     let join = SortMergeJoinExec::try_new(
@@ -6050,216 +6061,103 @@ async fn left_join_with_filter_preserves_streamed_order() -> Result<()> {
         right,
         on,
         Some(filter),
-        Left,
+        join_type,
         vec![SortOptions::default()],
         NullEquality::NullEqualsNothing,
     )?;
 
+    // A small batch size keeps the gate firing often enough to interleave the
+    // two output paths.
     let task_ctx = Arc::new(
         TaskContext::default()
             .with_session_config(SessionConfig::default().with_batch_size(8)),
     );
     let batches = common::collect(join.execute(0, task_ctx)?).await?;
 
-    let streamed_keys: Vec<i32> = batches
+    // Output is always [left cols.., right cols..], so the streamed key is
+    // `a2` at index 3 for RIGHT and `a1` at index 0 otherwise.
+    let key_col = if join_type == Right { 3 } else { 0 };
+    Ok(batches
         .iter()
         .flat_map(|b| {
-            b.column(0)
+            b.column(key_col)
                 .as_any()
                 .downcast_ref::<Int32Array>()
                 .unwrap()
                 .values()
                 .to_vec()
         })
-        .collect();
+        .collect())
+}
+
+/// `a1 < 0`, which never passes — so every streamed row is emitted
+/// null-joined by the deferred-filtering pipeline.
+fn never_passing_filter() -> (ColumnIndex, i32) {
+    (
+        ColumnIndex {
+            index: 0,
+            side: JoinSide::Left,
+        },
+        0,
+    )
+}
+
+/// Regression test: deferred-filtered outer joins must not reorder their
+/// output.
+///
+/// `LEFT JOIN` advertises `maintains_input_order = [true, false]`, so the
+/// output must stay ordered on the streamed side. The final flush used to
+/// emit its batch directly instead of through the `output` coalescer, so any
+/// rows still buffered there from an earlier flush were emitted *after* it.
+#[tokio::test]
+async fn left_join_with_filter_preserves_streamed_order() -> Result<()> {
+    let (filter_column, filter_lt) = never_passing_filter();
+    let streamed_keys = collect_streamed_keys(Left, filter_column, filter_lt).await?;
 
     assert_eq!(
-        streamed_keys, keys,
+        streamed_keys,
+        (0..ORDER_KEYS).collect::<Vec<_>>(),
         "LEFT JOIN output must stay ordered on the streamed side"
     );
     Ok(())
 }
 
 /// Mirror of [`left_join_with_filter_preserves_streamed_order`] for
-/// `RIGHT JOIN`, which advertises `maintains_input_order = [false, true]`
-/// and therefore streams its *right* input. The buffered (left) side carries
-/// the duplicate groups here, and the streamed key lands in the output after
-/// the buffered columns.
+/// `RIGHT JOIN`, which advertises `maintains_input_order = [false, true]` and
+/// therefore streams its *right* input.
 #[tokio::test]
 async fn right_join_with_filter_preserves_streamed_order() -> Result<()> {
-    let num_keys = 7i32;
-
-    // Buffered (left) side: large groups for the first five keys, so the
-    // deferred-filter gate fires once per key; single rows for the last two,
-    // whose output only leaves through the final flush.
-    let mut l_a = vec![];
-    let mut l_b = vec![];
-    let mut l_c = vec![];
-    for k in 0..num_keys {
-        let dup = if k < 5 { 20 } else { 1 };
-        for j in 0..dup {
-            l_a.push(k * 100 + j);
-            l_b.push(k);
-            l_c.push(j);
-        }
-    }
-    let left = build_table_i32(("a1", &l_a), ("b1", &l_b), ("c1", &l_c));
-
-    // Streamed (right) side: one row per key, in key order.
-    let keys: Vec<i32> = (0..num_keys).collect();
-    let right = build_table_i32(("a2", &keys), ("b2", &keys), ("c2", &keys));
-
-    let left_schema = left.schema();
-    let right_schema = right.schema();
-    let left = TestMemoryExec::try_new_exec(&[vec![left]], left_schema, None)?;
-    let right = TestMemoryExec::try_new_exec(&[vec![right]], right_schema, None)?;
-
-    let on: JoinOn = vec![(
-        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
-        Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
-    )];
-
-    // A filter that never passes, so every streamed row is emitted
-    // null-joined by the deferred-filtering pipeline.
-    let intermediate_schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
-    let filter = JoinFilter::new(
-        Arc::new(BinaryExpr::new(
-            Arc::new(Column::new("x", 0)),
-            Operator::Lt,
-            Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
-        )) as PhysicalExprRef,
-        vec![ColumnIndex {
-            index: 0,
-            side: JoinSide::Left,
-        }],
-        Arc::new(intermediate_schema),
-    );
-
-    let join = SortMergeJoinExec::try_new(
-        left,
-        right,
-        on,
-        Some(filter),
-        Right,
-        vec![SortOptions::default()],
-        NullEquality::NullEqualsNothing,
-    )?;
-
-    let task_ctx = Arc::new(
-        TaskContext::default()
-            .with_session_config(SessionConfig::default().with_batch_size(8)),
-    );
-    let batches = common::collect(join.execute(0, task_ctx)?).await?;
-
-    // Output layout for RIGHT JOIN is [left cols.., right cols..], so the
-    // streamed key `a2` sits at index 3.
-    let streamed_keys: Vec<i32> = batches
-        .iter()
-        .flat_map(|b| {
-            b.column(3)
-                .as_any()
-                .downcast_ref::<Int32Array>()
-                .unwrap()
-                .values()
-                .to_vec()
-        })
-        .collect();
+    let (filter_column, filter_lt) = never_passing_filter();
+    let streamed_keys = collect_streamed_keys(Right, filter_column, filter_lt).await?;
 
     assert_eq!(
-        streamed_keys, keys,
+        streamed_keys,
+        (0..ORDER_KEYS).collect::<Vec<_>>(),
         "RIGHT JOIN output must stay ordered on the streamed side"
     );
     Ok(())
 }
 
-/// Same shape as [`left_join_with_filter_preserves_streamed_order`], but with
-/// a filter that passes for *some* rows. The all-fail case only exercises the
-/// null-joined path; here matched rows survive the filter too, so the output
-/// mixes filter-passing and null-joined rows and must still be non-decreasing
-/// on the streamed key.
+/// Same shape, but with a filter that passes for *some* rows. The all-fail
+/// cases above only exercise the null-joined path; here matched rows survive
+/// the filter too, so the output mixes filter-passing and null-joined rows.
 #[tokio::test]
 async fn left_join_with_partial_filter_preserves_streamed_order() -> Result<()> {
-    let num_keys = 7i32;
-
-    let keys: Vec<i32> = (0..num_keys).collect();
-    let left = build_table_i32(("a1", &keys), ("b1", &keys), ("c1", &keys));
-
-    let mut r_a = vec![];
-    let mut r_b = vec![];
-    let mut r_c = vec![];
-    for k in 0..num_keys {
-        let dup = if k < 5 { 20 } else { 1 };
-        for j in 0..dup {
-            r_a.push(k * 100 + j);
-            r_b.push(k);
-            r_c.push(j);
-        }
-    }
-    let right = build_table_i32(("a2", &r_a), ("b2", &r_b), ("c2", &r_c));
-
-    let left_schema = left.schema();
-    let right_schema = right.schema();
-    let left = TestMemoryExec::try_new_exec(&[vec![left]], left_schema, None)?;
-    let right = TestMemoryExec::try_new_exec(&[vec![right]], right_schema, None)?;
-
-    let on: JoinOn = vec![(
-        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
-        Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
-    )];
-
     // `c2 < 3`: keys 0..5 keep three of their twenty buffered rows, keys 5
     // and 6 keep their single row.
-    let intermediate_schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
-    let filter = JoinFilter::new(
-        Arc::new(BinaryExpr::new(
-            Arc::new(Column::new("x", 0)),
-            Operator::Lt,
-            Arc::new(Literal::new(ScalarValue::Int32(Some(3)))),
-        )) as PhysicalExprRef,
-        vec![ColumnIndex {
-            index: 2,
-            side: JoinSide::Right,
-        }],
-        Arc::new(intermediate_schema),
-    );
-
-    let join = SortMergeJoinExec::try_new(
-        left,
-        right,
-        on,
-        Some(filter),
-        Left,
-        vec![SortOptions::default()],
-        NullEquality::NullEqualsNothing,
-    )?;
-
-    let task_ctx = Arc::new(
-        TaskContext::default()
-            .with_session_config(SessionConfig::default().with_batch_size(8)),
-    );
-    let batches = common::collect(join.execute(0, task_ctx)?).await?;
-
-    let streamed_keys: Vec<i32> = batches
-        .iter()
-        .flat_map(|b| {
-            b.column(0)
-                .as_any()
-                .downcast_ref::<Int32Array>()
-                .unwrap()
-                .values()
-                .to_vec()
-        })
-        .collect();
+    let filter_column = ColumnIndex {
+        index: 2,
+        side: JoinSide::Right,
+    };
+    let streamed_keys = collect_streamed_keys(Left, filter_column, 3).await?;
 
-    assert!(
-        streamed_keys.windows(2).all(|w| w[0] <= w[1]),
-        "LEFT JOIN output must stay ordered on the streamed side, got {streamed_keys:?}"
-    );
-    // Every streamed key must still be represented exactly once per
-    // surviving match: 3 per key for keys 0..5, 1 each for keys 5 and 6.
-    let expected: Vec<i32> = (0..num_keys)
+    let expected: Vec<i32> = (0..ORDER_KEYS)
         .flat_map(|k| std::iter::repeat_n(k, if k < 5 { 3 } else { 1 }))
         .collect();
-    assert_eq!(streamed_keys, expected);
+    assert_eq!(
+        streamed_keys, expected,
+        "LEFT JOIN output must stay ordered on the streamed side, \
+         with every surviving match present exactly once"
+    );
     Ok(())
 }

From af8078c8e50d4b4724d3808cfe0b516e96739157 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Mon, 24 Aug 2026 20:10:45 +0800
Subject: [PATCH 4/7] cargo fmt

---
 datafusion/physical-plan/src/joins/sort_merge_join/tests.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 5673e3c3e18f9..b4415546f42e3 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -6356,4 +6356,4 @@ async fn left_join_with_partial_filter_preserves_streamed_order() -> Result<()>
          with every surviving match present exactly once"
     );
     Ok(())
-}
\ No newline at end of file
+}

From cf3e63539979f6469af30a617a27967ca0e3b6a8 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Tue, 25 Aug 2026 21:11:08 +0800
Subject: [PATCH 5/7] feat: optimize source index mapping in SortMergeJoin and
 add test for group spanning batches

---
 .../sort_merge_join/materializing_stream.rs   | 104 +++++++++---------
 .../src/joins/sort_merge_join/tests.rs        |  84 ++++++++++++++
 2 files changed, 134 insertions(+), 54 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 c3e04f09bb203..9bb54b1c942ea 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
@@ -1677,69 +1677,65 @@ impl MaterializingSortMergeJoinStream {
             .any(|(_, _, right)| right.null_count() > 0);
         let source_offset = usize::from(needs_null_sentinel);
 
-        // A group spans only a handful of buffered batches, so a linear
-        // scan beats hashing here. Measured over 8192 rows in 2048 chunks,
-        // against a `HashMap<usize, usize>` built in one pass and read back
-        // in a second (what this used to do):
+        // Map each distinct `buffered_batch_idx` to a contiguous source
+        // index for `interleave`. The keys are positions in
+        // `self.buffered_data.batches`, and `BufferedData::scanning_advance`
+        // walks that deque in order, so they form a dense run: a
+        // direct-addressed table over `min..=max` resolves every chunk in
+        // O(1), with no hashing and no key comparison.
         //
-        //   distinct sources |  hashmap  |  linear scan
-        //   -----------------+-----------+-------------
-        //                  4 |  21.5 us  |   5.0 us
-        //                 16 |  22.0 us  |   9.4 us
-        //                 32 |  22.4 us  |  13.6 us
-        //                 64 |  22.9 us  |  23.5 us
-        //                128 |  24.1 us  |  44.8 us
+        // A linear `position()` scan over `source_batches` is not enough
+        // here, even though a freeze holds at most `batch_size` pairs.
+        // `pair_streamed_row_with_group` restarts the buffered scan at batch
+        // 0 for *every* streamed row of the key group (`scanning_reset`), so
+        // the chunk sequence cycles `0,1,..,S-1,0,1,..` and the chunk count
+        // is not bounded by the distinct-source count `S`. The scan is then
+        // O(chunks * S), and nothing bounds `S`: `SortMergeJoinExec` accepts
+        // arbitrary `ExecutionPlan` children, so one emitting tiny batches
+        // pushes `S` towards `batch_size`.
         //
-        // `std::collections::HashMap` hashes with SipHash-1-3, so a single
-        // `usize` lookup costs several ns of serial latency before the probe
-        // begins, while a scan over a handful of `usize` is one L1-resident
-        // cache line with a perfectly predicted trip count. The map is also
-        // purely additive state: `source_batches` has to be built regardless
-        // (`source_data` is gathered from it), so hashing means maintaining
-        // two containers holding the same keys.
+        // Measured over 8192 rows in 2048 chunks, against a
+        // `HashMap<usize, usize>` built in one pass and read back in a
+        // second:
         //
-        // The crossover is ~32 distinct sources. That bound follows from how
-        // pairs accumulate, not from any assumption about key skew:
+        //   distinct sources |  hashmap  |  linear scan  |  direct table
+        //   -----------------+-----------+---------------+--------------
+        //                  4 |   19.7 us |       4.5 us  |      4.8 us
+        //                 32 |   20.7 us |      13.0 us  |      5.0 us
+        //                128 |   23.5 us |      42.7 us  |      5.1 us
+        //               1024 |   48.1 us |     281.7 us  |      5.8 us
+        //               8192 |  293.3 us |    8347.6 us  |     16.9 us
         //
-        //   1. `pair_streamed_row_with_group` appends exactly one pair per
-        //      buffered row and re-checks `num_unfrozen_pairs() < batch_size`
-        //      before each append, so at most `batch_size` pairs accumulate
-        //      between two `freeze_streamed()` calls.
-        //   2. `BufferedData::scanning_advance` walks the group's rows in
-        //      order, so those pairs cover a *contiguous run* of buffered
-        //      rows.
-        //   3. So the distinct `buffered_batch_idx` values seen here are the
-        //      batches spanned by at most `batch_size` consecutive buffered
-        //      rows: `len(source_batches) <= batch_size / R + 1`, where `R`
-        //      is the smallest buffered batch in that run.
+        // The last row is the degenerate shape a one-row-per-batch child
+        // produces: 8192 chunks of a single row each, all from distinct
+        // buffered batches. 8.3 ms of index construction, in one freeze.
         //
-        // The assumption is therefore not "key groups are narrow" — a group
-        // of any width still only contributes `batch_size` rows per freeze —
-        // but "buffered batches are not tiny relative to `batch_size`".
-        // Exceeding 32 sources needs `R < batch_size / 31`, i.e. under ~264
-        // rows per batch at the default `batch_size` of 8192. The buffered
-        // side of a merge join is sorted input, and every operator that
-        // normally feeds it emits ~`batch_size` batches: `SortExec` chunks
-        // its output with `sort_batch_chunked(.., batch_size)`, and
-        // `FilterExec` and `RepartitionExec` each embed a
-        // `LimitedBatchCoalescer` targeting `batch_size`.
-        //
-        // If something does feed tiny batches, this degrades gradually rather
-        // than falling off a cliff, and never affects correctness: at 4
-        // sources this loop is ~13% of the cost of the `interleave` calls it
-        // feeds (3 columns, 8192 rows), so even the 128-source case above
-        // leaves `interleave` the dominant term.
+        // The table ties the scan where the scan is at its best (a handful
+        // of sources): both stay in L1 and neither hashes, whereas
+        // `std::collections::HashMap` uses SipHash-1-3 and pays several ns
+        // of serial latency before each probe begins. Unlike the scan, it
+        // stays flat. `source_batches` has to be built regardless
+        // (`source_data` is gathered from it), so the table is the only
+        // added state, sized by the span of buffered batches this freeze
+        // touches rather than by the whole buffer.
+        let (min_batch_idx, max_batch_idx) = matched_chunks
+            .iter()
+            .fold((usize::MAX, 0usize), |(lo, hi), (batch_idx, _, _)| {
+                (lo.min(*batch_idx), hi.max(*batch_idx))
+            });
+        // Sentinel for "no source index assigned to this buffered batch yet".
+        const UNSEEN: usize = usize::MAX;
+        let mut source_of_batch = vec![UNSEEN; max_batch_idx - min_batch_idx + 1];
         let mut source_batches: Vec<usize> = Vec::new();
         let mut interleave_indices: Vec<(usize, usize)> =
             Vec::with_capacity(total_matched_rows);
         for (batch_idx, _, right) in matched_chunks {
-            let source = match source_batches.iter().position(|b| b == batch_idx) {
-                Some(pos) => pos + source_offset,
-                None => {
-                    source_batches.push(*batch_idx);
-                    source_batches.len() - 1 + source_offset
-                }
-            };
+            let slot = &mut source_of_batch[batch_idx - min_batch_idx];
+            if *slot == UNSEEN {
+                *slot = source_batches.len();
+                source_batches.push(*batch_idx);
+            }
+            let source = *slot + source_offset;
             if right.null_count() == 0 {
                 // Hot path: no per-row null check, and `values()` avoids
                 // the bounds check `value(i)` would repeat.
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 b4415546f42e3..6e9ab5e8ae9d6 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -4137,6 +4137,90 @@ async fn join_filtered_with_multiple_buffered_batches() -> Result<()> {
     Ok(())
 }
 
+/// A single key group spanning many buffered batches, re-scanned once per
+/// streamed row.
+///
+/// `pair_streamed_row_with_group` walks the group from buffered batch 0 for
+/// *every* streamed row (`scanning_reset`), and freezes whenever `batch_size`
+/// pairs have accumulated -- which happens mid-scan when `batch_size` is not a
+/// multiple of the group size. So one `freeze_streamed()` can see chunks whose
+/// `buffered_batch_idx` wraps (`.. 4, 5, 0, 1 ..`) or never reaches 0 at all,
+/// rather than a single ascending run. `materialize_right_columns` maps those
+/// indices to `interleave` source slots, so it must not assume either.
+///
+/// 6 one-row buffered batches x 2 streamed rows at `batch_size` 5 produces
+/// freezes covering batches `[0,1,2,3,4]`, `[5,0,1,2,3]` (wrapped) and
+/// `[4,5]` (no zero).
+#[tokio::test]
+async fn join_with_group_spanning_batches_rescanned_per_streamed_row() -> Result<()> {
+    let left_schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int32, false),
+        Field::new("val_l", DataType::Int32, false),
+    ]));
+    let right_schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int32, false),
+        Field::new("val_r", DataType::Int32, false),
+    ]));
+
+    // Two streamed rows sharing one key, so the buffered group is scanned twice.
+    let left = build_table_from_batches(vec![RecordBatch::try_new(
+        Arc::clone(&left_schema),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 1])),
+            Arc::new(Int32Array::from(vec![10, 20])),
+        ],
+    )?]);
+
+    // One row per batch, all the same key: the group spans all 6 batches.
+    let right_batches: Vec<RecordBatch> = (1..=6)
+        .map(|i| {
+            RecordBatch::try_new(
+                Arc::clone(&right_schema),
+                vec![
+                    Arc::new(Int32Array::from(vec![1])),
+                    Arc::new(Int32Array::from(vec![i * 100])),
+                ],
+            )
+            .unwrap()
+        })
+        .collect();
+    let right = build_table_from_batches(right_batches);
+
+    let on: JoinOn = vec![(
+        Arc::new(Column::new_with_schema("key", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("key", &right.schema())?) as _,
+    )];
+
+    // 5 does not divide the 6-row group, so freezes land mid-scan.
+    let task_ctx = Arc::new(
+        TaskContext::default()
+            .with_session_config(SessionConfig::new().with_batch_size(5)),
+    );
+    let join = join(left, right, on, Inner)?;
+    let batches = common::collect(join.execute(0, task_ctx)?).await?;
+
+    assert_snapshot!(batches_to_sort_string(&batches), @r"
+    +-----+-------+-----+-------+
+    | key | val_l | key | val_r |
+    +-----+-------+-----+-------+
+    | 1   | 10    | 1   | 100   |
+    | 1   | 10    | 1   | 200   |
+    | 1   | 10    | 1   | 300   |
+    | 1   | 10    | 1   | 400   |
+    | 1   | 10    | 1   | 500   |
+    | 1   | 10    | 1   | 600   |
+    | 1   | 20    | 1   | 100   |
+    | 1   | 20    | 1   | 200   |
+    | 1   | 20    | 1   | 300   |
+    | 1   | 20    | 1   | 400   |
+    | 1   | 20    | 1   | 500   |
+    | 1   | 20    | 1   | 600   |
+    +-----+-------+-----+-------+
+    ");
+
+    Ok(())
+}
+
 /// Returns the column names on the schema
 fn columns(schema: &Schema) -> Vec<String> {
     schema.fields().iter().map(|f| f.name().clone()).collect()

From 1e8ba6d848fdea9a6a3829fbaffdce0654a8384e Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Tue, 25 Aug 2026 21:54:15 +0800
Subject: [PATCH 6/7] fix: improve comments for clarity in
 MaterializingSortMergeJoinStream

---
 .../sort_merge_join/materializing_stream.rs   | 31 ++++++++++++++-----
 1 file changed, 24 insertions(+), 7 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 9bb54b1c942ea..43306248d039d 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
@@ -1678,11 +1678,22 @@ impl MaterializingSortMergeJoinStream {
         let source_offset = usize::from(needs_null_sentinel);
 
         // Map each distinct `buffered_batch_idx` to a contiguous source
-        // index for `interleave`. The keys are positions in
-        // `self.buffered_data.batches`, and `BufferedData::scanning_advance`
-        // walks that deque in order, so they form a dense run: a
-        // direct-addressed table over `min..=max` resolves every chunk in
-        // O(1), with no hashing and no key comparison.
+        // index for `interleave`. The keys are not opaque: they are
+        // positions in `self.buffered_data.batches`, so the key space is
+        // dense and bounded by the deque length. A direct-addressed table
+        // over `min..=max` resolves every chunk in O(1), with no hashing and
+        // 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
+        // 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
+        // by the sources present. That costs O(group) for O(batch_size) of
+        // work -- but only once per pass, against the O(group) of useful
+        // work the rest of the pass does, so it stays O(1) amortized per
+        // pair. Measured over a 524288-batch group at `batch_size` 8192,
+        // a full pass costs 1.17 ms here against 11.25 ms for the hashmap.
         //
         // A linear `position()` scan over `source_batches` is not enough
         // here, even though a freeze holds at most `batch_size` pairs.
@@ -1716,13 +1727,19 @@ impl MaterializingSortMergeJoinStream {
         // of serial latency before each probe begins. Unlike the scan, it
         // stays flat. `source_batches` has to be built regardless
         // (`source_data` is gathered from it), so the table is the only
-        // added state, sized by the span of buffered batches this freeze
-        // touches rather than by the whole buffer.
+        // added state, and it is transient: sized to the span this freeze
+        // touches rather than held across freezes.
         let (min_batch_idx, max_batch_idx) = matched_chunks
             .iter()
             .fold((usize::MAX, 0usize), |(lo, hi), (batch_idx, _, _)| {
                 (lo.min(*batch_idx), hi.max(*batch_idx))
             });
+        // Every key indexes the live buffered deque -- this is what keeps
+        // the key space dense, and what makes `source_data` below safe.
+        debug_assert!(
+            max_batch_idx < self.buffered_data.batches.len(),
+            "buffered batch index {max_batch_idx} outside the buffered deque"
+        );
         // Sentinel for "no source index assigned to this buffered batch yet".
         const UNSEEN: usize = usize::MAX;
         let mut source_of_batch = vec![UNSEEN; max_batch_idx - min_batch_idx + 1];

From 2e303b684e1c70c1c15b357a7aa27b1807ee7d10 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Wed, 26 Aug 2026 22:46:53 +0800
Subject: [PATCH 7/7] feat: add benchmark for multi-source interleave with null
 buffered index in SortMergeJoin

---
 .../physical-plan/benches/sort_merge_join.rs  | 50 +++++++++++-
 .../src/joins/sort_merge_join/tests.rs        | 79 +++++++++++++++++++
 2 files changed, 128 insertions(+), 1 deletion(-)

diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs
index 26522136c2e30..d5f87a04f0a28 100644
--- a/datafusion/physical-plan/benches/sort_merge_join.rs
+++ b/datafusion/physical-plan/benches/sort_merge_join.rs
@@ -44,6 +44,21 @@ fn build_sorted_batches(
     num_rows: usize,
     key_mod: usize,
     schema: &SchemaRef,
+) -> Vec<RecordBatch> {
+    build_sorted_batches_with_size(num_rows, key_mod, 8192, schema)
+}
+
+/// Like [`build_sorted_batches`], but with an explicit output batch size.
+///
+/// `SortMergeJoinExec` takes arbitrary children, so the buffered side is not
+/// guaranteed to arrive in `batch_size`-sized batches. Small `batch_size`
+/// values make a single key group span many buffered batches, which is what
+/// drives `materialize_right_columns` onto its multi-source `interleave` path.
+fn build_sorted_batches_with_size(
+    num_rows: usize,
+    key_mod: usize,
+    batch_size: usize,
+    schema: &SchemaRef,
 ) -> Vec<RecordBatch> {
     let mut rows: Vec<(i64, i64)> = (0..num_rows)
         .map(|i| ((i % key_mod) as i64, i as i64))
@@ -64,7 +79,6 @@ fn build_sorted_batches(
     )
     .unwrap();
 
-    let batch_size = 8192;
     let mut batches = Vec::new();
     let mut offset = 0;
     while offset < batch.num_rows() {
@@ -197,6 +211,40 @@ fn bench_smj(c: &mut Criterion) {
         });
     }
 
+    // Multi-source interleave path — one buffered key group spanning many
+    // small buffered batches.
+    //
+    // Every other case here keeps a key group inside a single buffered batch,
+    // so `materialize_right_columns` takes its single-source `take` fast path
+    // and never reaches `interleave`. Shrinking the buffered batch size makes
+    // a group span `group_rows / rows_per_batch` batches, which is what the
+    // source-index mapping is actually paid for. Four streamed rows share each
+    // key, so the buffered scan is re-walked per streamed row and freezes wrap
+    // mid-group.
+    {
+        let keys = 8;
+        let group_rows = 8192;
+        let left_batches = build_sorted_batches(keys * 4, keys, &s);
+        for rows_per_batch in [512, 64, 8] {
+            let right_batches = build_sorted_batches_with_size(
+                keys * group_rows,
+                keys,
+                rows_per_batch,
+                &s,
+            );
+            group.bench_function(
+                BenchmarkId::new("inner_group_spans_buffered_batches", rows_per_batch),
+                |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)
+                    })
+                },
+            );
+        }
+    }
+
     group.finish();
 }
 
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 6e9ab5e8ae9d6..f0c18fd7cf084 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -4221,6 +4221,85 @@ async fn join_with_group_spanning_batches_rescanned_per_streamed_row() -> Result
     Ok(())
 }
 
+/// A wrapped multi-source freeze that also carries a null buffered index.
+///
+/// `materialize_right_columns` has two independent offsets in play on the
+/// interleave path: `batch_idx - min_batch_idx` addresses the source table,
+/// and `+ source_offset` shifts past the null sentinel that occupies
+/// `interleave` slot 0. Only their combination is interesting, and the two
+/// halves are awkward to get into the same freeze: `freeze_dequeuing_buffered`
+/// freezes before popping consumed batches, so a null-joined streamed row
+/// normally lands in its own single-source freeze.
+///
+/// The one shape that combines them puts the unmatched streamed row *before*
+/// a key group spanning several batches, with two streamed rows matching that
+/// group so the scan wraps:
+///
+///   chunk sequence [0, 1, 2, 0, 1, 2], chunk 0 carrying the null
+///
+/// Streamed key 5 finds no buffered match, so `null_join_streamed_row` appends
+/// a null pair at scan position 0; the two streamed 10s then each re-walk
+/// batches 0..2 (`scanning_reset`), wrapping inside the same freeze.
+#[tokio::test]
+async fn join_wrapped_multi_source_freeze_with_null_buffered_index() -> Result<()> {
+    let left_schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int32, false),
+        Field::new("val_l", DataType::Int32, false),
+    ]));
+    let right_schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int32, false),
+        Field::new("val_r", DataType::Int32, false),
+    ]));
+
+    // Key 5 has no buffered match; the two 10s share one group.
+    let left = build_table_from_batches(vec![RecordBatch::try_new(
+        Arc::clone(&left_schema),
+        vec![
+            Arc::new(Int32Array::from(vec![5, 10, 10])),
+            Arc::new(Int32Array::from(vec![50, 101, 102])),
+        ],
+    )?]);
+
+    // One row per batch, all key 10: the group spans all three batches.
+    let right_batches: Vec<RecordBatch> = [1000, 2000, 3000]
+        .into_iter()
+        .map(|v| {
+            RecordBatch::try_new(
+                Arc::clone(&right_schema),
+                vec![
+                    Arc::new(Int32Array::from(vec![10])),
+                    Arc::new(Int32Array::from(vec![v])),
+                ],
+            )
+            .unwrap()
+        })
+        .collect();
+    let right = build_table_from_batches(right_batches);
+
+    let on: JoinOn = vec![(
+        Arc::new(Column::new_with_schema("key", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("key", &right.schema())?) as _,
+    )];
+
+    let (_, batches) = join_collect(left, right, on, Left).await?;
+
+    assert_snapshot!(batches_to_sort_string(&batches), @r"
+    +-----+-------+-----+-------+
+    | key | val_l | key | val_r |
+    +-----+-------+-----+-------+
+    | 10  | 101   | 10  | 1000  |
+    | 10  | 101   | 10  | 2000  |
+    | 10  | 101   | 10  | 3000  |
+    | 10  | 102   | 10  | 1000  |
+    | 10  | 102   | 10  | 2000  |
+    | 10  | 102   | 10  | 3000  |
+    | 5   | 50    |     |       |
+    +-----+-------+-----+-------+
+    ");
+
+    Ok(())
+}
+
 /// Returns the column names on the schema
 fn columns(schema: &Schema) -> Vec<String> {
     schema.fields().iter().map(|f| f.name().clone()).collect()