From 9d916ab4744f2e534b2167d0276094fe4d84bf40 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 9 Nov 2026 15:53:14 -0800 Subject: [PATCH 1/5] bench: add filtered inner sort-merge join cases Adds `sort_merge_join` benchmark cases for Inner joins with a filter: - `inner_1to10_filtered` and `inner_1to100_filtered`, in which every streamed row matches 10 or 100 buffered rows and `left.data < right.data` keeps about half of the pairs. - `inner_1to10_high_pass_filtered`, in which `left.data <> right.data` keeps 90% of the pairs, and its `_narrow` variant over just the `key` and `data` columns, where the filter reads half of the output columns. - `inner_1to1000_selective_filtered`, in which every streamed row matches 1000 buffered rows and `left.data > right.data` keeps about 5% of the 2M pairs. This is the shape of TPC-DS q72's `catalog_sales` to `inventory` join. --- .../physical-plan/benches/sort_merge_join.rs | 128 +++++++++++++++++- 1 file changed, 127 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index c46ad9500b6ba..40ffe4d97903a 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)) @@ -232,6 +261,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; From 58325af3eb8b0be12ac8e4d5feae21fdd2a370ea Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 9 Nov 2026 15:53:14 -0800 Subject: [PATCH 2/5] perf: filter sort-merge inner join pairs before materializing them An Inner sort-merge join with a filter materialized every column of every pair that matched on the join keys, and only then evaluated the filter and dropped the pairs that failed it. When a key matches many rows on each side and the filter keeps few of the pairs, most of the join's time went into copying rows that were then dropped. An Inner join with a filter now evaluates it before materializing: - A run of at least 256 buffered rows is filtered as it is paired with a streamed row. The filter reads the run as a slice of its buffered batch and the streamed row's values repeated, and only the pairs that pass are recorded, in chunks marked as filtered. - The pairs of shorter runs are filtered at freeze on just the columns the filter reads, and the output columns are materialized for the pairs that pass. - When earlier freezes kept most of their pairs and the filter reads most of the output columns, the freeze materializes the batch whole and filters it, as before. Outer joins keep the deferred filtering they need for null-joined rows, and joins without a filter keep the original pairing loop. --- .../sort_merge_join/materializing_stream.rs | 642 +++++++++++++++--- .../src/joins/sort_merge_join/tests.rs | 355 ++++++++++ 2 files changed, 920 insertions(+), 77 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs index 65bc6afcc3609..7030ec3fb5d6e 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, }; 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, ScalarValue, exec_err, + internal_err, }; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::runtime_env::RuntimeEnv; @@ -47,10 +49,63 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use futures::StreamExt; +/// The shortest run of buffered rows that an Inner join with a filter filters as +/// it pairs the run with a streamed row. Each such evaluation has a fixed cost +/// that a long run amortizes. Shorter runs are paired unfiltered, and the +/// freeze filters a whole batch of their pairs at once. +const FILTERED_RUN_MIN_ROWS: usize = 256; + +/// The key that [`StreamedBatch::buffered_batch_idx`] holds while the current chunk +/// holds filtered pairs. No buffered batch has this index, so the next unfiltered +/// pair starts a chunk of its own. +const FILTERED_CHUNK_KEY: Option = Some(usize::MAX); + +/// The matched pairs of one [`StreamedJoinedChunk`], with the indices finished +/// into arrays. +#[derive(Clone)] +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, + /// Whether the pairs already passed the join filter + filtered: bool, +} + +impl MatchedChunk { + fn len(&self) -> usize { + self.streamed_indices.len() + } + + /// Keeps the pairs whose bit is set in `mask`. + fn filter(&self, mask: &BooleanBuffer) -> Result { + let mask = BooleanArray::new(mask.clone(), None); + Ok(Self { + buffered_batch_idx: self.buffered_batch_idx, + streamed_indices: as_uint64_array(&compute::filter( + &self.streamed_indices, + &mask, + )?)? + .clone(), + buffered_indices: as_uint64_array(&compute::filter( + &self.buffered_indices, + &mask, + )?)? + .clone(), + filtered: self.filtered, + }) + } +} + /// Represents a chunk of joined data from streamed and buffered side pub(super) struct StreamedJoinedChunk { /// Index of batch in buffered_data buffered_batch_idx: Option, + /// Whether the join filter already ran on these pairs, so that the chunk + /// holds only the pairs that passed it + filtered: bool, /// Streamed indices, which are never null streamed_indices: Vec, /// Buffered indices. The slot of a null-joined pair holds 0. @@ -74,7 +129,8 @@ pub(super) struct StreamedBatch { pub output_indices: Vec, /// Total number of output rows across all chunks in `output_indices` pub num_output_rows: usize, - /// Index of currently scanned batch from buffered data + /// Index of currently scanned batch from buffered data, or + /// [`FILTERED_CHUNK_KEY`] while the current chunk holds filtered pairs pub buffered_batch_idx: Option, } @@ -181,6 +237,55 @@ impl StreamedBatch { self.num_output_rows += len; } + /// Appends a pair of the current streamed index and each row + /// `buffered_start + i` of the buffered batch with `buffered_batch_idx` index + /// whose bit `i` is set in `passed`, the join filter's result for those pairs. + /// The pairs go to a chunk of filtered pairs, which the freeze does not filter + /// again. Returns the number of pairs appended. + fn append_filtered_pairs( + &mut self, + buffered_batch_idx: usize, + buffered_start: usize, + passed: &BooleanBuffer, + batch_size: usize, + ) -> usize { + let count = passed.count_set_bits(); + if count == 0 { + return 0; + } + debug_assert!( + self.num_output_rows + count <= batch_size, + "appending {count} pairs to {} with batch_size {batch_size}", + self.num_output_rows + ); + let streamed_idx = self.idx as u64; + let reuse = matches!( + self.output_indices.last(), + Some(chunk) if chunk.filtered + && chunk.buffered_batch_idx == Some(buffered_batch_idx) + ); + if !reuse { + self.start_filtered_chunk(buffered_batch_idx, batch_size); + } + let current_chunk = self.output_indices.last_mut().unwrap(); + current_chunk + .streamed_indices + .extend(repeat_n(streamed_idx, count)); + let start = buffered_start as u64; + if count == passed.len() { + current_chunk + .buffered_indices + .extend(start..start + count as u64); + } else { + current_chunk + .buffered_indices + .extend(passed.set_indices().map(|i| start + i as u64)); + } + current_chunk.buffered_nulls.append_n_non_nulls(count); + self.num_output_rows += count; + count + } + /// 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)] @@ -211,12 +316,28 @@ impl StreamedBatch { let capacity = batch_size - self.num_output_rows; self.output_indices.push(StreamedJoinedChunk { buffered_batch_idx, + filtered: false, streamed_indices: Vec::with_capacity(capacity), buffered_indices: Vec::with_capacity(capacity), buffered_nulls: NullBufferBuilder::new(capacity), }); self.buffered_batch_idx = buffered_batch_idx; } + + /// Starts a chunk of filtered pairs for `buffered_batch_idx`. It sets the key + /// that [`Self::chunk_for`] compares to [`FILTERED_CHUNK_KEY`], so that the next + /// unfiltered pair starts a chunk of its own. + fn start_filtered_chunk(&mut self, buffered_batch_idx: usize, batch_size: usize) { + let capacity = batch_size - self.num_output_rows; + self.output_indices.push(StreamedJoinedChunk { + buffered_batch_idx: Some(buffered_batch_idx), + filtered: true, + streamed_indices: Vec::with_capacity(capacity), + buffered_indices: Vec::with_capacity(capacity), + buffered_nulls: NullBufferBuilder::new(capacity), + }); + self.buffered_batch_idx = FILTERED_CHUNK_KEY; + } } /// Per-row filter outcome tracking for full outer joins. @@ -360,6 +481,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 filter runs before any output column is + /// materialized: on a long run of buffered rows as it is paired with a + /// streamed row (see [`Self::pair_filtered_run`]), and at freeze for the + /// remaining pairs (see [`Self::freeze_inner_filtered`]). Outer joins keep + /// a rejected pair for the deferred pipeline, which may null-join its row. + pub early_filtering: bool, + /// The join input whose rows are streamed: the right one for a Right join, + /// the left one otherwise. The join filter's column indices name join inputs. + pub streamed_side: JoinSide, + /// The indices of every buffered column, for materializing all of them + pub buffered_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: usize, + /// How many of the pairs in `freeze_filter_evaluated` passed the filter + pub freeze_filter_passed: usize, /// Target output batch size pub batch_size: usize, @@ -616,12 +754,14 @@ impl MaterializingSortMergeJoinStream { semi/anti/mark joins use BitwiseSortMergeJoinStream" ); let join_time = join_metrics.join_time(); - let buffered_side = if join_type == JoinType::Right { - JoinSide::Left + let (streamed_side, buffered_side) = if join_type == JoinType::Right { + (JoinSide::Right, JoinSide::Left) } else { - JoinSide::Right + (JoinSide::Left, 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 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 +797,11 @@ impl MaterializingSortMergeJoinStream { on_buffered, deferred_filtering, filter_schema, + early_filtering, + streamed_side, + buffered_columns, + freeze_filter_evaluated: 0, + freeze_filter_passed: 0, filter, joined_record_batches: JoinedRecordBatches { joined_batches: new_output_coalescer(Arc::clone(&schema), batch_size), @@ -749,8 +894,14 @@ impl MaterializingSortMergeJoinStream { // batch of pairs accumulates — then advance streamed. // The group stays for the next streamed row. Ordering::Equal => { - while !self.pair_streamed_row_with_group() { - self.freeze_and_emit(emitter).await?; + if self.early_filtering { + while !self.pair_streamed_row_with_group_filtered()? { + self.freeze_and_emit(emitter).await?; + } + } else { + while !self.pair_streamed_row_with_group() { + self.freeze_and_emit(emitter).await?; + } } if !self.try_advance_streamed_row() { self.load_next_streamed_batch().await?; @@ -805,6 +956,91 @@ impl MaterializingSortMergeJoinStream { true } + /// [`Self::pair_streamed_row_with_group`] for an Inner join with a filter, + /// which keeps only the pairs that pass it. For a run of at least + /// [`FILTERED_RUN_MIN_ROWS`] buffered rows, the filter runs as the run is + /// paired, and the pairs it rejects are never recorded. + fn pair_streamed_row_with_group_filtered(&mut self) -> Result { + 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); + let filtered = if len >= FILTERED_RUN_MIN_ROWS { + self.pair_filtered_run(rows.start, len)? + } else { + None + }; + match filtered { + Some(appended) => room -= appended, + None => { + self.streamed_batch.append_output_pairs( + self.buffered_data.scanning_batch_idx, + rows.start, + len, + self.batch_size, + ); + room -= len; + } + } + self.buffered_data.scanning_advance_by(len); + } + if room == 0 { + return Ok(false); + } + + self.buffered_group_matched = true; + self.buffered_data.scanning_reset(); + Ok(true) + } + + /// Pairs the current streamed row with the `len` rows of the scanning + /// buffered batch from `buffered_start` on, keeping only the pairs that pass + /// the join filter. The filter reads the buffered rows as a slice of their + /// batch and the streamed row as a single value repeated, so no index is + /// recorded and no column is gathered for a pair it rejects. + /// + /// Returns the number of pairs appended, or `None` without pairing anything + /// when the scanning batch is spilled; the caller then pairs every row and + /// the freeze filters them once the batch is read back. + fn pair_filtered_run( + &mut self, + buffered_start: usize, + len: usize, + ) -> Result> { + let Some(filter) = &self.filter else { + return internal_err!("early filtering requires a join filter"); + }; + let BufferedBatchState::InMemory(buffered) = + &self.buffered_data.scanning_batch().batch + else { + return Ok(None); + }; + let streamed = &self.streamed_batch.batch; + let streamed_idx = self.streamed_batch.idx; + let columns = filter + .column_indices() + .iter() + .map(|column_index| match column_index.side { + JoinSide::None => internal_err!("Unexpected JoinSide::None in filter"), + side if side == self.streamed_side => ScalarValue::try_from_array( + streamed.column(column_index.index), + streamed_idx, + )? + .to_array_of_size(len), + _ => Ok(buffered + .column(column_index.index) + .slice(buffered_start, len)), + }) + .collect::>>()?; + let passed = evaluate_join_filter(filter, filter.schema(), columns, len)?; + Ok(Some(self.streamed_batch.append_filtered_pairs( + self.buffered_data.scanning_batch_idx, + buffered_start, + &passed, + self.batch_size, + ))) + } + /// `Less` (outer joins): no buffered row matches the current streamed /// row — emit it joined to NULLs. Inner joins emit nothing. fn null_join_streamed_row(&mut self) { @@ -1520,6 +1756,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 +1793,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 +1812,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 +1855,12 @@ 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, + filtered: chunk.filtered, + }); } if !matched_chunks.is_empty() { @@ -1639,7 +1878,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 +1886,39 @@ 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); + } + debug_assert!( + matched_chunks.iter().all(|chunk| !chunk.filtered), + "only early filtering pairs filtered chunks" + ); + + 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 chunk_refs: Vec<&MatchedChunk> = matched_chunks.iter().collect(); + let right_columns = self.materialize_right_columns( + &chunk_refs, + 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,7 +1938,10 @@ 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), @@ -1706,8 +1953,7 @@ impl MaterializingSortMergeJoinStream { .evaluate(&filter_batch)? .into_array(filter_batch.num_rows())?; - let filter_result_mask = - datafusion_common::cast::as_boolean_array(&filter_result)?; + let filter_result_mask = as_boolean_array(&filter_result)?; // Convert NULL filter results to false — NULL means "not satisfied" // per SQL semantics, same as Left/Right outer joins. @@ -1717,27 +1963,23 @@ impl MaterializingSortMergeJoinStream { filter_result_mask.clone() }; - 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 +2013,197 @@ 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, for the pairs of the chunks that + /// [`Self::pair_filtered_run`] did not filter already, 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 no chunk was filtered already and earlier freezes + /// kept enough of their pairs, [`Self::filter_first_pays`] says the opposite, + /// and 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" + ); + + // Evaluate the filter over the pairs of the chunks not filtered yet. + let unfiltered: Vec<&MatchedChunk> = + matched_chunks.iter().filter(|c| !c.filtered).collect(); + if unfiltered.len() == matched_chunks.len() && !self.filter_first_pays(filter) { + return self.freeze_inner_filtered_whole(matched_chunks, total_matched_rows); + } + let passed = if unfiltered.is_empty() { + None + } else { + let num_rows = unfiltered.iter().map(|c| c.len()).sum(); + let streamed_indices = + concat_indices(unfiltered.iter().map(|c| &c.streamed_indices))?; + let buffered_filter_columns: Vec = filter + .column_indices() + .iter() + .filter(|c| c.side != self.streamed_side && c.side != JoinSide::None) + .map(|c| c.index) + .collect(); + let mut buffered_arrays = self + .materialize_right_columns( + &unfiltered, + num_rows, + &buffered_filter_columns, + )? + .into_iter(); + let columns = filter + .column_indices() + .iter() + .map(|column_index| match column_index.side { + JoinSide::None => { + internal_err!("Unexpected JoinSide::None in filter") + } + side if side == self.streamed_side => take_column( + self.streamed_batch.batch.column(column_index.index), + &streamed_indices, + ), + _ => buffered_arrays.next().ok_or_else(|| { + DataFusionError::Internal( + "missing buffered filter column".to_string(), + ) + }), + }) + .collect::>>()?; + let passed = + evaluate_join_filter(filter, filter.schema(), columns, num_rows)?; + self.freeze_filter_evaluated += num_rows; + self.freeze_filter_passed += passed.count_set_bits(); + Some(passed) + }; + + // Keep the pairs that passed, in order. + let mut kept: Vec = Vec::with_capacity(matched_chunks.len()); + let mut kept_rows = 0; + let mut offset = 0; + for chunk in matched_chunks { + let len = chunk.len(); + match (&passed, chunk.filtered) { + (Some(passed), false) => { + let mask = passed.slice(offset, len); + offset += len; + let count = mask.count_set_bits(); + if count == len { + kept.push(chunk.clone()); + } else if count > 0 { + kept.push(chunk.filter(&mask)?); + } + kept_rows += count; + } + _ => { + kept.push(chunk.clone()); + kept_rows += len; + } + } + } + if kept_rows == 0 { + return Ok(()); + } + + let streamed_indices = concat_indices(kept.iter().map(|c| &c.streamed_indices))?; + let streamed_columns = + materialize_left_columns(&self.streamed_batch.batch, &streamed_indices)?; + let kept_refs: Vec<&MatchedChunk> = kept.iter().collect(); + let buffered_columns = self.materialize_right_columns( + &kept_refs, + 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. None of the pairs was filtered already. + fn freeze_inner_filtered_whole( + &mut self, + matched_chunks: &[MatchedChunk], + total_matched_rows: usize, + ) -> Result<()> { + let streamed_indices = + concat_indices(matched_chunks.iter().map(|c| &c.streamed_indices))?; + let streamed_columns = + materialize_left_columns(&self.streamed_batch.batch, &streamed_indices)?; + let chunk_refs: Vec<&MatchedChunk> = matched_chunks.iter().collect(); + let buffered_columns = self.materialize_right_columns( + &chunk_refs, + 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; + self.freeze_filter_passed += kept_rows; + 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. + 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 +2211,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 +2241,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 +2298,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 +2314,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 { + 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 +2339,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 +2350,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 +2363,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 +2495,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 +2547,18 @@ 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 +/// 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 +2566,21 @@ 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: Vec = columns + .iter() + .map(|&c| Arc::clone(batch.column(c))) + .collect(); if let Some(range) = is_contiguous_range(buffered_indices) { - Ok(batch.slice(range.start, range.len()).columns().to_vec()) + Ok(arrays + .iter() + .map(|a| a.slice(range.start, range.len())) + .collect()) } else { - Ok(take_arrays(batch.columns(), buffered_indices, None)?) + Ok(take_arrays(&arrays, buffered_indices, None)?) } } BufferedBatchState::Spilled(_) => { @@ -2147,6 +2591,50 @@ fn fetch_right_columns_from_batch_by_idxs( } } +/// Concatenates the index arrays of several chunks into one. +fn concat_indices<'a>( + indices: impl ExactSizeIterator, +) -> 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..459b9cdafe408 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -8395,3 +8395,358 @@ 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, long enough to be filtered as it is paired, and the 250-row +/// buffered batches split each such group. The other keys have groups of one or +/// two buffered rows, whose pairs are filtered at freeze. 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 `b2 <> b1 OR c2 < c1`, which nearly every pair passes, with its +/// intermediate columns in the order `b2, b1, c2, c1`. It reads four of the six +/// output columns, so after the first freeze the join 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("b2", 0)), + Operator::NotEq, + Arc::new(Column::new("b1", 1)), + )), + Operator::Or, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c2", 2)), + Operator::Lt, + Arc::new(Column::new("c1", 3)), + )), + )); + 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 key column of `batches` is ascending, the order of the +/// streamed side that the join maintains. +fn assert_streamed_order(batches: &[RecordBatch]) { + let keys: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert!( + keys.is_sorted(), + "Inner join output must stay ordered on the streamed side" + ); +} + +/// An Inner join with a filter filters long runs of buffered rows as it pairs them +/// and the pairs of shorter runs at freeze. It must keep the same pairs as a hash +/// join whatever the batch size, which decides where a batch of pairs cuts a run, +/// 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. The +/// high-pass filter also covers freezes that materialize whole batches. +#[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)?; + let filters = [ + build_c1_lt_c2_filter(&left_schema, &right_schema), + interleaved_early_filter(&left_schema, &right_schema), + high_pass_early_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, + ) + }; + + for filter 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!(expected.num_rows() > 0, "the filter must keep some pairs"); + assert_streamed_order(&output); + assert_eq!( + prepare_record_batches_for_cmp(output), + expected, + "filter {filter} with batch_size {batch_size}" + ); + } + } + Ok(()) +} + +/// A filter that reads no columns still runs for every pair, both on a long run +/// as it is paired and at freeze. +#[tokio::test] +async fn join_inner_early_filtering_column_free_filter() -> Result<()> { + let ((left_batches, left_schema), (right_batches, right_schema)) = + early_filtering_inputs(); + let on = early_filtering_on(&left_schema, &right_schema)?; + let left: Arc = TestMemoryExec::try_new_exec( + std::slice::from_ref(&left_batches), + Arc::clone(&left_schema), + None, + )?; + let right: Arc = TestMemoryExec::try_new_exec( + std::slice::from_ref(&right_batches), + Arc::clone(&right_schema), + None, + )?; + + let (_, unfiltered) = + join_collect(Arc::clone(&left), Arc::clone(&right), on.clone(), Inner).await?; + let (_, kept) = join_collect_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + column_free_early_filter(true), + Inner, + ) + .await?; + assert_eq!( + batches_to_sort_string(&kept), + batches_to_sort_string(&unfiltered) + ); + + let (_, dropped) = + join_collect_with_filter(left, right, on, column_free_early_filter(false), Inner) + .await?; + assert_eq!(dropped.iter().map(|b| b.num_rows()).sum::(), 0); + Ok(()) +} + +/// A long run in a spilled buffered batch is paired unfiltered, and its pairs are +/// filtered at freeze once the batch is 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(()) +} From 378d5250b3058242dd5672126d48a0e4e9ca4cdd Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 9 Nov 2026 16:40:20 -0800 Subject: [PATCH 3/5] refactor: simplify sort-merge join early filtering - Drop the `streamed_side` field: only an Inner join filters early, so the streamed side is always the left input. - Compute the buffered columns the filter reads once, not on every freeze. - Repeat the streamed row with `take` instead of a `ScalarValue` round trip. - Make `materialize_right_columns` generic over `Borrow`, so callers pass chunks without collecting a `Vec<&MatchedChunk>` and kept chunks are borrowed rather than cloned. - Reuse `evaluate_join_filter` in the outer join path. - Build one `FilterPredicate` for both index arrays in `MatchedChunk::filter`. - Fold the column-free filter test into the hash join comparison test. --- .../sort_merge_join/materializing_stream.rs | 185 ++++++++---------- .../src/joins/sort_merge_join/tests.rs | 56 +----- 2 files changed, 93 insertions(+), 148 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs index 7030ec3fb5d6e..8f0b5b1cef719 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 @@ -9,6 +9,7 @@ //! *streamed* side and the *buffered* side), compares join keys, and //! produces joined `RecordBatch`es. +use std::borrow::{Borrow, Cow}; use std::cmp::Ordering; use std::collections::VecDeque; use std::fmt::Debug; @@ -31,16 +32,16 @@ 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_boolean_array, as_uint64_array}; use datafusion_common::instant::Instant; use datafusion_common::{ - DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue, 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; @@ -81,18 +82,18 @@ impl MatchedChunk { /// Keeps the pairs whose bit is set in `mask`. fn filter(&self, mask: &BooleanBuffer) -> Result { - let mask = BooleanArray::new(mask.clone(), None); + let predicate = FilterBuilder::new(&BooleanArray::new(mask.clone(), None)) + .optimize() + .build(); Ok(Self { buffered_batch_idx: self.buffered_batch_idx, - streamed_indices: as_uint64_array(&compute::filter( - &self.streamed_indices, - &mask, - )?)? + streamed_indices: as_uint64_array( + &predicate.filter(&self.streamed_indices)?, + )? .clone(), - buffered_indices: as_uint64_array(&compute::filter( - &self.buffered_indices, - &mask, - )?)? + buffered_indices: as_uint64_array( + &predicate.filter(&self.buffered_indices)?, + )? .clone(), filtered: self.filtered, }) @@ -487,12 +488,13 @@ pub(super) struct MaterializingSortMergeJoinStream { /// streamed row (see [`Self::pair_filtered_run`]), and at freeze for the /// remaining pairs (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 join input whose rows are streamed: the right one for a Right join, - /// the left one otherwise. The join filter's column indices name join inputs. - pub streamed_side: JoinSide, /// 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: usize, @@ -754,14 +756,23 @@ impl MaterializingSortMergeJoinStream { semi/anti/mark joins use BitwiseSortMergeJoinStream" ); let join_time = join_metrics.join_time(); - let (streamed_side, buffered_side) = if join_type == JoinType::Right { - (JoinSide::Right, JoinSide::Left) + let buffered_side = if join_type == JoinType::Right { + JoinSide::Left } else { - (JoinSide::Left, JoinSide::Right) + 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 @@ -798,8 +809,8 @@ impl MaterializingSortMergeJoinStream { deferred_filtering, filter_schema, early_filtering, - streamed_side, buffered_columns, + buffered_filter_columns, freeze_filter_evaluated: 0, freeze_filter_passed: 0, filter, @@ -1016,20 +1027,20 @@ impl MaterializingSortMergeJoinStream { return Ok(None); }; let streamed = &self.streamed_batch.batch; - let streamed_idx = self.streamed_batch.idx; + let streamed_idx = UInt64Array::from_value(self.streamed_batch.idx as u64, len); let columns = filter .column_indices() .iter() .map(|column_index| match column_index.side { - JoinSide::None => internal_err!("Unexpected JoinSide::None in filter"), - side if side == self.streamed_side => ScalarValue::try_from_array( + JoinSide::Left => Ok(compute::take( streamed.column(column_index.index), - streamed_idx, - )? - .to_array_of_size(len), - _ => Ok(buffered + &streamed_idx, + None, + )?), + JoinSide::Right => Ok(buffered .column(column_index.index) .slice(buffered_start, len)), + JoinSide::None => internal_err!("Unexpected JoinSide::None in filter"), }) .collect::>>()?; let passed = evaluate_join_filter(filter, filter.schema(), columns, len)?; @@ -1913,9 +1924,8 @@ impl MaterializingSortMergeJoinStream { let left_columns = materialize_left_columns(&self.streamed_batch.batch, &combined_left_indices)?; - let chunk_refs: Vec<&MatchedChunk> = matched_chunks.iter().collect(); let right_columns = self.materialize_right_columns( - &chunk_refs, + matched_chunks, total_matched_rows, &self.buffered_columns, )?; @@ -1943,25 +1953,10 @@ impl MaterializingSortMergeJoinStream { 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 = 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, + ); self.joined_record_batches.push_batch_with_filter_metadata( output_batch, @@ -2048,46 +2043,40 @@ impl MaterializingSortMergeJoinStream { let num_rows = unfiltered.iter().map(|c| c.len()).sum(); let streamed_indices = concat_indices(unfiltered.iter().map(|c| &c.streamed_indices))?; - let buffered_filter_columns: Vec = filter - .column_indices() - .iter() - .filter(|c| c.side != self.streamed_side && c.side != JoinSide::None) - .map(|c| c.index) - .collect(); let mut buffered_arrays = self .materialize_right_columns( &unfiltered, num_rows, - &buffered_filter_columns, + &self.buffered_filter_columns, )? .into_iter(); let columns = filter .column_indices() .iter() .map(|column_index| match column_index.side { - JoinSide::None => { - internal_err!("Unexpected JoinSide::None in filter") - } - side if side == self.streamed_side => take_column( + JoinSide::Left => take_column( self.streamed_batch.batch.column(column_index.index), &streamed_indices, ), - _ => buffered_arrays.next().ok_or_else(|| { - DataFusionError::Internal( - "missing buffered filter column".to_string(), - ) + 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, num_rows)?; self.freeze_filter_evaluated += num_rows; - self.freeze_filter_passed += passed.count_set_bits(); - Some(passed) + Some(evaluate_join_filter( + filter, + filter.schema(), + columns, + num_rows, + )?) }; // Keep the pairs that passed, in order. - let mut kept: Vec = Vec::with_capacity(matched_chunks.len()); + let mut kept: Vec> = Vec::with_capacity(matched_chunks.len()); let mut kept_rows = 0; let mut offset = 0; for chunk in matched_chunks { @@ -2098,14 +2087,15 @@ impl MaterializingSortMergeJoinStream { offset += len; let count = mask.count_set_bits(); if count == len { - kept.push(chunk.clone()); + kept.push(Cow::Borrowed(chunk)); } else if count > 0 { - kept.push(chunk.filter(&mask)?); + kept.push(Cow::Owned(chunk.filter(&mask)?)); } + self.freeze_filter_passed += count; kept_rows += count; } _ => { - kept.push(chunk.clone()); + kept.push(Cow::Borrowed(chunk)); kept_rows += len; } } @@ -2117,12 +2107,8 @@ impl MaterializingSortMergeJoinStream { let streamed_indices = concat_indices(kept.iter().map(|c| &c.streamed_indices))?; let streamed_columns = materialize_left_columns(&self.streamed_batch.batch, &streamed_indices)?; - let kept_refs: Vec<&MatchedChunk> = kept.iter().collect(); - let buffered_columns = self.materialize_right_columns( - &kept_refs, - kept_rows, - &self.buffered_columns, - )?; + 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(), @@ -2144,9 +2130,8 @@ impl MaterializingSortMergeJoinStream { concat_indices(matched_chunks.iter().map(|c| &c.streamed_indices))?; let streamed_columns = materialize_left_columns(&self.streamed_batch.batch, &streamed_indices)?; - let chunk_refs: Vec<&MatchedChunk> = matched_chunks.iter().collect(); let buffered_columns = self.materialize_right_columns( - &chunk_refs, + matched_chunks, total_matched_rows, &self.buffered_columns, )?; @@ -2209,20 +2194,23 @@ impl MaterializingSortMergeJoinStream { /// for a single fetch. When multiple batches are involved, `interleave` /// gathers columns across sources. A null-row sentinel at source index 0 /// handles null right indices (unmatched streamed rows). - fn materialize_right_columns( + fn materialize_right_columns>( &self, - matched_chunks: &[&MatchedChunk], + matched_chunks: &[C], total_matched_rows: usize, columns: &[usize], ) -> Result> { - 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); + let chunks = || { + matched_chunks + .iter() + .map(|c| -> &MatchedChunk { c.borrow() }) + }; + let first_batch_idx = matched_chunks[0].borrow().buffered_batch_idx; + let single_source = chunks().all(|c| c.buffered_batch_idx == first_batch_idx); if single_source { let combined_right_indices = - concat_indices(matched_chunks.iter().map(|c| &c.buffered_indices))?; + concat_indices(chunks().map(|c| &c.buffered_indices))?; return fetch_right_columns_by_idxs( &self.buffered_data, @@ -2239,9 +2227,7 @@ impl MaterializingSortMergeJoinStream { // `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(|c| c.buffered_indices.null_count() > 0); + let needs_null_sentinel = chunks().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 @@ -2296,8 +2282,7 @@ impl MaterializingSortMergeJoinStream { // (`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() + let (min_batch_idx, max_batch_idx) = chunks() .fold((usize::MAX, 0usize), |(lo, hi), c| { (lo.min(c.buffered_batch_idx), hi.max(c.buffered_batch_idx)) }); @@ -2314,7 +2299,7 @@ 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 chunk in matched_chunks { + for chunk in chunks() { let batch_idx = chunk.buffered_batch_idx; let right = &chunk.buffered_indices; let slot = &mut source_of_batch[batch_idx - min_batch_idx]; @@ -2570,17 +2555,13 @@ fn fetch_right_columns_from_batch_by_idxs( ) -> Result> { match &buffered_batch.batch { BufferedBatchState::InMemory(batch) => { - let arrays: Vec = columns - .iter() - .map(|&c| Arc::clone(batch.column(c))) - .collect(); + let arrays = columns.iter().map(|&c| batch.column(c)); if let Some(range) = is_contiguous_range(buffered_indices) { - Ok(arrays - .iter() - .map(|a| a.slice(range.start, range.len())) - .collect()) + Ok(arrays.map(|a| a.slice(range.start, range.len())).collect()) } else { - Ok(take_arrays(&arrays, buffered_indices, None)?) + arrays + .map(|a| Ok(compute::take(a.as_ref(), buffered_indices, None)?)) + .collect() } } BufferedBatchState::Spilled(_) => { @@ -2593,7 +2574,7 @@ fn fetch_right_columns_from_batch_by_idxs( /// Concatenates the index arrays of several chunks into one. fn concat_indices<'a>( - indices: impl ExactSizeIterator, + indices: impl Iterator, ) -> Result { let arrays: Vec<&dyn Array> = indices.map(|a| a as &dyn Array).collect(); match arrays.as_slice() { 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 459b9cdafe408..062a2b2cf1342 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -8602,16 +8602,20 @@ fn assert_streamed_order(batches: &[RecordBatch]) { /// join whatever the batch size, which decides where a batch of pairs cuts a run, /// 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. The -/// high-pass filter also covers freezes that materialize whole batches. +/// high-pass filter also covers freezes that materialize whole batches, and 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), - interleaved_early_filter(&left_schema, &right_schema), - high_pass_early_filter(&left_schema, &right_schema), + (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( @@ -8621,7 +8625,7 @@ async fn join_inner_early_filtering_matches_hash_join() -> Result<()> { ) }; - for filter in filters { + 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)?, @@ -8635,7 +8639,7 @@ async fn join_inner_early_filtering_matches_hash_join() -> Result<()> { ) .await?; let output = common::collect(stream).await?; - assert!(expected.num_rows() > 0, "the filter must keep some pairs"); + assert_eq!(expected.num_rows() > 0, keeps_pairs, "filter {filter}"); assert_streamed_order(&output); assert_eq!( prepare_record_batches_for_cmp(output), @@ -8647,46 +8651,6 @@ async fn join_inner_early_filtering_matches_hash_join() -> Result<()> { Ok(()) } -/// A filter that reads no columns still runs for every pair, both on a long run -/// as it is paired and at freeze. -#[tokio::test] -async fn join_inner_early_filtering_column_free_filter() -> Result<()> { - let ((left_batches, left_schema), (right_batches, right_schema)) = - early_filtering_inputs(); - let on = early_filtering_on(&left_schema, &right_schema)?; - let left: Arc = TestMemoryExec::try_new_exec( - std::slice::from_ref(&left_batches), - Arc::clone(&left_schema), - None, - )?; - let right: Arc = TestMemoryExec::try_new_exec( - std::slice::from_ref(&right_batches), - Arc::clone(&right_schema), - None, - )?; - - let (_, unfiltered) = - join_collect(Arc::clone(&left), Arc::clone(&right), on.clone(), Inner).await?; - let (_, kept) = join_collect_with_filter( - Arc::clone(&left), - Arc::clone(&right), - on.clone(), - column_free_early_filter(true), - Inner, - ) - .await?; - assert_eq!( - batches_to_sort_string(&kept), - batches_to_sort_string(&unfiltered) - ); - - let (_, dropped) = - join_collect_with_filter(left, right, on, column_free_early_filter(false), Inner) - .await?; - assert_eq!(dropped.iter().map(|b| b.num_rows()).sum::(), 0); - Ok(()) -} - /// A long run in a spilled buffered batch is paired unfiltered, and its pairs are /// filtered at freeze once the batch is read back. #[tokio::test] From e2f465717a78d3d935f40d5b07a30def18cf756c Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 9 Nov 2026 17:24:12 -0800 Subject: [PATCH 4/5] test: handle empty join output in prepare_record_batches_for_cmp The always-false column-free filter case makes the hash join oracle return no batches, so take the schema as a parameter instead of reading it from the first batch. --- .../src/joins/sort_merge_join/tests.rs | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index 062a2b2cf1342..14bc8a5d0be84 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()); } @@ -8642,7 +8646,7 @@ async fn join_inner_early_filtering_matches_hash_join() -> Result<()> { assert_eq!(expected.num_rows() > 0, keeps_pairs, "filter {filter}"); assert_streamed_order(&output); assert_eq!( - prepare_record_batches_for_cmp(output), + prepare_record_batches_for_cmp(&expected.schema(), output), expected, "filter {filter} with batch_size {batch_size}" ); From 8012102c4a0880b1de6e96f551332dceeb9115b1 Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 11 Nov 2026 08:45:18 -0800 Subject: [PATCH 5/5] refactor: filter sort-merge inner join pairs at freeze only Address review feedback: - Drop filtering while pairing. It wins only for selective filters on runs of 512+ rows and regresses narrow high-pass joins. This removes the run threshold, `FILTERED_CHUNK_KEY`, the `filtered` flag, `Cow` and the spill fallback. - Take the whole-batch path when both sides' indices span a range, as in a 1:1 join, where the columns materialize as slices. - Filter the concatenated index arrays once and slice them per chunk. - Count the cost model's pairs in `u64`, and document its limits. - Assert the order within a key, and make the high-pass test filter reject some pairs so that the whole-batch filter is covered. - Add a filtered 1:1 Inner join benchmark. --- .../physical-plan/benches/sort_merge_join.rs | 15 + .../sort_merge_join/materializing_stream.rs | 423 +++++------------- .../src/joins/sort_merge_join/tests.rs | 70 +-- 3 files changed, 168 insertions(+), 340 deletions(-) diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index 40ffe4d97903a..0604f83603e9f 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -203,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 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 8f0b5b1cef719..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 @@ -9,7 +9,6 @@ //! *streamed* side and the *buffered* side), compares join keys, and //! produces joined `RecordBatch`es. -use std::borrow::{Borrow, Cow}; use std::cmp::Ordering; use std::collections::VecDeque; use std::fmt::Debug; @@ -50,20 +49,8 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use futures::StreamExt; -/// The shortest run of buffered rows that an Inner join with a filter filters as -/// it pairs the run with a streamed row. Each such evaluation has a fixed cost -/// that a long run amortizes. Shorter runs are paired unfiltered, and the -/// freeze filters a whole batch of their pairs at once. -const FILTERED_RUN_MIN_ROWS: usize = 256; - -/// The key that [`StreamedBatch::buffered_batch_idx`] holds while the current chunk -/// holds filtered pairs. No buffered batch has this index, so the next unfiltered -/// pair starts a chunk of its own. -const FILTERED_CHUNK_KEY: Option = Some(usize::MAX); - /// The matched pairs of one [`StreamedJoinedChunk`], with the indices finished /// into arrays. -#[derive(Clone)] struct MatchedChunk { /// Index of the batch in `buffered_data` that the buffered rows come from buffered_batch_idx: usize, @@ -71,42 +58,18 @@ struct MatchedChunk { streamed_indices: UInt64Array, /// The buffered row of each pair buffered_indices: UInt64Array, - /// Whether the pairs already passed the join filter - filtered: bool, } impl MatchedChunk { fn len(&self) -> usize { self.streamed_indices.len() } - - /// Keeps the pairs whose bit is set in `mask`. - fn filter(&self, mask: &BooleanBuffer) -> Result { - let predicate = FilterBuilder::new(&BooleanArray::new(mask.clone(), None)) - .optimize() - .build(); - Ok(Self { - buffered_batch_idx: self.buffered_batch_idx, - streamed_indices: as_uint64_array( - &predicate.filter(&self.streamed_indices)?, - )? - .clone(), - buffered_indices: as_uint64_array( - &predicate.filter(&self.buffered_indices)?, - )? - .clone(), - filtered: self.filtered, - }) - } } /// Represents a chunk of joined data from streamed and buffered side pub(super) struct StreamedJoinedChunk { /// Index of batch in buffered_data buffered_batch_idx: Option, - /// Whether the join filter already ran on these pairs, so that the chunk - /// holds only the pairs that passed it - filtered: bool, /// Streamed indices, which are never null streamed_indices: Vec, /// Buffered indices. The slot of a null-joined pair holds 0. @@ -130,8 +93,7 @@ pub(super) struct StreamedBatch { pub output_indices: Vec, /// Total number of output rows across all chunks in `output_indices` pub num_output_rows: usize, - /// Index of currently scanned batch from buffered data, or - /// [`FILTERED_CHUNK_KEY`] while the current chunk holds filtered pairs + /// Index of currently scanned batch from buffered data pub buffered_batch_idx: Option, } @@ -238,55 +200,6 @@ impl StreamedBatch { self.num_output_rows += len; } - /// Appends a pair of the current streamed index and each row - /// `buffered_start + i` of the buffered batch with `buffered_batch_idx` index - /// whose bit `i` is set in `passed`, the join filter's result for those pairs. - /// The pairs go to a chunk of filtered pairs, which the freeze does not filter - /// again. Returns the number of pairs appended. - fn append_filtered_pairs( - &mut self, - buffered_batch_idx: usize, - buffered_start: usize, - passed: &BooleanBuffer, - batch_size: usize, - ) -> usize { - let count = passed.count_set_bits(); - if count == 0 { - return 0; - } - debug_assert!( - self.num_output_rows + count <= batch_size, - "appending {count} pairs to {} with batch_size {batch_size}", - self.num_output_rows - ); - let streamed_idx = self.idx as u64; - let reuse = matches!( - self.output_indices.last(), - Some(chunk) if chunk.filtered - && chunk.buffered_batch_idx == Some(buffered_batch_idx) - ); - if !reuse { - self.start_filtered_chunk(buffered_batch_idx, batch_size); - } - let current_chunk = self.output_indices.last_mut().unwrap(); - current_chunk - .streamed_indices - .extend(repeat_n(streamed_idx, count)); - let start = buffered_start as u64; - if count == passed.len() { - current_chunk - .buffered_indices - .extend(start..start + count as u64); - } else { - current_chunk - .buffered_indices - .extend(passed.set_indices().map(|i| start + i as u64)); - } - current_chunk.buffered_nulls.append_n_non_nulls(count); - self.num_output_rows += count; - count - } - /// 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)] @@ -317,28 +230,12 @@ impl StreamedBatch { let capacity = batch_size - self.num_output_rows; self.output_indices.push(StreamedJoinedChunk { buffered_batch_idx, - filtered: false, streamed_indices: Vec::with_capacity(capacity), buffered_indices: Vec::with_capacity(capacity), buffered_nulls: NullBufferBuilder::new(capacity), }); self.buffered_batch_idx = buffered_batch_idx; } - - /// Starts a chunk of filtered pairs for `buffered_batch_idx`. It sets the key - /// that [`Self::chunk_for`] compares to [`FILTERED_CHUNK_KEY`], so that the next - /// unfiltered pair starts a chunk of its own. - fn start_filtered_chunk(&mut self, buffered_batch_idx: usize, batch_size: usize) { - let capacity = batch_size - self.num_output_rows; - self.output_indices.push(StreamedJoinedChunk { - buffered_batch_idx: Some(buffered_batch_idx), - filtered: true, - streamed_indices: Vec::with_capacity(capacity), - buffered_indices: Vec::with_capacity(capacity), - buffered_nulls: NullBufferBuilder::new(capacity), - }); - self.buffered_batch_idx = FILTERED_CHUNK_KEY; - } } /// Per-row filter outcome tracking for full outer joins. @@ -483,12 +380,11 @@ pub(super) struct MaterializingSortMergeJoinStream { /// `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 filter runs before any output column is - /// materialized: on a long run of buffered rows as it is paired with a - /// streamed row (see [`Self::pair_filtered_run`]), and at freeze for the - /// remaining pairs (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. + /// 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, @@ -497,9 +393,9 @@ pub(super) struct MaterializingSortMergeJoinStream { 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: usize, + pub freeze_filter_evaluated: u64, /// How many of the pairs in `freeze_filter_evaluated` passed the filter - pub freeze_filter_passed: usize, + pub freeze_filter_passed: u64, /// Target output batch size pub batch_size: usize, @@ -905,14 +801,8 @@ impl MaterializingSortMergeJoinStream { // batch of pairs accumulates — then advance streamed. // The group stays for the next streamed row. Ordering::Equal => { - if self.early_filtering { - while !self.pair_streamed_row_with_group_filtered()? { - self.freeze_and_emit(emitter).await?; - } - } else { - while !self.pair_streamed_row_with_group() { - self.freeze_and_emit(emitter).await?; - } + while !self.pair_streamed_row_with_group() { + self.freeze_and_emit(emitter).await?; } if !self.try_advance_streamed_row() { self.load_next_streamed_batch().await?; @@ -967,91 +857,6 @@ impl MaterializingSortMergeJoinStream { true } - /// [`Self::pair_streamed_row_with_group`] for an Inner join with a filter, - /// which keeps only the pairs that pass it. For a run of at least - /// [`FILTERED_RUN_MIN_ROWS`] buffered rows, the filter runs as the run is - /// paired, and the pairs it rejects are never recorded. - fn pair_streamed_row_with_group_filtered(&mut self) -> Result { - 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); - let filtered = if len >= FILTERED_RUN_MIN_ROWS { - self.pair_filtered_run(rows.start, len)? - } else { - None - }; - match filtered { - Some(appended) => room -= appended, - None => { - self.streamed_batch.append_output_pairs( - self.buffered_data.scanning_batch_idx, - rows.start, - len, - self.batch_size, - ); - room -= len; - } - } - self.buffered_data.scanning_advance_by(len); - } - if room == 0 { - return Ok(false); - } - - self.buffered_group_matched = true; - self.buffered_data.scanning_reset(); - Ok(true) - } - - /// Pairs the current streamed row with the `len` rows of the scanning - /// buffered batch from `buffered_start` on, keeping only the pairs that pass - /// the join filter. The filter reads the buffered rows as a slice of their - /// batch and the streamed row as a single value repeated, so no index is - /// recorded and no column is gathered for a pair it rejects. - /// - /// Returns the number of pairs appended, or `None` without pairing anything - /// when the scanning batch is spilled; the caller then pairs every row and - /// the freeze filters them once the batch is read back. - fn pair_filtered_run( - &mut self, - buffered_start: usize, - len: usize, - ) -> Result> { - let Some(filter) = &self.filter else { - return internal_err!("early filtering requires a join filter"); - }; - let BufferedBatchState::InMemory(buffered) = - &self.buffered_data.scanning_batch().batch - else { - return Ok(None); - }; - let streamed = &self.streamed_batch.batch; - let streamed_idx = UInt64Array::from_value(self.streamed_batch.idx as u64, len); - let columns = filter - .column_indices() - .iter() - .map(|column_index| match column_index.side { - JoinSide::Left => Ok(compute::take( - streamed.column(column_index.index), - &streamed_idx, - None, - )?), - JoinSide::Right => Ok(buffered - .column(column_index.index) - .slice(buffered_start, len)), - JoinSide::None => internal_err!("Unexpected JoinSide::None in filter"), - }) - .collect::>>()?; - let passed = evaluate_join_filter(filter, filter.schema(), columns, len)?; - Ok(Some(self.streamed_batch.append_filtered_pairs( - self.buffered_data.scanning_batch_idx, - buffered_start, - &passed, - self.batch_size, - ))) - } - /// `Less` (outer joins): no buffered row matches the current streamed /// row — emit it joined to NULLs. Inner joins emit nothing. fn null_join_streamed_row(&mut self) { @@ -1870,7 +1675,6 @@ impl MaterializingSortMergeJoinStream { buffered_batch_idx: chunk.buffered_batch_idx.unwrap(), streamed_indices: left_indices, buffered_indices: right_indices, - filtered: chunk.filtered, }); } @@ -1913,10 +1717,6 @@ impl MaterializingSortMergeJoinStream { if self.early_filtering { return self.freeze_inner_filtered(matched_chunks, total_matched_rows); } - debug_assert!( - matched_chunks.iter().all(|chunk| !chunk.filtered), - "only early filtering pairs filtered chunks" - ); let combined_left_indices = concat_indices(matched_chunks.iter().map(|c| &c.streamed_indices))?; @@ -2009,14 +1809,14 @@ impl MaterializingSortMergeJoinStream { } /// Freezes the matched pairs of an Inner join with a filter. The filter runs - /// on just the columns it reads, for the pairs of the chunks that - /// [`Self::pair_filtered_run`] did not filter already, 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 no chunk was filtered already and earlier freezes - /// kept enough of their pairs, [`Self::filter_first_pays`] says the opposite, - /// and the batch is materialized whole and then filtered. + /// 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], @@ -2031,82 +1831,80 @@ impl MaterializingSortMergeJoinStream { "only an Inner join filters early, so the streamed side is the left one" ); - // Evaluate the filter over the pairs of the chunks not filtered yet. - let unfiltered: Vec<&MatchedChunk> = - matched_chunks.iter().filter(|c| !c.filtered).collect(); - if unfiltered.len() == matched_chunks.len() && !self.filter_first_pays(filter) { - return self.freeze_inner_filtered_whole(matched_chunks, total_matched_rows); + 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 passed = if unfiltered.is_empty() { - None - } else { - let num_rows = unfiltered.iter().map(|c| c.len()).sum(); - let streamed_indices = - concat_indices(unfiltered.iter().map(|c| &c.streamed_indices))?; - let mut buffered_arrays = self - .materialize_right_columns( - &unfiltered, - num_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::>>()?; - self.freeze_filter_evaluated += num_rows; - Some(evaluate_join_filter( - filter, - filter.schema(), - columns, - num_rows, - )?) - }; - // Keep the pairs that passed, in order. - let mut kept: Vec> = Vec::with_capacity(matched_chunks.len()); - let mut kept_rows = 0; - let mut offset = 0; - for chunk in matched_chunks { - let len = chunk.len(); - match (&passed, chunk.filtered) { - (Some(passed), false) => { - let mask = passed.slice(offset, len); - offset += len; - let count = mask.count_set_bits(); - if count == len { - kept.push(Cow::Borrowed(chunk)); - } else if count > 0 { - kept.push(Cow::Owned(chunk.filter(&mask)?)); - } - self.freeze_filter_passed += count; - kept_rows += count; - } - _ => { - kept.push(Cow::Borrowed(chunk)); - kept_rows += len; - } - } - } + 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(()); } - let streamed_indices = concat_indices(kept.iter().map(|c| &c.streamed_indices))?; + // 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, &streamed_indices)?; + 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( @@ -2120,16 +1918,15 @@ impl MaterializingSortMergeJoinStream { /// 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. None of the pairs was filtered already. + /// and filtering the batch. fn freeze_inner_filtered_whole( &mut self, matched_chunks: &[MatchedChunk], total_matched_rows: usize, + streamed_indices: &UInt64Array, ) -> Result<()> { - let streamed_indices = - concat_indices(matched_chunks.iter().map(|c| &c.streamed_indices))?; let streamed_columns = - materialize_left_columns(&self.streamed_batch.batch, &streamed_indices)?; + materialize_left_columns(&self.streamed_batch.batch, streamed_indices)?; let buffered_columns = self.materialize_right_columns( matched_chunks, total_matched_rows, @@ -2147,8 +1944,8 @@ impl MaterializingSortMergeJoinStream { total_matched_rows, )?; let kept_rows = passed.count_set_bits(); - self.freeze_filter_evaluated += total_matched_rows; - self.freeze_filter_passed += kept_rows; + self.freeze_filter_evaluated += total_matched_rows as u64; + self.freeze_filter_passed += kept_rows as u64; if kept_rows == 0 { return Ok(()); } @@ -2177,6 +1974,10 @@ impl MaterializingSortMergeJoinStream { /// 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; @@ -2194,23 +1995,20 @@ impl MaterializingSortMergeJoinStream { /// for a single fetch. When multiple batches are involved, `interleave` /// gathers columns across sources. A null-row sentinel at source index 0 /// handles null right indices (unmatched streamed rows). - fn materialize_right_columns>( + fn materialize_right_columns( &self, - matched_chunks: &[C], + matched_chunks: &[MatchedChunk], total_matched_rows: usize, columns: &[usize], ) -> Result> { - let chunks = || { - matched_chunks - .iter() - .map(|c| -> &MatchedChunk { c.borrow() }) - }; - let first_batch_idx = matched_chunks[0].borrow().buffered_batch_idx; - let single_source = chunks().all(|c| c.buffered_batch_idx == 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 = - concat_indices(chunks().map(|c| &c.buffered_indices))?; + concat_indices(matched_chunks.iter().map(|c| &c.buffered_indices))?; return fetch_right_columns_by_idxs( &self.buffered_data, @@ -2227,7 +2025,9 @@ impl MaterializingSortMergeJoinStream { // `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 = chunks().any(|c| c.buffered_indices.null_count() > 0); + let needs_null_sentinel = matched_chunks + .iter() + .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 @@ -2282,7 +2082,8 @@ impl MaterializingSortMergeJoinStream { // (`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) = chunks() + let (min_batch_idx, max_batch_idx) = matched_chunks + .iter() .fold((usize::MAX, 0usize), |(lo, hi), c| { (lo.min(c.buffered_batch_idx), hi.max(c.buffered_batch_idx)) }); @@ -2299,7 +2100,7 @@ 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 chunk in 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]; @@ -2532,6 +2333,16 @@ fn is_contiguous_range(indices: &UInt64Array) -> Option> { Some(start as usize..(start + len) as usize) } +/// 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( 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 14bc8a5d0be84..a48a038730dde 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -8452,10 +8452,9 @@ fn early_filtering_batches( } /// Inputs for the early-filtering tests. Every third key has a group of 600 -/// buffered rows, long enough to be filtered as it is paired, and the 250-row -/// buffered batches split each such group. The other keys have groups of one or -/// two buffered rows, whose pairs are filtered at freeze. Keys repeat up to four -/// times on the streamed side. +/// 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 = @@ -8520,22 +8519,23 @@ fn interleaved_early_filter(left_schema: &Schema, right_schema: &Schema) -> Join JoinFilter::new(expression, column_indices, Arc::new(schema)) } -/// The filter `b2 <> b1 OR c2 < c1`, which nearly every pair passes, with its -/// intermediate columns in the order `b2, b1, c2, c1`. It reads four of the six -/// output columns, so after the first freeze the join materializes whole batches -/// and filters them (see `filter_first_pays`). +/// 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("b2", 0)), + Arc::new(Column::new("c2", 2)), Operator::NotEq, - Arc::new(Column::new("b1", 1)), + Arc::new(Column::new("c1", 3)), )), Operator::Or, Arc::new(BinaryExpr::new( - Arc::new(Column::new("c2", 2)), + Arc::new(Column::new("b2", 0)), Operator::Lt, - Arc::new(Column::new("c1", 3)), + Arc::new(Column::new("b1", 1)), )), )); let column_indices = vec![ @@ -8581,33 +8581,36 @@ fn early_filtering_on(left_schema: &Schema, right_schema: &Schema) -> Result = batches + 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| { - b.column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec() - }) + .flat_map(|b| column(b, 1).into_iter().zip(column(b, 4))) .collect(); assert!( - keys.is_sorted(), - "Inner join output must stay ordered on the streamed side" + pairs.is_sorted(), + "Inner join output must keep the order of both inputs" ); } -/// An Inner join with a filter filters long runs of buffered rows as it pairs them -/// and the pairs of shorter runs at freeze. It must keep the same pairs as a hash -/// join whatever the batch size, which decides where a batch of pairs cuts a run, -/// 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. The -/// high-pass filter also covers freezes that materialize whole batches, and a -/// filter that reads no columns must still run for every pair. +/// 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)) = @@ -8655,8 +8658,7 @@ async fn join_inner_early_filtering_matches_hash_join() -> Result<()> { Ok(()) } -/// A long run in a spilled buffered batch is paired unfiltered, and its pairs are -/// filtered at freeze once the batch is read back. +/// 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)) =