From 16b386d1a83ed9eb1207865f3f54b64f3c5d574d Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Sat, 5 Sep 2026 11:55:50 +0530 Subject: [PATCH 1/2] bench: add PWMJ vs NestedLoopJoin criterion coverage for LeftMark/RigMark --- datafusion/core/benches/pwmj_semi_anti_sql.rs | 241 +++++++++++++++++- 1 file changed, 231 insertions(+), 10 deletions(-) diff --git a/datafusion/core/benches/pwmj_semi_anti_sql.rs b/datafusion/core/benches/pwmj_semi_anti_sql.rs index d86d241419baa..ce31feb93ca3a 100644 --- a/datafusion/core/benches/pwmj_semi_anti_sql.rs +++ b/datafusion/core/benches/pwmj_semi_anti_sql.rs @@ -43,7 +43,8 @@ //! //! ## Axes //! - **join type**: `EXISTS` (LeftSemi) and `NOT EXISTS` (LeftAnti), plus `RIGHT SEMI` and -//! `RIGHT ANTI` written out as joins. +//! `RIGHT ANTI` written out as joins, and `EXISTS` inside an always-false disjunction +//! (LeftMark). //! - **match regime**: the fraction of rows on the marked side that have at least one match //! on the other, set by shifting the right-side key range relative to the left one: //! `all_match` (100%), `no_match` (0%) and `half_match` (~50%, where the buffered side @@ -65,6 +66,26 @@ //! The **same three key offsets** serve both halves, each read against the other side's //! extreme: a left row survives `EXISTS` iff `lhs.key < max(rhs.key)`, a right row survives //! `RIGHT SEMI` iff `min(lhs.key) < rhs.key`. +//! +//! ## Mark joins +//! `LeftMark` reuses the exact same watermark-marking path as `LeftSemi` -- the only +//! difference is the final pass, which appends a `mark` column instead of slicing the batch +//! -- so it is folded into the same SQL sweep above via an `EXISTS` wrapped in an +//! always-false `OR`, the one shape that decorrelates to it (see `Kind::LeftMark`'s doc). It +//! needs no build-dependent handling beyond what `Kind::LeftMark` already gets from +//! `PWMJ_OR_NLJ`: on a build that does not yet route `LeftMark` to PWMJ, the planner itself +//! falls back to `NestedLoopJoinExec`, same as every other case here before its dependency +//! landed. +//! +//! `RightMark` has no SQL surface at all: no keyword parses to it, and no optimizer rule +//! constructs it (`decorrelate_predicate_subquery.rs` only ever builds `LeftMark`), so +//! `ctx.sql(...)` can never plan one. `bench_pwmj_right_mark_hand_built`, in a separate +//! benchmark group below, builds `PiecewiseMergeJoinExec` and `NestedLoopJoinExec` directly +//! instead -- the one place in this file that deviates from planning through SQL, forced by +//! there being no SQL to plan. That also means it has no planner fallback to lean on if the +//! build does not support `RightMark` yet, so it probes `try_new` up front and skips the +//! whole group with a note instead of panicking, the hand-built equivalent of what the +//! planner does automatically for every SQL-planned case in this file. use std::sync::Arc; @@ -75,8 +96,15 @@ use criterion::{ BatchSize, BenchmarkGroup, BenchmarkId, Criterion, criterion_group, criterion_main, }; use datafusion::datasource::MemTable; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::logical_expr::{JoinType, Operator}; +use datafusion::physical_expr::expressions::{BinaryExpr, Column}; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion::physical_plan::joins::{NestedLoopJoinExec, PiecewiseMergeJoinExec}; use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::JoinSide; use tokio::runtime::Runtime; const LEFT_ROWS: usize = 20_000; @@ -171,14 +199,19 @@ enum Kind { RightSemi, /// `RIGHT ANTI`, written out. RightAnti, + /// `EXISTS` inside an always-false disjunction, decorrelated to `LeftMark`: the outer + /// filter then needs the subquery's match result as a value rather than as a row filter + /// (see `decorrelate_predicate_subquery.rs`). No SQL syntax reaches `LeftMark` directly. + LeftMark, } impl Kind { /// Whether this join keeps the marked rows that matched, rather than the ones that did not. /// - /// The two are complements, which is why one expected row count covers both. + /// The two are complements, which is why one expected row count covers both. `LeftMark`'s + /// disjunct is always false, so its final row set is the same as `LeftSemi`'s. fn keeps_matched(self) -> bool { - matches!(self, Kind::LeftSemi | Kind::RightSemi) + matches!(self, Kind::LeftSemi | Kind::RightSemi | Kind::LeftMark) } /// The query, over the range relation `lhs.key < rhs.key`, projecting whichever side this @@ -201,6 +234,13 @@ impl Kind { "SELECT rhs.key, rhs.payload FROM lhs \ RIGHT ANTI JOIN rhs ON lhs.key < rhs.key" } + // `lhs.payload` is always >= 0 (see `build_batches`), so `< 0` is always false and + // the `OR` degenerates to the `EXISTS` alone -- but only after decorrelating to a + // `LeftMark` join whose `mark` the outer `Filter` reads. + Kind::LeftMark => { + "SELECT lhs.key, lhs.payload FROM lhs \ + WHERE lhs.payload < 0 OR EXISTS (SELECT 1 FROM rhs WHERE lhs.key < rhs.key)" + } } } @@ -212,14 +252,19 @@ impl Kind { Kind::LeftAnti => "anti", Kind::RightSemi => "right_semi", Kind::RightAnti => "right_anti", + Kind::LeftMark => "left_mark", } } - /// Plan fragment every arm must contain. Only the Semi/Anti half of the join type is pinned, - /// not the side: the planner is free to swap the nested-loop inputs, so the arms can disagree - /// on the side while computing the same thing. + /// Plan fragment every arm must contain. Only the Semi/Anti/Mark half of the join type is + /// pinned, not the side: the planner is free to swap the nested-loop inputs, so the arms + /// can disagree on the side while computing the same thing. fn plan_fragment(self) -> &'static str { - if self.keeps_matched() { "Semi" } else { "Anti" } + match self { + Kind::LeftSemi | Kind::RightSemi => "Semi", + Kind::LeftAnti | Kind::RightAnti => "Anti", + Kind::LeftMark => "Mark", + } } /// Rows on the marked side: this join emits either those that matched or the rest. Both @@ -227,18 +272,19 @@ impl Kind { /// whole difference between the two halves. fn marked_rows(self) -> usize { match self { - Kind::LeftSemi | Kind::LeftAnti => LEFT_ROWS, + Kind::LeftSemi | Kind::LeftAnti | Kind::LeftMark => LEFT_ROWS, Kind::RightSemi | Kind::RightAnti => RIGHT_ROWS, } } } /// Every join type covered, in the order cases run. -const KINDS: [Kind; 4] = [ +const KINDS: [Kind; 5] = [ Kind::LeftSemi, Kind::LeftAnti, Kind::RightSemi, Kind::RightAnti, + Kind::LeftMark, ]; /// How much of the marked side has a match, set by where the `rhs` key range sits relative to @@ -472,5 +518,180 @@ fn bench_pwmj_semi_anti_sql(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_pwmj_semi_anti_sql); +/// A `JoinFilter` for `lhs.key < rhs.key`, the same range relation every case in this file +/// uses, for the `NestedLoopJoinExec` arm below -- which needs one built by hand since it is +/// constructed directly rather than planned from SQL. +fn key_lt_key_filter(s: &SchemaRef) -> JoinFilter { + let expr = Arc::new(BinaryExpr::new( + Arc::new(Column::new("key", 0)), + Operator::Lt, + Arc::new(Column::new("key", 1)), + )) as _; + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let key_field = s.field_with_name("key").unwrap().clone(); + let intermediate_schema = Schema::new(vec![key_field.clone(), key_field]); + JoinFilter::new(expr, column_indices, Arc::new(intermediate_schema)) +} + +/// `RightMark` has no SQL surface (see the module doc's "Mark joins" section), so both arms +/// are hand-built here instead of planned from SQL text: `PiecewiseMergeJoinExec` and +/// `NestedLoopJoinExec`, over the same data, the same `lhs.key < rhs.key` relation, and the +/// same join type -- with no risk of comparing an operator against itself, since which +/// operator each arm uses is fixed by construction rather than read back off a plan. +/// +/// Neither arm needs a `SortExec`: `RightMark`, like `RightSemi`/`RightAnti`, folds the +/// buffered side to one key regardless of its order, and `NestedLoopJoinExec` never needs +/// either side ordered. +fn bench_pwmj_right_mark_hand_built(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let s = schema(); + let ctx = SessionContext::new(); + + // Every other case in this file falls back to `NestedLoopJoinExec` through the planner + // itself when a build does not yet support the join type -- there is no such fallback + // here, since `RightMark` has no SQL surface to plan through in the first place. Probe + // `try_new` directly and skip the whole group rather than panic, so this benchmark stays + // runnable (as a no-op) against a build that has not merged `RightMark` support yet, and + // starts measuring on its own once that support lands. + if let Err(err) = PiecewiseMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&s))), + Arc::new(EmptyExec::new(Arc::clone(&s))), + ( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + ), + Operator::Lt, + JoinType::RightMark, + 1, + ) { + println!( + "note: pwmj_vs_nlj_right_mark_hand_built skipped -- this build's \ + PiecewiseMergeJoinExec does not support RightMark yet: {err}" + ); + return; + } + + let mut group = c.benchmark_group("pwmj_vs_nlj_right_mark_hand_built"); + group.sample_size(10); + + for (regime, right_offset, _fraction) in REGIMES { + let lhs_batches = build_batches(LEFT_ROWS, 0, &s); + let rhs_batches = build_batches(RIGHT_ROWS, right_offset, &s); + + let pwmj_plan = { + let (lhs_batches, rhs_batches, s) = + (lhs_batches.clone(), rhs_batches.clone(), Arc::clone(&s)); + move || -> Arc { + let lhs = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&lhs_batches), + Arc::clone(&s), + None, + ) + .unwrap(); + let rhs = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&rhs_batches), + Arc::clone(&s), + None, + ) + .unwrap(); + Arc::new( + PiecewiseMergeJoinExec::try_new( + lhs, + rhs, + ( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + ), + Operator::Lt, + JoinType::RightMark, + 1, + ) + .unwrap(), + ) + } + }; + let nlj_plan = { + let (lhs_batches, rhs_batches, s) = + (lhs_batches.clone(), rhs_batches.clone(), Arc::clone(&s)); + move || -> Arc { + let lhs = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&lhs_batches), + Arc::clone(&s), + None, + ) + .unwrap(); + let rhs = MemorySourceConfig::try_new_exec( + std::slice::from_ref(&rhs_batches), + Arc::clone(&s), + None, + ) + .unwrap(); + Arc::new( + NestedLoopJoinExec::try_new( + lhs, + rhs, + Some(key_lt_key_filter(&s)), + &JoinType::RightMark, + None, + ) + .unwrap(), + ) + } + }; + + // `RightMark` keeps every streamed row, matched or not, so both arms must return + // exactly `RIGHT_ROWS` regardless of the regime -- unlike `RightSemi`/`RightAnti`, + // where the regime changes the row count. The regime still matters to what is timed + // below: it changes how much of the comparison work each arm actually does (`mark` + // true vs false), even though the row count it returns cannot show that. + let pwmj_rows = run(pwmj_plan(), &ctx, &rt); + let nlj_rows = run(nlj_plan(), &ctx, &rt); + assert_eq!( + pwmj_rows, nlj_rows, + "right_mark_{regime}: pwmj and nlj disagree ({pwmj_rows} vs {nlj_rows} rows)" + ); + assert_eq!( + pwmj_rows, RIGHT_ROWS, + "right_mark_{regime}: RightMark must keep every streamed row" + ); + + group.bench_function( + BenchmarkId::new("pwmj", format!("{regime}_{RIGHT_ROWS}")), + |b| { + b.iter_batched( + pwmj_plan.clone(), + |plan| run(plan, &ctx, &rt), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("nlj", format!("{regime}_{RIGHT_ROWS}")), + |b| { + b.iter_batched( + nlj_plan.clone(), + |plan| run(plan, &ctx, &rt), + BatchSize::SmallInput, + ) + }, + ); + } + + group.finish(); +} + +criterion_group!( + benches, + bench_pwmj_semi_anti_sql, + bench_pwmj_right_mark_hand_built +); criterion_main!(benches); From 89a1400f2ea1ad3e35d8be19133fade2e716a871 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Sat, 5 Sep 2026 21:19:10 +0530 Subject: [PATCH 2/2] address review feedback --- datafusion/core/benches/pwmj_semi_anti_sql.rs | 136 ++++++++++++++---- 1 file changed, 111 insertions(+), 25 deletions(-) diff --git a/datafusion/core/benches/pwmj_semi_anti_sql.rs b/datafusion/core/benches/pwmj_semi_anti_sql.rs index ce31feb93ca3a..55827115b9df7 100644 --- a/datafusion/core/benches/pwmj_semi_anti_sql.rs +++ b/datafusion/core/benches/pwmj_semi_anti_sql.rs @@ -77,19 +77,23 @@ //! falls back to `NestedLoopJoinExec`, same as every other case here before its dependency //! landed. //! -//! `RightMark` has no SQL surface at all: no keyword parses to it, and no optimizer rule -//! constructs it (`decorrelate_predicate_subquery.rs` only ever builds `LeftMark`), so -//! `ctx.sql(...)` can never plan one. `bench_pwmj_right_mark_hand_built`, in a separate -//! benchmark group below, builds `PiecewiseMergeJoinExec` and `NestedLoopJoinExec` directly -//! instead -- the one place in this file that deviates from planning through SQL, forced by -//! there being no SQL to plan. That also means it has no planner fallback to lean on if the -//! build does not support `RightMark` yet, so it probes `try_new` up front and skips the -//! whole group with a note instead of panicking, the hand-built equivalent of what the -//! planner does automatically for every SQL-planned case in this file. +//! `RightMark` has no *direct* SQL surface: no keyword parses to it, and no logical rule +//! constructs it (`decorrelate_predicate_subquery.rs` only ever builds `LeftMark`). The +//! physical optimizer can still land on one indirectly -- `JoinSelection` may call +//! `NestedLoopJoinExec::swap_inputs`, which swaps `LeftMark` to `RightMark` via +//! `JoinType::swap` -- but there is no query text or table layout in this file that drives +//! that swap for a `LeftMark` plan, so `ctx.sql(...)` still never plans a `RightMark` here in +//! practice. `bench_pwmj_right_mark_hand_built`, in a separate benchmark group below, builds +//! `PiecewiseMergeJoinExec` and `NestedLoopJoinExec` directly instead -- a deliberate +//! direct-operator microbenchmark, not a stand-in for a SQL plan, and the one place in this +//! file that deviates from planning through SQL. That also means it has no planner fallback +//! to lean on if the build does not support `RightMark` yet, so it probes `try_new` up front +//! and skips the whole group with a note instead of panicking, the hand-built equivalent of +//! what the planner does automatically for every SQL-planned case in this file. use std::sync::Arc; -use arrow::array::{Int32Array, RecordBatch}; +use arrow::array::{Array, Int32Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::measurement::WallTime; use criterion::{ @@ -104,7 +108,7 @@ use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion::physical_plan::joins::{NestedLoopJoinExec, PiecewiseMergeJoinExec}; use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_common::JoinSide; +use datafusion_common::{DataFusionError, JoinSide}; use tokio::runtime::Runtime; const LEFT_ROWS: usize = 20_000; @@ -381,6 +385,48 @@ fn run(plan: Arc, ctx: &SessionContext, rt: &Runtime) -> usiz }) } +/// Counts of the `mark` column's values across every output batch: `(true, false, null)`. +/// +/// A row count alone cannot tell a correct mark join from a broken one: `RightMark` keeps +/// every streamed row regardless of whether it matched, so a buggy `PiecewiseMergeJoinExec` +/// that always marks `true` (or always `false`, or emits `lhs`'s row count under a different +/// column) would still pass a row-count-only check, since `LEFT_ROWS == RIGHT_ROWS` and every +/// regime here returns exactly `RIGHT_ROWS` rows either way. Reading the actual booleans back +/// is what catches that. +fn mark_counts( + plan: Arc, + ctx: &SessionContext, + rt: &Runtime, +) -> (usize, usize, usize) { + rt.block_on(async { + let batches = collect(plan, ctx.task_ctx()).await.unwrap(); + let mut true_count = 0; + let mut false_count = 0; + let mut null_count = 0; + for batch in &batches { + let mark_idx = batch + .schema() + .index_of("mark") + .expect("mark join output is expected to carry a `mark` column"); + let mark = batch + .column(mark_idx) + .as_any() + .downcast_ref::() + .expect("`mark` column is expected to be boolean"); + for i in 0..mark.len() { + if mark.is_null(i) { + null_count += 1; + } else if mark.value(i) { + true_count += 1; + } else { + false_count += 1; + } + } + } + (true_count, false_count, null_count) + }) +} + /// One point in a sweep: which join, over what data, and what it must emit. struct Case<'a> { kind: Kind, @@ -542,8 +588,8 @@ fn key_lt_key_filter(s: &SchemaRef) -> JoinFilter { JoinFilter::new(expr, column_indices, Arc::new(intermediate_schema)) } -/// `RightMark` has no SQL surface (see the module doc's "Mark joins" section), so both arms -/// are hand-built here instead of planned from SQL text: `PiecewiseMergeJoinExec` and +/// `RightMark` has no direct SQL surface (see the module doc's "Mark joins" section), so both +/// arms are hand-built here instead of planned from SQL text: `PiecewiseMergeJoinExec` and /// `NestedLoopJoinExec`, over the same data, the same `lhs.key < rhs.key` relation, and the /// same join type -- with no risk of comparing an operator against itself, since which /// operator each arm uses is fixed by construction rather than read back off a plan. @@ -562,7 +608,12 @@ fn bench_pwmj_right_mark_hand_built(c: &mut Criterion) { // `try_new` directly and skip the whole group rather than panic, so this benchmark stays // runnable (as a no-op) against a build that has not merged `RightMark` support yet, and // starts measuring on its own once that support lands. - if let Err(err) = PiecewiseMergeJoinExec::try_new( + // + // Only the specific `NotImplemented` this build's `try_new` returns for an unsupported + // join type is treated as "not landed yet"; any other error (a schema mistake here, or a + // real regression in `try_new`) is a bug in this benchmark or in PWMJ and must panic + // rather than be swallowed as a skip. + match PiecewiseMergeJoinExec::try_new( Arc::new(EmptyExec::new(Arc::clone(&s))), Arc::new(EmptyExec::new(Arc::clone(&s))), ( @@ -573,17 +624,24 @@ fn bench_pwmj_right_mark_hand_built(c: &mut Criterion) { JoinType::RightMark, 1, ) { - println!( - "note: pwmj_vs_nlj_right_mark_hand_built skipped -- this build's \ - PiecewiseMergeJoinExec does not support RightMark yet: {err}" - ); - return; + Ok(_) => {} + Err(err @ DataFusionError::NotImplemented(_)) => { + println!( + "note: pwmj_vs_nlj_right_mark_hand_built skipped -- this build's \ + PiecewiseMergeJoinExec does not support RightMark yet: {err}" + ); + return; + } + Err(err) => panic!( + "pwmj_vs_nlj_right_mark_hand_built: unexpected error probing RightMark \ + support, not the NotImplemented this benchmark skips on: {err}" + ), } let mut group = c.benchmark_group("pwmj_vs_nlj_right_mark_hand_built"); group.sample_size(10); - for (regime, right_offset, _fraction) in REGIMES { + for (regime, right_offset, fraction) in REGIMES { let lhs_batches = build_batches(LEFT_ROWS, 0, &s); let rhs_batches = build_batches(RIGHT_ROWS, right_offset, &s); @@ -648,11 +706,12 @@ fn bench_pwmj_right_mark_hand_built(c: &mut Criterion) { } }; - // `RightMark` keeps every streamed row, matched or not, so both arms must return - // exactly `RIGHT_ROWS` regardless of the regime -- unlike `RightSemi`/`RightAnti`, - // where the regime changes the row count. The regime still matters to what is timed - // below: it changes how much of the comparison work each arm actually does (`mark` - // true vs false), even though the row count it returns cannot show that. + // `RightMark` keeps every streamed row, matched or not, so a row count alone cannot + // tell a correct implementation from a broken one -- both arms return exactly + // `RIGHT_ROWS` regardless of the regime, unlike `RightSemi`/`RightAnti`, where the + // regime changes the row count. What the regime actually changes here is which rows + // get marked `true` vs `false`, so that is what gets checked: the `mark` column's own + // value distribution, from both arms, against what the regime claims. let pwmj_rows = run(pwmj_plan(), &ctx, &rt); let nlj_rows = run(nlj_plan(), &ctx, &rt); assert_eq!( @@ -664,6 +723,33 @@ fn bench_pwmj_right_mark_hand_built(c: &mut Criterion) { "right_mark_{regime}: RightMark must keep every streamed row" ); + let (pwmj_true, pwmj_false, pwmj_null) = mark_counts(pwmj_plan(), &ctx, &rt); + let (nlj_true, nlj_false, nlj_null) = mark_counts(nlj_plan(), &ctx, &rt); + assert_eq!( + (pwmj_true, pwmj_false, pwmj_null), + (nlj_true, nlj_false, nlj_null), + "right_mark_{regime}: pwmj and nlj disagree on mark values \ + (true/false/null): pwmj={pwmj_true}/{pwmj_false}/{pwmj_null}, \ + nlj={nlj_true}/{nlj_false}/{nlj_null}" + ); + assert_eq!( + pwmj_null, 0, + "right_mark_{regime}: RightMark's mark column must never be null" + ); + match fraction.expected_matched_rows(RIGHT_ROWS) { + Some(expected_true) => assert_eq!( + pwmj_true, expected_true, + "right_mark_{regime}: expected {expected_true} marked true, got {pwmj_true} \ + (a uniformly true/false mark column would be caught here)" + ), + None => assert!( + pwmj_true > 0 && pwmj_true < RIGHT_ROWS, + "right_mark_{regime}: {pwmj_true} rows marked true is all-or-nothing, \ + adjust the data shape (a uniformly true/false mark column would be caught \ + here)" + ), + } + group.bench_function( BenchmarkId::new("pwmj", format!("{regime}_{RIGHT_ROWS}")), |b| {