fix(df-repr): enable projection rules (#223)

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

Author: skyzh

optd-datafusion-repr/src/lib.rs

@@ -100,9 +100,9 @@ impl DatafusionOptimizer {
         for rule in rules {
             rule_wrappers.push(RuleWrapper::new_cascades(rule));
         }
-        // rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
-        //     rules::FilterProjectTransposeRule::new(),
-        // )));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
+            rules::FilterProjectTransposeRule::new(),
+        )));
         rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
             rules::FilterCrossJoinTransposeRule::new(),
         )));
@@ -130,9 +130,9 @@ impl DatafusionOptimizer {
         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(
+            rules::ProjectionPullUpJoin::new(),
+        )));
         rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
             rules::EliminateProjectRule::new(),
         )));
@@ -148,9 +148,9 @@ impl DatafusionOptimizer {
         rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
             rules::EliminateFilterRule::new(),
         )));
-        // rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
-        //     rules::ProjectFilterTransposeRule::new(),
-        // )));
+        rule_wrappers.push(RuleWrapper::new_cascades(Arc::new(
+            rules::ProjectFilterTransposeRule::new(),
+        )));
         rule_wrappers
     }
 

optd-datafusion-repr/src/plan_nodes/predicates.rs

@@ -24,6 +24,7 @@ pub use data_type_pred::DataTypePred;
 pub use extern_column_ref_pred::ExternColumnRefPred;
 pub use func_pred::{FuncPred, FuncType};
 pub use in_list_pred::InListPred;
+use itertools::Itertools;
 pub use like_pred::LikePred;
 pub use list_pred::ListPred;
 pub use log_op_pred::{LogOpPred, LogOpType};
@@ -44,6 +45,8 @@ pub trait PredExt {
     fn rewrite_column_refs(&self, rewrite_fn: impl FnMut(usize) -> Option<usize>) -> Option<Self>
     where
         Self: Sized;
+
+    fn get_column_refs(&self) -> Vec<ColumnRefPred>;
 }
 
 impl<P: DfReprPredNode> PredExt for P {
@@ -53,6 +56,15 @@ impl<P: DfReprPredNode> PredExt for P {
     ) -> Option<Self> {
         rewrite_column_refs_inner(self, &mut rewrite_fn)
     }
+
+    fn get_column_refs(&self) -> Vec<ColumnRefPred> {
+        if let Some(col_ref) = ColumnRefPred::from_pred_node(self.clone().into_pred_node()) {
+            return vec![col_ref];
+        }
+        let node = self.clone().into_pred_node();
+        let children = node.children.iter().map(|child| child.get_column_refs());
+        children.collect_vec().concat()
+    }
 }
 
 fn rewrite_column_refs_inner<P: DfReprPredNode>(

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

@@ -1,6 +1,8 @@
-// pub mod project_filter_transpose;
-// pub mod project_join_transpose;
+pub mod project_filter_transpose;
+pub mod project_join_transpose;
 pub mod project_merge;
 pub mod project_transpose_common;
 
+pub use project_filter_transpose::*;
+pub use project_join_transpose::*;
 pub use project_merge::*;