Feat: simplify join cond rule and accepting all kinds of inner join conditions including non equal ones (#115)

AveryQi115 · web-flow · commit 1da620f752cd · 2024-03-20T23:11:48.000-04:00
# Major Changes - SimplifyJoinCondRule, just like SimplifyFilter in pr https://github.com/cmu-db/optd/pull/108 - Accepting all kinds of conditions for inner join from datafusion - before this pr, optd only supports inner join on equal conditions and inner join on True/False - after this pr, optd now supports queries like `select x inner join y on x.a=y.a and x.a<5`. These are good testcases for push filter into join and filter push down. P.S.: the core fix for merge group in this pr is based on pr https://github.com/cmu-db/optd/pull/114 --------- Signed-off-by: AveryQi115 <averyqi115@gmail.com>
diff --git a/optd-datafusion-bridge/src/into_optd.rs b/optd-datafusion-bridge/src/into_optd.rs
@@ -4,7 +4,6 @@ use datafusion::{
     logical_expr::{self, logical_plan, LogicalPlan, Operator},
     scalar::ScalarValue,
 };
-use datafusion_expr::Expr as DFExpr;
 use optd_core::rel_node::RelNode;
 use optd_datafusion_repr::plan_nodes::{
     BetweenExpr, BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, Expr, ExprList,
@@ -314,35 +313,25 @@ impl OptdPlanContext<'_> {
             let expr = BinOpExpr::new(left, right, op).into_expr();
             log_ops.push(expr);
         }
+        if node.filter.is_some() {
+            let filter =
+                self.conv_into_optd_expr(node.filter.as_ref().unwrap(), node.schema.as_ref())?;
+            log_ops.push(filter);
+        }
 
         if log_ops.is_empty() {
-            // optd currently only supports
-            // 1. normal equal condition join
-            //    select * from a join b on a.id = b.id
-            // 2. join on false/true
-            //    select * from a join b on false/true
-            // 3. join on other literals or other filters are not supported
-            //  instead of converting them to a join on true, we bail out
-
-            match node.filter {
-                Some(DFExpr::Literal(ScalarValue::Boolean(Some(val)))) => Ok(LogicalJoin::new(
-                    left,
-                    right,
-                    ConstantExpr::bool(val).into_expr(),
-                    join_type,
-                )),
-                None => Ok(LogicalJoin::new(
-                    left,
-                    right,
-                    ConstantExpr::bool(true).into_expr(),
-                    join_type,
-                )),
-                _ => bail!("unsupported join filter: {:?}", node.filter),
-            }
+            Ok(LogicalJoin::new(
+                left,
+                right,
+                ConstantExpr::bool(true).into_expr(),
+                join_type,
+            ))
         } else if log_ops.len() == 1 {
             Ok(LogicalJoin::new(left, right, log_ops.remove(0), join_type))
         } else {
             let expr_list = ExprList::new(log_ops);
+            // the expr from filter is already flattened in conv_into_optd_expr
+            let expr_list = flatten_nested_logical(LogOpType::And, expr_list);
             Ok(LogicalJoin::new(
                 left,
                 right,
diff --git a/optd-datafusion-repr/src/lib.rs b/optd-datafusion-repr/src/lib.rs
@@ -18,7 +18,7 @@ use properties::{
 use rules::{
     EliminateDuplicatedAggExprRule, EliminateDuplicatedSortExprRule, EliminateFilterRule,
     EliminateJoinRule, EliminateLimitRule, HashJoinRule, JoinAssocRule, JoinCommuteRule,
-    PhysicalConversionRule, ProjectionPullUpJoin, SimplifyFilterRule,
+    PhysicalConversionRule, ProjectionPullUpJoin, SimplifyFilterRule, SimplifyJoinCondRule,
 };
 
 pub use optd_core::rel_node::Value;
@@ -57,6 +57,7 @@ impl DatafusionOptimizer {
         let rules = PhysicalConversionRule::all_conversions();
         let mut rule_wrappers = vec![
             RuleWrapper::new_heuristic(Arc::new(SimplifyFilterRule::new())),
+            RuleWrapper::new_heuristic(Arc::new(SimplifyJoinCondRule::new())),
             RuleWrapper::new_heuristic(Arc::new(EliminateFilterRule::new())),
             RuleWrapper::new_heuristic(Arc::new(EliminateJoinRule::new())),
             RuleWrapper::new_heuristic(Arc::new(EliminateLimitRule::new())),
@@ -66,8 +67,8 @@ impl DatafusionOptimizer {
         for rule in rules {
             rule_wrappers.push(RuleWrapper::new_cascades(rule));
         }
-        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(HashJoinRule::new())));
-        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(JoinCommuteRule::new())));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(HashJoinRule::new()))); // 17
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(JoinCommuteRule::new()))); // 18
         rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(JoinAssocRule::new())));
         rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
             ProjectionPullUpJoin::new(),
diff --git a/optd-datafusion-repr/src/rules.rs b/optd-datafusion-repr/src/rules.rs
@@ -11,7 +11,7 @@ pub use eliminate_duplicated_expr::{
     EliminateDuplicatedAggExprRule, EliminateDuplicatedSortExprRule,
 };
 pub use eliminate_limit::EliminateLimitRule;
-pub use filter::{EliminateFilterRule, SimplifyFilterRule};
+pub use filter::{EliminateFilterRule, SimplifyFilterRule, SimplifyJoinCondRule};
 pub use joins::{
     EliminateJoinRule, HashJoinRule, JoinAssocRule, JoinCommuteRule, ProjectionPullUpJoin,
 };
diff --git a/optd-datafusion-repr/src/rules/filter.rs b/optd-datafusion-repr/src/rules/filter.rs
@@ -3,8 +3,8 @@ use std::sync::Arc;
 
 use super::macros::define_rule;
 use crate::plan_nodes::{
-    ConstantExpr, ConstantType, Expr, ExprList, LogOpExpr, LogOpType, LogicalEmptyRelation,
-    OptRelNode, OptRelNodeTyp,
+    ConstantExpr, ConstantType, Expr, ExprList, JoinType, LogOpExpr, LogOpType,
+    LogicalEmptyRelation, LogicalJoin, OptRelNode, OptRelNodeTyp, PlanNode,
 };
 use crate::properties::schema::SchemaPropertyBuilder;
 use crate::OptRelNodeRef;
@@ -120,6 +120,38 @@ fn apply_simplify_filter(
     }
 }
 
+// Same as SimplifyFilterRule, but for innerJoin conditions
+define_rule!(
+    SimplifyJoinCondRule,
+    apply_simplify_join_cond,
+    (Join(JoinType::Inner), left, right, [cond])
+);
+
+fn apply_simplify_join_cond(
+    _optimizer: &impl Optimizer<OptRelNodeTyp>,
+    SimplifyJoinCondRulePicks { left, right, cond }: SimplifyJoinCondRulePicks,
+) -> Vec<RelNode<OptRelNodeTyp>> {
+    match cond.typ {
+        OptRelNodeTyp::LogOp(_) => {
+            let mut changed = false;
+            let new_log_expr = simplify_log_expr(Arc::new(cond), &mut changed);
+            if changed {
+                let join_node = LogicalJoin::new(
+                    PlanNode::from_group(left.into()),
+                    PlanNode::from_group(right.into()),
+                    Expr::from_rel_node(new_log_expr).unwrap(),
+                    JoinType::Inner,
+                );
+                return vec![join_node.into_rel_node().as_ref().clone()];
+            }
+            vec![]
+        }
+        _ => {
+            vec![]
+        }
+    }
+}
+
 define_rule!(
     EliminateFilterRule,
     apply_eliminate_filter,
diff --git a/optd-sqlplannertest/tests/filter.planner.sql b/optd-sqlplannertest/tests/filter.planner.sql
@@ -183,3 +183,140 @@ PhysicalProjection { exprs: [ #0, #1, #2, #3 ] }
 2 2 2 202
 */
 
+-- Test SimplifyJoinCondRule and EliminateJoinRule (false filter to empty relation)
+select * from t1 inner join t2 on t1v1 = t2v1 and false;
+
+/*
+LogicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── LogicalJoin
+    ├── join_type: Inner
+    ├── cond:And
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #2
+    │   └── false
+    ├── LogicalScan { table: t1 }
+    └── LogicalScan { table: t2 }
+PhysicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── PhysicalEmptyRelation { produce_one_row: false }
+*/
+
+-- Test SimplifyJoinCondRule (skip true filter for and)
+select * from t1 inner join t2 on t1v1 = t2v1 and t1v1 = t2v3 and true;
+
+/*
+LogicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── LogicalJoin
+    ├── join_type: Inner
+    ├── cond:And
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #2
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #3
+    │   └── true
+    ├── LogicalScan { table: t1 }
+    └── LogicalScan { table: t2 }
+PhysicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── PhysicalHashJoin { join_type: Inner, left_keys: [ #0, #0 ], right_keys: [ #0, #1 ] }
+    ├── PhysicalScan { table: t1 }
+    └── PhysicalScan { table: t2 }
+*/
+
+-- Test SimplifyJoinCondRule (skip true filter for and)
+select * from t1 inner join t2 on t1v1 = t2v1 or t1v1 = t2v3 and true;
+
+/*
+LogicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── LogicalJoin
+    ├── join_type: Inner
+    ├── cond:Or
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #2
+    │   └── And
+    │       ├── Eq
+    │       │   ├── #0
+    │       │   └── #3
+    │       └── true
+    ├── LogicalScan { table: t1 }
+    └── LogicalScan { table: t2 }
+PhysicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── PhysicalNestedLoopJoin
+    ├── join_type: Inner
+    ├── cond:Or
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #2
+    │   └── Eq
+    │       ├── #0
+    │       └── #3
+    ├── PhysicalScan { table: t1 }
+    └── PhysicalScan { table: t2 }
+0 0 0 200
+1 1 1 201
+2 2 2 202
+*/
+
+-- Test SimplifyJoinCondRule, EliminateFilter (repace true filter for or)
+select * from t1 inner join t2 on t1v1 = t2v1 or t1v1 = t2v3 or true;
+
+/*
+LogicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── LogicalJoin
+    ├── join_type: Inner
+    ├── cond:Or
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #2
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #3
+    │   └── true
+    ├── LogicalScan { table: t1 }
+    └── LogicalScan { table: t2 }
+PhysicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── PhysicalNestedLoopJoin { join_type: Cross, cond: true }
+    ├── PhysicalScan { table: t1 }
+    └── PhysicalScan { table: t2 }
+0 0 0 200
+0 0 1 201
+0 0 2 202
+1 1 0 200
+1 1 1 201
+1 1 2 202
+2 2 0 200
+2 2 1 201
+2 2 2 202
+*/
+
+-- Test SimplifyJoinCondRule (remove duplicates)
+select * from t1 inner join t2 on t1v1 = t2v1 or t1v1 = t2v1 and t1v1 = t2v1;
+
+/*
+LogicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── LogicalJoin
+    ├── join_type: Inner
+    ├── cond:Or
+    │   ├── Eq
+    │   │   ├── #0
+    │   │   └── #2
+    │   └── And
+    │       ├── Eq
+    │       │   ├── #0
+    │       │   └── #2
+    │       └── Eq
+    │           ├── #0
+    │           └── #2
+    ├── LogicalScan { table: t1 }
+    └── LogicalScan { table: t2 }
+PhysicalProjection { exprs: [ #0, #1, #2, #3 ] }
+└── PhysicalHashJoin { join_type: Inner, left_keys: [ #0 ], right_keys: [ #0 ] }
+    ├── PhysicalScan { table: t1 }
+    └── PhysicalScan { table: t2 }
+0 0 0 200
+1 1 1 201
+2 2 2 202
+*/
+
diff --git a/optd-sqlplannertest/tests/filter.yml b/optd-sqlplannertest/tests/filter.yml
@@ -44,6 +44,36 @@
 - sql: |
     select * from t1, t2 where t1v1 = t2v1 or t1v1 = t2v1 and t1v1 = t2v1;
   desc: Test SimplifyFilterRule (remove duplicates)
+  tasks:
+    - explain:logical_optd,physical_optd
+    - execute
+- sql: |
+    select * from t1 inner join t2 on t1v1 = t2v1 and false;
+  desc: Test SimplifyJoinCondRule and EliminateJoinRule (false filter to empty relation)
+  tasks:
+    - explain:logical_optd,physical_optd
+    - execute
+- sql: |
+    select * from t1 inner join t2 on t1v1 = t2v1 and t1v1 = t2v3 and true;
+  desc: Test SimplifyJoinCondRule (skip true filter for and)
+  tasks:
+    - explain:logical_optd,physical_optd
+    - execute
+- sql: |
+    select * from t1 inner join t2 on t1v1 = t2v1 or t1v1 = t2v3 and true;
+  desc: Test SimplifyJoinCondRule (skip true filter for and)
+  tasks:
+    - explain:logical_optd,physical_optd
+    - execute
+- sql: |
+    select * from t1 inner join t2 on t1v1 = t2v1 or t1v1 = t2v3 or true;
+  desc: Test SimplifyJoinCondRule, EliminateFilter (repace true filter for or)
+  tasks:
+    - explain:logical_optd,physical_optd
+    - execute
+- sql: |
+    select * from t1 inner join t2 on t1v1 = t2v1 or t1v1 = t2v1 and t1v1 = t2v1;
+  desc: Test SimplifyJoinCondRule (remove duplicates)
   tasks:
     - explain:logical_optd,physical_optd
     - execute
diff --git a/optd-sqlplannertest/tests/tpch.planner.sql b/optd-sqlplannertest/tests/tpch.planner.sql
@@ -1937,9 +1937,13 @@ LogicalProjection
     └── LogicalProjection { exprs: [ #1 ] }
         └── LogicalJoin
             ├── join_type: Inner
-            ├── cond:Eq
-            │   ├── #2
-            │   └── #4
+            ├── cond:And
+            │   ├── Eq
+            │   │   ├── #2
+            │   │   └── #4
+            │   └── Lt
+            │       ├── Cast { cast_to: Decimal128(30, 15), expr: #0 }
+            │       └── #3
             ├── LogicalProjection { exprs: [ #1, #2, #3 ] }
             │   └── LogicalJoin
             │       ├── join_type: Inner
@@ -1986,7 +1990,15 @@ PhysicalProjection
     │   └── [ #0 ]
     ├── groups: []
     └── PhysicalProjection { exprs: [ #1 ] }
-        └── PhysicalHashJoin { join_type: Inner, left_keys: [ #2 ], right_keys: [ #1 ] }
+        └── PhysicalNestedLoopJoin
+            ├── join_type: Inner
+            ├── cond:And
+            │   ├── Eq
+            │   │   ├── #2
+            │   │   └── #0
+            │   └── Lt
+            │       ├── Cast { cast_to: Decimal128(30, 15), expr: #0 }
+            │       └── #3
             ├── PhysicalProjection { exprs: [ #1, #2, #3 ] }
             │   └── PhysicalHashJoin { join_type: Inner, left_keys: [ #0 ], right_keys: [ #0 ] }
             │       ├── PhysicalProjection { exprs: [ #1, #4, #5 ] }
