diff --git a/datafusion/functions/src/core/nvl.rs b/datafusion/functions/src/core/nvl.rs index 1516f12ed9654..393de99221306 100644 --- a/datafusion/functions/src/core/nvl.rs +++ b/datafusion/functions/src/core/nvl.rs @@ -15,12 +15,15 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use crate::core::coalesce::CoalesceFunc; -use arrow::datatypes::{DataType, FieldRef}; -use datafusion_common::Result; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::{Result, utils::take_function_args}; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; +use datafusion_expr::type_coercion::functions::fields_with_udf; use datafusion_expr::{ - ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, + ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; use datafusion_macros::user_doc; @@ -59,9 +62,9 @@ pub struct NVLFunc { aliases: Vec, } -/// Currently supported types by the nvl/ifnull function. -/// The order of these types correspond to the order on which coercion applies -/// This should thus be from least informative to most informative +/// Argument types that keep nvl's original coercion, ordered from least to +/// most informative. Other argument types are coerced the way coalesce does, +/// falling back to the original coercion when coalesce cannot unify them. static SUPPORTED_NVL_TYPES: &[DataType] = &[ DataType::Boolean, DataType::UInt8, @@ -79,6 +82,26 @@ static SUPPORTED_NVL_TYPES: &[DataType] = &[ DataType::LargeUtf8, ]; +/// nvl's original coercion: both arguments become the first type in +/// [`SUPPORTED_NVL_TYPES`] that both can be coerced to. +fn uniform_coercion(arg_types: &[DataType]) -> Result> { + let uniform = ScalarUDF::new_from_impl(CoalesceFunc { + signature: Signature::uniform( + 2, + SUPPORTED_NVL_TYPES.to_vec(), + Volatility::Immutable, + ), + }); + let fields: Vec = arg_types + .iter() + .map(|t| Arc::new(Field::new("", t.clone(), true))) + .collect(); + Ok(fields_with_udf(&fields, &uniform)? + .iter() + .map(|f| f.data_type().clone()) + .collect()) +} + impl Default for NVLFunc { fn default() -> Self { Self::new() @@ -89,11 +112,7 @@ impl NVLFunc { pub fn new() -> Self { Self { coalesce: CoalesceFunc { - signature: Signature::uniform( - 2, - SUPPORTED_NVL_TYPES.to_vec(), - Volatility::Immutable, - ), + signature: Signature::user_defined(Volatility::Immutable), }, aliases: vec![String::from("ifnull")], } @@ -140,6 +159,21 @@ impl ScalarUDFImpl for NVLFunc { self.coalesce.short_circuits() } + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [_, _] = take_function_args(self.name(), arg_types)?; + + // Preserve accepted pairs: (Int8, UInt64) stays Float32, and + // (Int64, Boolean) stays Utf8View instead of being rejected by COALESCE. + let listed = arg_types + .iter() + .all(|t| t.is_null() || SUPPORTED_NVL_TYPES.contains(t)); + if !listed && let Ok(coerced) = self.coalesce.coerce_types(arg_types) { + return Ok(coerced); + } + + uniform_coercion(arg_types) + } + fn aliases(&self) -> &[String] { &self.aliases } @@ -148,3 +182,27 @@ impl ScalarUDFImpl for NVLFunc { self.doc() } } + +#[cfg(test)] +mod tests { + use arrow::datatypes::DataType; + + #[test] + fn rejects_wrong_argument_counts_through_public_udf() { + let nvl = crate::core::nvl(); + for count in [0, 1, 3, 4, 8] { + let error = nvl.coerce_types(&vec![DataType::Int64; count]).unwrap_err(); + assert_eq!( + error.strip_backtrace(), + format!( + "Execution error: nvl function requires 2 arguments, got {count}" + ) + ); + } + assert_eq!( + nvl.coerce_types(&[DataType::Int64, DataType::Int64]) + .unwrap(), + vec![DataType::Int64, DataType::Int64] + ); + } +} diff --git a/datafusion/sqllogictest/test_files/nvl.slt b/datafusion/sqllogictest/test_files/nvl.slt index a6637cadf3b3d..27165cb618137 100644 --- a/datafusion/sqllogictest/test_files/nvl.slt +++ b/datafusion/sqllogictest/test_files/nvl.slt @@ -183,3 +183,79 @@ physical_plan 01)ProjectionExec: expr=[CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE 9999 END as nvl(test.int_field,Int64(9999))] 02)--ProjectionExec: expr=[CAST(int_field@0 AS Int64) as __common_expr_1] 03)----DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +CREATE TABLE nvl_type_coercion AS SELECT + CAST('12345678901234567890123456789012345678' AS DECIMAL(38,0)) AS big, + CAST(1.23 AS DECIMAL(5,2)) AS d, + TIMESTAMP '2024-01-01 10:00:00' AS ts, + DATE '2024-01-01' AS dt +; + +query TR +SELECT + arrow_typeof(nvl(big, 0)), + nvl(big, 0) +FROM nvl_type_coercion; +---- +Decimal128(38, 0) 12345678901234567890123456789012345678 + +query TTT +SELECT + arrow_typeof(nvl(d, d)), + arrow_typeof(ifnull(d, d)), + arrow_typeof(coalesce(d, d)) +FROM nvl_type_coercion; +---- +Decimal128(5, 2) Decimal128(5, 2) Decimal128(5, 2) + +query TTT +SELECT + arrow_typeof(nvl(ts, ts)), + arrow_typeof(ifnull(ts, ts)), + arrow_typeof(coalesce(ts, ts)) +FROM nvl_type_coercion; +---- +Timestamp(ns) Timestamp(ns) Timestamp(ns) + +query TT +SELECT + arrow_typeof(nvl(dt, dt)), + arrow_typeof(coalesce(dt, dt)) +FROM nvl_type_coercion; +---- +Date32 Date32 + +query P +SELECT ifnull(ts, ts) + INTERVAL '1 hour' FROM nvl_type_coercion; +---- +2024-01-01T11:00:00 + +# Types the original signature listed keep its coercion, and types coalesce +# cannot unify fall back to it. +query TT +SELECT + arrow_typeof(nvl(1, false)), + arrow_typeof(nvl(ts, '2024-01-01T00:00:00')) +FROM nvl_type_coercion; +---- +Utf8View Utf8View + +# Both names use the shared fixed-arity validation. +query error nvl function requires 2 arguments, got 0 +SELECT nvl(); + +query error nvl function requires 2 arguments, got 1 +SELECT nvl(1); + +query error nvl function requires 2 arguments, got 3 +SELECT nvl(1, 2, 3); + +query error nvl function requires 2 arguments, got 0 +SELECT ifnull(); + +query error nvl function requires 2 arguments, got 1 +SELECT ifnull(1); + +query error nvl function requires 2 arguments, got 3 +SELECT ifnull(1, 2, 3);