diff --git a/datafusion/functions/src/core/coalesce.rs b/datafusion/functions/src/core/coalesce.rs index 9cf3536443e6c..be93df6a710c6 100644 --- a/datafusion/functions/src/core/coalesce.rs +++ b/datafusion/functions/src/core/coalesce.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +use arrow::array::{BooleanArray, new_null_array}; +use arrow::compute::kernels::zip::zip; +use arrow::compute::{and, is_not_null, is_null}; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::{Result, exec_err, internal_err, plan_err}; use datafusion_expr::binary::try_type_union_resolution; @@ -29,7 +32,7 @@ use itertools::Itertools; #[user_doc( doc_section(label = "Conditional Functions"), - description = "Returns the first of its arguments that is not _null_. Returns _null_ if all arguments are _null_. This function is often used to substitute a default value for _null_ values.", + description = "Returns the first of its arguments that is not _null_. Returns _null_ if all arguments are _null_. Arguments after the first non-_null_ one are normally not evaluated, but when any argument other than the last is volatile every argument is evaluated, so that each is evaluated exactly once. This function is often used to substitute a default value for _null_ values.", syntax_example = "coalesce(expression1[, ..., expression_n])", sql_example = r#"```sql > select coalesce(null, null, 'datafusion'); @@ -104,6 +107,25 @@ impl ScalarUDFImpl for CoalesceFunc { } let n = args.len(); + + // The `CASE WHEN a IS NOT NULL THEN a ELSE b END` rewrite below mentions + // every non-final argument *twice* (once in the `WHEN` predicate and once + // in the `THEN` result). That is fine for deterministic arguments, but for + // a volatile argument the two mentions are two independent draws, so + // `coalesce(random_nullable_expr, default)` can take the `THEN` branch + // after the `WHEN` draw was non-null and still produce NULL. + // + // See https://github.com/apache/datafusion/issues/25477. For volatile + // arguments we keep `coalesce` intact and let the runtime kernel in + // `invoke_with_args` evaluate each argument exactly once. + // + // Only the non-final arguments are duplicated: the last argument becomes + // the `ELSE` branch, which names it once, so a volatile last argument is + // safe to rewrite and keeps its laziness. + if args[..n - 1].iter().any(|arg| arg.is_volatile()) { + return Ok(ExprSimplifyResult::Original(args)); + } + let (init, last_elem) = args.split_at(n - 1); let whens = init .iter() @@ -117,8 +139,63 @@ impl ScalarUDFImpl for CoalesceFunc { } /// coalesce evaluates to the first value which is not NULL - fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { - internal_err!("coalesce should have been simplified to case") + /// + /// This kernel is only reached when [`Self::simplify`] declined to rewrite the + /// call into a `CASE` expression, which today only happens when one of the + /// arguments is volatile. It evaluates each argument exactly once. + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let args = args.args; + // do not accept 0 arguments. + if args.is_empty() { + return exec_err!( + "coalesce was called with {} arguments. It requires at least 1.", + args.len() + ); + } + + let return_type = args[0].data_type(); + let mut return_array = args.iter().filter_map(|x| match x { + ColumnarValue::Array(array) => Some(array.len()), + _ => None, + }); + + if let Some(size) = return_array.next() { + // start with nulls as default output + let mut current_value = new_null_array(&return_type, size); + let mut remainder = BooleanArray::from(vec![true; size]); + + for arg in args { + match arg { + ColumnarValue::Array(ref array) => { + let to_apply = and(&remainder, &is_not_null(array.as_ref())?)?; + current_value = zip(&to_apply, array, ¤t_value)?; + remainder = and(&remainder, &is_null(array)?)?; + } + ColumnarValue::Scalar(value) => { + if value.is_null() { + continue; + } else { + let last_value = value.to_scalar()?; + current_value = zip(&remainder, &last_value, ¤t_value)?; + break; + } + } + } + if remainder.iter().all(|x| x == Some(false)) { + break; + } + } + Ok(ColumnarValue::Array(current_value)) + } else { + let result = args + .iter() + .find_map(|x| match x { + ColumnarValue::Scalar(s) if !s.is_null() => Some(x.clone()), + _ => None, + }) + .unwrap_or_else(|| args[0].clone()); + Ok(result) + } } fn conditional_arguments<'a>( @@ -146,3 +223,211 @@ impl ScalarUDFImpl for CoalesceFunc { self.doc() } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Array, ArrayRef, Int32Array, StringArray}; + use arrow::datatypes::DataType; + use datafusion_common::ScalarValue; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, lit}; + + use super::*; + + fn invoke(args: Vec, number_rows: usize) -> ColumnarValue { + let return_type = args[0].data_type(); + let arg_fields = args + .iter() + .enumerate() + .map(|(i, a)| Field::new(format!("a{i}"), a.data_type(), true).into()) + .collect(); + CoalesceFunc::new() + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Field::new("f", return_type, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }) + .expect("coalesce kernel failed") + } + + /// The restored runtime kernel picks the first non-null value per row. + #[test] + fn coalesce_kernel_arrays() { + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, None, None])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![Some(9), Some(2), None, None])); + let c: ArrayRef = + Arc::new(Int32Array::from(vec![Some(8), Some(8), Some(3), None])); + + let result = invoke( + vec![ + ColumnarValue::Array(a), + ColumnarValue::Array(b), + ColumnarValue::Array(c), + ], + 4, + ) + .into_array(4) + .unwrap(); + + let expected: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3), None])); + assert_eq!(&result, &expected); + } + + /// A trailing non-null scalar fills every remaining row, so the output has no nulls. + #[test] + fn coalesce_kernel_array_with_scalar_fallback() { + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3), None])); + + let result = invoke( + vec![ + ColumnarValue::Array(a), + ColumnarValue::Scalar(ScalarValue::Int32(Some(-1))), + ], + 4, + ) + .into_array(4) + .unwrap(); + + let expected: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), Some(-1), Some(3), Some(-1)])); + assert_eq!(&result, &expected); + assert_eq!(result.null_count(), 0); + } + + /// All-scalar input short-circuits to the first non-null scalar. + #[test] + fn coalesce_kernel_all_scalars() { + let result = invoke( + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(None)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("datafusion".into()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("unused".into()))), + ], + 1, + ); + match result { + ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => { + assert_eq!(s, "datafusion") + } + other => panic!("expected Utf8 scalar, got {other:?}"), + } + } + + /// All-null scalars produce a null scalar rather than an error. + #[test] + fn coalesce_kernel_all_null_scalars() { + let result = invoke( + vec![ + ColumnarValue::Scalar(ScalarValue::Int32(None)), + ColumnarValue::Scalar(ScalarValue::Int32(None)), + ], + 1, + ); + match result { + ColumnarValue::Scalar(v) => assert!(v.is_null()), + other => panic!("expected null scalar, got {other:?}"), + } + } + + /// All-null arrays produce an all-null array of the right type and length. + #[test] + fn coalesce_kernel_all_null_arrays() { + let a: ArrayRef = Arc::new(StringArray::from(vec![None as Option<&str>, None])); + let b: ArrayRef = Arc::new(StringArray::from(vec![None as Option<&str>, None])); + + let result = invoke(vec![ColumnarValue::Array(a), ColumnarValue::Array(b)], 2) + .into_array(2) + .unwrap(); + + assert_eq!(result.len(), 2); + assert_eq!(result.null_count(), 2); + assert_eq!(result.data_type(), &DataType::Utf8); + } + + #[test] + fn coalesce_kernel_rejects_empty_args() { + let err = CoalesceFunc::new() + .invoke_with_args(ScalarFunctionArgs { + args: vec![], + arg_fields: vec![], + number_rows: 1, + return_field: Field::new("f", DataType::Int32, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap_err(); + assert!(err.to_string().contains("It requires at least 1")); + } + + fn simplify(args: Vec) -> ExprSimplifyResult { + CoalesceFunc::new() + .simplify(args, &SimplifyContext::default()) + .expect("simplify failed") + } + + /// Non-volatile arguments keep the lazy `CASE` rewrite. + #[test] + fn simplify_non_volatile_expands_to_case() { + let result = simplify(vec![lit(1i64), lit(2i64)]); + match result { + ExprSimplifyResult::Simplified(Expr::Case(_)) => {} + other => panic!("expected a CASE expression, got {other:?}"), + } + } + + /// A volatile argument must NOT be expanded, because the `CASE` rewrite would + /// name it twice and evaluate it twice. See issue #25477. + #[test] + fn simplify_volatile_is_left_intact() { + let random = + datafusion_expr::expr::ScalarFunction::new_udf(crate::math::random(), vec![]); + let args = vec![Expr::ScalarFunction(random), lit(-1i64)]; + match simplify(args.clone()) { + ExprSimplifyResult::Original(original) => assert_eq!(original, args), + other => panic!("expected the original coalesce args, got {other:?}"), + } + } + + /// A volatile argument nested inside a larger expression is also caught. + #[test] + fn simplify_nested_volatile_is_left_intact() { + let random = Expr::ScalarFunction( + datafusion_expr::expr::ScalarFunction::new_udf(crate::math::random(), vec![]), + ); + let args = vec![random + lit(1.0f64), lit(-1.0f64)]; + match simplify(args.clone()) { + ExprSimplifyResult::Original(original) => assert_eq!(original, args), + other => panic!("expected the original coalesce args, got {other:?}"), + } + } + + /// A volatile *last* argument is named once (it becomes the `ELSE` branch), + /// so the rewrite is safe and must still happen -- otherwise the other + /// arguments needlessly lose their laziness. + #[test] + fn simplify_volatile_last_arg_still_expands_to_case() { + let random = Expr::ScalarFunction( + datafusion_expr::expr::ScalarFunction::new_udf(crate::math::random(), vec![]), + ); + match simplify(vec![lit(1.0f64), random]) { + ExprSimplifyResult::Simplified(Expr::Case(_)) => {} + other => panic!("expected a CASE expression, got {other:?}"), + } + } + + /// A single argument is still unwrapped, volatile or not (it is named once). + #[test] + fn simplify_single_volatile_arg_is_unwrapped() { + let random = Expr::ScalarFunction( + datafusion_expr::expr::ScalarFunction::new_udf(crate::math::random(), vec![]), + ); + match simplify(vec![random.clone()]) { + ExprSimplifyResult::Simplified(e) => assert_eq!(e, random), + other => panic!("expected the argument itself, got {other:?}"), + } + } +} diff --git a/datafusion/functions/src/core/nvl.rs b/datafusion/functions/src/core/nvl.rs index 1516f12ed9654..55f93af96523f 100644 --- a/datafusion/functions/src/core/nvl.rs +++ b/datafusion/functions/src/core/nvl.rs @@ -27,7 +27,7 @@ use datafusion_macros::user_doc; #[user_doc( doc_section(label = "Conditional Functions"), - description = "Returns _expression2_ if _expression1_ is NULL otherwise it returns _expression1_ and _expression2_ is not evaluated. This function can be used to substitute a default value for NULL values.", + description = "Returns _expression2_ if _expression1_ is NULL otherwise it returns _expression1_. _expression2_ is normally not evaluated, but when _expression1_ is volatile both arguments are evaluated, so that _expression1_ is evaluated exactly once. This function can be used to substitute a default value for NULL values.", syntax_example = "nvl(expression1, expression2)", sql_example = r#"```sql > select nvl(null, 'a'); diff --git a/datafusion/sqllogictest/test_files/coalesce.slt b/datafusion/sqllogictest/test_files/coalesce.slt index 43d53485692b7..766ff3dd10db6 100644 --- a/datafusion/sqllogictest/test_files/coalesce.slt +++ b/datafusion/sqllogictest/test_files/coalesce.slt @@ -458,3 +458,78 @@ world statement ok drop table ree_t; + +########## +# Volatile arguments must be evaluated exactly once +# https://github.com/apache/datafusion/issues/25477 +########## + +# `coalesce` on a volatile argument must NOT be rewritten to +# `CASE WHEN a IS NOT NULL THEN a ELSE b END`, because that names `a` twice and +# therefore draws it twice. It must stay a `coalesce` call so the runtime kernel +# evaluates each argument exactly once. +query TT +explain select coalesce(nullif(floor(random() * 2), 0), -1) as c; +---- +logical_plan +01)Projection: coalesce(nullif(floor(random() * Float64(2)), Float64(0)), Float64(-1)) AS c +02)--EmptyRelation: rows=1 +physical_plan +01)ProjectionExec: expr=[coalesce(nullif(floor(random() * 2), 0), -1) as c] +02)--PlaceholderRowExec + +statement ok +create table volatile_t as select * from generate_series(1, 100000) as t(i); + +# The fallback `-1` makes the result non-nullable, and with a single evaluation +# per row the result really never is null: every row is either the non-null +# volatile draw or the literal. Before the fix this both produced ~25% spurious +# NULLs and tripped the non-nullable assertion in Arrow. +query II +select count(*), count(coalesce(nullif(floor(random() * 2), 0), -1)) +from volatile_t; +---- +100000 100000 + +# Every produced value must come from one of the two operands. +query I +select count(*) +from ( + select coalesce(nullif(floor(random() * 2), 0), -1) as c from volatile_t +) +where c not in (1, -1); +---- +0 + +# `nvl` delegates to `coalesce`, so it inherits the fix. +query TT +explain select nvl(nullif(floor(random() * 2), 0), -1) as c; +---- +logical_plan +01)Projection: nvl(nullif(floor(random() * Float64(2)), Float64(0)), Float64(-1)) AS c +02)--EmptyRelation: rows=1 +physical_plan +01)ProjectionExec: expr=[nvl(nullif(floor(random() * 2), 0), -1) as c] +02)--PlaceholderRowExec + +query II +select count(*), count(nvl(nullif(floor(random() * 2), 0), -1)) +from volatile_t; +---- +100000 100000 + +statement ok +drop table volatile_t; + +# Non-volatile `coalesce` keeps the lazy CASE rewrite (see also +# select.slt, which asserts `coalesce(1, y/x)` never divides by zero). +query TT +explain select coalesce(column1, column2) as c from (values (1, 2), (null, 3)) t; +---- +logical_plan +01)Projection: CASE WHEN t.column1 IS NOT NULL THEN t.column1 ELSE t.column2 END AS c +02)--SubqueryAlias: t +03)----Values: (Int64(1), Int64(2)), (Int64(NULL), Int64(3)) +physical_plan +01)ProjectionExec: expr=[CASE WHEN column1@0 IS NOT NULL THEN column1@0 ELSE column2@1 END as c] +02)--DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index e11024449ac33..a49798fd419a7 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -577,3 +577,29 @@ Wire compatibility is directional: - A 55.0 reader must not consume an alias-preserving 56.0 MERGE payload. It ignores the unknown field but cannot preserve the qualifier required by the expressions, which can cause resolution failure or incorrect rebinding. +### A volatile `coalesce` or `nvl` argument is now evaluated eagerly + +`coalesce` and `nvl` are normally rewritten to `CASE WHEN a IS NOT NULL THEN a +ELSE b END`, which names every argument but the last one twice. For a volatile +argument those two mentions were two independent draws, so +`coalesce(nullif(floor(random() * 2), 0), -1)` could return `NULL` — a value a +single evaluation can never produce — or fail at run time with +`Column 'c' is declared as non-nullable but contains null values`. + +The rewrite is now skipped when any argument other than the last is volatile, +and the call is evaluated by a kernel that reads each argument exactly once. +That fixes the wrong results, but it also means the remaining arguments are +evaluated rather than skipped: + +```sql +-- previously returned rows, because `y / x` was never evaluated; +-- now raises `Divide by zero error` +SELECT coalesce(random(), y / x) FROM t; +``` + +Only calls that contain a volatile argument before the last one are affected. +A volatile *last* argument still takes the lazy rewrite, and non-volatile +`coalesce` and `nvl` are unchanged. + +To force conditional evaluation, rewrite using `CASE`, which has standardized +short-circuit semantics. diff --git a/docs/source/user-guide/sql/scalar_functions.md b/docs/source/user-guide/sql/scalar_functions.md index 11f50b84e3fd6..78ff6d399d448 100644 --- a/docs/source/user-guide/sql/scalar_functions.md +++ b/docs/source/user-guide/sql/scalar_functions.md @@ -962,7 +962,7 @@ trunc(numeric_expression[, decimal_places]) ### `coalesce` -Returns the first of its arguments that is not _null_. Returns _null_ if all arguments are _null_. This function is often used to substitute a default value for _null_ values. +Returns the first of its arguments that is not _null_. Returns _null_ if all arguments are _null_. Arguments after the first non-_null_ one are normally not evaluated, but when any argument other than the last is volatile every argument is evaluated, so that each is evaluated exactly once. This function is often used to substitute a default value for _null_ values. ```sql coalesce(expression1[, ..., expression_n]) @@ -1066,7 +1066,7 @@ nullif(expression1, expression2) ### `nvl` -Returns _expression2_ if _expression1_ is NULL otherwise it returns _expression1_ and _expression2_ is not evaluated. This function can be used to substitute a default value for NULL values. +Returns _expression2_ if _expression1_ is NULL otherwise it returns _expression1_. _expression2_ is normally not evaluated, but when _expression1_ is volatile both arguments are evaluated, so that _expression1_ is evaluated exactly once. This function can be used to substitute a default value for NULL values. ```sql nvl(expression1, expression2)