fix: plan aggregates over CASE conditions with multiple columns

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

@@ -19,7 +19,7 @@ mod literal_lookup_table;
 
 use super::{Column, Literal};
 use crate::expressions::{
-    CastExpr, LambdaVariable, NegativeExpr, NotExpr, lit, try_cast,
+    BinaryExpr, CastExpr, LambdaVariable, NegativeExpr, NotExpr, lit, try_cast,
 };
 use crate::{PhysicalExpr, ScalarFunctionExpr};
 use arrow::array::*;
@@ -35,6 +35,8 @@ use datafusion_common::{
     internal_datafusion_err, internal_err,
 };
 use datafusion_expr::ColumnarValue;
+use datafusion_expr_common::interval_arithmetic::NullableInterval;
+use datafusion_expr_common::operator::Operator;
 use indexmap::IndexMap;
 use std::borrow::Cow;
 use std::collections::BTreeSet;
@@ -1292,19 +1294,14 @@ impl PhysicalExpr for CaseExpr {
                 Ok(e) => e,
             };
 
-            // Try to const evaluate the modified `when` expression.
-            let predicate_result = match evaluate_predicate(&with_null) {
+            let can_be_true = match evaluate_predicate(&with_null).and_then(|bounds| {
+                bounds.contains_value(ScalarValue::Boolean(Some(true)))
+            }) {
                 Err(e) => return Some(Err(e)),
                 Ok(b) => b,
             };
 
-            match predicate_result {
-                // Evaluation was inconclusive or true, so the 'then' expression is reachable
-                None | Some(true) => Some(Ok(())),
-                // Evaluation proves the branch will never be taken.
-                // The most common pattern for this is `WHEN x IS NOT NULL THEN x`.
-                Some(false) => None,
-            }
+            if can_be_true { Some(Ok(())) } else { None }
         });
 
         if let Some(nullable_then) = nullable_then {
@@ -1516,12 +1513,8 @@ impl CaseExpr {
     }
 }
 
-/// Attempts to const evaluate the given `predicate`.
-/// Returns:
-/// - `Some(true)` if the predicate evaluates to a truthy value.
-/// - `Some(false)` if the predicate evaluates to a falsy value.
-/// - `None` if the predicate could not be evaluated.
-fn evaluate_predicate(predicate: &Arc<dyn PhysicalExpr>) -> Result<Option<bool>> {
+/// Uses three-valued bounds when input columns prevent constant evaluation.
+fn evaluate_predicate(predicate: &Arc<dyn PhysicalExpr>) -> Result<NullableInterval> {
     // Create a dummy record with no columns and one row
     let batch = RecordBatch::try_new_with_options(
         Arc::new(Schema::empty()),
@@ -1531,15 +1524,24 @@ fn evaluate_predicate(predicate: &Arc<dyn PhysicalExpr>) -> Result<Option<bool>>
 
     // Evaluate the predicate and interpret the result as a boolean
     let result = match predicate.evaluate(&batch) {
-        // An error during evaluation means we couldn't const evaluate the predicate, so return `None`
-        Err(_) => None,
-        Ok(ColumnarValue::Array(array)) => Some(
+        Err(_) => {
+            if let Some(expr) = predicate.downcast_ref::<BinaryExpr>()
+                && matches!(expr.op(), Operator::And | Operator::Or)
+            {
+                return evaluate_predicate(expr.left())?
+                    .apply_operator(expr.op(), &evaluate_predicate(expr.right())?);
+            }
+            if let Some(expr) = predicate.downcast_ref::<NotExpr>() {
+                return evaluate_predicate(expr.arg())?.not();
+            }
+            return Ok(NullableInterval::ANY_TRUTH_VALUE);
+        }
+        Ok(ColumnarValue::Array(array)) => {
             ScalarValue::try_from_array(array.as_ref(), 0)?
-                .cast_to(&DataType::Boolean)?,
-        ),
-        Ok(ColumnarValue::Scalar(scalar)) => Some(scalar.cast_to(&DataType::Boolean)?),
+        }
+        Ok(ColumnarValue::Scalar(scalar)) => scalar,
     };
-    Ok(result.map(|v| matches!(v, ScalarValue::Boolean(Some(true)))))
+    Ok(result.cast_to(&DataType::Boolean)?.into())
 }
 
 fn replace_with_null(
@@ -1594,7 +1596,7 @@ mod tests {
     use super::*;
 
     use crate::expressions;
-    use crate::expressions::{BinaryExpr, binary, cast, col, is_not_null};
+    use crate::expressions::{binary, cast, col, is_not_null};
     use arrow::buffer::Buffer;
     use arrow::datatypes::DataType::Float64;
     use arrow::datatypes::Field;
@@ -1603,7 +1605,6 @@ mod tests {
     use datafusion_common::plan_err;
     use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
     use datafusion_expr::type_coercion::binary::type_union_coercion;
-    use datafusion_expr_common::operator::Operator;
     use datafusion_physical_expr_common::physical_expr::fmt_sql;
     use half::f16;
 
@@ -2796,6 +2797,145 @@ mod tests {
         case_expression_nullability(false)
     }
 
+    #[rstest::rstest]
+    #[case::and(Operator::And, false, false)]
+    #[case::or(Operator::Or, false, true)]
+    #[case::not_and(Operator::And, true, true)]
+    #[case::not_or(Operator::Or, true, false)]
+    fn test_case_expression_nullability_with_independent_columns(
+        #[case] op: Operator,
+        #[case] negated: bool,
+        #[case] nullable: bool,
+        #[values(false, true)] reverse: bool,
+    ) -> Result<()> {
+        let schema = Schema::new(vec![
+            Field::new("a", DataType::Int32, true),
+            Field::new("b", DataType::Int32, true),
+        ]);
+        let a: Arc<dyn PhysicalExpr> = col("a", &schema)?;
+        let left: Arc<dyn PhysicalExpr> =
+            binary(Arc::clone(&a), Operator::Gt, lit(1), &schema)?;
+        let right: Arc<dyn PhysicalExpr> =
+            binary(col("b", &schema)?, Operator::Lt, lit(3), &schema)?;
+        let predicate: Arc<dyn PhysicalExpr> = if reverse {
+            binary(right, op, left, &schema)?
+        } else {
+            binary(left, op, right, &schema)?
+        };
+        let predicate: Arc<dyn PhysicalExpr> = if negated {
+            expressions::not(predicate)?
+        } else {
+            predicate
+        };
+
+        assert_nullability(when_then_else(&predicate, &a, &lit(0))?, &schema, nullable);
+        Ok(())
+    }
+
+    #[rstest::rstest]
+    #[case::and_over_or(false, false)]
+    #[case::not_and_over_or(true, true)]
+    fn test_case_expression_nullability_with_nested_predicates(
+        #[case] negated: bool,
+        #[case] nullable: bool,
+        #[values(false, true)] reverse: bool,
+    ) -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, true),
+            Field::new("b", DataType::Int32, true),
+        ]));
+        let a: Arc<dyn PhysicalExpr> = col("a", &schema)?;
+        let b: Arc<dyn PhysicalExpr> = col("b", &schema)?;
+
+        // a > 1 AND (b < 3 OR b IS NULL)
+        let a_gt_1: Arc<dyn PhysicalExpr> =
+            binary(Arc::clone(&a), Operator::Gt, lit(1), &schema)?;
+        let b_lt_3_or_null: Arc<dyn PhysicalExpr> = binary(
+            binary(Arc::clone(&b), Operator::Lt, lit(3), &schema)?,
+            Operator::Or,
+            expressions::is_null(b)?,
+            &schema,
+        )?;
+        let predicate: Arc<dyn PhysicalExpr> = if reverse {
+            binary(b_lt_3_or_null, Operator::And, a_gt_1, &schema)?
+        } else {
+            binary(a_gt_1, Operator::And, b_lt_3_or_null, &schema)?
+        };
+        let predicate: Arc<dyn PhysicalExpr> = if negated {
+            expressions::not(predicate)?
+        } else {
+            predicate
+        };
+        let expr = when_then_else(&predicate, &a, &lit(0))?;
+
+        assert_nullability(Arc::clone(&expr), &schema, nullable);
+
+        // Cross-check the inferred nullability against evaluation on every combination
+        // of NULL and non-NULL inputs. With a = NULL and b = 3 the negated predicate is
+        // TRUE, so the THEN branch yields NULL.
+        let rows: Vec<(Option<i32>, Option<i32>)> = [None, Some(1), Some(2)]
+            .into_iter()
+            .flat_map(|a| [None, Some(2), Some(3)].into_iter().map(move |b| (a, b)))
+            .collect();
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int32Array::from_iter(rows.iter().map(|(a, _)| *a))),
+                Arc::new(Int32Array::from_iter(rows.iter().map(|(_, b)| *b))),
+            ],
+        )?;
+        let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?;
+        assert_eq!(result.null_count() > 0, nullable);
+        Ok(())
+    }
+
+    #[rstest::rstest]
+    #[case::literal(lit("invalid"))]
+    #[case::and_left(Arc::new(BinaryExpr::new(
+        lit("invalid"),
+        Operator::And,
+        make_col("flag", 0)
+    )))]
+    #[case::and_right(Arc::new(BinaryExpr::new(
+        make_col("flag", 0),
+        Operator::And,
+        lit("invalid")
+    )))]
+    #[case::or_left(Arc::new(BinaryExpr::new(
+        lit("invalid"),
+        Operator::Or,
+        make_col("flag", 0)
+    )))]
+    #[case::or_right(Arc::new(BinaryExpr::new(
+        make_col("flag", 0),
+        Operator::Or,
+        lit("invalid")
+    )))]
+    #[case::not(Arc::new(NotExpr::new(lit("invalid"))))]
+    fn test_case_expression_nullability_with_invalid_predicate(
+        #[case] predicate: Arc<dyn PhysicalExpr>,
+    ) -> Result<()> {
+        let batch = RecordBatch::try_from_iter(vec![
+            (
+                "flag",
+                Arc::new(BooleanArray::from(vec![Some(true), Some(false), None]))
+                    as ArrayRef,
+            ),
+            (
+                "value",
+                Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef,
+            ),
+        ])?;
+        let schema = batch.schema();
+        let expr = when_then_else(&predicate, &col("value", &schema)?, &lit(0))?;
+        let error = expr
+            .nullable(&schema)
+            .expect_err("invalid predicates must fail nullability inference");
+        assert!(error.to_string().contains("invalid"), "{error}");
+        assert!(expr.evaluate(&batch).is_err());
+        Ok(())
+    }
+
     fn case_expression_nullability(col_is_nullable: bool) -> Result<()> {
         let schema =
             Schema::new(vec![Field::new("foo", DataType::Int32, col_is_nullable)]);

datafusion/sqllogictest/test_files/cse.slt

@@ -214,6 +214,71 @@ physical_plan
 02)--ProjectionExec: expr=[a@0 = 1 as __common_expr_1, a@0 = 2 as __common_expr_2, a@0 + 3 as __common_expr_3, a@0 + 4 as __common_expr_4]
 03)----DataSourceExec: partitions=1, partition_sizes=[0]
 
+statement ok
+CREATE TABLE cse_nullable_case(a INT, b INT) AS VALUES (1, 1), (2, 2), (3, 3);
+
+query I
+SELECT sum(CASE WHEN (CASE WHEN a > 1 AND b < 3 THEN a ELSE 0 END)
+                BETWEEN 1 AND 2 THEN 1 ELSE 0 END)
+FROM cse_nullable_case;
+----
+1
+
+query I
+SELECT count(CASE WHEN (CASE WHEN a > 1 AND b < 3 THEN a ELSE 0 END)
+                  BETWEEN 1 AND 2 THEN 1 ELSE 0 END)
+FROM cse_nullable_case;
+----
+3
+
+query I
+SELECT sum(CASE WHEN (CASE WHEN a > 1 AND b < 3 THEN a ELSE 0 END) >= 1
+                AND (CASE WHEN a > 1 AND b < 3 THEN a ELSE 0 END) <= 2
+                THEN 1 ELSE 0 END)
+FROM cse_nullable_case;
+----
+1
+
+statement ok
+INSERT INTO cse_nullable_case VALUES (NULL, 2), (2, NULL), (NULL, NULL);
+
+query II
+SELECT sum(v), count(v)
+FROM (SELECT CASE WHEN a > 1 AND b < 3 THEN a ELSE 0 END AS v FROM cse_nullable_case);
+----
+2 6
+
+query II
+SELECT sum(v), count(v)
+FROM (SELECT CASE WHEN a > 1 OR b < 3 THEN a ELSE 0 END AS v FROM cse_nullable_case);
+----
+8 5
+
+query II
+SELECT sum(v), count(v)
+FROM (SELECT CASE WHEN NOT (a > 1 OR b < 3) THEN a ELSE 0 END AS v FROM cse_nullable_case);
+----
+0 6
+
+query II
+SELECT sum(v), count(v)
+FROM (SELECT CASE WHEN a > 1 AND b < 3 THEN a END AS v FROM cse_nullable_case);
+----
+2 1
+
+query II
+SELECT sum(v), count(v)
+FROM (SELECT CASE WHEN a > 1 AND b < 3 THEN a ELSE b END AS v FROM cse_nullable_case);
+----
+8 4
+
+query II
+SELECT sum(v), count(v)
+FROM (SELECT CASE WHEN a > 1 AND b < 3 THEN a ELSE 0 END AS v FROM cse_nullable_case)
+WHERE false;
+----
+NULL 0
+
 # Only conditionally evaluated expressions
 query TT
 EXPLAIN SELECT