feat: Eliminate Filter Rule (#54)

Author: jurplel

optd-core/src/cascades/memo.rs

@@ -109,15 +109,37 @@ impl<T: RelNodeTyp> Memo<T> {
         ExprId(id)
     }
 
-    fn merge_group(&mut self, group_a: ReducedGroupId, group_b: ReducedGroupId) -> ReducedGroupId {
+    fn merge_group_inner(
+        &mut self,
+        group_a: ReducedGroupId,
+        group_b: ReducedGroupId,
+    ) -> ReducedGroupId {
         if group_a == group_b {
             return group_a;
         }
         self.merged_groups
             .insert(group_a.as_group_id(), group_b.as_group_id());
+
+        // Copy all expressions from group a to group b
+        let group_a_exprs = self.get_all_exprs_in_group(group_a.as_group_id());
+        for expr_id in group_a_exprs {
+            let expr_node = self.expr_id_to_expr_node.get(&expr_id).unwrap();
+            self.add_expr_to_group(expr_id, group_b, expr_node.as_ref().clone());
+        }
+
+        // Remove all expressions from group a (so we don't accidentally access it)
+        self.clear_exprs_in_group(group_a);
+
         group_b
     }
 
+    pub fn merge_group(&mut self, group_a: GroupId, group_b: GroupId) -> GroupId {
+        let group_a_reduced = self.get_reduced_group_id(group_a);
+        let group_b_reduced = self.get_reduced_group_id(group_b);
+        self.merge_group_inner(group_a_reduced, group_b_reduced)
+            .as_group_id()
+    }
+
     fn get_group_id_of_expr_id(&self, expr_id: ExprId) -> GroupId {
         self.expr_id_to_group_id[&expr_id]
     }
@@ -136,9 +158,11 @@ impl<T: RelNodeTyp> Memo<T> {
         rel_node: RelNodeRef<T>,
         add_to_group_id: Option<GroupId>,
     ) -> (GroupId, ExprId) {
-        if rel_node.typ.extract_group().is_some() {
-            unreachable!();
-        }
+        let node_current_group = rel_node.typ.extract_group();
+        if let (Some(grp_a), Some(grp_b)) = (add_to_group_id, node_current_group) {
+            self.merge_group(grp_a, grp_b);
+        };
+
         let (group_id, expr_id) = self.add_new_group_expr_inner(
             rel_node,
             add_to_group_id.map(|x| self.get_reduced_group_id(x)),
@@ -198,6 +222,10 @@ impl<T: RelNodeTyp> Memo<T> {
         props
     }
 
+    fn clear_exprs_in_group(&mut self, group_id: ReducedGroupId) {
+        self.groups.remove(&group_id);
+    }
+
     fn add_expr_to_group(
         &mut self,
         expr_id: ExprId,
@@ -243,7 +271,7 @@ impl<T: RelNodeTyp> Memo<T> {
             let group_id = self.get_group_id_of_expr_id(expr_id);
             let group_id = self.get_reduced_group_id(group_id);
             if let Some(add_to_group_id) = add_to_group_id {
-                self.merge_group(add_to_group_id, group_id);
+                self.merge_group_inner(add_to_group_id, group_id);
             }
             return (group_id, expr_id);
         }

optd-core/src/cascades/optimizer.rs

@@ -290,6 +290,10 @@ impl<T: RelNodeTyp> CascadesOptimizer<T> {
         self.memo.update_group_info(group_id, group_info)
     }
 
+    pub(super) fn merge_group(&mut self, group_a: GroupId, group_b: GroupId) {
+        self.memo.merge_group(group_a, group_b);
+    }
+
     pub fn get_property_by_group<P: PropertyBuilder<T>>(
         &self,
         group_id: GroupId,

optd-core/src/cascades/tasks/apply_rule.rs

@@ -197,8 +197,10 @@ impl<T: RelNodeTyp> Task<T> for ApplyRuleTask {
             let applied = rule.apply(optimizer, expr);
             for expr in applied {
                 let RelNode { typ, .. } = &expr;
-                if typ.extract_group().is_some() {
-                    unreachable!();
+                if let Some(group_id_2) = typ.extract_group() {
+                    // If this is a group, merge the groups!
+                    optimizer.merge_group(group_id, group_id_2);
+                    continue;
                 }
                 let expr_typ = typ.clone();
                 let (_, expr_id) = optimizer.add_group_expr(expr.into(), Some(group_id));

optd-datafusion-bridge/src/into_optd.rs

@@ -110,6 +110,10 @@ impl OptdPlanContext<'_> {
                     let x = x.as_ref().unwrap();
                     Ok(ConstantExpr::decimal(*x as f64).into_expr())
                 }
+                ScalarValue::Boolean(x) => {
+                    let x = x.as_ref().unwrap();
+                    Ok(ConstantExpr::bool(*x).into_expr())
+                }
                 _ => bail!("{:?}", x),
             },
             Expr::Alias(x) => self.conv_into_optd_expr(x.expr.as_ref(), context),

optd-datafusion-repr/src/lib.rs

@@ -8,8 +8,8 @@ use optd_core::cascades::{CascadesOptimizer, GroupId, OptimizerProperties};
 use plan_nodes::{OptRelNode, OptRelNodeRef, OptRelNodeTyp, PlanNode};
 use properties::schema::{Catalog, SchemaPropertyBuilder};
 use rules::{
-    EliminateJoinRule, HashJoinRule, JoinAssocRule, JoinCommuteRule, PhysicalConversionRule,
-    ProjectionPullUpJoin,
+    EliminateFilterRule, EliminateJoinRule, HashJoinRule, JoinAssocRule, JoinCommuteRule,
+    PhysicalConversionRule, ProjectionPullUpJoin,
 };
 
 pub use adaptive::PhysicalCollector;
@@ -48,6 +48,7 @@ impl DatafusionOptimizer {
         rules.push(Arc::new(JoinAssocRule::new()));
         rules.push(Arc::new(ProjectionPullUpJoin::new()));
         rules.push(Arc::new(EliminateJoinRule::new()));
+        rules.push(Arc::new(EliminateFilterRule::new()));
 
         let cost_model = AdaptiveCostModel::new(50);
         Self {
@@ -72,6 +73,8 @@ impl DatafusionOptimizer {
         rules.insert(0, Arc::new(JoinCommuteRule::new()));
         rules.insert(1, Arc::new(JoinAssocRule::new()));
         rules.insert(2, Arc::new(ProjectionPullUpJoin::new()));
+        rules.insert(3, Arc::new(EliminateFilterRule::new()));
+
         let cost_model = AdaptiveCostModel::new(1000); // very large decay
         let runtime_statistics = cost_model.get_runtime_map();
         let optimizer = CascadesOptimizer::new(

optd-datafusion-repr/src/rules.rs

@@ -1,9 +1,11 @@
 // mod filter_join;
+mod eliminate_filter;
 mod joins;
 mod macros;
 mod physical;
 
 // pub use filter_join::FilterJoinPullUpRule;
+pub use eliminate_filter::EliminateFilterRule;
 pub use joins::{
     EliminateJoinRule, HashJoinRule, JoinAssocRule, JoinCommuteRule, ProjectionPullUpJoin,
 };

optd-datafusion-repr/src/rules/eliminate_filter.rs

@@ -0,0 +1,38 @@
+use std::collections::HashMap;
+
+use optd_core::rules::{Rule, RuleMatcher};
+use optd_core::{optimizer::Optimizer, rel_node::RelNode};
+
+use crate::plan_nodes::{ConstantType, LogicalEmptyRelation, OptRelNode, OptRelNodeTyp};
+
+use super::macros::define_rule;
+
+define_rule!(
+    EliminateFilterRule,
+    apply_eliminate_filter,
+    (Filter, child, [cond])
+);
+
+/// Transformations:
+///     - Filter node w/ false pred -> EmptyRelation
+///     - Filter node w/ true pred  -> Eliminate from the tree
+fn apply_eliminate_filter(
+    _optimizer: &impl Optimizer<OptRelNodeTyp>,
+    EliminateFilterRulePicks { child, cond }: EliminateFilterRulePicks,
+) -> Vec<RelNode<OptRelNodeTyp>> {
+    if let OptRelNodeTyp::Constant(ConstantType::Bool) = cond.typ {
+        if let Some(data) = cond.data {
+            if data.as_bool() {
+                // If the condition is true, eliminate the filter node, as it
+                // will yield everything from below it.
+                return vec![child];
+            } else {
+                // If the condition is false, replace this node with the empty relation,
+                // since it will never yield tuples.
+                let node = LogicalEmptyRelation::new(false);
+                return vec![node.into_rel_node().as_ref().clone()];
+            }
+        }
+    }
+    vec![]
+}

optd-sqlplannertest/tests/eliminate_filter.planner.sql

@@ -0,0 +1,36 @@
+-- (no id or description)
+create table t1(t1v1 int, t1v2 int);
+create table t2(t2v1 int, t2v3 int);
+insert into t1 values (0, 0), (1, 1), (2, 2);
+insert into t2 values (0, 200), (1, 201), (2, 202);
+
+/*
+3
+3
+*/
+
+-- Test EliminateFilterRule (false filter to empty relation)
+select * from t1 where false;
+
+/*
+LogicalProjection { exprs: [ #0, #1 ] }
+└── LogicalFilter { cond: false }
+    └── LogicalScan { table: t1 }
+PhysicalProjection { exprs: [ #0, #1 ] }
+└── PhysicalEmptyRelation { produce_one_row: false }
+*/
+
+-- Test EliminateFilterRule (replace true filter with child)
+select * from t1 where true;
+
+/*
+LogicalProjection { exprs: [ #0, #1 ] }
+└── LogicalFilter { cond: true }
+    └── LogicalScan { table: t1 }
+PhysicalProjection { exprs: [ #0, #1 ] }
+└── PhysicalScan { table: t1 }
+0 0
+1 1
+2 2
+*/
+

optd-sqlplannertest/tests/eliminate_filter.yml

@@ -0,0 +1,19 @@
+- sql: |
+    create table t1(t1v1 int, t1v2 int);
+    create table t2(t2v1 int, t2v3 int);
+    insert into t1 values (0, 0), (1, 1), (2, 2);
+    insert into t2 values (0, 200), (1, 201), (2, 202);
+  tasks:
+    - execute
+- sql: |
+    select * from t1 where false;
+  desc: Test EliminateFilterRule (false filter to empty relation)
+  tasks:
+    - explain:logical_optd,physical_optd
+    - execute
+- sql: |
+    select * from t1 where true;
+  desc: Test EliminateFilterRule (replace true filter with child)
+  tasks:
+    - explain:logical_optd,physical_optd
+    - execute

optd-sqlplannertest/tests/empty_relation.yml

@@ -18,4 +18,4 @@
   desc: Test whether the optimizer eliminates join to empty relation
   tasks:
     - explain:logical_optd,physical_optd
-    - execute
+    - execute