Support User Defined Window Functions #6703

alamb · 2023-06-16T20:03:07Z

Which issue does this PR close?

Closes #5781
Closes #6617

Rationale for this change

This is the final PR that integrates all the work we have done this week with @mustafarepo and others to support User Defined Window Functions

It includes the parts of #6617 that have not been merged already (connecting the window functions to the various registries), and adds tests.

I tried to keep the size of this PR as small as possible -- 2/3 of the change is tests or examples -- but it is still large and I apologize to reviewers.

What changes are included in this PR?

Introduce WindowUDF and connect it through datafusion_expr and contexts
Add simple_udwf.rs example
Add tests in user_defined_functions.rs integration test

Are these changes tested?

yes

Are there any user-facing changes?

yes, new feature

alamb · 2023-06-20T20:39:51Z

Ok, I think this PR is basically ready for review (once #6690 is merged).

cc @stuartcarnie

alamb · 2023-06-20T20:40:26Z

datafusion/core/tests/user_defined/user_defined_window_functions.rs

+    PartitionEvaluator, ReturnTypeFunction, Signature, Volatility, WindowUDF,
+};
+
+/// A query with a window function evaluated over the entire partition


Here are the new tests

alamb · 2023-06-20T20:41:09Z

datafusion/core/tests/user_defined/user_defined_window_functions.rs

+
+/// Basic stateful user defined window function
+#[tokio::test]
+async fn test_stateful_udwf() {


I couldn't think of anything else to test with a stateful function other than update_state is called. I would be interested in anyone else's opinions in this matter

alamb · 2023-06-21T13:29:20Z

datafusion/core/tests/user_defined_integration.rs

+// specific language governing permissions and limitations
+// under the License.
+
+/// Run all tests that are found in the `user_defined` directory


All the user defined tests now run as part of the same test binaries user_defined_integration

alamb · 2023-06-21T13:30:01Z

datafusion-examples/examples/simple_udwf.rs

@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one


This is a user facing example of how to use User Defined Window Functions

alamb · 2023-06-21T13:30:25Z

datafusion/core/src/physical_plan/windows/mod.rs

@@ -206,6 +212,71 @@ fn create_built_in_window_expr(
    })
 }

+/// Creates a `BuiltInWindowFunctionExpr` suitable for a user defined window function


This mirrors the code and structure for AggregateUDF

datafusion-examples/examples/simple_udwf.rs

mustafasrepo · 2023-06-22T11:30:45Z

datafusion-examples/examples/simple_udwf.rs

+        values: &[ArrayRef],
+        range: &std::ops::Range<usize>,
+    ) -> Result<ScalarValue> {
+        //println!("evaluate_inside_range(). range: {range:#?}, values: {values:#?}");


I think this is leftover

mustafasrepo · 2023-06-22T11:34:29Z

datafusion-examples/examples/simple_udwf.rs

+    WindowUDF {
+        name: String::from("smooth_it"),
+        // it will take 1 arguments -- the column to smooth
+        signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),


It doesn't seem to be effecting result, but I don't understand why signature takes DataType::Int32?

mustafasrepo · 2023-06-22T11:36:31Z

datafusion/core/src/dataframe.rs

@@ -218,6 +218,14 @@ impl DataFrame {
        Ok(DataFrame::new(self.session_state, plan))
    }

+    /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+    pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {


Maybe we can add an example usage to the docstring. such as in the pub fn aggregate. We can do so in following PRs also.

Good call -- I tried to make one, and it turns out to be non trivial. I will do it in a follow on PR

Here is what I was thinking: #6746 if you like it I will polish it up

datafusion/core/src/execution/context.rs

datafusion/core/tests/user_defined/user_defined_window_functions.rs

mustafasrepo · 2023-06-22T12:17:24Z

datafusion/execution/src/task.rs

@@ -92,13 +100,15 @@ impl TaskContext {
            config.set(&k, &v)?;
        }
        let session_config = SessionConfig::from(config);
+        let window_functions = HashMap::new();


I guess, since this API is deprecated anyway, you didn't update its arguments.

Yes, exactly. I also tried to lead people to SessionContext::task_ctx with another doc comment above.

datafusion/expr/src/expr_fn.rs

datafusion/expr/src/udwf.rs

mustafasrepo

Thanks @alamb. I left some inline comments, related to comments, etc. Other than these, This PR is LGTM!. Thanks for this feature.

Co-authored-by: Mustafa Akur <[email protected]>

…on into alamb/udfw_for_real

Co-authored-by: Mustafa Akur <[email protected]>

…on into alamb/udfw_for_real

alamb

Thank you for the thorough review @mustafasrepo.

alamb · 2023-06-22T15:33:55Z

datafusion-examples/examples/simple_udwf.rs

+    WindowUDF {
+        name: String::from("smooth_it"),
+        // it will take 1 arguments -- the column to smooth
+        signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),


alamb · 2023-06-22T15:37:44Z

datafusion/execution/src/task.rs

@@ -92,13 +100,15 @@ impl TaskContext {
            config.set(&k, &v)?;
        }
        let session_config = SessionConfig::from(config);
+        let window_functions = HashMap::new();


Yes, exactly. I also tried to lead people to SessionContext::task_ctx with another doc comment above.

datafusion/expr/src/expr_fn.rs

alamb · 2023-06-22T15:44:26Z

datafusion/core/src/dataframe.rs

@@ -218,6 +218,14 @@ impl DataFrame {
        Ok(DataFrame::new(self.session_state, plan))
    }

+    /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+    pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {


Good call -- I tried to make one, and it turns out to be non trivial. I will do it in a follow on PR

stuartcarnie

Great stuff, @alamb – apologies, I didn't realise my comments were tied to a review, so they are late, but all just giving praise 👍🏻

stuartcarnie · 2023-06-22T05:16:49Z

datafusion-examples/README.md

@@ -57,6 +57,7 @@ cargo run --example csv_sql
 - [`rewrite_expr.rs`](examples/rewrite_expr.rs): Define and invoke a custom Query Optimizer pass
 - [`simple_udaf.rs`](examples/simple_udaf.rs): Define and invoke a User Defined Aggregate Function (UDAF)
 - [`simple_udf.rs`](examples/simple_udf.rs): Define and invoke a User Defined (scalar) Function (UDF)
+- [`simple_udfw.rs`](examples/simple_udwf.rs): Define and invoke a User Defined Window Function (UDWF)


These examples are 💯

stuartcarnie · 2023-06-22T05:19:12Z

datafusion-examples/examples/simple_udwf.rs

The documentation in this example is excellent 👏🏻

stuartcarnie · 2023-06-22T05:20:14Z

datafusion/core/src/dataframe.rs

+    /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+    pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {


I was just looking for a similar API today – great to know it is coming

@stuartcarnie I wonder if you have any thoughts about the API here for making window functions: #6746 (it is somewhat complex at the moment)

alamb force-pushed the alamb/udfw_for_real branch from 18a034c to 6f87009 Compare June 20, 2023 19:56

github-actions bot added core Core DataFusion crate logical-expr Logical plan and expressions physical-expr Physical Expressions sql SQL Planner labels Jun 20, 2023

alamb added the api change Changes the API exposed to users of the crate label Jun 20, 2023

alamb mentioned this pull request Jun 20, 2023

Implement serialization for UDWF as window function in plan protobuf #4584 #6733

Closed

alamb force-pushed the alamb/udfw_for_real branch from 6f87009 to 5fe955b Compare June 20, 2023 20:18

github-actions bot added optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) labels Jun 20, 2023

alamb force-pushed the alamb/udfw_for_real branch from 5fe955b to 1427b1c Compare June 20, 2023 20:39

alamb commented Jun 20, 2023

View reviewed changes

Support User Defined Window Functions

fa0e40f

alamb force-pushed the alamb/udfw_for_real branch from 1427b1c to fa0e40f Compare June 21, 2023 13:27

github-actions bot removed the physical-expr Physical Expressions label Jun 21, 2023

alamb commented Jun 21, 2023

View reviewed changes

alamb marked this pull request as ready for review June 21, 2023 13:33

alamb mentioned this pull request Jun 21, 2023

RFC: User Defined Window Functions #6617

Closed