Add user defined window function support #880

timsaucer · 2024-09-21T00:58:07Z

Which issue does this PR close?

Closes #845

Rationale for this change

Currently users can only create user defined scalar functions and user defined window functions. This enables authoring user defined window functions. It is a core functionality that has been missing on the python library.

What changes are included in this PR?

Designed similar to the UDAF, we have an abstract class for UDWF that users must inherit from.

Are there any user-facing changes?

Adds new udwf() function and WindowUDF class.

timsaucer · 2024-09-21T00:58:18Z

~~I'll move it out of draft after adding unit tests.~~

timsaucer · 2024-09-21T00:58:26Z

~~Also should add an example.~~ done

timsaucer · 2024-09-21T15:15:06Z

~~Also need to update online documentation~~ done

…requires python side work.

…eterize it

… and also multiple columns

mesejo · 2024-09-26T16:11:49Z

examples/python-udwf.py

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


As a side comment: Should we have automated tests of the examples? Nothing fancy; simple execution of the scripts should be enough

That's a great idea. We already do it for the tpc-h examples. How about we open an issue to run all the examples in CI?

examples/python-udwf.py

python/datafusion/tests/test_udwf.py

docs/source/user-guide/common-operations/udf-and-udfa.rst

Michael-J-Ward · 2024-09-26T19:49:49Z

python/datafusion/udf.py

+    |``uses_window_frame``|``supports_bounded_execution``|``include_rank``|function_to_implement|
+    |---|---|----|----|
+    |False (default)      |False (default)               |False (default) | ``evaluate_all``           |
+    |False                |True                          |False           | ``evaluate``               |
+    |False                |True/False                    |True            | ``evaluate_all_with_rank`` |
+    |True                 |True/False                    |True/False      | ``evaluate``               |


Does the table render properly for you? I had to update to this suggestion in order for the table to render properly.

Suggested change

|``uses_window_frame``|``supports_bounded_execution``|``include_rank``|function_to_implement|

|---|---|----|----|

|False (default) |False (default) |False (default) | ``evaluate_all`` |

|False |True |False | ``evaluate`` |

|False |True/False |True | ``evaluate_all_with_rank`` |

|True |True/False |True/False | ``evaluate`` |

+------------------------+--------------------------------+------------------+---------------------------+

| ``uses_window_frame`` | ``supports_bounded_execution`` | ``include_rank`` | function_to_implement |

+========================+================================+==================+===========================+

| False (default) | False (default) | False (default) | ``evaluate_all`` |

+------------------------+--------------------------------+------------------+---------------------------+

| False | True | False | ``evaluate`` |

+------------------------+--------------------------------+------------------+---------------------------+

| False | True/False | True | ``evaluate_all_with_rank``|

+------------------------+--------------------------------+------------------+---------------------------+

| True | True/False | True/False | ``evaluate`` |

+------------------------+--------------------------------+------------------+---------------------------+

The tables rendered for me in VS code but they didn't render in the generated online documentation. I've adopted your suggestions in my latest push.

python/datafusion/udf.py

Michael-J-Ward · 2024-09-26T20:13:04Z

src/udwf.rs

+use crate::utils::parse_volatility;
+
+#[derive(Debug)]
+struct RustPartitionEvaluator {


Why RustPartitionEvaluator instead of PyPartitionEvaluator?

To be completely honest, I just copied the naming convention from the aggregate

Michael-J-Ward · 2024-09-26T20:17:49Z

python/datafusion/udf.py

+    def memoize(self) -> None:
+        """Perform a memoize operation to improve performance.
+
+        When the window frame has a fixed beginning (e.g UNBOUNDED


Q for my own understanding:

Should this be "When the window frame has a fixed bound (beginning or end)..."?

I ask because the docs use last_value as an example, and I don't see how last_value could benefit from memoize because of a fixed beginning.

(If I'm correct, I'll go fix the docs upstream)

https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.PartitionEvaluator.html#method.memoize

You're right! last_value does not memoize the results, as evidenced here: https://github.com/apache/datafusion/blob/main/datafusion/physical-expr/src/window/nth_value.rs#L200

Michael-J-Ward · 2024-09-26T20:27:27Z

src/udwf.rs

+    }
+}
+
+pub struct MultiColumnWindowUDF {


2 Qs:

Does this need to be public?

Why motivation behind the name? The Multi makes me expect to see a SingleColumnWindowUDF counterpart.

I wanted to make it public in case someone does build extension libraries on top of datafusion-python, but I haven't thought too deeply about which structs need to be public vs private.

I called it MultiColumnWindowUDF as a contrast to SimpleWindowUDF which only takes a single column of input. I'm open to suggestion on a better name.

docs/source/user-guide/common-operations/udf-and-udfa.rst

timsaucer · 2024-09-27T17:22:30Z

Thank you for the reviews! I’m pretty slammed right now but should get to this early next week

Michael-J-Ward

Looks great, thanks @timsaucer!

timsaucer marked this pull request as ready for review September 22, 2024 17:28

timsaucer mentioned this pull request Sep 22, 2024

Enhance udf to take additional non-expr arguments #822

Closed

timsaucer added 9 commits September 23, 2024 08:34

Adding PyWindowUDF and implementing PartitionEvaluator for it. Still …

94470d8

…requires python side work.

Add python wrappers for UDWF

f306dde

adding unit tests for user defined window functions

a661ef9

Change udwf() to take an instance rather than a class so we can param…

f530f0e

…eterize it

Pass multiple arrays for udwf evaluate so we can capture the order_by…

d63f84f

… and also multiple columns

Update udwf to take multiple input columns

b8baad4

Add user exampe for UDWF

326dacd

Update template for how values are passed to update

d03ce13

Add user documentation for UDWF

3dfa330

timsaucer force-pushed the feature/udwf branch from 3fb3848 to 3dfa330 Compare September 23, 2024 12:34

mesejo approved these changes Sep 26, 2024

View reviewed changes

Michael-J-Ward reviewed Sep 26, 2024

View reviewed changes

docs/source/user-guide/common-operations/udf-and-udfa.rst Outdated Show resolved Hide resolved

Updating documentation per PR review

42092f7

timsaucer mentioned this pull request Sep 29, 2024

Run all user examples in CI #888

Open

Small documentation update

0fc0b74

Michael-J-Ward approved these changes Sep 30, 2024

View reviewed changes

timsaucer merged commit 022e4b3 into apache:main Sep 30, 2024
15 checks passed

timsaucer deleted the feature/udwf branch September 30, 2024 11:43

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add user defined window function support #880

Add user defined window function support #880

timsaucer commented Sep 21, 2024

timsaucer commented Sep 21, 2024 •

edited

Loading

timsaucer commented Sep 21, 2024 •

edited

Loading

timsaucer commented Sep 21, 2024 •

edited

Loading

mesejo Sep 26, 2024

timsaucer Sep 29, 2024

timsaucer Sep 29, 2024

mesejo Sep 29, 2024

Michael-J-Ward Sep 26, 2024 •

edited

Loading

timsaucer Sep 29, 2024

Michael-J-Ward Sep 26, 2024

timsaucer Sep 29, 2024

Michael-J-Ward Sep 26, 2024 •

edited

Loading

timsaucer Sep 29, 2024

Michael-J-Ward Sep 26, 2024

timsaucer Sep 29, 2024

timsaucer commented Sep 27, 2024

Michael-J-Ward left a comment

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

Add user defined window function support #880

Add user defined window function support #880

Conversation

timsaucer commented Sep 21, 2024

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

timsaucer commented Sep 21, 2024 • edited Loading

timsaucer commented Sep 21, 2024 • edited Loading

timsaucer commented Sep 21, 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

Michael-J-Ward Sep 26, 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

Michael-J-Ward Sep 26, 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

timsaucer commented Sep 27, 2024

Michael-J-Ward left a comment

Choose a reason for hiding this comment

timsaucer commented Sep 21, 2024 •

edited

Loading

timsaucer commented Sep 21, 2024 •

edited

Loading

timsaucer commented Sep 21, 2024 •

edited

Loading

Michael-J-Ward Sep 26, 2024 •

edited

Loading

Michael-J-Ward Sep 26, 2024 •

edited

Loading