perf: avoid re-flattening join chains in EliminateCrossJoin

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

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::<Vec<_>>(),
+    );
+    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);

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<bool> {
+        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::<Vec<_>>();
+            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(())
+    }
 }

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

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<Arc<dyn OptimizerRule + Send + Sync>> =
+        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);

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::<Result<Vec<_>>>()?;
+
         // 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 + (<subquery>) = 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")?)