From b3d430c4ad8a3987f41fe3e5296816ec2cc15984 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sat, 22 Aug 2026 15:59:02 +0800 Subject: [PATCH 1/3] 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]) -> Self { - let join_arrays = join_arrays(&batch, on_column); - StreamedBatch { + fn try_new(batch: RecordBatch, on_column: &[Arc]) -> Result { + 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, on_column: &[PhysicalExprRef], - ) -> Self { - let join_arrays = join_arrays(&batch, on_column); + ) -> Result { + let join_arrays = join_arrays(&batch, on_column)?; // Estimation is calculated as // inner batch size @@ -238,7 +238,7 @@ impl BufferedBatch { + size_of::(); 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, + ) { + 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> { + /// 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 = HashMap::new(); + // A group spans only a handful of buffered batches, so a linear + // scan beats hashing here. let mut source_batches: Vec = 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 { +/// 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> { + 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 = (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 = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .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 Date: Sat, 22 Aug 2026 21:20:15 +0800 Subject: [PATCH 2/3] 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> { @@ -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` 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 = 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> = 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::>()?; - 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 = 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 = (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 = batches + .iter() + .flat_map(|b| { + b.column(3) + .as_any() + .downcast_ref::() + .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 = (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 = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .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 = (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 Date: Sat, 22 Aug 2026 21:26:08 +0800 Subject: [PATCH 3/3] 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 = (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 = (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_lt` over the intermediate schema. +async fn collect_streamed_keys( + join_type: JoinType, + filter_column: ColumnIndex, + filter_lt: i32, +) -> Result> { + // 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 = 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::() .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::>(), "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 = (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 = batches - .iter() - .flat_map(|b| { - b.column(3) - .as_any() - .downcast_ref::() - .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::>(), "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 = (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 = batches - .iter() - .flat_map(|b| { - b.column(0) - .as_any() - .downcast_ref::() - .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 = (0..num_keys) + let expected: Vec = (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(()) }