diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index c46ad9500b6ba..0604f83603e9f 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -78,6 +78,22 @@ fn build_sorted_batches_with_size( batches } +/// Like [`build_sorted_batches`], but with only the `key` and `data` columns of +/// [`narrow_schema`]. +fn build_narrow_sorted_batches(num_rows: usize, key_mod: usize) -> Vec { + let s = schema(); + let narrow = narrow_schema(); + build_sorted_batches(num_rows, key_mod, &s) + .into_iter() + .map(|b| { + b.project(&[0, 1]) + .unwrap() + .with_schema(Arc::clone(&narrow)) + .unwrap() + }) + .collect() +} + fn make_exec( batches: &[RecordBatch], schema: &SchemaRef, @@ -93,8 +109,21 @@ fn schema() -> SchemaRef { ])) } +/// The `key` and `data` columns of [`schema`]. +fn narrow_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("data", DataType::Int64, false), + ])) +} + /// The join filter `left.data < right.data`. fn data_lt_filter() -> JoinFilter { + data_filter(Operator::Lt) +} + +/// The join filter `left.data right.data`. +fn data_filter(op: Operator) -> JoinFilter { let column_indices = vec![ ColumnIndex { index: 1, @@ -111,7 +140,7 @@ fn data_lt_filter() -> JoinFilter { ]); let expression = Arc::new(BinaryExpr::new( Arc::new(Column::new("left_data", 0)), - Operator::Lt, + op, Arc::new(Column::new("right_data", 1)), )); JoinFilter::new(expression, column_indices, Arc::new(intermediate_schema)) @@ -174,6 +203,21 @@ fn bench_smj(c: &mut Criterion) { do_join(left, right, datafusion_common::JoinType::Inner, &rt) }) }); + // `left.data <= right.data` keeps every pair, and the indices of each + // freeze still form ranges, so the columns materialize as slices. + group.bench_function(BenchmarkId::new("inner_1to1_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::Inner, + Some(data_filter(Operator::LtEq)), + &rt, + ) + }) + }); } // 1:10 Inner Join — 100K left, 100K right, 10K distinct keys @@ -232,6 +276,103 @@ fn bench_smj(c: &mut Criterion) { }); } + // 1:10 and 1:100 Inner Joins with a filter. `left.data < right.data` keeps + // about half of the pairs. A streamed row meets a run of 10 or 100 buffered + // rows, so its pairs are filtered at freeze. + for (n, key_mod, name) in [ + (100_000, 10_000, "inner_1to10_filtered"), + (20_000, 200, "inner_1to100_filtered"), + ] { + 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(name, 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::Inner, + Some(data_lt_filter()), + &rt, + ) + }) + }); + } + + // 1:10 Inner Joins with a filter that keeps most pairs — 100K left, 100K right, + // 10K distinct keys. `left.data <> right.data` drops only the pair of a row with + // itself, keeping 90% of the pairs. The narrow case joins just `key` and `data`, + // so the filter reads half of the output columns. + { + let n = 100_000; + let key_mod = 10_000; + 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_1to10_high_pass_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::Inner, + Some(data_filter(Operator::NotEq)), + &rt, + ) + }) + }, + ); + let narrow = narrow_schema(); + let left_batches = build_narrow_sorted_batches(n, key_mod); + let right_batches = build_narrow_sorted_batches(n, key_mod); + group.bench_function( + BenchmarkId::new("inner_1to10_high_pass_filtered_narrow", n), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &narrow); + let right = make_exec(&right_batches, &narrow); + do_join_with_filter( + left, + right, + datafusion_common::JoinType::Inner, + Some(data_filter(Operator::NotEq)), + &rt, + ) + }) + }, + ); + } + + // 1:1000 Inner Join with a selective filter — 2K left, 20K right, 20 keys + // Every streamed row meets a run of 1000 buffered rows, the shape of TPC-DS + // q72's `catalog_sales` to `inventory` join. Each key's 100 streamed rows take + // the smallest of its `data` values, so `left.data > right.data` keeps about + // 5% of the 2M pairs. Such runs are filtered as they are paired. + { + let keys = 20; + let left_batches = build_sorted_batches(keys * 100, keys, &s); + let right_batches = build_sorted_batches(keys * 1000, keys, &s); + group.bench_function( + BenchmarkId::new("inner_1to1000_selective_filtered", keys * 1000), + |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::Inner, + Some(data_filter(Operator::Gt)), + &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 65bc6afcc3609..929961f30a271 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 @@ -29,16 +29,18 @@ use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAda use crate::{PhysicalExpr, SendableRecordBatchStream}; use arrow::array::{types::UInt64Type, *}; +use arrow::buffer::BooleanBuffer; use arrow::compute::{ - self, BatchCoalescer, SortOptions, concat_batches, filter_record_batch, interleave, - take_arrays, + self, BatchCoalescer, FilterBuilder, SortOptions, concat_batches, + filter_record_batch, interleave, take_arrays, }; use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatchOptions; -use datafusion_common::cast::as_uint64_array; +use datafusion_common::cast::{as_boolean_array, as_uint64_array}; use datafusion_common::instant::Instant; use datafusion_common::{ - DataFusionError, JoinSide, JoinType, NullEquality, Result, exec_err, internal_err, + DataFusionError, JoinSide, JoinType, NullEquality, Result, exec_err, + internal_datafusion_err, internal_err, }; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::runtime_env::RuntimeEnv; @@ -47,6 +49,23 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use futures::StreamExt; +/// The matched pairs of one [`StreamedJoinedChunk`], with the indices finished +/// into arrays. +struct MatchedChunk { + /// Index of the batch in `buffered_data` that the buffered rows come from + buffered_batch_idx: usize, + /// The streamed row of each pair + streamed_indices: UInt64Array, + /// The buffered row of each pair + buffered_indices: UInt64Array, +} + +impl MatchedChunk { + fn len(&self) -> usize { + self.streamed_indices.len() + } +} + /// Represents a chunk of joined data from streamed and buffered side pub(super) struct StreamedJoinedChunk { /// Index of batch in buffered_data @@ -360,6 +379,23 @@ pub(super) struct MaterializingSortMergeJoinStream { /// since NULL-padded rows reach filter evaluation. `None` unless /// `deferred_filtering`. pub filter_schema: Option, + /// Whether this is an Inner join with a filter. Such a join drops every pair + /// the filter rejects, so the freeze can run the filter before it + /// materializes the output columns (see [`Self::freeze_inner_filtered`]). + /// Outer joins keep a rejected pair for the deferred pipeline, which may + /// null-join its row. Only an Inner join filters early, so its streamed side + /// is the left input. + pub early_filtering: bool, + /// The indices of every buffered column, for materializing all of them + pub buffered_columns: Vec, + /// The buffered columns that the join filter reads, in the filter's order. + /// Empty unless `early_filtering`. + pub buffered_filter_columns: Vec, + /// How many pairs the Inner join's filter has evaluated at freeze so far + /// (see [`Self::filter_first_pays`]) + pub freeze_filter_evaluated: u64, + /// How many of the pairs in `freeze_filter_evaluated` passed the filter + pub freeze_filter_passed: u64, /// Target output batch size pub batch_size: usize, @@ -622,6 +658,17 @@ impl MaterializingSortMergeJoinStream { JoinSide::Right }; let deferred_filtering = needs_deferred_filtering(filter.as_ref(), join_type); + let early_filtering = join_type == JoinType::Inner && filter.is_some(); + let buffered_columns = (0..buffered_schema.fields().len()).collect(); + let buffered_filter_columns = match &filter { + Some(f) if early_filtering => f + .column_indices() + .iter() + .filter(|c| c.side == JoinSide::Right) + .map(|c| c.index) + .collect(), + _ => vec![], + }; let filter_schema = if deferred_filtering && let Some(f) = &filter { let mut fields = Vec::with_capacity(f.schema().fields().len()); for (column_index, field) in @@ -657,6 +704,11 @@ impl MaterializingSortMergeJoinStream { on_buffered, deferred_filtering, filter_schema, + early_filtering, + buffered_columns, + buffered_filter_columns, + freeze_filter_evaluated: 0, + freeze_filter_passed: 0, filter, joined_record_batches: JoinedRecordBatches { joined_batches: new_output_coalescer(Arc::clone(&schema), batch_size), @@ -1520,6 +1572,7 @@ impl MaterializingSortMergeJoinStream { &self.streamed_schema, &buffered_indices, buffered_batch, + &self.buffered_columns, )? { self.joined_record_batches .push_batch_with_null_metadata(record_batch, self.join_type); @@ -1556,6 +1609,7 @@ impl MaterializingSortMergeJoinStream { &self.streamed_schema, &buffered_indices, buffered_batch, + &self.buffered_columns, )? { self.joined_record_batches .push_batch_with_null_metadata(record_batch, self.join_type); @@ -1574,7 +1628,7 @@ impl MaterializingSortMergeJoinStream { // Matched chunks are collected and processed together in // freeze_streamed_matched() to amortize filter evaluation overhead. fn freeze_streamed(&mut self) -> Result<()> { - let mut matched_chunks: Vec<(usize, UInt64Array, UInt64Array)> = + let mut matched_chunks: Vec = Vec::with_capacity(self.streamed_batch.output_indices.len()); let mut total_matched_rows: usize = 0; @@ -1617,11 +1671,11 @@ impl MaterializingSortMergeJoinStream { } total_matched_rows += left_indices.len(); - matched_chunks.push(( - chunk.buffered_batch_idx.unwrap(), - left_indices, - right_indices, - )); + matched_chunks.push(MatchedChunk { + buffered_batch_idx: chunk.buffered_batch_idx.unwrap(), + streamed_indices: left_indices, + buffered_indices: right_indices, + }); } if !matched_chunks.is_empty() { @@ -1639,7 +1693,7 @@ impl MaterializingSortMergeJoinStream { /// cost when keys are near-unique (1 row per chunk). fn freeze_streamed_matched( &mut self, - matched_chunks: &[(usize, UInt64Array, UInt64Array)], + matched_chunks: &[MatchedChunk], total_matched_rows: usize, ) -> Result<()> { debug_assert!( @@ -1647,34 +1701,34 @@ impl MaterializingSortMergeJoinStream { "caller guards this with an is_empty check before calling" ); debug_assert!( - matched_chunks.iter().all(|(idx, left, right)| { - left.len() == right.len() && *idx < self.buffered_data.batches.len() + matched_chunks.iter().all(|chunk| { + chunk.streamed_indices.len() == chunk.buffered_indices.len() + && chunk.buffered_batch_idx < self.buffered_data.batches.len() }), "left/right indices are built in pairs from the same streamed×buffered cross, \ and batch_idx comes from iterating buffered_data.batches" ); debug_assert_eq!( - matched_chunks - .iter() - .map(|(_, l, _)| l.len()) - .sum::(), + matched_chunks.iter().map(MatchedChunk::len).sum::(), total_matched_rows, "total_matched_rows is accumulated from the same chunks in freeze_streamed" ); - let combined_left_indices = if matched_chunks.len() == 1 { - matched_chunks[0].1.clone() - } else { - let refs: Vec<&dyn Array> = - matched_chunks.iter().map(|c| &c.1 as &dyn Array).collect(); - as_uint64_array(&compute::concat(&refs)?)?.clone() - }; + if self.early_filtering { + return self.freeze_inner_filtered(matched_chunks, total_matched_rows); + } + + let combined_left_indices = + concat_indices(matched_chunks.iter().map(|c| &c.streamed_indices))?; let left_columns = materialize_left_columns(&self.streamed_batch.batch, &combined_left_indices)?; - let right_columns = - self.materialize_right_columns(matched_chunks, total_matched_rows)?; + let right_columns = self.materialize_right_columns( + matched_chunks, + total_matched_rows, + &self.buffered_columns, + )?; // `get_filter_columns` takes arrays in join-side order. `left_columns` // / `right_columns` here are streamed / buffered, and for a Right join @@ -1694,50 +1748,33 @@ impl MaterializingSortMergeJoinStream { // A filter that reads no columns (e.g. a bound parameter or a volatile // function) still has to be evaluated, so key on the filter itself. + // An Inner join's filter ran in `freeze_inner_filtered` above, so a + // filter here belongs to an outer join, which filters deferred. if let Some(f) = &self.filter { + debug_assert!(self.deferred_filtering); let schema = self.filter_schema.as_ref().unwrap_or_else(|| f.schema()); - let filter_batch = RecordBatch::try_new_with_options( - Arc::clone(schema), - filter_columns, - &RecordBatchOptions::new().with_row_count(Some(total_matched_rows)), - )?; - let filter_result = f - .expression() - .evaluate(&filter_batch)? - .into_array(filter_batch.num_rows())?; - - let filter_result_mask = - datafusion_common::cast::as_boolean_array(&filter_result)?; - - // Convert NULL filter results to false — NULL means "not satisfied" - // per SQL semantics, same as Left/Right outer joins. - let mask = if filter_result_mask.null_count() > 0 { - compute::prep_null_mask_filter(filter_result_mask) - } else { - filter_result_mask.clone() - }; + let mask = BooleanArray::new( + evaluate_join_filter(f, schema, filter_columns, total_matched_rows)?, + None, + ); - if self.deferred_filtering { - self.joined_record_batches.push_batch_with_filter_metadata( - output_batch, - &combined_left_indices, - &mask, - self.streamed_batch_counter, - self.join_type, - ); - } else { - let filtered_batch = filter_record_batch(&output_batch, &mask)?; - self.joined_record_batches - .push_batch_without_metadata(filtered_batch); - } + self.joined_record_batches.push_batch_with_filter_metadata( + output_batch, + &combined_left_indices, + &mask, + self.streamed_batch_counter, + self.join_type, + ); // Track which buffered rows had all filter matches fail, // so full join can emit them as null-joined later. if self.join_type == JoinType::Full { let mut offset = 0usize; - for (batch_idx, _left, right) in matched_chunks { + for chunk in matched_chunks { + let right = &chunk.buffered_indices; let chunk_len = right.len(); - let buffered_batch = &mut self.buffered_data.batches[*batch_idx]; + let buffered_batch = + &mut self.buffered_data.batches[chunk.buffered_batch_idx]; for i in 0..chunk_len { if right.is_null(i) { @@ -1771,7 +1808,188 @@ impl MaterializingSortMergeJoinStream { Ok(()) } - /// Materializes right-side columns across all matched chunks. + /// Freezes the matched pairs of an Inner join with a filter. The filter runs + /// on just the columns it reads, and only the pairs that pass have their + /// output columns materialized. A selective filter over a many-to-many join + /// rejects most pairs, so materializing every column of every pair before + /// filtering would spend most of the freeze on rows that are then dropped. + /// When the pairs' indices form ranges on both sides, as in a 1:1 join, the + /// columns materialize as slices and filtering first cannot save anything. + /// Then, or when [`Self::filter_first_pays`] says earlier freezes kept too + /// many of their pairs, the batch is materialized whole and then filtered. + fn freeze_inner_filtered( + &mut self, + matched_chunks: &[MatchedChunk], + total_matched_rows: usize, + ) -> Result<()> { + let Some(filter) = &self.filter else { + return internal_err!("early filtering requires a join filter"); + }; + debug_assert_eq!( + self.join_type, + JoinType::Inner, + "only an Inner join filters early, so the streamed side is the left one" + ); + + let streamed_indices = + concat_indices(matched_chunks.iter().map(|c| &c.streamed_indices))?; + let single_source = matched_chunks + .iter() + .all(|c| c.buffered_batch_idx == matched_chunks[0].buffered_batch_idx); + let buffered_indices = + concat_indices(matched_chunks.iter().map(|c| &c.buffered_indices))?; + let slices = single_source + && spans_range(&streamed_indices) + && spans_range(&buffered_indices); + if slices || !self.filter_first_pays(filter) { + return self.freeze_inner_filtered_whole( + matched_chunks, + total_matched_rows, + &streamed_indices, + ); + } + + let mut buffered_arrays = self + .materialize_right_columns( + matched_chunks, + total_matched_rows, + &self.buffered_filter_columns, + )? + .into_iter(); + let columns = filter + .column_indices() + .iter() + .map(|column_index| match column_index.side { + JoinSide::Left => take_column( + self.streamed_batch.batch.column(column_index.index), + &streamed_indices, + ), + JoinSide::Right => buffered_arrays.next().ok_or_else(|| { + internal_datafusion_err!("missing buffered filter column") + }), + JoinSide::None => internal_err!("Unexpected JoinSide::None in filter"), + }) + .collect::>>()?; + let passed = + evaluate_join_filter(filter, filter.schema(), columns, total_matched_rows)?; + let kept_rows = passed.count_set_bits(); + self.freeze_filter_evaluated += total_matched_rows as u64; + self.freeze_filter_passed += kept_rows as u64; + if kept_rows == 0 { + return Ok(()); + } + + // Filter the index arrays of all chunks at once, and give each chunk + // with a pair that passed a slice of the result. + let predicate = FilterBuilder::new(&BooleanArray::new(passed.clone(), None)) + .optimize() + .build(); + let kept_streamed = + as_uint64_array(&predicate.filter(&streamed_indices)?)?.clone(); + let kept_buffered = + as_uint64_array(&predicate.filter(&buffered_indices)?)?.clone(); + let mut kept = Vec::with_capacity(matched_chunks.len()); + let (mut offset, mut kept_offset) = (0, 0); + for chunk in matched_chunks { + let count = passed.slice(offset, chunk.len()).count_set_bits(); + offset += chunk.len(); + if count > 0 { + kept.push(MatchedChunk { + buffered_batch_idx: chunk.buffered_batch_idx, + streamed_indices: kept_streamed.slice(kept_offset, count), + buffered_indices: kept_buffered.slice(kept_offset, count), + }); + kept_offset += count; + } + } + + let streamed_columns = + materialize_left_columns(&self.streamed_batch.batch, &kept_streamed)?; + let buffered_columns = + self.materialize_right_columns(&kept, kept_rows, &self.buffered_columns)?; + let output_batch = RecordBatch::try_new( + Arc::clone(&self.schema), + [streamed_columns, buffered_columns].concat(), + )?; + self.joined_record_batches + .push_batch_without_metadata(output_batch); + Ok(()) + } + + /// Freezes the matched pairs of an Inner join with a filter by materializing + /// every column of every pair, evaluating the filter on the columns it reads, + /// and filtering the batch. + fn freeze_inner_filtered_whole( + &mut self, + matched_chunks: &[MatchedChunk], + total_matched_rows: usize, + streamed_indices: &UInt64Array, + ) -> Result<()> { + let streamed_columns = + materialize_left_columns(&self.streamed_batch.batch, streamed_indices)?; + let buffered_columns = self.materialize_right_columns( + matched_chunks, + total_matched_rows, + &self.buffered_columns, + )?; + let Some(filter) = &self.filter else { + return internal_err!("early filtering requires a join filter"); + }; + let filter_columns = + get_filter_columns(Some(filter), &streamed_columns, &buffered_columns)?; + let passed = evaluate_join_filter( + filter, + filter.schema(), + filter_columns, + total_matched_rows, + )?; + let kept_rows = passed.count_set_bits(); + self.freeze_filter_evaluated += total_matched_rows as u64; + self.freeze_filter_passed += kept_rows as u64; + if kept_rows == 0 { + return Ok(()); + } + + let output_batch = RecordBatch::try_new( + Arc::clone(&self.schema), + [streamed_columns, buffered_columns].concat(), + )?; + let output_batch = if kept_rows == total_matched_rows { + output_batch + } else { + filter_record_batch(&output_batch, &BooleanArray::new(passed, None))? + }; + self.joined_record_batches + .push_batch_without_metadata(output_batch); + Ok(()) + } + + /// Whether filtering the pairs of a freeze before materializing them is cheaper + /// than materializing every column of every pair and filtering the batch. + /// + /// Filtering first gathers the filter's `f` columns for every pair, filters the + /// pairs' two index arrays, and gathers all `k` output columns for the share + /// `p` of pairs that pass. The other way gathers the `k` columns for every + /// pair and copies the share `p` of them to filter the batch. Counting a + /// gathered value as 1 and a copied one as 1/2, filtering first is cheaper + /// while `f + 1 + k·p < k + k·p/2`, that is while `p < 2·(k - f - 1)/k`. `p` is + /// the share that passed in the earlier freezes, or 0 before the first. + /// + /// The model counts columns rather than bytes, so a string column weighs the + /// same as an integer one, and `p` averages over the life of the stream + /// rather than tracking recent freezes. + fn filter_first_pays(&self, filter: &JoinFilter) -> bool { + let k = self.schema.fields().len() as f64; + let f = filter.column_indices().len() as f64; + let p = if self.freeze_filter_evaluated == 0 { + 0.0 + } else { + self.freeze_filter_passed as f64 / self.freeze_filter_evaluated as f64 + }; + p < 2.0 * (k - f - 1.0) / k + } + + /// Materializes the buffered columns at `columns` across all matched chunks. /// /// When chunks reference a single buffered batch, indices are concatenated /// for a single fetch. When multiple batches are involved, `interleave` @@ -1779,25 +1997,24 @@ impl MaterializingSortMergeJoinStream { /// handles null right indices (unmatched streamed rows). fn materialize_right_columns( &self, - matched_chunks: &[(usize, UInt64Array, UInt64Array)], + matched_chunks: &[MatchedChunk], total_matched_rows: usize, + columns: &[usize], ) -> Result> { - let first_batch_idx = matched_chunks[0].0; - let single_source = matched_chunks.iter().all(|c| c.0 == first_batch_idx); + let first_batch_idx = matched_chunks[0].buffered_batch_idx; + let single_source = matched_chunks + .iter() + .all(|c| c.buffered_batch_idx == first_batch_idx); if single_source { - let combined_right_indices = if matched_chunks.len() == 1 { - matched_chunks[0].2.clone() - } else { - let refs: Vec<&dyn Array> = - matched_chunks.iter().map(|c| &c.2 as &dyn Array).collect(); - as_uint64_array(&compute::concat(&refs)?)?.clone() - }; + let combined_right_indices = + concat_indices(matched_chunks.iter().map(|c| &c.buffered_indices))?; return fetch_right_columns_by_idxs( &self.buffered_data, first_batch_idx, &combined_right_indices, + columns, ); } @@ -1810,7 +2027,7 @@ impl MaterializingSortMergeJoinStream { // common all-matched case. let needs_null_sentinel = matched_chunks .iter() - .any(|(_, _, right)| right.null_count() > 0); + .any(|c| c.buffered_indices.null_count() > 0); let source_offset = usize::from(needs_null_sentinel); // Map each distinct `buffered_batch_idx` to a contiguous source @@ -1867,8 +2084,8 @@ impl MaterializingSortMergeJoinStream { // 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)) + .fold((usize::MAX, 0usize), |(lo, hi), c| { + (lo.min(c.buffered_batch_idx), hi.max(c.buffered_batch_idx)) }); // Every key indexes the live buffered deque -- this is what keeps // the key space dense, and what makes `source_data` below safe. @@ -1883,11 +2100,13 @@ impl MaterializingSortMergeJoinStream { Vec::with_capacity(matched_chunks.len().min(source_of_batch.len())); let mut interleave_indices: Vec<(usize, usize)> = Vec::with_capacity(total_matched_rows); - for (batch_idx, _, right) in matched_chunks { + for chunk in matched_chunks.iter() { + let batch_idx = chunk.buffered_batch_idx; + let right = &chunk.buffered_indices; let slot = &mut source_of_batch[batch_idx - min_batch_idx]; if *slot == UNSEEN { *slot = source_batches.len(); - source_batches.push(*batch_idx); + source_batches.push(batch_idx); } let source = *slot + source_offset; if right.null_count() == 0 { @@ -1906,8 +2125,6 @@ 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: Vec<&RecordBatch> = source_batches .iter() @@ -1919,13 +2136,12 @@ impl MaterializingSortMergeJoinStream { }) .collect::>()?; - // One single-row null array per column, built up front so the + // One single-row null array per selected 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() + columns .iter() - .map(|f| new_null_array(f.data_type(), 1)) + .map(|&c| new_null_array(self.buffered_schema.field(c).data_type(), 1)) .collect() } else { vec![] @@ -1933,10 +2149,10 @@ impl MaterializingSortMergeJoinStream { 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 mut right_columns = Vec::with_capacity(columns.len()); + for (position, &col_idx) in columns.iter().enumerate() { source_arrays.clear(); - source_arrays.extend(null_arrays.get(col_idx).map(|a| a.as_ref())); + source_arrays.extend(null_arrays.get(position).map(|a| a.as_ref())); source_arrays.extend(source_data.iter().map(|d| d.column(col_idx).as_ref())); right_columns.push(interleave(&source_arrays, &interleave_indices)?); @@ -2065,14 +2281,18 @@ fn produce_buffered_null_batch( streamed_schema: &SchemaRef, buffered_indices: &PrimitiveArray, buffered_batch: &BufferedBatch, + buffered_columns: &[usize], ) -> Result> { if buffered_indices.is_empty() { return Ok(None); } // Take buffered (right) columns - let right_columns = - fetch_right_columns_from_batch_by_idxs(buffered_batch, buffered_indices)?; + let right_columns = fetch_right_columns_from_batch_by_idxs( + buffered_batch, + buffered_indices, + buffered_columns, + )?; // Create null streamed (left) columns let mut left_columns = streamed_schema @@ -2113,16 +2333,28 @@ fn is_contiguous_range(indices: &UInt64Array) -> Option> { Some(start as usize..(start + len) as usize) } -/// Get `buffered_indices` rows for `buffered_data[buffered_batch_idx]` by specific column indices +/// Whether `indices` may form a range, judging by their ends only. A freeze uses +/// it to choose its path, and the path checks the range before it slices. +fn spans_range(indices: &UInt64Array) -> bool { + let values = indices.values(); + !values.is_empty() + && indices.null_count() == 0 + && values[values.len() - 1].checked_sub(values[0]) + == Some(values.len() as u64 - 1) +} + +/// Get `buffered_indices` rows of the `columns` of `buffered_data[buffered_batch_idx]` #[inline(always)] fn fetch_right_columns_by_idxs( buffered_data: &BufferedData, buffered_batch_idx: usize, buffered_indices: &UInt64Array, + columns: &[usize], ) -> Result> { fetch_right_columns_from_batch_by_idxs( &buffered_data.batches[buffered_batch_idx], buffered_indices, + columns, ) } @@ -2130,13 +2362,17 @@ fn fetch_right_columns_by_idxs( fn fetch_right_columns_from_batch_by_idxs( buffered_batch: &BufferedBatch, buffered_indices: &UInt64Array, + columns: &[usize], ) -> Result> { match &buffered_batch.batch { BufferedBatchState::InMemory(batch) => { + let arrays = columns.iter().map(|&c| batch.column(c)); if let Some(range) = is_contiguous_range(buffered_indices) { - Ok(batch.slice(range.start, range.len()).columns().to_vec()) + Ok(arrays.map(|a| a.slice(range.start, range.len())).collect()) } else { - Ok(take_arrays(batch.columns(), buffered_indices, None)?) + arrays + .map(|a| Ok(compute::take(a.as_ref(), buffered_indices, None)?)) + .collect() } } BufferedBatchState::Spilled(_) => { @@ -2147,6 +2383,50 @@ fn fetch_right_columns_from_batch_by_idxs( } } +/// Concatenates the index arrays of several chunks into one. +fn concat_indices<'a>( + indices: impl Iterator, +) -> Result { + let arrays: Vec<&dyn Array> = indices.map(|a| a as &dyn Array).collect(); + match arrays.as_slice() { + [single] => Ok(as_uint64_array(*single)?.clone()), + _ => Ok(as_uint64_array(&compute::concat(&arrays)?)?.clone()), + } +} + +/// The rows of `array` at `indices`, sliced when the indices form a range. +fn take_column(array: &ArrayRef, indices: &UInt64Array) -> Result { + if let Some(range) = is_contiguous_range(indices) { + Ok(array.slice(range.start, range.len())) + } else { + Ok(compute::take(array.as_ref(), indices, None)?) + } +} + +/// Evaluates `filter` over `num_rows` pairs whose `columns` follow the filter's +/// intermediate `schema`, and returns which pairs pass. A NULL result does not +/// pass, as in SQL. +fn evaluate_join_filter( + filter: &JoinFilter, + schema: &SchemaRef, + columns: Vec, + num_rows: usize, +) -> Result { + // A filter that reads no columns (e.g. a bound parameter or a volatile + // function) still has to be evaluated, so the batch carries its row count. + let batch = RecordBatch::try_new_with_options( + Arc::clone(schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + )?; + let result = filter.expression().evaluate(&batch)?.into_array(num_rows)?; + let mask = as_boolean_array(&result)?; + Ok(match mask.nulls() { + Some(nulls) => mask.values() & nulls.inner(), + None => mask.values().clone(), + }) +} + /// Buffered data contains all buffered batches with one unique join key #[derive(Debug, Default)] pub(super) struct BufferedData { 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 babc9851a985c..a48a038730dde 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -4881,8 +4881,11 @@ fn build_batched_finish_barrier_table( } /// Concat and sort batches by all the columns to make sure we can compare them with different join -fn prepare_record_batches_for_cmp(output: Vec) -> RecordBatch { - let output_batch = arrow::compute::concat_batches(output[0].schema_ref(), &output) +fn prepare_record_batches_for_cmp( + schema: &SchemaRef, + output: Vec, +) -> RecordBatch { + let output_batch = arrow::compute::concat_batches(schema, &output) .expect("failed to concat batches"); // Sort on all columns to make sure we have a deterministic order for the assertion @@ -4937,7 +4940,7 @@ async fn join_get_stream_and_get_expected( let batches = common::collect(stream).await?; - prepare_record_batches_for_cmp(batches) + prepare_record_batches_for_cmp(&oracle.schema(), batches) }; let join = SortMergeJoinExec::try_new( @@ -5150,7 +5153,8 @@ async fn test_should_emit_early_when_have_enough_data_to_emit() -> Result<()> { // Just sanity test to make sure we are still producing valid output { let output = [output_batched, output_batches_after_finish].concat(); - let actual_prepared = prepare_record_batches_for_cmp(output); + let actual_prepared = + prepare_record_batches_for_cmp(&expected.schema(), output); assert_eq!(actual_prepared.columns(), expected.columns()); } @@ -8395,3 +8399,324 @@ async fn left_join_with_partial_filter_preserves_streamed_order() -> Result<()> ); Ok(()) } + +// ============================================================================ +// Early filtering of Inner joins +// ============================================================================ + +/// Rows for the early-filtering tests, sorted by the key column `names[0]`: key +/// `k` of `0..keys` has `rows_per_key(k)` rows. `names[1]` is the row number and +/// `names[2]` a pseudo-random value in `0..1000` that is NULL in every 37th row. +/// The rows are split into batches of `rows_per_batch`. +fn early_filtering_batches( + names: [&str; 3], + keys: i32, + rows_per_key: impl Fn(i32) -> usize, + seed: u64, + rows_per_batch: usize, +) -> (Vec, SchemaRef) { + let mut key_values = vec![]; + let mut row_numbers = vec![]; + let mut values = vec![]; + let mut state = seed; + for key in 0..keys { + for _ in 0..rows_per_key(key) { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let row = key_values.len() as i32; + key_values.push(key); + row_numbers.push(row); + values.push((row % 37 != 0).then_some(((state >> 33) % 1000) as i32)); + } + } + let schema = Arc::new(Schema::new(vec![ + Field::new(names[0], DataType::Int32, false), + Field::new(names[1], DataType::Int32, false), + Field::new(names[2], DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(key_values)), + Arc::new(Int32Array::from(row_numbers)), + Arc::new(Int32Array::from(values)), + ], + ) + .unwrap(); + let batches = (0..batch.num_rows()) + .step_by(rows_per_batch) + .map(|offset| batch.slice(offset, rows_per_batch.min(batch.num_rows() - offset))) + .collect(); + (batches, schema) +} + +/// Inputs for the early-filtering tests. Every third key has a group of 600 +/// buffered rows, which the 250-row buffered batches split, so that a freeze +/// gathers its pairs from several buffered batches. The other keys have groups of +/// one or two buffered rows. Keys repeat up to four times on the streamed side. +fn early_filtering_inputs() +-> ((Vec, SchemaRef), (Vec, SchemaRef)) { + let left = + early_filtering_batches(["a1", "b1", "c1"], 40, |k| 1 + (k % 4) as usize, 1, 7); + let right = early_filtering_batches( + ["a2", "b2", "c2"], + 40, + |k| if k % 3 == 0 { 600 } else { (k % 3) as usize }, + 2, + 250, + ); + (left, right) +} + +/// The filter `c2 > c1 AND b1 % 2 = b2 % 2`, whose intermediate columns are in +/// the order `c2, c1, b1, b2`, so that they alternate between the join sides. +fn interleaved_early_filter(left_schema: &Schema, right_schema: &Schema) -> JoinFilter { + let parity = |name: &str, index: usize| -> PhysicalExprRef { + Arc::new(BinaryExpr::new( + Arc::new(Column::new(name, index)), + Operator::Modulo, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )) + }; + let expression = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c2", 0)), + Operator::Gt, + Arc::new(Column::new("c1", 1)), + )), + Operator::And, + Arc::new(BinaryExpr::new( + parity("b1", 2), + Operator::Eq, + parity("b2", 3), + )), + )); + let column_indices = vec![ + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + let schema = Schema::new(vec![ + right_schema.field_with_name("c2").unwrap().clone(), + left_schema.field_with_name("c1").unwrap().clone(), + left_schema.field_with_name("b1").unwrap().clone(), + right_schema.field_with_name("b2").unwrap().clone(), + ]); + JoinFilter::new(expression, column_indices, Arc::new(schema)) +} + +/// The filter `c2 <> c1 OR b2 < b1`, with its intermediate columns in the order +/// `b2, b1, c2, c1`. It rejects only the few pairs with a NULL or equal `c`, so +/// it reads four of the six output columns and keeps most pairs. After the first +/// freeze the join then materializes whole batches and filters them (see +/// `filter_first_pays`). +fn high_pass_early_filter(left_schema: &Schema, right_schema: &Schema) -> JoinFilter { + let expression = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c2", 2)), + Operator::NotEq, + Arc::new(Column::new("c1", 3)), + )), + Operator::Or, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b2", 0)), + Operator::Lt, + Arc::new(Column::new("b1", 1)), + )), + )); + let column_indices = vec![ + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ]; + let schema = Schema::new(vec![ + right_schema.field_with_name("b2").unwrap().clone(), + left_schema.field_with_name("b1").unwrap().clone(), + right_schema.field_with_name("c2").unwrap().clone(), + left_schema.field_with_name("c1").unwrap().clone(), + ]); + JoinFilter::new(expression, column_indices, Arc::new(schema)) +} + +/// A filter that reads no columns and always evaluates to `value`. +fn column_free_early_filter(value: bool) -> JoinFilter { + JoinFilter::new( + Arc::new(Literal::new(ScalarValue::Boolean(Some(value)))), + vec![], + Arc::new(Schema::empty()), + ) +} + +fn early_filtering_on(left_schema: &Schema, right_schema: &Schema) -> Result { + Ok(vec![( + Arc::new(Column::new_with_schema("a1", left_schema)?) as _, + Arc::new(Column::new_with_schema("a2", right_schema)?) as _, + )]) +} + +/// Asserts that the pairs of `batches` ascend on `(b1, b2)`, the row numbers of +/// the two inputs: an Inner join keeps the streamed order, and the buffered order +/// within each streamed row. +fn assert_streamed_order(batches: &[RecordBatch]) { + let column = |b: &RecordBatch, i: usize| { + b.column(i) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }; + let pairs: Vec<(i32, i32)> = batches + .iter() + .flat_map(|b| column(b, 1).into_iter().zip(column(b, 4))) + .collect(); + assert!( + pairs.is_sorted(), + "Inner join output must keep the order of both inputs" + ); +} + +/// An Inner join with a filter filters the pairs of each freeze before it +/// materializes them. It must keep the same pairs as a hash join whatever the +/// batch size, which decides where a freeze cuts a group, and whatever the order +/// of the filter's columns. The pseudo-random values pass `c1 < c2` for about +/// half of the pairs, and their NULLs pass for none. Freezes that materialize +/// whole batches are covered by the high-pass filter and by the single pairs of +/// `batch_size` 1, whose indices form ranges. A filter that reads no columns must +/// still run for every pair. +#[tokio::test] +async fn join_inner_early_filtering_matches_hash_join() -> Result<()> { + let ((left_batches, left_schema), (right_batches, right_schema)) = + early_filtering_inputs(); + let on = early_filtering_on(&left_schema, &right_schema)?; + // Each filter with whether it keeps any pair of these inputs. + let filters = [ + (build_c1_lt_c2_filter(&left_schema, &right_schema), true), + (interleaved_early_filter(&left_schema, &right_schema), true), + (high_pass_early_filter(&left_schema, &right_schema), true), + (column_free_early_filter(true), true), + (column_free_early_filter(false), false), + ]; + let memory_exec = |batches: &Vec, schema: &SchemaRef| { + TestMemoryExec::try_new_exec( + std::slice::from_ref(batches), + Arc::clone(schema), + None, + ) + }; + + for (filter, keeps_pairs) in filters { + for batch_size in [1, 7, 300, 8192] { + let (stream, expected) = join_get_stream_and_get_expected( + memory_exec(&left_batches, &left_schema)?, + memory_exec(&right_batches, &right_schema)?, + memory_exec(&left_batches, &left_schema)?, + memory_exec(&right_batches, &right_schema)?, + on.clone(), + Inner, + Some(filter.clone()), + batch_size, + ) + .await?; + let output = common::collect(stream).await?; + assert_eq!(expected.num_rows() > 0, keeps_pairs, "filter {filter}"); + assert_streamed_order(&output); + assert_eq!( + prepare_record_batches_for_cmp(&expected.schema(), output), + expected, + "filter {filter} with batch_size {batch_size}" + ); + } + } + Ok(()) +} + +/// The freeze filters pairs from spilled buffered batches once they are read back. +#[tokio::test] +async fn join_inner_early_filtering_with_spilled_buffered_batches() -> Result<()> { + let ((left_batches, left_schema), (right_batches, right_schema)) = + early_filtering_inputs(); + let on = early_filtering_on(&left_schema, &right_schema)?; + let filter = build_c1_lt_c2_filter(&left_schema, &right_schema); + let memory_exec = |batches: &Vec, schema: &SchemaRef| { + TestMemoryExec::try_new_exec( + std::slice::from_ref(batches), + Arc::clone(schema), + None, + ) + }; + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + let session_config = SessionConfig::default().with_batch_size(8192); + let sort_options = vec![SortOptions::default(); on.len()]; + + let join = join_with_filter( + memory_exec(&left_batches, &left_schema)?, + memory_exec(&right_batches, &right_schema)?, + on.clone(), + filter.clone(), + Inner, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(runtime), + ); + let spilled = common::collect(join.execute(0, task_ctx)?).await?; + assert!( + join.metrics().unwrap().spill_count().unwrap() > 0, + "expected the buffered batches to spill" + ); + + let join_no_spill = join_with_filter( + memory_exec(&left_batches, &left_schema)?, + memory_exec(&right_batches, &right_schema)?, + on, + filter, + Inner, + sort_options, + NullEquality::NullEqualsNothing, + )?; + let task_ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let no_spill = common::collect(join_no_spill.execute(0, task_ctx)?).await?; + + assert_streamed_order(&spilled); + assert_eq!( + batches_to_sort_string(&spilled), + batches_to_sort_string(&no_spill) + ); + Ok(()) +}