Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add user defined window function support #880

Merged
merged 11 commits into from
Sep 30, 2024
Merged

Conversation

timsaucer
Copy link
Contributor

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
Copy link
Contributor Author

timsaucer commented Sep 21, 2024

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

@timsaucer
Copy link
Contributor Author

timsaucer commented Sep 21, 2024

Also should add an example. done

@timsaucer
Copy link
Contributor Author

timsaucer commented Sep 21, 2024

Also need to update online documentation done

@timsaucer timsaucer marked this pull request as ready for review September 22, 2024 17:28
@@ -0,0 +1,270 @@
# Licensed to the Apache Software Foundation (ASF) under one
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added #888

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

examples/python-udwf.py Outdated Show resolved Hide resolved
python/datafusion/tests/test_udwf.py Outdated Show resolved Hide resolved
docs/source/user-guide/common-operations/udf-and-udfa.rst Outdated Show resolved Hide resolved
docs/source/user-guide/common-operations/udf-and-udfa.rst Outdated Show resolved Hide resolved
Comment on lines 256 to 261
|``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`` |
Copy link
Contributor

@Michael-J-Ward Michael-J-Ward Sep 26, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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`` |
+------------------------+--------------------------------+------------------+---------------------------+

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Outdated Show resolved Hide resolved
python/datafusion/udf.py Outdated Show resolved Hide resolved
python/datafusion/udf.py Outdated Show resolved Hide resolved
use crate::utils::parse_volatility;

#[derive(Debug)]
struct RustPartitionEvaluator {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why RustPartitionEvaluator instead of PyPartitionEvaluator?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

def memoize(self) -> None:
"""Perform a memoize operation to improve performance.

When the window frame has a fixed beginning (e.g UNBOUNDED
Copy link
Contributor

@Michael-J-Ward Michael-J-Ward Sep 26, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

}
}

pub struct MultiColumnWindowUDF {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 Qs:

  1. Does this need to be public?
  2. Why motivation behind the name? The Multi makes me expect to see a SingleColumnWindowUDF counterpart.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. 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.
  2. 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.

@timsaucer
Copy link
Contributor Author

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

Copy link
Contributor

@Michael-J-Ward Michael-J-Ward left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great, thanks @timsaucer!

@timsaucer timsaucer merged commit 022e4b3 into apache:main Sep 30, 2024
15 checks passed
@timsaucer timsaucer deleted the feature/udwf branch September 30, 2024 11:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add support for user defined window functions
3 participants