Introduce user-defined signature #10439

jayzhan211 · 2024-05-10T01:24:59Z

Which issue does this PR close?

Closes #10423 .

Rationale for this change

What changes are included in this PR?

Add VariadicCoercion, UniformCoercion
Remove VariadicEqual (coercion rule does not require the result type to be the same for all)

Add the datafusion error aside from the signature error message for better debugging usage that is why many test files are updated.

Are these changes tested?

UDAF code is also included but no test is covered since no UDAF requires coercion there yet (We will cover it after moving builtin to UDAF)

Are there any user-facing changes?

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 · 2024-05-10T01:34:09Z

datafusion/expr/src/expr_schema.rs

                            plan_datafusion_err!(
-                                "{}",
+                                "{} and {}",


keep error for debugging

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 · 2024-05-10T02:33:18Z

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -833,12 +865,8 @@ mod test {
            signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable),
        })
        .call(vec![lit("Apple")]);
-        let plan_err = Projection::try_new(vec![udf], empty)
+        Projection::try_new(vec![udf], empty)


I think matching err message is unnecessary

As long as the message is tested elsewhere I think this is fine

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

alamb

Thank you @jayzhan211 -- this looks great to me

Prior to merge, I think it is worth considering (but not required):

Using a single new variant for TypeSignature (rather than 2)
Improving the doc strings (I left some suggestions)
Moving the semantic changes to coalsce into their own PR

alamb · 2024-05-10T18:50:05Z

datafusion/expr/src/signature.rs

-    ///
-    /// `make_array(i32, i64) -> make_array(i64, i64)`
-    VariadicEqual,
+    /// One or more arguments of an arbitrary type and coerced with user-defined coercion rules.


Since both TypeSignature::VariadicCoercion and TypeCoercion::UniformCoercion call into ScalarUDFImpl::coerce_types I wonder if there is a reason to have both.

Maybe instead we could add a single variant for both cases. Perhaps something like:

pub enum TypeSignature { ... /// The acceptable signature and coercions rules to coerce arguments to this /// signature are special for this function. If this signature is specified, /// Datafusion will call [`ScalarUDFImpl::coerce_types`] to prepare argument types. FunctionDefined, ... }

alamb · 2024-05-10T18:50:48Z

datafusion/expr/src/type_coercion/functions.rs

-            }
+        TypeSignature::VariadicCoercion | TypeSignature::UniformCoercion(_) => {
+            return internal_err!(
+                "Coercion signature is handled in function-specific get_valid_types."


Suggested change

"Coercion signature is handled in function-specific get_valid_types."

"Coercion signature should be handled by in function-specific coerce_types."

alamb · 2024-05-10T18:52:57Z

datafusion/functions-array/src/make_array.rs

@@ -111,6 +111,25 @@ impl ScalarUDFImpl for MakeArray {
    fn aliases(&self) -> &[String] {
        &self.aliases
    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {


alamb · 2024-05-10T18:54:12Z

datafusion/functions/src/core/nvl2.rs

        Ok(arg_types[1].clone())
    }

    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
        nvl2_func(args)
    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        let new_type = arg_types.iter().skip(1).try_fold(


this pattern to coerce all arguments to the type of the first one is repeated several times -- maybe we could make some sort of utility function to do it

As a follow on PR perhaps

alamb · 2024-05-10T18:55:10Z

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -833,12 +865,8 @@ mod test {
            signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable),
        })
        .call(vec![lit("Apple")]);
-        let plan_err = Projection::try_new(vec![udf], empty)
+        Projection::try_new(vec![udf], empty)


As long as the message is tested elsewhere I think this is fine

alamb · 2024-05-10T18:55:58Z

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -914,7 +942,7 @@ mod test {
            .err()
            .unwrap();
        assert_eq!(
-            "type_coercion\ncaused by\nError during planning: Coercion from [Utf8] to the signature Uniform(1, [Float64]) failed.",
+            "type_coercion\ncaused by\nError during planning: [data_types_with_aggregate_udf] Coercion from [Utf8] to the signature Uniform(1, [Float64]) failed.",


Maybe we could make this error say "User defined coercsion rule failed" so it was clearer what happened

alamb · 2024-05-10T18:56:40Z

datafusion/sqllogictest/test_files/coalesce.slt

@@ -23,7 +23,7 @@ select coalesce(1, 2, 3);
 1

 # test with first null
-query IT
+query ?T


I would still prefer to make the change to the coalsce coercion rules as its own PR, but I am also fine if we include it in this PR

We need to keep the old signature VariadicEqual for coalesce to split the change for coalesce :(. I prefer to make the change at once.

alamb · 2024-05-10T19:05:32Z

datafusion/expr/src/udf.rs

@@ -420,6 +425,11 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
    fn short_circuits(&self) -> bool {
        false
    }
+
+    /// Coerce the types of the arguments to the types that the function expects


Nice -- I think it would be nice to expand this documentation to help people understand the API some more

Perhaps something like

Suggested change

/// Coerce the types of the arguments to the types that the function expects

/// Coerce arguments of a function call to types that the function can evaluate.

///

/// This function is only called if [`Self::signature`] returns [`TypeSignature::UserDefined`]. Most

/// UDFs should return one of the other variants of `TypeSignature` which handle common

/// cases

///

/// See the [type coercion module](datafusion_expr::type_coercion)

/// documentation for more details on type coercion

///

/// For example, if your function requires a floating point arguments, but the user calls

/// it like `my_func(1::int)` (aka with `1` as an integer), coerce_types could return `[DataType::Float64]`

/// to ensure the argument was cast to `1::double`

///

/// # Parameters

/// * `arg_types`: The argument types of the arguments this function with

///

/// # Return value

/// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call

/// arguments to these specific types.

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

alamb

Looks great to me -- thanks @jayzhan211

I had one question on a comment in the docstrings but we could fix that as a follow on as well

alamb · 2024-05-11T11:14:54Z

datafusion/expr/src/type_coercion/functions.rs

 /// Performs type coercion for function arguments.
 ///
 /// Returns the data types to which each argument must be coerced to
 /// match `signature`.
 ///
 /// For more details on coercion in general, please see the
 /// [`type_coercion`](crate::type_coercion) module.
+///
+/// This function will be replaced with [data_types_with_scalar_udf],


I am not sure we want to replace this function over time -- I think having the basic simple Signatures that handle most common coercions makes sense to have in DataFusion core (even if it could be done purely in a udf) as it will make creating UDFs easier for users

Agree with you

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

alamb · 2024-05-11T19:22:57Z

🎉

jayzhan211 added 7 commits May 10, 2024 07:50

introduce new sig

ff441ea

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

add udfimpl

8b6016b

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

replace fun

36abe3e

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

replace array

0988efe

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

coalesce

a64b813

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

nvl2

5bbd2a0

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

rm variadic equal

6f0a90b

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

github-actions bot added logical-expr Logical plan and expressions physical-expr Physical Expressions optimizer Optimizer rules sqllogictest labels May 10, 2024

jayzhan211 commented May 10, 2024

View reviewed changes

fix test

5cc047b

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 commented May 10, 2024

View reviewed changes

rm err msg to fix ci

a515aad

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 force-pushed the sig-coerce branch from 560196b to a515aad Compare May 10, 2024 11:47

jayzhan211 marked this pull request as ready for review May 10, 2024 12:35

alamb added the api change Changes the API exposed to users of the crate label May 10, 2024

alamb approved these changes May 10, 2024

View reviewed changes

jayzhan211 added 6 commits May 11, 2024 08:28

user defined sig

17e6ec1

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

add err msg

5eaacc8

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fmt

a606581

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

cleanup

40d5444

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fix ci

ca7f942

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fix ci

5cb0d0f

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

alamb approved these changes May 11, 2024

View reviewed changes

upd comment

68fdf52

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 changed the title ~~Introduce coercion signature VariadicCoercion and UniformCoercion~~ Introduce user-defined signature May 11, 2024

jayzhan211 merged commit 76f9e2e into apache:main May 11, 2024
24 checks passed

alamb mentioned this pull request May 13, 2024

DataFusion weekly project plan (Andrew Lamb) - May 13, 2024 #10482

Closed

8 tasks

jayzhan211 mentioned this pull request May 14, 2024

select array_concat([]) panicked #10200

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduce user-defined signature #10439

Introduce user-defined signature #10439

jayzhan211 commented May 10, 2024 •

edited

jayzhan211 May 10, 2024

jayzhan211 May 10, 2024

alamb May 10, 2024

alamb left a comment

alamb May 10, 2024

alamb May 10, 2024

alamb May 10, 2024

alamb May 10, 2024

alamb May 10, 2024

alamb May 10, 2024

alamb May 10, 2024

jayzhan211 May 10, 2024

alamb May 10, 2024

alamb left a comment

alamb May 11, 2024

jayzhan211 May 11, 2024

alamb commented May 11, 2024

	"Coercion signature is handled in function-specific get_valid_types."
	"Coercion signature should be handled by in function-specific coerce_types."

-    /// Coerce the types of the arguments to the types that the function expects
+    /// Coerce arguments of a function call to types that the function can evaluate.
+    ///
+    /// This function is only called if [`Self::signature`] returns [`TypeSignature::UserDefined`]. Most
+    /// UDFs should return one of the other variants of `TypeSignature` which handle common
+    /// cases
+    ///
+    /// See the [type coercion module](datafusion_expr::type_coercion)
+    /// documentation for more details on type coercion
+    ///
+    /// For example, if your function requires a floating point arguments, but the user calls
+    /// it like `my_func(1::int)` (aka with `1` as an integer), coerce_types could return `[DataType::Float64]`
+    /// to ensure the argument was cast to `1::double`
+    ///
+    /// # Parameters
+    /// * `arg_types`: The argument types of the arguments  this function with
+    ///
+    /// # Return value
+    /// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call
+    /// arguments to these specific types.

Introduce user-defined signature #10439

Introduce user-defined signature #10439

Conversation

jayzhan211 commented May 10, 2024 • edited

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?

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

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 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 commented May 11, 2024

jayzhan211 commented May 10, 2024 •

edited