Support compute return types from argument values (not just their DataTypes) #8985

yyy1000 · 2024-01-24T19:20:44Z

Which issue does this PR close?

Rationale for this change

This PR will extend ScalarFunctionImpl, so that when users want to create a ScalarUDF by implementing ScalarUDFImpl, they can implement function return_type_from_exprs in trait ScalarUDFImpl. It also provides a default implementation that wil call return_type.

What changes are included in this PR?

Extending the trait ScalarUDFImpl by adding a new fn and default implementation
An example that shows how to resolve the problem in UDF: Can I get a return type from a constant parameter in args? #7657

Are these changes tested?

Yes, by adding an example and I tested it.

Weijun-H · 2024-01-26T03:50:06Z

Thanks @yyy1000 👍 , could you check the output schema in your test like main...JasonLi-cn:arrow-datafusion:feature/refactor_udf

datafusion/expr/src/udf.rs

yyy1000 · 2024-01-26T05:29:46Z

Thanks @yyy1000 👍 , could you check the output schema in your test like main...JasonLi-cn:arrow-datafusion:feature/refactor_udf

@Weijun-H Thanks for your review!
Aha, it seems that my current implementation still use return_type to get the schema.
Maybe here changing return_type to return_type_from_exprs would be a good idea. But a problem
is here the schema has type S: ExprSchema while return_type_from_exprs can't accept a generic type parameter S because it's in a trait. 🤔 What do you think?
https://github.com/apache/arrow-datafusion/blob/4d02cc0114908d4f805b2323f20751b1f6d9c2f4/datafusion/expr/src/expr_schema.rs#L106-L108

yyy1000 · 2024-01-26T05:33:49Z

datafusion/expr/src/udf.rs

+    fn return_type_from_exprs(
+        &self,
+        arg_exprs: &[Expr],
+        schema: &DFSchema,
+    ) -> Result<DataType> {


Suggested change

fn return_type_from_exprs(

&self,

arg_exprs: &[Expr],

schema: &DFSchema,

) -> Result<DataType> {

fn return_type_from_exprs<S: ExprSchema>(

&self,

arg_exprs: &[Expr],

schema: S,

) -> Result<DataType> {

@Weijun-H If I'd like to make a change like this, there's an error 🥲

for a trait to be "object safe" it needs to allow building a vtable to allow the call to be resolvable dynamically;

I am unfamiliar with this part. It seems you need to extend ScalarUDFImpl to allow ExprSchema.

Maybe here changing return_type to return_type_from_exprs would be a good idea. But a problem
is here the schema has type S: ExprSchema while return_type_from_exprs can't accept a generic type parameter S because it's in a trait. 🤔 What do you think?

Maybe you can use a trait object like &dyn ExprSchema 🤔

I tried &dyn ExprSchema before, there will be an issue when calling get_type, like below. 🥲. I also tried the fix method recommended by the complier(adding +?Sized but failed)

@alamb Sorry for the delay.
I think the only problem that blocks me is the argument type of the function. I had some issue when passing schema as a param.
I can also search for some ways to do it.🤔

Cool -- thanks -- I'll try and find some time to play around with it but it may be a while

@yyy1000 the following changes worked for me. If anyone else have a better solution please suggest it. I am still a newbie in rust but trying to help to learn.

expr_schema.rs

return_types_udf.rs

Those look like great changes to me @brayanjuls -- 👌

What do you think @yyy1000 ?

Thanks for help! @alamb @brayanjuls
However, I have tried this method before and it will not work.
I will show you the reason on my side. :)

alamb

Thank you @yyy1000 -- this looks like a great start. Thank you for the help with reviewing @Weijun-H

alamb · 2024-01-26T21:15:22Z

datafusion-examples/examples/return_types_udf.rs

+        &self.signature
+    }
+    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Int32)


This is pretty confusing I think -- as it seems inconsistent with the return_type_from_exprs

alamb · 2024-01-26T21:21:17Z

datafusion-examples/examples/return_types_udf.rs

+}
+
+#[tokio::main]
+async fn main() -> Result<()> {


I think this example is missing actually using the function in a query / dataframe. As @Weijun-H pointed out the logic added to ScalarUDFImpl doesn't seem to be connected anywhere else 🤔

What I think the example needs to do is someething like

Create a ScalarUDF

register the function with a SessionContext

Run a query that uses that function (ideally both with SQL and dataframe APIs)

So for example, a good example function might be a function that takes a string argument select my_cast(<arg>, 'string') that converts the argument based on the value of the string

Then for example run queries like

select my_cast(c1, 'i32'), arrow_typeof(my_cast(c1, 'i32')); -- returns value and DataType::Int32 select my_cast(c1, 'i64'), arrow_typeof(my_cast(c1, 'i64')); -- returns value and DataType::Int64

Does that make sense?

Yeah, I realized that it missing using the function.
I tried to make changes in https://github.com/apache/arrow-datafusion/blob/4d02cc0114908d4f805b2323f20751b1f6d9c2f4/datafusion/expr/src/expr_schema.rs#L106-L108
to replace return_type with return_type_from_exprs. But passing schema as a param is a problem cause the type of schema is a generic type S. 🤔

alamb · 2024-02-05T15:40:14Z

What is the status of this PR -- are you waiting for more feedback or help @yyy1000 ?

yyy1000 · 2024-02-09T03:34:11Z

datafusion/expr/src/expr_schema.rs

@@ -136,7 +136,7 @@ impl ExprSchemable for Expr {
                        fun.return_type(&arg_data_types)
                    }
                    ScalarFunctionDefinition::UDF(fun) => {
-                        Ok(fun.return_type(&arg_data_types)?)
+                        fun.return_type_from_exprs(args, schema)


@alamb @brayanjuls . It's here that it gets the return type of ScalarUDF, and for UDF, it should be changed to this.
However, this can't work because schema needs to be Sized. Here is the error from my complier.

universalmind303 · 2024-02-09T04:18:58Z

datafusion/expr/src/udf.rs

@@ -249,6 +261,22 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
    /// the arguments
    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;

+    /// What [`DataType`] will be returned by this function, given the types of
+    /// the expr arguments
+    fn return_type_from_exprs(


Since we need to use the ExprSchema here, but there's issues going from <S: ExprSchema> to &dyn ExprSchema inside the trait, what if we moved up the default implementation out of the trait and into ScalarUDF?

we could change the trait impl to something like this

pub trait ScalarUDFImpl: Debug + Send + Sync { /// What [`DataType`] will be returned by this function, given the types of /// the expr arguments fn return_type_from_exprs( &self, arg_exprs: &[Expr], schema: &dyn ExprSchema, ) -> Option<Result<DataType>> { // The default implementation returns None // so that people don't have to implement `return_type_from_exprs` if they dont want to None } }

then change the ScalarUDF impl to this

impl ScalarUDF /// The datatype this function returns given the input argument input types. /// This function is used when the input arguments are [`Expr`]s. /// See [`ScalarUDFImpl::return_type_from_exprs`] for more details. pub fn return_type_from_exprs<S: ExprSchema>( &self, args: &[Expr], schema: &S, ) -> Result<DataType> { // If the implementation provides a return_type_from_exprs, use it if let Some(return_type) = self.inner.return_type_from_exprs(args, schema) { return_type // Otherwise, use the return_type function } else { let arg_types = args .iter() .map(|arg| arg.get_type(schema)) .collect::<Result<Vec<_>>>()?; self.return_type(&arg_types) } } }

this way we don't need to constrain the ExprSchemable functions to ?Sized, and we can update the get_type function to use the return_type_from_exprs without any compile time errors.

ScalarFunctionDefinition::UDF(fun) => { Ok(fun.return_type_from_exprs(&args, schema)?) }

and it still makes return_type_from_exprs an opt-in method.

It does make it very slightly less ergonomic as end users now need to wrap their body in an Option, but overall i think it's a decent compromise.

This works well on my side. Thanks!

Another question for me is, how can user implement return_type_from_exprs when using schema. For example, arg_exprs.get(take_idx).unwrap().get_type(schema) will lead an error

the size for values of type dyn ExprSchema cannot be known at compilation time
the trait Sized is not implemented for dyn ExprSchema

Hmm yeah @yyy1000 you still run into the same error then.

I'm wondering if it'd be easiest to just change the type signature on ExprSchemable to be generic over the trait instead of the functions.

pub trait ExprSchemable<S: ExprSchema> { /// given a schema, return the type of the expr fn get_type(&self, schema: &S) -> Result<DataType>; /// given a schema, return the nullability of the expr fn nullable(&self, input_schema: &S) -> Result<bool>; /// given a schema, return the expr's optional metadata fn metadata(&self, schema: &S) -> Result<HashMap<String, String>>; /// convert to a field with respect to a schema fn to_field(&self, input_schema: &DFSchema) -> Result<DFField>; /// cast to a type with respect to a schema fn cast_to(self, cast_to_type: &DataType, schema: &S) -> Result<Expr>; } impl ExprSchemable<DFSchema> for Expr { //... }

then the trait can just go back to the original implementation you had using &DFSchema

fn return_type_from_exprs( &self, arg_exprs: &[Expr], schema: &DFSchema, ) -> Result<DataType> { let arg_types = arg_exprs .iter() .map(|e| e.get_type(schema)) .collect::<Result<Vec<_>>>()?; self.return_type(&arg_types) }

I tried this locally and was able to get things to compile locally, and was able to implement a udf using the trait.

It does make it a little less flexible as it's expecting a DFSchema, but i think thats ok?

I think the only other approach would be to make ScalarUDFImpl dynamic over <S: ExprSchema>, but I feel like that's much less ideal than just using a concrete type.

Thanks for your help! @universalmind303
It looks good and I think I can try it to see. :)

Update is I changed the signature to take &dyn ExprSchema which seems to have worked just fine

yyy1000 · 2024-02-11T05:55:47Z

datafusion/physical-expr/src/udf.rs

@@ -38,7 +38,7 @@ pub fn create_physical_expr(
        fun.name(),
        fun.fun(),
        input_phy_exprs.to_vec(),
-        fun.return_type(&input_exprs_types)?,
+        fun.return_type_from_exprs(&input_phy_exprs, input_schema)?,


Thanks @universalmind303 's method helped a lot and it's getting closer!
I think maybe the last issue is here how can we convert input_phy_exprs of type &[Arc] to &[Expr]. 🤔

It isn't possible in general to go from physical expressions back to Exprs.

I recommend adding a return_type argument to this function and thus require the caller to pass it in.

I believe the caller of udf::create_physical_expr has this information:

/// Create a physical expression of the UDF. /// /// Note that the return type must match the type reported by fun pub fn create_physical_expr( fun: &ScalarUDF, input_phy_exprs: &[Arc<dyn PhysicalExpr>], return_type: DataType, ) -> Result<Arc<dyn PhysicalExpr>> {

alamb

Thank you for this PR @yyy1000 and for the suggestions @universalmind303 and @Weijun-H

I think this is a pretty important strategic feature, so I want to help. Here is a PR with changes to this one that I think solves the last little bit: yyy1000#1

Let me know what you think

alamb · 2024-02-12T16:59:55Z

datafusion/physical-expr/src/udf.rs

@@ -38,7 +38,7 @@ pub fn create_physical_expr(
        fun.name(),
        fun.fun(),
        input_phy_exprs.to_vec(),
-        fun.return_type(&input_exprs_types)?,
+        fun.return_type_from_exprs(&input_phy_exprs, input_schema)?,


It isn't possible in general to go from physical expressions back to Exprs.

I recommend adding a return_type argument to this function and thus require the caller to pass it in.

I believe the caller of udf::create_physical_expr has this information:

/// Create a physical expression of the UDF. /// /// Note that the return type must match the type reported by fun pub fn create_physical_expr( fun: &ScalarUDF, input_phy_exprs: &[Arc<dyn PhysicalExpr>], return_type: DataType, ) -> Result<Arc<dyn PhysicalExpr>> {

* Improve return_type_from_args * Rework example * Update datafusion/core/tests/user_defined/user_defined_scalar_functions.rs --------- Co-authored-by: Junhao Liu <junhaoliu2023@gmail.com>

Co-authored-by: Alex Huang <huangweijun1001@gmail.com>

alamb · 2024-02-12T22:02:46Z

I am going to try and fix the CI checks here

alamb · 2024-02-12T22:30:47Z

datafusion/expr/src/expr_schema.rs

 }

-impl ExprSchemable for Expr {
+impl ExprSchemable<DFSchema> for Expr {


I think this means the only way to use Expr::get_type is now if you have a DFSchema with requires bringing in the entire core crate

I think the reason this trait was the way it was orginally was to allow things like dask-sql to use just the logical expr parts.

I am looking into this more carefully

alamb · 2024-02-12T22:44:30Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -494,6 +495,127 @@ async fn test_user_defined_functions_zero_argument() -> Result<()> {
    Ok(())
 }

+#[derive(Debug)]
+struct TakeUDF {


Here is an example of the feature working

alamb · 2024-02-12T22:45:47Z

datafusion/expr/src/expr_schema.rs

@@ -37,19 +37,19 @@ use std::sync::Arc;
 /// trait to allow expr to typable with respect to a schema
 pub trait ExprSchemable {
    /// given a schema, return the type of the expr
-    fn get_type<S: ExprSchema>(&self, schema: &S) -> Result<DataType>;
+    fn get_type(&self, schema: &dyn ExprSchema) -> Result<DataType>;


I had to change the traits to use dyn dispatch rather than generics so that the UDF could use the same object (and e.g. not have to create its own copy of these methods for Expr)

I expect this to have 0 performance impact, but I will run the planning benchmarks to be sure if this acceptable

I ran

cargo bench --bench sql_planner

And the results looked good ( within the noise threshold / reported 1% slower which I don't attribute to this change)

alamb · 2024-02-12T22:46:31Z

datafusion/expr/src/udf.rs

@@ -249,6 +261,22 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
    /// the arguments
    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;

+    /// What [`DataType`] will be returned by this function, given the types of
+    /// the expr arguments
+    fn return_type_from_exprs(


Update is I changed the signature to take &dyn ExprSchema which seems to have worked just fine

alamb

I think this PR is looking ready to go. However, given I contributed a bunch of code I think it would be good for another committer to give it a review prior to merging.

Thanks again @yyy1000 @Weijun-H and @universalmind303

alamb · 2024-02-14T21:51:43Z

@Jefffrey do you possibly have time to review this PR?

Jefffrey

Seems reasonable to me 👍

Just some minor doc comments

Jefffrey · 2024-02-15T10:10:28Z

datafusion/expr/src/udf.rs

+    /// What [`DataType`] will be returned by this function, given the
+    /// arguments?
+    ///
+    /// Note most UDFs should implement [`Self::return_type`] and not this
+    /// function. The output type for most functions only depends on the types
+    /// of their inputs (e.g. `sqrt(f32)` is always `f32`).
+    ///
+    /// By default, this function calls [`Self::return_type`] with the
+    /// types of each argument.
+    ///
+    /// This method can be overridden for functions that return different
+    /// *types* based on the *values* of their arguments.
+    ///
+    /// For example, the following two function calls get the same argument
+    /// types (something and a `Utf8` string) but return different types based
+    /// on the value of the second argument:
+    ///
+    /// * `arrow_cast(x, 'Int16')` --> `Int16`
+    /// * `arrow_cast(x, 'Float32')` --> `Float32`
+    ///
+    /// # Notes:
+    ///
+    /// This function must consistently return the same type for the same
+    /// logical input even if the input is simplified (e.g. it must return the same
+    /// value for `('foo' | 'bar')` as it does for ('foobar').


Maybe add some documentation about what would happen if a user tries to implement both return_type() and return_type_from_exprs()? (Which takes priority, etc.)

And what the suggested implementation for return_type() be if they choose to implement return_type_from_exprs() instead

Excellent idea - added in 653577f

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

Co-authored-by: Jeffrey Vo <jeffrey.vo.australia@gmail.com>

…to return-types

…type_from_exprs` are called

alamb

Thank you @yyy1000 and @universalmind303 @Weijun-H , @Dandandan and @Jefffrey for the great reviews.

I am just running a quick sql planning benchmark and then I plan to merge this PR in

cc @brayanjuls

alamb · 2024-02-15T16:28:39Z

I don't actually think this is an API change so removing the label

yyy1000 added 2 commits January 24, 2024 13:12

ScalarValue return types from argument values

4e06013

change file name

17a2c91

github-actions bot added the logical-expr Logical plan and expressions label Jan 24, 2024

yyy1000 mentioned this pull request Jan 24, 2024

Support computing ScalarValue return types from argument values (e.g. constants) rather than just types #8624

Closed

alamb mentioned this pull request Jan 24, 2024

DataFusion weekly project plan (Andrew Lamb) - Jan 22, 2024 #8933

Closed

9 tasks

Weijun-H reviewed Jan 26, 2024

View reviewed changes

datafusion/expr/src/udf.rs Show resolved Hide resolved

yyy1000 commented Jan 26, 2024

View reviewed changes

alamb reviewed Jan 26, 2024

View reviewed changes

alamb mentioned this pull request Jan 29, 2024

DataFusion weekly project plan (Andrew Lamb) - Jan 29, 2024 #9030

Closed

6 tasks

Merge branch 'apache:main' into return-types

02f2284

alamb mentioned this pull request Feb 4, 2024

DataFusion weekly project plan (Andrew Lamb) - Feb 5, 2024 #9121

Closed

6 tasks

alamb mentioned this pull request Feb 7, 2024

Issue using arrow_cast in ORDER BY expressions #9143

Closed

yyy1000 added 2 commits February 8, 2024 21:18

Merge branch 'apache:main' into return-types

0c9acdd

try using ?Sized

56b71ae

yyy1000 commented Feb 9, 2024

View reviewed changes

use Ok

3dbc0c7

universalmind303 reviewed Feb 9, 2024

View reviewed changes

yyy1000 added 2 commits February 9, 2024 15:50

move method default impl outside trait

491a4a1

Use type trait for ExprSchemable

468b38f

github-actions bot added the physical-expr Physical Expressions label Feb 11, 2024

fix nit

5772d9f

yyy1000 commented Feb 11, 2024

View reviewed changes

alamb mentioned this pull request Feb 12, 2024

DataFusion weekly project plan (Andrew Lamb) - Feb 12, 2024 #9200

Closed

8 tasks

alamb changed the title ~~ScalarValue return types from argument values~~ Support compute return types from argument values (not just their DataTypes) Feb 12, 2024

alamb mentioned this pull request Feb 12, 2024

Proposed Return Type from Expr suggestions yyy1000/arrow-datafusion#1

Merged

alamb reviewed Feb 12, 2024

View reviewed changes

Proposed Return Type from Expr suggestions (#1)

59b3958

* Improve return_type_from_args * Rework example * Update datafusion/core/tests/user_defined/user_defined_scalar_functions.rs --------- Co-authored-by: Junhao Liu <junhaoliu2023@gmail.com>

Apply suggestions from code review

f195fba

Co-authored-by: Alex Huang <huangweijun1001@gmail.com>

Fix tests + clippy

21d495f

github-actions bot added the optimizer Optimizer rules label Feb 12, 2024

alamb reviewed Feb 12, 2024

View reviewed changes

rework types to use dyn trait

b2e8457

alamb added the api change Changes the API exposed to users of the crate label Feb 12, 2024

alamb reviewed Feb 12, 2024

View reviewed changes

alamb added 2 commits February 12, 2024 19:17

fmt

4efb395

docs

a9546ee

alamb approved these changes Feb 13, 2024

View reviewed changes

Dandandan approved these changes Feb 15, 2024

View reviewed changes

Jefffrey approved these changes Feb 15, 2024

View reviewed changes

alamb mentioned this pull request Feb 15, 2024

DRAFT: arrow_cast function as UDF #9235

Closed

alamb and others added 4 commits February 15, 2024 07:06

Merge remote-tracking branch 'apache/main' into return-types

040d319

Apply suggestions from code review

93b72ee

Co-authored-by: Jeffrey Vo <jeffrey.vo.australia@gmail.com>

Merge branch 'return-types' of github.com:yyy1000/arrow-datafusion in…

e0add48

…to return-types

Add docs explaining what happens when both return_type and `return_…

653577f

…type_from_exprs` are called

alamb reviewed Feb 15, 2024

View reviewed changes

alamb added 3 commits February 15, 2024 08:44

clippy

7993af8

Merge remote-tracking branch 'apache/main' into return-types

a3b9648

fix doc -- comedy of errors

2121770

alamb merged commit 92d9274 into apache:main Feb 15, 2024
23 checks passed

alamb removed the api change Changes the API exposed to users of the crate label Feb 15, 2024

alamb mentioned this pull request Feb 20, 2024

Move arrow_cast to datafusion-functions crate #9287

Closed

yyy1000 deleted the return-types branch February 21, 2024 23:03

jayzhan211 mentioned this pull request Mar 10, 2024

Extend argument types for udf return_type_from_exprs #9522

Merged

gstvg mentioned this pull request Mar 22, 2024

Support custom struct field names with new scalar function named_struct #9743

Merged

Support compute return types from argument values (not just their DataTypes) #8985

Support compute return types from argument values (not just their DataTypes) #8985

Conversation

yyy1000 commented Jan 24, 2024

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Weijun-H commented Jan 26, 2024

yyy1000 commented Jan 26, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb Feb 6, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Feb 5, 2024

yyy1000 Feb 9, 2024 • edited Loading

Choose a reason for hiding this comment

universalmind303 Feb 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yyy1000 Feb 11, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Feb 12, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

alamb commented Feb 14, 2024

Jefffrey left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

alamb commented Feb 15, 2024

alamb Feb 6, 2024 •

edited

Loading

yyy1000 Feb 9, 2024 •

edited

Loading

universalmind303 Feb 9, 2024 •

edited

Loading

yyy1000 Feb 11, 2024 •

edited

Loading

alamb left a comment •

edited

Loading