Skip to content

Support User Defined Window Functions - #6703

Merged
alamb merged 9 commits into
apache:mainfrom
alamb:alamb/udfw_for_real
Jun 22, 2023
Merged

Support User Defined Window Functions#6703
alamb merged 9 commits into
apache:mainfrom
alamb:alamb/udfw_for_real

Conversation

@alamb

@alambalamb commented Jun 16, 2023•
edited
Loading

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes#5781
Closes#6617

Rationale for this change

See #5781

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?

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

Are these changes tested?

yes

Are there any user-facing changes?

yes, new feature

@alamb
alambforce-pushed the alamb/udfw_for_real branch from 18a034c to 6f87009CompareJune 20, 2023 19:56
@github-actionsgithub-actionsBot added core Core DataFusion crate logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates sql SQL Planner labels Jun 20, 2023
@alambalamb added the api change Changes the API exposed to users of the crate label Jun 20, 2023
@alambalamb mentioned this pull request Jun 20, 2023
Closed
@alamb
alambforce-pushed the alamb/udfw_for_real branch from 6f87009 to 5fe955bCompareJune 20, 2023 20:18
@github-actionsgithub-actionsBot added optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) labels Jun 20, 2023
@alamb
alambforce-pushed the alamb/udfw_for_real branch from 5fe955b to 1427b1cCompareJune 20, 2023 20:39
@alamb

alamb commented Jun 20, 2023•
edited
Loading

Copy link
Copy Markdown
ContributorAuthor

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

cc @stuartcarnie

alamb
alamb commented Jun 20, 2023
View reviewed changes
Comment threaddatafusion/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

@alambalambJun 20, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

Here are the new tests

Comment threaddatafusion/core/tests/user_defined/user_defined_window_functions.rs

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

@alambalambJun 20, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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
alambforce-pushed the alamb/udfw_for_real branch from 1427b1c to fa0e40fCompareJune 21, 2023 13:27
@github-actionsgithub-actionsBot removed the physical-expr Changes to the physical-expr crates label Jun 21, 2023
alamb
alamb commented Jun 21, 2023
View reviewed changes
Comment threaddatafusion/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

@alambalambJun 21, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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

Comment threaddatafusion-examples/examples/simple_udwf.rs
@@ -0,0 +1,213 @@
// Licensed to the Apache Software Foundation (ASF) under one

@alambalambJun 21, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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

Comment threaddatafusion/core/src/physical_plan/windows/mod.rs
})
}

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

@alambalambJun 21, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

This mirrors the code and structure for AggregateUDF

@alamb
alamb marked this pull request as ready for review June 21, 2023 13:33
@alambalamb mentioned this pull request Jun 21, 2023
Closed
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion-examples/examples/simple_udwf.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion-examples/examples/simple_udwf.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion-examples/examples/simple_udwf.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion-examples/examples/simple_udwf.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion-examples/examples/simple_udwf.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion-examples/examples/simple_udwf.rs Outdated
values: &[ArrayRef],
range: &std::ops::Range<usize>,
) -> Result<ScalarValue> {
//println!("evaluate_inside_range(). range: {range:#?}, values: {values:#?}");

@mustafasrepomustafasrepoJun 22, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I think this is leftover

mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion-examples/examples/simple_udwf.rs Outdated
WindowUDF {
name: String::from("smooth_it"),
// it will take 1 arguments -- the column to smooth
signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),

@mustafasrepomustafasrepoJun 22, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

@alambalambJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

I thought it was the default type that the CSV parser picked, but apparently not:

❯ describe './datafusion/core/tests/data/cars.csv';
describe './datafusion/core/tests/data/cars.csv';
+-------------+-----------------------------+-------------+
| column_name | data_type | is_nullable |
+-------------+-----------------------------+-------------+
| car | Utf8 | YES |
| speed | Float64 | YES |
| time | Timestamp(Nanosecond, None) | YES |
+-------------+-----------------------------+-------------+3 rows inset. Query took 0.031 seconds.

I'lll update this to take a Float instead.

mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion/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> {

@mustafasrepomustafasrepoJun 22, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

@alambalambJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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

@alambalambJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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

mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion/core/src/execution/context.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion/core/src/execution/context.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion/core/tests/user_defined/user_defined_window_functions.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion/execution/src/task.rs
config.set(&k, &v)?;
}
let session_config = SessionConfig::from(config);
let window_functions = HashMap::new();

@mustafasrepomustafasrepoJun 22, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

@alambalambJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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

mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion/expr/src/expr_fn.rs Outdated
mustafasrepo
mustafasrepo reviewed Jun 22, 2023
View reviewed changes
Comment threaddatafusion/expr/src/udwf.rs Outdated
mustafasrepo
mustafasrepo approved these changes Jun 22, 2023
View reviewed changes

@mustafasrepomustafasrepo left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

alamband others added 7 commits June 22, 2023 11:34
@alamb@mustafasrepo
Apply suggestions from code review
a29e431
Co-authored-by: Mustafa Akur <106137913+mustafasrepo@users.noreply.github.com>
@alamb
Remove left over println
209144a
@alamb@mustafasrepo
Apply suggestions from code review
b23703e
Co-authored-by: Mustafa Akur <106137913+mustafasrepo@users.noreply.github.com>
@alamb
Merge branch 'alamb/udfw_for_real' of github.com:alamb/arrow-datafusi…
2ad2e76
…on into alamb/udfw_for_real
@alamb@mustafasrepo
Apply suggestions from code review
cd94076
Co-authored-by: Mustafa Akur <106137913+mustafasrepo@users.noreply.github.com>
@alamb
Merge branch 'alamb/udfw_for_real' of github.com:alamb/arrow-datafusi…
a8f848f
…on into alamb/udfw_for_real
@alamb
Merge remote-tracking branch 'apache/main' into alamb/udfw_for_real
d3e62e3
alamb
alamb commented Jun 22, 2023
View reviewed changes

@alambalamb left a comment

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

Thank you for the thorough review @mustafasrepo.

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

@alambalambJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

I thought it was the default type that the CSV parser picked, but apparently not:

❯ describe './datafusion/core/tests/data/cars.csv';
describe './datafusion/core/tests/data/cars.csv';
+-------------+-----------------------------+-------------+
| column_name | data_type | is_nullable |
+-------------+-----------------------------+-------------+
| car | Utf8 | YES |
| speed | Float64 | YES |
| time | Timestamp(Nanosecond, None) | YES |
+-------------+-----------------------------+-------------+3 rows inset. Query took 0.031 seconds.

I'lll update this to take a Float instead.

Comment threaddatafusion/execution/src/task.rs
config.set(&k, &v)?;
}
let session_config = SessionConfig::from(config);
let window_functions = HashMap::new();

@alambalambJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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

Comment threaddatafusion/expr/src/expr_fn.rs Outdated
Comment threaddatafusion/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> {

@alambalambJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

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

@alamb
alamb merged commit b1b8c9c into apache:mainJun 22, 2023
@alambalamb mentioned this pull request Jun 22, 2023
Closed
stuartcarnie
stuartcarnie reviewed Jun 23, 2023
View reviewed changes

@stuartcarniestuartcarnie left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

Comment threaddatafusion-examples/README.md
- [`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)

@stuartcarniestuartcarnieJun 22, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

These examples are 💯

Comment threaddatafusion-examples/examples/simple_udwf.rs

@stuartcarniestuartcarnieJun 22, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The documentation in this example is excellent 👏🏻

Comment threaddatafusion/core/src/dataframe.rs
Comment on lines +221 to +222
/// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {

@stuartcarniestuartcarnieJun 22, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

@alambalambJun 23, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

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

@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
alamb deleted the alamb/udfw_for_real branch June 23, 2023 14:57
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

2 more reviewers
@stuartcarniestuartcarniestuartcarnie left review comments
@mustafasrepomustafasrepomustafasrepo approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

api changeChanges the API exposed to users of the cratecoreCore DataFusion cratelogical-exprLogical plan and expressionsoptimizerOptimizer rulessqlSQL PlannersqllogictestSQL Logic Tests (.slt)

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

User defined window functions

3 participants

@alamb@stuartcarnie@mustafasrepo