From 5ec1790897998f9ef1b21dd710f685dfd93fe771 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 24 Jul 2026 21:28:40 +0530 Subject: [PATCH] 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;