diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 1e4d106837cc..e833d616f445 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/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) -> Result> { +/// Uses three-valued bounds when input columns prevent constant evaluation. +fn evaluate_predicate(predicate: &Arc) -> Result { // 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) -> Result> // 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::() + && 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::() { + 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 = col("a", &schema)?; + let left: Arc = + binary(Arc::clone(&a), Operator::Gt, lit(1), &schema)?; + let right: Arc = + binary(col("b", &schema)?, Operator::Lt, lit(3), &schema)?; + let predicate: Arc = if reverse { + binary(right, op, left, &schema)? + } else { + binary(left, op, right, &schema)? + }; + let predicate: Arc = 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 = col("a", &schema)?; + let b: Arc = col("b", &schema)?; + + // a > 1 AND (b < 3 OR b IS NULL) + let a_gt_1: Arc = + binary(Arc::clone(&a), Operator::Gt, lit(1), &schema)?; + let b_lt_3_or_null: Arc = binary( + binary(Arc::clone(&b), Operator::Lt, lit(3), &schema)?, + Operator::Or, + expressions::is_null(b)?, + &schema, + )?; + let predicate: Arc = 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 = 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, Option)> = [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, + ) -> 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)]); diff --git a/datafusion/sqllogictest/test_files/cse.slt b/datafusion/sqllogictest/test_files/cse.slt index 6f8abd8085d3..5681c16d14eb 100644 --- a/datafusion/sqllogictest/test_files/cse.slt +++ b/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