Add Expression Simplification API

datafusion/src/logical_plan/expr.rs

@@ -20,11 +20,13 @@
 
 pub use super::Operator;
 use crate::error::{DataFusionError, Result};
+use crate::execution::context::ExecutionProps;
 use crate::field_util::get_indexed_field;
 use crate::logical_plan::{
     plan::Aggregate, window_frames, DFField, DFSchema, LogicalPlan,
 };
-use crate::physical_plan::functions::Volatility;
+use crate::optimizer::simplify_expressions::{ConstEvaluator, Simplifier};
+use crate::physical_plan::functions::{BuiltinScalarFunction, Volatility};
 use crate::physical_plan::{
     aggregates, expressions::binary_operator_data_type, functions, udf::ScalarUDF,
     window_functions,
@@ -971,6 +973,58 @@ impl Expr {
             Ok(expr)
         }
     }
+
+    /// Simplifies this [`Expr`]`s as much as possible, evaluating
+    /// constants and applying algebraic simplifications
+    ///
+    /// # Example:
+    /// `b > 2 AND b > 2`
+    /// can be written to
+    /// `b > 2`
+    ///
+    /// ```
+    /// use datafusion::logical_plan::*;
+    /// use datafusion::error::Result;
+    /// use datafusion::execution::context::ExecutionProps;
+    ///
+    /// /// Simple implementation that provides `Simplifier` the information it needs
+    /// #[derive(Default)]
+    /// struct Info {
+    ///   execution_props: ExecutionProps,
+    /// };
+    ///
+    /// impl SimplifyInfo for Info {
+    ///   fn is_boolean_type(&self, expr: &Expr) -> Result<bool> {
+    ///     Ok(false)
+    ///   }
+    ///   fn nullable(&self, expr: &Expr) -> Result<bool> {
+    ///     Ok(true)
+    ///   }
+    ///   fn execution_props(&self) -> &ExecutionProps {
+    ///     &self.execution_props
+    ///   }
+    /// }
+    ///
+    /// // b < 2
+    /// let b_lt_2 = col("b").gt(lit(2));
+    ///
+    /// // (b < 2) OR (b < 2)
+    /// let expr = b_lt_2.clone().or(b_lt_2.clone());
+    ///
+    /// // (b < 2) OR (b < 2) --> (b < 2)
+    /// let expr = expr.simplify(&Info::default()).unwrap();
+    /// assert_eq!(expr, b_lt_2);
+    /// ```
+    pub fn simplify<S: SimplifyInfo>(self, info: &S) -> Result<Self> {
+        let mut rewriter = Simplifier::new(info);
+        let mut const_evaluator = ConstEvaluator::new(info.execution_props());
+
+        // TODO iterate until no changes are made during rewrite
+        // (evaluating constants can enable new simplifications and
+        // simplifications can enable new constant evaluation)
+        // https://github.com/apache/arrow-datafusion/issues/1160
+        self.rewrite(&mut const_evaluator)?.rewrite(&mut rewriter)
+    }
 }
 
 impl Not for Expr {
@@ -1092,6 +1146,20 @@ pub trait ExprRewriter: Sized {
     fn mutate(&mut self, expr: Expr) -> Result<Expr>;
 }
 
+/// The information necessary to apply algebraic simplification to an
+/// [Expr]. See [SimplifyContext] for one implementation
+pub trait SimplifyInfo {
+    /// returns true if this Expr has boolean type
+    fn is_boolean_type(&self, expr: &Expr) -> Result<bool>;
+
+    /// returns true of this expr is nullable (could possibly be NULL)
+    fn nullable(&self, expr: &Expr) -> Result<bool>;
+
+    /// Returns details needed for partial expression evaluation
+    fn execution_props(&self) -> &ExecutionProps;
+}
+
+/// Helper struct for building [Expr::Case]
 pub struct CaseBuilder {
     expr: Option<Box<Expr>>,
     when_expr: Vec<Expr>,

datafusion/src/logical_plan/mod.rs

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

datafusion/src/optimizer/simplify_expressions.rs

@@ -23,17 +23,68 @@ use arrow::record_batch::RecordBatch;
 
 use crate::error::DataFusionError;
 use crate::execution::context::ExecutionProps;
-use crate::logical_plan::{lit, DFSchemaRef, Expr};
-use crate::logical_plan::{DFSchema, ExprRewriter, LogicalPlan, RewriteRecursion};
+use crate::logical_plan::{
+    lit, DFSchema, DFSchemaRef, Expr, ExprRewriter, LogicalPlan, RewriteRecursion,
+    SimplifyInfo,
+};
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::utils;
 use crate::physical_plan::functions::Volatility;
 use crate::physical_plan::planner::create_physical_expr;
 use crate::scalar::ScalarValue;
 use crate::{error::Result, logical_plan::Operator};
 
-/// Simplifies plans by rewriting [`Expr`]`s evaluating constants
-/// and applying algebraic simplifications
+/// Provides simplification information based on schema and properties
+struct SimplifyContext<'a, 'b> {
+    schemas: Vec<&'a DFSchemaRef>,
+    props: &'b ExecutionProps,
+}
+
+impl<'a, 'b> SimplifyContext<'a, 'b> {
+    /// Create a new SimplifyContext
+    pub fn new(schemas: Vec<&'a DFSchemaRef>, props: &'b ExecutionProps) -> Self {
+        Self { schemas, props }
+    }
+}
+
+impl<'a, 'b> SimplifyInfo for SimplifyContext<'a, 'b> {
+    /// returns true if this Expr has boolean type
+    fn is_boolean_type(&self, expr: &Expr) -> Result<bool> {
+        for schema in &self.schemas {
+            if let Ok(DataType::Boolean) = expr.get_type(schema) {
+                return Ok(true);
+            }
+        }
+
+        Ok(false)
+    }
+    /// Returns true if expr is nullable
+    fn nullable(&self, expr: &Expr) -> Result<bool> {
+        self.schemas
+            .iter()
+            .find_map(|schema| {
+                // expr may be from another input, so ignore errors
+                // by converting to None to keep trying
+                expr.nullable(schema.as_ref()).ok()
+            })
+            .ok_or_else(|| {
+                // This means we weren't able to compute `Expr::nullable` with
+                // *any* input schemas, signalling a problem
+                DataFusionError::Internal(format!(
+                    "Could not find find columns in '{}' during simplify",
+                    expr
+                ))
+            })
+    }
+
+    fn execution_props(&self) -> &ExecutionProps {
+        self.props
+    }
+}
+
+/// Optimizer Pass that simplifies [`LogicalPlan`]s by rewriting
+/// [`Expr`]`s evaluating constants and applying algebraic
+/// simplifications
 ///
 /// # Introduction
 /// It uses boolean algebra laws to simplify or reduce the number of terms in expressions.
@@ -44,7 +95,7 @@ use crate::{error::Result, logical_plan::Operator};
 /// `Filter: b > 2`
 ///
 #[derive(Default)]
-pub struct SimplifyExpressions {}
+pub(crate) struct SimplifyExpressions {}
 
 /// returns true if `needle` is found in a chain of search_op
 /// expressions. Such as: (A AND B) AND C
@@ -150,9 +201,7 @@ impl OptimizerRule for SimplifyExpressions {
         // projected columns. With just the projected schema, it's not possible to infer types for
         // expressions that references non-projected columns within the same project plan or its
         // children plans.
-        let mut simplifier = Simplifier::new(plan.all_schemas());
-
-        let mut const_evaluator = ConstEvaluator::new(execution_props);
+        let info = SimplifyContext::new(plan.all_schemas(), execution_props);
 
         let new_inputs = plan
             .inputs()
@@ -168,15 +217,8 @@ impl OptimizerRule for SimplifyExpressions {
                 // Constant folding should not change expression name.
                 let name = &e.name(plan.schema());
 
-                // TODO iterate until no changes are made
-                // during rewrite (evaluating constants can
-                // enable new simplifications and
-                // simplifications can enable new constant
-                // evaluation)
-                let new_e = e
-                    // fold constants and then simplify
-                    .rewrite(&mut const_evaluator)?
-                    .rewrite(&mut simplifier)?;
+                // Apply the actual simplification logic
+                let new_e = e.simplify(&info)?;
 
                 let new_name = &new_e.name(plan.schema());
 
@@ -389,52 +431,23 @@ impl<'a> ConstEvaluator<'a> {
 /// * `false = true` and `true = false` to `false`
 /// * `!!expr` to `expr`
 /// * `expr = null` and `expr != null` to `null`
-pub(crate) struct Simplifier<'a> {
-    /// input schemas
-    schemas: Vec<&'a DFSchemaRef>,
+pub(crate) struct Simplifier<'a, S> {
+    info: &'a S,
 }
 
-impl<'a> Simplifier<'a> {
-    pub fn new(schemas: Vec<&'a DFSchemaRef>) -> Self {
-        Self { schemas }
-    }
-
-    fn is_boolean_type(&self, expr: &Expr) -> bool {
-        for schema in &self.schemas {
-            if let Ok(DataType::Boolean) = expr.get_type(schema) {
-                return true;
-            }
-        }
-
-        false
-    }
-
-    /// Returns true if expr is nullable
-    fn nullable(&self, expr: &Expr) -> Result<bool> {
-        self.schemas
-            .iter()
-            .find_map(|schema| {
-                // expr may be from another input, so ignore errors
-                // by converting to None to keep trying
-                expr.nullable(schema.as_ref()).ok()
-            })
-            .ok_or_else(|| {
-                // This means we weren't able to compute `Expr::nullable` with
-                // *any* input schemas, signalling a problem
-                DataFusionError::Internal(format!(
-                    "Could not find find columns in '{}' during simplify",
-                    expr
-                ))
-            })
+impl<'a, S> Simplifier<'a, S> {
+    pub fn new(info: &'a S) -> Self {
+        Self { info }
     }
 }
 
-impl<'a> ExprRewriter for Simplifier<'a> {
+impl<'a, S: SimplifyInfo> ExprRewriter for Simplifier<'a, S> {
     /// rewrite the expression simplifying any constant expressions
     fn mutate(&mut self, expr: Expr) -> Result<Expr> {
         use Expr::*;
         use Operator::{And, Divide, Eq, Multiply, NotEq, Or};
 
+        let info = self.info;
         let new_expr = match expr {
             //
             // Rules for Eq
@@ -447,7 +460,7 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                 left,
                 op: Eq,
                 right,
-            } if is_bool_lit(&left) && self.is_boolean_type(&right) => {
+            } if is_bool_lit(&left) && info.is_boolean_type(&right)? => {
                 match as_bool_lit(*left) {
                     Some(true) => *right,
                     Some(false) => Not(right),
@@ -461,7 +474,7 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                 left,
                 op: Eq,
                 right,
-            } if is_bool_lit(&right) && self.is_boolean_type(&left) => {
+            } if is_bool_lit(&right) && info.is_boolean_type(&left)? => {
                 match as_bool_lit(*right) {
                     Some(true) => *left,
                     Some(false) => Not(left),
@@ -480,7 +493,7 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                 left,
                 op: NotEq,
                 right,
-            } if is_bool_lit(&left) && self.is_boolean_type(&right) => {
+            } if is_bool_lit(&left) && info.is_boolean_type(&right)? => {
                 match as_bool_lit(*left) {
                     Some(true) => Not(right),
                     Some(false) => *right,
@@ -494,7 +507,7 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                 left,
                 op: NotEq,
                 right,
-            } if is_bool_lit(&right) && self.is_boolean_type(&left) => {
+            } if is_bool_lit(&right) && info.is_boolean_type(&left)? => {
                 match as_bool_lit(*right) {
                     Some(true) => Not(left),
                     Some(false) => *left,
@@ -547,13 +560,13 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                 left,
                 op: Or,
                 right,
-            } if !self.nullable(&right)? && is_op_with(And, &right, &left) => *left,
+            } if !info.nullable(&right)? && is_op_with(And, &right, &left) => *left,
             // (A AND B) OR A --> A (if B not null)
             BinaryExpr {
                 left,
                 op: Or,
                 right,
-            } if !self.nullable(&left)? && is_op_with(And, &left, &right) => *right,
+            } if !info.nullable(&left)? && is_op_with(And, &left, &right) => *right,
 
             //
             // Rules for AND
@@ -600,13 +613,13 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                 left,
                 op: And,
                 right,
-            } if !self.nullable(&right)? && is_op_with(Or, &right, &left) => *left,
+            } if !info.nullable(&right)? && is_op_with(Or, &right, &left) => *left,
             // (A OR B) AND A --> A (if B not null)
             BinaryExpr {
                 left,
                 op: And,
                 right,
-            } if !self.nullable(&left)? && is_op_with(Or, &left, &right) => *right,
+            } if !info.nullable(&left)? && is_op_with(Or, &left, &right) => *right,
 
             //
             // Rules for Multiply
@@ -643,7 +656,7 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                 left,
                 op: Divide,
                 right,
-            } if !self.nullable(&left)? && left == right => lit(1),
+            } if !info.nullable(&left)? && left == right => lit(1),
 
             //
             // Rules for Not
@@ -1160,15 +1173,9 @@ mod tests {
 
     fn simplify(expr: Expr) -> Expr {
         let schema = expr_test_schema();
-        let mut rewriter = Simplifier::new(vec![&schema]);
-
         let execution_props = ExecutionProps::new();
-        let mut const_evaluator = ConstEvaluator::new(&execution_props);
-
-        expr.rewrite(&mut rewriter)
-            .expect("expected to simplify")
-            .rewrite(&mut const_evaluator)
-            .expect("expected to const evaluate")
+        let info = SimplifyContext::new(vec![&schema], &execution_props);
+        expr.simplify(&info).unwrap()
     }
 
     fn expr_test_schema() -> DFSchemaRef {