fix(df-repr): enable eliminate limit rule (#223)

Signed-off-by: Alex Chi <[email protected]>

Author: skyzh

optd-datafusion-repr/src/lib.rs

@@ -20,7 +20,6 @@ use properties::{
     column_ref::ColumnRefPropertyBuilder,
     schema::{Catalog, SchemaPropertyBuilder},
 };
-use rules::*;
 
 pub use memo_ext::{LogicalJoinOrder, MemoExt};
 pub use optimizer_ext::OptimizerExt;
@@ -96,7 +95,7 @@ impl DatafusionOptimizer {
 
     pub fn default_cascades_rules(
     ) -> Vec<Arc<RuleWrapper<DfNodeType, CascadesOptimizer<DfNodeType>>>> {
-        let rules = PhysicalConversionRule::all_conversions();
+        let rules = rules::PhysicalConversionRule::all_conversions();
         let mut rule_wrappers = vec![];
         for rule in rules {
             rule_wrappers.push(RuleWrapper::new_cascades(rule));
@@ -117,22 +116,33 @@ impl DatafusionOptimizer {
         // rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
         //     FilterAggTransposeRule::new(),
         // )));
-        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(
-            InnerCrossJoinRule::new(),
+            rules::HashJoinRule::new(),
         )));
-        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(JoinAssocRule::new())));
         rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
-            JoinAbsorbFilterRule::new(),
+            rules::JoinCommuteRule::new(),
+        )));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
+            rules::InnerCrossJoinRule::new(),
+        )));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
+            rules::JoinAssocRule::new(),
+        )));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
+            rules::JoinAbsorbFilterRule::new(),
         )));
         // rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
         //     ProjectionPullUpJoin::new(),
         // )));
         rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
-            EliminateProjectRule::new(),
+            rules::EliminateProjectRule::new(),
+        )));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
+            rules::ProjectMergeRule::new(),
+        )));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
+            rules::EliminateLimitRule::new(),
         )));
-        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(ProjectMergeRule::new())));
         // rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
         //     EliminateFilterRule::new(),
         // )));
@@ -188,7 +198,7 @@ impl DatafusionOptimizer {
 
     /// The optimizer settings for three-join demo as a perfect optimizer.
     pub fn new_alternative_physical_for_demo(catalog: Arc<dyn Catalog>) -> Self {
-        let rules = PhysicalConversionRule::all_conversions();
+        let rules = rules::PhysicalConversionRule::all_conversions();
         let mut rule_wrappers = Vec::new();
         for rule in rules {
             rule_wrappers.push(RuleWrapper::new_cascades(rule));

optd-datafusion-repr/src/rules.rs

@@ -1,5 +1,5 @@
 // mod eliminate_duplicated_expr;
-// mod eliminate_limit;
+mod eliminate_limit;
 // mod filter;
 // mod filter_pushdown;
 mod joins;
@@ -11,7 +11,7 @@ mod project_transpose;
 // pub use eliminate_duplicated_expr::{
 //     EliminateDuplicatedAggExprRule, EliminateDuplicatedSortExprRule,
 // };
-// pub use eliminate_limit::EliminateLimitRule;
+pub use eliminate_limit::*;
 // pub use filter::{EliminateFilterRule, SimplifyFilterRule, SimplifyJoinCondRule};
 // pub use filter_pushdown::{
 //     FilterAggTransposeRule, FilterCrossJoinTransposeRule, FilterInnerJoinTransposeRule,

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

@@ -1,52 +1,45 @@
-use optd_core::rules::{Rule, RuleMatcher};
-use optd_core::{nodes::PlanNode, optimizer::Optimizer};
-use std::collections::HashMap;
-use std::sync::Arc;
-
 use crate::plan_nodes::{
-    ConstantPred, ConstantType, DfNodeType, DfReprPlanNode, LogicalEmptyRelation,
+    ArcDfPlanNode, ConstantPred, ConstantType, DfNodeType, DfPredType, DfReprPlanNode,
+    DfReprPredNode, LogicalEmptyRelation, LogicalLimit,
 };
-
-use crate::properties::schema::SchemaPropertyBuilder;
+use crate::OptimizerExt;
+use optd_core::nodes::PlanNodeOrGroup;
+use optd_core::optimizer::Optimizer;
+use optd_core::rules::{Rule, RuleMatcher};
 
 use super::macros::define_rule;
 
-define_rule!(
-    EliminateLimitRule,
-    apply_eliminate_limit,
-    (Limit, [child], [skip, fetch])
-);
+define_rule!(EliminateLimitRule, apply_eliminate_limit, (Limit, child));
 
 /// Transformations:
 ///     - Limit with skip 0 and no fetch -> Eliminate from the tree
 ///     - Limit with limit 0 -> EmptyRelation
 fn apply_eliminate_limit(
     optimizer: &impl Optimizer<DfNodeType>,
-    EliminateLimitRulePicks { child, skip, fetch }: EliminateLimitRulePicks,
+    binding: ArcDfPlanNode,
 ) -> Vec<PlanNodeOrGroup<DfNodeType>> {
-    if let DfNodeType::Constant(ConstantType::UInt64) = skip.typ {
-        if let DfNodeType::Constant(ConstantType::UInt64) = fetch.typ {
-            let skip_val = ConstantPred::from_rel_node(skip.into())
-                .unwrap()
-                .value()
-                .as_u64();
-
-            let fetch_val = ConstantPred::from_rel_node(fetch.into())
+    let limit = LogicalLimit::from_plan_node(binding).unwrap();
+    let skip = limit.skip();
+    let fetch = limit.fetch();
+    let child = limit.child();
+    if let DfPredType::Constant(ConstantType::UInt64) = skip.typ {
+        if let DfPredType::Constant(ConstantType::UInt64) = fetch.typ {
+            let skip_val = ConstantPred::from_pred_node(skip).unwrap().value().as_u64();
+
+            let fetch_val = ConstantPred::from_pred_node(fetch)
                 .unwrap()
                 .value()
                 .as_u64();
 
             // Bad convention to have u64 max represent None
             let fetch_is_none = fetch_val == u64::MAX;
 
-            let schema =
-                optimizer.get_property::<SchemaPropertyBuilder>(Arc::new(child.clone()), 0);
-
+            let schema = optimizer.get_schema_of(child.clone());
             if fetch_is_none && skip_val == 0 {
                 return vec![child];
             } else if fetch_val == 0 {
                 let node = LogicalEmptyRelation::new(false, schema);
-                return vec![node.into_rel_node().as_ref().clone()];
+                return vec![node.into_plan_node().into()];
             }
         }
     }