From 770924033da52e233391aba91fa371f67214ae55 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 27 Oct 2026 09:37:22 +0800 Subject: [PATCH 1/2] perf: avoid re-flattening join chains in EliminateCrossJoin After flattening a chain of inner/cross joins and reassembling it, `EliminateCrossJoin` called `rewrite_children` on the reassembled tree. Its left child is the reassembled join chain minus one input, which the rule flattened and rebuilt again, and so on at every nesting level, making the rule quadratic in the number of joined inputs (with each rebuild also searching the accumulated left schema). Instead, rewrite each flattened input once before reassembling the joins. `rewrite` preserves the schema of its input, which is all that `find_inner_join` looks at, so the result is unchanged. Uncorrelated subqueries in the join keys, previously visited by `rewrite_children`, are now rewritten explicitly. Also check the right schema first in `find_valid_equijoin_key_pair`: the left schema accumulates all previous inputs and is much more expensive to search, while most candidate keys do not reference the right input. Add an `eliminate_cross_join` benchmark. --- datafusion/expr/src/utils.rs | 12 +- datafusion/optimizer/Cargo.toml | 4 + .../optimizer/benches/eliminate_cross_join.rs | 114 ++++++++++++++++++ .../optimizer/src/eliminate_cross_join.rs | 55 ++++++++- 4 files changed, 178 insertions(+), 7 deletions(-) create mode 100644 datafusion/optimizer/benches/eliminate_cross_join.rs diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 4517c52c72270..d41771feabbbe 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -1191,12 +1191,16 @@ pub fn find_valid_equijoin_key_pair( return Ok(None); } - if check_all_columns_from_schema(&left_using_columns, left_schema)? - && check_all_columns_from_schema(&right_using_columns, right_schema)? + // Check `right_schema` first: when building a left-deep join tree, + // `left_schema` is the accumulated schema of all previous inputs, which is + // much more expensive to search, and most candidate keys do not reference + // the right input at all. + if check_all_columns_from_schema(&right_using_columns, right_schema)? + && check_all_columns_from_schema(&left_using_columns, left_schema)? { return Ok(Some((left_key.clone(), right_key.clone()))); - } else if check_all_columns_from_schema(&right_using_columns, left_schema)? - && check_all_columns_from_schema(&left_using_columns, right_schema)? + } else if check_all_columns_from_schema(&left_using_columns, right_schema)? + && check_all_columns_from_schema(&right_using_columns, left_schema)? { return Ok(Some((right_key.clone(), left_key.clone()))); } diff --git a/datafusion/optimizer/Cargo.toml b/datafusion/optimizer/Cargo.toml index 31a73062e119e..e74bffee26848 100644 --- a/datafusion/optimizer/Cargo.toml +++ b/datafusion/optimizer/Cargo.toml @@ -91,3 +91,7 @@ harness = false [[bench]] name = "simplify_function_over_case" harness = false + +[[bench]] +name = "eliminate_cross_join" +harness = false diff --git a/datafusion/optimizer/benches/eliminate_cross_join.rs b/datafusion/optimizer/benches/eliminate_cross_join.rs new file mode 100644 index 0000000000000..4f86a4b36589e --- /dev/null +++ b/datafusion/optimizer/benches/eliminate_cross_join.rs @@ -0,0 +1,114 @@ +// 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. + +//! Microbenchmarks for the [`EliminateCrossJoin`] optimizer rule. +//! +//! Two scenarios are covered, each over `n` three-column tables: +//! +//! 1. **comma_join** – `FROM t0, t1, ..., tn WHERE t0.a = t1.a AND t1.a = +//! t2.a AND ...`: a `Filter` over a chain of cross joins that the rule +//! rewrites into a chain of inner joins. +//! +//! 2. **inner_join** – `t0 JOIN t1 ON t0.a = t1.a JOIN t2 ON t1.a = t2.a ...`: +//! a chain of inner joins that already carry their keys. The rule still +//! flattens and reassembles it. + +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_expr::logical_plan::{JoinType, table_scan}; +use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder, and, col}; +use datafusion_optimizer::OptimizerContext; +use datafusion_optimizer::eliminate_cross_join::EliminateCrossJoin; +use datafusion_optimizer::{Optimizer, OptimizerRule}; +use std::hint::black_box; +use std::sync::Arc; + +/// Build a three-column table scan for `name`. +fn scan(name: &str) -> LogicalPlan { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ]); + table_scan(Some(name), &schema, None) + .unwrap() + .build() + .unwrap() +} + +/// `t{i - 1}.a = t{i}.a` +fn chain_key(i: usize) -> Expr { + col(format!("t{}.a", i - 1)).eq(col(format!("t{i}.a"))) +} + +/// `SELECT * FROM t0, t1, ... WHERE t0.a = t1.a AND t1.a = t2.a AND ...` +fn build_comma_join_plan(n: usize) -> LogicalPlan { + let mut builder = LogicalPlanBuilder::from(scan("t0")); + for i in 1..n { + builder = builder.cross_join(scan(&format!("t{i}"))).unwrap(); + } + let predicate = (1..n).map(chain_key).reduce(and).unwrap(); + builder.filter(predicate).unwrap().build().unwrap() +} + +/// `SELECT * FROM t0 JOIN t1 ON t0.a = t1.a JOIN t2 ON t1.a = t2.a ...` +fn build_inner_join_plan(n: usize) -> LogicalPlan { + let mut builder = LogicalPlanBuilder::from(scan("t0")); + for i in 1..n { + builder = builder + .join( + scan(&format!("t{i}")), + JoinType::Inner, + (vec![format!("t{}.a", i - 1)], vec![format!("t{i}.a")]), + None, + ) + .unwrap(); + } + builder.build().unwrap() +} + +fn run_optimizer(plan: &LogicalPlan, ctx: &OptimizerContext) -> LogicalPlan { + let rules: Vec> = + vec![Arc::new(EliminateCrossJoin::new())]; + Optimizer::with_rules(rules) + .optimize(plan.clone(), ctx, |_, _| {}) + .unwrap() +} + +fn bench_eliminate_cross_join(c: &mut Criterion) { + let ctx = OptimizerContext::new().with_max_passes(1); + + for (name, build) in [ + ( + "comma_join", + build_comma_join_plan as fn(usize) -> LogicalPlan, + ), + ("inner_join", build_inner_join_plan), + ] { + let mut group = c.benchmark_group(format!("eliminate_cross_join/{name}")); + for n in [8, 32, 64, 128] { + let plan = build(n); + group.bench_with_input(BenchmarkId::from_parameter(n), &plan, |b, p| { + b.iter(|| black_box(run_optimizer(p, &ctx))); + }); + } + group.finish(); + } +} + +criterion_group!(benches, bench_eliminate_cross_join); +criterion_main!(benches); diff --git a/datafusion/optimizer/src/eliminate_cross_join.rs b/datafusion/optimizer/src/eliminate_cross_join.rs index 6464b6a6fd679..1200fa922df5d 100644 --- a/datafusion/optimizer/src/eliminate_cross_join.rs +++ b/datafusion/optimizer/src/eliminate_cross_join.rs @@ -155,6 +155,16 @@ impl OptimizerRule for EliminateCrossJoin { } }; + // Rewrite each flattened input once, before reassembling the joins. + // `rewrite` preserves its input's schema, which is all that + // `find_inner_join` looks at, so this is equivalent to rewriting the + // children of the reassembled join tree, but avoids flattening and + // rebuilding the rest of the join chain again at every nesting level. + let mut all_inputs = all_inputs + .into_iter() + .map(|input| self.rewrite(input, config).map(|t| t.data)) + .collect::>>()?; + // Join keys are handled locally: let mut all_join_keys = JoinKeySet::new(); let mut left = all_inputs.remove(0); @@ -165,11 +175,12 @@ impl OptimizerRule for EliminateCrossJoin { &possible_join_keys, &mut all_join_keys, null_equality, - )?; + )? + // Join keys may contain uncorrelated subqueries to rewrite + .map_uncorrelated_subqueries(|input| self.rewrite(input, config))? + .data; } - left = rewrite_children(self, left, config)?.data; - if &plan_schema != left.schema() { left = LogicalPlan::Projection(Projection::new_from_schema( Arc::new(left), @@ -460,6 +471,7 @@ mod tests { Operator::{And, Or}, binary_expr, col, lit, logical_plan::builder::LogicalPlanBuilder, + scalar_subquery, }; use insta::assert_snapshot; @@ -1355,6 +1367,43 @@ mod tests { ) } + #[test] + fn eliminate_cross_join_in_join_key_subquery() -> Result<()> { + // An uncorrelated subquery containing a cross join that can be eliminated + let subquery = LogicalPlanBuilder::from(test_table_scan_with_name("s1")?) + .cross_join(test_table_scan_with_name("s2")?)? + .filter(col("s1.a").eq(col("s2.a")))? + .project(vec![col("s1.b")])? + .build()?; + + // The subquery ends up in the join key of the innermost join + let plan = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) + .cross_join(test_table_scan_with_name("t2")?)? + .cross_join(test_table_scan_with_name("t3")?)? + .filter(binary_expr( + (col("t1.a") + scalar_subquery(Arc::new(subquery))).eq(col("t2.a")), + And, + col("t2.b").eq(col("t3.b")), + ))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @ r" + Inner Join: t2.b = t3.b [a:UInt32, b:UInt32, c:UInt32, a:UInt32, b:UInt32, c:UInt32, a:UInt32, b:UInt32, c:UInt32] + Inner Join: t1.a + () = t2.a [a:UInt32, b:UInt32, c:UInt32, a:UInt32, b:UInt32, c:UInt32] + Subquery: [b:UInt32] + Projection: s1.b [b:UInt32] + Inner Join: s1.a = s2.a [a:UInt32, b:UInt32, c:UInt32, a:UInt32, b:UInt32, c:UInt32] + TableScan: s1 [a:UInt32, b:UInt32, c:UInt32] + TableScan: s2 [a:UInt32, b:UInt32, c:UInt32] + TableScan: t1 [a:UInt32, b:UInt32, c:UInt32] + TableScan: t2 [a:UInt32, b:UInt32, c:UInt32] + TableScan: t3 [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + #[test] fn preserve_outer_join_boundary() -> Result<()> { let plan = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) From 148bf92d17eb37925422c54539631fb432b0968b Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 27 Oct 2026 23:16:32 +0800 Subject: [PATCH 2/2] perf: check the narrower schema first in find_valid_equijoin_key_pair The previous commit made `find_valid_equijoin_key_pair` always check the right schema first, which assumes the left-deep join trees built by `EliminateCrossJoin`. Other callers, such as `ExtractEquijoinPredicate` and `LogicalPlanBuilder` joins, may see a narrow left input and a wide right one (e.g. `t JOIN (t1 JOIN t2 ...)`), which that order made slower. Check the narrower schema first instead. Looking up a column scans the schema, so a key that does not belong to the join inputs is then usually rejected without scanning the wider one. `EliminateCrossJoin` keeps its speedup, as it joins the accumulated left input with narrower inputs. Add a `find_valid_equijoin_key_pair` benchmark covering both a wide left and a wide right input, and a unit test covering both orders. --- datafusion/expr/Cargo.toml | 4 + .../benches/find_valid_equijoin_key_pair.rs | 78 +++++++++++++++++++ datafusion/expr/src/utils.rs | 67 +++++++++++++--- 3 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 datafusion/expr/benches/find_valid_equijoin_key_pair.rs diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index 45b66ee89885e..0a0d9cb96c858 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -75,6 +75,10 @@ ctor = { workspace = true } env_logger = { workspace = true } insta = { workspace = true } +[[bench]] +name = "find_valid_equijoin_key_pair" +harness = false + [[bench]] name = "normalize_columns" harness = false diff --git a/datafusion/expr/benches/find_valid_equijoin_key_pair.rs b/datafusion/expr/benches/find_valid_equijoin_key_pair.rs new file mode 100644 index 0000000000000..2b6d30becf337 --- /dev/null +++ b/datafusion/expr/benches/find_valid_equijoin_key_pair.rs @@ -0,0 +1,78 @@ +// 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. + +//! Microbenchmarks for [`find_valid_equijoin_key_pair`] with a candidate key +//! that does not belong to the join, when one join input is much wider than +//! the other: +//! +//! 1. **wide_left** – e.g. the accumulated left input of a left-deep join +//! tree, as built by `EliminateCrossJoin`. +//! +//! 2. **wide_right** – e.g. `t JOIN (t1 JOIN t2 ...)`, as seen by +//! `ExtractEquijoinPredicate`. + +use std::hint::black_box; + +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::DFSchema; +use datafusion_expr::col; +use datafusion_expr::utils::find_valid_equijoin_key_pair; + +/// Schema of table `name` with columns `c0`, `c1`, ..., `c{width - 1}`. +fn schema(name: &str, width: usize) -> DFSchema { + let schema = Schema::new( + (0..width) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, false)) + .collect::>(), + ); + DFSchema::try_from_qualified_schema(name, &schema).unwrap() +} + +fn bench_find_valid_equijoin_key_pair(c: &mut Criterion) { + let mut group = c.benchmark_group("find_valid_equijoin_key_pair"); + let narrow = schema("narrow", 3); + for width in [16, 128, 1024] { + let wide = schema("wide", width); + // `wide.c{width / 2} = other.c0`, where `other` is in neither input + let left_key = col(format!("wide.c{}", width / 2)); + let right_key = col("other.c0"); + + for (name, left_schema, right_schema) in [ + ("wide_left", &wide, &narrow), + ("wide_right", &narrow, &wide), + ] { + group.bench_function(BenchmarkId::new(name, width), |b| { + b.iter(|| { + black_box( + find_valid_equijoin_key_pair( + &left_key, + &right_key, + left_schema, + right_schema, + ) + .unwrap(), + ) + }) + }); + } + } + group.finish(); +} + +criterion_group!(benches, bench_find_valid_equijoin_key_pair); +criterion_main!(benches); diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index d41771feabbbe..a6b68626c1cd6 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -1191,17 +1191,29 @@ pub fn find_valid_equijoin_key_pair( return Ok(None); } - // Check `right_schema` first: when building a left-deep join tree, - // `left_schema` is the accumulated schema of all previous inputs, which is - // much more expensive to search, and most candidate keys do not reference - // the right input at all. - if check_all_columns_from_schema(&right_using_columns, right_schema)? - && check_all_columns_from_schema(&left_using_columns, left_schema)? - { + // Returns whether all `left_columns` are from `left_schema` and all + // `right_columns` are from `right_schema`. Looking up a column scans the + // schema, so check the narrower schema first: a key that does not belong + // to the join inputs is then usually rejected without scanning the wider + // one, such as the accumulated left input of a left-deep join tree. + let left_is_narrower = left_schema.fields().len() <= right_schema.fields().len(); + let check = |left_columns: &HashSet<&Column>, + right_columns: &HashSet<&Column>| + -> Result { + let left = (left_columns, left_schema); + let right = (right_columns, right_schema); + let (first, second) = if left_is_narrower { + (left, right) + } else { + (right, left) + }; + Ok(check_all_columns_from_schema(first.0, first.1)? + && check_all_columns_from_schema(second.0, second.1)?) + }; + + if check(&left_using_columns, &right_using_columns)? { return Ok(Some((left_key.clone(), right_key.clone()))); - } else if check_all_columns_from_schema(&left_using_columns, right_schema)? - && check_all_columns_from_schema(&right_using_columns, left_schema)? - { + } else if check(&right_using_columns, &left_using_columns)? { return Ok(Some((right_key.clone(), left_key.clone()))); } @@ -2405,4 +2417,39 @@ mod tests { Ok(()) } + + #[test] + fn find_valid_equijoin_key_pair_inputs_of_different_widths() -> Result<()> { + let schema = |name: &str, width: usize| { + let fields = (0..width) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, false)) + .collect::>(); + DFSchema::try_from_qualified_schema(name, &Schema::new(fields)) + }; + let narrow = schema("n", 1)?; + let wide = schema("w", 3)?; + let other = col("o.c0"); + + // The narrower input is checked first, so cover both orders + for (left_schema, right_schema, left, right) in [ + (&narrow, &wide, col("n.c0"), col("w.c2")), + (&wide, &narrow, col("w.c2"), col("n.c0")), + ] { + let find = |left_key: &Expr, right_key: &Expr| { + find_valid_equijoin_key_pair( + left_key, + right_key, + left_schema, + right_schema, + ) + }; + let expected = Some((left.clone(), right.clone())); + assert_eq!(find(&left, &right)?, expected); + assert_eq!(find(&right, &left)?, expected); + assert_eq!(find(&left, &other)?, None); + assert_eq!(find(&other, &right)?, None); + assert_eq!(find(&(left.clone() + right.clone()), &right)?, None); + } + Ok(()) + } }