Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]

["yinli-systems (via GitHub)" <gi...@apache.org>]

yinli-systems commented on code in PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#discussion_r3749239040


##########
datafusion/optimizer/src/analyzer/type_coercion.rs:
##########
@@ -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)

Review Comment:
   Yes. This fallback intentionally preserves the previous error timing and message path when an implicit literal cast is invalid: the value remains an `Expr::Cast`, so the existing UDF validation still happens during collection/execution instead of being replaced by a new planning-time Arrow cast error. Because it is not a literal at the coerced type, `ReturnFieldArgs::scalar_arguments` correctly reports `None`. The invalid-literal and `arrow_cast` regressions cover this behavior.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org