[arrow-datafusion] branch master updated: Improve simplification by running it twice (#3811)

[agrove]

This is an automated email from the ASF dual-hosted git repository.

agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 37fe93826 Improve simplification by running it twice (#3811)
37fe93826 is described below

commit 37fe938261636bb7710b26cf26e2bbd010e1dbf0
Author: Andrew Lamb <and...@nerdnetworks.org>
AuthorDate: Wed Oct 12 16:13:00 2022 -0400

    Improve simplification by running it twice (#3811)
---
 datafusion-examples/examples/expr_api.rs    |  9 +++---
 datafusion/optimizer/src/expr_simplifier.rs | 44 +++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/datafusion-examples/examples/expr_api.rs 
b/datafusion-examples/examples/expr_api.rs
index 35a508b0f..fe7a7936d 100644
--- a/datafusion-examples/examples/expr_api.rs
+++ b/datafusion-examples/examples/expr_api.rs
@@ -106,12 +106,11 @@ fn simplify_demo() -> Result<()> {
         col("i") + lit(3)
     );
 
-    // TODO uncomment when 
https://github.com/apache/arrow-datafusion/issues/1160 is done
     // (i * 0) > 5 --> false (only if null)
-    // assert_eq!(
-    //     simplifier.simplify((col("i") * lit(0)).gt(lit(5)))?,
-    //     lit(false)
-    // );
+    assert_eq!(
+        simplifier.simplify((col("i") * lit(0)).gt(lit(5)))?,
+        lit(false)
+    );
 
     // Logical simplification
 
diff --git a/datafusion/optimizer/src/expr_simplifier.rs 
b/datafusion/optimizer/src/expr_simplifier.rs
index 6e56ff715..a30fa9c3b 100644
--- a/datafusion/optimizer/src/expr_simplifier.rs
+++ b/datafusion/optimizer/src/expr_simplifier.rs
@@ -111,14 +111,18 @@ impl<S: SimplifyInfo> ExprSimplifier<S> {
     /// assert_eq!(expr, b_lt_2);
     /// ```
     pub fn simplify(&self, expr: Expr) -> Result<Expr> {
-        let mut rewriter = Simplifier::new(&self.info);
+        let mut simplifier = Simplifier::new(&self.info);
         let mut const_evaluator = 
ConstEvaluator::try_new(self.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
-        expr.rewrite(&mut const_evaluator)?.rewrite(&mut rewriter)
+        expr.rewrite(&mut const_evaluator)?
+            .rewrite(&mut simplifier)?
+            // run both passes twice to try an minimize simplifications that 
we missed
+            .rewrite(&mut const_evaluator)?
+            .rewrite(&mut simplifier)
     }
 
     /// Apply type coercion to an [`Expr`] so that it can be
@@ -231,7 +235,7 @@ mod tests {
     use super::*;
     use arrow::datatypes::{Field, Schema};
     use datafusion_common::ToDFSchema;
-    use datafusion_expr::{col, lit};
+    use datafusion_expr::{col, lit, when};
 
     #[test]
     fn api_basic() {
@@ -274,4 +278,38 @@ mod tests {
         .to_dfschema_ref()
         .unwrap()
     }
+
+    #[test]
+    fn simplify_and_constant_prop() {
+        let props = ExecutionProps::new();
+        let simplifier =
+            
ExprSimplifier::new(SimplifyContext::new(&props).with_schema(test_schema()));
+
+        // should be able to simplify to false
+        // (i * (1 - 2)) > 0
+        let expr = (col("i") * (lit(1) - lit(1))).gt(lit(0));
+        let expected = lit(false);
+        assert_eq!(expected, simplifier.simplify(expr).unwrap());
+    }
+
+    #[test]
+    fn simplify_and_constant_prop_with_case() {
+        let props = ExecutionProps::new();
+        let simplifier =
+            
ExprSimplifier::new(SimplifyContext::new(&props).with_schema(test_schema()));
+
+        //   CASE
+        //     WHEN i>5 AND false THEN i > 5
+        //     WHEN i<5 AND true THEN i < 5
+        //     ELSE false
+        //   END
+        //
+        // Can be simplified to `i < 5`
+        let expr = when(col("i").gt(lit(5)).and(lit(false)), 
col("i").gt(lit(5)))
+            .when(col("i").lt(lit(5)).and(lit(true)), col("i").lt(lit(5)))
+            .otherwise(lit(false))
+            .unwrap();
+        let expected = col("i").lt(lit(5));
+        assert_eq!(expected, simplifier.simplify(expr).unwrap());
+    }
 }