You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "Kevin-Li-2025 (via GitHub)" <gi...@apache.org> on 2026/06/29 06:59:50 UTC
[PR] Align scalar UDF literal args with coerced return fields [datafusion]
Kevin-Li-2025 opened a new pull request, #23232:
URL: https://github.com/apache/datafusion/pull/23232
## Which issue does this PR close?
- Closes #19982.
## Rationale for this change
`ReturnFieldArgs` exposes both coerced `arg_fields` and optional scalar literal arguments to scalar UDFs. For coerced literal calls such as `test_udf(1)` where the UDF signature coerces the argument to `Int16`, the field metadata is `Int16` but the scalar argument was not aligned with that field. On current `main`, the coerced expression can also appear as `Cast(Literal)`, causing `scalar_arguments` to be `None` even though the argument is still a literal.
## What changes are included in this PR?
- Build scalar UDF `ReturnFieldArgs.scalar_arguments` from literals after aligning them to the verified/coerced argument fields.
- Recognize literals wrapped in `Cast` / `TryCast` so coerced constants are still reported as scalar arguments.
- Add a regression test that asserts `arg_fields[0].data_type()` matches the scalar argument data type for a coercible scalar UDF.
## Are these changes tested?
- `cargo fmt --all --check`
- `CARGO_HOME=/Users/yinxiaogou/Documents/github/.cargo-home CARGO_TARGET_DIR=/Users/yinxiaogou/Documents/github/datafusion-target cargo test -p datafusion --test user_defined_integration user_defined::user_defined_scalar_functions::test_return_field_args_scalar_argument_types_match_arg_fields -- --exact`
I also ran `cargo test -p datafusion --test user_defined_integration user_defined::user_defined_scalar_functions`; the new test and 22 other tests passed, while 3 existing tests failed locally because `ARROW_TEST_DATA` / `PARQUET_TEST_DATA` submodule data was unavailable in this checkout.
--
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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "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(¤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: 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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#issuecomment-5391131918
thanks @yinli-systems
--
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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#issuecomment-5147906113
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/23232?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `97.34513%` with `3 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 80.86%. Comparing base ([`2fb472a`](https://app.codecov.io/gh/apache/datafusion/commit/2fb472a2173baeab28f6b498e0aa63ebd893ee77?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`a88feaf`](https://app.codecov.io/gh/apache/datafusion/commit/a88feaf6a18251e70da27c937102df99c372d894?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 31 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/23232?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [datafusion/optimizer/src/analyzer/type\_coercion.rs](https://app.codecov.io/gh/apache/datafusion/pull/23232?src=pr&el=tree&filepath=datafusion%2Foptimizer%2Fsrc%2Fanalyzer%2Ftype_coercion.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9vcHRpbWl6ZXIvc3JjL2FuYWx5emVyL3R5cGVfY29lcmNpb24ucnM=) | 94.73% | [0 Missing and 3 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/23232?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #23232 +/- ##
==========================================
+ Coverage 80.75% 80.86% +0.10%
==========================================
Files 1096 1101 +5
Lines 373605 375025 +1420
Branches 373605 375025 +1420
==========================================
+ Hits 301713 303250 +1537
+ Misses 53890 53675 -215
- Partials 18002 18100 +98
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/23232?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#issuecomment-4884571084
This seems to follow a similar approach to a previous PR that attempted this issue:
- https://github.com/apache/datafusion/pull/20012
There was some good discussion on there, perhaps can double check that since we found some difficulties with such an approach
--
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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "Kevin-Li-2025 (via GitHub)" <gi...@apache.org>.
Kevin-Li-2025 commented on PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#issuecomment-4899397169
Thanks for pointing me back to #20012. I went through that discussion and tightened this PR so it does not take the broad approach that caused the earlier test churn.
The latest commit narrows scalar literal recovery to the implicit-cast shape inserted by type coercion, `Cast(Literal -> expected arg type)`. It deliberately does not recursively unwrap user-written casts/try_casts, so cases such as `arrow_cast(1, cast('Utf8' as varchar))` still fail as before instead of treating the cast expression as a constant type argument. It also preserves the original scalar value when the planning-time scalar cast fails, so builtins that validate literal values themselves keep their existing error path instead of surfacing an early Arrow cast error.
Local checks run:
- `cargo test -p datafusion --test user_defined_integration test_return_field_args_scalar_argument_types_match_arg_fields`
- `cargo test -p datafusion --test core_integration dataframe::dataframe_functions::test_fn_arrow_cast_requires_literal_type_arg`
- `cargo run -p datafusion-examples --example builtin_functions -- date_time`
- `git diff --check`
New CI is now queued on head `053b1c0e8`.
--
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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey merged PR #23232:
URL: https://github.com/apache/datafusion/pull/23232
--
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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
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: 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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#issuecomment-5307984501
ill take a look at this soon 👍
--
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
Re: [PR] Align scalar UDF literal args with coerced return fields [datafusion]
Posted by "yinli-systems (via GitHub)" <gi...@apache.org>.
yinli-systems commented on PR #23232:
URL: https://github.com/apache/datafusion/pull/23232#issuecomment-5239790123
All inline feedback is addressed on the current head and the six review threads are now resolved. The full CI is green. @Jefffrey, could you please re-review when convenient?
--
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