Jefffrey commented on code in PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#discussion_r3642373443
##########
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:
is this to maintain backwards compatibility?
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -2078,6 +2080,92 @@ 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 let Some(scalar) = args.scalar_arguments[0] {
+ assert_eq!(args.arg_fields[0].data_type(),
&scalar.data_type());
+ }
Review Comment:
```suggestion
assert_eq!(args.arg_fields[0].data_type(),
&args.scalar_arguments[0].unwrap().data_type());
```
nit: we already assert it `is_some()` above
##########
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(),
Review Comment:
i suppose this is an easier way to ensure proper type without a wider
refactor 🤔
##########
datafusion/core/tests/dataframe/dataframe_functions.rs:
##########
@@ -214,6 +214,25 @@ async fn test_fn_arrow_cast() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn test_fn_arrow_cast_requires_literal_type_arg() -> Result<()> {
Review Comment:
we can probably just put this as an SLT which is less verbose
##########
datafusion/functions/src/math/round.rs:
##########
@@ -242,8 +242,8 @@ 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.
+ // `scalar_arguments` uses the coerced argument type, so an integer
literal
+ // is Int32 here when the signature coerces it to Int32.
Review Comment:
we can just remove the note here if we fix this behaviour, no need to
explain it again (since `scalar_arguments` docstring should explain it)
##########
datafusion/expr/src/udf.rs:
##########
@@ -456,6 +456,10 @@ pub struct ReturnFieldArgs<'a> {
///
/// If the argument `i` is not a scalar, it will be None
///
+ /// When present, the scalar value has the same data type as the
corresponding
+ /// entry in [`Self::arg_fields`], including after implicit type coercion.
+ /// User-written `Cast` and `TryCast` expressions are not scalar arguments.
Review Comment:
```suggestion
```
this just seems confusing, can remove it
--
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]
