Implement basic common subexpression eliminate optimization

datafusion/src/execution/context.rs

@@ -57,6 +57,7 @@ use crate::execution::dataframe_impl::DataFrameImpl;
 use crate::logical_plan::{
  FunctionRegistry, LogicalPlan, LogicalPlanBuilder, UNNAMED_TABLE,
 };
+use crate::optimizer::common_subexpr_eliminate::CommonSubexprEliminate;
 use crate::optimizer::constant_folding::ConstantFolding;
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::limit_push_down::LimitPushDown;
@@ -709,6 +710,7 @@ impl Default for ExecutionConfig {
  batch_size: 8192,
  optimizers: vec![
  Arc::new(ConstantFolding::new()),
+ Arc::new(CommonSubexprEliminate::new()),
  Arc::new(EliminateLimit::new()),
  Arc::new(ProjectionPushDown::new()),
  Arc::new(FilterPushDown::new()),
@@ -988,6 +990,7 @@ impl FunctionRegistry for ExecutionContextState {
 mod tests {
 
  use super::*;
+ use crate::logical_plan::{binary_expr, lit, Operator};
  use crate::physical_plan::functions::make_scalar_function;
  use crate::physical_plan::{collect, collect_partitioned};
  use crate::test;
@@ -1966,6 +1969,27 @@ mod tests {
  Ok(())
  }
 
+ #[tokio::test]
+ async fn aggregate_avg_add() -> Result<()> {
+ let results = execute(
+ "SELECT AVG(c1), AVG(c1) + 1, AVG(c1) + 2, 1 + AVG(c1) FROM test",
+ 4,
+ )
+ .await?;
+ assert_eq!(results.len(), 1);
+
+ let expected = vec![
+ "+--------------+----------------------------+----------------------------+----------------------------+",
+ "| AVG(test.c1) | AVG(test.c1) Plus Int64(1) | AVG(test.c1) Plus Int64(2) | Int64(1) Plus AVG(test.c1) |",
+ "+--------------+----------------------------+----------------------------+----------------------------+",
+ "| 1.5 | 2.5 | 3.5 | 2.5 |",
+ "+--------------+----------------------------+----------------------------+----------------------------+",
+ ];
+ assert_batches_sorted_eq!(expected, &results);
+
+ Ok(())
+ }
+
  #[tokio::test]
  async fn join_partitioned() -> Result<()> {
  // self join on partition id (workaround for duplicate column name)
@@ -2134,6 +2158,30 @@ mod tests {
  }
  }
 
+ #[tokio::test]
+ async fn unprojected_filter() {
+ let mut ctx = ExecutionContext::new();
+ let df = ctx
+ .read_table(test::table_with_sequence(1, 3).unwrap())
+ .unwrap();
+
+ let df = df
+ .select(vec![binary_expr(col("i"), Operator::Plus, col("i"))])
+ .unwrap()
+ .filter(col("i").gt(lit(2)))
+ .unwrap();
+ let results = df.collect().await.unwrap();
+
+ let expected = vec![
+ "+--------------------------+",
+ "| ?table?.i Plus ?table?.i |",
+ "+--------------------------+",
+ "| 6 |",
+ "+--------------------------+",
+ ];
+ assert_batches_sorted_eq!(expected, &results);
+ }
+
  #[tokio::test]
  async fn group_by_dictionary() {
  async fn run_test_case<K: ArrowDictionaryKeyType>() {

datafusion/src/logical_plan/expr.rs

@@ -804,8 +804,11 @@ impl Expr {
  where
  R: ExprRewriter,
  {
- if !rewriter.pre_visit(&self)? {
- return Ok(self);
+ let need_mutate = match rewriter.pre_visit(&self)? {
+ RewriteRecursion::Mutate => return rewriter.mutate(self),
+ RewriteRecursion::Stop => return Ok(self),
+ RewriteRecursion::Continue => true,
+ RewriteRecursion::Skip => false,
  };
 
  // recurse into all sub expressions(and cover all expression types)
@@ -915,14 +918,18 @@ impl Expr {
  negated,
  } => Expr::InList {
  expr: rewrite_boxed(expr, rewriter)?,
- list,
+ list: rewrite_vec(list, rewriter)?,
  negated,
  },
  Expr::Wildcard => Expr::Wildcard,
  };
 
  // now rewrite this expression itself
- rewriter.mutate(expr)
+ if need_mutate {
+ rewriter.mutate(expr)
+ } else {
+ Ok(expr)
+ }
  }
 }
 
@@ -990,15 +997,27 @@ 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,
+ /// Keep recursive but skip mutate on this expression
+ Skip,
+}
+
 /// 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
@@ -1721,13 +1740,17 @@ fn create_name(e: &Expr, input_schema: &DFSchema) -> Result<String> {
  } => {
  let mut name = "CASE ".to_string();
  if let Some(e) = expr {
- name += &format!("{:?} ", e);
+ let e = create_name(e, input_schema)?;
+ name += &format!("{} ", e);
  }
  for (w, t) in when_then_expr {
- name += &format!("WHEN {:?} THEN {:?} ", w, t);
+ let when = create_name(w, input_schema)?;
+ let then = create_name(t, input_schema)?;
+ name += &format!("WHEN {} THEN {} ", when, then);
  }
  if let Some(e) = else_expr {
- name += &format!("ELSE {:?} ", e);
+ let e = create_name(e, input_schema)?;
+ name += &format!("ELSE {} ", e);
  }
  name += "END";
  Ok(name)
@@ -1887,9 +1910,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,7 +45,7 @@ pub use expr::{
  right, round, rpad, rtrim, 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,
+ ExpressionVisitor, Literal, Recursion, RewriteRecursion,
 };
 pub use extension::UserDefinedLogicalNode;
 pub use operators::Operator;