You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/24 05:54:30 UTC
(datafusion) branch gh-readonly-queue/main/pr-23232-ad8a1476bf75324da71f13319debb6d5bf4edae9 created (now c8573822ed)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23232-ad8a1476bf75324da71f13319debb6d5bf4edae9
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at c8573822ed Align scalar UDF literal args with coerced return fields (#23232)
This branch includes the following new commits:
new c8573822ed Align scalar UDF literal args with coerced return fields (#23232)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Align scalar UDF literal args with coerced return fields (#23232)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23232-ad8a1476bf75324da71f13319debb6d5bf4edae9
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit c8573822ed9d48fac01e6d8a635ed0e10e091084
Author: Yin Li <kx...@student.bham.ac.uk>
AuthorDate: Mon Aug 24 05:54:13 2026 +0000
Align scalar UDF literal args with coerced return fields (#23232)
## Which issue does this PR close?
- Closes #19982.
## Rationale for this change
`ReturnFieldArgs` exposes both coerced `arg_fields` and optional scalar
literal arguments. A coercible UDF could receive an `Int16` field but an
original `Int64` literal, or lose the literal behind an implicit cast.
That makes return-field decisions inconsistent and forced UDF
implementations to compensate for planner internals.
## What changes are included in this PR?
- Materialize successful implicit casts of scalar-function literals
during type coercion, so literal values and `arg_fields` always have the
same type.
- Keep user-written `Cast` and `TryCast` expressions as expressions
rather than reporting them as literals.
- Preserve execution-time errors for invalid literal casts.
- Simplify return-field scalar extraction to accept only genuine
post-coercion literals.
- Update `round` to use the strengthened contract and document it on
`ReturnFieldArgs`.
- Add unit and integration regressions for coerced integers, typed
nulls, list literals, invalid values, explicit casts, expression
arguments, and `arrow_cast`'s required literal type argument.
This follows the approach discussed in #20012: materialize implicit
literal coercion in the type-coercion pass without changing
explicit-cast semantics.
## Are these changes tested?
- `cargo fmt --all --check`
- `cargo test -p datafusion --test user_defined_integration
test_return_field_args_scalar_argument_types_match_arg_fields` (1
passed)
- `cargo test -p datafusion --test core_integration
test_fn_arrow_cast_requires_literal_type_arg` (1 passed)
- `cargo test -p datafusion-expr scalar_arguments_` (2 passed)
- `cargo clippy --all-targets --all-features -- -D warnings`
- `git diff --check`
The validation process caught and fixed two semantic regressions before
push: explicit casts being folded into scalar literals, and `arrow_cast`
reporting its existing error during planning rather than collection.
## AI assistance
AI tooling assisted with implementation and test iteration. I reviewed
the final diff, the related #20012 design discussion, and the
scalar-function coercion/return-field flow end to end, and ran the
validations listed above.
---------
Signed-off-by: Kevin-Li-2025 <22...@qq.com>
Signed-off-by: Kevin-Li-2025 <kx...@student.bham.ac.uk>
Co-authored-by: Kevin-Li-2025 <22...@qq.com>
Co-authored-by: Jeffrey Vo <je...@gmail.com>
---
.../user_defined/user_defined_scalar_functions.rs | 95 +++++++++++++++++++++-
datafusion/expr/src/expr_schema.rs | 86 ++++++++++++++++++--
datafusion/functions/src/math/round.rs | 3 -
datafusion/optimizer/src/analyzer/type_coercion.rs | 88 ++++++++++++++++----
.../test_files/string_numeric_coercion.slt | 4 +
5 files changed, 248 insertions(+), 28 deletions(-)
diff --git a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
index a4d8af5e0c..4f8078542e 100644
--- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
+++ b/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;
@@ -2073,6 +2075,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<()> {
diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs
index 07fa6852a1..21c68089df 100644
--- a/datafusion/expr/src/expr_schema.rs
+++ b/datafusion/expr/src/expr_schema.rs
@@ -87,6 +87,26 @@ 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, _) => Some(
+ sv.cast_to(arg_field.data_type())
+ .unwrap_or_else(|_| sv.clone()),
+ ),
+ _ => None,
+ }
+}
+
impl ExprSchemable for Expr {
/// Returns the [arrow::datatypes::DataType] of the expression
/// based on [ExprSchema]
@@ -585,16 +605,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)
@@ -827,6 +843,60 @@ 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_preserve_invalid_values() {
+ let int16_field: FieldRef = Field::new("arg", DataType::Int16, true).into();
+ let int32_field: FieldRef = Field::new("arg", DataType::Int32, 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));
+ let out_of_i32_range = i64::from(i32::MAX) + 1;
+
+ 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),
+ Some(ScalarValue::Utf8(Some("not an integer".to_string())))
+ );
+ assert_eq!(
+ scalar_argument_for_field(&lit(out_of_i32_range), &int32_field),
+ Some(ScalarValue::Int64(Some(out_of_i32_range)))
+ );
+ }
+
#[test]
fn expr_schema_nullability() {
let expr = col("foo").eq(lit(1));
diff --git a/datafusion/functions/src/math/round.rs b/datafusion/functions/src/math/round.rs
index 10500810a5..120b07f4eb 100644
--- a/datafusion/functions/src/math/round.rs
+++ b/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)
diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs
index feadc0370b..2ca679bc51 100644
--- a/datafusion/optimizer/src/analyzer/type_coercion.rs
+++ b/datafusion/optimizer/src/analyzer/type_coercion.rs
@@ -813,8 +813,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),
)))
@@ -1147,23 +1150,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)
+}
+
fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result<Case> {
// Given expressions like:
//
@@ -2034,7 +2094,7 @@ mod test {
assert_analyzed_plan_eq!(
plan,
@r"
- Projection: TestScalarUDF(CAST(Int32(123) AS Float32))
+ Projection: TestScalarUDF(Float32(123)) AS TestScalarUDF(Int32(123))
EmptyRelation: rows=0
"
)
@@ -2071,7 +2131,7 @@ mod test {
assert_analyzed_plan_eq!(
plan,
@r"
- Projection: TestScalarUDF(CAST(Int64(10) AS Float32))
+ Projection: TestScalarUDF(Float32(10)) AS TestScalarUDF(Int64(10))
EmptyRelation: rows=0
"
)
@@ -2611,7 +2671,7 @@ mod test {
assert_analyzed_plan_eq!(
plan,
@r#"
- Projection: TestScalarUDF(a, Utf8("b"), CAST(Boolean(true) AS Utf8), CAST(Boolean(false) AS Utf8), CAST(Int32(13) AS Utf8))
+ Projection: TestScalarUDF(a, Utf8("b"), Utf8("true"), Utf8("false"), Utf8("13")) AS TestScalarUDF(a,Utf8("b"),Boolean(true),Boolean(false),Int32(13))
EmptyRelation: rows=0
"#
)
diff --git a/datafusion/sqllogictest/test_files/string_numeric_coercion.slt b/datafusion/sqllogictest/test_files/string_numeric_coercion.slt
index 1a3ef93f7e..196da416a2 100644
--- a/datafusion/sqllogictest/test_files/string_numeric_coercion.slt
+++ b/datafusion/sqllogictest/test_files/string_numeric_coercion.slt
@@ -573,6 +573,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
# -------------------------------------------------
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org