From 41088840c610fa37866adf329764f789c4c16f38 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 10 Nov 2026 01:48:07 +0000 Subject: [PATCH 1/5] bench: cover residual semi/anti sort-merge join witnesses and latest versions --- .../smj/benchmarks/q27.benchmark | 13 + benchmarks/sql_benchmarks/smj/queries/q27.sql | 11 + benchmarks/src/smj.rs | 25 +- .../physical-plan/benches/sort_merge_join.rs | 298 ++++++++++++++++-- 4 files changed, 304 insertions(+), 43 deletions(-) create mode 100644 benchmarks/sql_benchmarks/smj/benchmarks/q27.benchmark create mode 100644 benchmarks/sql_benchmarks/smj/queries/q27.sql diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q27.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q27.benchmark new file mode 100644 index 0000000000000..2e931be04bd9e --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q27.benchmark @@ -0,0 +1,13 @@ +name Q27 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec +expect_plan Anti + +run sql_benchmarks/smj/queries/q27.sql + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/queries/q27.sql b/benchmarks/sql_benchmarks/smj/queries/q27.sql new file mode 100644 index 0000000000000..fa3267d0362ba --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/queries/q27.sql @@ -0,0 +1,11 @@ +-- Keep the newest of 128 versions of each of 1024 documents. +WITH documents AS ( + SELECT value / 128 AS document_id, value % 128 AS version + FROM range(131072) +) +SELECT d.document_id, d.version +FROM documents d +WHERE NOT EXISTS ( + SELECT 1 FROM documents newer + WHERE newer.document_id = d.document_id AND newer.version > d.version +) diff --git a/benchmarks/src/smj.rs b/benchmarks/src/smj.rs index f4923341cac67..761f58bf7440f 100644 --- a/benchmarks/src/smj.rs +++ b/benchmarks/src/smj.rs @@ -26,7 +26,7 @@ use futures::StreamExt; #[derive(Debug, Args, Clone)] #[command(verbatim_doc_comment)] pub struct RunOpt { - /// Query number (between 1 and 26). If not specified, runs all queries + /// One-based query number. If not specified, runs all queries #[arg(short, long)] query: Option, @@ -509,6 +509,8 @@ const SMJ_QUERIES: &[&str] = &[ AND t2_sorted.data % 10 <> 0 ) "#, + // Q27: LEFT ANTI | 1024 documents × 128 versions | one latest version per key + include_str!("../sql_benchmarks/smj/queries/q27.sql"), ]; impl RunOpt { @@ -543,10 +545,7 @@ impl RunOpt { let sql = SMJ_QUERIES[query_index]; benchmark_run.start_new_case(&format!("Query {query_id}")); - let expect_mark = query_id >= 24; - let query_run = self - .benchmark_query(sql, &query_id.to_string(), expect_mark, &ctx) - .await; + let query_run = self.benchmark_query(sql, query_id, &ctx).await; match query_run { Ok(query_results) => { for iter in query_results { @@ -569,8 +568,7 @@ impl RunOpt { async fn benchmark_query( &self, sql: &str, - query_name: &str, - expect_mark: bool, + query_id: usize, ctx: &SessionContext, ) -> Result> { let mut query_results = vec![]; @@ -582,13 +580,13 @@ impl RunOpt { if !plan_string.contains("SortMergeJoinExec") { return Err(exec_datafusion_err!( - "Query {query_name} does not use Sort Merge Join. Physical plan: {plan_string}" + "Query {query_id} does not use Sort Merge Join. Physical plan: {plan_string}" )); } - if expect_mark && !plan_string.contains("LeftMark") { + if (24..=26).contains(&query_id) && !plan_string.contains("LeftMark") { return Err(exec_datafusion_err!( - "Query {query_name} expected LeftMark join. Physical plan: {plan_string}" + "Query {query_id} expected LeftMark join. Physical plan: {plan_string}" )); } @@ -598,9 +596,14 @@ impl RunOpt { let row_count = Self::execute_sql_without_result_buffering(sql, ctx).await?; let elapsed = start.elapsed(); + if query_id == 27 && row_count != 1024 { + return exec_err!( + "Query {query_id} returned {row_count} latest versions, expected 1024" + ); + } println!( - "Query {query_name} iteration {i} returned {row_count} rows in {elapsed:?}" + "Query {query_id} iteration {i} returned {row_count} rows in {elapsed:?}" ); query_results.push(QueryResult { elapsed, row_count }); diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index c46ad9500b6ba..792f874e6d4db 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -10,18 +10,24 @@ use std::sync::Arc; -use arrow::array::{Int64Array, RecordBatch, StringArray}; +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; -use datafusion_common::{JoinSide, NullEquality}; -use datafusion_execution::TaskContext; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::{JoinSide, JoinType, NullEquality}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::{ + TaskContext, config::SessionConfig, memory_pool::FairSpillPool, + runtime_env::RuntimeEnvBuilder, +}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Column, col}; -use datafusion_physical_plan::collect; -use datafusion_physical_plan::joins::SortMergeJoinExec; -use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; +use datafusion_physical_plan::joins::{ + SortMergeJoinExec, + utils::{ColumnIndex, JoinFilter, JoinOn}, +}; use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, collect}; use tokio::runtime::Runtime; /// Build pre-sorted RecordBatches (split into ~8192-row chunks). @@ -78,10 +84,7 @@ fn build_sorted_batches_with_size( batches } -fn make_exec( - batches: &[RecordBatch], - schema: &SchemaRef, -) -> Arc { +fn make_exec(batches: &[RecordBatch], schema: &SchemaRef) -> Arc { TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() } @@ -93,8 +96,7 @@ fn schema() -> SchemaRef { ])) } -/// The join filter `left.data < right.data`. -fn data_lt_filter() -> JoinFilter { +fn data_filter(op: Operator, data_type: &DataType) -> JoinFilter { let column_indices = vec![ ColumnIndex { index: 1, @@ -106,30 +108,30 @@ fn data_lt_filter() -> JoinFilter { }, ]; let intermediate_schema = Schema::new(vec![ - Field::new("left_data", DataType::Int64, false), - Field::new("right_data", DataType::Int64, false), + Field::new("left_data", data_type.clone(), false), + Field::new("right_data", data_type.clone(), false), ]); 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)) } fn do_join( - left: Arc, - right: Arc, - join_type: datafusion_common::JoinType, + left: Arc, + right: Arc, + join_type: JoinType, rt: &Runtime, ) -> usize { do_join_with_filter(left, right, join_type, None, rt) } fn do_join_with_filter( - left: Arc, - right: Arc, - join_type: datafusion_common::JoinType, + left: Arc, + right: Arc, + join_type: JoinType, filter: Option, rt: &Runtime, ) -> usize { @@ -171,7 +173,7 @@ fn bench_smj(c: &mut Criterion) { b.iter(|| { let left = make_exec(&left_batches, &s); let right = make_exec(&right_batches, &s); - do_join(left, right, datafusion_common::JoinType::Inner, &rt) + do_join(left, right, JoinType::Inner, &rt) }) }); } @@ -186,7 +188,7 @@ fn bench_smj(c: &mut Criterion) { b.iter(|| { let left = make_exec(&left_batches, &s); let right = make_exec(&right_batches, &s); - do_join(left, right, datafusion_common::JoinType::Inner, &rt) + do_join(left, right, JoinType::Inner, &rt) }) }); } @@ -202,7 +204,7 @@ fn bench_smj(c: &mut Criterion) { b.iter(|| { let left = make_exec(&left_batches, &s); let right = make_exec(&right_batches, &s); - do_join(left, right, datafusion_common::JoinType::Inner, &rt) + do_join(left, right, JoinType::Inner, &rt) }) }); } @@ -224,8 +226,8 @@ fn bench_smj(c: &mut Criterion) { do_join_with_filter( left, right, - datafusion_common::JoinType::Left, - Some(data_lt_filter()), + JoinType::Left, + Some(data_filter(Operator::Lt, &DataType::Int64)), &rt, ) }) @@ -241,7 +243,7 @@ fn bench_smj(c: &mut Criterion) { b.iter(|| { let left = make_exec(&left_batches, &s); let right = make_exec(&right_batches, &s); - do_join(left, right, datafusion_common::JoinType::Left, &rt) + do_join(left, right, JoinType::Left, &rt) }) }); } @@ -256,7 +258,7 @@ fn bench_smj(c: &mut Criterion) { b.iter(|| { let left = make_exec(&left_batches, &s); let right = make_exec(&right_batches, &s); - do_join(left, right, datafusion_common::JoinType::LeftSemi, &rt) + do_join(left, right, JoinType::LeftSemi, &rt) }) }); } @@ -270,7 +272,7 @@ fn bench_smj(c: &mut Criterion) { b.iter(|| { let left = make_exec(&left_batches, &s); let right = make_exec(&right_batches, &s); - do_join(left, right, datafusion_common::JoinType::LeftAnti, &rt) + do_join(left, right, JoinType::LeftAnti, &rt) }) }); } @@ -302,7 +304,7 @@ fn bench_smj(c: &mut Criterion) { b.iter(|| { let left = make_exec(&left_batches, &s); let right = make_exec(&right_batches, &s); - do_join(left, right, datafusion_common::JoinType::Inner, &rt) + do_join(left, right, JoinType::Inner, &rt) }) }, ); @@ -312,5 +314,237 @@ fn bench_smj(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_smj); +/// Rows before `first_witness` equal the outer value. All later rows are greater. +fn residual_input( + groups: usize, + rows_per_group: usize, + first_witness: Option, + data_type: &DataType, +) -> Arc { + let rows = groups * rows_per_group; + let batches: Vec<_> = (0..rows) + .step_by(127) + .map(|start| { + let indices = start..(start + 127).min(rows); + let values = indices.clone().map(|i| { + 7 + i64::from( + first_witness.is_some_and(|first| i % rows_per_group >= first), + ) + }); + let data: ArrayRef = match data_type { + DataType::Int64 => Arc::new(Int64Array::from_iter_values(values)), + DataType::Utf8 => Arc::new(StringArray::from_iter_values( + values.map(|v| format!("{:x>120}{v:08}", "")), + )), + _ => unreachable!(), + }; + let keys = Int64Array::from_iter_values( + indices.map(|i| (i / rows_per_group) as i64), + ); + RecordBatch::try_from_iter(vec![ + ("key", Arc::new(keys) as ArrayRef), + ("data", data), + ]) + .unwrap() + }) + .collect(); + make_exec(&batches, &batches[0].schema()) +} + +fn bench_semi_anti_filter(c: &mut Criterion) { + let mut group = c.benchmark_group("sort_merge_join_semi_anti_filter"); + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .unwrap(); + + // Only the seven-row case runs concurrently: it stresses per-group pool traffic. + for (name, groups, outer_rows, inner_rows, first_witness, op, kind, compare_pools) in [ + ( + "not_equal_no_match", + 1, + 256, + 4096, + None, + Operator::NotEq, + JoinType::LeftAnti, + false, + ), + ( + "range_no_match", + 1, + 256, + 4096, + None, + Operator::Gt, + JoinType::LeftSemi, + false, + ), + ( + "range_early_match_many_probes", + 16, + 64, + 4096, + Some(0), + Operator::Lt, + JoinType::LeftSemi, + false, + ), + ( + "not_equal_third_witness", + 16, + 64, + 4096, + Some(2), + Operator::NotEq, + JoinType::LeftSemi, + false, + ), + ( + "range_early_match_one_probe", + 16, + 1, + 4096, + Some(0), + Operator::Lt, + JoinType::LeftSemi, + false, + ), + ( + "singleton_no_match", + 16, + 1, + 4096, + None, + Operator::Lt, + JoinType::LeftAnti, + false, + ), + ( + "small_group_seven_no_match", + 4096, + 7, + 7, + None, + Operator::NotEq, + JoinType::LeftAnti, + true, + ), + ] { + for data_type in [DataType::Int64, DataType::Utf8] { + for (pool_mode, workers, shared_pool) in [ + ("", 1, false), + ("_shared_fair_pool_4_tasks", 4, true), + ("_isolated_fair_pools_4_tasks", 4, false), + ] { + if workers > 1 && !compare_pools { + continue; + } + let id = BenchmarkId::new(name, format!("{data_type}{pool_mode}")); + group.bench_function(id, |b| { + let left = residual_input(groups, outer_rows, None, &data_type); + let right = + residual_input(groups, inner_rows, first_witness, &data_type); + let new_pool = || Arc::new(FairSpillPool::new(256 * 1024 * 1024)); + let shared = new_pool(); + let contexts: Vec<_> = (0..workers) + .map(|_| { + let ctx = TaskContext::default().with_session_config( + SessionConfig::new().with_batch_size(512), + ); + if workers == 1 { + return Arc::new(ctx); + } + let pool = if shared_pool { + Arc::clone(&shared) + } else { + new_pool() + }; + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc() + .unwrap(); + Arc::new(ctx.with_runtime(runtime)) + }) + .collect(); + + let make_plans = || { + (0..workers) + .map(|_| { + let on: JoinOn = vec![( + col("key", &left.schema()).unwrap(), + col("key", &right.schema()).unwrap(), + )]; + let plan = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on, + Some(data_filter(op, &data_type)), + kind, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + Arc::new(plan) as Arc + }) + .collect::>() + }; + let execute = |plans: Vec>| { + rt.block_on(async { + let count = |plan, ctx| async move { + collect(plan, ctx) + .await + .unwrap() + .iter() + .map(RecordBatch::num_rows) + .sum::() + }; + if workers == 1 { + return count( + Arc::clone(&plans[0]), + Arc::clone(&contexts[0]), + ) + .await; + } + let barrier = Arc::new(tokio::sync::Barrier::new(workers)); + let handles: Vec<_> = plans + .into_iter() + .zip(&contexts) + .map(|(plan, ctx)| { + let ctx = Arc::clone(ctx); + let barrier = Arc::clone(&barrier); + SpawnedTask::spawn(async move { + barrier.wait().await; + count(plan, ctx).await + }) + }) + .collect(); + let mut rows = 0; + for handle in handles { + rows += handle.await.unwrap(); + } + rows + }) + }; + let verification = make_plans(); + let emits = (kind == JoinType::LeftSemi) == first_witness.is_some(); + let expected_rows = if emits { + workers * groups * outer_rows + } else { + 0 + }; + assert_eq!(execute(verification.clone()), expected_rows); + for plan in verification { + assert_eq!(plan.metrics().unwrap().spill_count(), Some(0)); + } + b.iter_batched(make_plans, execute, BatchSize::PerIteration); + }); + } + } + } + group.finish(); +} + +criterion_group!(benches, bench_smj, bench_semi_anti_filter); criterion_main!(benches); From 9608e5bd01fa2834b99999fc5fa5459a48d71f9e Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 10 Nov 2026 01:48:07 +0000 Subject: [PATCH 2/5] perf: bound residual semi/anti sort-merge probes with incremental extrema --- .../joins/sort_merge_join/bitwise_stream.rs | 145 +++++- .../src/joins/sort_merge_join/mod.rs | 3 + .../sort_merge_join/semi_anti_summary.rs | 295 +++++++++++ .../semi_anti_summary_tests.rs | 469 ++++++++++++++++++ .../test_files/sort_merge_join.slt | 55 ++ .../test_files/sort_merge_join_matrix.slt | 359 ++++++++++++++ 6 files changed, 1306 insertions(+), 20 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs create mode 100644 datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index e313594d6b07c..c69d480f87ca0 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -57,6 +57,11 @@ //! outer key group as a batch. Results are OR'd into the matched bitset. A //! short-circuit exits early when all outer rows in the group are matched. //! +//! **With one cross-side column comparison**: Ordinary semi/anti joins can +//! summarize an in-memory inner group using min/max kernels. The summary is +//! computed lazily after initial probes, then reused across outer batches. +//! Spilled groups, complex predicates, and mark joins use ordinary evaluation. +//! //! ```text //! matched bitset: [0, 0, 1, 0, 0, ...] //! ▲── one bit per outer row ──▲ @@ -85,6 +90,8 @@ //! - Inner key group buffer: only for filtered joins, one key group at a time. //! Tracked via `MemoryReservation`; spilled to disk when the memory pool //! limit is exceeded. +//! - A completed summary replaces the buffered group in its reservation. +//! Temporary comparison masks and reduction values are local to a probe. //! - `BatchCoalescer`: output buffering to target batch size //! //! # Degenerate cases @@ -109,6 +116,7 @@ use std::cmp::Ordering; use std::sync::Arc; +use super::semi_anti_summary::{GroupSummary, SemiAntiComparison}; use crate::EmptyRecordBatchStream; use crate::joins::sort_merge_join::exec::NullAwareMode; use crate::joins::utils::{JoinFilter, JoinKeyComparator}; @@ -276,6 +284,7 @@ pub(crate) struct BitwiseSortMergeJoinStream { on_outer: Vec, on_inner: Vec, filter: Option, + semi_anti_comparison: Option, sort_options: Vec, null_equality: NullEquality, null_aware: Option, @@ -318,6 +327,10 @@ pub(crate) struct BitwiseSortMergeJoinStream { inner_self_cmp: Option, } +/// Below seven inner rows the ordinary filter costs less than building extrema +/// in the small-group benchmarks. +const MIN_SUMMARY_ROWS: usize = 7; + impl BitwiseSortMergeJoinStream { /// Creates a sort-merge join stream for semi, anti, and mark joins. /// @@ -363,6 +376,21 @@ impl BitwiseSortMergeJoinStream { JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark ); + let semi_anti_comparison = match &filter { + Some(filter) + if null_aware.is_none() + && !matches!(join_type, JoinType::LeftMark | JoinType::RightMark) => + { + SemiAntiComparison::try_new( + filter, + outer_is_left, + outer.schema().as_ref(), + inner.schema().as_ref(), + )? + } + _ => None, + }; + let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition); let input_batches = MetricBuilder::new(metrics).counter("input_batches", partition); @@ -386,6 +414,7 @@ impl BitwiseSortMergeJoinStream { on_outer, on_inner, filter, + semi_anti_comparison, sort_options, null_equality, null_aware, @@ -779,6 +808,51 @@ impl BitwiseSortMergeJoinStream { } } + /// Apply the recognized comparison to this slice and cache a completed + /// summary if its owned values fit in the group's reservation. Returns + /// false if later slices of this key must use the original filter. + fn apply_summary( + &mut self, + comparison: SemiAntiComparison, + outer: &RecordBatch, + summary: &mut Option>, + ) -> Result { + let matchable = + outer.num_rows() - outer.column(comparison.outer_column).null_count(); + if matchable == 0 { + return Ok(true); + } + if let Some(values) = summary { + merge_matches( + &mut self.matched, + self.outer_offset, + &comparison.evaluate(values, outer)?, + ); + return Ok(true); + } + + let start = self.outer_offset; + let matched = &mut self.matched; + let values = comparison.scan(&self.inner_key_buffer, outer, |mask| { + merge_matches(matched, start, mask); + UnalignedBitChunk::new(matched.as_slice(), start, outer.num_rows()) + .count_ones() + == matchable + })?; + if let Some(values) = values { + // A sliced string group can own extrema larger than its prorated + // buffer charge. Keep the buffer for later slices if resizing fails. + // The scan already produced the exact mask for this outer slice. + if self.reservation.try_resize(values.memory_size()).is_err() { + return Ok(false); + } + self.peak_mem_used.set_max(self.reservation.size()); + self.clear_inner_key_group(); + *summary = Some(Box::new(values)); + } + Ok(true) + } + /// Process a key match with a filter. For each inner row in the buffered /// key group — the spilled slices in `spill` plus the in-memory /// `inner_key_buffer` — evaluates the filter against the outer key group @@ -787,12 +861,14 @@ impl BitwiseSortMergeJoinStream { async fn process_key_match_with_filter( &mut self, spill: Option<&Arc>, + summary: &mut Option>, + group_comparison: &mut Option, ) -> Result<()> { let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); // buffer_inner_key_group must be called before this function debug_assert!( - !self.inner_key_buffer.is_empty() || spill.is_some(), + !self.inner_key_buffer.is_empty() || spill.is_some() || summary.is_some(), "process_key_match_with_filter called with no inner key data" ); debug_assert!( @@ -809,10 +885,19 @@ impl BitwiseSortMergeJoinStream { find_key_group_end(self.get_outer_self_cmp()?, outer_group_start, num_outer); let outer_group_len = outer_group_end - outer_group_start; - let filter = self.filter.as_ref().unwrap(); let outer_batch = self.outer_batch.as_ref().unwrap(); let outer_slice = outer_batch.slice(outer_group_start, outer_group_len); + if let Some(comparison) = *group_comparison { + if !self.apply_summary(comparison, &outer_slice, summary)? { + *group_comparison = None; + } + self.outer_offset = outer_group_end; + return Ok(()); + } + + let filter = self.filter.as_ref().unwrap(); + // Count already-matched bits using popcnt on u64 chunks (zero-copy). let mut matched_count = UnalignedBitChunk::new( self.matched.as_slice(), @@ -895,8 +980,24 @@ impl BitwiseSortMergeJoinStream { spill: Option>, emitter: &mut TryEmitter, ) -> Result<()> { + let mut summary = None; + let mut group_comparison = self.semi_anti_comparison.filter(|_| { + spill.is_none() + && self + .inner_key_buffer + .iter() + .take(MIN_SUMMARY_ROWS) + .map(RecordBatch::num_rows) + .sum::() + >= MIN_SUMMARY_ROWS + }); loop { - self.process_key_match_with_filter(spill.as_ref()).await?; + self.process_key_match_with_filter( + spill.as_ref(), + &mut summary, + &mut group_comparison, + ) + .await?; let outer_batch = self.outer_batch.as_ref().unwrap(); if self.outer_offset < outer_batch.num_rows() { @@ -924,6 +1025,9 @@ impl BitwiseSortMergeJoinStream { } } + if summary.is_some() { + self.reservation.free(); + } self.clear_inner_key_group(); Ok(()) } @@ -1281,23 +1385,7 @@ fn eval_filter_for_inner_slice( inner_row, )?; - // OR filter results into the matched bitset. Both sides are - // bit-packed [u8] buffers, so apply_bitwise_binary_op - // processes 64 bits per loop iteration (not 1 bit at a time). - // - // The offsets handle alignment: outer_offset is the bit - // position within matched where this key group starts, - // and filter_buf.offset() is the BooleanBuffer's internal - // bit offset (usually 0, but not guaranteed by Arrow). - let filter_buf = filter_result.values(); - apply_bitwise_binary_op( - matched.as_slice_mut(), - outer_offset, - filter_buf.inner().as_slice(), - filter_buf.offset(), - outer_group_len, - |a, b| a | b, - ); + merge_matches(matched, outer_offset, &filter_result); // Recount matched bits after the OR. UnalignedBitChunk is // zero-copy — it reads the bytes in place and uses popcnt. @@ -1308,6 +1396,23 @@ fn eval_filter_for_inner_slice( Ok(matched_count) } +/// OR a filter mask into the corresponding range of the outer batch's bitset. +fn merge_matches( + matched: &mut BooleanBufferBuilder, + outer_offset: usize, + result: &BooleanArray, +) { + let values = result.values(); + apply_bitwise_binary_op( + matched.as_slice_mut(), + outer_offset, + values.inner().as_slice(), + values.offset(), + result.len(), + |a, b| a | b, + ); +} + /// Slice each key array to a single row at `idx`. fn slice_keys(keys: &[ArrayRef], idx: usize) -> Vec { keys.iter().map(|a| a.slice(idx, 1)).collect() diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs b/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs index d17a74b5a1899..5b7c762fb6a22 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs @@ -11,6 +11,9 @@ mod exec; mod filter; pub(crate) mod materializing_stream; mod metrics; +mod semi_anti_summary; +#[cfg(test)] +mod semi_anti_summary_tests; #[cfg(test)] mod tests; diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs new file mode 100644 index 0000000000000..8a953934ef788 --- /dev/null +++ b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs @@ -0,0 +1,295 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Exact existence checks for one cross-side column comparison. +//! +//! A buffered key group witnesses `outer <[=] inner` iff it does so against +//! the inner maximum, and `outer >[=] inner` iff it does so against the minimum. +//! For `!=`, two distinct non-null values match every non-null outer value. +//! Complex expressions and floating-point types retain ordinary evaluation. + +use std::cmp::Ordering; + +use arrow::array::{Array, ArrayRef, BooleanArray, Datum, RecordBatch, Scalar}; +use arrow::compute::kernels::cmp; +use arrow::compute::or; +use arrow::datatypes::{DataType, Schema}; +use datafusion_common::{JoinSide, Result, ScalarValue}; +use datafusion_expr::Operator; +use datafusion_functions_aggregate_common::min_max::{max_batch, min_batch}; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, NotExpr}; + +use crate::joins::utils::{JoinFilter, boolean_mask_from_filter}; + +/// Direct probes avoid setting up min/max when either of the first two rows +/// witnesses the comparison, as in the common early-match benchmark cases. +const DIRECT_PROBES: usize = 2; +/// Check the first small vector of rows before doubling the prefix length. +const FIRST_PREFIX_ROWS: usize = 8; + +/// A recognized comparison, normalized to `outer OP inner`. +#[derive(Debug, Clone, Copy)] +pub(super) struct SemiAntiComparison { + pub(super) outer_column: usize, + inner_column: usize, + op: Operator, +} + +/// Cached once per buffered key group and reused across outer batches. +#[derive(Debug)] +pub(super) struct GroupSummary { + extreme: Scalar, + // A not-equal comparison needs both min and max. + other: Option>, +} + +impl GroupSummary { + fn from_extrema(extreme: &ScalarValue, other: Option<&ScalarValue>) -> Result { + Ok(Self { + extreme: extreme.to_scalar()?, + other: other.map(ScalarValue::to_scalar).transpose()?, + }) + } + + /// Inline state and the compact Arrow arrays retained for cached comparisons. + pub(super) fn memory_size(&self) -> usize { + let array_size = |value: &Scalar| value.get().0.get_array_memory_size(); + size_of::() + + array_size(&self.extreme) + + self.other.as_ref().map_or(0, array_size) + } +} + +impl SemiAntiComparison { + pub(super) fn try_new( + filter: &JoinFilter, + outer_is_left: bool, + outer_schema: &Schema, + inner_schema: &Schema, + ) -> Result> { + let (binary, mut op) = if let Some(not) = + filter.expression().downcast_ref::() + { + let Some(binary) = not.arg().downcast_ref::() else { + return Ok(None); + }; + let Some(op) = binary.op().negate() else { + return Ok(None); + }; + (binary, op) + } else if let Some(binary) = filter.expression().downcast_ref::() { + (binary, *binary.op()) + } else { + return Ok(None); + }; + if !matches!( + op, + Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + ) { + return Ok(None); + } + let (Some(left), Some(right)) = ( + binary.left().downcast_ref::(), + binary.right().downcast_ref::(), + ) else { + return Ok(None); + }; + let (Some(left), Some(right)) = ( + filter.column_indices().get(left.index()), + filter.column_indices().get(right.index()), + ) else { + return Ok(None); + }; + let outer_side = if outer_is_left { + JoinSide::Left + } else { + JoinSide::Right + }; + let inner_side = outer_side.negate(); + let (outer_column, inner_column) = + if left.side == outer_side && right.side == inner_side { + (left.index, right.index) + } else if right.side == outer_side && left.side == inner_side { + op = op.swap().unwrap(); + (right.index, left.index) + } else { + return Ok(None); + }; + let (Some(outer), Some(inner)) = ( + outer_schema.fields().get(outer_column), + inner_schema.fields().get(inner_column), + ) else { + return Ok(None); + }; + if outer.data_type() != inner.data_type() || !supported_type(inner.data_type()) { + return Ok(None); + } + Ok(Some(Self { + outer_column, + inner_column, + op, + })) + } + + /// Apply two direct row probes, then check the extrema after 8, 16, 32, ... + /// rows. Only the newly reached rows enter each reduction, including when a + /// checkpoint crosses an input batch boundary. A witness at row k costs + /// O(k) inner rows even when the full group is much larger. + /// + /// The callback records the mask for this outer slice and returns true once + /// every matchable row has a witness. A complete scan returns the exact + /// summary for later outer slices. An early stop leaves the input buffered. + pub(super) fn scan( + &self, + batches: &[RecordBatch], + outer: &RecordBatch, + mut matches_all: impl FnMut(&BooleanArray) -> bool, + ) -> Result> { + let mut probes_left = DIRECT_PROBES; + for batch in batches { + let len = probes_left.min(batch.num_rows()); + for row in 0..len { + if matches_all(&self.evaluate_inner_row(batch, row, outer)?) { + return Ok(None); + } + } + probes_left -= len; + if probes_left == 0 { + break; + } + } + + // The probes neither copy string payloads nor build reduction state. + // Include their values in the first reduction for later outer slices. + let data_type = batches[0].column(self.inner_column).data_type(); + let mut extreme = ScalarValue::try_from(data_type)?; + let mut other = (self.op == Operator::NotEq).then(|| extreme.clone()); + type Reduce = fn(&ArrayRef) -> Result; + let (reduce, order): (Reduce, Ordering) = match self.op { + Operator::Lt | Operator::LtEq => (max_batch, Ordering::Greater), + _ => (min_batch, Ordering::Less), + }; + + let mut scanned = 0; + let mut checkpoint = FIRST_PREFIX_ROWS; + let mut changed = false; + for (index, batch) in batches.iter().enumerate() { + let column = batch.column(self.inner_column); + let mut start = 0; + while start < column.len() { + let len = (checkpoint - scanned).min(column.len() - start); + let finished = index + 1 == batches.len() && start + len == column.len(); + let values = column.slice(start, len); + changed |= update_extreme(&mut extreme, reduce(&values)?, order); + if let Some(max) = &mut other { + changed |= + update_extreme(max, max_batch(&values)?, Ordering::Greater); + } + start += len; + scanned += len; + if scanned == checkpoint && !finished { + // The prior mask already included any unchanged extrema. + if changed { + let partial = + GroupSummary::from_extrema(&extreme, other.as_ref())?; + if matches_all(&self.evaluate(&partial, outer)?) { + return Ok(None); + } + } + changed = false; + checkpoint = checkpoint.saturating_mul(2); + } + } + } + // Reuse compact Arrow scalars across outer batches. A cached comparison + // with a long string must not copy that string into a new Arrow buffer. + let summary = GroupSummary::from_extrema(&extreme, other.as_ref())?; + if changed { + matches_all(&self.evaluate(&summary, outer)?); + } + Ok(Some(summary)) + } + + pub(super) fn evaluate( + &self, + summary: &GroupSummary, + outer: &RecordBatch, + ) -> Result { + let matched = self.compare(&summary.extreme, outer)?; + match &summary.other { + Some(other) => Ok(or(&matched, &self.compare(other, outer)?)?), + None => Ok(matched), + } + } + + /// Borrow the input for a single-row Arrow scalar, including string payloads. + pub(super) fn evaluate_inner_row( + &self, + inner: &RecordBatch, + row: usize, + outer: &RecordBatch, + ) -> Result { + let value = Scalar::new(inner.column(self.inner_column).slice(row, 1)); + self.compare(&value, outer) + } + + fn compare(&self, value: &dyn Datum, outer: &RecordBatch) -> Result { + let values = outer.column(self.outer_column).as_ref(); + // Only ordinary comparisons on non-floating primitive or string columns + // reach here, so these are the same Arrow kernels used by apply_cmp. + let compared = match self.op { + Operator::Lt => cmp::lt(&values, value)?, + Operator::LtEq => cmp::lt_eq(&values, value)?, + Operator::Gt => cmp::gt(&values, value)?, + Operator::GtEq => cmp::gt_eq(&values, value)?, + Operator::NotEq => cmp::neq(&values, value)?, + _ => unreachable!("invalid semi/anti summary comparison"), + }; + Ok(boolean_mask_from_filter(&compared)) + } +} + +/// Move a winning batch extremum into place. A losing batch never clones the +/// running winner, which may be a long string accumulated many batches ago. +fn update_extreme(current: &mut ScalarValue, next: ScalarValue, order: Ordering) -> bool { + if !next.is_null() && (current.is_null() || next.partial_cmp(current) == Some(order)) + { + *current = next; + true + } else { + false + } +} + +fn supported_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Date32 + | DataType::Date64 + | DataType::Time32(_) + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::Duration(_) + | DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + ) +} diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs new file mode 100644 index 0000000000000..e99a735a5c672 --- /dev/null +++ b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs @@ -0,0 +1,469 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Memory, stream, and predicate-eligibility coverage for semi/anti join summaries. +//! SQL results, data types, and batch boundaries are covered in +//! `sort_merge_join_matrix.slt`. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, StringArray}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::{ + DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue, + assert_contains, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, NotExpr}; +use futures::StreamExt; + +use super::semi_anti_summary::SemiAntiComparison; +use crate::joins::SortMergeJoinExec; +use crate::joins::utils::{ColumnIndex, JoinFilter}; +use crate::test::TestMemoryExec; +use crate::test::exec::MockExec; +use crate::{ExecutionPlan, PhysicalExpr, common}; + +fn column(index: usize) -> Arc { + Arc::new(Column::new(["left_value", "right_value"][index], index)) +} + +fn binary( + left: Arc, + op: Operator, + right: Arc, +) -> Arc { + Arc::new(BinaryExpr::new(left, op, right)) +} + +fn filter(expr: Arc, value_type: DataType) -> JoinFilter { + JoinFilter::new( + expr, + vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("left_value", value_type.clone(), true), + Field::new("right_value", value_type, true), + ])), + ) +} + +fn batch(values: ArrayRef) -> Result { + keyed_batch(vec![1; values.len()], values) +} + +fn keyed_batch(keys: Vec, values: ArrayRef) -> Result { + RecordBatch::try_from_iter(vec![ + ("key", Arc::new(Int32Array::from(keys)) as ArrayRef), + ("value", Arc::clone(&values)), + ( + "id", + Arc::new(Int32Array::from_iter_values(0..values.len() as i32)), + ), + ]) + .map_err(Into::into) +} + +fn input(batch: &RecordBatch, chunk: usize) -> Result> { + let batches = (0..batch.num_rows()) + .step_by(chunk) + .map(|offset| batch.slice(offset, chunk.min(batch.num_rows() - offset))) + .collect::>(); + Ok(TestMemoryExec::try_new_exec( + &[batches], + batch.schema(), + None, + )?) +} + +fn join( + left: Arc, + right: Arc, + join_type: JoinType, + op: Operator, +) -> Result { + let value_type = left.schema().field(1).data_type().clone(); + SortMergeJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + )], + Some(filter(binary(column(0), op, column(1)), value_type)), + join_type, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + ) +} + +fn row_ids(batch: &RecordBatch) -> Vec { + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() +} + +fn single_row_context() -> Arc { + Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(1)), + ) +} + +async fn ids(plan: &dyn ExecutionPlan, ctx: Arc) -> Result> { + let mut ids = common::collect(plan.execute(0, ctx)?) + .await? + .iter() + .flat_map(row_ids) + .collect::>(); + ids.sort_unstable(); + Ok(ids) +} + +#[test] +fn only_single_cross_side_comparisons_are_eligible() -> Result<()> { + let schema = batch(Arc::new(Int32Array::from(vec![1])))?.schema(); + // Exercise the physical predicate directly: SQL optimization can rewrite + // NOT or push a side-local guard below the join before this parser sees it. + for outer_is_left in [false, true] { + for op in [ + Operator::NotEq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + ] { + let direct = binary(column(0), op, column(1)); + assert!( + SemiAntiComparison::try_new( + &filter(direct, DataType::Int32), + outer_is_left, + &schema, + &schema, + )? + .is_some() + ); + } + for op in [ + Operator::Eq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + ] { + let negated = Arc::new(NotExpr::new(binary(column(0), op, column(1)))); + let comparison = SemiAntiComparison::try_new( + &filter(negated, DataType::Int32), + outer_is_left, + &schema, + &schema, + )? + .unwrap(); + let inner = batch(Arc::new(Int32Array::from(vec![None, Some(1), Some(3)])))?; + let outer = batch(Arc::new(Int32Array::from(vec![ + None, + Some(1), + Some(2), + Some(3), + ])))?; + assert_eq!( + comparison.evaluate_inner_row(&inner, 0, &outer)?, + BooleanArray::from(vec![false; 4]), + "NULL inner row: NOT {op:?}, outer_is_left={outer_is_left}" + ); + let expected_probe = match (op, outer_is_left) { + (Operator::Eq, _) | (Operator::LtEq, true) | (Operator::GtEq, false) => { + vec![false, false, true, true] + } + (Operator::Lt, true) | (Operator::Gt, false) => { + vec![false, true, true, true] + } + (Operator::Gt, true) | (Operator::Lt, false) => { + vec![false, true, false, false] + } + (Operator::LtEq, false) | (Operator::GtEq, true) => { + vec![false, false, false, false] + } + _ => unreachable!(), + }; + assert_eq!( + comparison.evaluate_inner_row(&inner, 1, &outer)?, + BooleanArray::from(expected_probe), + "inner value 1: NOT {op:?}, outer_is_left={outer_is_left}" + ); + let summary = comparison + .scan(&[inner.slice(0, 1), inner.slice(1, 2)], &outer, |_| false)? + .unwrap(); + let expected = match (op, outer_is_left) { + (Operator::LtEq, true) | (Operator::GtEq, false) => { + vec![false, false, true, true] + } + (Operator::GtEq, true) | (Operator::LtEq, false) => { + vec![false, true, true, false] + } + _ => vec![false, true, true, true], + }; + assert_eq!( + comparison.evaluate(&summary, &outer)?, + BooleanArray::from(expected), + "NOT {op:?}, outer_is_left={outer_is_left}" + ); + } + let comparison = binary(column(0), Operator::Lt, column(1)); + let guard = binary( + column(1), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + ); + for expr in [ + binary(column(0), Operator::Eq, column(1)), + binary(column(0), Operator::NotEq, column(0)), + binary(Arc::clone(&comparison), Operator::And, guard), + binary(Arc::clone(&comparison), Operator::Or, comparison), + binary( + binary(column(0), Operator::Plus, column(0)), + Operator::Lt, + column(1), + ), + ] { + assert!( + SemiAntiComparison::try_new( + &filter(expr, DataType::Int32), + outer_is_left, + &schema, + &schema, + )? + .is_none() + ); + } + } + Ok(()) +} + +#[tokio::test] +async fn large_strings_produce_the_same_results_when_the_buffer_spills() -> Result<()> { + let values = [ + "a".repeat(16 * 1024), + "a".repeat(16 * 1024), + "z".repeat(16 * 1024), + "b".repeat(16 * 1024), + "b".repeat(16 * 1024), + "b".repeat(16 * 1024), + "b".repeat(16 * 1024), + ]; + let outer = batch(Arc::new(StringArray::from(vec![values[0].as_str(), "m"])))?; + let inner = batch(Arc::new(StringArray::from_iter_values(values.iter())))?; + // The two initial probes cross input batches. For != the first probe + // already matches "m", so its matched bit must survive the second probe. + for op in [Operator::NotEq, Operator::Lt] { + for (limit, spill) in [(None, false), (Some(32 * 1024), true)] { + let mut builder = RuntimeEnvBuilder::new(); + if let Some(limit) = limit { + builder = builder + .with_memory_limit(limit, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory), + ); + } + let runtime = builder.build_arc()?; + let ctx = Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let plan = + join(input(&outer, 2)?, input(&inner, 1)?, JoinType::LeftSemi, op)?; + assert_eq!(ids(&plan, ctx).await?, vec![0, 1]); + assert_eq!(plan.metrics().unwrap().spill_count().unwrap() > 0, spill); + assert_eq!(runtime.memory_pool.reserved(), 0); + } + } + Ok(()) +} + +#[tokio::test] +async fn not_equal_reserves_both_compact_string_extrema() -> Result<()> { + let min = "a".repeat(16 * 1024); + let max = "z".repeat(32 * 1024); + let mut values = vec![min.as_str(); 7]; + values[6] = max.as_str(); + let inner = batch(Arc::new(StringArray::from(values)))?; + let outer = batch(Arc::new(StringArray::from(vec![min.as_str()])))?; + let plan = join( + input(&outer, 1)?, + input(&inner, 7)?, + JoinType::LeftSemi, + Operator::NotEq, + )?; + let ctx = single_row_context(); + let pool = Arc::clone(ctx.memory_pool()); + let mut stream = plan.execute(0, ctx)?; + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + + let min_bytes = ScalarValue::Utf8(Some(min)) + .to_array()? + .get_array_memory_size(); + let max_bytes = ScalarValue::Utf8(Some(max)) + .to_array()? + .get_array_memory_size(); + let retained = pool.reserved(); + assert!( + retained >= min_bytes + max_bytes, + "both Arrow scalars are live" + ); + assert!( + retained < inner.get_array_memory_size(), + "the parent buffer was released" + ); + drop(stream); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn oversized_extrema_fall_back_to_the_prorated_buffer() -> Result<()> { + // Group 1 ends inside the batch. Its seven rows fit at the group's + // per-row share, but even one of its two owned extrema exceeds the limit. + let long = "x".repeat(16 * 1024); + let mut keys = vec![2; 256]; + keys[..7].fill(1); + let mut values = vec![""; keys.len()]; + values[..7].fill(long.as_str()); + let inner = keyed_batch(keys, Arc::new(StringArray::from(values)))?; + let outer = batch(Arc::new(StringArray::from(vec![long.as_str(), "a"])))?; + let buffered_size = inner + .slice(0, 7) + .get_array_memory_size() + .div_ceil(inner.num_rows()) + * 7; + assert!(buffered_size < long.len()); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(buffered_size, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let ctx = Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let plan = join( + input(&outer, 1)?, + input(&inner, inner.num_rows())?, + JoinType::LeftAnti, + Operator::NotEq, + )?; + assert_eq!(ids(&plan, ctx).await?, vec![0]); + assert_eq!(plan.metrics().unwrap().spill_count().unwrap(), 0); + assert_eq!(runtime.memory_pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn partial_prefix_keeps_the_tail_for_the_next_outer_batch() -> Result<()> { + let outer = batch(Arc::new(Int32Array::from(vec![7, 16])))?; + let mut values = vec![1; 32]; + values[2] = 9; + values[31] = 17; + let inner = batch(Arc::new(Int32Array::from(values)))?; + let plan = join( + input(&outer, 1)?, + input(&inner, 32)?, + JoinType::LeftSemi, + Operator::Lt, + )?; + let ctx = single_row_context(); + let pool = Arc::clone(ctx.memory_pool()); + let mut stream = plan.execute(0, ctx)?; + + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + // The third inner row satisfied the first singleton. An early exit must + // retain the whole buffer because the next singleton needs the last row. + assert_eq!(pool.reserved(), inner.get_array_memory_size()); + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [1]); + assert!(stream.next().await.is_none()); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn singleton_outer_reuses_all_null_summary_until_key_end() -> Result<()> { + let outer = keyed_batch(vec![1, 1, 2], Arc::new(Int32Array::from(vec![7; 3])))?; + let mut keys = vec![1; 7]; + keys.push(3); + let inner = keyed_batch(keys, Arc::new(Int32Array::from(vec![None; 8])))?; + let plan = join( + input(&outer, 1)?, + input(&inner, 7)?, + JoinType::LeftAnti, + Operator::Lt, + )?; + let ctx = single_row_context(); + let pool = Arc::clone(ctx.memory_pool()); + let mut stream = plan.execute(0, ctx)?; + + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + let summary_bytes = pool.reserved(); + let group_bytes = inner.slice(0, 7).get_array_memory_size(); + assert!(summary_bytes > 0 && summary_bytes < group_bytes); + assert_eq!( + plan.metrics() + .unwrap() + .sum_by_name("peak_mem_used") + .unwrap() + .as_usize(), + group_bytes, + "summary admission must reuse the buffer reservation" + ); + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [1]); + assert_eq!(pool.reserved(), summary_bytes); + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [2]); + // The cursor is still at inner key 3, so draining the inner input cannot + // account for releasing the summary from key 1. + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn cached_summary_is_released_on_outer_error_or_drop() -> Result<()> { + let outer = batch(Arc::new(Int32Array::from(vec![7])))?; + let inner = batch(Arc::new(Int32Array::from(vec![1; 7])))?; + for read_error in [false, true] { + let source: Arc = Arc::new( + MockExec::new( + vec![ + Ok(outer.clone()), + Err(DataFusionError::Execution("injected outer failure".into())), + ], + outer.schema(), + ) + .with_use_task(false) + .with_unknown_statistics(), + ); + let plan = join(source, input(&inner, 7)?, JoinType::LeftAnti, Operator::Lt)?; + let ctx = single_row_context(); + let pool = Arc::clone(ctx.memory_pool()); + let mut stream = plan.execute(0, ctx)?; + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + assert!(pool.reserved() > 0 && pool.reserved() < inner.get_array_memory_size()); + if read_error { + let error = stream.next().await.unwrap().unwrap_err(); + assert_contains!(error.to_string(), "injected outer failure"); + } + drop(stream); + assert_eq!(pool.reserved(), 0); + } + Ok(()) +} diff --git a/datafusion/sqllogictest/test_files/sort_merge_join.slt b/datafusion/sqllogictest/test_files/sort_merge_join.slt index be394db7dc650..3664b5f73659e 100644 --- a/datafusion/sqllogictest/test_files/sort_merge_join.slt +++ b/datafusion/sqllogictest/test_files/sort_merge_join.slt @@ -1129,6 +1129,61 @@ DROP TABLE smj_key_r; statement ok RESET datafusion.execution.batch_size; +# Complex predicates retain the ordinary residual path. OR clauses can have +# different witnesses, and a missing inner group must survive an anti join. +statement ok +CREATE TABLE identities(id BIGINT, k BIGINT, user_id VARCHAR, org_id VARCHAR) AS VALUES + (1, 1, 'alice', 'x'), (2, 2, 'alice', 'x'), + (3, 3, NULL, 'x'), (4, 4, 'alice', NULL), (5, 5, 'carol', 'x'), + (6, 6, 'alice', 'x'); + +statement ok +CREATE TABLE observations(k BIGINT, user_id VARCHAR, org_id VARCHAR, active BOOLEAN) AS VALUES + (1, 'alice', 'x', true), (1, 'alice', 'x', true), + (1, 'bob', 'y', false), (2, 'alice', 'y', true), + (2, 'bob', 'x', true), (2, 'alice', 'x', true), (2, 'alice', 'x', true), + (2, 'alice', 'x', true), (2, 'alice', 'x', true), (2, 'alice', 'x', true), + (3, 'bob', NULL, true), (4, 'bob', NULL, true), (6, 'bob', 'y', true); + +query I rowsort +SELECT l.id FROM identities l LEFT ANTI JOIN observations r +ON l.k = r.k AND r.active AND (l.user_id <> r.user_id OR l.org_id <> r.org_id); +---- +1 +3 +5 + +query I rowsort +SELECT l.id FROM identities l LEFT SEMI JOIN observations r +ON l.k = r.k AND r.active AND (l.user_id <> r.user_id OR l.org_id <> r.org_id); +---- +2 +4 +6 + +# Keeping the side-local guard within one OR arm prevents it from being +# pushed below the join; the full residual must use generic evaluation. +query I rowsort +SELECT l.id FROM identities l LEFT ANTI JOIN observations r +ON l.k = r.k AND ((r.active AND l.user_id <> r.user_id) OR l.org_id <> r.org_id); +---- +3 +5 + +# Independent summaries cannot prove that the same observation has both +# differences: group 2 has different user and organization on separate rows. +query I rowsort +SELECT l.id FROM identities l LEFT SEMI JOIN observations r +ON l.k = r.k AND r.active AND l.user_id <> r.user_id AND l.org_id <> r.org_id; +---- +6 + +statement ok +DROP TABLE identities; + +statement ok +DROP TABLE observations; + # Config reset statement ok RESET datafusion.optimizer.prefer_hash_join; diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt b/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt index 766aa7b3b031a..4123d6ee11043 100644 --- a/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt +++ b/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt @@ -300,3 +300,362 @@ query I SELECT count(*) FROM smj_dup_l l WHERE NOT EXISTS (SELECT 1 FROM smj_dup_r r WHERE l.k = r.k); ---- 1 + +# ------------------------------------------------------------------ +# Semi/anti residual comparisons: duplicate extrema, all-NULL inner groups, +# unmatched groups, and NULL equi keys. Every case also runs with HashJoin and +# with one-row batches, exercising equal-key groups spanning several batches. +# Seven-row groups include NULL padding so larger batches exercise summaries; +# those NULL rows never witness a residual comparison and survive anti joins. +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_cmp_l(id INT, k INT, v INT) AS VALUES + (0, NULL, NULL), (1, NULL, 4), (2, 0, 7), + (3, 1, NULL), (4, 1, 2), (5, 1, 2), (6, 1, 6), + (7, 2, 3), (8, 4, 5), (9, 1, NULL), (10, 1, NULL), (11, 1, NULL); + +statement ok +CREATE TABLE smj_cmp_r(id INT, k INT, v INT) AS VALUES + (0, NULL, 1), (1, 1, NULL), (2, 1, 2), (3, 1, 2), + (4, 1, 6), (5, 2, NULL), (6, 3, 4), + (7, 1, NULL), (8, 1, NULL), (9, 1, NULL); + +# Residual <>: compare all four preserved-side/polarity choices. +query TI rowsort +SELECT 'left_semi', l.id FROM smj_cmp_l l LEFT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <> r.v +UNION ALL +SELECT 'left_anti', l.id FROM smj_cmp_l l LEFT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <> r.v +UNION ALL +SELECT 'right_semi', r.id FROM smj_cmp_l l RIGHT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <> r.v +UNION ALL +SELECT 'right_anti', r.id FROM smj_cmp_l l RIGHT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <> r.v; +---- +left_anti 0 +left_anti 1 +left_anti 10 +left_anti 11 +left_anti 2 +left_anti 3 +left_anti 7 +left_anti 8 +left_anti 9 +left_semi 4 +left_semi 5 +left_semi 6 +right_anti 0 +right_anti 1 +right_anti 5 +right_anti 6 +right_anti 7 +right_anti 8 +right_anti 9 +right_semi 2 +right_semi 3 +right_semi 4 + +# Residual <: compare all four preserved-side/polarity choices. +query TI rowsort +SELECT 'left_semi', l.id FROM smj_cmp_l l LEFT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v < r.v +UNION ALL +SELECT 'left_anti', l.id FROM smj_cmp_l l LEFT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v < r.v +UNION ALL +SELECT 'right_semi', r.id FROM smj_cmp_l l RIGHT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v < r.v +UNION ALL +SELECT 'right_anti', r.id FROM smj_cmp_l l RIGHT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v < r.v; +---- +left_anti 0 +left_anti 1 +left_anti 10 +left_anti 11 +left_anti 2 +left_anti 3 +left_anti 6 +left_anti 7 +left_anti 8 +left_anti 9 +left_semi 4 +left_semi 5 +right_anti 0 +right_anti 1 +right_anti 2 +right_anti 3 +right_anti 5 +right_anti 6 +right_anti 7 +right_anti 8 +right_anti 9 +right_semi 4 + +# Residual <=: compare all four preserved-side/polarity choices. +query TI rowsort +SELECT 'left_semi', l.id FROM smj_cmp_l l LEFT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <= r.v +UNION ALL +SELECT 'left_anti', l.id FROM smj_cmp_l l LEFT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <= r.v +UNION ALL +SELECT 'right_semi', r.id FROM smj_cmp_l l RIGHT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <= r.v +UNION ALL +SELECT 'right_anti', r.id FROM smj_cmp_l l RIGHT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v <= r.v; +---- +left_anti 0 +left_anti 1 +left_anti 10 +left_anti 11 +left_anti 2 +left_anti 3 +left_anti 7 +left_anti 8 +left_anti 9 +left_semi 4 +left_semi 5 +left_semi 6 +right_anti 0 +right_anti 1 +right_anti 5 +right_anti 6 +right_anti 7 +right_anti 8 +right_anti 9 +right_semi 2 +right_semi 3 +right_semi 4 + +# Residual >: compare all four preserved-side/polarity choices. +query TI rowsort +SELECT 'left_semi', l.id FROM smj_cmp_l l LEFT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v > r.v +UNION ALL +SELECT 'left_anti', l.id FROM smj_cmp_l l LEFT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v > r.v +UNION ALL +SELECT 'right_semi', r.id FROM smj_cmp_l l RIGHT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v > r.v +UNION ALL +SELECT 'right_anti', r.id FROM smj_cmp_l l RIGHT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v > r.v; +---- +left_anti 0 +left_anti 1 +left_anti 10 +left_anti 11 +left_anti 2 +left_anti 3 +left_anti 4 +left_anti 5 +left_anti 7 +left_anti 8 +left_anti 9 +left_semi 6 +right_anti 0 +right_anti 1 +right_anti 4 +right_anti 5 +right_anti 6 +right_anti 7 +right_anti 8 +right_anti 9 +right_semi 2 +right_semi 3 + +# Residual >=: compare all four preserved-side/polarity choices. +query TI rowsort +SELECT 'left_semi', l.id FROM smj_cmp_l l LEFT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v >= r.v +UNION ALL +SELECT 'left_anti', l.id FROM smj_cmp_l l LEFT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v >= r.v +UNION ALL +SELECT 'right_semi', r.id FROM smj_cmp_l l RIGHT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.v >= r.v +UNION ALL +SELECT 'right_anti', r.id FROM smj_cmp_l l RIGHT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND l.v >= r.v; +---- +left_anti 0 +left_anti 1 +left_anti 10 +left_anti 11 +left_anti 2 +left_anti 3 +left_anti 7 +left_anti 8 +left_anti 9 +left_semi 4 +left_semi 5 +left_semi 6 +right_anti 0 +right_anti 1 +right_anti 5 +right_anti 6 +right_anti 7 +right_anti 8 +right_anti 9 +right_semi 2 +right_semi 3 +right_semi 4 + +# NULL equi keys may match when requested, while a NULL residual is never true. +query I rowsort +SELECT l.id FROM smj_cmp_l l LEFT SEMI JOIN smj_cmp_r r +ON l.k IS NOT DISTINCT FROM r.k AND l.v > r.v; +---- +1 +6 + +# Negated equality has the same existence semantics as <>. +query I rowsort +SELECT l.id FROM smj_cmp_l l LEFT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND NOT (l.v = r.v); +---- +0 +1 +10 +11 +2 +3 +7 +8 +9 + +# For outer IDs 4 and 5 the value comparison needs inner ID 4. Later inner +# rows satisfy ID 5's ID comparison but have NULL values. Only ID 4 has a +# single witness for both predicates. +query I rowsort +SELECT l.id FROM smj_cmp_l l LEFT SEMI JOIN smj_cmp_r r +ON l.k = r.k AND l.id <= r.id AND l.v < r.v; +---- +4 + +# Complex expressions use the ordinary residual path, including normalization +# of NULL string values to empty strings. +statement ok +CREATE TABLE smj_values AS SELECT id, 1 AS k, + arrow_cast(s, 'Utf8') AS utf8, + arrow_cast(s, 'LargeUtf8') AS large_utf8, + arrow_cast(s, 'Utf8View') AS utf8_view, + CAST(d AS DATE) AS date_value, + arrow_cast(ts, 'Timestamp(Microsecond, Some("UTC"))') AS timestamp_value, + CAST(n AS DECIMAL(20, 2)) AS decimal_value +FROM (VALUES + (0, NULL, NULL, NULL, NULL), + (1, '', '1969-12-31', '1969-12-31T23:59:59Z', -1.23), + (2, 'a', '1970-01-01', '1970-01-01T00:00:00Z', 0.00), + (3, 'out-of-line-string-z', '1970-01-02', '1970-01-01T00:00:01Z', 4.56), + (4, 'é', '1970-01-02', '1970-01-01T00:00:01Z', 4.56), + (5, NULL, NULL, NULL, NULL), (6, NULL, NULL, NULL, NULL) +) AS t(id, s, d, ts, n); + +# Semi join with < over string layouts and temporal/decimal types. +query TI rowsort +SELECT 'utf8', l.id FROM smj_values l LEFT SEMI JOIN smj_values r +ON l.k = r.k AND l.utf8 < r.utf8 +UNION ALL +SELECT 'large_utf8', l.id FROM smj_values l LEFT SEMI JOIN smj_values r +ON l.k = r.k AND l.large_utf8 < r.large_utf8 +UNION ALL +SELECT 'utf8_view', l.id FROM smj_values l LEFT SEMI JOIN smj_values r +ON l.k = r.k AND l.utf8_view < r.utf8_view +UNION ALL +SELECT 'date_value', l.id FROM smj_values l LEFT SEMI JOIN smj_values r +ON l.k = r.k AND l.date_value < r.date_value +UNION ALL +SELECT 'timestamp_value', l.id FROM smj_values l LEFT SEMI JOIN smj_values r +ON l.k = r.k AND l.timestamp_value < r.timestamp_value +UNION ALL +SELECT 'decimal_value', l.id FROM smj_values l LEFT SEMI JOIN smj_values r +ON l.k = r.k AND l.decimal_value < r.decimal_value; +---- +date_value 1 +date_value 2 +decimal_value 1 +decimal_value 2 +large_utf8 1 +large_utf8 2 +large_utf8 3 +timestamp_value 1 +timestamp_value 2 +utf8 1 +utf8 2 +utf8 3 +utf8_view 1 +utf8_view 2 +utf8_view 3 + +# Anti join with <> over string layouts and temporal/decimal types. +query TI rowsort +SELECT 'utf8', l.id FROM smj_values l LEFT ANTI JOIN smj_values r +ON l.k = r.k AND l.utf8 <> r.utf8 +UNION ALL +SELECT 'large_utf8', l.id FROM smj_values l LEFT ANTI JOIN smj_values r +ON l.k = r.k AND l.large_utf8 <> r.large_utf8 +UNION ALL +SELECT 'utf8_view', l.id FROM smj_values l LEFT ANTI JOIN smj_values r +ON l.k = r.k AND l.utf8_view <> r.utf8_view +UNION ALL +SELECT 'date_value', l.id FROM smj_values l LEFT ANTI JOIN smj_values r +ON l.k = r.k AND l.date_value <> r.date_value +UNION ALL +SELECT 'timestamp_value', l.id FROM smj_values l LEFT ANTI JOIN smj_values r +ON l.k = r.k AND l.timestamp_value <> r.timestamp_value +UNION ALL +SELECT 'decimal_value', l.id FROM smj_values l LEFT ANTI JOIN smj_values r +ON l.k = r.k AND l.decimal_value <> r.decimal_value; +---- +date_value 0 +date_value 5 +date_value 6 +decimal_value 0 +decimal_value 5 +decimal_value 6 +large_utf8 0 +large_utf8 5 +large_utf8 6 +timestamp_value 0 +timestamp_value 5 +timestamp_value 6 +utf8 0 +utf8 5 +utf8 6 +utf8_view 0 +utf8_view 5 +utf8_view 6 + +query I rowsort +SELECT l.id FROM smj_values l LEFT ANTI JOIN smj_values r +ON l.k = r.k AND COALESCE(l.utf8, '') <> COALESCE(r.utf8, ''); +---- + +# OR can have a witness from either comparison, but NULL never supplies one. +query I rowsort +SELECT l.id FROM smj_cmp_l l LEFT ANTI JOIN smj_cmp_r r +ON l.k = r.k AND (l.v < r.v OR l.v > r.v); +---- +0 +1 +10 +11 +2 +3 +7 +8 +9 + +statement ok +DROP TABLE smj_cmp_l; + +statement ok +DROP TABLE smj_cmp_r; + +statement ok +DROP TABLE smj_values; From ef9dbae2bc4f339d5f23972ac1b1a9ac105f5acf Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 10 Nov 2026 04:00:05 +0000 Subject: [PATCH 3/5] Borrow string extrema until cache admission and verify Q27 output --- benchmarks/src/smj.rs | 19 ++ .../joins/sort_merge_join/bitwise_stream.rs | 16 +- .../sort_merge_join/semi_anti_summary.rs | 180 +++++++++++++----- .../semi_anti_summary_tests.rs | 102 +++++++--- 4 files changed, 235 insertions(+), 82 deletions(-) diff --git a/benchmarks/src/smj.rs b/benchmarks/src/smj.rs index 761f58bf7440f..8ac4109ba5536 100644 --- a/benchmarks/src/smj.rs +++ b/benchmarks/src/smj.rs @@ -6,6 +6,7 @@ use crate::util::{BenchmarkRun, CommonOpt, QueryResult}; use clap::Args; use datafusion::physical_plan::execute_stream; use datafusion::{error::Result, prelude::SessionContext}; +use datafusion_common::cast::as_int64_array; use datafusion_common::instant::Instant; use datafusion_common::{DataFusionError, exec_datafusion_err, exec_err}; @@ -609,6 +610,24 @@ impl RunOpt { query_results.push(QueryResult { elapsed, row_count }); } + if query_id == 27 { + let mut rows = Vec::with_capacity(1024); + let mut stream = execute_stream(physical_plan, ctx.task_ctx())?; + while let Some(batch) = stream.next().await { + let batch = batch?; + let ids = as_int64_array(batch.column(0).as_ref())?; + let versions = as_int64_array(batch.column(1).as_ref())?; + rows.extend(ids.iter().zip(versions.iter())); + } + rows.sort_unstable(); + let expected = (0..1024).map(|id| (Some(id), Some(127))); + if rows.into_iter().ne(expected) { + return exec_err!( + "Query 27 expected version 127 exactly once for each document ID from 0 through 1023" + ); + } + } + Ok(query_results) } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index c69d480f87ca0..c6b9c94d52774 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -840,15 +840,21 @@ impl BitwiseSortMergeJoinStream { == matchable })?; if let Some(values) = values { - // A sliced string group can own extrema larger than its prorated - // buffer charge. Keep the buffer for later slices if resizing fails. - // The scan already produced the exact mask for this outer slice. - if self.reservation.try_resize(values.memory_size()).is_err() { + let retained = values.memory_size(); + // A string prefix borrows the input. Admit its compact copy alongside + // the buffer before allocating it. On refusal the mask is already + // exact for this slice and later slices can still use the buffer. + let admitted = match values { + GroupSummary::Strings { .. } => self.reservation.try_grow(retained), + GroupSummary::Scalars { .. } => self.reservation.try_resize(retained), + }; + if admitted.is_err() { return Ok(false); } self.peak_mem_used.set_max(self.reservation.size()); + *summary = Some(Box::new(values.into_owned())); self.clear_inner_key_group(); - *summary = Some(Box::new(values)); + self.reservation.resize(retained); } Ok(true) } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs index 8a953934ef788..a2f926381a67c 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs @@ -11,7 +11,8 @@ use std::cmp::Ordering; -use arrow::array::{Array, ArrayRef, BooleanArray, Datum, RecordBatch, Scalar}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Datum, RecordBatch, Scalar}; +use arrow::buffer::BooleanBuffer; use arrow::compute::kernels::cmp; use arrow::compute::or; use arrow::datatypes::{DataType, Schema}; @@ -38,26 +39,57 @@ pub(super) struct SemiAntiComparison { /// Cached once per buffered key group and reused across outer batches. #[derive(Debug)] -pub(super) struct GroupSummary { - extreme: Scalar, - // A not-equal comparison needs both min and max. - other: Option>, +pub(super) enum GroupSummary, N = Scalar> { + Scalars { + extreme: N, + other: Option, + }, + Strings { + extreme: Option, + other: Option, + }, } -impl GroupSummary { - fn from_extrema(extreme: &ScalarValue, other: Option<&ScalarValue>) -> Result { - Ok(Self { - extreme: extreme.to_scalar()?, - other: other.map(ScalarValue::to_scalar).transpose()?, +/// While scanning, fixed-width values use DataFusion's min/max kernels and +/// string values borrow directly from the buffered input. +type RunningExtrema<'a> = GroupSummary<&'a str, ScalarValue>; + +impl<'a> RunningExtrema<'a> { + fn snapshot(&self) -> Result> { + Ok(match self { + Self::Scalars { extreme, other } => GroupSummary::Scalars { + extreme: extreme.to_scalar()?, + other: other.as_ref().map(ScalarValue::to_scalar).transpose()?, + }, + &Self::Strings { extreme, other } => GroupSummary::Strings { extreme, other }, }) } +} - /// Inline state and the compact Arrow arrays retained for cached comparisons. +impl GroupSummary<&str> { pub(super) fn memory_size(&self) -> usize { let array_size = |value: &Scalar| value.get().0.get_array_memory_size(); - size_of::() - + array_size(&self.extreme) - + self.other.as_ref().map_or(0, array_size) + size_of::() + + match self { + Self::Scalars { extreme, other } => { + array_size(extreme) + other.as_ref().map_or(0, array_size) + } + Self::Strings { extreme, other } => { + extreme.map_or(0, str::len) + other.map_or(0, str::len) + } + } + } + + /// String winners borrow the buffered input until the reservation admits + /// their exact byte lengths. Fixed-width Arrow scalars are already owned. + pub(super) fn into_owned(self) -> GroupSummary { + match self { + Self::Scalars { extreme, other } => GroupSummary::Scalars { extreme, other }, + Self::Strings { extreme, other } => GroupSummary::Strings { + extreme: extreme.map(Box::from), + other: other.map(Box::from), + }, + } } } @@ -144,12 +176,12 @@ impl SemiAntiComparison { /// The callback records the mask for this outer slice and returns true once /// every matchable row has a witness. A complete scan returns the exact /// summary for later outer slices. An early stop leaves the input buffered. - pub(super) fn scan( + pub(super) fn scan<'a>( &self, - batches: &[RecordBatch], + batches: &'a [RecordBatch], outer: &RecordBatch, mut matches_all: impl FnMut(&BooleanArray) -> bool, - ) -> Result> { + ) -> Result>> { let mut probes_left = DIRECT_PROBES; for batch in batches { let len = probes_left.min(batch.num_rows()); @@ -167,12 +199,21 @@ impl SemiAntiComparison { // The probes neither copy string payloads nor build reduction state. // Include their values in the first reduction for later outer slices. let data_type = batches[0].column(self.inner_column).data_type(); - let mut extreme = ScalarValue::try_from(data_type)?; - let mut other = (self.op == Operator::NotEq).then(|| extreme.clone()); - type Reduce = fn(&ArrayRef) -> Result; - let (reduce, order): (Reduce, Ordering) = match self.op { - Operator::Lt | Operator::LtEq => (max_batch, Ordering::Greater), - _ => (min_batch, Ordering::Less), + let wants_max = matches!(self.op, Operator::Lt | Operator::LtEq); + let order = if wants_max { + Ordering::Greater + } else { + Ordering::Less + }; + let mut state = if data_type.is_string() { + RunningExtrema::Strings { + extreme: None, + other: None, + } + } else { + let extreme = ScalarValue::try_from(data_type)?; + let other = (self.op == Operator::NotEq).then(|| extreme.clone()); + RunningExtrema::Scalars { extreme, other } }; let mut scanned = 0; @@ -184,46 +225,85 @@ impl SemiAntiComparison { while start < column.len() { let len = (checkpoint - scanned).min(column.len() - start); let finished = index + 1 == batches.len() && start + len == column.len(); - let values = column.slice(start, len); - changed |= update_extreme(&mut extreme, reduce(&values)?, order); - if let Some(max) = &mut other { - changed |= - update_extreme(max, max_batch(&values)?, Ordering::Greater); + match &mut state { + RunningExtrema::Scalars { extreme, other } => { + let values = column.slice(start, len); + let next = if wants_max { + max_batch(&values)? + } else { + min_batch(&values)? + }; + changed |= update_extreme(extreme, next, order); + if let Some(max) = other { + changed |= update_extreme( + max, + max_batch(&values)?, + Ordering::Greater, + ); + } + } + RunningExtrema::Strings { extreme, other } => { + for value in strings(column).skip(start).take(len).flatten() { + if extreme.is_none_or(|old| value.cmp(old) == order) { + *extreme = Some(value); + changed = true; + } + if self.op == Operator::NotEq + && other.is_none_or(|old| value > old) + { + *other = Some(value); + changed = true; + } + } + } } start += len; scanned += len; if scanned == checkpoint && !finished { - // The prior mask already included any unchanged extrema. - if changed { - let partial = - GroupSummary::from_extrema(&extreme, other.as_ref())?; - if matches_all(&self.evaluate(&partial, outer)?) { - return Ok(None); - } + if changed && matches_all(&self.evaluate(&state.snapshot()?, outer)?) + { + return Ok(None); } changed = false; checkpoint = checkpoint.saturating_mul(2); } } } - // Reuse compact Arrow scalars across outer batches. A cached comparison - // with a long string must not copy that string into a new Arrow buffer. - let summary = GroupSummary::from_extrema(&extreme, other.as_ref())?; + let summary = state.snapshot()?; if changed { matches_all(&self.evaluate(&summary, outer)?); } Ok(Some(summary)) } - pub(super) fn evaluate( + pub(super) fn evaluate>( &self, - summary: &GroupSummary, + summary: &GroupSummary, outer: &RecordBatch, ) -> Result { - let matched = self.compare(&summary.extreme, outer)?; - match &summary.other { - Some(other) => Ok(or(&matched, &self.compare(other, outer)?)?), - None => Ok(matched), + match summary { + GroupSummary::Scalars { extreme, other } => { + let matched = self.compare(extreme, outer)?; + match other { + Some(other) => Ok(or(&matched, &self.compare(other, outer)?)?), + None => Ok(matched), + } + } + GroupSummary::Strings { extreme, other } => { + let extreme = extreme.as_ref().map(AsRef::as_ref); + let other = other.as_ref().map(AsRef::as_ref); + let hits = strings(outer.column(self.outer_column)).map(|value| { + value.zip(extreme).is_some_and(|(x, y)| match self.op { + Operator::Lt => x < y, + Operator::LtEq => x <= y, + Operator::Gt => x > y, + Operator::GtEq => x >= y, + Operator::NotEq => x != y || other.is_some_and(|z| x != z), + _ => unreachable!("invalid semi/anti summary comparison"), + }) + }); + Ok(BooleanArray::new(BooleanBuffer::from_iter(hits), None)) + } } } @@ -254,8 +334,16 @@ impl SemiAntiComparison { } } -/// Move a winning batch extremum into place. A losing batch never clones the -/// running winner, which may be a long string accumulated many batches ago. +fn strings(column: &ArrayRef) -> Box> + '_> { + match column.data_type() { + DataType::Utf8 => Box::new(column.as_string::().iter()), + DataType::LargeUtf8 => Box::new(column.as_string::().iter()), + DataType::Utf8View => Box::new(column.as_string_view().iter()), + _ => unreachable!("expected a string column"), + } +} + +/// Update a fixed-width extremum, ignoring NULLs. fn update_extreme(current: &mut ScalarValue, next: ScalarValue, order: Ordering) -> bool { if !next.is_null() && (current.is_null() || next.partial_cmp(current) == Some(order)) { diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs index e99a735a5c672..146dce594735a 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs @@ -23,7 +23,7 @@ use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, NotExpr}; use futures::StreamExt; -use super::semi_anti_summary::SemiAntiComparison; +use super::semi_anti_summary::{GroupSummary, SemiAntiComparison}; use crate::joins::SortMergeJoinExec; use crate::joins::utils::{ColumnIndex, JoinFilter}; use crate::test::TestMemoryExec; @@ -209,9 +209,8 @@ fn only_single_cross_side_comparisons_are_eligible() -> Result<()> { BooleanArray::from(expected_probe), "inner value 1: NOT {op:?}, outer_is_left={outer_is_left}" ); - let summary = comparison - .scan(&[inner.slice(0, 1), inner.slice(1, 2)], &outer, |_| false)? - .unwrap(); + let inner_batches = [inner.slice(0, 1), inner.slice(1, 2)]; + let summary = comparison.scan(&inner_batches, &outer, |_| false)?.unwrap(); let expected = match (op, outer_is_left) { (Operator::LtEq, true) | (Operator::GtEq, false) => { vec![false, false, true, true] @@ -300,10 +299,36 @@ async fn large_strings_produce_the_same_results_when_the_buffer_spills() -> Resu async fn not_equal_reserves_both_compact_string_extrema() -> Result<()> { let min = "a".repeat(16 * 1024); let max = "z".repeat(32 * 1024); - let mut values = vec![min.as_str(); 7]; + let middle = "m".repeat(1024); + let mut values = vec![middle.as_str(); 7]; + values[2] = min.as_str(); values[6] = max.as_str(); - let inner = batch(Arc::new(StringArray::from(values)))?; - let outer = batch(Arc::new(StringArray::from(vec![min.as_str()])))?; + let values = Arc::new(StringArray::from(values)); + let inner = batch(Arc::clone(&values) as ArrayRef)?; + let outer = batch(Arc::new(StringArray::from(vec![middle.as_str()])))?; + + // Until the cache is admitted, even extrema from different batches + // borrow their original Arrow payloads. + let comparison = SemiAntiComparison::try_new( + &filter( + binary(column(0), Operator::NotEq, column(1)), + DataType::Utf8, + ), + true, + &outer.schema(), + &inner.schema(), + )? + .unwrap(); + let batches = [inner.slice(0, 4), inner.slice(4, 3)]; + let Some(GroupSummary::Strings { + extreme: Some(lo), + other: Some(hi), + }) = comparison.scan(&batches, &outer, |_| false)? + else { + panic!("expected two string extrema"); + }; + assert!(std::ptr::eq(lo, values.value(2))); + assert!(std::ptr::eq(hi, values.value(6))); let plan = join( input(&outer, 1)?, input(&inner, 7)?, @@ -315,16 +340,10 @@ async fn not_equal_reserves_both_compact_string_extrema() -> Result<()> { let mut stream = plan.execute(0, ctx)?; assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); - let min_bytes = ScalarValue::Utf8(Some(min)) - .to_array()? - .get_array_memory_size(); - let max_bytes = ScalarValue::Utf8(Some(max)) - .to_array()? - .get_array_memory_size(); let retained = pool.reserved(); assert!( - retained >= min_bytes + max_bytes, - "both Arrow scalars are live" + retained >= min.len() + max.len(), + "both owned string payloads must be charged" ); assert!( retained < inner.get_array_memory_size(), @@ -352,22 +371,43 @@ async fn oversized_extrema_fall_back_to_the_prorated_buffer() -> Result<()> { .div_ceil(inner.num_rows()) * 7; assert!(buffered_size < long.len()); - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(buffered_size, 1.0) - .with_disk_manager_builder( - DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), - ) - .build_arc()?; - let ctx = Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); - let plan = join( - input(&outer, 1)?, - input(&inner, inner.num_rows())?, - JoinType::LeftAnti, - Operator::NotEq, - )?; - assert_eq!(ids(&plan, ctx).await?, vec![0]); - assert_eq!(plan.metrics().unwrap().spill_count().unwrap(), 0); - assert_eq!(runtime.memory_pool.reserved(), 0); + for (limit, expect_cache) in [(buffered_size, false), (long.len() * 8, true)] { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(1)) + .with_runtime(Arc::clone(&runtime)), + ); + let plan = join( + input(&outer, 1)?, + input(&inner, inner.num_rows())?, + JoinType::LeftAnti, + Operator::NotEq, + )?; + let mut stream = plan.execute(0, ctx)?; + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + // At the yield the same key may continue in the next outer batch. + let retained = runtime.memory_pool.reserved(); + if expect_cache { + assert!( + retained >= long.len(), + "owned string extrema must be charged" + ); + } else { + assert_eq!( + retained, buffered_size, + "keep the original group on refusal" + ); + } + assert!(stream.next().await.is_none()); + assert_eq!(plan.metrics().unwrap().spill_count().unwrap(), 0); + assert_eq!(runtime.memory_pool.reserved(), 0); + } Ok(()) } From cf73b48d001e15fdf577e69657f65724a85af82f Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 10 Nov 2026 05:54:43 +0000 Subject: [PATCH 4/5] Bound string witness scans and retained semi/anti join memory --- .../joins/sort_merge_join/bitwise_stream.rs | 55 ++-- .../sort_merge_join/semi_anti_summary.rs | 216 +++++++++------- .../semi_anti_summary_tests.rs | 234 ++++++++++++------ 3 files changed, 322 insertions(+), 183 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index c6b9c94d52774..cefdf7805f84a 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -809,18 +809,18 @@ impl BitwiseSortMergeJoinStream { } /// Apply the recognized comparison to this slice and cache a completed - /// summary if its owned values fit in the group's reservation. Returns - /// false if later slices of this key must use the original filter. + /// summary when a later slice can reuse it within the group's reservation. fn apply_summary( &mut self, comparison: SemiAntiComparison, outer: &RecordBatch, summary: &mut Option>, - ) -> Result { + may_continue: bool, + ) -> Result<()> { let matchable = outer.num_rows() - outer.column(comparison.outer_column).null_count(); if matchable == 0 { - return Ok(true); + return Ok(()); } if let Some(values) = summary { merge_matches( @@ -828,7 +828,7 @@ impl BitwiseSortMergeJoinStream { self.outer_offset, &comparison.evaluate(values, outer)?, ); - return Ok(true); + return Ok(()); } let start = self.outer_offset; @@ -839,24 +839,28 @@ impl BitwiseSortMergeJoinStream { .count_ones() == matchable })?; - if let Some(values) = values { + if let Some(values) = values + && may_continue + { let retained = values.memory_size(); - // A string prefix borrows the input. Admit its compact copy alongside - // the buffer before allocating it. On refusal the mask is already - // exact for this slice and later slices can still use the buffer. - let admitted = match values { - GroupSummary::Strings { .. } => self.reservation.try_grow(retained), - GroupSummary::Scalars { .. } => self.reservation.try_resize(retained), - }; - if admitted.is_err() { - return Ok(false); + // A skewed string can exceed this group's prorated buffer charge. + // Keep the buffer in that case so a later outer batch cannot hold + // more of the shared pool than the buffered join would have held. + if retained > self.reservation.size() { + return Ok(()); + } + if matches!(values, GroupSummary::String(Some(_))) { + // The input and its owned copy coexist until the buffer is cleared. + if self.reservation.try_grow(retained).is_err() { + return Ok(()); + } + self.peak_mem_used.set_max(self.reservation.size()); } - self.peak_mem_used.set_max(self.reservation.size()); *summary = Some(Box::new(values.into_owned())); self.clear_inner_key_group(); self.reservation.resize(retained); } - Ok(true) + Ok(()) } /// Process a key match with a filter. For each inner row in the buffered @@ -868,7 +872,7 @@ impl BitwiseSortMergeJoinStream { &mut self, spill: Option<&Arc>, summary: &mut Option>, - group_comparison: &mut Option, + group_comparison: Option, ) -> Result<()> { let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); @@ -894,10 +898,13 @@ impl BitwiseSortMergeJoinStream { let outer_batch = self.outer_batch.as_ref().unwrap(); let outer_slice = outer_batch.slice(outer_group_start, outer_group_len); - if let Some(comparison) = *group_comparison { - if !self.apply_summary(comparison, &outer_slice, summary)? { - *group_comparison = None; - } + if let Some(comparison) = group_comparison { + self.apply_summary( + comparison, + &outer_slice, + summary, + outer_group_end == num_outer, + )?; self.outer_offset = outer_group_end; return Ok(()); } @@ -987,7 +994,7 @@ impl BitwiseSortMergeJoinStream { emitter: &mut TryEmitter, ) -> Result<()> { let mut summary = None; - let mut group_comparison = self.semi_anti_comparison.filter(|_| { + let group_comparison = self.semi_anti_comparison.filter(|_| { spill.is_none() && self .inner_key_buffer @@ -1001,7 +1008,7 @@ impl BitwiseSortMergeJoinStream { self.process_key_match_with_filter( spill.as_ref(), &mut summary, - &mut group_comparison, + group_comparison, ) .await?; diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs index a2f926381a67c..05f6a673ce7cc 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary.rs @@ -14,7 +14,7 @@ use std::cmp::Ordering; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Datum, RecordBatch, Scalar}; use arrow::buffer::BooleanBuffer; use arrow::compute::kernels::cmp; -use arrow::compute::or; +use arrow::compute::{is_not_null, or}; use arrow::datatypes::{DataType, Schema}; use datafusion_common::{JoinSide, Result, ScalarValue}; use datafusion_expr::Operator; @@ -28,6 +28,9 @@ use crate::joins::utils::{JoinFilter, boolean_mask_from_filter}; const DIRECT_PROBES: usize = 2; /// Check the first small vector of rows before doubling the prefix length. const FIRST_PREFIX_ROWS: usize = 8; +/// Also check a pending string winner before the next value crosses this budget. +/// A single string can exceed the budget on its own. +const STRING_PREFIX_BYTES: usize = 64 * 1024; /// A recognized comparison, normalized to `outer OP inner`. #[derive(Debug, Clone, Copy)] @@ -39,34 +42,25 @@ pub(super) struct SemiAntiComparison { /// Cached once per buffered key group and reused across outer batches. #[derive(Debug)] -pub(super) enum GroupSummary, N = Scalar> { +pub(super) enum GroupSummary> { Scalars { - extreme: N, - other: Option, - }, - Strings { - extreme: Option, - other: Option, + extreme: Scalar, + other: Option>, }, + /// The range extremum or the single distinct non-null string for `!=`. + String(Option), + /// Two distinct non-null strings witness `!=` for every non-null outer row. + AllNonNull, } -/// While scanning, fixed-width values use DataFusion's min/max kernels and -/// string values borrow directly from the buffered input. -type RunningExtrema<'a> = GroupSummary<&'a str, ScalarValue>; - -impl<'a> RunningExtrema<'a> { - fn snapshot(&self) -> Result> { - Ok(match self { - Self::Scalars { extreme, other } => GroupSummary::Scalars { - extreme: extreme.to_scalar()?, - other: other.as_ref().map(ScalarValue::to_scalar).transpose()?, - }, - &Self::Strings { extreme, other } => GroupSummary::Strings { extreme, other }, +impl GroupSummary<&str> { + fn from_scalars(extreme: &ScalarValue, other: Option<&ScalarValue>) -> Result { + Ok(Self::Scalars { + extreme: extreme.to_scalar()?, + other: other.map(ScalarValue::to_scalar).transpose()?, }) } -} -impl GroupSummary<&str> { pub(super) fn memory_size(&self) -> usize { let array_size = |value: &Scalar| value.get().0.get_array_memory_size(); size_of::() @@ -74,9 +68,8 @@ impl GroupSummary<&str> { Self::Scalars { extreme, other } => { array_size(extreme) + other.as_ref().map_or(0, array_size) } - Self::Strings { extreme, other } => { - extreme.map_or(0, str::len) + other.map_or(0, str::len) - } + Self::String(value) => value.map_or(0, str::len), + Self::AllNonNull => 0, } } @@ -85,10 +78,8 @@ impl GroupSummary<&str> { pub(super) fn into_owned(self) -> GroupSummary { match self { Self::Scalars { extreme, other } => GroupSummary::Scalars { extreme, other }, - Self::Strings { extreme, other } => GroupSummary::Strings { - extreme: extreme.map(Box::from), - other: other.map(Box::from), - }, + Self::String(value) => GroupSummary::String(value.map(Box::from)), + Self::AllNonNull => GroupSummary::AllNonNull, } } } @@ -168,20 +159,25 @@ impl SemiAntiComparison { })) } - /// Apply two direct row probes, then check the extrema after 8, 16, 32, ... - /// rows. Only the newly reached rows enter each reduction, including when a - /// checkpoint crosses an input batch boundary. A witness at row k costs - /// O(k) inner rows even when the full group is much larger. + /// Probe the first two rows directly, then check extrema after 8, 16, 32, ... + /// rows, or earlier for wide strings. Each reduction visits only new rows + /// even when a checkpoint crosses an input batch. String `!=` probes the + /// first non-null value; a second distinct value determines every outer slice. /// /// The callback records the mask for this outer slice and returns true once - /// every matchable row has a witness. A complete scan returns the exact - /// summary for later outer slices. An early stop leaves the input buffered. + /// every matchable row has a witness. An exact summary can serve later + /// outer slices; an early stop without one leaves the input buffered. pub(super) fn scan<'a>( &self, batches: &'a [RecordBatch], outer: &RecordBatch, mut matches_all: impl FnMut(&BooleanArray) -> bool, ) -> Result>> { + let data_type = batches[0].column(self.inner_column).data_type(); + if data_type.is_string() && self.op == Operator::NotEq { + return self.scan_string_not_equal(batches, outer, matches_all); + } + let mut probes_left = DIRECT_PROBES; for batch in batches { let len = probes_left.min(batch.num_rows()); @@ -198,24 +194,17 @@ impl SemiAntiComparison { // The probes neither copy string payloads nor build reduction state. // Include their values in the first reduction for later outer slices. - let data_type = batches[0].column(self.inner_column).data_type(); let wants_max = matches!(self.op, Operator::Lt | Operator::LtEq); let order = if wants_max { Ordering::Greater } else { Ordering::Less }; - let mut state = if data_type.is_string() { - RunningExtrema::Strings { - extreme: None, - other: None, - } - } else { - let extreme = ScalarValue::try_from(data_type)?; - let other = (self.op == Operator::NotEq).then(|| extreme.clone()); - RunningExtrema::Scalars { extreme, other } - }; - + if data_type.is_string() { + return self.scan_string_range(batches, outer, order, matches_all); + } + let mut extreme = ScalarValue::try_from(data_type)?; + let mut other = (self.op == Operator::NotEq).then(|| extreme.clone()); let mut scanned = 0; let mut checkpoint = FIRST_PREFIX_ROWS; let mut changed = false; @@ -225,57 +214,112 @@ impl SemiAntiComparison { while start < column.len() { let len = (checkpoint - scanned).min(column.len() - start); let finished = index + 1 == batches.len() && start + len == column.len(); - match &mut state { - RunningExtrema::Scalars { extreme, other } => { - let values = column.slice(start, len); - let next = if wants_max { - max_batch(&values)? - } else { - min_batch(&values)? - }; - changed |= update_extreme(extreme, next, order); - if let Some(max) = other { - changed |= update_extreme( - max, - max_batch(&values)?, - Ordering::Greater, - ); - } - } - RunningExtrema::Strings { extreme, other } => { - for value in strings(column).skip(start).take(len).flatten() { - if extreme.is_none_or(|old| value.cmp(old) == order) { - *extreme = Some(value); - changed = true; - } - if self.op == Operator::NotEq - && other.is_none_or(|old| value > old) - { - *other = Some(value); - changed = true; - } - } - } + let values = column.slice(start, len); + let next = if wants_max { + max_batch(&values)? + } else { + min_batch(&values)? + }; + changed |= update_extreme(&mut extreme, next, order); + if let Some(max) = &mut other { + changed |= + update_extreme(max, max_batch(&values)?, Ordering::Greater); } start += len; scanned += len; if scanned == checkpoint && !finished { - if changed && matches_all(&self.evaluate(&state.snapshot()?, outer)?) - { - return Ok(None); + if changed { + let summary = + GroupSummary::from_scalars(&extreme, other.as_ref())?; + if matches_all(&self.evaluate(&summary, outer)?) { + return Ok(None); + } } changed = false; checkpoint = checkpoint.saturating_mul(2); } } } - let summary = state.snapshot()?; + let summary = GroupSummary::from_scalars(&extreme, other.as_ref())?; + if changed { + matches_all(&self.evaluate(&summary, outer)?); + } + Ok(Some(summary)) + } + + fn scan_string_range<'a>( + &self, + batches: &'a [RecordBatch], + outer: &RecordBatch, + order: Ordering, + mut matches_all: impl FnMut(&BooleanArray) -> bool, + ) -> Result>> { + let mut extreme = None; + let mut changed = false; + let mut checkpoint = FIRST_PREFIX_ROWS; + let mut bytes = 0usize; + let values = batches + .iter() + .flat_map(|b| strings(b.column(self.inner_column))); + for (row, value) in values.enumerate() { + let len = value.map_or(0, str::len); + if row == checkpoint || len > STRING_PREFIX_BYTES.saturating_sub(bytes) { + if changed + && matches_all(&self.evaluate(&GroupSummary::String(extreme), outer)?) + { + return Ok(None); + } + changed = false; + bytes = 0; + if row == checkpoint { + checkpoint = checkpoint.saturating_mul(2); + } + } + bytes = bytes.saturating_add(len); + if let Some(value) = value + && extreme.is_none_or(|old| value.cmp(old) == order) + { + extreme = Some(value); + // The direct probes already recorded matches for these rows. + changed |= row >= DIRECT_PROBES; + } + } + let summary = GroupSummary::String(extreme); if changed { matches_all(&self.evaluate(&summary, outer)?); } Ok(Some(summary)) } + fn scan_string_not_equal<'a>( + &self, + batches: &'a [RecordBatch], + outer: &RecordBatch, + mut matches_all: impl FnMut(&BooleanArray) -> bool, + ) -> Result>> { + let mut representative = None; + for batch in batches { + for (row, value) in strings(batch.column(self.inner_column)).enumerate() { + let Some(value) = value else { continue }; + match representative { + None => { + representative = Some(value); + if matches_all(&self.evaluate_inner_row(batch, row, outer)?) { + return Ok(None); + } + } + Some(first) if value != first => { + let summary = GroupSummary::AllNonNull; + matches_all(&self.evaluate(&summary, outer)?); + return Ok(Some(summary)); + } + _ => {} + } + } + } + Ok(Some(GroupSummary::String(representative))) + } + pub(super) fn evaluate>( &self, summary: &GroupSummary, @@ -289,16 +333,16 @@ impl SemiAntiComparison { None => Ok(matched), } } - GroupSummary::Strings { extreme, other } => { - let extreme = extreme.as_ref().map(AsRef::as_ref); - let other = other.as_ref().map(AsRef::as_ref); + GroupSummary::AllNonNull => Ok(is_not_null(outer.column(self.outer_column))?), + GroupSummary::String(value) => { + let extreme = value.as_ref().map(AsRef::as_ref); let hits = strings(outer.column(self.outer_column)).map(|value| { value.zip(extreme).is_some_and(|(x, y)| match self.op { Operator::Lt => x < y, Operator::LtEq => x <= y, Operator::Gt => x > y, Operator::GtEq => x >= y, - Operator::NotEq => x != y || other.is_some_and(|z| x != z), + Operator::NotEq => x != y, _ => unreachable!("invalid semi/anti summary comparison"), }) }); diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs index 146dce594735a..6fd4c26e272d5 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs @@ -18,6 +18,7 @@ use datafusion_common::{ use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::MemoryConsumer; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, NotExpr}; @@ -270,8 +271,8 @@ async fn large_strings_produce_the_same_results_when_the_buffer_spills() -> Resu ]; let outer = batch(Arc::new(StringArray::from(vec![values[0].as_str(), "m"])))?; let inner = batch(Arc::new(StringArray::from_iter_values(values.iter())))?; - // The two initial probes cross input batches. For != the first probe - // already matches "m", so its matched bit must survive the second probe. + // For != the first inner value matches "m". The third inner value + // also supplies a witness for the outer row equal to the first. for op in [Operator::NotEq, Operator::Lt] { for (limit, spill) in [(None, false), (Some(32 * 1024), true)] { let mut builder = RuntimeEnvBuilder::new(); @@ -296,19 +297,16 @@ async fn large_strings_produce_the_same_results_when_the_buffer_spills() -> Resu } #[tokio::test] -async fn not_equal_reserves_both_compact_string_extrema() -> Result<()> { - let min = "a".repeat(16 * 1024); - let max = "z".repeat(32 * 1024); - let middle = "m".repeat(1024); - let mut values = vec![middle.as_str(); 7]; - values[2] = min.as_str(); - values[6] = max.as_str(); - let values = Arc::new(StringArray::from(values)); +async fn uniform_not_equal_caches_one_string_when_the_copy_fits() -> Result<()> { + let long = "x".repeat(16 * 1024); + let values = Arc::new(StringArray::from(vec![long.as_str(); 7])); let inner = batch(Arc::clone(&values) as ArrayRef)?; - let outer = batch(Arc::new(StringArray::from(vec![middle.as_str()])))?; + let outer = batch(Arc::new(StringArray::from(vec![ + long.as_str(), + long.as_str(), + "different", + ])))?; - // Until the cache is admitted, even extrema from different batches - // borrow their original Arrow payloads. let comparison = SemiAntiComparison::try_new( &filter( binary(column(0), Operator::NotEq, column(1)), @@ -320,18 +318,62 @@ async fn not_equal_reserves_both_compact_string_extrema() -> Result<()> { )? .unwrap(); let batches = [inner.slice(0, 4), inner.slice(4, 3)]; - let Some(GroupSummary::Strings { - extreme: Some(lo), - other: Some(hi), - }) = comparison.scan(&batches, &outer, |_| false)? + let Some(GroupSummary::String(Some(value))) = + comparison.scan(&batches, &outer, |_| false)? else { - panic!("expected two string extrema"); + panic!("expected one borrowed representative"); }; - assert!(std::ptr::eq(lo, values.value(2))); - assert!(std::ptr::eq(hi, values.value(6))); + assert!((0..inner.num_rows()).any(|row| std::ptr::eq(value, values.value(row)))); + let buffered = inner.get_array_memory_size(); + for (headroom, expect_cache) in [(0, false), (long.len() * 3 / 2, true)] { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(buffered + headroom, 1.0) + .build_arc()?; + let ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(1)) + .with_runtime(Arc::clone(&runtime)), + ); + let plan = join( + input(&outer, 1)?, + input(&inner, 7)?, + JoinType::LeftAnti, + Operator::NotEq, + )?; + let mut stream = plan.execute(0, ctx)?; + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + let retained = runtime.memory_pool.reserved(); + if expect_cache { + assert!((long.len()..2 * long.len()).contains(&retained)); + } else { + assert_eq!(retained, buffered); + } + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [1]); + assert_eq!(runtime.memory_pool.reserved(), retained); + assert!(stream.next().await.is_none()); + assert_eq!(runtime.memory_pool.reserved(), 0); + } + Ok(()) +} + +#[tokio::test] +async fn mixed_not_equal_matches_all_later_non_nulls() -> Result<()> { + let first = "a".repeat(16 * 1024); + let second = "z".repeat(16 * 1024); + let mut values = vec![first.as_str(); 7]; + values[2] = second.as_str(); + let inner = batch(Arc::new(StringArray::from(values)))?; + let outer = batch(Arc::new(StringArray::from(vec![ + Some(first.as_str()), + None, + None, + Some(first.as_str()), + Some(second.as_str()), + None, + ])))?; let plan = join( - input(&outer, 1)?, - input(&inner, 7)?, + input(&outer, 2)?, + input(&inner, 2)?, JoinType::LeftSemi, Operator::NotEq, )?; @@ -339,25 +381,18 @@ async fn not_equal_reserves_both_compact_string_extrema() -> Result<()> { let pool = Arc::clone(ctx.memory_pool()); let mut stream = plan.execute(0, ctx)?; assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); - - let retained = pool.reserved(); - assert!( - retained >= min.len() + max.len(), - "both owned string payloads must be charged" - ); - assert!( - retained < inner.get_array_memory_size(), - "the parent buffer was released" - ); - drop(stream); + assert!(pool.reserved() < first.len()); + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [3]); + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [4]); + assert!(stream.next().await.is_none()); assert_eq!(pool.reserved(), 0); Ok(()) } #[tokio::test] -async fn oversized_extrema_fall_back_to_the_prorated_buffer() -> Result<()> { - // Group 1 ends inside the batch. Its seven rows fit at the group's - // per-row share, but even one of its two owned extrema exceeds the limit. +async fn oversized_string_leaves_headroom_for_another_consumer() -> Result<()> { + // Group 1 ends inside the batch. Its share is smaller than a single copy + // of the string, although there is room in the pool to make that copy. let long = "x".repeat(16 * 1024); let mut keys = vec![2; 256]; keys[..7].fill(1); @@ -371,43 +406,96 @@ async fn oversized_extrema_fall_back_to_the_prorated_buffer() -> Result<()> { .div_ceil(inner.num_rows()) * 7; assert!(buffered_size < long.len()); - for (limit, expect_cache) in [(buffered_size, false), (long.len() * 8, true)] { - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(limit, 1.0) - .with_disk_manager_builder( - DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), - ) - .build_arc()?; - let ctx = Arc::new( - TaskContext::default() - .with_session_config(SessionConfig::new().with_batch_size(1)) - .with_runtime(Arc::clone(&runtime)), - ); - let plan = join( - input(&outer, 1)?, - input(&inner, inner.num_rows())?, - JoinType::LeftAnti, - Operator::NotEq, - )?; - let mut stream = plan.execute(0, ctx)?; - assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); - // At the yield the same key may continue in the next outer batch. - let retained = runtime.memory_pool.reserved(); - if expect_cache { - assert!( - retained >= long.len(), - "owned string extrema must be charged" - ); - } else { - assert_eq!( - retained, buffered_size, - "keep the original group on refusal" - ); - } - assert!(stream.next().await.is_none()); - assert_eq!(plan.metrics().unwrap().spill_count().unwrap(), 0); - assert_eq!(runtime.memory_pool.reserved(), 0); - } + let limit = 8 * long.len(); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(1)) + .with_runtime(Arc::clone(&runtime)), + ); + let plan = join( + input(&outer, 1)?, + input(&inner, inner.num_rows())?, + JoinType::LeftAnti, + Operator::NotEq, + )?; + let mut stream = plan.execute(0, ctx)?; + assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + assert_eq!(runtime.memory_pool.reserved(), buffered_size); + let competitor = MemoryConsumer::new("other operator").register(&runtime.memory_pool); + competitor.try_grow(limit - buffered_size)?; + assert_eq!(runtime.memory_pool.reserved(), limit); + assert!(stream.next().await.is_none()); + assert_eq!(plan.metrics().unwrap().spill_count().unwrap(), 0); + assert_eq!(runtime.memory_pool.reserved(), competitor.size()); + drop(competitor); + assert_eq!(runtime.memory_pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn mid_batch_key_boundary_avoids_a_string_cache_copy() -> Result<()> { + let long = "z".repeat(16 * 1024); + let outer = keyed_batch( + vec![1, 2], + Arc::new(StringArray::from(vec![long.as_str(); 2])), + )?; + let inner = batch(Arc::new(StringArray::from(vec![long.as_str(); 7])))?; + let plan = join( + input(&outer, 2)?, + input(&inner, 7)?, + JoinType::LeftAnti, + Operator::Lt, + )?; + assert_eq!(ids(&plan, single_row_context()).await?, [0, 1]); + assert_eq!( + plan.metrics() + .unwrap() + .sum_by_name("peak_mem_used") + .unwrap() + .as_usize(), + inner.get_array_memory_size(), + "the last slice of this key has no later use for a copied extremum" + ); + Ok(()) +} + +#[test] +fn range_checks_the_third_string_before_a_wide_tail() -> Result<()> { + let wide = "z".repeat(256 * 1024); + let inner = batch(Arc::new(StringArray::from(vec![ + "x", "x", "y", &wide, &wide, &wide, &wide, &wide, + ])))?; + let outer = batch(Arc::new(StringArray::from(vec!["x", "z"])))?; + let comparison = SemiAntiComparison::try_new( + &filter(binary(column(0), Operator::Lt, column(1)), DataType::Utf8), + true, + &outer.schema(), + &inner.schema(), + )? + .unwrap(); + let mut masks = Vec::new(); + let batches = [inner]; + let summary = comparison.scan(&batches, &outer, |mask| { + masks.push(mask.clone()); + mask.value(0) + })?; + // "x" matches the third row. "z" would also match if any wide string + // had entered the reduction before this checkpoint. + assert_eq!( + masks, + [ + BooleanArray::from(vec![false, false]), + BooleanArray::from(vec![false, false]), + BooleanArray::from(vec![true, false]), + ] + ); + assert!(summary.is_none()); Ok(()) } From 3e676a9640cbf68f99a705c5b341a5b71a49293a Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 10 Nov 2026 17:53:53 +0000 Subject: [PATCH 5/5] Release exhausted merge join keys and remember cache refusals --- .../joins/sort_merge_join/bitwise_stream.rs | 34 +++++++++---- .../semi_anti_summary_tests.rs | 49 ++++++++++++++++++- 2 files changed, 72 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index cefdf7805f84a..95289068953d7 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -331,6 +331,13 @@ pub(crate) struct BitwiseSortMergeJoinStream { /// in the small-group benchmarks. const MIN_SUMMARY_ROWS: usize = 7; +enum SummaryCache { + Pending, + Cached(Box), + /// Keep scanning the buffered group without retrying the reservation. + Declined, +} + impl BitwiseSortMergeJoinStream { /// Creates a sort-merge join stream for semi, anti, and mark joins. /// @@ -605,6 +612,10 @@ impl BitwiseSortMergeJoinStream { // Release the inner input pipeline's resources. let inner_schema = self.inner.schema(); self.inner = Box::pin(EmptyRecordBatchStream::new(inner_schema)); + // The exhausted cursor and its comparators own the last batch's keys. + self.inner_key_arrays.clear(); + self.inner_self_cmp = None; + self.outer_inner_cmp = None; return Ok(false); } Some(Err(e)) => return Err(e), @@ -814,7 +825,7 @@ impl BitwiseSortMergeJoinStream { &mut self, comparison: SemiAntiComparison, outer: &RecordBatch, - summary: &mut Option>, + cache: &mut SummaryCache, may_continue: bool, ) -> Result<()> { let matchable = @@ -822,7 +833,7 @@ impl BitwiseSortMergeJoinStream { if matchable == 0 { return Ok(()); } - if let Some(values) = summary { + if let SummaryCache::Cached(values) = cache { merge_matches( &mut self.matched, self.outer_offset, @@ -841,22 +852,25 @@ impl BitwiseSortMergeJoinStream { })?; if let Some(values) = values && may_continue + && matches!(cache, SummaryCache::Pending) { let retained = values.memory_size(); // A skewed string can exceed this group's prorated buffer charge. // Keep the buffer in that case so a later outer batch cannot hold // more of the shared pool than the buffered join would have held. if retained > self.reservation.size() { + *cache = SummaryCache::Declined; return Ok(()); } if matches!(values, GroupSummary::String(Some(_))) { // The input and its owned copy coexist until the buffer is cleared. if self.reservation.try_grow(retained).is_err() { + *cache = SummaryCache::Declined; return Ok(()); } self.peak_mem_used.set_max(self.reservation.size()); } - *summary = Some(Box::new(values.into_owned())); + *cache = SummaryCache::Cached(Box::new(values.into_owned())); self.clear_inner_key_group(); self.reservation.resize(retained); } @@ -871,14 +885,16 @@ impl BitwiseSortMergeJoinStream { async fn process_key_match_with_filter( &mut self, spill: Option<&Arc>, - summary: &mut Option>, + cache: &mut SummaryCache, group_comparison: Option, ) -> Result<()> { let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); // buffer_inner_key_group must be called before this function debug_assert!( - !self.inner_key_buffer.is_empty() || spill.is_some() || summary.is_some(), + !self.inner_key_buffer.is_empty() + || spill.is_some() + || matches!(cache, SummaryCache::Cached(_)), "process_key_match_with_filter called with no inner key data" ); debug_assert!( @@ -902,7 +918,7 @@ impl BitwiseSortMergeJoinStream { self.apply_summary( comparison, &outer_slice, - summary, + cache, outer_group_end == num_outer, )?; self.outer_offset = outer_group_end; @@ -993,7 +1009,7 @@ impl BitwiseSortMergeJoinStream { spill: Option>, emitter: &mut TryEmitter, ) -> Result<()> { - let mut summary = None; + let mut cache = SummaryCache::Pending; let group_comparison = self.semi_anti_comparison.filter(|_| { spill.is_none() && self @@ -1007,7 +1023,7 @@ impl BitwiseSortMergeJoinStream { loop { self.process_key_match_with_filter( spill.as_ref(), - &mut summary, + &mut cache, group_comparison, ) .await?; @@ -1038,7 +1054,7 @@ impl BitwiseSortMergeJoinStream { } } - if summary.is_some() { + if matches!(cache, SummaryCache::Cached(_)) { self.reservation.free(); } self.clear_inner_key_group(); diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs index 6fd4c26e272d5..25e60d8087bf6 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/semi_anti_summary_tests.rs @@ -297,7 +297,7 @@ async fn large_strings_produce_the_same_results_when_the_buffer_spills() -> Resu } #[tokio::test] -async fn uniform_not_equal_caches_one_string_when_the_copy_fits() -> Result<()> { +async fn uniform_not_equal_caches_one_string_when_the_first_copy_fits() -> Result<()> { let long = "x".repeat(16 * 1024); let values = Arc::new(StringArray::from(vec![long.as_str(); 7])); let inner = batch(Arc::clone(&values) as ArrayRef)?; @@ -325,10 +325,16 @@ async fn uniform_not_equal_caches_one_string_when_the_copy_fits() -> Result<()> }; assert!((0..inner.num_rows()).any(|row| std::ptr::eq(value, values.value(row)))); let buffered = inner.get_array_memory_size(); - for (headroom, expect_cache) in [(0, false), (long.len() * 3 / 2, true)] { + let headroom = long.len() * 3 / 2; + for expect_cache in [false, true] { let runtime = RuntimeEnvBuilder::new() .with_memory_limit(buffered + headroom, 1.0) .build_arc()?; + let competitor = + MemoryConsumer::new("other operator").register(&runtime.memory_pool); + if !expect_cache { + competitor.try_grow(headroom)?; + } let ctx = Arc::new( TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(1)) @@ -342,6 +348,9 @@ async fn uniform_not_equal_caches_one_string_when_the_copy_fits() -> Result<()> )?; let mut stream = plan.execute(0, ctx)?; assert_eq!(row_ids(&stream.next().await.transpose()?.unwrap()), [0]); + // A refused cache would now fit. The next outer batch of the same + // key must keep using the buffer without retrying the reservation. + drop(competitor); let retained = runtime.memory_pool.reserved(); if expect_cache { assert!((long.len()..2 * long.len()).contains(&retained)); @@ -595,3 +604,39 @@ async fn cached_summary_is_released_on_outer_error_or_drop() -> Result<()> { } Ok(()) } + +#[tokio::test] +async fn cached_summary_releases_exhausted_inner_keys_before_yielding() -> Result<()> { + use arrow::array::AsArray; + + let key = "x".repeat(16 * 1024); + let make_batch = |rows| { + RecordBatch::try_from_iter([ + ( + "key", + Arc::new(StringArray::from(vec![key.as_str(); rows])) as ArrayRef, + ), + ( + "value", + Arc::new(Int32Array::from(vec![0; rows])) as ArrayRef, + ), + ]) + }; + let outer = make_batch(2)?; + let inner = make_batch(7)?; + let key_values = inner.column(0).as_string::().values().clone(); + let mut stream = join( + input(&outer, 1)?, + input(&inner, 7)?, + JoinType::LeftAnti, + Operator::Lt, + )? + .execute(0, single_row_context())?; + drop(inner); + + assert_eq!(stream.next().await.transpose()?.unwrap().num_rows(), 1); + assert_eq!(key_values.strong_count(), 1); + assert_eq!(stream.next().await.transpose()?.unwrap().num_rows(), 1); + assert!(stream.next().await.is_none()); + Ok(()) +}