#17801 Improve nullability reporting of case expressions

pepijnve · pepijnve · commit 482d0bea868e · 2025-09-29T12:51:22.000+02:00
diff --git a/datafusion/core/tests/tpcds_planning.rs b/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)?;
         }
     }
 
diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
@@ -340,6 +340,11 @@ pub fn is_null(expr: Expr) -> Expr {
     Expr::IsNull(Box::new(expr))
 }
 
+/// Create is not null expression
+pub fn is_not_null(expr: Expr) -> Expr {
+    Expr::IsNotNull(Box::new(expr))
+}
+
 /// Create is true expression
 pub fn is_true(expr: Expr) -> Expr {
     Expr::IsTrue(Box::new(expr))
diff --git a/datafusion/expr/src/expr_schema.rs b/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,11 @@ impl ExprSchemable for Expr {
                 let then_nullable = case
                     .when_then_expr
                     .iter()
+                    .filter(|(w, t)| {
+                        // Disregard branches where we can determine statically that the predicate
+                        // is always false when the then expression would evaluate to null
+                        const_result_when_value_is_null(w, t).unwrap_or(true)
+                    })
                     .map(|(_, t)| t.nullable(input_schema))
                     .collect::<Result<Vec<_>>>()?;
                 if then_nullable.contains(&true) {
@@ -647,6 +653,38 @@ impl ExprSchemable for Expr {
     }
 }
 
+/// Determines if the given `predicate` can be const evaluated if `value` were to evaluate to `NULL`.
+/// Returns a `Some` value containing the const result if so; otherwise returns `None`.
+fn const_result_when_value_is_null(predicate: &Expr, value: &Expr) -> Option<bool> {
+    match predicate {
+        Expr::IsNotNull(e) => if e.as_ref().eq(value) { Some(false) } else { None },
+        Expr::IsNull(e) => if e.as_ref().eq(value) { Some(true) } else { None },
+        Expr::Not(e) => const_result_when_value_is_null(e, value).map(|b| !b),
+        Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
+            Operator::And => {
+                let l = const_result_when_value_is_null(left, value);
+                let r = const_result_when_value_is_null(right, value);
+                match (l, r) {
+                    (Some(l), Some(r)) => Some(l && r),
+                    (Some(l), None) => Some(l),
+                    (None, Some(r)) => Some(r),
+                    _ => None,
+                }
+            }
+            Operator::Or => {
+                let l = const_result_when_value_is_null(left, value);
+                let r = const_result_when_value_is_null(right, value);
+                match (l, r) {
+                    (Some(l), Some(r)) => Some(l || r),
+                    _ => 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
@@ -777,7 +815,7 @@ pub fn cast_subquery(subquery: Subquery, cast_to_type: &DataType) -> Result<Subq
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{col, lit, out_ref_col_with_metadata};
+    use crate::{and, binary_expr, col, is_not_null, is_null, lit, not, or, out_ref_col_with_metadata, when};
 
     use datafusion_common::{internal_err, DFSchema, HashMap, ScalarValue};
 
@@ -830,6 +868,153 @@ mod tests {
         assert!(expr.nullable(&get_schema(false)).unwrap());
     }
 
+    fn check_nullability(
+        expr: Expr,
+        nullable: bool,
+        get_schema: fn(bool) -> MockExprSchema,
+    ) -> Result<()> {
+        assert_eq!(
+            expr.nullable(&get_schema(true))?,
+            nullable,
+            "Nullability of '{}' should be {} when column is nullable",
+            expr,
+            nullable
+        );
+        assert!(
+            !expr.nullable(&get_schema(false))?,
+            "Nullability of '{}' should be false when column is not nullable",
+            expr
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_case_expression_nullability() -> Result<()> {
+        let get_schema = |nullable| {
+            MockExprSchema::new()
+                .with_data_type(DataType::Int32)
+                .with_nullable(nullable)
+        };
+
+        check_nullability(
+            when(is_not_null(col("foo")), col("foo")).otherwise(lit(0))?,
+            false,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(not(is_null(col("foo"))), col("foo")).otherwise(lit(0))?,
+            false,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(binary_expr(col("foo"), Operator::Eq, lit(5)), col("foo"))
+                .otherwise(lit(0))?,
+            true,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(
+                and(
+                    is_not_null(col("foo")),
+                    binary_expr(col("foo"), Operator::Eq, lit(5)),
+                ),
+                col("foo"),
+            )
+            .otherwise(lit(0))?,
+            false,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(
+                and(
+                    binary_expr(col("foo"), Operator::Eq, lit(5)),
+                    is_not_null(col("foo")),
+                ),
+                col("foo"),
+            )
+            .otherwise(lit(0))?,
+            false,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(
+                or(
+                    is_not_null(col("foo")),
+                    binary_expr(col("foo"), Operator::Eq, lit(5)),
+                ),
+                col("foo"),
+            )
+            .otherwise(lit(0))?,
+            true,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(
+                or(
+                    binary_expr(col("foo"), Operator::Eq, lit(5)),
+                    is_not_null(col("foo")),
+                ),
+                col("foo"),
+            )
+            .otherwise(lit(0))?,
+            true,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(
+                or(
+                    is_not_null(col("foo")),
+                    binary_expr(col("foo"), Operator::Eq, lit(5)),
+                ),
+                col("foo"),
+            )
+            .otherwise(lit(0))?,
+            true,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(
+                or(
+                    binary_expr(col("foo"), Operator::Eq, lit(5)),
+                    is_not_null(col("foo")),
+                ),
+                col("foo"),
+            )
+            .otherwise(lit(0))?,
+            true,
+            get_schema,
+        )?;
+
+        check_nullability(
+            when(
+                or(
+                    and(
+                        binary_expr(col("foo"), Operator::Eq, lit(5)),
+                        is_not_null(col("foo")),
+                    ),
+                    and(
+                        binary_expr(col("foo"), Operator::Eq, col("bar")),
+                        is_not_null(col("foo")),
+                    ),
+                ),
+                col("foo"),
+            )
+            .otherwise(lit(0))?,
+            false,
+            get_schema,
+        )?;
+
+        Ok(())
+    }
+
     #[test]
     fn test_inlist_nullability() {
         let get_schema = |nullable| {
diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs

Original file line number	Diff line number	Diff line change
`@@ -1052,9 +1052,10 @@ async fn regression_test(query_no: u8, create_physical: bool) -> Result<()> {`
`1052`	`1052`	`for sql in &sql {`
`1053`	`1053`	`let df = ctx.sql(sql).await?;`
`1054`	`1054`	`let (state, plan) = df.into_parts();`
`1055`		`- let plan = state.optimize(&plan)?;`
`1056`	`1055`	`if create_physical {`
`1057`	`1056`	`let _ = state.create_physical_plan(&plan).await?;`
	`1057`	`+ } else {`
	`1058`	`+ let _ = state.optimize(&plan)?;`
`1058`	`1059`	`}`
`1059`	`1060`	`}`
`1060`	`1061`