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(¤t_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(¤t_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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
