Introduce `return_type_from_args` for ScalarFunction. #14094

jayzhan211 · 2025-01-12T09:25:15Z

Which issue does this PR close?

Rationale for this change

return_type_from_args that has less dependencies on Expr itself but the computed properties of Expr and Schema including data_type and nullability

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

TODO

Combine return_type_from_args and is_nullable_from_args_nullable

Signed-off-by: Jay Zhan <[email protected]>

jayzhan211 · 2025-01-12T09:27:07Z

datafusion/expr/src/udf.rs

    pub fn is_nullable(&self, args: &[Expr], schema: &dyn ExprSchema) -> bool {
        self.inner.is_nullable(args, schema)
    }

+    pub fn is_nullable_from_args_nullable(&self, args_nullables: &[bool]) -> bool {


Remove Expr dependency

jayzhan211 · 2025-01-12T09:28:11Z

datafusion/expr/src/udf.rs

+    /// The data types of the arguments to the function
+    pub arg_types: &'a [DataType],
+    /// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string
+    pub arguments: &'a [String],


better name 🤔 ?

Would it be possible to unify the argument handling so that both return type and nullability are returned the same?

I wonder if it would somehow be possible to add the input nullable information here too 🤔

I am also not sure about only supporting string args, that is likely a regression in behavior for some users (For example, maybe they look for constant integers as well)

datafusion/functions/src/core/arrow_cast.rs

jayzhan211 · 2025-01-12T09:29:51Z

datafusion/functions/src/core/named_struct.rs

            let name_column = &chunk[0];
            let name = match name_column {
-                ColumnarValue::Scalar(ScalarValue::Utf8(Some(name_scalar))) => name_scalar,
-                _ => return exec_err!("named_struct even arguments must be string literals, got {name_column:?} instead at position {}", i * 2)


name_column output less readable array in this change, remove it for now.

findepi · 2025-01-12T10:09:23Z

As stated in #13717 (comment) , this new method doesn't necessarily simplify anything.

Can you please fill "Rationale for this change"? What problem are we solving?

alamb

Thanks @jayzhan211 -- I think this is a step in the right direction, but I am worried it just makes the API more complicated (adds as many functions as it deprecates)

Challenge: Exprs / constants

It seems to me one challenge is that different information is known for computing return types at different points in the plan (e.g. sometimes we have Expr and sometimes we don't)

What would you think about making this more explicit in ReturnTypeArgs by making it an enum:

#[derive(Debug)]
pub enum ReturnTypeArgs<'a> {
    /// information known at logical planning time
    /// Note you can get get type and nullability for each arg
   // using the specified ExprSchema
    Planning {
       pub args: &'a[Expr],
       pub schema: &'a dyn ExprSchema
    },
    /// Information known during Execution
    Execution {
    /// The data types of the arguments to the function
      pub arg_types: &'a [DataType],
      pub arg_nullability: [bool],
    }
}

Challenge: Multiple APIs (Nullability and return type)

It is somewhat akward to have two functions, one for nullability and one for return type. Also I can imagine that the nullability calculation depends on the input type of arguments too (not just the input nullability) I wonder if we can combine them into a single API:

Maybe something like

/// Information about the output of the function
/// including the data type and nullability:
struct ReturnTypeInfo {
  data_type: DataType,
  nullable: bool,
}

trait ScalarUDFImpl {
    /// Returns the 
    pub fn return_type_from_args(&self, args: ReturnTypeArgs) -> Result<ReturnTypeInfo> 
}

alamb · 2025-01-12T11:46:57Z

datafusion/expr/src/udf.rs

+    /// The data types of the arguments to the function
+    pub arg_types: &'a [DataType],
+    /// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string
+    pub arguments: &'a [String],


Would it be possible to unify the argument handling so that both return type and nullability are returned the same?

I wonder if it would somehow be possible to add the input nullable information here too 🤔

alamb · 2025-01-12T11:48:20Z

datafusion/expr/src/udf.rs

+    /// The data types of the arguments to the function
+    pub arg_types: &'a [DataType],
+    /// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string
+    pub arguments: &'a [String],


I am also not sure about only supporting string args, that is likely a regression in behavior for some users (For example, maybe they look for constant integers as well)

jayzhan211 · 2025-01-12T14:10:48Z

Multiple APIs (Nullability and return type)

Great, I also want this too.

jayzhan211 · 2025-01-12T14:14:37Z

I am also not sure about only supporting string args, that is likely a regression in behavior for some users (For example, maybe they look for constant integers as well)

Yes, it might be, since I assume we don't really need Expr but String instead. However, in theory they can achieve what they need with String, but of course this is breaking change.

If the constant integer is the only concern, ScalarValue or ColumnarValue looks good for me too!

jayzhan211 · 2025-01-12T14:17:11Z

#[derive(Debug)]
pub enum ReturnTypeArgs<'a> {
    /// information known at logical planning time
    /// Note you can get get type and nullability for each arg
   // using the specified ExprSchema
    Planning {
       pub args: &'a[Expr],
       pub schema: &'a dyn ExprSchema
    },
    /// Information known during Execution
    Execution {
    /// The data types of the arguments to the function
      pub arg_types: &'a [DataType],
      pub arg_nullability: [bool],
    }
}

One good thing in this PR is that we don't need Expr anymore, we compute data type and nullable in datafusion core and they are not "public" for customization.

Do we really need Planning? My thought is that whenever we have Expr and Schema we can compute corresponding DataType and Nullability. Therefore, even for Planning stage, we can still get the information DataType + Nullability

Signed-off-by: Jay Zhan <[email protected]>

alamb · 2025-01-13T19:09:41Z

Do we really need Planning? My thought is that whenever we have Expr and Schema we can compute corresponding DataType and Nullability. Therefore, even for Planning stage, we can still get the information DataType + Nullability

I think the usecase is for expressions like arrow_cast whose output type depends on the actual value of the arguments (not just the type of the argument)

So like arrow-cast(String) can return Int64 or Float64 depending on the value of the argument -- the value is passed as an Expr

jayzhan211 · 2025-01-14T00:21:16Z

So like arrow-cast(String

For arrow_cast, the output type is determined by the 2nd argument which is string.

arrow_cast(a, 'Int64') returns I64, arrow_cast(a, 'List(Int64)') returns List(Int64)

All the functions in datafusion have constant string type, the only concern is that other users do the crazy things with Expr, but I don't think it is practical since String is enough in my opionion

Signed-off-by: Jay Zhan <[email protected]>

alamb · 2025-01-16T22:24:49Z

This is on my review queue, I just haven't quite had time to get to it yet. I will do so tomorrow if not later tonight

alamb

Thank you @jayzhan211 -- I think this is a very nice improvement and makes the argument handling in user defined functions better

I also like how most functions will be unaffected as we still have ScalarUDFImpl::return_type

I really do think that we should change

    pub arguments: &'a [String],

to this before the release:

    pub arguments: &'a [Option<&'a ScalarValue>]>`,

But I believe this strongly enough that I will make a proposed follow on PR to this one if you don't want to make the change here

Thanks again

datafusion/expr/src/udf.rs

alamb · 2025-01-17T15:31:53Z

datafusion/functions/src/core/coalesce.rs

@@ -174,39 +177,3 @@ fn get_coalesce_doc() -> &'static Documentation {
            .build()
    })
 }
-
-#[cfg(test)]


why remove these tests?

It is covered in slt, remove duplicated

alamb · 2025-01-17T17:16:46Z

datafusion/functions/src/core/named_struct.rs

+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(name_scalar))) => {
+                    name_scalar
+                }
+                // TODO: Implement Display for ColumnarValue


That is a good todo -- maybe we can make a follow on ticket for it

alamb · 2025-01-17T17:17:36Z

datafusion/expr/src/udf.rs

@@ -490,6 +547,40 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
        self.return_type(arg_types)
    }

+    /// What [`DataType`] will be returned by this function, given the


I also recommend removing the comments from return_type_from_exprs now that it is deprecated and you have moved the content here

datafusion/expr/src/udf.rs

alamb · 2025-01-17T17:21:25Z

datafusion/physical-expr/src/scalar_function.rs

+                e.as_any()
+                    .downcast_ref::<Literal>()
+                    .map(|literal| match literal.value() {
+                        ScalarValue::Utf8(Some(s)) => s.clone(),


It would be really nice to avoid this clone here -- if we changed the signature to &[Option<&str>] I think that would be easier to use and more efficient

alamb · 2025-01-17T17:25:18Z

FYI @andygrove

jayzhan211 · 2025-01-18T08:17:55Z

I plan to change to ScalarValue after #14167, so we can parse the string easily.

If anyone interested in this change can also take it.

Signed-off-by: Jay Zhan <[email protected]>

alamb · 2025-01-18T13:21:57Z

I plan to change to ScalarValue after #14167, so we can parse the string easily.

If anyone interested in this change can also take it.

I have just merged in #14167

Signed-off-by: Jay Zhan <[email protected]>

jayzhan211 · 2025-01-19T03:19:33Z

We can start return_type deprecation after this PR

Signed-off-by: Jay Zhan <[email protected]>

…UDFs

… UDFs

alamb · 2025-01-19T13:06:08Z

We can start return_type deprecation after this PR

I recommend we do not deprecate return_type - after this PR I think we have a nice API. Namely most simple uses can use return_type and ones that need more control (like returning nullability, or getting scalar args) can use return_type_from_args

alamb

Thank you @jayzhan211 -- I think this looks really nice ❤️

I left some comment suggestions, but I can also make them as a follow on PR as well

alamb · 2025-01-19T13:06:28Z

datafusion/common/src/utils/mod.rs

+        let a = Arc::new(Int32Array::from(vec![3; 200])) as ArrayRef;
+        println!(
+            "display {}",
+            pretty_format_columns("ColumnarValue(ArrayRef)", &[a]).unwrap()


Oh, I forgot to remove this 😆

alamb · 2025-01-19T13:08:35Z

datafusion/expr/src/expr_schema.rs

+
+                let (return_type, nullable) =
+                    func.return_type_from_args(args)?.into_parts();
+                Ok((return_type, nullable))


alamb · 2025-01-19T13:10:31Z

datafusion/expr/src/udf.rs

+    /// The data types of the arguments to the function
+    pub arg_types: &'a [DataType],
+    /// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string
+    pub arguments: &'a [Option<&'a ScalarValue>],


I also recommend naming this scalar_arguments

Suggested change

pub arguments: &'a [Option<&'a ScalarValue>],

pub scalar_arguments: &'a [Option<&'a ScalarValue>],

alamb · 2025-01-19T13:11:08Z

datafusion/expr/src/udf.rs

+pub struct ReturnTypeArgs<'a> {
+    /// The data types of the arguments to the function
+    pub arg_types: &'a [DataType],
+    /// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string


Suggested change

/// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string

/// Is argument `i` to the function a scalar (constant)

///

/// If argument `i` is not a scalar, it will be None

///

/// For example, if a function is called like `my_function(column_a, 5)`

/// this field will be `[None, Some(ScalarValue::Int32(Some(5)))]`

alamb · 2025-01-19T13:12:43Z

datafusion/expr/src/udf.rs

+#[derive(Debug)]
+pub struct ReturnInfo {
+    return_type: DataType,
+    nullable: bool,
+}


Suggested change

#[derive(Debug)]

pub struct ReturnInfo {

return_type: DataType,

nullable: bool,

}

/// Return metadata for this function.

///

/// See [`ScalarUDFImpl::return_type_from_args`] for more information

#[derive(Debug)]

pub struct ReturnInfo {

/// The type of the value return by the function

return_type: DataType,

/// Can the returned value (ever) be null?

nullable: bool,

}

alamb · 2025-01-19T13:14:05Z

datafusion/expr/src/udf.rs

@@ -342,6 +348,56 @@ pub struct ScalarFunctionArgs<'a> {
    pub return_type: &'a DataType,
 }

+#[derive(Debug)]


Suggested change

#[derive(Debug)]

/// Information about arguments passed to the function

///

/// This structure contains metadata about how the function was called

/// such as the type of the arguments, any scalar arguments and if the

/// arguments can (ever) be null

///

/// See [`ScalarUDFImpl::return_type_from_args`] for more information

#[derive(Debug)]

alamb · 2025-01-19T13:14:32Z

datafusion/expr/src/udf.rs

+    pub arg_types: &'a [DataType],
+    /// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string
+    pub arguments: &'a [Option<&'a ScalarValue>],
+    pub nullables: &'a [bool],


Suggested change

pub nullables: &'a [bool],

/// Can argument `i` (ever) null?

pub nullables: &'a [bool],

alamb · 2025-01-19T13:15:56Z

datafusion/functions/src/core/arrow_cast.rs

+        debug_assert_eq!(args.arguments.len(), 2);
+
+        args.arguments[1]
+            .and_then(|sv| sv.try_as_str().flatten().filter(|s| !s.is_empty()))


this is pretty fancy functional programming

Signed-off-by: Jay Zhan <[email protected]>

…turnTypeArgs

Signed-off-by: Jay Zhan <[email protected]>

alamb · 2025-01-20T22:13:05Z

Love it. And love that this isn't an API change. Thank you so much @jayzhan211 -- this really looks great

jayzhan211 added 9 commits January 12, 2025 13:57

switch func

6b00b9a

Signed-off-by: Jay Zhan <[email protected]>

fix test

b079be3

Signed-off-by: Jay Zhan <[email protected]>

fix test

8c9ee8c

Signed-off-by: Jay Zhan <[email protected]>

deprecate old

6df7476

Signed-off-by: Jay Zhan <[email protected]>

add try new

fe7f6a5

Signed-off-by: Jay Zhan <[email protected]>

deprecate

4da4c71

Signed-off-by: Jay Zhan <[email protected]>

rm deprecate

de4b484

Signed-off-by: Jay Zhan <[email protected]>

reaplce deprecated func

02a64ce

Signed-off-by: Jay Zhan <[email protected]>

cleanup

f26ce70

Signed-off-by: Jay Zhan <[email protected]>

github-actions bot added logical-expr Logical plan and expressions physical-expr Physical Expressions core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) functions labels Jan 12, 2025

jayzhan211 commented Jan 12, 2025

View reviewed changes

datafusion/functions/src/core/arrow_cast.rs Show resolved Hide resolved

jayzhan211 commented Jan 12, 2025

View reviewed changes

alamb reviewed Jan 12, 2025

View reviewed changes

jayzhan211 added 3 commits January 13, 2025 22:17

combine type and nullable

b967034

Signed-off-by: Jay Zhan <[email protected]>

fix slowdown

50cac9e

Signed-off-by: Jay Zhan <[email protected]>

clippy

7909231

Signed-off-by: Jay Zhan <[email protected]>

jayzhan211 changed the title ~~Add two new methods in ScalarFunction return_type_from_args and is_nullable_from_args_nullable~~ Introduce return_type_from_args for ScalarFunction. Jan 13, 2025

fix take

9a95659

Signed-off-by: Jay Zhan <[email protected]>

alamb approved these changes Jan 17, 2025

View reviewed changes

andygrove mentioned this pull request Jan 17, 2025

Prepare for DataFusion 45 apache/datafusion-comet#1304

Open

jayzhan211 added 2 commits January 18, 2025 16:07

Merge branch 'main' of github.com:apache/datafusion into ret-ty

fd2f35d

refactor: remove unused documentation sections from scalar functions

26e6346

upd doc

3f2ae5c

Signed-off-by: Jay Zhan <[email protected]>

jayzhan211 added 3 commits January 19, 2025 10:40

Merge branch 'main' of github.com:apache/datafusion into ret-ty

727ba44

use scalar value

5ad7b5c

Signed-off-by: Jay Zhan <[email protected]>

fix test

0545181

Signed-off-by: Jay Zhan <[email protected]>

github-actions bot added the common Related to common crate label Jan 19, 2025

jayzhan211 marked this pull request as draft January 19, 2025 03:32

jayzhan211 added 5 commits January 19, 2025 11:43

fix test

3014267

Signed-off-by: Jay Zhan <[email protected]>

use try_as_str

8463698

Signed-off-by: Jay Zhan <[email protected]>

refactor: improve error handling for constant string arguments in UDFs

40dfc6c

refactor: enhance error messages for constant string requirements in …

c321ff8

…UDFs

refactor: streamline argument validation in return_type_from_args for…

8ea6cef

… UDFs

jayzhan211 marked this pull request as ready for review January 19, 2025 07:16

alamb approved these changes Jan 19, 2025

View reviewed changes

jayzhan211 added 4 commits January 20, 2025 08:49

rename and doc

486a3b6

Signed-off-by: Jay Zhan <[email protected]>

refactor: add documentation for nullability of scalar arguments in Re…

78b8173

…turnTypeArgs

rm test

a72f116

Signed-off-by: Jay Zhan <[email protected]>

refactor: remove unused import of Int32Array in utils tests

61abb93

github-actions bot removed the common Related to common crate label Jan 20, 2025

alamb merged commit d3f1c9a into apache:main Jan 20, 2025
25 checks passed

	pub arguments: &'a [Option<&'a ScalarValue>],
	pub scalar_arguments: &'a [Option<&'a ScalarValue>],

-    /// The Utf8 arguments to the function, if the expression is not Utf8, it will be empty string
+    /// Is argument `i` to the function a scalar (constant)
+    ///
+    /// If argument `i` is not a scalar, it will be None
+    ///
+    /// For example, if a function is called like `my_function(column_a, 5)`
+    /// this field will be `[None, Some(ScalarValue::Int32(Some(5)))]`

-#[derive(Debug)]
+/// Information about arguments passed to the function
+///
+/// This structure contains metadata about how the function was called
+/// such as the type of the arguments, any scalar arguments and if the
+/// arguments can (ever) be null
+///
+/// See [`ScalarUDFImpl::return_type_from_args`] for more information
+#[derive(Debug)]

	pub nullables: &'a [bool],
	/// Can argument `i` (ever) null?
	pub nullables: &'a [bool],

Introduce return_type_from_args for ScalarFunction. #14094

Introduce return_type_from_args for ScalarFunction. #14094

Conversation

jayzhan211 commented Jan 12, 2025 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

TODO

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

findepi commented Jan 12, 2025

alamb left a comment

Choose a reason for hiding this comment

Challenge: Exprs / constants

Challenge: Multiple APIs (Nullability and return type)

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jayzhan211 commented Jan 12, 2025

jayzhan211 commented Jan 12, 2025 • edited Loading

jayzhan211 commented Jan 12, 2025 • edited Loading

alamb commented Jan 13, 2025

jayzhan211 commented Jan 14, 2025 • edited Loading

alamb commented Jan 16, 2025

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Jan 17, 2025

jayzhan211 commented Jan 18, 2025

alamb commented Jan 18, 2025

jayzhan211 commented Jan 19, 2025

alamb commented Jan 19, 2025

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

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

Choose a reason for hiding this comment

alamb commented Jan 20, 2025

Introduce `return_type_from_args` for ScalarFunction. #14094

Introduce `return_type_from_args` for ScalarFunction. #14094

jayzhan211 commented Jan 12, 2025 •

edited

Loading

jayzhan211 commented Jan 12, 2025 •

edited

Loading

jayzhan211 commented Jan 12, 2025 •

edited

Loading

jayzhan211 commented Jan 14, 2025 •

edited

Loading