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 { + 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 { 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/materializing_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs index 3baa0c4a3e792..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 @@ -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; } @@ -1640,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> { @@ -1664,26 +1666,105 @@ 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(); - 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 + // 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); + + // Map each distinct `buffered_batch_idx` to a contiguous source + // 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. + // `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`. + // + // Measured over 8192 rows in 2048 chunks, against a + // `HashMap` built in one pass and read back in a + // second: + // + // 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 + // + // 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 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, 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]; + 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 = batch_idx_to_source[batch_idx]; - 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)); + 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. + 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)); + } } } } @@ -1691,33 +1772,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)?); } @@ -1987,14 +2071,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 7a0c9e8562a98..f0c18fd7cf084 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,169 @@ 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 = (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(()) +} + +/// 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 = [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 { schema.fields().iter().map(|f| f.name().clone()).collect() @@ -6181,3 +6344,179 @@ async fn an_empty_projection_keeps_the_rows() -> Result<()> { Ok(()) } + +/// 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. +/// +/// 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); + } + } + 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, 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)?; + + let on: JoinOn = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(filter_lt)))), + )) as PhysicalExprRef, + vec![filter_column], + Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)])), + ); + + let join = SortMergeJoinExec::try_new( + left, + right, + on, + Some(filter), + 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?; + + // 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(key_col) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .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, + (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. +#[tokio::test] +async fn right_join_with_filter_preserves_streamed_order() -> Result<()> { + let (filter_column, filter_lt) = never_passing_filter(); + let streamed_keys = collect_streamed_keys(Right, filter_column, filter_lt).await?; + + assert_eq!( + streamed_keys, + (0..ORDER_KEYS).collect::>(), + "RIGHT JOIN output must stay ordered on the streamed side" + ); + Ok(()) +} + +/// 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<()> { + // `c2 < 3`: keys 0..5 keep three of their twenty buffered rows, keys 5 + // and 6 keep their single row. + let filter_column = ColumnIndex { + index: 2, + side: JoinSide::Right, + }; + let streamed_keys = collect_streamed_keys(Left, filter_column, 3).await?; + + 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, + "LEFT JOIN output must stay ordered on the streamed side, \ + with every surviving match present exactly once" + ); + Ok(()) +}