From 5ec1790897998f9ef1b21dd710f685dfd93fe771 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 24 Jul 2026 21:28:40 +0530 Subject: [PATCH 1/6] feat(pwmj): support LeftSemi/LeftAnti existence joins via classic scan --- datafusion/core/src/physical_planner.rs | 8 +- datafusion/physical-plan/Cargo.toml | 5 + .../benches/piecewise_merge_join_semi_anti.rs | 228 +++++ .../piecewise_merge_join/classic_join.rs | 896 ++++++++++++++++-- .../src/joins/piecewise_merge_join/exec.rs | 65 +- .../src/joins/piecewise_merge_join/utils.rs | 19 +- datafusion/sqllogictest/test_files/pwmj.slt | 89 ++ 7 files changed, 1223 insertions(+), 87 deletions(-) create mode 100644 datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 4e914556b4cc0..749afad9da7a9 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1678,11 +1678,13 @@ impl DefaultPhysicalPlanner { Arc::new(CrossJoinExec::new(physical_left, physical_right)) } else if num_range_filters == 1 && total_filters == 1 + // PWMJ supports classic joins and Left Semi/Anti existence joins. + // Right Semi/Anti and Mark joins are not implemented yet (they + // would require swapping the inputs so the marked side is buffered), + // so exclude them here and let them fall back to NestedLoopJoin. && !matches!( join_type, - JoinType::LeftSemi - | JoinType::RightSemi - | JoinType::LeftAnti + JoinType::RightSemi | JoinType::RightAnti | JoinType::LeftMark | JoinType::RightMark diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 58c2f0d7da537..8c22a2741c6be 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -139,6 +139,11 @@ harness = false name = "hash_join_semi_anti" required-features = ["test_utils"] +[[bench]] +harness = false +name = "piecewise_merge_join_semi_anti" +required-features = ["test_utils"] + [[bench]] harness = false name = "multi_group_by" diff --git a/datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs b/datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs new file mode 100644 index 0000000000000..f2f053e24e656 --- /dev/null +++ b/datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs @@ -0,0 +1,228 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Criterion benchmark comparing existence (LeftSemi / LeftAnti) joins over a single +//! range predicate (`left.key < right.key`) evaluated two ways: +//! +//! - `PiecewiseMergeJoinExec` (with the required `SortExec` on the buffered/left side, +//! as the physical planner would insert), and +//! - `NestedLoopJoinExec`, which is the fallback used when +//! `enable_piecewise_merge_join` is off. +//! +//! Both plans compute the same result, so this measures the win from routing an +//! inequality-correlated `EXISTS` / `NOT EXISTS` to PWMJ instead of the O(n*m) +//! nested-loop join. The `SortExec` is included on the PWMJ side because it is a real +//! cost of that plan. +//! +//! ## Axes +//! - **join type**: LeftSemi (`EXISTS`) and LeftAnti (`NOT EXISTS`). +//! - **selectivity**: the fraction of left rows that have at least one matching right +//! row, controlled by shifting the right-side key range. Semi output size grows with +//! selectivity; Anti output size shrinks. + +use std::sync::Arc; + +use arrow::array::{Int32Array, RecordBatch}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::JoinSide; +use datafusion_common::JoinType; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column}; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion_physical_plan::joins::{NestedLoopJoinExec, PiecewiseMergeJoinExec}; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr, collect}; +use tokio::runtime::Runtime; + +/// Two-column schema: (`key`, `payload`). +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("payload", DataType::Int32, false), + ])) +} + +/// Build a single-partition input of `num_rows` rows. Keys are drawn from +/// `[key_offset, key_offset + key_span)` in a fixed, reproducible pattern (no RNG so +/// the benchmark is deterministic). +fn build_exec( + num_rows: usize, + key_offset: i32, + key_span: i32, + schema: &SchemaRef, +) -> Arc { + let keys: Vec = (0..num_rows) + .map(|i| key_offset + (i as i32 * 2_654_435_761u32 as i32).rem_euclid(key_span)) + .collect(); + let payload: Vec = (0..num_rows as i32).collect(); + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(payload)), + ], + ) + .unwrap(); + + // Slice into 8192-row batches to mirror a realistic streamed input. + let batch_size = 8192; + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + batches.push(batch.slice(offset, len)); + offset += len; + } + TestMemoryExec::try_new_exec(&[batches], Arc::clone(schema), None).unwrap() +} + +/// `PiecewiseMergeJoinExec` over `left.key < right.key`, with the required `SortExec` +/// on the buffered (left) side. `<` requires the buffered side sorted descending. +fn pwmj_plan( + left: Arc, + right: Arc, + join_type: JoinType, +) -> Arc { + let sort = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("key", 0)), + SortOptions::new(true, true), + )]) + .unwrap(); + let sorted_left = Arc::new(SortExec::new(sort, left)); + + let on: (Arc, Arc) = ( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + ); + Arc::new( + PiecewiseMergeJoinExec::try_new( + sorted_left, + right, + on, + Operator::Lt, + join_type, + 1, + ) + .unwrap(), + ) +} + +/// `NestedLoopJoinExec` over the same `left.key < right.key` predicate. +fn nlj_plan( + left: Arc, + right: Arc, + join_type: JoinType, +) -> Arc { + let intermediate_schema = Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("key", DataType::Int32, false), + ]); + let expr = Arc::new(BinaryExpr::new( + Arc::new(Column::new("key", 0)), + Operator::Lt, + Arc::new(Column::new("key", 1)), + )) as Arc; + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let filter = JoinFilter::new(expr, column_indices, Arc::new(intermediate_schema)); + Arc::new( + NestedLoopJoinExec::try_new(left, right, Some(filter), &join_type, None).unwrap(), + ) +} + +fn run(plan: Arc, rt: &Runtime) -> usize { + let task_ctx = Arc::new(TaskContext::default()); + rt.block_on(async { + let batches = collect(plan, task_ctx).await.unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + }) +} + +fn bench_pwmj_semi_anti(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let s = schema(); + + // Left (buffered) is deliberately smaller than right (streamed); the streamed side + // drives the loop in both operators. + let left_rows = 20_000; + let right_rows = 20_000; + let key_span = 10_000; + + // Selectivity is set by how far the right key range sits above the left range. + // - "high": right keys mostly above left keys -> most left rows match (Semi large) + // - "low": right keys mostly below left keys -> few left rows match (Anti large) + let regimes: [(&str, i32); 2] = [("sel_high", key_span), ("sel_low", -key_span)]; + + let mut group = c.benchmark_group("pwmj_vs_nlj_semi_anti"); + // Nested-loop is O(n*m); keep sample counts modest so the suite finishes. + group.sample_size(10); + + for (regime, right_offset) in regimes { + for join_type in [JoinType::LeftSemi, JoinType::LeftAnti] { + let jt = match join_type { + JoinType::LeftSemi => "semi", + JoinType::LeftAnti => "anti", + _ => unreachable!(), + }; + + let build_inputs = || { + ( + build_exec(left_rows, 0, key_span, &s), + build_exec(right_rows, right_offset, key_span, &s), + ) + }; + + group.bench_function( + BenchmarkId::new(format!("pwmj_{jt}_{regime}"), right_rows), + |b| { + b.iter(|| { + let (left, right) = build_inputs(); + run(pwmj_plan(left, right, join_type), &rt) + }) + }, + ); + + group.bench_function( + BenchmarkId::new(format!("nlj_{jt}_{regime}"), right_rows), + |b| { + b.iter(|| { + let (left, right) = build_inputs(); + run(nlj_plan(left, right, join_type), &rt) + }) + }, + ); + } + } + + group.finish(); +} + +criterion_group!(benches, bench_pwmj_semi_anti); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs index 50ef78f18bf65..0ea641ab3c8c0 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs @@ -36,7 +36,9 @@ use std::{sync::Arc, task::Poll}; use crate::handle_state; use crate::joins::piecewise_merge_join::exec::{BufferedSide, BufferedSideReadyState}; -use crate::joins::piecewise_merge_join::utils::need_produce_result_in_final; +use crate::joins::piecewise_merge_join::utils::{ + is_supported_existence_join, need_produce_result_in_final, +}; use crate::joins::utils::{BuildProbeJoinMetrics, StatefulStreamResult}; use crate::joins::utils::{JoinKeyComparator, get_final_indices_from_shared_bitmap}; use crate::stream::EmptyRecordBatchStream; @@ -358,14 +360,18 @@ impl ClassicPWMJStream { take_record_batch(buffered_data.batch(), &buffered_indices)?; let mut buffered_columns = new_buffered_batch.columns().to_vec(); - let streamed_columns: Vec = self - .streamed_schema - .fields() - .iter() - .map(|f| new_null_array(f.data_type(), new_buffered_batch.num_rows())) - .collect(); - - buffered_columns.extend(streamed_columns); + // LeftSemi/LeftAnti emit only the buffered (left) columns; classic outer joins + // (Left/Full) additionally emit null-padded streamed columns for unmatched rows. + if !is_supported_existence_join(self.join_type) { + let streamed_columns: Vec = self + .streamed_schema + .fields() + .iter() + .map(|f| new_null_array(f.data_type(), new_buffered_batch.num_rows())) + .collect(); + + buffered_columns.extend(streamed_columns); + } let batch = RecordBatch::try_new(Arc::clone(&self.schema), buffered_columns)?; @@ -415,6 +421,14 @@ struct BatchProcessState { continue_process: bool, // Skip nulls processed_null_count: bool, + // For existence joins (LeftSemi/LeftAnti): the lowest buffered index marked in the + // bitmap so far, across all stream batches processed by this stream. Buffered rows + // `[existence_min_marked..len)` are already marked, so later batches only need to mark + // the new prefix `[buffer_idx..existence_min_marked)`. This bounds total marking work + // to O(buffered) per partition instead of O(num_batches * buffered). `usize::MAX` + // means nothing has been marked yet. It intentionally persists across batches and is + // not cleared by `reset`. + existence_min_marked: usize, } impl BatchProcessState { @@ -427,6 +441,7 @@ impl BatchProcessState { found: false, continue_process: true, processed_null_count: false, + existence_min_marked: usize::MAX, } } @@ -437,6 +452,7 @@ impl BatchProcessState { self.found = false; self.continue_process = true; self.processed_null_count = false; + // `existence_min_marked` is deliberately not reset: it accumulates across batches. } } @@ -485,67 +501,15 @@ fn resolve_classic_join( // Our buffer_idx variable allows us to start probing on the buffered side where we last matched // in the previous stream row. - for row_idx in stream_idx..stream_batch.batch.num_rows() { + 'stream_rows: for row_idx in stream_idx..stream_batch.batch.num_rows() { while buffer_idx < buffered_len { let compare = cmp.compare(row_idx, buffer_idx); - // If we find a match we append all indices and move to the next stream row index - match operator { - Operator::Gt | Operator::Lt => { - if compare == Ordering::Less { - batch_process_state.found = true; - let count = buffered_len - buffer_idx; - - let batch = build_matched_indices_and_set_buffered_bitmap( - (buffer_idx, count), - (row_idx, count), - buffered_side, - stream_batch, - join_type, - join_schema, - )?; - - batch_process_state.output_batches.push_batch(batch)?; - - // Flush batch and update pointers if we have a completed batch - if let Some(batch) = - batch_process_state.output_batches.next_completed_batch() - { - batch_process_state.found = false; - batch_process_state.start_buffer_idx = buffer_idx; - batch_process_state.start_stream_idx = row_idx + 1; - return Ok(batch); - } - - break; - } - } + // Determine whether the current stream row matches the current buffered row. + let is_match = match operator { + Operator::Gt | Operator::Lt => compare == Ordering::Less, Operator::GtEq | Operator::LtEq => { - if matches!(compare, Ordering::Equal | Ordering::Less) { - batch_process_state.found = true; - let count = buffered_len - buffer_idx; - let batch = build_matched_indices_and_set_buffered_bitmap( - (buffer_idx, count), - (row_idx, count), - buffered_side, - stream_batch, - join_type, - join_schema, - )?; - - // Flush batch and update pointers if we have a completed batch - batch_process_state.output_batches.push_batch(batch)?; - if let Some(batch) = - batch_process_state.output_batches.next_completed_batch() - { - batch_process_state.found = false; - batch_process_state.start_buffer_idx = buffer_idx; - batch_process_state.start_stream_idx = row_idx + 1; - return Ok(batch); - } - - break; - } + matches!(compare, Ordering::Equal | Ordering::Less) } _ => { return internal_err!( @@ -555,6 +519,62 @@ fn resolve_classic_join( } }; + if is_match { + batch_process_state.found = true; + + // Existence joins (LeftSemi/LeftAnti) only need to know which buffered + // rows have at least one match. `buffer_idx` advances monotonically within + // a batch, so the first match sits at the smallest buffered index reached; + // because the buffered side is sorted, every row in `[buffer_idx..end]` + // matches too. Marking that suffix once therefore covers every match this + // batch can produce, so we mark it and stop scanning the batch entirely + // (any later stream row would only re-mark a subset). Output is produced + // later from the bitmap. + // + // Across batches we only mark the not-yet-marked prefix + // `[buffer_idx..existence_min_marked)`: rows from `existence_min_marked` + // onward were already marked by an earlier batch. This keeps total marking + // work at O(buffered) per partition even with many stream batches. + if is_supported_existence_join(join_type) { + let upper = + batch_process_state.existence_min_marked.min(buffered_len); + if buffer_idx < upper { + let mut bitmap = + buffered_side.buffered_data.visited_indices_bitmap.lock(); + for i in buffer_idx..upper { + bitmap.set_bit(i, true); + } + batch_process_state.existence_min_marked = buffer_idx; + } + break 'stream_rows; + } + + let count = buffered_len - buffer_idx; + + let batch = build_matched_indices_and_set_buffered_bitmap( + (buffer_idx, count), + (row_idx, count), + buffered_side, + stream_batch, + join_type, + join_schema, + )?; + + batch_process_state.output_batches.push_batch(batch)?; + + // Flush batch and update pointers if we have a completed batch + if let Some(batch) = + batch_process_state.output_batches.next_completed_batch() + { + batch_process_state.found = false; + batch_process_state.start_buffer_idx = buffer_idx; + batch_process_state.start_stream_idx = row_idx + 1; + return Ok(batch); + } + + break; + } + // Increment buffer_idx after every row buffer_idx += 1; } @@ -1543,4 +1563,748 @@ mod tests { "); Ok(()) } + + // Builds a table whose middle (`b`) column is nullable so existence-join NULL + // semantics can be exercised. + fn build_table_nullable_b( + a: (&str, &Vec), + b: (&str, &Vec>), + c: (&str, &Vec), + ) -> Arc { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Int32, false), + Field::new(b.0, DataType::Int32, true), + Field::new(c.0, DataType::Int32, false), + ]); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(arrow::array::Int32Array::from(a.1.clone())), + Arc::new(arrow::array::Int32Array::from(b.1.clone())), + Arc::new(arrow::array::Int32Array::from(c.1.clone())), + ], + ) + .unwrap(); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + // LeftSemi keeps buffered (left) rows that have at least one match, and outputs only + // the left columns. Buffered side is pre-sorted in the operator's required order + // (descending for `<`) because these unit tests bypass the optimizer's SortExec. + #[tokio::test] + async fn join_left_semi_less_than() -> Result<()> { + // left.b1 < right.b1 ; left is buffered, sorted descending for `<` + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![5, 2, 1]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 2 | 2 | 8 | + | 3 | 1 | 9 | + +----+----+----+ + "); + Ok(()) + } + + // LeftAnti keeps buffered (left) rows that have NO match. + #[tokio::test] + async fn join_left_anti_less_than() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![5, 2, 1]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 5 | 7 | + +----+----+----+ + "); + Ok(()) + } + + // `>` uses ascending buffered order. + #[tokio::test] + async fn join_left_semi_greater_than() -> Result<()> { + // left.b1 > right.b1 ; left is buffered, sorted ascending for `>` + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 5]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 3 | 5 | 9 | + +----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_left_anti_greater_than() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 5]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::LeftAnti).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 7 | + | 2 | 2 | 8 | + +----+----+----+ + "); + Ok(()) + } + + // `<=` includes equal values. + #[tokio::test] + async fn join_left_semi_less_than_equal() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![5, 4, 2]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::LtEq, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 2 | 4 | 8 | + | 3 | 2 | 9 | + +----+----+----+ + "); + Ok(()) + } + + // `>=` includes equal values (ascending buffered order). + #[tokio::test] + async fn join_left_semi_greater_than_equal() -> Result<()> { + // Keep left rows with some right b1 <= left b1. left b1 = {1,3,5}, + // right b1 = {3,4}. 1 has none; 3>=3; 5>={3,4}. + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 3, 5]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20]), + ("b1", &vec![3, 4]), + ("c2", &vec![70, 80]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::GtEq, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 2 | 3 | 8 | + | 3 | 5 | 9 | + +----+----+----+ + "); + Ok(()) + } + + // Empty buffered (left) side: LeftSemi produces nothing, LeftAnti produces nothing. + #[tokio::test] + async fn join_left_semi_empty_left() -> Result<()> { + let left = build_table( + ("a1", &Vec::::new()), + ("b1", &Vec::::new()), + ("c1", &Vec::::new()), + ); + let right = build_table( + ("a2", &vec![1, 2]), + ("b1", &vec![1, 2]), + ("c2", &vec![1, 2]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + +----+----+----+ + "); + Ok(()) + } + + // Empty streamed (right) side: no right row can satisfy the predicate, so LeftSemi is + // empty and LeftAnti returns all buffered rows. + #[tokio::test] + async fn join_left_anti_empty_right() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![3, 2, 1]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &Vec::::new()), + ("b1", &Vec::::new()), + ("c2", &Vec::::new()), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 3 | 7 | + | 2 | 2 | 8 | + | 3 | 1 | 9 | + +----+----+----+ + "); + Ok(()) + } + + // NULL join keys never satisfy a comparison predicate, so a null-keyed left row is + // excluded from LeftSemi and included in LeftAnti. Buffered side sorted descending + // with nulls first for `<`. + #[tokio::test] + async fn join_left_semi_less_than_left_nulls() -> Result<()> { + let left = build_table_nullable_b( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![None, Some(5), Some(1)]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20]), + ("b1", &vec![2, 4]), + ("c2", &vec![70, 80]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 3 | 1 | 9 | + +----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_left_anti_less_than_left_nulls() -> Result<()> { + let left = build_table_nullable_b( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![None, Some(5), Some(1)]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20]), + ("b1", &vec![2, 4]), + ("c2", &vec![70, 80]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | | 7 | + | 2 | 5 | 8 | + +----+----+----+ + "); + Ok(()) + } + + // Existence join over a streamed side split across multiple partitions: the final + // pass that emits the bitmap result runs once, on the last streamed partition to + // finish. This exercises that coordination for LeftSemi. + #[tokio::test] + async fn join_left_semi_multi_partition_stream() -> Result<()> { + // Buffered (left) side, sorted ascending for `>`. + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 3, 5, 7]), + ("c1", &vec![10, 20, 30, 40]), + ); + + // Streamed (right) side split across two partitions. + let right_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + let right_p0 = build_table_i32( + ("a2", &vec![10, 20]), + ("b1", &vec![2, 4]), + ("c2", &vec![70, 80]), + ); + let right_p1 = + build_table_i32(("a2", &vec![30]), ("b1", &vec![6]), ("c2", &vec![90])); + let right = TestMemoryExec::try_new_exec( + &[vec![right_p0], vec![right_p1]], + Arc::new(right_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + // Keep buffered rows with some smaller streamed b1: streamed b1 = {2,4,6}, + // buffered b1 = {1,3,5,7}. 3>2, 5>{2,4}, 7>{2,4,6}. 1 has none. + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftSemi, + 2, + )?; + + let task_ctx = Arc::new(TaskContext::default()); + let mut batches = Vec::new(); + for p in 0..2 { + let stream = join.execute(p, Arc::clone(&task_ctx))?; + batches.extend(common::collect(stream).await?); + } + // Sort output for a stable comparison (partitions may interleave). + let sorted = arrow::compute::concat_batches(&join.schema(), batches.iter())?; + let indices = sort_to_indices(sorted.column(1).as_ref(), None, None)?; + let sorted = take_record_batch(&sorted, &indices)?; + + assert_snapshot!(batches_to_string(&[sorted]), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 2 | 3 | 20 | + | 3 | 5 | 30 | + | 4 | 7 | 40 | + +----+----+----+ + "); + Ok(()) + } + + // LeftAnti `>=` — complement of the LeftSemi `>=` case, closes the operator matrix. + #[tokio::test] + async fn join_left_anti_greater_than_equal() -> Result<()> { + // Keep left rows with NO right b1 <= left b1. left b1 = {1,3,5}, + // right b1 = {3,4}. Only 1 has no smaller-or-equal right value. + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 3, 5]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20]), + ("b1", &vec![3, 4]), + ("c2", &vec![70, 80]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::GtEq, JoinType::LeftAnti).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 7 | + +----+----+----+ + "); + Ok(()) + } + + // Every streamed (right) join key is NULL, so no comparison can ever match: + // LeftAnti must return all buffered rows. + #[tokio::test] + async fn join_left_anti_all_right_nulls() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![3, 2, 1]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table_nullable_b( + ("a2", &vec![10, 20]), + ("b1", &vec![None, None]), + ("c2", &vec![70, 80]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 3 | 7 | + | 2 | 2 | 8 | + | 3 | 1 | 9 | + +----+----+----+ + "); + Ok(()) + } + + // Non-integer key type: existence joins reuse the same comparator as classic joins, + // so a Date32 key should behave identically. + #[tokio::test] + async fn join_left_semi_date32_greater_than() -> Result<()> { + // left dates > some right date. left b1 = {19100, 19105, 19110}, + // right b1 = {19102, 19108}. 19100 has none; 19105>19102; 19110>{19102,19108}. + let left = build_date_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![19100, 19105, 19110]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_date_table( + ("a2", &vec![10, 20]), + ("b1", &vec![19102, 19108]), + ("c2", &vec![70, 80]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +------------+------------+------------+ + | a1 | b1 | c1 | + +------------+------------+------------+ + | 1970-01-03 | 2022-04-23 | 1970-01-09 | + | 1970-01-04 | 2022-04-28 | 1970-01-10 | + +------------+------------+------------+ + "); + Ok(()) + } + + // Streamed (right) side delivered as multiple batches within a single partition. The + // bitmap must accumulate matches across batches; here the second batch contributes + // matches the first did not, so dropping either batch would change the result. + #[tokio::test] + async fn join_left_semi_multi_batch_stream() -> Result<()> { + // Buffered (left) side, sorted ascending for `>`. + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 3, 5, 7]), + ("c1", &vec![10, 20, 30, 40]), + ); + + let right_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // Batch 1 (b1=6) marks only b1=7; batch 2 (b1=2) additionally marks b1={3,5,7}. + let batch1 = + build_table_i32(("a2", &vec![10]), ("b1", &vec![6]), ("c2", &vec![70])); + let batch2 = + build_table_i32(("a2", &vec![20]), ("b1", &vec![2]), ("c2", &vec![80])); + // Single partition containing two batches. + let right = TestMemoryExec::try_new_exec( + &[vec![batch1, batch2]], + Arc::new(right_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 2 | 3 | 20 | + | 3 | 5 | 30 | + | 4 | 7 | 40 | + +----+----+----+ + "); + Ok(()) + } + + // Exercises the cross-batch low-water mark, including its skip branch. Batch order is + // chosen so that: + // batch 1 (b1=2) matches low -> marks buffered b1={3,5,7}, watermark drops to idx 1 + // batch 2 (b1=6) matches high -> would only mark b1=7 (already marked) -> skipped + // batch 3 (b1=0) matches even lower -> marks the remaining b1=1 + // Result is all four buffered rows; a change that either re-marked greedily or + // skipped batch 3 would be caught here. + #[tokio::test] + async fn join_left_semi_multi_batch_watermark_skip() -> Result<()> { + // Buffered (left) side, sorted ascending for `>`. + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 3, 5, 7]), + ("c1", &vec![10, 20, 30, 40]), + ); + + let right_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + let batch1 = + build_table_i32(("a2", &vec![10]), ("b1", &vec![2]), ("c2", &vec![70])); + let batch2 = + build_table_i32(("a2", &vec![20]), ("b1", &vec![6]), ("c2", &vec![80])); + let batch3 = + build_table_i32(("a2", &vec![30]), ("b1", &vec![0]), ("c2", &vec![90])); + // Single partition, three batches processed in order. + let right = TestMemoryExec::try_new_exec( + &[vec![batch1, batch2, batch3]], + Arc::new(right_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 10 | + | 2 | 3 | 20 | + | 3 | 5 | 30 | + | 4 | 7 | 40 | + +----+----+----+ + "); + Ok(()) + } + + // LeftAnti over a streamed side split across multiple partitions AND multiple batches. + // Only LeftSemi had multi-partition coverage; the anti final pass emits the *unmarked* + // buffered rows, so this exercises the complementary bitmap read plus the multi-partition + // final-pass counter. Buffered b1 = {1,3,5,7}; streamed b1 = {2,4,6} spread across two + // partitions. For `>`, buffered rows with some smaller streamed value match (3,5,7); only + // b1=1 has none, so LeftAnti must return exactly that row. + #[tokio::test] + async fn join_left_anti_multi_partition_multi_batch_stream() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 3, 5, 7]), + ("c1", &vec![10, 20, 30, 40]), + ); + + let right_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // Partition 0 delivers b1=2 then b1=6 as two separate batches; partition 1 delivers b1=4. + let p0_b0 = + build_table_i32(("a2", &vec![10]), ("b1", &vec![2]), ("c2", &vec![70])); + let p0_b1 = + build_table_i32(("a2", &vec![20]), ("b1", &vec![6]), ("c2", &vec![80])); + let p1_b0 = + build_table_i32(("a2", &vec![30]), ("b1", &vec![4]), ("c2", &vec![90])); + let right = TestMemoryExec::try_new_exec( + &[vec![p0_b0, p0_b1], vec![p1_b0]], + Arc::new(right_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftAnti, + 2, + )?; + + let task_ctx = Arc::new(TaskContext::default()); + let mut batches = Vec::new(); + for p in 0..2 { + let stream = join.execute(p, Arc::clone(&task_ctx))?; + batches.extend(common::collect(stream).await?); + } + let out = arrow::compute::concat_batches(&join.schema(), batches.iter())?; + + assert_snapshot!(batches_to_string(&[out]), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 10 | + +----+----+----+ + "); + Ok(()) + } + + // Existence output is produced entirely in the final pass, which runs exactly once on + // the last streamed partition to finish. That is coordinated by a counter seeded from + // the streamed side's partition count. Here the streamed side is single-partition while + // `target_partitions` is 4: `execute` is called once, so the counter must be seeded + // from the streamed partition count (1) for the final pass to run — otherwise `LeftSemi` + // would emit nothing. + #[tokio::test] + async fn left_semi_multi_partition_final_pass() -> Result<()> { + // left.b1 > right.b1 ; buffered sorted ascending for `>`. Keep left rows with some + // smaller right b1: left b1 = {1,3,4}, right b1 = {3,2,1}. 3>{2,1}, 4>{3,2,1}; 1 has none. + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 3, 4]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![3, 2, 1]), + ("c2", &vec![70, 80, 90]), + ); + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + // target_partitions = 4, but the streamed side is single-partition. + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftSemi, + 4, + )?; + + let task_ctx = Arc::new(TaskContext::default()); + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + let out = arrow::compute::concat_batches(&join.schema(), batches.iter())?; + // The two matching left rows (b1=3, b1=4) must be emitted by the final pass. + assert_eq!(out.num_rows(), 2); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 5ec564295ece1..a91fd4a5b2c23 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -48,6 +48,7 @@ use crate::joins::piecewise_merge_join::classic_join::{ }; use crate::joins::piecewise_merge_join::utils::{ build_visited_indices_map, is_existence_join, is_right_existence_join, + is_supported_existence_join, }; use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; @@ -162,17 +163,34 @@ use crate::{ /// ``` /// /// ## Existence Joins (Semi, Anti, Mark) -/// Existence joins are made magnitudes of times faster with a `PiecewiseMergeJoin` as we only need to find -/// the min/max value of the streamed side to be able to emit all matches on the buffered side. By putting -/// the side we need to mark onto the sorted buffer side, we can emit all these matches at once. +/// Currently only `LeftSemi` and `LeftAnti` are supported. For these the marked side is +/// already the left (buffered) side, so no input swap is needed. `RightSemi`/`RightAnti` +/// and `Mark` joins are not yet implemented and are rejected in [`Self::try_new`]. /// -/// For less than operations (`<`) both inputs are to be sorted in descending order and vice versa for greater -/// than (`>`) operations. `SortExec` is used to enforce sorting on the buffered side and streamed side does not -/// need to be sorted due to only needing to find the min/max. +/// Existence joins reuse the classic-join scan (see above) rather than a dedicated +/// algorithm: for each streamed row we advance the pointer on the sorted buffered side to +/// the first match, and because the buffered side is sorted, every buffered row from that +/// position onward also matches. Instead of materializing the joined output, we simply +/// mark those buffered rows in the visited-indices bitmap and continue. Once all streamed +/// partitions are processed, the final pass emits the result directly from the bitmap: +/// `LeftSemi` emits the marked buffered rows, `LeftAnti` emits the unmarked ones (rows +/// whose join key is NULL are never marked, so they are correctly excluded from `Semi` +/// and included in `Anti`). Only the buffered (left) columns are produced. /// -/// For Left Semi, Anti, and Mark joins we swap the inputs so that the marked side is on the buffered side. +/// Reusing the classic scan keeps a single code path and works for the same predicates +/// classic joins support, at the cost of still sorting the streamed side. /// -/// The pseudocode for the algorithm looks like this: +/// The pseudocode looks like this: +/// +/// ```text +/// for stream_row in sorted_stream_batch: +/// advance buffer_idx to the first buffered row that matches stream_row +/// if a match is found: +/// mark buffered[buffer_idx..end] in the bitmap +/// // final pass, once all partitions finish: +/// // LeftSemi -> emit buffered rows where bit == 1 +/// // LeftAnti -> emit buffered rows where bit == 0 +/// ``` /// /// ```text /// // Using the example of a less than `<` operation @@ -294,10 +312,12 @@ impl PiecewiseMergeJoinExec { join_type: JoinType, num_partitions: usize, ) -> Result { - // TODO: Implement existence joins for PiecewiseMergeJoin - if is_existence_join(join_type) { + // Left Semi/Anti reuse the classic scan (marked side is already the buffered + // side, no input swap needed). Right existence joins and Mark joins are not yet + // supported. + if is_existence_join(join_type) && !is_supported_existence_join(join_type) { return not_impl_err!( - "Existence Joins are currently not supported for PiecewiseMergeJoin" + "Existence join {join_type} is currently not supported for PiecewiseMergeJoin" ); } @@ -569,6 +589,14 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { let on_streamed = Arc::clone(&self.on.1); let metrics = BuildProbeJoinMetrics::new(partition, &self.metrics); + // The final pass over unmatched/existence rows must run exactly once, on the + // last streamed partition to finish. That is coordinated by an atomic counter + // seeded with the number of streamed partitions that will actually call + // `execute`, which is the streamed side's output partition count — not the + // planner's `target_partitions` (they can differ, e.g. when the streamed input + // has a single partition), otherwise the counter never reaches 1 and the final + // pass is skipped. + let streamed_partitions = self.streamed.output_partitioning().partition_count(); let buffered_fut = self.buffered_fut.try_once(|| { let reservation = MemoryConsumer::new("PiecewiseMergeJoinInput") .register(context.memory_pool()); @@ -580,7 +608,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { metrics.clone(), reservation, build_visited_indices_map(self.join_type), - self.num_partitions, + streamed_partitions, )) })?; @@ -588,9 +616,16 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { let batch_size = context.session_config().batch_size(); - // TODO: Add existence joins + this is guarded at physical planner - if is_existence_join(self.join_type()) { - unreachable!() + // Right existence joins and Mark joins are rejected in `try_new`; every other + // join type (classic + supported Left Semi/Anti existence) runs the classic scan, + // which marks the buffered bitmap and emits the existence result from it. + if is_existence_join(self.join_type()) + && !is_supported_existence_join(self.join_type()) + { + internal_err!( + "PiecewiseMergeJoin does not support existence join {} (should have been rejected in try_new)", + self.join_type() + ) } else { Ok(Box::pin(ClassicPWMJStream::try_new( Arc::clone(&self.schema), diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs index 5bbb496322b5f..88a4b9e81ab62 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs @@ -38,10 +38,23 @@ pub(super) fn is_existence_join(join_type: JoinType) -> bool { ) } -// Returns boolean to check if the join type needs to record -// buffered side matches for classic joins +// Returns boolean for whether the join is a left existence join that is currently +// supported by `PiecewiseMergeJoin`. These do not require swapping the inputs: the +// marked (left) side is already the buffered side, so the classic scan can mark the +// buffered bitmap and the final indices are emitted from it. +pub(super) fn is_supported_existence_join(join_type: JoinType) -> bool { + matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti) +} + +// Returns true if the join type produces its output in the final pass over the buffered +// side (using the visited-indices bitmap) rather than while scanning stream batches: +// `Left`/`Full` emit unmatched buffered rows, and `LeftSemi`/`LeftAnti` emit the +// matched/unmatched buffered rows respectively. pub(super) fn need_produce_result_in_final(join_type: JoinType) -> bool { - matches!(join_type, JoinType::Full | JoinType::Left) + matches!( + join_type, + JoinType::Full | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti + ) } // Returns boolean for whether or not we need to build the buffered side diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 9789c0e4e5392..4f8ae04bc8d98 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -342,5 +342,94 @@ ORDER BY 1,2; 1 3 2 3 +# ------------------------------------------------------------------ +# Existence joins (LeftSemi / LeftAnti) via PiecewiseMergeJoin +# ------------------------------------------------------------------ + +# EXISTS with a range correlation -> LeftSemi. Keep t1 rows that have at least one +# smaller t2_id: 22>11, 33>11, 44>{11,22}. 11 has none. +query I +SELECT t1.t1_id +FROM join_t1 t1 +WHERE EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +22 +33 +44 + +query TT +EXPLAIN +SELECT t1.t1_id +FROM join_t1 t1 +WHERE EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +logical_plan +01)Sort: t1.t1_id ASC NULLS LAST +02)--LeftSemi Join: Filter: t1.t1_id > __correlated_sq_1.t2_id +03)----SubqueryAlias: t1 +04)------TableScan: join_t1 projection=[t1_id] +05)----SubqueryAlias: __correlated_sq_1 +06)------SubqueryAlias: t2 +07)--------TableScan: join_t2 projection=[t2_id] +physical_plan +01)SortExec: expr=[t1_id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(t1_id > t2_id) +03)----SortExec: expr=[t1_id@0 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] + +# NOT EXISTS with a range correlation -> LeftAnti. Complement of the above. +query I +SELECT t1.t1_id +FROM join_t1 t1 +WHERE NOT EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +11 + +query TT +EXPLAIN +SELECT t1.t1_id +FROM join_t1 t1 +WHERE NOT EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +logical_plan +01)Sort: t1.t1_id ASC NULLS LAST +02)--LeftAnti Join: Filter: t1.t1_id > __correlated_sq_1.t2_id +03)----SubqueryAlias: t1 +04)------TableScan: join_t1 projection=[t1_id] +05)----SubqueryAlias: __correlated_sq_1 +06)------SubqueryAlias: t2 +07)--------TableScan: join_t2 projection=[t2_id] +physical_plan +01)SortExec: expr=[t1_id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--PiecewiseMergeJoin: operator=Gt, join_type=LeftAnti, on=(t1_id > t2_id) +03)----SortExec: expr=[t1_id@0 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] + +# NULL join key on the semi/anti (buffered) side never matches: excluded from EXISTS, +# included in NOT EXISTS. null_join_t1 = {1, 2, NULL}, null_join_t2 = {1, NULL, 3}. +# EXISTS t1.id > t2.id : 2>1 -> {2}. 1 and NULL have no smaller t2. +query I +SELECT t1.id +FROM null_join_t1 t1 +WHERE EXISTS (SELECT 1 FROM null_join_t2 t2 WHERE t1.id > t2.id) +ORDER BY 1; +---- +2 + +query I +SELECT t1.id +FROM null_join_t1 t1 +WHERE NOT EXISTS (SELECT 1 FROM null_join_t2 t2 WHERE t1.id > t2.id) +ORDER BY 1 NULLS FIRST; +---- +NULL +1 + statement ok set datafusion.optimizer.enable_piecewise_merge_join = false; From 86c70ce0dc578c80a104bacac7d8f9dfb1bcdbf3 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 7 Aug 2026 14:18:51 +0530 Subject: [PATCH 2/6] Remove bench code --- datafusion/physical-plan/Cargo.toml | 5 - .../benches/piecewise_merge_join_semi_anti.rs | 228 ------------------ datafusion/sqllogictest/test_files/pwmj.slt | 89 ------- 3 files changed, 322 deletions(-) delete mode 100644 datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 8c22a2741c6be..58c2f0d7da537 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -139,11 +139,6 @@ harness = false name = "hash_join_semi_anti" required-features = ["test_utils"] -[[bench]] -harness = false -name = "piecewise_merge_join_semi_anti" -required-features = ["test_utils"] - [[bench]] harness = false name = "multi_group_by" diff --git a/datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs b/datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs deleted file mode 100644 index f2f053e24e656..0000000000000 --- a/datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs +++ /dev/null @@ -1,228 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Criterion benchmark comparing existence (LeftSemi / LeftAnti) joins over a single -//! range predicate (`left.key < right.key`) evaluated two ways: -//! -//! - `PiecewiseMergeJoinExec` (with the required `SortExec` on the buffered/left side, -//! as the physical planner would insert), and -//! - `NestedLoopJoinExec`, which is the fallback used when -//! `enable_piecewise_merge_join` is off. -//! -//! Both plans compute the same result, so this measures the win from routing an -//! inequality-correlated `EXISTS` / `NOT EXISTS` to PWMJ instead of the O(n*m) -//! nested-loop join. The `SortExec` is included on the PWMJ side because it is a real -//! cost of that plan. -//! -//! ## Axes -//! - **join type**: LeftSemi (`EXISTS`) and LeftAnti (`NOT EXISTS`). -//! - **selectivity**: the fraction of left rows that have at least one matching right -//! row, controlled by shifting the right-side key range. Semi output size grows with -//! selectivity; Anti output size shrinks. - -use std::sync::Arc; - -use arrow::array::{Int32Array, RecordBatch}; -use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; -use datafusion_common::JoinSide; -use datafusion_common::JoinType; -use datafusion_execution::TaskContext; -use datafusion_expr::Operator; -use datafusion_physical_expr::expressions::{BinaryExpr, Column}; -use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; -use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; -use datafusion_physical_plan::joins::{NestedLoopJoinExec, PiecewiseMergeJoinExec}; -use datafusion_physical_plan::sorts::sort::SortExec; -use datafusion_physical_plan::test::TestMemoryExec; -use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr, collect}; -use tokio::runtime::Runtime; - -/// Two-column schema: (`key`, `payload`). -fn schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("key", DataType::Int32, false), - Field::new("payload", DataType::Int32, false), - ])) -} - -/// Build a single-partition input of `num_rows` rows. Keys are drawn from -/// `[key_offset, key_offset + key_span)` in a fixed, reproducible pattern (no RNG so -/// the benchmark is deterministic). -fn build_exec( - num_rows: usize, - key_offset: i32, - key_span: i32, - schema: &SchemaRef, -) -> Arc { - let keys: Vec = (0..num_rows) - .map(|i| key_offset + (i as i32 * 2_654_435_761u32 as i32).rem_euclid(key_span)) - .collect(); - let payload: Vec = (0..num_rows as i32).collect(); - let batch = RecordBatch::try_new( - Arc::clone(schema), - vec![ - Arc::new(Int32Array::from(keys)), - Arc::new(Int32Array::from(payload)), - ], - ) - .unwrap(); - - // Slice into 8192-row batches to mirror a realistic streamed input. - let batch_size = 8192; - let mut batches = Vec::new(); - let mut offset = 0; - while offset < batch.num_rows() { - let len = (batch.num_rows() - offset).min(batch_size); - batches.push(batch.slice(offset, len)); - offset += len; - } - TestMemoryExec::try_new_exec(&[batches], Arc::clone(schema), None).unwrap() -} - -/// `PiecewiseMergeJoinExec` over `left.key < right.key`, with the required `SortExec` -/// on the buffered (left) side. `<` requires the buffered side sorted descending. -fn pwmj_plan( - left: Arc, - right: Arc, - join_type: JoinType, -) -> Arc { - let sort = LexOrdering::new(vec![PhysicalSortExpr::new( - Arc::new(Column::new("key", 0)), - SortOptions::new(true, true), - )]) - .unwrap(); - let sorted_left = Arc::new(SortExec::new(sort, left)); - - let on: (Arc, Arc) = ( - Arc::new(Column::new("key", 0)), - Arc::new(Column::new("key", 0)), - ); - Arc::new( - PiecewiseMergeJoinExec::try_new( - sorted_left, - right, - on, - Operator::Lt, - join_type, - 1, - ) - .unwrap(), - ) -} - -/// `NestedLoopJoinExec` over the same `left.key < right.key` predicate. -fn nlj_plan( - left: Arc, - right: Arc, - join_type: JoinType, -) -> Arc { - let intermediate_schema = Schema::new(vec![ - Field::new("key", DataType::Int32, false), - Field::new("key", DataType::Int32, false), - ]); - let expr = Arc::new(BinaryExpr::new( - Arc::new(Column::new("key", 0)), - Operator::Lt, - Arc::new(Column::new("key", 1)), - )) as Arc; - let column_indices = vec![ - ColumnIndex { - index: 0, - side: JoinSide::Left, - }, - ColumnIndex { - index: 0, - side: JoinSide::Right, - }, - ]; - let filter = JoinFilter::new(expr, column_indices, Arc::new(intermediate_schema)); - Arc::new( - NestedLoopJoinExec::try_new(left, right, Some(filter), &join_type, None).unwrap(), - ) -} - -fn run(plan: Arc, rt: &Runtime) -> usize { - let task_ctx = Arc::new(TaskContext::default()); - rt.block_on(async { - let batches = collect(plan, task_ctx).await.unwrap(); - batches.iter().map(|b| b.num_rows()).sum() - }) -} - -fn bench_pwmj_semi_anti(c: &mut Criterion) { - let rt = Runtime::new().unwrap(); - let s = schema(); - - // Left (buffered) is deliberately smaller than right (streamed); the streamed side - // drives the loop in both operators. - let left_rows = 20_000; - let right_rows = 20_000; - let key_span = 10_000; - - // Selectivity is set by how far the right key range sits above the left range. - // - "high": right keys mostly above left keys -> most left rows match (Semi large) - // - "low": right keys mostly below left keys -> few left rows match (Anti large) - let regimes: [(&str, i32); 2] = [("sel_high", key_span), ("sel_low", -key_span)]; - - let mut group = c.benchmark_group("pwmj_vs_nlj_semi_anti"); - // Nested-loop is O(n*m); keep sample counts modest so the suite finishes. - group.sample_size(10); - - for (regime, right_offset) in regimes { - for join_type in [JoinType::LeftSemi, JoinType::LeftAnti] { - let jt = match join_type { - JoinType::LeftSemi => "semi", - JoinType::LeftAnti => "anti", - _ => unreachable!(), - }; - - let build_inputs = || { - ( - build_exec(left_rows, 0, key_span, &s), - build_exec(right_rows, right_offset, key_span, &s), - ) - }; - - group.bench_function( - BenchmarkId::new(format!("pwmj_{jt}_{regime}"), right_rows), - |b| { - b.iter(|| { - let (left, right) = build_inputs(); - run(pwmj_plan(left, right, join_type), &rt) - }) - }, - ); - - group.bench_function( - BenchmarkId::new(format!("nlj_{jt}_{regime}"), right_rows), - |b| { - b.iter(|| { - let (left, right) = build_inputs(); - run(nlj_plan(left, right, join_type), &rt) - }) - }, - ); - } - } - - group.finish(); -} - -criterion_group!(benches, bench_pwmj_semi_anti); -criterion_main!(benches); diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 4f8ae04bc8d98..9789c0e4e5392 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -342,94 +342,5 @@ ORDER BY 1,2; 1 3 2 3 -# ------------------------------------------------------------------ -# Existence joins (LeftSemi / LeftAnti) via PiecewiseMergeJoin -# ------------------------------------------------------------------ - -# EXISTS with a range correlation -> LeftSemi. Keep t1 rows that have at least one -# smaller t2_id: 22>11, 33>11, 44>{11,22}. 11 has none. -query I -SELECT t1.t1_id -FROM join_t1 t1 -WHERE EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) -ORDER BY 1; ----- -22 -33 -44 - -query TT -EXPLAIN -SELECT t1.t1_id -FROM join_t1 t1 -WHERE EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) -ORDER BY 1; ----- -logical_plan -01)Sort: t1.t1_id ASC NULLS LAST -02)--LeftSemi Join: Filter: t1.t1_id > __correlated_sq_1.t2_id -03)----SubqueryAlias: t1 -04)------TableScan: join_t1 projection=[t1_id] -05)----SubqueryAlias: __correlated_sq_1 -06)------SubqueryAlias: t2 -07)--------TableScan: join_t2 projection=[t2_id] -physical_plan -01)SortExec: expr=[t1_id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(t1_id > t2_id) -03)----SortExec: expr=[t1_id@0 ASC], preserve_partitioning=[false] -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)----DataSourceExec: partitions=1, partition_sizes=[1] - -# NOT EXISTS with a range correlation -> LeftAnti. Complement of the above. -query I -SELECT t1.t1_id -FROM join_t1 t1 -WHERE NOT EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) -ORDER BY 1; ----- -11 - -query TT -EXPLAIN -SELECT t1.t1_id -FROM join_t1 t1 -WHERE NOT EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) -ORDER BY 1; ----- -logical_plan -01)Sort: t1.t1_id ASC NULLS LAST -02)--LeftAnti Join: Filter: t1.t1_id > __correlated_sq_1.t2_id -03)----SubqueryAlias: t1 -04)------TableScan: join_t1 projection=[t1_id] -05)----SubqueryAlias: __correlated_sq_1 -06)------SubqueryAlias: t2 -07)--------TableScan: join_t2 projection=[t2_id] -physical_plan -01)SortExec: expr=[t1_id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--PiecewiseMergeJoin: operator=Gt, join_type=LeftAnti, on=(t1_id > t2_id) -03)----SortExec: expr=[t1_id@0 ASC], preserve_partitioning=[false] -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)----DataSourceExec: partitions=1, partition_sizes=[1] - -# NULL join key on the semi/anti (buffered) side never matches: excluded from EXISTS, -# included in NOT EXISTS. null_join_t1 = {1, 2, NULL}, null_join_t2 = {1, NULL, 3}. -# EXISTS t1.id > t2.id : 2>1 -> {2}. 1 and NULL have no smaller t2. -query I -SELECT t1.id -FROM null_join_t1 t1 -WHERE EXISTS (SELECT 1 FROM null_join_t2 t2 WHERE t1.id > t2.id) -ORDER BY 1; ----- -2 - -query I -SELECT t1.id -FROM null_join_t1 t1 -WHERE NOT EXISTS (SELECT 1 FROM null_join_t2 t2 WHERE t1.id > t2.id) -ORDER BY 1 NULLS FIRST; ----- -NULL -1 - statement ok set datafusion.optimizer.enable_piecewise_merge_join = false; From a16a4895e229ef08e6e0b314d6da04b0cdd0ef32 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 7 Aug 2026 15:28:27 +0530 Subject: [PATCH 3/6] Adds slt test --- datafusion/sqllogictest/test_files/pwmj.slt | 89 +++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 9789c0e4e5392..4f8ae04bc8d98 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -342,5 +342,94 @@ ORDER BY 1,2; 1 3 2 3 +# ------------------------------------------------------------------ +# Existence joins (LeftSemi / LeftAnti) via PiecewiseMergeJoin +# ------------------------------------------------------------------ + +# EXISTS with a range correlation -> LeftSemi. Keep t1 rows that have at least one +# smaller t2_id: 22>11, 33>11, 44>{11,22}. 11 has none. +query I +SELECT t1.t1_id +FROM join_t1 t1 +WHERE EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +22 +33 +44 + +query TT +EXPLAIN +SELECT t1.t1_id +FROM join_t1 t1 +WHERE EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +logical_plan +01)Sort: t1.t1_id ASC NULLS LAST +02)--LeftSemi Join: Filter: t1.t1_id > __correlated_sq_1.t2_id +03)----SubqueryAlias: t1 +04)------TableScan: join_t1 projection=[t1_id] +05)----SubqueryAlias: __correlated_sq_1 +06)------SubqueryAlias: t2 +07)--------TableScan: join_t2 projection=[t2_id] +physical_plan +01)SortExec: expr=[t1_id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(t1_id > t2_id) +03)----SortExec: expr=[t1_id@0 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] + +# NOT EXISTS with a range correlation -> LeftAnti. Complement of the above. +query I +SELECT t1.t1_id +FROM join_t1 t1 +WHERE NOT EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +11 + +query TT +EXPLAIN +SELECT t1.t1_id +FROM join_t1 t1 +WHERE NOT EXISTS (SELECT 1 FROM join_t2 t2 WHERE t1.t1_id > t2.t2_id) +ORDER BY 1; +---- +logical_plan +01)Sort: t1.t1_id ASC NULLS LAST +02)--LeftAnti Join: Filter: t1.t1_id > __correlated_sq_1.t2_id +03)----SubqueryAlias: t1 +04)------TableScan: join_t1 projection=[t1_id] +05)----SubqueryAlias: __correlated_sq_1 +06)------SubqueryAlias: t2 +07)--------TableScan: join_t2 projection=[t2_id] +physical_plan +01)SortExec: expr=[t1_id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--PiecewiseMergeJoin: operator=Gt, join_type=LeftAnti, on=(t1_id > t2_id) +03)----SortExec: expr=[t1_id@0 ASC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] + +# NULL join key on the semi/anti (buffered) side never matches: excluded from EXISTS, +# included in NOT EXISTS. null_join_t1 = {1, 2, NULL}, null_join_t2 = {1, NULL, 3}. +# EXISTS t1.id > t2.id : 2>1 -> {2}. 1 and NULL have no smaller t2. +query I +SELECT t1.id +FROM null_join_t1 t1 +WHERE EXISTS (SELECT 1 FROM null_join_t2 t2 WHERE t1.id > t2.id) +ORDER BY 1; +---- +2 + +query I +SELECT t1.id +FROM null_join_t1 t1 +WHERE NOT EXISTS (SELECT 1 FROM null_join_t2 t2 WHERE t1.id > t2.id) +ORDER BY 1 NULLS FIRST; +---- +NULL +1 + statement ok set datafusion.optimizer.enable_piecewise_merge_join = false; From 1dda519bdf46ccac2466a0ccee714a49d6f43398 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Mon, 10 Aug 2026 19:42:58 +0530 Subject: [PATCH 4/6] split existence joins into their own ExistencePWMJStream --- .../piecewise_merge_join/classic_join.rs | 896 ++---------------- .../src/joins/piecewise_merge_join/exec.rs | 80 +- .../piecewise_merge_join/existence_join.rs | 811 ++++++++++++++++ .../src/joins/piecewise_merge_join/mod.rs | 1 + .../src/joins/piecewise_merge_join/utils.rs | 21 +- datafusion/sqllogictest/test_files/pwmj.slt | 615 ++++++++++++ 6 files changed, 1544 insertions(+), 880 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs index 0ea641ab3c8c0..50ef78f18bf65 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs @@ -36,9 +36,7 @@ use std::{sync::Arc, task::Poll}; use crate::handle_state; use crate::joins::piecewise_merge_join::exec::{BufferedSide, BufferedSideReadyState}; -use crate::joins::piecewise_merge_join::utils::{ - is_supported_existence_join, need_produce_result_in_final, -}; +use crate::joins::piecewise_merge_join::utils::need_produce_result_in_final; use crate::joins::utils::{BuildProbeJoinMetrics, StatefulStreamResult}; use crate::joins::utils::{JoinKeyComparator, get_final_indices_from_shared_bitmap}; use crate::stream::EmptyRecordBatchStream; @@ -360,18 +358,14 @@ impl ClassicPWMJStream { take_record_batch(buffered_data.batch(), &buffered_indices)?; let mut buffered_columns = new_buffered_batch.columns().to_vec(); - // LeftSemi/LeftAnti emit only the buffered (left) columns; classic outer joins - // (Left/Full) additionally emit null-padded streamed columns for unmatched rows. - if !is_supported_existence_join(self.join_type) { - let streamed_columns: Vec = self - .streamed_schema - .fields() - .iter() - .map(|f| new_null_array(f.data_type(), new_buffered_batch.num_rows())) - .collect(); - - buffered_columns.extend(streamed_columns); - } + let streamed_columns: Vec = self + .streamed_schema + .fields() + .iter() + .map(|f| new_null_array(f.data_type(), new_buffered_batch.num_rows())) + .collect(); + + buffered_columns.extend(streamed_columns); let batch = RecordBatch::try_new(Arc::clone(&self.schema), buffered_columns)?; @@ -421,14 +415,6 @@ struct BatchProcessState { continue_process: bool, // Skip nulls processed_null_count: bool, - // For existence joins (LeftSemi/LeftAnti): the lowest buffered index marked in the - // bitmap so far, across all stream batches processed by this stream. Buffered rows - // `[existence_min_marked..len)` are already marked, so later batches only need to mark - // the new prefix `[buffer_idx..existence_min_marked)`. This bounds total marking work - // to O(buffered) per partition instead of O(num_batches * buffered). `usize::MAX` - // means nothing has been marked yet. It intentionally persists across batches and is - // not cleared by `reset`. - existence_min_marked: usize, } impl BatchProcessState { @@ -441,7 +427,6 @@ impl BatchProcessState { found: false, continue_process: true, processed_null_count: false, - existence_min_marked: usize::MAX, } } @@ -452,7 +437,6 @@ impl BatchProcessState { self.found = false; self.continue_process = true; self.processed_null_count = false; - // `existence_min_marked` is deliberately not reset: it accumulates across batches. } } @@ -501,15 +485,67 @@ fn resolve_classic_join( // Our buffer_idx variable allows us to start probing on the buffered side where we last matched // in the previous stream row. - 'stream_rows: for row_idx in stream_idx..stream_batch.batch.num_rows() { + for row_idx in stream_idx..stream_batch.batch.num_rows() { while buffer_idx < buffered_len { let compare = cmp.compare(row_idx, buffer_idx); - // Determine whether the current stream row matches the current buffered row. - let is_match = match operator { - Operator::Gt | Operator::Lt => compare == Ordering::Less, + // If we find a match we append all indices and move to the next stream row index + match operator { + Operator::Gt | Operator::Lt => { + if compare == Ordering::Less { + batch_process_state.found = true; + let count = buffered_len - buffer_idx; + + let batch = build_matched_indices_and_set_buffered_bitmap( + (buffer_idx, count), + (row_idx, count), + buffered_side, + stream_batch, + join_type, + join_schema, + )?; + + batch_process_state.output_batches.push_batch(batch)?; + + // Flush batch and update pointers if we have a completed batch + if let Some(batch) = + batch_process_state.output_batches.next_completed_batch() + { + batch_process_state.found = false; + batch_process_state.start_buffer_idx = buffer_idx; + batch_process_state.start_stream_idx = row_idx + 1; + return Ok(batch); + } + + break; + } + } Operator::GtEq | Operator::LtEq => { - matches!(compare, Ordering::Equal | Ordering::Less) + if matches!(compare, Ordering::Equal | Ordering::Less) { + batch_process_state.found = true; + let count = buffered_len - buffer_idx; + let batch = build_matched_indices_and_set_buffered_bitmap( + (buffer_idx, count), + (row_idx, count), + buffered_side, + stream_batch, + join_type, + join_schema, + )?; + + // Flush batch and update pointers if we have a completed batch + batch_process_state.output_batches.push_batch(batch)?; + if let Some(batch) = + batch_process_state.output_batches.next_completed_batch() + { + batch_process_state.found = false; + batch_process_state.start_buffer_idx = buffer_idx; + batch_process_state.start_stream_idx = row_idx + 1; + return Ok(batch); + } + + break; + } } _ => { return internal_err!( @@ -519,62 +555,6 @@ fn resolve_classic_join( } }; - if is_match { - batch_process_state.found = true; - - // Existence joins (LeftSemi/LeftAnti) only need to know which buffered - // rows have at least one match. `buffer_idx` advances monotonically within - // a batch, so the first match sits at the smallest buffered index reached; - // because the buffered side is sorted, every row in `[buffer_idx..end]` - // matches too. Marking that suffix once therefore covers every match this - // batch can produce, so we mark it and stop scanning the batch entirely - // (any later stream row would only re-mark a subset). Output is produced - // later from the bitmap. - // - // Across batches we only mark the not-yet-marked prefix - // `[buffer_idx..existence_min_marked)`: rows from `existence_min_marked` - // onward were already marked by an earlier batch. This keeps total marking - // work at O(buffered) per partition even with many stream batches. - if is_supported_existence_join(join_type) { - let upper = - batch_process_state.existence_min_marked.min(buffered_len); - if buffer_idx < upper { - let mut bitmap = - buffered_side.buffered_data.visited_indices_bitmap.lock(); - for i in buffer_idx..upper { - bitmap.set_bit(i, true); - } - batch_process_state.existence_min_marked = buffer_idx; - } - break 'stream_rows; - } - - let count = buffered_len - buffer_idx; - - let batch = build_matched_indices_and_set_buffered_bitmap( - (buffer_idx, count), - (row_idx, count), - buffered_side, - stream_batch, - join_type, - join_schema, - )?; - - batch_process_state.output_batches.push_batch(batch)?; - - // Flush batch and update pointers if we have a completed batch - if let Some(batch) = - batch_process_state.output_batches.next_completed_batch() - { - batch_process_state.found = false; - batch_process_state.start_buffer_idx = buffer_idx; - batch_process_state.start_stream_idx = row_idx + 1; - return Ok(batch); - } - - break; - } - // Increment buffer_idx after every row buffer_idx += 1; } @@ -1563,748 +1543,4 @@ mod tests { "); Ok(()) } - - // Builds a table whose middle (`b`) column is nullable so existence-join NULL - // semantics can be exercised. - fn build_table_nullable_b( - a: (&str, &Vec), - b: (&str, &Vec>), - c: (&str, &Vec), - ) -> Arc { - let schema = Schema::new(vec![ - Field::new(a.0, DataType::Int32, false), - Field::new(b.0, DataType::Int32, true), - Field::new(c.0, DataType::Int32, false), - ]); - let batch = RecordBatch::try_new( - Arc::new(schema), - vec![ - Arc::new(arrow::array::Int32Array::from(a.1.clone())), - Arc::new(arrow::array::Int32Array::from(b.1.clone())), - Arc::new(arrow::array::Int32Array::from(c.1.clone())), - ], - ) - .unwrap(); - let schema = batch.schema(); - TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() - } - - // LeftSemi keeps buffered (left) rows that have at least one match, and outputs only - // the left columns. Buffered side is pre-sorted in the operator's required order - // (descending for `<`) because these unit tests bypass the optimizer's SortExec. - #[tokio::test] - async fn join_left_semi_less_than() -> Result<()> { - // left.b1 < right.b1 ; left is buffered, sorted descending for `<` - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![5, 2, 1]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![2, 3, 4]), - ("c2", &vec![70, 80, 90]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Lt, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 2 | 2 | 8 | - | 3 | 1 | 9 | - +----+----+----+ - "); - Ok(()) - } - - // LeftAnti keeps buffered (left) rows that have NO match. - #[tokio::test] - async fn join_left_anti_less_than() -> Result<()> { - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![5, 2, 1]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![2, 3, 4]), - ("c2", &vec![70, 80, 90]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | 5 | 7 | - +----+----+----+ - "); - Ok(()) - } - - // `>` uses ascending buffered order. - #[tokio::test] - async fn join_left_semi_greater_than() -> Result<()> { - // left.b1 > right.b1 ; left is buffered, sorted ascending for `>` - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![1, 2, 5]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![2, 3, 4]), - ("c2", &vec![70, 80, 90]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 3 | 5 | 9 | - +----+----+----+ - "); - Ok(()) - } - - #[tokio::test] - async fn join_left_anti_greater_than() -> Result<()> { - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![1, 2, 5]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![2, 3, 4]), - ("c2", &vec![70, 80, 90]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Gt, JoinType::LeftAnti).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | 1 | 7 | - | 2 | 2 | 8 | - +----+----+----+ - "); - Ok(()) - } - - // `<=` includes equal values. - #[tokio::test] - async fn join_left_semi_less_than_equal() -> Result<()> { - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![5, 4, 2]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![2, 3, 4]), - ("c2", &vec![70, 80, 90]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::LtEq, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 2 | 4 | 8 | - | 3 | 2 | 9 | - +----+----+----+ - "); - Ok(()) - } - - // `>=` includes equal values (ascending buffered order). - #[tokio::test] - async fn join_left_semi_greater_than_equal() -> Result<()> { - // Keep left rows with some right b1 <= left b1. left b1 = {1,3,5}, - // right b1 = {3,4}. 1 has none; 3>=3; 5>={3,4}. - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![1, 3, 5]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20]), - ("b1", &vec![3, 4]), - ("c2", &vec![70, 80]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::GtEq, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 2 | 3 | 8 | - | 3 | 5 | 9 | - +----+----+----+ - "); - Ok(()) - } - - // Empty buffered (left) side: LeftSemi produces nothing, LeftAnti produces nothing. - #[tokio::test] - async fn join_left_semi_empty_left() -> Result<()> { - let left = build_table( - ("a1", &Vec::::new()), - ("b1", &Vec::::new()), - ("c1", &Vec::::new()), - ); - let right = build_table( - ("a2", &vec![1, 2]), - ("b1", &vec![1, 2]), - ("c2", &vec![1, 2]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Lt, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - +----+----+----+ - "); - Ok(()) - } - - // Empty streamed (right) side: no right row can satisfy the predicate, so LeftSemi is - // empty and LeftAnti returns all buffered rows. - #[tokio::test] - async fn join_left_anti_empty_right() -> Result<()> { - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![3, 2, 1]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &Vec::::new()), - ("b1", &Vec::::new()), - ("c2", &Vec::::new()), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | 3 | 7 | - | 2 | 2 | 8 | - | 3 | 1 | 9 | - +----+----+----+ - "); - Ok(()) - } - - // NULL join keys never satisfy a comparison predicate, so a null-keyed left row is - // excluded from LeftSemi and included in LeftAnti. Buffered side sorted descending - // with nulls first for `<`. - #[tokio::test] - async fn join_left_semi_less_than_left_nulls() -> Result<()> { - let left = build_table_nullable_b( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![None, Some(5), Some(1)]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20]), - ("b1", &vec![2, 4]), - ("c2", &vec![70, 80]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Lt, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 3 | 1 | 9 | - +----+----+----+ - "); - Ok(()) - } - - #[tokio::test] - async fn join_left_anti_less_than_left_nulls() -> Result<()> { - let left = build_table_nullable_b( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![None, Some(5), Some(1)]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20]), - ("b1", &vec![2, 4]), - ("c2", &vec![70, 80]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | | 7 | - | 2 | 5 | 8 | - +----+----+----+ - "); - Ok(()) - } - - // Existence join over a streamed side split across multiple partitions: the final - // pass that emits the bitmap result runs once, on the last streamed partition to - // finish. This exercises that coordination for LeftSemi. - #[tokio::test] - async fn join_left_semi_multi_partition_stream() -> Result<()> { - // Buffered (left) side, sorted ascending for `>`. - let left = build_table( - ("a1", &vec![1, 2, 3, 4]), - ("b1", &vec![1, 3, 5, 7]), - ("c1", &vec![10, 20, 30, 40]), - ); - - // Streamed (right) side split across two partitions. - let right_schema = Schema::new(vec![ - Field::new("a2", DataType::Int32, false), - Field::new("b1", DataType::Int32, false), - Field::new("c2", DataType::Int32, false), - ]); - let right_p0 = build_table_i32( - ("a2", &vec![10, 20]), - ("b1", &vec![2, 4]), - ("c2", &vec![70, 80]), - ); - let right_p1 = - build_table_i32(("a2", &vec![30]), ("b1", &vec![6]), ("c2", &vec![90])); - let right = TestMemoryExec::try_new_exec( - &[vec![right_p0], vec![right_p1]], - Arc::new(right_schema), - None, - )?; - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - // Keep buffered rows with some smaller streamed b1: streamed b1 = {2,4,6}, - // buffered b1 = {1,3,5,7}. 3>2, 5>{2,4}, 7>{2,4,6}. 1 has none. - let join = PiecewiseMergeJoinExec::try_new( - left, - right, - on, - Operator::Gt, - JoinType::LeftSemi, - 2, - )?; - - let task_ctx = Arc::new(TaskContext::default()); - let mut batches = Vec::new(); - for p in 0..2 { - let stream = join.execute(p, Arc::clone(&task_ctx))?; - batches.extend(common::collect(stream).await?); - } - // Sort output for a stable comparison (partitions may interleave). - let sorted = arrow::compute::concat_batches(&join.schema(), batches.iter())?; - let indices = sort_to_indices(sorted.column(1).as_ref(), None, None)?; - let sorted = take_record_batch(&sorted, &indices)?; - - assert_snapshot!(batches_to_string(&[sorted]), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 2 | 3 | 20 | - | 3 | 5 | 30 | - | 4 | 7 | 40 | - +----+----+----+ - "); - Ok(()) - } - - // LeftAnti `>=` — complement of the LeftSemi `>=` case, closes the operator matrix. - #[tokio::test] - async fn join_left_anti_greater_than_equal() -> Result<()> { - // Keep left rows with NO right b1 <= left b1. left b1 = {1,3,5}, - // right b1 = {3,4}. Only 1 has no smaller-or-equal right value. - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![1, 3, 5]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20]), - ("b1", &vec![3, 4]), - ("c2", &vec![70, 80]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::GtEq, JoinType::LeftAnti).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | 1 | 7 | - +----+----+----+ - "); - Ok(()) - } - - // Every streamed (right) join key is NULL, so no comparison can ever match: - // LeftAnti must return all buffered rows. - #[tokio::test] - async fn join_left_anti_all_right_nulls() -> Result<()> { - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![3, 2, 1]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table_nullable_b( - ("a2", &vec![10, 20]), - ("b1", &vec![None, None]), - ("c2", &vec![70, 80]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Lt, JoinType::LeftAnti).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | 3 | 7 | - | 2 | 2 | 8 | - | 3 | 1 | 9 | - +----+----+----+ - "); - Ok(()) - } - - // Non-integer key type: existence joins reuse the same comparator as classic joins, - // so a Date32 key should behave identically. - #[tokio::test] - async fn join_left_semi_date32_greater_than() -> Result<()> { - // left dates > some right date. left b1 = {19100, 19105, 19110}, - // right b1 = {19102, 19108}. 19100 has none; 19105>19102; 19110>{19102,19108}. - let left = build_date_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![19100, 19105, 19110]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_date_table( - ("a2", &vec![10, 20]), - ("b1", &vec![19102, 19108]), - ("c2", &vec![70, 80]), - ); - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +------------+------------+------------+ - | a1 | b1 | c1 | - +------------+------------+------------+ - | 1970-01-03 | 2022-04-23 | 1970-01-09 | - | 1970-01-04 | 2022-04-28 | 1970-01-10 | - +------------+------------+------------+ - "); - Ok(()) - } - - // Streamed (right) side delivered as multiple batches within a single partition. The - // bitmap must accumulate matches across batches; here the second batch contributes - // matches the first did not, so dropping either batch would change the result. - #[tokio::test] - async fn join_left_semi_multi_batch_stream() -> Result<()> { - // Buffered (left) side, sorted ascending for `>`. - let left = build_table( - ("a1", &vec![1, 2, 3, 4]), - ("b1", &vec![1, 3, 5, 7]), - ("c1", &vec![10, 20, 30, 40]), - ); - - let right_schema = Schema::new(vec![ - Field::new("a2", DataType::Int32, false), - Field::new("b1", DataType::Int32, false), - Field::new("c2", DataType::Int32, false), - ]); - // Batch 1 (b1=6) marks only b1=7; batch 2 (b1=2) additionally marks b1={3,5,7}. - let batch1 = - build_table_i32(("a2", &vec![10]), ("b1", &vec![6]), ("c2", &vec![70])); - let batch2 = - build_table_i32(("a2", &vec![20]), ("b1", &vec![2]), ("c2", &vec![80])); - // Single partition containing two batches. - let right = TestMemoryExec::try_new_exec( - &[vec![batch1, batch2]], - Arc::new(right_schema), - None, - )?; - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 2 | 3 | 20 | - | 3 | 5 | 30 | - | 4 | 7 | 40 | - +----+----+----+ - "); - Ok(()) - } - - // Exercises the cross-batch low-water mark, including its skip branch. Batch order is - // chosen so that: - // batch 1 (b1=2) matches low -> marks buffered b1={3,5,7}, watermark drops to idx 1 - // batch 2 (b1=6) matches high -> would only mark b1=7 (already marked) -> skipped - // batch 3 (b1=0) matches even lower -> marks the remaining b1=1 - // Result is all four buffered rows; a change that either re-marked greedily or - // skipped batch 3 would be caught here. - #[tokio::test] - async fn join_left_semi_multi_batch_watermark_skip() -> Result<()> { - // Buffered (left) side, sorted ascending for `>`. - let left = build_table( - ("a1", &vec![1, 2, 3, 4]), - ("b1", &vec![1, 3, 5, 7]), - ("c1", &vec![10, 20, 30, 40]), - ); - - let right_schema = Schema::new(vec![ - Field::new("a2", DataType::Int32, false), - Field::new("b1", DataType::Int32, false), - Field::new("c2", DataType::Int32, false), - ]); - let batch1 = - build_table_i32(("a2", &vec![10]), ("b1", &vec![2]), ("c2", &vec![70])); - let batch2 = - build_table_i32(("a2", &vec![20]), ("b1", &vec![6]), ("c2", &vec![80])); - let batch3 = - build_table_i32(("a2", &vec![30]), ("b1", &vec![0]), ("c2", &vec![90])); - // Single partition, three batches processed in order. - let right = TestMemoryExec::try_new_exec( - &[vec![batch1, batch2, batch3]], - Arc::new(right_schema), - None, - )?; - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let (_, batches) = - join_collect(left, right, on, Operator::Gt, JoinType::LeftSemi).await?; - - assert_snapshot!(batches_to_string(&batches), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | 1 | 10 | - | 2 | 3 | 20 | - | 3 | 5 | 30 | - | 4 | 7 | 40 | - +----+----+----+ - "); - Ok(()) - } - - // LeftAnti over a streamed side split across multiple partitions AND multiple batches. - // Only LeftSemi had multi-partition coverage; the anti final pass emits the *unmarked* - // buffered rows, so this exercises the complementary bitmap read plus the multi-partition - // final-pass counter. Buffered b1 = {1,3,5,7}; streamed b1 = {2,4,6} spread across two - // partitions. For `>`, buffered rows with some smaller streamed value match (3,5,7); only - // b1=1 has none, so LeftAnti must return exactly that row. - #[tokio::test] - async fn join_left_anti_multi_partition_multi_batch_stream() -> Result<()> { - let left = build_table( - ("a1", &vec![1, 2, 3, 4]), - ("b1", &vec![1, 3, 5, 7]), - ("c1", &vec![10, 20, 30, 40]), - ); - - let right_schema = Schema::new(vec![ - Field::new("a2", DataType::Int32, false), - Field::new("b1", DataType::Int32, false), - Field::new("c2", DataType::Int32, false), - ]); - // Partition 0 delivers b1=2 then b1=6 as two separate batches; partition 1 delivers b1=4. - let p0_b0 = - build_table_i32(("a2", &vec![10]), ("b1", &vec![2]), ("c2", &vec![70])); - let p0_b1 = - build_table_i32(("a2", &vec![20]), ("b1", &vec![6]), ("c2", &vec![80])); - let p1_b0 = - build_table_i32(("a2", &vec![30]), ("b1", &vec![4]), ("c2", &vec![90])); - let right = TestMemoryExec::try_new_exec( - &[vec![p0_b0, p0_b1], vec![p1_b0]], - Arc::new(right_schema), - None, - )?; - - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - - let join = PiecewiseMergeJoinExec::try_new( - left, - right, - on, - Operator::Gt, - JoinType::LeftAnti, - 2, - )?; - - let task_ctx = Arc::new(TaskContext::default()); - let mut batches = Vec::new(); - for p in 0..2 { - let stream = join.execute(p, Arc::clone(&task_ctx))?; - batches.extend(common::collect(stream).await?); - } - let out = arrow::compute::concat_batches(&join.schema(), batches.iter())?; - - assert_snapshot!(batches_to_string(&[out]), @r" - +----+----+----+ - | a1 | b1 | c1 | - +----+----+----+ - | 1 | 1 | 10 | - +----+----+----+ - "); - Ok(()) - } - - // Existence output is produced entirely in the final pass, which runs exactly once on - // the last streamed partition to finish. That is coordinated by a counter seeded from - // the streamed side's partition count. Here the streamed side is single-partition while - // `target_partitions` is 4: `execute` is called once, so the counter must be seeded - // from the streamed partition count (1) for the final pass to run — otherwise `LeftSemi` - // would emit nothing. - #[tokio::test] - async fn left_semi_multi_partition_final_pass() -> Result<()> { - // left.b1 > right.b1 ; buffered sorted ascending for `>`. Keep left rows with some - // smaller right b1: left b1 = {1,3,4}, right b1 = {3,2,1}. 3>{2,1}, 4>{3,2,1}; 1 has none. - let left = build_table( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![1, 3, 4]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![3, 2, 1]), - ("c2", &vec![70, 80, 90]), - ); - let on = ( - Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, - Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, - ); - // target_partitions = 4, but the streamed side is single-partition. - let join = PiecewiseMergeJoinExec::try_new( - left, - right, - on, - Operator::Gt, - JoinType::LeftSemi, - 4, - )?; - - let task_ctx = Arc::new(TaskContext::default()); - let stream = join.execute(0, task_ctx)?; - let batches = common::collect(stream).await?; - let out = arrow::compute::concat_batches(&join.schema(), batches.iter())?; - // The two matching left rows (b1=3, b1=4) must be emitted by the final pass. - assert_eq!(out.num_rows(), 2); - Ok(()) - } } diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 588196f59e68b..8e2e39b902886 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -47,6 +47,7 @@ use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::joins::piecewise_merge_join::classic_join::{ ClassicPWMJStream, PiecewiseMergeJoinStreamState, }; +use crate::joins::piecewise_merge_join::existence_join::ExistencePWMJStream; use crate::joins::piecewise_merge_join::utils::{ build_visited_indices_map, is_existence_join, is_right_existence_join, is_supported_existence_join, @@ -165,33 +166,14 @@ use crate::{ /// /// ## Existence Joins (Semi, Anti, Mark) /// Currently only `LeftSemi` and `LeftAnti` are supported. For these the marked side is -/// already the left (buffered) side, so no input swap is needed. `RightSemi`/`RightAnti` -/// and `Mark` joins are not yet implemented and are rejected in [`Self::try_new`]. +/// already the left (buffered) side, so no input swap is needed. The rest are rejected in +/// [`Self::try_new`]: `RightSemi`/`RightAnti`/`RightMark` mark the right side and need an +/// input swap, and `LeftMark` needs an extra boolean column rather than a filtered slice. /// -/// Existence joins reuse the classic-join scan (see above) rather than a dedicated -/// algorithm: for each streamed row we advance the pointer on the sorted buffered side to -/// the first match, and because the buffered side is sorted, every buffered row from that -/// position onward also matches. Instead of materializing the joined output, we simply -/// mark those buffered rows in the visited-indices bitmap and continue. Once all streamed -/// partitions are processed, the final pass emits the result directly from the bitmap: -/// `LeftSemi` emits the marked buffered rows, `LeftAnti` emits the unmarked ones (rows -/// whose join key is NULL are never marked, so they are correctly excluded from `Semi` -/// and included in `Anti`). Only the buffered (left) columns are produced. -/// -/// Reusing the classic scan keeps a single code path and works for the same predicates -/// classic joins support, at the cost of still sorting the streamed side. -/// -/// The pseudocode looks like this: -/// -/// ```text -/// for stream_row in sorted_stream_batch: -/// advance buffer_idx to the first buffered row that matches stream_row -/// if a match is found: -/// mark buffered[buffer_idx..end] in the bitmap -/// // final pass, once all partitions finish: -/// // LeftSemi -> emit buffered rows where bit == 1 -/// // LeftAnti -> emit buffered rows where bit == 0 -/// ``` +/// `LeftSemi`/`LeftAnti` are served by a dedicated stream, `ExistencePWMJStream` (see +/// `existence_join.rs`). Instead of materializing row pairs it records the matched set as a +/// single index -- the start of the matched suffix of the buffered side -- and slices the +/// buffered batch at that index once every streamed partition has been consumed. /// /// ```text /// // Using the example of a less than `<` operation @@ -313,9 +295,9 @@ impl PiecewiseMergeJoinExec { join_type: JoinType, num_partitions: usize, ) -> Result { - // Left Semi/Anti reuse the classic scan (marked side is already the buffered - // side, no input swap needed). Right existence joins and Mark joins are not yet - // supported. + // Left Semi/Anti are handled by `ExistencePWMJStream` (the marked side is + // already the buffered side, so no input swap is needed). Right existence joins + // and Mark joins are not yet supported. if is_existence_join(join_type) && !is_supported_existence_join(join_type) { return not_impl_err!( "Existence join {join_type} is currently not supported for PiecewiseMergeJoin" @@ -525,7 +507,12 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { fn required_input_ordering(&self) -> Vec> { // Existence joins don't need to be sorted on one side. if is_right_existence_join(self.join_type) { - unimplemented!() + // Unreachable: `try_new` rejects right existence joins, and this signature + // cannot return `Result`. They swap the inputs, so whoever implements them + // must require the order on the streamed side instead. + unimplemented!( + "required_input_ordering for right existence joins; guarded by try_new" + ) } else { // Sort the right side in memory, so we do not need to enforce any sorting vec![ @@ -625,15 +612,26 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { let batch_size = context.session_config().batch_size(); - // Right existence joins and Mark joins are rejected in `try_new`; every other - // join type (classic + supported Left Semi/Anti existence) runs the classic scan, - // which marks the buffered bitmap and emits the existence result from it. - if is_existence_join(self.join_type()) - && !is_supported_existence_join(self.join_type()) - { + let buffered_side = + BufferedSide::Initial(BufferedSideInitialState { buffered_fut }); + + if is_supported_existence_join(self.join_type) { + Ok(Box::pin(ExistencePWMJStream::try_new( + Arc::clone(&self.schema), + on_streamed, + self.join_type, + self.operator, + streamed, + buffered_side, + self.sort_options, + metrics, + batch_size, + ))) + } else if is_existence_join(self.join_type) { + // Right existence joins and Mark joins are rejected in `try_new`. internal_err!( "PiecewiseMergeJoin does not support existence join {} (should have been rejected in try_new)", - self.join_type() + self.join_type ) } else { Ok(Box::pin(ClassicPWMJStream::try_new( @@ -642,7 +640,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { self.join_type, self.operator, streamed, - BufferedSide::Initial(BufferedSideInitialState { buffered_fut }), + buffered_side, PiecewiseMergeJoinStreamState::WaitBufferedSide, self.sort_options, metrics, @@ -755,6 +753,11 @@ pub(super) struct BufferedSideData { values: ArrayRef, pub(super) visited_indices_bitmap: SharedBitmapBuilder, pub(super) remaining_partitions: AtomicUsize, + /// Existence joins only: the start of the matched suffix of the buffered side, or + /// `usize::MAX` before the first match. `[existence_min_marked, len)` *is* the matched + /// set -- no bitmap is allocated. Shared so each partition benefits from what the + /// others have marked; it only ever decreases, so a stale read is safe. + pub(super) existence_min_marked: AtomicUsize, _reservation: MemoryReservation, } @@ -771,6 +774,7 @@ impl BufferedSideData { values, visited_indices_bitmap, remaining_partitions: AtomicUsize::new(remaining_partitions), + existence_min_marked: AtomicUsize::new(usize::MAX), _reservation: reservation, } } diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs new file mode 100644 index 0000000000000..b66f6e01fed71 --- /dev/null +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs @@ -0,0 +1,811 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! PiecewiseMergeJoin stream specialized for existence joins. +//! +//! Instantiated by [`PiecewiseMergeJoinExec`] when the join type is `LeftSemi` or `LeftAnti`. +//! The other existence joins are rejected in `PiecewiseMergeJoinExec::try_new`: +//! `RightSemi`/`RightAnti`/`RightMark` mark the right side and so need an input swap, while +//! `LeftMark` marks the left side but needs an extra boolean column rather than a slice. +//! +//! # Motivation +//! +//! `ClassicPWMJStream` (see `classic_join.rs`) materializes `(buffered, streamed)` row +//! pairs: on the first matching buffered row it emits the whole matching suffix joined +//! with the current streamed row. Existence joins only need a boolean per buffered row — +//! does any match exist? — so pair materialization is pure waste, and the resume state it +//! requires (partial output batches, per-row buffered/streamed cursors) is dead weight. +//! +//! This stream instead records matches as a single index -- the start of the matched +//! suffix -- and emits nothing at all while scanning. Output is produced once, at the end, +//! by slicing the buffered batch at that index. +//! +//! # Algorithm +//! +//! The buffered (left) side arrives globally sorted, enforced by `required_input_ordering` +//! (the streamed side carries no ordering requirement). For a streamed key, a binary search +//! finds the first matching buffered row; because the buffered side is sorted, every buffered +//! row from that position onward matches too: +//! +//! ```text +//! buffered (sorted): [1, 3, 5, 7] +//! ▲ +//! first match for this streamed row +//! → mark [1..4), i.e. buffered 3, 5, 7 +//! ``` +//! +//! Marking that suffix covers every match the batch can produce, so only one row of the +//! batch is ever compared against the buffered side: the extreme key, which reaches the +//! smallest matching `buffer_idx`, while every other row matches a subset of that suffix. +//! +//! The search stops at `min_marked`: rows from there on were already marked, by this +//! partition or another, so a match found among them would record nothing. Marking is then a +//! single `fetch_min`, and finding the first match is a binary search, so a batch costs +//! `O(log buffered)` rather than `O(buffered)`. +//! +//! `min_marked` lives in the shared buffered data, so partitions benefit from each other's +//! marking rather than each rediscovering it. It only ever decreases, so a partition that +//! reads a stale value scans a wider range than it had to -- never a narrower one. +//! +//! Once `min_marked` reaches the first non-null buffered row nothing can ever be marked +//! again, and every partition stops reading the streamed side -- checked before each poll, +//! so a partition that starts late reads nothing at all. +//! +//! Rows whose join key is NULL never satisfy a comparison predicate. Buffered NULLs sort to +//! the front, so the scan starts past them and null-keyed buffered rows are left unmarked — +//! correctly excluded from `LeftSemi` and included in `LeftAnti`. Each streamed batch is +//! sorted NULLs-*last* instead, so the extreme key taken from it is non-null unless the whole +//! batch is. +//! +//! # Output +//! +//! Marking only ever covers a suffix, and each mark lowers the watermark to its own start, +//! so the matched set is always exactly `[min_marked, buffered_len)`. A bitmap would be a +//! less compact encoding of that one index, so none is allocated (see +//! `build_visited_indices_map`). +//! +//! Once every streamed partition has been consumed, the last one to finish slices the +//! buffered batch: `LeftSemi` takes `[min_marked, len)`, `LeftAnti` the complementary +//! prefix `[0, min_marked)`, which is where the null-keyed rows live. Only the buffered +//! (left) columns are produced. +//! +//! [`PiecewiseMergeJoinExec`]: super::PiecewiseMergeJoinExec + +use std::cmp::Ordering; +use std::sync::Arc; +use std::sync::atomic::Ordering as AtomicOrdering; +use std::task::{Poll, ready}; + +use arrow::array::{Array, ArrayRef, RecordBatch}; +use arrow::compute::{BatchCoalescer, sort_to_indices, take}; +use arrow_schema::{SchemaRef, SortOptions}; +use datafusion_common::{NullEquality, Result, internal_err}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion_expr::{JoinType, Operator}; +use datafusion_physical_expr::PhysicalExprRef; +use futures::{Stream, StreamExt}; + +use crate::handle_state; +use crate::joins::piecewise_merge_join::exec::{BufferedSide, BufferedSideReadyState}; +use crate::joins::utils::{ + BuildProbeJoinMetrics, JoinKeyComparator, StatefulStreamResult, +}; +use crate::stream::EmptyRecordBatchStream; + +pub(super) enum ExistencePWMJStreamState { + /// Load the buffered side into memory. + WaitBufferedSide, + /// Fetch, sort and scan streamed batches, lowering the watermark. Emits nothing. + ScanStreamBatches, + /// Emit the result. Reached only by the last streamed partition. + EmitMatched, + Completed, +} + +pub(super) struct ExistencePWMJStream { + /// Output schema, which for `LeftSemi`/`LeftAnti` is the buffered side's schema + schema: SchemaRef, + /// Physical expression evaluated on the streamed side. The buffered side's + /// equivalent is already evaluated when the buffered side is collected. + on_streamed: PhysicalExprRef, + /// `LeftSemi` or `LeftAnti` + join_type: JoinType, + /// Comparison operator + operator: Operator, + streamed: SendableRecordBatchStream, + buffered_side: BufferedSide, + state: ExistencePWMJStreamState, + /// Whether the buffered side is sorted ascending or descending, per the operator + sort_option: SortOptions, + join_metrics: BuildProbeJoinMetrics, + /// Chunks the single final-pass batch to `batch_size` + output_batches: Box, + /// Whether the final pass has already pushed its result into `output_batches` + emitted: bool, +} + +impl ExistencePWMJStream { + #[expect(clippy::too_many_arguments)] + pub(super) fn try_new( + schema: SchemaRef, + on_streamed: PhysicalExprRef, + join_type: JoinType, + operator: Operator, + streamed: SendableRecordBatchStream, + buffered_side: BufferedSide, + sort_option: SortOptions, + join_metrics: BuildProbeJoinMetrics, + batch_size: usize, + ) -> Self { + Self { + output_batches: Box::new(BatchCoalescer::new( + Arc::clone(&schema), + batch_size, + )), + schema, + on_streamed, + join_type, + operator, + streamed, + buffered_side, + state: ExistencePWMJStreamState::WaitBufferedSide, + sort_option, + join_metrics, + emitted: false, + } + } + + fn poll_next_impl( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>> { + loop { + return match self.state { + ExistencePWMJStreamState::WaitBufferedSide => { + handle_state!(ready!(self.collect_buffered_side(cx))) + } + ExistencePWMJStreamState::ScanStreamBatches => { + handle_state!(ready!(self.scan_stream_batch(cx))) + } + ExistencePWMJStreamState::EmitMatched => { + handle_state!(self.emit_matched()) + } + ExistencePWMJStreamState::Completed => Poll::Ready(None), + }; + } + } + + /// Collects the buffered side into memory. + fn collect_buffered_side( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>>> { + let build_timer = self.join_metrics.build_time.timer(); + let buffered_data = ready!( + self.buffered_side + .try_as_initial_mut()? + .buffered_fut + .get_shared(cx) + )?; + build_timer.done(); + + self.buffered_side = + BufferedSide::Ready(BufferedSideReadyState { buffered_data }); + self.state = ExistencePWMJStreamState::ScanStreamBatches; + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Fetches one streamed batch, sorts it, and marks the buffered rows it matches. + /// Never produces output: existence results come from the watermark in `emit_matched`. + fn scan_stream_batch( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll>>> { + // Every buffered row that can ever be marked is already marked -- by this + // partition or any other, since the watermark is shared -- so no batch can lower + // the watermark further. Stop reading rather than scanning batches that provably + // cannot contribute. Checked before polling so a partition that starts after another + // has saturated the watermark reads nothing at all. + if self.nothing_left_to_mark()? { + self.finish_streamed_side()?; + return Poll::Ready(Ok(StatefulStreamResult::Continue)); + } + + match ready!(self.streamed.poll_next_unpin(cx)) { + None => self.finish_streamed_side()?, + Some(Ok(batch)) => { + let stream_values: ArrayRef = self + .on_streamed + .evaluate(&batch)? + .into_array(batch.num_rows())?; + + self.join_metrics.input_batches.add(1); + self.join_metrics.input_rows.add(batch.num_rows()); + + // Only the batch's extreme key is ever compared against the buffered + // side, so ask arrow for just that one index. `nulls_first: false` makes + // the chosen row non-null whenever the batch holds any non-null key, and + // leaves a null row only for an all-null batch, which marks nothing. + // + // The limit truncates the returned indices, not the work: `sort_impl` + // feeds it into the sort only when `nulls_first` is true, so this still + // fully sorts the batch's non-null keys. + let indices = sort_to_indices( + stream_values.as_ref(), + Some(SortOptions { + descending: self.sort_option.descending, + nulls_first: false, + }), + Some(1), + )?; + let stream_values = take(stream_values.as_ref(), &indices, None)?; + + self.mark_matched_buffered_rows(&stream_values)?; + } + Some(Err(err)) => return Poll::Ready(Err(err)), + } + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Whether no buffered row can ever be marked again, in which case the streamed side + /// no longer needs reading. Reads the shared watermark, so one partition saturating it + /// lets every other partition stop too. + fn nothing_left_to_mark(&self) -> Result { + let buffered_data = &self.buffered_side.try_as_ready()?.buffered_data; + let min_marked = buffered_data + .existence_min_marked + .load(AtomicOrdering::SeqCst); + let buffered_values = buffered_data.values(); + + Ok(min_marked.min(buffered_values.len()) <= buffered_values.null_count()) + } + + /// Marks this partition done with the streamed side: releases the input pipeline and, + /// if this is the last streamed partition to finish, moves on to the final pass. + fn finish_streamed_side(&mut self) -> Result<()> { + // Release the streamed input pipeline's resources. + let streamed_schema = self.streamed.schema(); + self.streamed = Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + + // The final pass must run exactly once, on the last streamed partition to finish. + if self + .buffered_side + .try_as_ready()? + .buffered_data + .remaining_partitions + .fetch_sub(1, AtomicOrdering::SeqCst) + == 1 + { + self.state = ExistencePWMJStreamState::EmitMatched; + } else { + self.state = ExistencePWMJStreamState::Completed; + } + + Ok(()) + } + + /// Marks every buffered row matched by `stream_values`, a one-row array holding the + /// batch's extreme compare key (null only if the whole batch was null). + fn mark_matched_buffered_rows(&mut self, stream_values: &ArrayRef) -> Result<()> { + let operator = self.operator; + let sort_option = self.sort_option; + + { + let buffered_data = &self.buffered_side.try_as_ready()?.buffered_data; + let buffered_values = buffered_data.values(); + let buffered_len = buffered_values.len(); + + // NULL keys can never match, and `sort_options` uses `nulls_first` for every + // operator (see `try_new`), so buffered nulls sit at the front -- skip past them. + let first_non_null_buffered = buffered_values.null_count(); + + // `[min_marked, buffered_len)` was already marked, by this partition or + // another, so a match found there would write nothing. Stop the scan at the + // watermark: that bounds the comparisons this batch performs, not just the + // bits it writes. + let scan_limit = buffered_data + .existence_min_marked + .load(AtomicOrdering::SeqCst) + .min(buffered_len); + + // The extreme key is the only one that can decide anything: it reaches the + // smallest matching `buffer_idx`, and every other row in the batch matches a + // subset of the same buffered suffix, so could only re-mark it. `null_count()` + // is 0 for a real key and 1 for an all-null batch, which skips the scan. + let row_idx = stream_values.null_count(); + + // `<=`/`>=` also match on equality; validated once here rather than inside + // the search below. + let match_on_equal = match operator { + Operator::Gt | Operator::Lt => false, + Operator::GtEq | Operator::LtEq => true, + _ => { + return internal_err!( + "PiecewiseMergeJoin should not contain operator, {}", + operator + ); + } + }; + + if row_idx < stream_values.len() && first_non_null_buffered < scan_limit { + let cmp = JoinKeyComparator::new( + &[Arc::clone(stream_values)], + &[Arc::clone(buffered_values)], + &[sort_option], + NullEquality::NullEqualsNothing, + )?; + let is_match = |buffer_idx: usize| { + let compare = cmp.compare(row_idx, buffer_idx); + compare == Ordering::Less + || (match_on_equal && compare == Ordering::Equal) + }; + + // Because the buffered side is sorted, `is_match` is monotone over it: + // false while the buffered key has not yet passed the streamed key, true + // from there on. So the first match is a partition point and can be found + // by binary search instead of a walk -- `O(log buffered)` per batch rather + // than `O(buffered)`. + let mut lo = first_non_null_buffered; + let mut hi = scan_limit; + while lo < hi { + let mid = lo + (hi - lo) / 2; + if is_match(mid) { + hi = mid; + } else { + lo = mid + 1; + } + } + + // `lo` is now the first matching buffered index, or `scan_limit` if this + // batch matches nothing new. + let buffer_idx = lo; + if buffer_idx < scan_limit { + // Everything from `buffer_idx` on matches, so lowering the + // watermark to it records the match: the marked set is exactly + // `[existence_min_marked, buffered_len)` and needs no bitmap. + // + // INVARIANT: sound only because the buffered side and each + // streamed batch are sorted the same way for this operator + // (`try_new` derives `sort_option`: descending for `<`/`<=`, + // ascending for `>`/`>=`). That makes this the smallest reachable + // `buffer_idx`, so the marked suffix is maximal. Only the ordering + // *within* a batch matters; batches themselves may arrive in any + // order, which is why the watermark takes a `min` rather than just + // decreasing. + buffered_data + .existence_min_marked + .fetch_min(buffer_idx, AtomicOrdering::SeqCst); + } + } + } + + Ok(()) + } + + /// Emits the existence result by slicing at the watermark: the marked buffered rows for + /// `LeftSemi`, the unmarked ones for `LeftAnti`. + fn emit_matched(&mut self) -> Result>> { + if !self.emitted { + self.emitted = true; + + let buffered_data = + Arc::clone(&self.buffered_side.try_as_ready()?.buffered_data); + let buffered_batch = buffered_data.batch(); + let buffered_len = buffered_batch.num_rows(); + + // The marked rows are always the contiguous suffix `[min_marked, len)`: each + // match covers `[k, previous min_marked)` and then lowers the watermark to + // `k`, so the union is `[k, len)`. The result is therefore a slice, with no + // index array to materialize and no `take`. + let min_marked = buffered_data + .existence_min_marked + .load(AtomicOrdering::SeqCst) + .min(buffered_len); + + let sliced = match self.join_type { + JoinType::LeftSemi => { + buffered_batch.slice(min_marked, buffered_len - min_marked) + } + // The unmarked prefix, which includes every null-keyed row: nulls sort + // first and the watermark never drops below the buffered null count. + _ => buffered_batch.slice(0, min_marked), + }; + + if sliced.num_rows() > 0 { + // Existence joins output the buffered (left) columns only; rebuild against + // the join's own schema, which keeps the slice zero-copy. + let batch = RecordBatch::try_new( + Arc::clone(&self.schema), + sliced.columns().to_vec(), + )?; + self.output_batches.push_batch(batch)?; + self.output_batches.finish_buffered_batch()?; + } + } + + // Drain one coalesced batch per poll; `emitted` keeps the block above from + // re-running, so this always terminates. + match self.output_batches.next_completed_batch() { + Some(batch) => Ok(StatefulStreamResult::Ready(Some(batch))), + None => { + self.state = ExistencePWMJStreamState::Completed; + Ok(StatefulStreamResult::Ready(None)) + } + } + } +} + +impl RecordBatchStream for ExistencePWMJStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for ExistencePWMJStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + // `record_poll` fills in `output_rows` and `end_time`; `elapsed_compute` is handled + // by `BuildProbeJoinMetrics::drop`. + let poll = self.poll_next_impl(cx); + self.join_metrics.baseline.record_poll(poll) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + ExecutionPlan, common, + joins::PiecewiseMergeJoinExec, + test::{TestMemoryExec, build_table_i32}, + }; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_common::test_util::batches_to_string; + use datafusion_execution::TaskContext; + use datafusion_execution::config::SessionConfig; + use datafusion_physical_expr::expressions::Column; + use insta::assert_snapshot; + + // Coverage for existence joins lives in `pwmj.slt`: operators, NULL handling, empty + // sides, key types, both correlation orientations, and the streamed-side batch and + // partition layouts (which SQL pins via one `INSERT` per batch plus an `EXPLAIN` + // asserting `partition_sizes`). The two tests here are demos of the operator itself. + + fn build_table( + a: (&str, &Vec), + b: (&str, &Vec), + c: (&str, &Vec), + ) -> Arc { + let batch = build_table_i32(a, b, c); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + async fn join_collect(join_type: JoinType) -> Result> { + // Buffered (left) side pre-sorted ascending, as `>` requires. These tests build + // the exec directly, so there is no `SortExec` to enforce it. + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 5]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let join = + PiecewiseMergeJoinExec::try_new(left, right, on, Operator::Gt, join_type, 1)?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + common::collect(stream).await + } + + /// `LeftSemi` keeps the buffered rows with at least one match, and outputs only the + /// buffered columns. Of b1 = {1,2,5} against streamed {2,3,4}, only 5 > some streamed + /// value. + #[tokio::test] + async fn join_left_semi() -> Result<()> { + let batches = join_collect(JoinType::LeftSemi).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 3 | 5 | 9 | + +----+----+----+ + "); + Ok(()) + } + + /// `LeftAnti` is the complement: the buffered rows with no match. + #[tokio::test] + async fn join_left_anti() -> Result<()> { + let batches = join_collect(JoinType::LeftAnti).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 7 | + | 2 | 2 | 8 | + +----+----+----+ + "); + Ok(()) + } + + /// Once every markable buffered row is marked, no later streamed batch can lower the + /// watermark, so the stream stops reading rather than scanning batches that provably + /// cannot contribute. The first batch below is smaller than every buffered value, so it + /// marks all four rows; the two batches after it must never be read. Asserted through + /// the `input_batches` metric, which SQL cannot observe. + #[tokio::test] + async fn early_exit_stops_reading_streamed_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 3, 5, 7]), + ("c1", &vec![10, 20, 30, 40]), + ); + + let streamed_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // b1=0 is below every buffered value, so batch 1 marks the whole buffered side. + let batch1 = + build_table_i32(("a2", &vec![10]), ("b1", &vec![0]), ("c2", &vec![70])); + let batch2 = + build_table_i32(("a2", &vec![20]), ("b1", &vec![2]), ("c2", &vec![80])); + let batch3 = + build_table_i32(("a2", &vec![30]), ("b1", &vec![6]), ("c2", &vec![90])); + let right = TestMemoryExec::try_new_exec( + &[vec![batch1, batch2, batch3]], + Arc::new(streamed_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftSemi, + 1, + )?; + + let stream = join.execute(0, Arc::new(TaskContext::default()))?; + let batches = common::collect(stream).await?; + + // All four buffered rows still come out. + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 10 | + | 2 | 3 | 20 | + | 3 | 5 | 30 | + | 4 | 7 | 40 | + +----+----+----+ + "); + + // ...but only the first of the three streamed batches was ever read. + let consumed = join + .metrics() + .unwrap() + .sum_by_name("input_batches") + .expect("input_batches metric") + .as_usize(); + assert_eq!(consumed, 1, "expected early exit after the first batch"); + Ok(()) + } + + /// The final pass pushes one slice of the buffered batch into a `BatchCoalescer`, so a + /// result wider than `batch_size` has to be drained over several polls. Also pins the + /// `output_rows` metric, which `record_poll` in `poll_next` is what supplies. + #[tokio::test] + async fn final_pass_chunks_output_and_records_output_rows() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 3, 5, 7]), + ("c1", &vec![10, 20, 30, 40]), + ); + + let streamed_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // b1=0 is below every buffered value, so every buffered row matches `>`. + let batch = + build_table_i32(("a2", &vec![10]), ("b1", &vec![0]), ("c2", &vec![70])); + let right = TestMemoryExec::try_new_exec( + &[vec![batch]], + Arc::new(streamed_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftSemi, + 1, + )?; + + // batch_size 2 over a 4-row result forces the coalescer to hand back two batches. + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(2)), + ); + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + assert_eq!(batches.len(), 2, "expected the final pass to be chunked"); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 4); + + let output_rows = join + .metrics() + .unwrap() + .output_rows() + .expect("output_rows metric"); + assert_eq!(output_rows, 4); + Ok(()) + } + + /// The watermark lives in the shared buffered data, so one partition saturating it + /// lets the others stop too. Partition 0's single batch marks every buffered row; + /// partition 1 is then driven and must read none of its three batches. Only the shared + /// watermark makes that possible -- with a per-partition watermark, partition 1 would + /// rescan all three. + #[tokio::test] + async fn early_exit_is_shared_across_streamed_partitions() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 3, 5, 7]), + ("c1", &vec![10, 20, 30, 40]), + ); + + let streamed_schema = Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // Partition 0: one batch below every buffered value -> marks the whole buffered side. + let p0 = build_table_i32(("a2", &vec![10]), ("b1", &vec![0]), ("c2", &vec![70])); + // Partition 1: three batches that can now contribute nothing. + let p1_b0 = + build_table_i32(("a2", &vec![20]), ("b1", &vec![2]), ("c2", &vec![80])); + let p1_b1 = + build_table_i32(("a2", &vec![30]), ("b1", &vec![4]), ("c2", &vec![90])); + let p1_b2 = + build_table_i32(("a2", &vec![40]), ("b1", &vec![6]), ("c2", &vec![100])); + let right = TestMemoryExec::try_new_exec( + &[vec![p0], vec![p1_b0, p1_b1, p1_b2]], + Arc::new(streamed_schema), + None, + )?; + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let join = PiecewiseMergeJoinExec::try_new( + left, + right, + on, + Operator::Gt, + JoinType::LeftSemi, + 2, + )?; + + // Driven in order so partition 0 saturates the watermark before partition 1 starts. + let task_ctx = Arc::new(TaskContext::default()); + let mut batches = Vec::new(); + for partition in 0..2 { + let stream = join.execute(partition, Arc::clone(&task_ctx))?; + batches.extend(common::collect(stream).await?); + } + let out = arrow::compute::concat_batches(&join.schema(), batches.iter())?; + + assert_snapshot!(batches_to_string(&[out]), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 10 | + | 2 | 3 | 20 | + | 3 | 5 | 30 | + | 4 | 7 | 40 | + +----+----+----+ + "); + + // Partition 0 read its one batch; partition 1 read none of its three. + let consumed = join + .metrics() + .unwrap() + .sum_by_name("input_batches") + .expect("input_batches metric") + .as_usize(); + assert_eq!(consumed, 1, "partition 1 should not have read any batch"); + Ok(()) + } + + /// The unsupported existence joins must be rejected at construction, not deeper in. + /// `required_input_ordering` still has an `unimplemented!()` for right existence joins + /// and cannot return an error, so this test is what keeps that panic unreachable: if + /// someone opens the gate for RightSemi/RightAnti without also supplying an ordering + /// requirement, this fails instead of panicking the optimizer at runtime. + #[test] + fn try_new_rejects_unsupported_existence_joins() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 5]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + for join_type in [ + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ] { + let err = PiecewiseMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Operator::Gt, + join_type, + 1, + ) + .expect_err(&format!("{join_type} should be rejected")) + .to_string(); + assert!( + err.contains("not supported for PiecewiseMergeJoin"), + "unexpected error for {join_type}: {err}" + ); + } + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs index c85a7cc16f657..8c6815ad6c631 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/mod.rs @@ -21,4 +21,5 @@ pub use exec::PiecewiseMergeJoinExec; mod classic_join; mod exec; +mod existence_join; mod utils; diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs index 88a4b9e81ab62..5093be0ca19be 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/utils.rs @@ -40,33 +40,30 @@ pub(super) fn is_existence_join(join_type: JoinType) -> bool { // Returns boolean for whether the join is a left existence join that is currently // supported by `PiecewiseMergeJoin`. These do not require swapping the inputs: the -// marked (left) side is already the buffered side, so the classic scan can mark the -// buffered bitmap and the final indices are emitted from it. +// marked (left) side is already the buffered side, so `ExistencePWMJStream` can track the +// matched suffix and slice the buffered batch at its start. pub(super) fn is_supported_existence_join(join_type: JoinType) -> bool { matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti) } -// Returns true if the join type produces its output in the final pass over the buffered -// side (using the visited-indices bitmap) rather than while scanning stream batches: -// `Left`/`Full` emit unmatched buffered rows, and `LeftSemi`/`LeftAnti` emit the -// matched/unmatched buffered rows respectively. +// Returns boolean to check if the join type needs to record +// buffered side matches for classic joins pub(super) fn need_produce_result_in_final(join_type: JoinType) -> bool { - matches!( - join_type, - JoinType::Full | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti - ) + matches!(join_type, JoinType::Full | JoinType::Left) } // Returns boolean for whether or not we need to build the buffered side // bitmap for marking matched rows on the buffered side. +// +// `LeftSemi`/`LeftAnti` are absent on purpose: `ExistencePWMJStream` only ever marks a +// contiguous suffix of the buffered side, so it tracks the boundary as a single index +// (`BufferedSideData::existence_min_marked`) and needs no bitmap. pub(super) fn build_visited_indices_map(join_type: JoinType) -> bool { matches!( join_type, JoinType::Full | JoinType::Left - | JoinType::LeftAnti | JoinType::RightAnti - | JoinType::LeftSemi | JoinType::RightSemi | JoinType::LeftMark | JoinType::RightMark diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 4f8ae04bc8d98..8ed4d27022898 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -431,5 +431,620 @@ ORDER BY 1 NULLS FIRST; NULL 1 +# ------------------------------------------------------------------ +# Existence joins: operator, NULL, empty-side and type coverage +# ------------------------------------------------------------------ + +statement ok +CREATE TABLE ex_l(id INT, v INT); + +statement ok +INSERT INTO ex_l VALUES (1, 5), (2, 4), (3, 2), (4, 1); + +statement ok +CREATE TABLE ex_r(v INT); + +statement ok +INSERT INTO ex_r VALUES (2), (3), (4); + +# `<` : left v = {5,4,2,1}, right v = {2,3,4}. 2<{3,4}; 1` : 5>{2,3,4}; 4>{2,3}. 2 and 1 have no smaller right value. +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v) ORDER BY 1; +---- +3 +4 + +# `>=` additionally admits v=2 via the equal right value, so it must differ from `>`. +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE l.v >= r.v) ORDER BY 1; +---- +1 +2 +3 + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_r r WHERE l.v >= r.v) ORDER BY 1; +---- +4 + +# Writing the correlation with the inner column first is the same join; the planner +# normalizes it, so the operator on the physical join flips to keep the buffered (left) +# side as the marked side. Same rows as `l.v > r.v` above. +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE r.v < l.v) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_r r WHERE r.v < l.v) ORDER BY 1; +---- +3 +4 + +# The physical operator is `Gt`, not `Lt`: the predicate is flipped so the marked (left) +# side stays on the buffered side and no input swap is needed. +query TT +EXPLAIN SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE r.v < l.v) ORDER BY 1; +---- +logical_plan +01)Sort: l.id ASC NULLS LAST +02)--Projection: l.id +03)----LeftSemi Join: Filter: __correlated_sq_1.v < l.v +04)------SubqueryAlias: l +05)--------TableScan: ex_l projection=[id, v] +06)------SubqueryAlias: __correlated_sq_1 +07)--------SubqueryAlias: r +08)----------TableScan: ex_r projection=[v] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--ProjectionExec: expr=[id@0 as id] +03)----PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(v > v) +04)------SortExec: expr=[v@1 ASC], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)------DataSourceExec: partitions=1, partition_sizes=[1] + +# A filter inside the subquery lets the streamed side be repartitioned, so the final pass +# has to be coordinated across more than one streamed partition. +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v AND r.v > 0) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v AND r.v > 0) ORDER BY 1; +---- +3 +4 + +# Confirms the streamed side really is repartitioned above (RoundRobinBatch(4)), which is +# what makes the multi-partition final-pass coordination reachable from SQL. +query TT +EXPLAIN SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v AND r.v > 0) ORDER BY 1; +---- +logical_plan +01)Sort: l.id ASC NULLS LAST +02)--Projection: l.id +03)----LeftSemi Join: Filter: l.v > __correlated_sq_1.v +04)------SubqueryAlias: l +05)--------TableScan: ex_l projection=[id, v] +06)------SubqueryAlias: __correlated_sq_1 +07)--------SubqueryAlias: r +08)----------Filter: ex_r.v > Int32(0) +09)------------TableScan: ex_r projection=[v] +physical_plan +01)SortPreservingMergeExec: [id@0 ASC NULLS LAST] +02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----ProjectionExec: expr=[id@0 as id] +04)------PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(v > v) +05)--------SortExec: expr=[v@1 ASC], preserve_partitioning=[false] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] +07)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +08)----------FilterExec: v@0 > 0 +09)------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +CREATE TABLE ex_empty(v INT); + +# Empty buffered (left) side: nothing to mark and nothing to emit either way. +query I +SELECT l.v FROM ex_empty l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v) ORDER BY 1; +---- + +query I +SELECT l.v FROM ex_empty l WHERE NOT EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v) ORDER BY 1; +---- + +# Empty streamed (right) side: the watermark never moves, so EXISTS is empty and NOT EXISTS keeps +# every buffered row. +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_empty r WHERE l.v > r.v) ORDER BY 1; +---- + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_empty r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +statement ok +CREATE TABLE ex_r_null(v INT); + +statement ok +INSERT INTO ex_r_null VALUES (NULL), (NULL); + +# Every streamed key is NULL, so no comparison can match: same outcome as an empty +# streamed side. +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_r_null r WHERE l.v > r.v) ORDER BY 1; +---- + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_r_null r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +statement ok +CREATE TABLE ex_dl(id INT, d DATE); + +statement ok +INSERT INTO ex_dl VALUES (1, DATE '2022-04-23'), (2, DATE '2022-04-28'), (3, DATE '2022-04-18'); + +statement ok +CREATE TABLE ex_dr(d DATE); + +statement ok +INSERT INTO ex_dr VALUES (DATE '2022-04-20'), (DATE '2022-04-26'); + +# Non-integer key: existence joins share the comparator with the matching joins, so a +# Date32 key behaves identically. 04-23 > 04-20; 04-28 > both; 04-18 has none. +query I +SELECT l.id FROM ex_dl l WHERE EXISTS (SELECT 1 FROM ex_dr r WHERE l.d > r.d) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_dl l WHERE NOT EXISTS (SELECT 1 FROM ex_dr r WHERE l.d > r.d) ORDER BY 1; +---- +3 + +# Stopping at the first matching buffered row is load-bearing: it relies on both sides +# being sorted in the same direction, so that the first match yields the maximal matching +# buffered suffix and later streamed rows could only re-mark a subset. This interleaved +# case (buffered {2,4,6,8} vs streamed {1,3,5,7}, all in one batch) is the shape that would +# expose an off-by-one there: every buffered row has some smaller streamed value, so all +# four must survive EXISTS and none may survive NOT EXISTS. +statement ok +CREATE TABLE ex_interleave_l(id INT, v INT); + +statement ok +INSERT INTO ex_interleave_l VALUES (1, 2), (2, 4), (3, 6), (4, 8); + +statement ok +CREATE TABLE ex_interleave_r(v INT); + +statement ok +INSERT INTO ex_interleave_r VALUES (1), (3), (5), (7); + +query I +SELECT l.id FROM ex_interleave_l l WHERE EXISTS (SELECT 1 FROM ex_interleave_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +query I +SELECT l.id FROM ex_interleave_l l WHERE NOT EXISTS (SELECT 1 FROM ex_interleave_r r WHERE l.v > r.v) ORDER BY 1; +---- + +# Only the extreme key of each streamed batch is compared against the buffered side, so +# that key must be picked correctly out of an unsorted batch that also contains a NULL. +# One INSERT -> one batch of {7, NULL, 9, 3}: the deciding key sits at neither end, and a +# NULL must never be chosen. For `>` the decider is min=3 (buffered 5,4 exceed it; 2,1 do +# not); for `<` it is max=9 (every buffered row is below it). +statement ok +CREATE TABLE ex_extreme_r(v INT); + +statement ok +INSERT INTO ex_extreme_r VALUES (7), (NULL), (9), (3); + +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_extreme_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_extreme_r r WHERE l.v > r.v) ORDER BY 1; +---- +3 +4 + +query I +SELECT l.id FROM ex_l l WHERE EXISTS (SELECT 1 FROM ex_extreme_r r WHERE l.v < r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +query I +SELECT l.id FROM ex_l l WHERE NOT EXISTS (SELECT 1 FROM ex_extreme_r r WHERE l.v < r.v) ORDER BY 1; +---- + +# ------------------------------------------------------------------ +# Existence joins: streamed-side batch and partition layout +# ------------------------------------------------------------------ +# The existence path lowers the shared watermark incrementally as streamed batches arrive, so +# the batch and partition layout of the streamed side is load-bearing. Each INSERT below +# adds one batch, and the EXPLAINs assert the layout the row expectations rely on. + +statement ok +CREATE TABLE ex_layout_l(id INT, v INT); + +statement ok +INSERT INTO ex_layout_l VALUES (1, 1), (2, 3), (3, 5), (4, 7); + +statement ok +CREATE TABLE ex_layout_r(v INT); + +# Only the ordering within a streamed batch is enforced, so batches may arrive in any +# order. This deliberately non-monotonic arrival order drives the cross-batch low-water +# mark through all three of its cases: +# batch 1 (v=2) matches low -> marks buffered v={3,5,7}, watermark drops to index 1 +# batch 2 (v=6) matches above the watermark -> scan stops at it, nothing written +# batch 3 (v=0) matches lower still-> marks the remaining v=1 +# So all four buffered rows match. Re-marking greedily or skipping batch 3 both change this. +statement ok +INSERT INTO ex_layout_r VALUES (2); + +statement ok +INSERT INTO ex_layout_r VALUES (6); + +statement ok +INSERT INTO ex_layout_r VALUES (0); + +query I +SELECT l.id FROM ex_layout_l l WHERE EXISTS (SELECT 1 FROM ex_layout_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +query I +SELECT l.id FROM ex_layout_l l WHERE NOT EXISTS (SELECT 1 FROM ex_layout_r r WHERE l.v > r.v) ORDER BY 1; +---- + +# Asserts the streamed side really arrives as three batches in a single partition +# (`partition_sizes=[3]`), which is what the batch ordering above depends on. +query TT +EXPLAIN SELECT l.id FROM ex_layout_l l WHERE EXISTS (SELECT 1 FROM ex_layout_r r WHERE l.v > r.v) ORDER BY 1; +---- +logical_plan +01)Sort: l.id ASC NULLS LAST +02)--Projection: l.id +03)----LeftSemi Join: Filter: l.v > __correlated_sq_1.v +04)------SubqueryAlias: l +05)--------TableScan: ex_layout_l projection=[id, v] +06)------SubqueryAlias: __correlated_sq_1 +07)--------SubqueryAlias: r +08)----------TableScan: ex_layout_r projection=[v] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--ProjectionExec: expr=[id@0 as id] +03)----PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(v > v) +04)------SortExec: expr=[v@1 ASC], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)------DataSourceExec: partitions=1, partition_sizes=[3] + +# Same three batches, now also repartitioned across streamed partitions. The filter admits +# {2,6}, so buffered 3,5,7 match and only v=1 does not. This covers marking across batches +# and the cross-partition final-pass counter together, for both semi and anti. +query I +SELECT l.id FROM ex_layout_l l WHERE EXISTS (SELECT 1 FROM ex_layout_r r WHERE l.v > r.v AND r.v > 0) ORDER BY 1; +---- +2 +3 +4 + +query I +SELECT l.id FROM ex_layout_l l WHERE NOT EXISTS (SELECT 1 FROM ex_layout_r r WHERE l.v > r.v AND r.v > 0) ORDER BY 1; +---- +1 + +query TT +EXPLAIN SELECT l.id FROM ex_layout_l l WHERE NOT EXISTS (SELECT 1 FROM ex_layout_r r WHERE l.v > r.v AND r.v > 0) ORDER BY 1; +---- +logical_plan +01)Sort: l.id ASC NULLS LAST +02)--Projection: l.id +03)----LeftAnti Join: Filter: l.v > __correlated_sq_1.v +04)------SubqueryAlias: l +05)--------TableScan: ex_layout_l projection=[id, v] +06)------SubqueryAlias: __correlated_sq_1 +07)--------SubqueryAlias: r +08)----------Filter: ex_layout_r.v > Int32(0) +09)------------TableScan: ex_layout_r projection=[v] +physical_plan +01)SortPreservingMergeExec: [id@0 ASC NULLS LAST] +02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----ProjectionExec: expr=[id@0 as id] +04)------PiecewiseMergeJoin: operator=Gt, join_type=LeftAnti, on=(v > v) +05)--------SortExec: expr=[v@1 ASC], preserve_partitioning=[false] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] +07)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +08)----------FilterExec: v@0 > 0 +09)------------DataSourceExec: partitions=1, partition_sizes=[3] + +# The first matching buffered row is located by binary search, so the boundary cases matter. +# `ex_layout_l` holds v = {1,3,5,7} and nothing has lowered the watermark for these tables, +# so each search runs over the full buffered range. + +# Boundary at the LAST buffered row: only 7 > 6. +statement ok +CREATE TABLE ex_bs_last(v INT); + +statement ok +INSERT INTO ex_bs_last VALUES (6); + +query I +SELECT l.id FROM ex_layout_l l WHERE EXISTS (SELECT 1 FROM ex_bs_last r WHERE l.v > r.v) ORDER BY 1; +---- +4 + +query I +SELECT l.id FROM ex_layout_l l WHERE NOT EXISTS (SELECT 1 FROM ex_bs_last r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 +3 + +# No boundary at all: no buffered row exceeds 9, so the search must run off the end. +statement ok +CREATE TABLE ex_bs_none(v INT); + +statement ok +INSERT INTO ex_bs_none VALUES (9); + +query I +SELECT l.id FROM ex_layout_l l WHERE EXISTS (SELECT 1 FROM ex_bs_none r WHERE l.v > r.v) ORDER BY 1; +---- + +query I +SELECT l.id FROM ex_layout_l l WHERE NOT EXISTS (SELECT 1 FROM ex_bs_none r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +# ------------------------------------------------------------------ +# Existence joins: duplicate buffered keys across the match boundary +# ------------------------------------------------------------------ +# Every fixture above uses distinct buffered keys, so the boundary always fell between two +# different values. The binary search must return the FIRST index of a run of equal keys -- +# landing on a later one silently drops the earlier duplicates from EXISTS. `ex_dup_l` holds +# v = {5,5,3,3,1} and the deciding streamed key is 3, so for `<=` and `>=` the boundary falls +# *inside* the run of 3s. +statement ok +CREATE TABLE ex_dup_l(id INT, v INT); + +statement ok +INSERT INTO ex_dup_l VALUES (1, 5), (2, 5), (3, 3), (4, 3), (5, 1); + +statement ok +CREATE TABLE ex_dup_r(v INT); + +statement ok +INSERT INTO ex_dup_r VALUES (3); + +# `<` : only v=1 is below 3. +query I +SELECT l.id FROM ex_dup_l l WHERE EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v < r.v) ORDER BY 1; +---- +5 + +query I +SELECT l.id FROM ex_dup_l l WHERE NOT EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v < r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +# `<=` : the boundary is the first of the two v=3 rows, so BOTH must appear. +query I +SELECT l.id FROM ex_dup_l l WHERE EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v <= r.v) ORDER BY 1; +---- +3 +4 +5 + +query I +SELECT l.id FROM ex_dup_l l WHERE NOT EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v <= r.v) ORDER BY 1; +---- +1 +2 + +# `>` : only the two v=5 rows exceed 3. +query I +SELECT l.id FROM ex_dup_l l WHERE EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_dup_l l WHERE NOT EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v > r.v) ORDER BY 1; +---- +3 +4 +5 + +# `>=` : the boundary is again inside the run of 3s, now from the ascending side. +query I +SELECT l.id FROM ex_dup_l l WHERE EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v >= r.v) ORDER BY 1; +---- +1 +2 +3 +4 + +query I +SELECT l.id FROM ex_dup_l l WHERE NOT EXISTS (SELECT 1 FROM ex_dup_r r WHERE l.v >= r.v) ORDER BY 1; +---- +5 + +# ------------------------------------------------------------------ +# Existence joins: NULL-only and multi-NULL buffered side +# ------------------------------------------------------------------ +# An all-NULL buffered side saturates the watermark before the first poll -- a different +# early-exit trigger from the "every row marked" case above -- so no streamed batch is read +# at all and EXISTS must still be empty. +statement ok +CREATE TABLE ex_l_allnull(id INT, v INT); + +statement ok +INSERT INTO ex_l_allnull VALUES (1, NULL), (2, NULL); + +query I +SELECT l.id FROM ex_l_allnull l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE l.v < r.v) ORDER BY 1; +---- + +query I +SELECT l.id FROM ex_l_allnull l WHERE NOT EXISTS (SELECT 1 FROM ex_r r WHERE l.v < r.v) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_l_allnull l WHERE EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v) ORDER BY 1; +---- + +query I +SELECT l.id FROM ex_l_allnull l WHERE NOT EXISTS (SELECT 1 FROM ex_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 + +# More than one NULL on the buffered side: the scan starts at `null_count()`, so a second +# NULL shifts that offset. Buffered v = {NULL,3,NULL,1} sorts to [NULL,NULL,3,1] for `<`. +statement ok +CREATE TABLE ex_multinull_l(id INT, v INT); + +statement ok +INSERT INTO ex_multinull_l VALUES (1, NULL), (2, 3), (3, NULL), (4, 1); + +statement ok +CREATE TABLE ex_multinull_r(v INT); + +statement ok +INSERT INTO ex_multinull_r VALUES (2); + +# `<` : only v=1 is below 2. Both NULLs stay in NOT EXISTS. +query I +SELECT l.id FROM ex_multinull_l l WHERE EXISTS (SELECT 1 FROM ex_multinull_r r WHERE l.v < r.v) ORDER BY 1; +---- +4 + +query I +SELECT l.id FROM ex_multinull_l l WHERE NOT EXISTS (SELECT 1 FROM ex_multinull_r r WHERE l.v < r.v) ORDER BY 1; +---- +1 +2 +3 + +# `>` : only v=3 exceeds 2. +query I +SELECT l.id FROM ex_multinull_l l WHERE EXISTS (SELECT 1 FROM ex_multinull_r r WHERE l.v > r.v) ORDER BY 1; +---- +2 + +query I +SELECT l.id FROM ex_multinull_l l WHERE NOT EXISTS (SELECT 1 FROM ex_multinull_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +3 +4 + +# ------------------------------------------------------------------ +# Existence joins: float keys and negative zero +# ------------------------------------------------------------------ +# The existence path shares `JoinKeyComparator`, which normalizes `-0.0` before building the +# comparator because arrow's `make_comparator` uses IEEE totalOrder and would rank `-0.0` +# below `+0.0`. SQL says they are equal, so `-0.0 < 0.0` must be false and `-0.0 <= 0.0` +# true. Skipping the normalization would put id 1 in the `<` result. +statement ok +CREATE TABLE ex_fl(id INT, v DOUBLE); + +statement ok +INSERT INTO ex_fl VALUES (1, -0.0), (2, 2.5); + +statement ok +CREATE TABLE ex_fr(v DOUBLE); + +statement ok +INSERT INTO ex_fr VALUES (0.0); + +query I +SELECT l.id FROM ex_fl l WHERE EXISTS (SELECT 1 FROM ex_fr r WHERE l.v < r.v) ORDER BY 1; +---- + +query I +SELECT l.id FROM ex_fl l WHERE NOT EXISTS (SELECT 1 FROM ex_fr r WHERE l.v < r.v) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_fl l WHERE EXISTS (SELECT 1 FROM ex_fr r WHERE l.v <= r.v) ORDER BY 1; +---- +1 + +query I +SELECT l.id FROM ex_fl l WHERE NOT EXISTS (SELECT 1 FROM ex_fr r WHERE l.v <= r.v) ORDER BY 1; +---- +2 + statement ok set datafusion.optimizer.enable_piecewise_merge_join = false; From 2b73c41200016c9da23ae429e5849d28a3b7131f Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Wed, 12 Aug 2026 18:49:50 +0530 Subject: [PATCH 5/6] pick the streamed extreme key with arrow min/max kernels --- .../piecewise_merge_join/existence_join.rs | 54 ++++---- datafusion/sqllogictest/test_files/pwmj.slt | 118 ++++++++++++++++++ 2 files changed, 149 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs index b66f6e01fed71..b2c5212999f3b 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/existence_join.rs @@ -67,9 +67,8 @@ //! //! Rows whose join key is NULL never satisfy a comparison predicate. Buffered NULLs sort to //! the front, so the scan starts past them and null-keyed buffered rows are left unmarked — -//! correctly excluded from `LeftSemi` and included in `LeftAnti`. Each streamed batch is -//! sorted NULLs-*last* instead, so the extreme key taken from it is non-null unless the whole -//! batch is. +//! correctly excluded from `LeftSemi` and included in `LeftAnti`. The extreme key is picked +//! from each streamed batch with NULLs ignored, so it is non-null unless the whole batch is. //! //! # Output //! @@ -91,11 +90,12 @@ use std::sync::atomic::Ordering as AtomicOrdering; use std::task::{Poll, ready}; use arrow::array::{Array, ArrayRef, RecordBatch}; -use arrow::compute::{BatchCoalescer, sort_to_indices, take}; +use arrow::compute::BatchCoalescer; use arrow_schema::{SchemaRef, SortOptions}; use datafusion_common::{NullEquality, Result, internal_err}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; use datafusion_expr::{JoinType, Operator}; +use datafusion_functions_aggregate_common::min_max::{max_batch, min_batch}; use datafusion_physical_expr::PhysicalExprRef; use futures::{Stream, StreamExt}; @@ -109,7 +109,7 @@ use crate::stream::EmptyRecordBatchStream; pub(super) enum ExistencePWMJStreamState { /// Load the buffered side into memory. WaitBufferedSide, - /// Fetch, sort and scan streamed batches, lowering the watermark. Emits nothing. + /// Fetch and scan streamed batches, lowering the watermark. Emits nothing. ScanStreamBatches, /// Emit the result. Reached only by the last streamed partition. EmitMatched, @@ -210,7 +210,8 @@ impl ExistencePWMJStream { Poll::Ready(Ok(StatefulStreamResult::Continue)) } - /// Fetches one streamed batch, sorts it, and marks the buffered rows it matches. + /// Fetches one streamed batch, reduces it to its extreme compare key, and marks the + /// buffered rows that key matches. /// Never produces output: existence results come from the watermark in `emit_matched`. fn scan_stream_batch( &mut self, @@ -237,23 +238,10 @@ impl ExistencePWMJStream { self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(batch.num_rows()); - // Only the batch's extreme key is ever compared against the buffered - // side, so ask arrow for just that one index. `nulls_first: false` makes - // the chosen row non-null whenever the batch holds any non-null key, and - // leaves a null row only for an all-null batch, which marks nothing. - // - // The limit truncates the returned indices, not the work: `sort_impl` - // feeds it into the sort only when `nulls_first` is true, so this still - // fully sorts the batch's non-null keys. - let indices = sort_to_indices( - stream_values.as_ref(), - Some(SortOptions { - descending: self.sort_option.descending, - nulls_first: false, - }), - Some(1), - )?; - let stream_values = take(stream_values.as_ref(), &indices, None)?; + // Only the batch's extreme key is ever compared against the buffered side, + // so reduce the batch to that one key. + let stream_values = + extreme_key(&stream_values, self.sort_option.descending)?; self.mark_matched_buffered_rows(&stream_values)?; } @@ -451,6 +439,26 @@ impl ExistencePWMJStream { } } +/// Reduces `values` to a one-row array holding the batch's extreme compare key: the maximum +/// when the batch's sort is `descending`, the minimum otherwise. Nulls are ignored, so that +/// row is null only when every key in the batch is (or the batch is empty), which marks +/// nothing. +/// +/// Ordered the same way as [`JoinKeyComparator`]: both use IEEE 754 totalOrder for floats, +/// and the comparator normalizes `-0.0` on either side of it. +/// +/// Numeric, temporal, string, binary and boolean keys get a typed arrow kernel -- a linear +/// scan that allocates nothing. Dictionary and nested keys fall to `min_max_batch_generic`, a +/// `ScalarValue`-per-row comparator loop; specializing those is left to a follow-up. +fn extreme_key(values: &ArrayRef, descending: bool) -> Result { + let extreme = if descending { + max_batch(values)? + } else { + min_batch(values)? + }; + extreme.to_array_of_size(1) +} + impl RecordBatchStream for ExistencePWMJStream { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 8ed4d27022898..493fd5afd1887 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -1046,5 +1046,123 @@ SELECT l.id FROM ex_fl l WHERE NOT EXISTS (SELECT 1 FROM ex_fr r WHERE l.v <= r. ---- 2 +# ------------------------------------------------------------------ +# Classic joins: final pass when the streamed side has fewer partitions than +# `target_partitions` +# ------------------------------------------------------------------ +# The final pass -- the only thing that emits unmatched buffered rows for `Left`/`Full` -- runs +# on the last streamed partition to finish, gated by a counter seeded with the number of +# partitions that will call `execute`. That is the streamed side's partition count, not +# `target_partitions`: the streamed side carries no distribution requirement, so a small input +# stays single-partition while `target_partitions` is 4 (asserted by the plan below). Seeding +# from `target_partitions` leaves the counter at 3, no partition ever emits, and `x = 10` +# disappears from both queries. +statement ok +CREATE TABLE pwmj_unmatched_l(x INT); + +statement ok +INSERT INTO pwmj_unmatched_l VALUES (1), (10); + +statement ok +CREATE TABLE pwmj_unmatched_r(y INT); + +statement ok +INSERT INTO pwmj_unmatched_r VALUES (5); + +query II +SELECT l.x, r.y FROM pwmj_unmatched_l l LEFT JOIN pwmj_unmatched_r r ON l.x < r.y ORDER BY 1; +---- +1 5 +10 NULL + +query II +SELECT l.x, r.y FROM pwmj_unmatched_l l FULL JOIN pwmj_unmatched_r r ON l.x < r.y ORDER BY 1; +---- +1 5 +10 NULL + +query TT +EXPLAIN SELECT l.x, r.y FROM pwmj_unmatched_l l LEFT JOIN pwmj_unmatched_r r ON l.x < r.y ORDER BY 1; +---- +logical_plan +01)Sort: l.x ASC NULLS LAST +02)--Left Join: Filter: l.x < r.y +03)----SubqueryAlias: l +04)------TableScan: pwmj_unmatched_l projection=[x] +05)----SubqueryAlias: r +06)------TableScan: pwmj_unmatched_r projection=[y] +physical_plan +01)SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--PiecewiseMergeJoin: operator=Lt, join_type=Left, on=(x < y) +03)----SortExec: expr=[x@0 DESC], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] + +# ------------------------------------------------------------------ +# Existence joins: dictionary-encoded keys +# ------------------------------------------------------------------ +# Each streamed batch is reduced to its extreme key. Dictionaries have no typed arrow min/max +# kernel, so they take the generic `ScalarValue`-per-row path -- this is the only coverage of +# that path. +statement ok +CREATE TABLE ex_dict_l AS + SELECT column1 AS id, arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS v + FROM (VALUES (1, 'a'), (2, 'c'), (3, 'e'), (4, NULL)); + +statement ok +CREATE TABLE ex_dict_r AS + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS v FROM (VALUES ('c')); + +query T +SELECT arrow_typeof(v) FROM ex_dict_l LIMIT 1; +---- +Dictionary(Int32, Utf8) + +# `>` : only 'e' exceeds 'c'. The NULL key is never marked, so it stays in NOT EXISTS. +query I +SELECT l.id FROM ex_dict_l l WHERE EXISTS (SELECT 1 FROM ex_dict_r r WHERE l.v > r.v) ORDER BY 1; +---- +3 + +query I +SELECT l.id FROM ex_dict_l l WHERE NOT EXISTS (SELECT 1 FROM ex_dict_r r WHERE l.v > r.v) ORDER BY 1; +---- +1 +2 +4 + +# `<=` : 'a' and 'c'. +query I +SELECT l.id FROM ex_dict_l l WHERE EXISTS (SELECT 1 FROM ex_dict_r r WHERE l.v <= r.v) ORDER BY 1; +---- +1 +2 + +query I +SELECT l.id FROM ex_dict_l l WHERE NOT EXISTS (SELECT 1 FROM ex_dict_r r WHERE l.v <= r.v) ORDER BY 1; +---- +3 +4 + +query TT +EXPLAIN SELECT l.id FROM ex_dict_l l WHERE EXISTS (SELECT 1 FROM ex_dict_r r WHERE l.v > r.v) ORDER BY 1; +---- +logical_plan +01)Sort: l.id ASC NULLS LAST +02)--Projection: l.id +03)----LeftSemi Join: Filter: l.v > __correlated_sq_1.v +04)------SubqueryAlias: l +05)--------TableScan: ex_dict_l projection=[id, v] +06)------SubqueryAlias: __correlated_sq_1 +07)--------SubqueryAlias: r +08)----------TableScan: ex_dict_r projection=[v] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--ProjectionExec: expr=[id@0 as id] +03)----PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(v > v) +04)------SortExec: expr=[v@1 ASC], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)------DataSourceExec: partitions=1, partition_sizes=[1] + statement ok set datafusion.optimizer.enable_piecewise_merge_join = false; From d3cec1334861894a69127190d13f9f9ffe26ce8c Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Wed, 12 Aug 2026 19:42:42 +0530 Subject: [PATCH 6/6] Adds fuzz test --- datafusion/core/tests/fuzz_cases/join_fuzz.rs | 223 +++++++++++++++++- 1 file changed, 222 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 81c7c9f83928e..fce6999fd0d77 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -34,18 +34,24 @@ use datafusion::physical_plan::collect; use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + HashJoinExec, NestedLoopJoinExec, PartitionMode, PiecewiseMergeJoinExec, + SortMergeJoinExec, }; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, common}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::{NullEquality, ScalarValue}; +use datafusion_common_runtime::SpawnedTask; use datafusion_execution::TaskContext; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::expressions::Literal; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use itertools::Itertools; use rand::Rng; +use rand::{SeedableRng, rngs::StdRng}; use test_utils::stagger_batch_with_seed; // Determines what Fuzz tests needs to run @@ -1347,3 +1353,218 @@ fn make_staggered_batches_binary( // preserve your existing randomized partitioning stagger_batch_with_seed(batch, 42) } + +// ---- Differential fuzz: PiecewiseMergeJoin existence joins vs NestedLoopJoin ---- +// +// `PiecewiseMergeJoinExec` takes a range predicate and no equi keys, so it cannot be added to +// `JoinFuzzTestCase` above (that harness joins on `a` and `b` and folds the equality into the +// NestedLoopJoin filter). These tests use `NestedLoopJoinExec` with an equivalent filter as +// the oracle instead. +// +// What only randomization reaches: `LeftSemi`/`LeftAnti` record matches in a shared +// `AtomicUsize` watermark and emit once, from whichever streamed partition finishes last. The +// streamed side below is spread round-robin over several partitions as one-row batches, so +// batches arrive in an order no static test pins down, and the counter that gates the final +// pass is seeded from a partition count that deliberately disagrees with the `num_partitions` +// argument. + +fn pwmj_kv_schema() -> Arc { + Arc::new(Schema::new(vec![ + arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false), + arrow::datatypes::Field::new("k", arrow::datatypes::DataType::Int32, true), + ])) +} + +fn pwmj_kv_batch(ids: &[i32], keys: &[Option]) -> RecordBatch { + RecordBatch::try_new( + pwmj_kv_schema(), + vec![ + Arc::new(Int32Array::from(ids.to_vec())), + Arc::new(Int32Array::from(keys.to_vec())), + ], + ) + .unwrap() +} + +/// Single-partition, single-batch input: the buffered side, and the oracle's probe side. +fn pwmj_single_exec(ids: &[i32], keys: &[Option]) -> Arc { + MemorySourceConfig::try_new_exec( + &[vec![pwmj_kv_batch(ids, keys)]], + pwmj_kv_schema(), + None, + ) + .unwrap() +} + +/// Streamed side spread round-robin across `nparts` partitions, one row per batch. +fn pwmj_parts_exec( + ids: &[i32], + keys: &[Option], + nparts: usize, +) -> Arc { + let nparts = nparts.max(1); + let mut partitions: Vec> = vec![Vec::new(); nparts]; + for (row, (&id, key)) in ids.iter().zip(keys.iter()).enumerate() { + partitions[row % nparts].push(pwmj_kv_batch(&[id], &[*key])); + } + for p in partitions.iter_mut() { + if p.is_empty() { + p.push(pwmj_kv_batch(&[], &[])); + } + } + MemorySourceConfig::try_new_exec(&partitions, pwmj_kv_schema(), None).unwrap() +} + +fn pwmj_existence_plan( + left: Arc, + right: Arc, + op: Operator, + join_type: JoinType, +) -> Arc { + // Matches `PiecewiseMergeJoinExec::required_input_ordering`: descending for `<`/`<=`, + // ascending for `>`/`>=`, NULLs first either way. + let sort_options = match op { + Operator::Lt | Operator::LtEq => SortOptions::new(true, true), + Operator::Gt | Operator::GtEq => SortOptions::new(false, true), + other => panic!("not a range operator: {other:?}"), + }; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("k", 1)), + sort_options, + )]) + .unwrap(); + let sorted_left = Arc::new(SortExec::new(ordering, left)); + let on: (PhysicalExprRef, PhysicalExprRef) = + (Arc::new(Column::new("k", 1)), Arc::new(Column::new("k", 1))); + // `num_partitions` is 1 while the streamed side has up to 3: the final-pass counter must + // come from the streamed side's partition count, not from this argument. + Arc::new( + PiecewiseMergeJoinExec::try_new(sorted_left, right, on, op, join_type, 1) + .unwrap(), + ) +} + +fn pwmj_nlj_oracle_plan( + left: Arc, + right: Arc, + op: Operator, + join_type: JoinType, +) -> Arc { + let intermediate_schema = Schema::new(vec![ + arrow::datatypes::Field::new("k", arrow::datatypes::DataType::Int32, true), + arrow::datatypes::Field::new("k", arrow::datatypes::DataType::Int32, true), + ]); + let expr = Arc::new(BinaryExpr::new( + Arc::new(Column::new("k", 0)), + op, + Arc::new(Column::new("k", 1)), + )) as PhysicalExprRef; + let column_indices = vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + let filter = JoinFilter::new(expr, column_indices, Arc::new(intermediate_schema)); + Arc::new( + NestedLoopJoinExec::try_new(left, right, Some(filter), &join_type, None).unwrap(), + ) +} + +/// Executes every output partition concurrently and returns the surviving left `id`s, sorted. +/// +/// Concurrent rather than one partition at a time: the partitions share the watermark and race +/// to be the one that runs the final pass, which is the part a sequential drain cannot reach. +async fn pwmj_collect_ids( + plan: Arc, + task_ctx: Arc, +) -> Vec { + let streams = (0..plan.output_partitioning().partition_count()) + .map(|partition| plan.execute(partition, Arc::clone(&task_ctx)).unwrap()) + .collect::>(); + let per_partition = + futures::future::join_all(streams.into_iter().map(|stream| { + SpawnedTask::spawn(async move { common::collect(stream).await }) + })) + .await; + + let mut ids = Vec::new(); + for batches in per_partition { + for batch in batches.unwrap().unwrap() { + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..col.len()).map(|i| col.value(i))); + } + } + ids.sort(); + ids +} + +#[tokio::test(flavor = "multi_thread")] +async fn fuzz_pwmj_existence_matches_nested_loop() { + // A small batch size splits the final-pass output across several coalesced batches even + // for these tiny inputs, covering that boundary too. + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(3)), + ); + let ops = [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq]; + + for seed in 0..60u64 { + let mut rng = StdRng::seed_from_u64(seed); + let left_len = rng.random_range(0..25usize); + let right_len = rng.random_range(0..25usize); + // A narrow key range forces duplicates and equal-boundary cases. + let key_range = rng.random_range(1..6i32); + let nparts = rng.random_range(1..4usize); + + let gen_keys = |n: usize, rng: &mut StdRng| -> Vec> { + (0..n) + .map(|_| (!rng.random_bool(0.2)).then(|| rng.random_range(0..key_range))) + .collect() + }; + + let left_ids: Vec = (0..left_len as i32).collect(); + let left_keys = gen_keys(left_len, &mut rng); + let right_ids: Vec = (0..right_len as i32).collect(); + let right_keys = gen_keys(right_len, &mut rng); + + for op in ops { + for join_type in [JoinType::LeftSemi, JoinType::LeftAnti] { + let got = pwmj_collect_ids( + pwmj_existence_plan( + pwmj_single_exec(&left_ids, &left_keys), + pwmj_parts_exec(&right_ids, &right_keys, nparts), + op, + join_type, + ), + Arc::clone(&task_ctx), + ) + .await; + let want = pwmj_collect_ids( + pwmj_nlj_oracle_plan( + pwmj_single_exec(&left_ids, &left_keys), + pwmj_single_exec(&right_ids, &right_keys), + op, + join_type, + ), + Arc::clone(&task_ctx), + ) + .await; + + assert_eq!( + got, want, + "mismatch seed={seed} op={op:?} join_type={join_type:?} \ + nparts={nparts} left_keys={left_keys:?} right_keys={right_keys:?}" + ); + } + } + } +}