fix: Enhance join handling in optimizer rules by introducing new join method and improving cross join elimination logic

Author: longshan.lu

qurious/src/execution/session.rs

@@ -10,7 +10,7 @@ use crate::common::table_relation::TableRelation;
 use crate::datasource::memory::MemoryTable;
 use crate::error::Error;
 use crate::functions::{all_builtin_functions, UserDefinedFunction};
-use crate::{internal_err, utils};
+use crate::internal_err;
 use crate::logical::plan::{
     CreateMemoryTable, DdlStatement, DmlOperator, DmlStatement, DropTable, Filter, LogicalPlan,
 };
@@ -98,7 +98,6 @@ impl ExecuteSession {
             LogicalPlan::Dml(stmt) => self.execute_dml(stmt),
             plan => {
                 let plan = self.optimizer.optimize(plan)?;
-                println!("{}", utils::format(&plan, 0));
                 self.planner.create_physical_plan(&plan)?.execute()
             }
         }

qurious/src/logical/builder.rs

@@ -85,9 +85,14 @@ impl LogicalPlanBuilder {
         })
     }
 
-    pub fn join_on(self, right: LogicalPlan, join_type: JoinType, filter: Option<LogicalExpr>) -> Result<Self> {
+    pub fn join(
+        self,
+        right: LogicalPlan,
+        join_type: JoinType,
+        on: Vec<(LogicalExpr, LogicalExpr)>,
+        filter: Option<LogicalExpr>,
+    ) -> Result<Self> {
         let schema = build_join_schema(join_type, &self.plan.table_schema(), &right.table_schema())?;
-        let on = vec![];
 
         if join_type != JoinType::Inner && filter.is_none() && on.is_empty() {
             return Err(Error::InternalError(format!("join condition should not be empty")));

qurious/src/logical/plan/mod.rs

@@ -294,11 +294,18 @@ impl TransformNode for LogicalPlan {
         })
     }
 
-    fn apply_children<'n, F>(&'n self, _f: F) -> Result<TreeNodeRecursion>
+    fn apply_children<'n, F>(&'n self, mut f: F) -> Result<TreeNodeRecursion>
     where
         F: FnMut(&'n LogicalPlan) -> Result<TreeNodeRecursion>,
     {
-        todo!()
+        for child in self.children().into_iter().flatten() {
+            match f(child)? {
+                TreeNodeRecursion::Continue => {}
+                TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop),
+            }
+        }
+
+        Ok(TreeNodeRecursion::Continue)
     }
 }
 

qurious/src/optimizer/rule/eliminate_cross_join.rs

@@ -1,17 +1,24 @@
-use indexmap::IndexSet;
+use std::{collections::HashMap, sync::Arc};
+
+use indexmap::{Equivalent, IndexSet};
 
 use crate::{
-    common::{join_type::JoinType, table_schema::TableSchemaRef, transformed::Transformed},
+    common::{
+        join_type::JoinType,
+        table_schema::TableSchemaRef,
+        transformed::{TransformNode, Transformed, TransformedResult, TreeNodeRecursion},
+    },
     datatypes::operator::Operator,
     error::Result,
     logical::{
         expr::{BinaryExpr, LogicalExpr},
-        plan::{Filter, Join, LogicalPlan},
+        plan::{Filter, LogicalPlan},
         LogicalPlanBuilder,
     },
     optimizer::rule::OptimizerRule,
 };
 
+/// Eliminate cross joins by rewriting them to inner joins when possible.
 pub struct EliminateCrossJoin;
 
 impl OptimizerRule for EliminateCrossJoin {
@@ -20,40 +27,74 @@ impl OptimizerRule for EliminateCrossJoin {
     }
 
     fn rewrite(&self, plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
-        match plan {
-            LogicalPlan::Filter(filter) if matches!(filter.input.as_ref(), LogicalPlan::CrossJoin(_)) => {
-                let LogicalPlan::CrossJoin(cross_join) = *filter.input else {
-                    return Ok(Transformed::no(LogicalPlan::Filter(filter)));
-                };
+        let LogicalPlan::Filter(Filter { input, expr: predicate }) = plan else {
+            return Ok(Transformed::no(plan));
+        };
 
-                let left_schema = cross_join.left.table_schema();
-                let right_schema = cross_join.right.table_schema();
+        let mut all_corss_joins = vec![];
 
-                let (join_keys, remaining_predicate) = extract_join_pairs(&filter.expr, &left_schema, &right_schema);
+        // collect all cross joins and filter predicates in order
+        input.apply(|plan| {
+            if let LogicalPlan::CrossJoin(cross_join) = plan {
+                all_corss_joins.push(cross_join);
+            }
+            Ok(TreeNodeRecursion::Continue)
+        })?;
 
-                if join_keys.is_empty() {
-                    Ok(Transformed::no(LogicalPlan::Filter(Filter {
-                        input: Box::new(LogicalPlan::CrossJoin(cross_join)),
-                        expr: filter.expr,
-                    })))
-                } else {
-                    let inner_join_plan = LogicalPlan::Join(Join {
-                        left: cross_join.left,
-                        right: cross_join.right,
-                        join_type: JoinType::Inner,
-                        on: join_keys.into_iter().map(|(l, r)| (l.clone(), r.clone())).collect(),
-                        filter: None,
-                        schema: cross_join.schema,
-                    });
-
-                    if let Some(predicate) = remaining_predicate {
-                        LogicalPlanBuilder::filter(inner_join_plan, predicate).map(Transformed::yes)
-                    } else {
-                        Ok(Transformed::yes(inner_join_plan))
-                    }
-                }
+        if all_corss_joins.is_empty() {
+            return Ok(Transformed::no(LogicalPlan::Filter(Filter { input, expr: predicate })));
+        }
+
+        let mut all_join_keys = IndexSet::new();
+        let mut replaced_cross_joins = HashMap::new();
+        let len = all_corss_joins.len();
+        // iteratively rewrite cross joins to inner joins from bottom to top
+        for (index, cross_join) in all_corss_joins.into_iter().rev().enumerate() {
+            let left_schema = cross_join.left.table_schema();
+            let right_schema = cross_join.right.table_schema();
+
+            let join_keys = extract_join_pairs(&predicate, &left_schema, &right_schema);
+
+            all_join_keys.extend(join_keys.clone());
+
+            if !join_keys.is_empty() {
+                let inner_join_plan = LogicalPlanBuilder::from(Arc::unwrap_or_clone(cross_join.left.clone()))
+                    .join(
+                        Arc::unwrap_or_clone(cross_join.right.clone()),
+                        JoinType::Inner,
+                        join_keys.into_iter().collect(),
+                        None,
+                    )?
+                    .build();
+
+                // this index should be from the top to the bottom
+                replaced_cross_joins.insert((len - 1) - index, inner_join_plan);
             }
-            _ => Ok(Transformed::no(plan)),
+        }
+
+        if replaced_cross_joins.is_empty() {
+            return Ok(Transformed::no(LogicalPlan::Filter(Filter { input, expr: predicate })));
+        }
+
+        // combine all predicates and replaced cross joins
+        let mut index = 0;
+        let new_input = input
+            .transform(|plan| {
+                let result = if let Some(replaced_join) = replaced_cross_joins.remove(&index) {
+                    Ok(Transformed::yes(replaced_join))
+                } else {
+                    Ok(Transformed::no(plan))
+                };
+                index += 1;
+                result
+            })
+            .data()?;
+
+        // remove all join keys from original predicates
+        if let Some(predicate) = remove_join_keys(predicate, &all_join_keys) {
+            LogicalPlanBuilder::filter(new_input, predicate).map(Transformed::yes)
+        } else {
+            Ok(Transformed::yes(new_input))
         }
     }
 }
@@ -62,7 +103,7 @@ fn extract_join_pairs<'a>(
     expr: &'a LogicalExpr,
     left_schema: &TableSchemaRef,
     right_schema: &TableSchemaRef,
-) -> (IndexSet<(&'a LogicalExpr, &'a LogicalExpr)>, Option<LogicalExpr>) {
+) -> IndexSet<(LogicalExpr, LogicalExpr)> {
     let mut join_keys = IndexSet::new();
 
     match expr {
@@ -75,69 +116,92 @@ fn extract_join_pairs<'a>(
             let right_col = right.try_as_column();
 
             if let (Some(left_col), Some(right_col)) = (left_col, right_col) {
-                if (left_schema.has_column(left_col) || right_schema.has_column(left_col))
-                    && (left_schema.has_column(right_col) || right_schema.has_column(right_col))
-                {
-                    join_keys.insert((left.as_ref(), right.as_ref()));
-
-                    return (join_keys, None);
+                if left_schema.has_column(left_col) && right_schema.has_column(right_col) {
+                    join_keys.insert((left.as_ref().clone(), right.as_ref().clone()));
+                } else if right_schema.has_column(left_col) && left_schema.has_column(right_col) {
+                    join_keys.insert((right.as_ref().clone(), left.as_ref().clone()));
                 }
             }
-
-            (join_keys, Some(expr.clone()))
         }
         LogicalExpr::BinaryExpr(BinaryExpr {
             left,
             op: Operator::And,
             right,
         }) => {
-            let (left_join_keys, left_predicate) = extract_join_pairs(left, left_schema, right_schema);
-            let (right_join_keys, right_predicate) = extract_join_pairs(right, left_schema, right_schema);
+            let left_join_keys = extract_join_pairs(left, left_schema, right_schema);
+            let right_join_keys = extract_join_pairs(right, left_schema, right_schema);
 
             join_keys.extend(left_join_keys);
             join_keys.extend(right_join_keys);
-
-            let predicate = match (left_predicate, right_predicate) {
-                (Some(left_predicate), Some(right_predicate)) => Some(LogicalExpr::BinaryExpr(BinaryExpr {
-                    left: Box::new(left_predicate),
-                    op: Operator::And,
-                    right: Box::new(right_predicate),
-                })),
-                (l, r) => l.or(r),
-            };
-
-            (join_keys, predicate)
         }
         LogicalExpr::BinaryExpr(BinaryExpr {
             left,
             op: Operator::Or,
             right,
         }) => {
-            let (left_join_keys, left_predicate) = extract_join_pairs(left, left_schema, right_schema);
-            let (right_join_keys, right_predicate) = extract_join_pairs(right, left_schema, right_schema);
+            let left_join_keys = extract_join_pairs(left, left_schema, right_schema);
+            let right_join_keys = extract_join_pairs(right, left_schema, right_schema);
 
             for (l, r) in left_join_keys {
-                if right_join_keys.contains(&(l, r)) || right_join_keys.contains(&(r, l)) {
+                if right_join_keys.contains(&ExprPair::new(&l, &r)) || right_join_keys.contains(&ExprPair::new(&r, &l))
+                {
                     join_keys.insert((l, r));
                 }
             }
+        }
+        _ => {}
+    }
 
-            let predicate = match (left_predicate, right_predicate) {
-                (Some(l), Some(r)) => Some(LogicalExpr::BinaryExpr(BinaryExpr {
-                    left: Box::new(l),
-                    op: Operator::Or,
-                    right: Box::new(r),
-                })),
-                (l, r) => l.or(r),
-            };
+    join_keys
+}
 
-            (join_keys, predicate)
+fn remove_join_keys(expr: LogicalExpr, join_keys: &IndexSet<(LogicalExpr, LogicalExpr)>) -> Option<LogicalExpr> {
+    match expr {
+        LogicalExpr::BinaryExpr(BinaryExpr {
+            left,
+            op: Operator::Eq,
+            right,
+        }) if join_keys.contains(&ExprPair::new(left.as_ref(), right.as_ref()))
+            || join_keys.contains(&ExprPair::new(right.as_ref(), left.as_ref())) =>
+        {
+            None
+        }
+        LogicalExpr::BinaryExpr(BinaryExpr { left, op, right }) if op == Operator::And => {
+            let l = remove_join_keys(*left, join_keys);
+            let r = remove_join_keys(*right, join_keys);
+            match (l, r) {
+                (Some(ll), Some(rr)) => Some(LogicalExpr::BinaryExpr(BinaryExpr::new(ll, op, rr))),
+                (Some(ll), _) => Some(ll),
+                (_, Some(rr)) => Some(rr),
+                _ => None,
+            }
         }
+        LogicalExpr::BinaryExpr(BinaryExpr { left, op, right }) if op == Operator::Or => {
+            let l = remove_join_keys(*left, join_keys);
+            let r = remove_join_keys(*right, join_keys);
+            match (l, r) {
+                (Some(ll), Some(rr)) => Some(LogicalExpr::BinaryExpr(BinaryExpr::new(ll, op, rr))),
+                _ => None,
+            }
+        }
+        _ => Some(expr),
+    }
+}
+
+#[derive(Debug, Eq, PartialEq, Hash)]
+struct ExprPair<'a>(&'a LogicalExpr, &'a LogicalExpr);
 
-        _ => (join_keys, Some(expr.clone())),
+impl<'a> ExprPair<'a> {
+    fn new(left: &'a LogicalExpr, right: &'a LogicalExpr) -> Self {
+        Self(left, right)
     }
 }
 
+impl Equivalent<(LogicalExpr, LogicalExpr)> for ExprPair<'_> {
+    fn equivalent(&self, other: &(LogicalExpr, LogicalExpr)) -> bool {
+        self.0 == &other.0 && self.1 == &other.1
+    }
+}
 #[cfg(test)]
 mod tests {
     use crate::{optimizer::rule::eliminate_cross_join::EliminateCrossJoin, test_utils::assert_after_optimizer};
@@ -231,4 +295,46 @@ mod tests {
             ],
         );
     }
+
+    #[test]
+    fn test_tpch_03() {
+        assert_after_optimizer(
+            "    select
+        l_orderkey,
+        sum(l_extendedprice * (1 - l_discount)) as revenue,
+        o_orderdate,
+        o_shippriority
+    from
+        customer,
+        orders,
+        lineitem
+    where
+            c_mktsegment = 'BUILDING'
+      and c_custkey = o_custkey
+      and l_orderkey = o_orderkey
+      and o_orderdate < date '1995-03-15'
+      and l_shipdate > date '1995-03-15'
+    group by
+        l_orderkey,
+        o_orderdate,
+        o_shippriority
+    order by
+        revenue desc,
+        o_orderdate
+     limit 10;",
+            vec![Box::new(EliminateCrossJoin)],
+            vec![
+                "Limit: fetch=10, skip=0",
+                "  Sort: revenue DESC, orders.o_orderdate ASC",
+                "    Projection: (lineitem.l_orderkey, SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS revenue, orders.o_orderdate, orders.o_shippriority)",
+                "      Aggregate: group_expr=[lineitem.l_orderkey,orders.o_orderdate,orders.o_shippriority], aggregat_expr=[SUM(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]",
+                "        Filter: customer.c_mktsegment = Utf8('BUILDING') AND orders.o_orderdate < CAST(Utf8('1995-03-15') AS Date32) AND lineitem.l_shipdate > CAST(Utf8('1995-03-15') AS Date32)",
+                "          Inner Join: On: (orders.o_orderkey, lineitem.l_orderkey)",
+                "            Inner Join: On: (customer.c_custkey, orders.o_custkey)",
+                "              TableScan: customer",
+                "              TableScan: orders",
+                "            TableScan: lineitem",
+            ],
+        );
+    }
 }

qurious/src/optimizer/rule/scalar_subquery_to_join.rs

@@ -84,7 +84,7 @@ impl OptimizerRule for ScalarSubqueryToJoin {
                                 })?;
 
                             LogicalPlanBuilder::from(cur_input)
-                                .join_on(new_subquery_plan, JoinType::Left, join_filter)?
+                                .join(new_subquery_plan, JoinType::Left, vec![], join_filter)?
                                 .build()
                         }
                     };