fix: keep decimal and temporal types in nvl and ifnull

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<String>,
 }
 
-/// 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<Vec<DataType>> {
+    let uniform = ScalarUDF::new_from_impl(CoalesceFunc {
+        signature: Signature::uniform(
+            2,
+            SUPPORTED_NVL_TYPES.to_vec(),
+            Volatility::Immutable,
+        ),
+    });
+    let fields: Vec<FieldRef> = 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<Vec<DataType>> {
+        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,29 @@ impl ScalarUDFImpl for NVLFunc {
         self.doc()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use arrow::datatypes::DataType;
+    use datafusion_common::DataFusionError;
+
+    #[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();
+            let DataFusionError::Execution(message) = error else {
+                panic!("Expected an execution error, got {error:?}");
+            };
+            assert_eq!(
+                message,
+                format!("nvl function requires 2 arguments, got {count}")
+            );
+        }
+        assert_eq!(
+            nvl.coerce_types(&[DataType::Int64, DataType::Int64])
+                .unwrap(),
+            vec![DataType::Int64, DataType::Int64]
+        );
+    }
+}

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);