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 4517c52c72270..a6b68626c1cd6 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -1191,13 +1191,29 @@ 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)? - { + // 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(&right_using_columns, left_schema)? - && check_all_columns_from_schema(&left_using_columns, right_schema)? - { + } else if check(&right_using_columns, &left_using_columns)? { return Ok(Some((right_key.clone(), left_key.clone()))); } @@ -2401,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(()) + } } 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")?)