Add `field` trait method to `WindowUDFImpl` #12374

jcsherin · 2024-09-07T13:35:28Z

Which issue does this PR close?

Rationale for this change

The result field from evaluating the user-defined window function is composed from the return_type and nullable trait methods in WindowUDFImpl.

This change explores folding both methods into a single trait method. The user-defined window functions have to implement only the field trait method which makes the intent more explicit.

The current implementation for a user-defined window function (without field trait method) looks like this:

impl WindowUDFImpl for RowNumber {
    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        Ok(DataType::UInt64)
    }

    fn nullable(&self) -> bool {
        true
    }
}

The implementation for a user-defined window function after this change:

impl WindowUDFImpl for RowNumber {
    fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field> {
        Ok(Field::new(
             field_args.name(), /* window function display name */
             DataType::UInt64,  /* result data type */
             false              /* row number is not nullable */
       ))
    }
}

What changes are included in this PR?

Add field trait method:

fn field(&self, field_args: WindowUDFFieldArgs) -> Result<Field>

Remove return_type trait method.

datafusion/datafusion/expr/src/udwf.rs

Lines 282 to 284 in a08f923

    
           /// What [`DataType`] will be returned by this function, given the types of 
        
           /// the arguments 
        
           fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;

Remove nullable trait method which was added in Convert built-in row_number to user-defined window function #12030.
Add WindowUDFFieldArgs:

/// Contains metadata necessary for defining the field which represents
/// the final result of evaluating a user-defined window function.
pub struct WindowUDFFieldArgs<'a> {
    /// The data types of input expressions to the user-defined window
    /// function.
    input_types: &'a [DataType],
    /// The display name of the user-defined window function.
    function_name: &'a str,
}

Are these changes tested?

Yes, against existing tests in CI.

Are there any user-facing changes?

Yes, this is a breaking change for user-defined window functions API.

datafusion/expr/src/udwf.rs

datafusion/physical-plan/src/windows/mod.rs

datafusion/expr/src/expr.rs

jcsherin · 2024-09-07T13:58:39Z

I lack permission to apply the api-change label to this PR.

jayzhan211 · 2024-09-07T13:53:50Z

datafusion/expr/src/expr_schema.rs

+                        .iter()
+                        .map(|e| e.get_type(input_schema))
+                        .collect::<Result<Vec<_>>>()?;
+                    let input_types = data_types_with_window_udf(&data_types, udwf)?;


I think we should handle this compute in to_field. data_types_with_window_udf could reuse to_field. data_types and nullable could reuse data_types_with_window_udf

jayzhan211 · 2024-09-07T13:59:30Z

datafusion/expr/src/expr_schema.rs

+                        .map(|e| e.get_type(input_schema))
+                        .collect::<Result<Vec<_>>>()?;
+                    let input_types = data_types_with_window_udf(&data_types, udwf)?;
+                    let function_name = self.schema_name().to_string();


Similar to self.qualified_name() in to_field, we might have different name for Column and Alias?

jcsherin · 2024-09-17T13:21:41Z

datafusion/expr/src/expr_schema.rs

+impl Expr {
+    /// Common method for window functions that applies type coercion
+    /// to all arguments of the window function to check if it matches
+    /// its signature.
+    ///
+    /// If successful, this method returns the data type and
+    /// nullability of the window function's result.
+    ///
+    /// Otherwise, returns an error if there's a type mismatch between
+    /// the window function's signature and the provided arguments.
+    fn data_type_and_nullable_with_window_function(
+        &self,
+        schema: &dyn ExprSchema,
+        window_function: &WindowFunction,
+    ) -> Result<(DataType, bool)> {


Extracted a common method to handle type coercion for all window function types (built-in, udaf and udwf) which is then reused by methods:,

data_type_and_nullable,

get_type and,

nullable

jcsherin · 2024-09-17T13:22:36Z

datafusion/expr/src/udwf.rs

-    /// Return the type of the function given its input types
-    ///
-    /// See [`WindowUDFImpl::return_type`] for more details.
-    pub fn return_type(&self, args: &[DataType]) -> Result<DataType> {
-        self.inner.return_type(args)
-    }
-


Removed return_type.

jcsherin · 2024-09-17T13:22:48Z

datafusion/expr/src/udwf.rs

-    /// Returns if column values are nullable for this window function.
+    /// Returns the field of the final result of evaluating this window function.
    ///
-    /// See [`WindowUDFImpl::nullable`] for more details.
-    pub fn nullable(&self) -> bool {
-        self.inner.nullable()


Removed nullable

jcsherin · 2024-09-17T13:24:40Z

datafusion/expr/src/expr_schema.rs

+                            )
+                        )
+                    })?;
+                let (_, function_name) = self.qualified_name();


Use Expr::qualified_name which also handles:

Expr::Column and,

Expr::Alias

jcsherin · 2024-09-17T13:31:52Z

datafusion/expr/src/expr.rs

+            WindowFunctionDefinition::WindowUDF(fun) => fun
+                .field(WindowUDFFieldArgs::new(input_expr_types, display_name))
+                .map(|field| field.data_type().clone()),


Return data type for udwf.

jcsherin added 30 commits September 3, 2024 16:43

Adds new library functions-window-common

ae69e1e

Adds FieldArgs struct for field of final result

f09383c

Adds field method to WindowUDFImpl trait

ca256fa

Minor: fixes formatting

9b21b7f

Fixes: udwf doc test

9324df6

Fixes: implements missing trait items

4cd1467

Updates datafusion-cli dependencies

deff13c

Fixes: formatting of Cargo.toml files

ac5ef0a

Fixes: implementation of field in udwf example

7c4293a

Pass FieldArgs argument to field

83d833b

Use field in place of return_type for udwf

927f1cc

Update field in udwf implementations

e29a493

Merge branch 'main' into field-args-for-window-udf

79d8372

Fixes: implementation of field in udwf example

fc969e8

Revert unrelated change

df2b1c8

Mark return_type for udwf as unreachable

e26aba6

Delete code

258986d

Uses schema name of udwf to construct FieldArgs

ba9e39a

Adds deprecated notice to return_type trait method

b671b36

Add doc comments to field trait method

712263e

Reify input_types when creating the udwf window expression

9426978

Rename name field to schema_name in FieldArgs

327db32

Make FieldArgs opaque

f5b62ef

Minor refactor

34565ff

Removes nullable trait method from WindowUDFImpl

ede7827

Add doc comments

25f3b50

Rename to WindowUDFResultArgs

f012bdd

Minor: fixes formatting

0e8795f

Copy edits for doc comments

4bd77d5

Renames field to function_name

a0f2ccb

jcsherin commented Sep 7, 2024

View reviewed changes

datafusion/expr/src/udwf.rs Outdated Show resolved Hide resolved

jcsherin commented Sep 7, 2024

View reviewed changes

datafusion/expr/src/udwf.rs Outdated Show resolved Hide resolved

jcsherin commented Sep 7, 2024

View reviewed changes

datafusion/physical-plan/src/windows/mod.rs Outdated Show resolved Hide resolved

jcsherin commented Sep 7, 2024

View reviewed changes

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

Merge branch 'main' into field-args-for-window-udf

acfa03e

jayzhan211 reviewed Sep 7, 2024

View reviewed changes

jcsherin added 2 commits September 8, 2024 00:07

Fixes: use Expr::qualified_name for window function name

a47b7a0

Fixes: apply previous fix to Expr::nullable

8da1364

jcsherin marked this pull request as draft September 7, 2024 21:32

jcsherin added 8 commits September 12, 2024 16:59

Merge branch 'main' into field-args-for-window-udf

9413a58

Refactor: reuse type coercion for window functions

275daea

Merge branch 'main' into field-args-for-window-udf

4c7157d

Fixes: clippy errors

045d352

Adds name parameter to WindowFunctionDefinition::return_type

5cc7d06

Removes return_type field from SimpleWindowUDF

0102334

Add doc comment for helper method

4bb799d

Rewrite doc comments

2530169

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Sep 17, 2024

Minor: remove empty comment

25b299e

github-actions bot removed the sqllogictest SQL Logic Tests (.slt) label Sep 17, 2024

jcsherin added 2 commits September 17, 2024 18:11

Remove WindowUDFImpl::return_type

ce7df3a

Fixes doc test

bb57f8f

jcsherin commented Sep 17, 2024

View reviewed changes

jcsherin marked this pull request as ready for review September 17, 2024 14:27

Merge branch 'main' into field-args-for-window-udf

e6a7113

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `field` trait method to `WindowUDFImpl` #12374

Add `field` trait method to `WindowUDFImpl` #12374

jcsherin commented Sep 7, 2024 •

edited

Loading

jcsherin commented Sep 7, 2024

jayzhan211 Sep 7, 2024

jayzhan211 Sep 7, 2024

jcsherin Sep 17, 2024

jcsherin Sep 17, 2024

jcsherin Sep 17, 2024

jcsherin Sep 17, 2024

jcsherin Sep 17, 2024

	/// What [`DataType`] will be returned by this function, given the types of
	/// the arguments
	fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;

Add field trait method to WindowUDFImpl #12374

Are you sure you want to change the base?

Add field trait method to WindowUDFImpl #12374

Conversation

jcsherin commented Sep 7, 2024 • 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?

jcsherin commented Sep 7, 2024

jayzhan211 Sep 7, 2024

Choose a reason for hiding this comment

jayzhan211 Sep 7, 2024

Choose a reason for hiding this comment

jcsherin Sep 17, 2024

Choose a reason for hiding this comment

jcsherin Sep 17, 2024

Choose a reason for hiding this comment

jcsherin Sep 17, 2024

Choose a reason for hiding this comment

jcsherin Sep 17, 2024

Choose a reason for hiding this comment

jcsherin Sep 17, 2024

Choose a reason for hiding this comment

Add `field` trait method to `WindowUDFImpl` #12374

Add `field` trait method to `WindowUDFImpl` #12374

jcsherin commented Sep 7, 2024 •

edited

Loading