Implement basic common subexpression eliminate optimization

datafusion/src/execution/context.rs

@@ -23,8 +23,9 @@ use crate::{
     },
     logical_plan::{PlanType, ToStringifiedPlan},
     optimizer::{
-        aggregate_statistics::AggregateStatistics, eliminate_limit::EliminateLimit,
-        hash_build_probe_order::HashBuildProbeOrder,
+        aggregate_statistics::AggregateStatistics,
+        common_subexpr_eliminate::CommonSubexprEliminate,
+        eliminate_limit::EliminateLimit, hash_build_probe_order::HashBuildProbeOrder,
     },
     physical_optimizer::optimizer::PhysicalOptimizerRule,
 };
@@ -684,6 +685,7 @@ impl Default for ExecutionConfig {
             batch_size: 8192,
             optimizers: vec![
                 Arc::new(ConstantFolding::new()),
+                Arc::new(CommonSubexprEliminate::new()),
                 Arc::new(EliminateLimit::new()),
                 Arc::new(AggregateStatistics::new()),
                 Arc::new(ProjectionPushDown::new()),

datafusion/src/logical_plan/expr.rs

@@ -803,9 +803,11 @@ impl Expr {
     where
         R: ExprRewriter,
     {
-        if !rewriter.pre_visit(&self)? {
-            return Ok(self);
-        };
+        match rewriter.pre_visit(&self)? {
+            RewriteRecursion::Mutate => return rewriter.mutate(self),
+            RewriteRecursion::Stop => return Ok(self),
+            RewriteRecursion::Continue => {}
+        }
 
         // recurse into all sub expressions(and cover all expression types)
         let expr = match self {
@@ -914,7 +916,7 @@ impl Expr {
                 negated,
             } => Expr::InList {
                 expr: rewrite_boxed(expr, rewriter)?,
-                list,
+                list: rewrite_vec(list, rewriter)?,
                 negated,
             },
             Expr::Wildcard => Expr::Wildcard,
@@ -981,15 +983,25 @@ pub trait ExpressionVisitor: Sized {
     }
 }
 
+/// Controls how the [ExprRewriter] recursion should proceed.
+pub enum RewriteRecursion {
+    /// Continue rewrite / visit this expression.
+    Continue,
+    /// Call [mutate()] immediately and return.
+    Mutate,
+    /// Do not rewrite / visit the children of this expression.
+    Stop,
+}
+
 /// Trait for potentially recursively rewriting an [`Expr`] expression
 /// tree. When passed to `Expr::rewrite`, `ExpressionVisitor::mutate` is
 /// invoked recursively on all nodes of an expression tree. See the
 /// comments on `Expr::rewrite` for details on its use
 pub trait ExprRewriter: Sized {
     /// Invoked before any children of `expr` are rewritten /
-    /// visited. Default implementation returns `Ok(true)`
-    fn pre_visit(&mut self, _expr: &Expr) -> Result<bool> {
-        Ok(true)
+    /// visited. Default implementation returns `Ok(RewriteRecursion::Continue)`
+    fn pre_visit(&mut self, _expr: &Expr) -> Result<RewriteRecursion> {
+        Ok(RewriteRecursion::Continue)
     }
 
     /// Invoked after all children of `expr` have been mutated and
@@ -1861,9 +1873,9 @@ mod tests {
             Ok(expr)
         }
 
-        fn pre_visit(&mut self, expr: &Expr) -> Result<bool> {
+        fn pre_visit(&mut self, expr: &Expr) -> Result<RewriteRecursion> {
             self.v.push(format!("Previsited {:?}", expr));
-            Ok(true)
+            Ok(RewriteRecursion::Continue)
         }
     }
 

datafusion/src/logical_plan/mod.rs

@@ -45,6 +45,7 @@ pub use expr::{
     sha224, sha256, sha384, sha512, signum, sin, split_part, sqrt, starts_with, strpos,
     substr, sum, tan, to_hex, translate, trim, trunc, unnormalize_col, unnormalize_cols,
     upper, when, Column, Expr, ExprRewriter, ExpressionVisitor, Literal, Recursion,
+    RewriteRecursion,
 };
 pub use extension::UserDefinedLogicalNode;
 pub use operators::Operator;