Align scalar UDF literal args with coerced return fields

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -20,8 +20,8 @@ use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
 use arrow::array::{
-    Array, ArrayRef, Float32Array, Float64Array, Int32Array, RecordBatch, StringArray,
-    builder::BooleanBuilder, cast::AsArray,
+    Array, ArrayRef, Float32Array, Float64Array, Int16Array, Int32Array, RecordBatch,
+    StringArray, builder::BooleanBuilder, cast::AsArray,
 };
 use arrow::array::{Int8Array, UInt64Array, as_string_array, create_array, record_batch};
 use arrow::compute::kernels::numeric::add;
@@ -40,13 +40,15 @@ use datafusion_common::{
     DFSchema, DataFusionError, Result, ScalarValue, assert_batches_eq,
     assert_batches_sorted_eq, assert_contains, exec_datafusion_err, exec_err,
     not_impl_err, plan_err,
+    types::{NativeType, logical_int16},
 };
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext};
 use datafusion_expr::{
     Accumulator, ColumnarValue, CreateFunction, CreateFunctionBody, LogicalPlanBuilder,
     OperateFunctionArg, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl,
-    Signature, Volatility, lit_with_metadata,
+    Signature, TypeSignatureClass, Volatility, lit_with_metadata,
 };
+use datafusion_expr_common::signature::Coercion;
 use datafusion_expr_common::signature::TypeSignature;
 use datafusion_functions_nested::range::range_udf;
 use parking_lot::Mutex;
@@ -2078,6 +2080,93 @@ AS t(string, extension)
     Ok(())
 }
 
+/// https://github.com/apache/datafusion/issues/19982
+#[tokio::test]
+async fn test_return_field_args_scalar_argument_types_match_arg_fields() -> Result<()> {
+    #[derive(Debug, PartialEq, Eq, Hash)]
+    struct TestUdf {
+        name: &'static str,
+        expect_literal: bool,
+        signature: Signature,
+    }
+
+    impl TestUdf {
+        fn new(name: &'static str, expect_literal: bool) -> Self {
+            Self {
+                name,
+                expect_literal,
+                signature: Signature::coercible(
+                    vec![Coercion::new_implicit(
+                        TypeSignatureClass::Native(logical_int16()),
+                        vec![TypeSignatureClass::Numeric],
+                        NativeType::Int16,
+                    )],
+                    Volatility::Immutable,
+                ),
+            }
+        }
+    }
+
+    impl ScalarUDFImpl for TestUdf {
+        fn name(&self) -> &str {
+            self.name
+        }
+
+        fn signature(&self) -> &Signature {
+            &self.signature
+        }
+
+        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+            unreachable!("return_field_from_args is implemented")
+        }
+
+        fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
+            assert_eq!(args.arg_fields.len(), 1);
+            assert_eq!(args.scalar_arguments.len(), 1);
+            assert_eq!(
+                args.scalar_arguments[0].is_some(),
+                self.expect_literal,
+                "unexpected scalar argument for {}",
+                self.name
+            );
+            if self.expect_literal {
+                let scalar = args.scalar_arguments[0].unwrap();
+                assert_eq!(args.arg_fields[0].data_type(), &scalar.data_type());
+            }
+            Ok(
+                Field::new(self.name(), args.arg_fields[0].data_type().clone(), true)
+                    .into(),
+            )
+        }
+
+        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+            assert_eq!(args.args[0].data_type(), DataType::Int16);
+            Ok(args.args[0].clone())
+        }
+    }
+
+    let ctx = SessionContext::new();
+    let coerced_literal_udf: ScalarUDF = TestUdf::new("coerced_literal_udf", true).into();
+    let expression_arg_udf: ScalarUDF = TestUdf::new("expression_arg_udf", false).into();
+    ctx.register_udf(coerced_literal_udf);
+    ctx.register_udf(expression_arg_udf.clone());
+
+    ctx.sql("select coerced_literal_udf(1), coerced_literal_udf(NULL)")
+        .await?
+        .collect()
+        .await?;
+    let batch = RecordBatch::try_new(
+        Arc::new(Schema::new(vec![Field::new("a", DataType::Int16, false)])),
+        vec![Arc::new(Int16Array::from(vec![1]))],
+    )?;
+    ctx.read_batch(batch)?
+        .select(vec![expression_arg_udf.call(vec![col("a")])])?
+        .collect()
+        .await?;
+
+    Ok(())
+}
+
 /// https://github.com/apache/datafusion/issues/17422
 #[tokio::test]
 async fn test_extension_metadata_preserve_in_subquery() -> Result<()> {

datafusion/expr/src/expr_schema.rs

@@ -87,6 +87,23 @@ fn cast_output_field(
     Arc::new(f)
 }
 
+fn scalar_arguments_for_fields(
+    args: &[Expr],
+    arg_fields: &[FieldRef],
+) -> Vec<Option<ScalarValue>> {
+    args.iter()
+        .zip(arg_fields)
+        .map(|(expr, field)| scalar_argument_for_field(expr, field))
+        .collect()
+}
+
+fn scalar_argument_for_field(expr: &Expr, arg_field: &FieldRef) -> Option<ScalarValue> {
+    match expr {
+        Expr::Literal(sv, _) => sv.cast_to(arg_field.data_type()).ok(),
+        _ => None,
+    }
+}
+
 impl ExprSchemable for Expr {
     /// Returns the [arrow::datatypes::DataType] of the expression
     /// based on [ExprSchema]
@@ -580,16 +597,12 @@ impl ExprSchemable for Expr {
                     .collect::<Result<Vec<_>>>()?;
                 let new_fields = verify_function_arguments(func.as_ref(), &fields)?;
 
-                let arguments = args
-                    .iter()
-                    .map(|e| match e {
-                        Expr::Literal(sv, _) => Some(sv),
-                        _ => None,
-                    })
-                    .collect::<Vec<_>>();
+                let arguments = scalar_arguments_for_fields(args, &new_fields);
+                let argument_refs =
+                    arguments.iter().map(Option::as_ref).collect::<Vec<_>>();
                 let args = ReturnFieldArgs {
                     arg_fields: &new_fields,
-                    scalar_arguments: &arguments,
+                    scalar_arguments: &argument_refs,
                 };
 
                 func.return_field_from_args(args)
@@ -807,6 +820,54 @@ mod tests {
         }};
     }
 
+    #[test]
+    fn scalar_arguments_match_coerced_fields() {
+        let int16_field: FieldRef = Field::new("arg", DataType::Int16, true).into();
+
+        assert_eq!(
+            scalar_argument_for_field(&lit(1_i64), &int16_field),
+            Some(ScalarValue::Int16(Some(1)))
+        );
+        assert_eq!(
+            scalar_argument_for_field(&lit(ScalarValue::Null), &int16_field),
+            Some(ScalarValue::Int16(None))
+        );
+
+        let int32_list = ScalarValue::List(ScalarValue::new_list(
+            &[ScalarValue::Int32(Some(1))],
+            &DataType::Int32,
+            true,
+        ));
+        let int64_list_type = DataType::new_list(DataType::Int64, true);
+        let int64_list_field: FieldRef =
+            Field::new("arg", int64_list_type.clone(), true).into();
+        assert_eq!(
+            scalar_argument_for_field(&lit(int32_list.clone()), &int64_list_field),
+            Some(int32_list.cast_to(&int64_list_type).unwrap())
+        );
+    }
+
+    #[test]
+    fn scalar_arguments_exclude_expression_casts_and_invalid_values() {
+        let int16_field: FieldRef = Field::new("arg", DataType::Int16, true).into();
+        let explicit_cast = Expr::Cast(Cast::new(Box::new(lit(1_i64)), DataType::Int16));
+        let explicit_try_cast =
+            Expr::TryCast(TryCast::new(Box::new(lit(1_i64)), DataType::Int16));
+
+        assert_eq!(
+            scalar_argument_for_field(&explicit_cast, &int16_field),
+            None
+        );
+        assert_eq!(
+            scalar_argument_for_field(&explicit_try_cast, &int16_field),
+            None
+        );
+        assert_eq!(
+            scalar_argument_for_field(&lit("not an integer"), &int16_field),
+            None
+        );
+    }
+
     #[test]
     fn expr_schema_nullability() {
         let expr = col("foo").eq(lit(1));

datafusion/functions/src/math/round.rs

@@ -241,9 +241,6 @@ impl ScalarUDFImpl for RoundFunc {
 
         // If decimal_places is a scalar literal, we can incorporate it into the output type
         // (scale reduction). Otherwise, keep the input scale as we can't pick a per-row scale.
-        //
-        // Note: `scalar_arguments` contains the original literal values (pre-coercion), so
-        // integer literals may appear as Int64 even though the signature coerces them to Int32.
         let decimal_places: Option<i32> = match args.scalar_arguments.get(1) {
             None => Some(0),    // No dp argument means default to 0
             Some(None) => None, // dp is not a literal (e.g. column)

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -745,8 +745,11 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
                 Ok(Transformed::yes(Expr::Case(case)))
             }
             Expr::ScalarFunction(ScalarFunction { func, args }) => {
-                let new_expr =
-                    coerce_arguments_for_signature(args, self.schema, func.as_ref())?;
+                let new_expr = coerce_scalar_function_arguments_for_signature(
+                    args,
+                    self.schema,
+                    func.as_ref(),
+                )?;
                 Ok(Transformed::yes(Expr::ScalarFunction(
                     ScalarFunction::new_udf(func, new_expr),
                 )))
@@ -1065,23 +1068,80 @@ fn coerce_arguments_for_signature<F: UDFCoercionExt>(
     schema: &DFSchema,
     func: &F,
 ) -> Result<Vec<Expr>> {
-    let current_fields = expressions
-        .iter()
-        .map(|e| e.to_field(schema).map(|(_, f)| f))
-        .collect::<Result<Vec<_>>>()?;
+    let coerced_types = coerced_argument_types(&expressions, schema, func)?;
 
-    let coerced_types = fields_with_udf(&current_fields, func)?
+    expressions
         .into_iter()
-        .map(|f| f.data_type().clone())
-        .collect::<Vec<_>>();
+        .zip(coerced_types)
+        .map(|(expr, data_type)| expr.cast_to(&data_type, schema))
+        .collect()
+}
+
+/// Coerces scalar function arguments while materializing successful implicit
+/// casts of literals. This preserves the literal for subsequent calls to
+/// `return_field_from_args` without treating user-written `Cast` or `TryCast`
+/// expressions as scalar arguments.
+fn coerce_scalar_function_arguments_for_signature<F: UDFCoercionExt>(
+    expressions: Vec<Expr>,
+    schema: &DFSchema,
+    func: &F,
+) -> Result<Vec<Expr>> {
+    let coerced_types = coerced_argument_types(&expressions, schema, func)?;
 
     expressions
         .into_iter()
-        .enumerate()
-        .map(|(i, expr)| expr.cast_to(&coerced_types[i], schema))
+        .zip(coerced_types)
+        .map(|(expr, data_type)| {
+            coerce_scalar_function_argument(expr, &data_type, schema)
+        })
         .collect()
 }
 
+fn coerced_argument_types<F: UDFCoercionExt>(
+    expressions: &[Expr],
+    schema: &DFSchema,
+    func: &F,
+) -> Result<Vec<DataType>> {
+    let current_fields = expressions
+        .iter()
+        .map(|e| e.to_field(schema).map(|(_, f)| f))
+        .collect::<Result<Vec<_>>>()?;
+
+    fields_with_udf(&current_fields, func).map(|fields| {
+        fields
+            .into_iter()
+            .map(|field| field.data_type().clone())
+            .collect()
+    })
+}
+
+fn coerce_scalar_function_argument(
+    expr: Expr,
+    data_type: &DataType,
+    schema: &DFSchema,
+) -> Result<Expr> {
+    if matches!(&expr, Expr::Cast(_) | Expr::TryCast(_))
+        && expr.get_type(schema)? == *data_type
+    {
+        return Ok(expr);
+    }
+
+    let Expr::Literal(value, metadata) = expr else {
+        return expr.cast_to(data_type, schema);
+    };
+
+    if value.data_type() != *data_type
+        && let Ok(value) = value.cast_to(data_type)
+    {
+        return Ok(Expr::Literal(value, metadata));
+    }
+
+    // A failed value cast remains an expression cast so execution produces the
+    // same error as before. Since it is no longer a literal at the coerced type,
+    // it is reported as `None` in `ReturnFieldArgs::scalar_arguments`.
+    Expr::Literal(value, metadata).cast_to(data_type, schema)
+}
+
 fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result<Case> {
     // Given expressions like:
     //

datafusion/sqllogictest/test_files/string_numeric_coercion.slt

@@ -553,6 +553,10 @@ SELECT arrow_cast('hello', 'RunEndEncoded("run_ends": non-null Int32, "values":
 ----
 true
 
+# The type argument to arrow_cast must remain a literal after planning.
+statement error arrow_cast requires its second argument
+SELECT arrow_cast(1, cast('Utf8' as varchar));
+
 # -------------------------------------------------
 # Cleanup
 # -------------------------------------------------