diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 296136fd7b3d1..7a30a56d590f7 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -662,6 +662,124 @@ async fn test_right_mark_join_1k_binary_filtered() { } } +// In the tests above, most key groups hold one row: the i32 tests draw the join keys +// `(a, b)` from 10,000 values for 1000 rows, and the binary ones from random bytes. +// The tests below draw them from 25 values, so each streamed row matches about 40 +// buffered rows. Sort-merge join pairs a streamed row with its matches in runs, which +// end where a buffered batch ends or where the output batch fills. + +#[tokio::test] +async fn test_inner_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Inner, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_inner_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Inner, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_left_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Left, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_left_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Left, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_right_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Right, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_right_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Right, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_full_join_1k_low_cardinality() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Full, + None, + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + +#[tokio::test] +async fn test_full_join_1k_low_cardinality_filtered() { + for (left_extra, right_extra) in [(true, true), (false, true), (true, false)] { + JoinFuzzTestCase::new( + make_staggered_batches_i32_with_key_range(1000, 5, left_extra), + make_staggered_batches_i32_with_key_range(1000, 5, right_extra), + JoinType::Full, + Some(Box::new(col_lt_col_filter)), + ) + .run_test(&[HjSmj, NljHj], false) + .await + } +} + type JoinFilterBuilder = Box, Arc) -> JoinFilter>; struct JoinFuzzTestCase { @@ -1392,12 +1510,25 @@ async fn test_nested_loop_join_spill_fuzz() { /// two sorted int32 columns 'a', 'b' ranged from 0..99 as join columns /// two random int32 columns 'x', 'y' as other columns fn make_staggered_batches_i32(len: usize, with_extra_column: bool) -> Vec { + make_staggered_batches_i32_with_key_range(len, 100, with_extra_column) +} + +/// Like [`make_staggered_batches_i32`], but with 'a' and 'b' drawn from +/// `0..key_range`, so that a small `key_range` makes few, large key groups +fn make_staggered_batches_i32_with_key_range( + len: usize, + key_range: i32, + with_extra_column: bool, +) -> Vec { let mut rng = rand::rng(); let mut input12: Vec<(i32, i32)> = vec![(0, 0); len]; let mut input3: Vec = vec![0; len]; let mut input4: Vec = vec![0; len]; for v in &mut input12 { - *v = (rng.random_range(0..100), rng.random_range(0..100)); + *v = ( + rng.random_range(0..key_range), + rng.random_range(0..key_range), + ); } rng.fill(&mut input3[..]); rng.fill(&mut input4[..]); diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index d5f87a04f0a28..9d750976fb909 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -27,11 +27,13 @@ use arrow::array::{Int64Array, RecordBatch, StringArray}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; -use datafusion_common::NullEquality; +use datafusion_common::{JoinSide, NullEquality}; use datafusion_execution::TaskContext; -use datafusion_physical_expr::expressions::col; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, col}; use datafusion_physical_plan::collect; -use datafusion_physical_plan::joins::{SortMergeJoinExec, utils::JoinOn}; +use datafusion_physical_plan::joins::SortMergeJoinExec; +use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; use datafusion_physical_plan::test::TestMemoryExec; use tokio::runtime::Runtime; @@ -104,11 +106,45 @@ fn schema() -> SchemaRef { ])) } +/// The join filter `left.data < right.data`. +fn data_lt_filter() -> JoinFilter { + let column_indices = vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + let intermediate_schema = Schema::new(vec![ + Field::new("left_data", DataType::Int64, false), + Field::new("right_data", DataType::Int64, false), + ]); + let expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("left_data", 0)), + Operator::Lt, + Arc::new(Column::new("right_data", 1)), + )); + JoinFilter::new(expression, column_indices, Arc::new(intermediate_schema)) +} + fn do_join( left: Arc, right: Arc, join_type: datafusion_common::JoinType, rt: &Runtime, +) -> usize { + do_join_with_filter(left, right, join_type, None, rt) +} + +fn do_join_with_filter( + left: Arc, + right: Arc, + join_type: datafusion_common::JoinType, + filter: Option, + rt: &Runtime, ) -> usize { let on: JoinOn = vec![( col("key", &left.schema()).unwrap(), @@ -118,7 +154,7 @@ fn do_join( left, right, on, - None, + filter, join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, @@ -168,6 +204,47 @@ fn bench_smj(c: &mut Criterion) { }); } + // 1:100 Inner Join — 20K left, 20K right, 200 distinct keys + // Every streamed row matches 100 buffered rows, so the join emits 2M rows. + { + let n = 20_000; + let key_mod = 200; + let left_batches = build_sorted_batches(n, key_mod, &s); + let right_batches = build_sorted_batches(n, key_mod, &s); + group.bench_function(BenchmarkId::new("inner_1to100", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join(left, right, datafusion_common::JoinType::Inner, &rt) + }) + }); + } + + // 1:100 Left Join with a filter — 20K left, 20K right, 200 distinct keys + // As in `inner_1to100`, every streamed row matches 100 buffered rows on the key, + // but here the pairs go through deferred filtering. `left.data < right.data` + // keeps about half of them, and none of the last streamed row of each key, which + // is null-joined. + { + let n = 20_000; + let key_mod = 200; + let left_batches = build_sorted_batches(n, key_mod, &s); + let right_batches = build_sorted_batches(n, key_mod, &s); + group.bench_function(BenchmarkId::new("left_1to100_filtered", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join_with_filter( + left, + right, + datafusion_common::JoinType::Left, + Some(data_lt_filter()), + &rt, + ) + }) + }); + } + // Left Join — 100K each, ~5% unmatched on left { let n = 100_000; diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs index 8027d83dbe124..892be4a2f39ec 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs @@ -25,7 +25,8 @@ use std::cmp::Ordering; use std::collections::VecDeque; use std::fmt::Debug; -use std::mem::size_of; +use std::iter::repeat_n; +use std::mem::{size_of, take}; use std::ops::Range; use std::sync::Arc; @@ -62,11 +63,12 @@ use futures::StreamExt; pub(super) struct StreamedJoinedChunk { /// Index of batch in buffered_data buffered_batch_idx: Option, - /// Array builder for streamed indices - streamed_indices: UInt64Builder, - /// Array builder for buffered indices - /// This could contain nulls if the join is null-joined - buffered_indices: UInt64Builder, + /// Streamed indices, which are never null + streamed_indices: Vec, + /// Buffered indices. The slot of a null-joined pair holds 0. + buffered_indices: Vec, + /// Validity of `buffered_indices`, null where the pair is null-joined + buffered_nulls: NullBufferBuilder, } /// Represents a record batch from streamed input. @@ -117,44 +119,115 @@ impl StreamedBatch { self.num_output_rows } - /// Appends new pair consisting of current streamed index and `buffered_idx` - /// index of buffered batch with `buffered_batch_idx` index. - fn append_output_pair( + /// Appends a pair of the current streamed index and a null buffered index, for a + /// streamed row that no buffered row matches. + fn append_null_joined_pair( &mut self, buffered_batch_idx: Option, - buffered_idx: Option, batch_size: usize, ) { - // If no current chunk exists or current chunk is not for current buffered batch, - // create a new chunk - if self.output_indices.is_empty() || self.buffered_batch_idx != buffered_batch_idx - { - // Compute capacity only when creating a new chunk (infrequent operation). - // The capacity is the remaining space to reach batch_size. - // This should always be >= 1 since we only call this when num_output_rows < batch_size. - debug_assert!( - batch_size > self.num_output_rows, - "batch_size ({batch_size}) must be > num_output_rows ({})", - self.num_output_rows - ); - let capacity = batch_size - self.num_output_rows; - self.output_indices.push(StreamedJoinedChunk { + let streamed_idx = self.idx as u64; + let current_chunk = self.chunk_for(buffered_batch_idx, batch_size); + current_chunk.streamed_indices.push(streamed_idx); + current_chunk.buffered_indices.push(0); + current_chunk.buffered_nulls.append_null(); + self.num_output_rows += 1; + } + + /// Appends a pair of the current streamed index and each of the `len` rows from + /// `buffered_start` on in the buffered batch with `buffered_batch_idx` index. The + /// indices are appended in bulk, which matters when a streamed row matches many + /// buffered rows. + fn append_output_pairs( + &mut self, + buffered_batch_idx: usize, + buffered_start: usize, + len: usize, + batch_size: usize, + ) { + debug_assert!( + len > 0 && self.num_output_rows + len <= batch_size, + "appending {len} pairs to {} with batch_size {batch_size}", + self.num_output_rows + ); + match len { + // A single match, as in a one-to-one join, is cheaper to append directly. + 1 => { + let streamed_idx = self.idx as u64; + let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); + current_chunk.streamed_indices.push(streamed_idx); + current_chunk.buffered_indices.push(buffered_start as u64); + current_chunk.buffered_nulls.append_non_null(); + self.num_output_rows += 1; + } + _ => self.append_output_range( buffered_batch_idx, - streamed_indices: UInt64Builder::with_capacity(capacity), - buffered_indices: UInt64Builder::with_capacity(capacity), - }); - self.buffered_batch_idx = buffered_batch_idx; + buffered_start, + len, + batch_size, + ), } - let current_chunk = self.output_indices.last_mut().unwrap(); + } - // Append index of streamed batch and index of buffered batch into current chunk - current_chunk.streamed_indices.append_value(self.idx as u64); - if let Some(idx) = buffered_idx { - current_chunk.buffered_indices.append_value(idx as u64); - } else { - current_chunk.buffered_indices.append_null(); + /// The bulk half of [`Self::append_output_pairs`], kept out of line so that the + /// one-row case, which one-to-one joins take for every row, does not pay for its + /// setup. + #[inline(never)] + fn append_output_range( + &mut self, + buffered_batch_idx: usize, + buffered_start: usize, + len: usize, + batch_size: usize, + ) { + let streamed_idx = self.idx as u64; + let current_chunk = self.chunk_for(Some(buffered_batch_idx), batch_size); + current_chunk + .streamed_indices + .extend(repeat_n(streamed_idx, len)); + let start = buffered_start as u64; + current_chunk + .buffered_indices + .extend(start..start + len as u64); + current_chunk.buffered_nulls.append_n_non_nulls(len); + self.num_output_rows += len; + } + + /// Returns the chunk that pairs for `buffered_batch_idx` go to, starting a new one + /// if there is no chunk yet or the current chunk is for another buffered batch. + #[inline(always)] + fn chunk_for( + &mut self, + buffered_batch_idx: Option, + batch_size: usize, + ) -> &mut StreamedJoinedChunk { + if self.output_indices.is_empty() || self.buffered_batch_idx != buffered_batch_idx + { + self.start_chunk(buffered_batch_idx, batch_size); } - self.num_output_rows += 1; + self.output_indices.last_mut().unwrap() + } + + /// Starts the chunk for `buffered_batch_idx`. Kept out of line because a chunk + /// usually takes many pairs, so the check in [`Self::chunk_for`] stays small. + #[cold] + #[inline(never)] + fn start_chunk(&mut self, buffered_batch_idx: Option, batch_size: usize) { + // The capacity is the remaining space to reach batch_size. + // This should always be >= 1 since we only call this when num_output_rows < batch_size. + debug_assert!( + batch_size > self.num_output_rows, + "batch_size ({batch_size}) must be > num_output_rows ({})", + self.num_output_rows + ); + let capacity = batch_size - self.num_output_rows; + self.output_indices.push(StreamedJoinedChunk { + buffered_batch_idx, + streamed_indices: Vec::with_capacity(capacity), + buffered_indices: Vec::with_capacity(capacity), + buffered_nulls: NullBufferBuilder::new(capacity), + }); + self.buffered_batch_idx = buffered_batch_idx; } } @@ -722,18 +795,20 @@ impl MaterializingSortMergeJoinStream { /// paused. Returns true when the group scan is complete and there is /// room for more pairs. fn pair_streamed_row_with_group(&mut self) -> bool { - while !self.buffered_data.scanning_finished() - && self.num_unfrozen_pairs() < self.batch_size - { - let scanning_idx = self.buffered_data.scanning_idx(); - self.streamed_batch.append_output_pair( - Some(self.buffered_data.scanning_batch_idx), - Some(scanning_idx), + let mut room = self.batch_size.saturating_sub(self.num_unfrozen_pairs()); + while !self.buffered_data.scanning_finished() && room > 0 { + let rows = self.buffered_data.scanning_remaining(); + let len = rows.len().min(room); + self.streamed_batch.append_output_pairs( + self.buffered_data.scanning_batch_idx, + rows.start, + len, self.batch_size, ); - self.buffered_data.scanning_advance(); + self.buffered_data.scanning_advance_by(len); + room -= len; } - if self.num_unfrozen_pairs() >= self.batch_size { + if room == 0 { return false; } @@ -754,11 +829,8 @@ impl MaterializingSortMergeJoinStream { } else { Some(self.buffered_data.scanning_batch_idx) }; - self.streamed_batch.append_output_pair( - scanning_batch_idx, - None, - self.batch_size, - ); + self.streamed_batch + .append_null_joined_pair(scanning_batch_idx, self.batch_size); } self.buffered_data.scanning_reset(); } @@ -774,7 +846,7 @@ impl MaterializingSortMergeJoinStream { .scanning_batch_mut() .null_joined .push(scanning_idx); - self.buffered_data.scanning_advance(); + self.buffered_data.scanning_advance_by(1); } } self.buffered_data.scanning_reset(); @@ -1519,11 +1591,15 @@ impl MaterializingSortMergeJoinStream { let mut total_matched_rows: usize = 0; for chunk in self.streamed_batch.output_indices.iter_mut() { - let left_indices = chunk.streamed_indices.finish(); + let left_indices = + UInt64Array::try_new(take(&mut chunk.streamed_indices).into(), None)?; if left_indices.is_empty() { continue; } - let right_indices: UInt64Array = chunk.buffered_indices.finish(); + let right_indices = UInt64Array::try_new( + take(&mut chunk.buffered_indices).into(), + chunk.buffered_nulls.finish(), + )?; if chunk.buffered_batch_idx.is_none() { let left_columns = @@ -1754,7 +1830,7 @@ impl MaterializingSortMergeJoinStream { // no key comparison. // // The keys a freeze sees are usually a contiguous run, since - // `scanning_advance` walks the deque in order. The exception is a + // `scanning_advance_by` walks the deque in order. The exception is a // freeze that straddles a `scanning_reset`: its window wraps (the // tail of one streamed row's pass, then the head of the next) and // leaves a gap, so the table is sized by the whole group rather than @@ -2113,8 +2189,11 @@ impl BufferedData { self.scanning_offset = 0; } - pub fn scanning_advance(&mut self) { - self.scanning_offset += 1; + /// Advances the scan by `n` rows, which must not pass the end of the scanning + /// batch, then skips to the next batch that has rows left. + pub fn scanning_advance_by(&mut self, n: usize) { + self.scanning_offset += n; + debug_assert!(self.scanning_offset <= self.scanning_batch().range.len()); while !self.scanning_finished() && self.scanning_batch_finished() { self.scanning_batch_idx += 1; self.scanning_offset = 0; @@ -2133,6 +2212,12 @@ impl BufferedData { self.scanning_batch().range.start + self.scanning_offset } + /// The rows of the scanning batch that the scan has not reached yet. + pub fn scanning_remaining(&self) -> Range { + let range = &self.scanning_batch().range; + range.start + self.scanning_offset..range.end + } + pub fn scanning_batch_finished(&self) -> bool { self.scanning_offset == self.scanning_batch().range.len() } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index f921812ec0619..9a77bc08f21ba 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -5084,6 +5084,70 @@ async fn join_with_group_spanning_batches_rescanned_per_streamed_row() -> Result Ok(()) } +/// A streamed row is paired with its buffered key group in runs, each ending where its +/// buffered batch ends or where the output batch fills. The test above uses one-row buffered +/// batches, so its runs have one row. Here the group of key 1 spans buffered batches of 3, 2 +/// and 1 rows and the batch size is 4, so runs of several rows are cut both ways. Every pair +/// must come out once, in the order of the buffered rows. +#[tokio::test] +async fn join_inner_group_spans_buffered_and_output_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 1, 2]), + ("b1", &vec![0, 0, 0]), + ("c1", &vec![10, 11, 20]), + ); + let right = build_table_from_batches(vec![ + build_table_i32( + ("a2", &vec![1, 1, 1]), + ("b2", &vec![0, 0, 0]), + ("c2", &vec![100, 101, 102]), + ), + build_table_i32( + ("a2", &vec![1, 1]), + ("b2", &vec![0, 0]), + ("c2", &vec![103, 104]), + ), + build_table_i32( + ("a2", &vec![1, 2]), + ("b2", &vec![0, 0]), + ("c2", &vec![105, 200]), + ), + ]); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(4)), + ); + let stream = join(left, right, on, Inner)?.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + assert!(batches.iter().all(|batch| batch.num_rows() <= 4)); + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+-----+ + | 1 | 0 | 10 | 1 | 0 | 100 | + | 1 | 0 | 10 | 1 | 0 | 101 | + | 1 | 0 | 10 | 1 | 0 | 102 | + | 1 | 0 | 10 | 1 | 0 | 103 | + | 1 | 0 | 10 | 1 | 0 | 104 | + | 1 | 0 | 10 | 1 | 0 | 105 | + | 1 | 0 | 11 | 1 | 0 | 100 | + | 1 | 0 | 11 | 1 | 0 | 101 | + | 1 | 0 | 11 | 1 | 0 | 102 | + | 1 | 0 | 11 | 1 | 0 | 103 | + | 1 | 0 | 11 | 1 | 0 | 104 | + | 1 | 0 | 11 | 1 | 0 | 105 | + | 2 | 0 | 20 | 2 | 0 | 200 | + +----+----+----+----+----+-----+ + "); + Ok(()) +} + /// A wrapped multi-source freeze that also carries a null buffered index. /// /// `materialize_right_columns` has two independent offsets in play on the