[arrow-datafusion] 03/04: test the outer function

[wj...@apache.org]

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

wjones127 pushed a commit to branch 6171-simplify-with-guarantee
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git

commit ff7ed70260d9aa5512ab1bc494bd885497d39479
Author: Will Jones <wi...@gmail.com>
AuthorDate: Mon Sep 4 11:34:01 2023 -0700

    test the outer function
---
 .../src/simplify_expressions/expr_simplifier.rs    | 66 +++++++++++++++++++++-
 1 file changed, 65 insertions(+), 1 deletion(-)

diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
index fc39aafe53..78a8d70635 100644
--- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
+++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
@@ -153,7 +153,7 @@ impl<S: SimplifyInfo> ExprSimplifier<S> {
     }
 
     /// Add guarantees
-    pub fn simplify_with_gurantee<'a>(
+    pub fn simplify_with_guarantees<'a>(
         &self,
         expr: Expr,
         guarantees: impl IntoIterator<Item = &'a (Expr, Guarantee)>,
@@ -1215,6 +1215,7 @@ mod tests {
     };
 
     use crate::simplify_expressions::{
+        guarantees::{GuaranteeBound, NullStatus},
         utils::for_test::{cast_to_int64_expr, now_expr, to_timestamp_expr},
         SimplifyContext,
     };
@@ -2693,6 +2694,17 @@ mod tests {
         try_simplify(expr).unwrap()
     }
 
+    fn simplify_with_guarantee(expr: Expr, guarantees: &[(Expr, Guarantee)]) -> Expr {
+        let schema = expr_test_schema();
+        let execution_props = ExecutionProps::new();
+        let simplifier = ExprSimplifier::new(
+            SimplifyContext::new(&execution_props).with_schema(schema),
+        );
+        simplifier
+            .simplify_with_guarantees(expr, guarantees)
+            .unwrap()
+    }
+
     fn expr_test_schema() -> DFSchemaRef {
         Arc::new(
             DFSchema::new_with_metadata(
@@ -3156,4 +3168,56 @@ mod tests {
         let expr = not_ilike(null, "%");
         assert_eq!(simplify(expr), lit_bool_null());
     }
+
+    #[test]
+    fn test_simplify_with_guarantee() {
+        // (x >= 3) AND (y + 2 < 10 OR (z NOT IN ("a", "b")))
+        let expr_x = col("c3").gt(lit(3_i64));
+        let expr_y = (col("c4") + lit(2_u32)).lt(lit(10_u32));
+        let expr_z = col("c1").in_list(vec![lit("a"), lit("b")], true);
+        let expr = expr_x.clone().and(expr_y.or(expr_z));
+
+        // All guaranteed null
+        let guarantees = vec![
+            (col("c3"), Guarantee::from(&ScalarValue::Int64(None))),
+            (col("c4"), Guarantee::from(&ScalarValue::UInt32(None))),
+            (col("c1"), Guarantee::from(&ScalarValue::Utf8(None))),
+        ];
+
+        let output = simplify_with_guarantee(expr.clone(), &guarantees);
+        assert_eq!(output, lit_bool_null());
+
+        // All guaranteed false
+        let guarantees = vec![
+            (
+                col("c3"),
+                Guarantee::new(
+                    Some(GuaranteeBound::new(ScalarValue::Int64(Some(0)), false)),
+                    Some(GuaranteeBound::new(ScalarValue::Int64(Some(2)), false)),
+                    NullStatus::NeverNull,
+                ),
+            ),
+            (col("c4"), Guarantee::from(&ScalarValue::UInt32(Some(9)))),
+            (
+                col("c1"),
+                Guarantee::from(&ScalarValue::Utf8(Some("a".to_string()))),
+            ),
+        ];
+        let output = simplify_with_guarantee(expr.clone(), &guarantees);
+        assert_eq!(output, lit(false));
+
+        // Sufficient true guarantees
+        let guarantees = vec![
+            (col("c3"), Guarantee::from(&ScalarValue::Int64(Some(9)))),
+            (col("c4"), Guarantee::from(&ScalarValue::UInt32(Some(3)))),
+        ];
+        let output = simplify_with_guarantee(expr.clone(), &guarantees);
+        assert_eq!(output, lit(true));
+
+        // Only partially simplify
+        let guarantees =
+            vec![(col("c4"), Guarantee::from(&ScalarValue::UInt32(Some(3))))];
+        let output = simplify_with_guarantee(expr.clone(), &guarantees);
+        assert_eq!(&output, &expr_x);
+    }
 }