From 7563cdbbcd75c252a5650eda59a5e420e2bd852d Mon Sep 17 00:00:00 2001
From: byteink <jonahgaox@gmail.com>
Date: Tue, 16 May 2023 00:18:20 +0800
Subject: [PATCH] Simplify IsNotNull and IsNull expression (#6345)

---
 .../simplify_expressions/expr_simplifier.rs   | 34 ++++++++++++++++++
 .../simplify_expressions/simplify_exprs.rs    | 35 ++++++++++++++++---
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
index d077c9f83a04..699e92a2085a 100644
--- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
+++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
@@ -1161,6 +1161,12 @@ impl<'a, S: SimplifyInfo> TreeNodeRewriter for Simplifier<'a, S> {
                 lit(!negated)
             }
 
+            // a IS NOT NULL --> true, if a is not nullable
+            Expr::IsNotNull(expr) if !info.nullable(&expr)? => lit(true),
+
+            // a IS NULL --> false, if a is not nullable
+            Expr::IsNull(expr) if !info.nullable(&expr)? => lit(false),
+
             // no additional rewrites possible
             expr => expr,
         };
@@ -2596,6 +2602,34 @@ mod tests {
         );
     }
 
+    #[test]
+    fn simplify_expr_is_not_null() {
+        assert_eq!(
+            simplify(Expr::IsNotNull(Box::new(col("c1")))),
+            Expr::IsNotNull(Box::new(col("c1")))
+        );
+
+        // 'c1_non_null IS NOT NULL' is always true
+        assert_eq!(
+            simplify(Expr::IsNotNull(Box::new(col("c1_non_null")))),
+            lit(true)
+        );
+    }
+
+    #[test]
+    fn simplify_expr_is_null() {
+        assert_eq!(
+            simplify(Expr::IsNull(Box::new(col("c1")))),
+            Expr::IsNull(Box::new(col("c1")))
+        );
+
+        // 'c1_non_null IS NULL' is always false
+        assert_eq!(
+            simplify(Expr::IsNull(Box::new(col("c1_non_null")))),
+            lit(false)
+        );
+    }
+
     #[test]
     fn simplify_expr_eq() {
         let schema = expr_test_schema();
diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
index 6717fe42bf20..83c7923f07d2 100644
--- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
+++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
@@ -167,6 +167,7 @@ mod tests {
             Field::new("b", DataType::Boolean, false),
             Field::new("c", DataType::Boolean, false),
             Field::new("d", DataType::UInt32, false),
+            Field::new("e", DataType::UInt32, true),
         ]);
         table_scan(Some("test"), &schema, None)
             .expect("creating scan")
@@ -623,9 +624,9 @@ mod tests {
         let table_scan = test_table_scan();
 
         let plan = LogicalPlanBuilder::from(table_scan)
-            .filter(col("d").is_null().not())?
+            .filter(col("e").is_null().not())?
             .build()?;
-        let expected = "Filter: test.d IS NOT NULL\
+        let expected = "Filter: test.e IS NOT NULL\
         \n  TableScan: test";
 
         assert_optimized_plan_eq(&plan, expected)
@@ -636,9 +637,9 @@ mod tests {
         let table_scan = test_table_scan();
 
         let plan = LogicalPlanBuilder::from(table_scan)
-            .filter(col("d").is_not_null().not())?
+            .filter(col("e").is_not_null().not())?
             .build()?;
-        let expected = "Filter: test.d IS NULL\
+        let expected = "Filter: test.e IS NULL\
         \n  TableScan: test";
 
         assert_optimized_plan_eq(&plan, expected)
@@ -848,4 +849,30 @@ mod tests {
             "TableScan: test, unsupported_filters=[g = f AS g = power(f,Float64(1))]";
         assert_optimized_plan_eq(&plan, expected)
     }
+
+    #[test]
+    fn simplify_is_not_null() -> Result<()> {
+        let table_scan = test_table_scan();
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .filter(col("d").is_not_null())?
+            .build()?;
+        let expected = "Filter: Boolean(true)\
+        \n  TableScan: test";
+
+        assert_optimized_plan_eq(&plan, expected)
+    }
+
+    #[test]
+    fn simplify_is_null() -> Result<()> {
+        let table_scan = test_table_scan();
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .filter(col("d").is_null())?
+            .build()?;
+        let expected = "Filter: Boolean(false)\
+        \n  TableScan: test";
+
+        assert_optimized_plan_eq(&plan, expected)
+    }
 }