#17801 Improve nullability reporting of case expressions

Author: pepijnve

datafusion/core/tests/tpcds_planning.rs

@@ -1052,9 +1052,10 @@ async fn regression_test(query_no: u8, create_physical: bool) -> Result<()> {
     for sql in &sql {
         let df = ctx.sql(sql).await?;
         let (state, plan) = df.into_parts();
-        let plan = state.optimize(&plan)?;
         if create_physical {
             let _ = state.create_physical_plan(&plan).await?;
+        } else {
+            let _ = state.optimize(&plan)?;
         }
     }
 

datafusion/expr/src/expr_schema.rs

@@ -32,6 +32,7 @@ use datafusion_common::{
     not_impl_err, plan_datafusion_err, plan_err, Column, DataFusionError, ExprSchema,
     Result, Spans, TableReference,
 };
+use datafusion_expr_common::operator::Operator;
 use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer;
 use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 use std::sync::Arc;
@@ -283,6 +284,9 @@ impl ExprSchemable for Expr {
                 let then_nullable = case
                     .when_then_expr
                     .iter()
+                    .filter(|(w, t)| {
+                        !always_false_when_value_is_null(w, t).unwrap_or(false)
+                    })
                     .map(|(_, t)| t.nullable(input_schema))
                     .collect::<Result<Vec<_>>>()?;
                 if then_nullable.contains(&true) {
@@ -647,6 +651,35 @@ impl ExprSchemable for Expr {
     }
 }
 
+fn always_false_when_value_is_null(predicate: &Expr, value: &Expr) -> Option<bool> {
+    match predicate {
+        Expr::IsNotNull(e) => Some(e.as_ref().eq(value)),
+        Expr::Not(e) => always_false_when_value_is_null(e, value).map(|b| !b),
+        Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
+            Operator::And => {
+                let l = always_false_when_value_is_null(left, value);
+                let r = always_false_when_value_is_null(right, value);
+                if l.is_some() && r.is_some() {
+                    Some(l.unwrap() || r.unwrap())
+                } else {
+                    None
+                }
+            }
+            Operator::Or => {
+                let l = always_false_when_value_is_null(left, value);
+                let r = always_false_when_value_is_null(right, value);
+                if l.is_some() && r.is_some() {
+                    Some(l.unwrap() && r.unwrap())
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
 impl Expr {
     /// Common method for window functions that applies type coercion
     /// to all arguments of the window function to check if it matches

datafusion/physical-expr/src/expressions/case.rs

@@ -15,12 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::expressions::try_cast;
+use crate::expressions::{try_cast, BinaryExpr, IsNotNullExpr, NotExpr};
 use crate::PhysicalExpr;
-use std::borrow::Cow;
-use std::hash::Hash;
-use std::{any::Any, sync::Arc};
-
 use arrow::array::*;
 use arrow::compute::kernels::zip::zip;
 use arrow::compute::{and, and_not, is_null, not, nullif, or, prep_null_mask_filter};
@@ -30,8 +26,13 @@ use datafusion_common::{
     exec_err, internal_datafusion_err, internal_err, DataFusionError, Result, ScalarValue,
 };
 use datafusion_expr::ColumnarValue;
+use std::borrow::Cow;
+use std::hash::Hash;
+use std::{any::Any, sync::Arc};
 
 use super::{Column, Literal};
+use datafusion_expr_common::dyn_eq::DynEq;
+use datafusion_expr_common::operator::Operator;
 use datafusion_physical_expr_common::datum::compare_with_eq;
 use itertools::Itertools;
 
@@ -481,6 +482,9 @@ impl PhysicalExpr for CaseExpr {
         let then_nullable = self
             .when_then_expr
             .iter()
+            .filter(|(w, t)| {
+                !always_false_when_value_is_null(w.as_ref(), t.as_ref()).unwrap_or(false)
+            })
             .map(|(_, t)| t.nullable(input_schema))
             .collect::<Result<Vec<_>>>()?;
         if then_nullable.contains(&true) {
@@ -588,6 +592,42 @@ impl PhysicalExpr for CaseExpr {
     }
 }
 
+fn always_false_when_value_is_null(
+    predicate: &dyn PhysicalExpr,
+    value: &dyn PhysicalExpr,
+) -> Option<bool> {
+    let predicate_any = predicate.as_any();
+    if let Some(not_null) = predicate_any.downcast_ref::<IsNotNullExpr>() {
+        Some(not_null.arg().as_ref().dyn_eq(value))
+    } else if let Some(not) = predicate_any.downcast_ref::<NotExpr>() {
+        always_false_when_value_is_null(not.arg().as_ref(), value).map(|b| !b)
+    } else if let Some(binary) = predicate_any.downcast_ref::<BinaryExpr>() {
+        match binary.op() {
+            Operator::And => {
+                let l = always_false_when_value_is_null(binary.left().as_ref(), value);
+                let r = always_false_when_value_is_null(binary.right().as_ref(), value);
+                if l.is_some() && r.is_some() {
+                    Some(l.unwrap() || r.unwrap())
+                } else {
+                    None
+                }
+            }
+            Operator::Or => {
+                let l = always_false_when_value_is_null(binary.left().as_ref(), value);
+                let r = always_false_when_value_is_null(binary.right().as_ref(), value);
+                if l.is_some() && r.is_some() {
+                    Some(l.unwrap() && r.unwrap())
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        }
+    } else {
+        None
+    }
+}
+
 /// Create a CASE expression
 pub fn case(
     expr: Option<Arc<dyn PhysicalExpr>>,