From 27b1559aba8552c8c669ad62a716c78290504ab3 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 7 Aug 2026 14:14:16 +0530 Subject: [PATCH 1/4] bench: pwmj left semi/anti join --- 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 insertions(+) create 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 0f72b74840d01..79dc8179f715f 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/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 6ae681d8e121f6c330f831a803bf4e1eb42c4829 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 7 Aug 2026 15:25:57 +0530 Subject: [PATCH 2/4] remove slt test --- datafusion/sqllogictest/test_files/pwmj.slt | 89 --------------------- 1 file changed, 89 deletions(-) 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 d52523e2778758326e74bb14733eca25ff28eb71 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 7 Aug 2026 22:05:47 +0530 Subject: [PATCH 3/4] Modify bench --- datafusion/core/Cargo.toml | 4 + datafusion/core/benches/pwmj_semi_anti_sql.rs | 232 ++++++++++++++++++ datafusion/physical-plan/Cargo.toml | 5 - .../benches/piecewise_merge_join_semi_anti.rs | 228 ----------------- 4 files changed, 236 insertions(+), 233 deletions(-) create mode 100644 datafusion/core/benches/pwmj_semi_anti_sql.rs delete mode 100644 datafusion/physical-plan/benches/piecewise_merge_join_semi_anti.rs diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 16db28a6d7600..474d8f272c3c6 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -205,6 +205,10 @@ name = "distinct_query_sql" harness = false name = "push_down_filter" +[[bench]] +harness = false +name = "pwmj_semi_anti_sql" + [[bench]] harness = false name = "sort_limit_query_sql" diff --git a/datafusion/core/benches/pwmj_semi_anti_sql.rs b/datafusion/core/benches/pwmj_semi_anti_sql.rs new file mode 100644 index 0000000000000..73017738f15f4 --- /dev/null +++ b/datafusion/core/benches/pwmj_semi_anti_sql.rs @@ -0,0 +1,232 @@ +// 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 for existence subqueries (`EXISTS` / `NOT EXISTS`) whose +//! correlation is a range predicate (`lhs.key < rhs.key`). The query is planned end to +//! end from SQL, with `datafusion.optimizer.enable_piecewise_merge_join` toggled to pick +//! the operator under test: +//! +//! - **on**: the subquery is decorrelated to a `LeftSemi` / `LeftAnti` join and planned as +//! `PiecewiseMergeJoin`, together with the `SortExec` the planner inserts on the +//! buffered side. The sort is included because it is a real cost of that plan. +//! - **off**: the same join falls back to `NestedLoopJoinExec`, which is O(n*m). +//! +//! Both arms compute the same result, so the comparison measures the win from routing an +//! inequality-correlated `EXISTS` / `NOT EXISTS` to PWMJ instead of the nested-loop join. +//! +//! Each arm asserts up front that the operator it means to measure is actually in the +//! physical plan. Without that check a planning change (or running this against a build +//! where PWMJ does not accept existence joins) would silently compare +//! `NestedLoopJoinExec` against itself and report a meaningless ~1.0x. +//! +//! ## Axes +//! - **join type**: `EXISTS` (LeftSemi) and `NOT EXISTS` (LeftAnti). +//! - **match regime**: the fraction of left rows that have at least one matching right +//! row, 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 +//! ends up only partially marked). Semi output size grows with that fraction; Anti +//! output size shrinks. + +use std::sync::Arc; + +use arrow::array::{Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion::datasource::MemTable; +use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use tokio::runtime::Runtime; + +const LEFT_ROWS: usize = 20_000; +const RIGHT_ROWS: usize = 20_000; +const KEY_SPAN: i32 = 10_000; + +/// 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 the batches for a single-partition table 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_batches( + num_rows: usize, + key_offset: i32, + schema: &SchemaRef, +) -> Vec { + let keys: Vec = (0..num_rows) + .map(|i| { + key_offset + + (i as i32) + .wrapping_mul(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 scan. + 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; + } + batches +} + +/// Register `lhs` and `rhs` in a context with PWMJ planning on or off. +fn create_context(right_offset: i32, pwmj: bool, schema: &SchemaRef) -> SessionContext { + let config = SessionConfig::new() + // Pinned so results are comparable across machines, and so the comparison + // isolates the join operator rather than how much repartitioning surrounds it. + .with_target_partitions(1) + .set_bool("datafusion.optimizer.enable_piecewise_merge_join", pwmj); + let ctx = SessionContext::new_with_config(config); + + for (name, key_offset, num_rows) in + [("lhs", 0, LEFT_ROWS), ("rhs", right_offset, RIGHT_ROWS)] + { + let table = MemTable::try_new( + Arc::clone(schema), + vec![build_batches(num_rows, key_offset, schema)], + ) + .unwrap(); + ctx.register_table(name, Arc::new(table)).unwrap(); + } + ctx +} + +/// `EXISTS` / `NOT EXISTS` over the range correlation `lhs.key < rhs.key`. +fn query(exists: bool) -> String { + let negation = if exists { "" } else { "NOT " }; + format!( + "SELECT lhs.key, lhs.payload FROM lhs \ + WHERE {negation}EXISTS (SELECT 1 FROM rhs WHERE lhs.key < rhs.key)" + ) +} + +fn physical_plan( + ctx: &SessionContext, + rt: &Runtime, + sql: &str, +) -> Arc { + rt.block_on(async { + ctx.sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap() + }) +} + +/// Fail loudly if the plan does not contain every expected fragment, so an arm can never +/// silently measure an operator other than the one it is named after. +fn assert_plan_contains(plan: &Arc, expected: &[&str], label: &str) { + let displayed = displayable(plan.as_ref()).indent(false).to_string(); + for fragment in expected { + assert!( + displayed.contains(fragment), + "{label}: expected `{fragment}` in the physical plan, got:\n{displayed}" + ); + } +} + +fn run(plan: Arc, ctx: &SessionContext, rt: &Runtime) -> usize { + rt.block_on(async { + let batches = collect(plan, ctx.task_ctx()).await.unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + }) +} + +fn bench_pwmj_semi_anti_sql(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let s = schema(); + + // An existence join needs only *one* match per left row, so all that matters is where + // `max(rhs.key)` falls inside the left range: a left row survives `EXISTS` iff + // `lhs.key < max(rhs.key)`. Shifting the right range up therefore saturates at + // all-match; only shifting it *down* leaves part of the left side unmatched. + let regimes: [(&str, i32); 3] = [ + // Right keys entirely above the left range: every left row matches. + ("all_match", KEY_SPAN), + // Right keys entirely below the left range: no left row matches, so the scan + // walks the whole buffered side without marking anything. + ("no_match", -KEY_SPAN), + // Ranges half-overlap: ~50% of left rows match, so only a suffix of the buffered + // side gets marked and the scan depth varies per streamed row. + ("half_match", -KEY_SPAN / 2), + ]; + + let mut group = c.benchmark_group("pwmj_vs_nlj_semi_anti_sql"); + // Nested-loop is O(n*m); keep sample counts modest so the suite finishes. + group.sample_size(10); + + for (regime, right_offset) in regimes { + // Only the Semi/Anti half of the join type is pinned, not the side: with PWMJ + // disabled the planner swaps the nested-loop inputs, so the same `EXISTS` plans as + // `RightSemi` there and `LeftSemi` under PWMJ. + for (exists, join_type) in [(true, "Semi"), (false, "Anti")] { + let sql = query(exists); + let label = if exists { "semi" } else { "anti" }; + + for (arm, pwmj, operator) in [ + ("pwmj", true, "PiecewiseMergeJoin"), + ("nlj", false, "NestedLoopJoinExec"), + ] { + let ctx = create_context(right_offset, pwmj, &s); + let name = format!("{arm}_{label}_{regime}"); + assert_plan_contains( + &physical_plan(&ctx, &rt, &sql), + &[operator, join_type], + &name, + ); + + // Plan afresh in the untimed setup rather than reusing one plan: the + // buffered side of `PiecewiseMergeJoinExec` (its visited-indices bitmap + // and final-pass partition counter) is cached in a `OnceAsync` on the + // exec, so a second `collect` over the same instance would not repeat + // the work. + group.bench_function(BenchmarkId::new(name, RIGHT_ROWS), |b| { + b.iter_batched( + || physical_plan(&ctx, &rt, &sql), + |plan| run(plan, &ctx, &rt), + BatchSize::SmallInput, + ) + }); + } + } + } + + group.finish(); +} + +criterion_group!(benches, bench_pwmj_semi_anti_sql); +criterion_main!(benches); diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 79dc8179f715f..0f72b74840d01 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); From cdd33d988b13735388df0270db795165d94f5ea6 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Mon, 10 Aug 2026 10:43:57 +0530 Subject: [PATCH 4/4] bench: let the pwmj_enabled arm accept either operator --- datafusion/core/benches/pwmj_semi_anti_sql.rs | 99 +++++++++++++++---- 1 file changed, 78 insertions(+), 21 deletions(-) diff --git a/datafusion/core/benches/pwmj_semi_anti_sql.rs b/datafusion/core/benches/pwmj_semi_anti_sql.rs index 73017738f15f4..f84b8bfced210 100644 --- a/datafusion/core/benches/pwmj_semi_anti_sql.rs +++ b/datafusion/core/benches/pwmj_semi_anti_sql.rs @@ -20,18 +20,26 @@ //! end from SQL, with `datafusion.optimizer.enable_piecewise_merge_join` toggled to pick //! the operator under test: //! -//! - **on**: the subquery is decorrelated to a `LeftSemi` / `LeftAnti` join and planned as -//! `PiecewiseMergeJoin`, together with the `SortExec` the planner inserts on the -//! buffered side. The sort is included because it is a real cost of that plan. -//! - **off**: the same join falls back to `NestedLoopJoinExec`, which is O(n*m). +//! - **`pwmj_enabled`**: on a build that routes existence joins to PWMJ, the subquery is +//! decorrelated to a `LeftSemi` / `LeftAnti` join and planned as `PiecewiseMergeJoin`, +//! together with the `SortExec` the planner inserts on the buffered side. The sort is +//! included because it is a real cost of that plan. +//! - **`nlj`**: the same join is planned as `NestedLoopJoinExec`, which is O(n*m). //! //! Both arms compute the same result, so the comparison measures the win from routing an //! inequality-correlated `EXISTS` / `NOT EXISTS` to PWMJ instead of the nested-loop join. //! -//! Each arm asserts up front that the operator it means to measure is actually in the -//! physical plan. Without that check a planning change (or running this against a build -//! where PWMJ does not accept existence joins) would silently compare -//! `NestedLoopJoinExec` against itself and report a meaningless ~1.0x. +//! The arms are named after the config they set, not the operator they get, because the +//! enabled arm's operator is build-dependent: while the planner still excludes semi/anti +//! join types from PWMJ, setting the flag changes nothing and both arms plan +//! `NestedLoopJoinExec`. So the enabled arm accepts either operator instead of aborting the +//! run. Benchmark ids stay the same either way, which is what makes such a run useful as a +//! baseline: a later build that does route these joins to PWMJ compares straight against it. +//! +//! What every arm does pin down is *which* operator it planned, printed before the +//! timings, and it says so outright when both arms landed on the same one. Without that a +//! planning change would quietly compare `NestedLoopJoinExec` against itself and the +//! resulting ~1.0x would read as "PWMJ is no faster". //! //! ## Axes //! - **join type**: `EXISTS` (LeftSemi) and `NOT EXISTS` (LeftAnti). @@ -55,6 +63,14 @@ const LEFT_ROWS: usize = 20_000; const RIGHT_ROWS: usize = 20_000; const KEY_SPAN: i32 = 10_000; +/// Operators the `pwmj_enabled` arm is allowed to plan. Two of them, because the planner +/// only hands existence joins to PWMJ once PWMJ accepts semi/anti join types; before that +/// the flag is a no-op and the nested-loop join stays. +const PWMJ_OR_NLJ: &[&str] = &["PiecewiseMergeJoin", "NestedLoopJoinExec"]; + +/// With the flag off, nothing but the nested-loop join can plan this query. +const NLJ_ONLY: &[&str] = &["NestedLoopJoinExec"]; + /// Two-column schema: (`key`, `payload`). fn schema() -> SchemaRef { Arc::new(Schema::new(vec![ @@ -147,16 +163,39 @@ fn physical_plan( }) } -/// Fail loudly if the plan does not contain every expected fragment, so an arm can never -/// silently measure an operator other than the one it is named after. -fn assert_plan_contains(plan: &Arc, expected: &[&str], label: &str) { +/// Check that the plan contains every fragment in `required` and exactly one of the +/// operators in `one_of`, returning the one it found. +/// +/// An arm named after a config flag says nothing about what ran, so each one reports the +/// operator it actually planned. `one_of` has two entries for the enabled arm because +/// whether existence joins reach PWMJ depends on the build; a single-entry list pins the +/// arm down completely. +fn assert_plan_operator<'a>( + plan: &Arc, + required: &[&str], + one_of: &[&'a str], + label: &str, +) -> &'a str { let displayed = displayable(plan.as_ref()).indent(false).to_string(); - for fragment in expected { + for fragment in required { assert!( displayed.contains(fragment), "{label}: expected `{fragment}` in the physical plan, got:\n{displayed}" ); } + + let found: Vec<&str> = one_of + .iter() + .copied() + .filter(|operator| displayed.contains(operator)) + .collect(); + assert_eq!( + found.len(), + 1, + "{label}: expected exactly one of {one_of:?} in the physical plan, found \ + {found:?} in:\n{displayed}" + ); + found[0] } fn run(plan: Arc, ctx: &SessionContext, rt: &Runtime) -> usize { @@ -197,27 +236,45 @@ fn bench_pwmj_semi_anti_sql(c: &mut Criterion) { let sql = query(exists); let label = if exists { "semi" } else { "anti" }; - for (arm, pwmj, operator) in [ - ("pwmj", true, "PiecewiseMergeJoin"), - ("nlj", false, "NestedLoopJoinExec"), - ] { + // Plan and check both arms before timing either, so the operators they picked + // are on screen ahead of the numbers those operators qualify. + let arms: Vec<(String, SessionContext, &str)> = [ + ("pwmj_enabled", true, PWMJ_OR_NLJ), + ("nlj", false, NLJ_ONLY), + ] + .into_iter() + .map(|(arm, pwmj, operators)| { let ctx = create_context(right_offset, pwmj, &s); let name = format!("{arm}_{label}_{regime}"); - assert_plan_contains( + let planned = assert_plan_operator( &physical_plan(&ctx, &rt, &sql), - &[operator, join_type], + &[join_type], + operators, &name, ); + println!("{name}: planned {planned}"); + (name, ctx, planned) + }) + .collect(); + + if arms.iter().all(|(_, _, planned)| *planned == arms[0].2) { + println!( + "note: {label}_{regime}: both arms planned {}, so the ratio between \ + them measures nothing about PWMJ", + arms[0].2 + ); + } + for (name, ctx, _) in &arms { // Plan afresh in the untimed setup rather than reusing one plan: the // buffered side of `PiecewiseMergeJoinExec` (its visited-indices bitmap // and final-pass partition counter) is cached in a `OnceAsync` on the // exec, so a second `collect` over the same instance would not repeat // the work. - group.bench_function(BenchmarkId::new(name, RIGHT_ROWS), |b| { + group.bench_function(BenchmarkId::new(name.as_str(), RIGHT_ROWS), |b| { b.iter_batched( - || physical_plan(&ctx, &rt, &sql), - |plan| run(plan, &ctx, &rt), + || physical_plan(ctx, &rt, &sql), + |plan| run(plan, ctx, &rt), BatchSize::SmallInput, ) });