Allow `AggregateUDF` to define retractable batch , implement sliding window functions #6671

alamb · 2023-06-14T19:56:22Z

Which issue does this PR close?

Closes #6611

Rationale for this change

When used as window functions, AggregateUDF can not use the retract_batch API for sliding window calculations

What changes are included in this PR?

Add a new supports_retract_batch in Accumulator trait (to mirror the API added by @mustafasrepo for Window Functions)
If the Accumulator reports it supports retract_batch, use the sliding accumulator code

Are these changes tested?

Yes

Are there any user-facing changes?

New API, though they should all be backwards compatible

Rationale: The default implementation of the `Accumulator` trait returns an error for the `retract_batch` API.

alamb · 2023-06-14T21:15:18Z

datafusion/core/src/physical_plan/udaf.rs

+
+    /// Returns true if this can support sliding accumulators
+    pub fn retractable(&self) -> Result<bool> {
+        Ok((self.fun.accumulator)(&self.data_type)?.supports_retract_batch())


it is non ideal that we need to instanatiate the accumulator just to check if it supports retract batch, but that is consistent with the window function API so I think that is ok

Maybe we can fix this and its window function counterpart in a refactor. We discussed this instantiation thing in a meeting recently and we are not super happy with it either. Let's get the basics in and then we can work on the rough edges.

Another other formulation I tried looked like the following. The upside is it avoided the instantiation, but the downside is that it an API change and is not consistent with the window functions

pub struct AggregateUDF { /// name pub name: String, /// Signature (input arguments) pub signature: Signature, /// Return type pub return_type: ReturnTypeFunction, /// Return an accumulator without retract batch pub accumulator: AccumulatorFunctionImplementation, /// Return an accumulator with retract batch pub retractable_accumulator: Option<AccumulatorFunctionImplementation>, <---- this is added /// the accumulator's state's description as a function of the return type pub state_type: StateTypeFunction, }

alamb · 2023-06-14T21:16:15Z

datafusion/core/tests/user_defined_aggregates.rs

+
+    let TestContext { ctx, test_state } = TestContext::new_with_test_state(test_state);
    let sql = "SELECT time_sum(time) OVER(ORDER BY time ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) as time_sum from t";
+    // TODO: It is not clear why this is a different value than when retract batch is used


I am not sure why the UDWF gets different answers when implemented with retract_batch 🤔 maybe @mustafasrepo can offer some hints or insights

Accumulators when they have window frame startings different than UNBOUNDED PRECEDING such as 1 preceding, need to implement retract_batch method, to be able to run correctly (If this method is lacking there is no way to calculate result). Consider the query

SELECT SUM(a) OVER(ORDER BY a ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) as sum_a from t

First sum value will be the sum of rows between [0, 1),
Second sum value will be the sum of rows between [0, 2),
Third sum value will be the sum of rows between [1, 3), etc.
Since accumulator keeps in its state, running sum
for first sum we add to the state sum value between [0, 1)
for second sum we add to the state sum value between [1, 2) ([0, 1) is already in the state sum, hence running sum will cover [0, 2) range)
for third sum we add to the state sum value between [2, 3) ([0, 2) is already in the state sum). Also we need to retract values between [0, 1) by this way we can obtain sum between [1, 3) which is indeed the apropriate range.

When we use UNBOUNDED PRECEDING in the query starting index will always be 0 for the desired range. Hence we will never call retract_batch method. In this case having retract_batch is not a requirement.

This approach is a a bit different than window function approach. In window function(when they use frame) they get all the desired range during evaluation. If we were to have a window_function sum we could have calculated sum from scratch each time inside given range. However, for accumulators this is the case, and above query should give error if retract_batch is not implemented. During review I have pinpointed the section where I think this error is prevented

Thank you -- I am glad I asked!

I thought about this some more and have made queries with unbounded windows without retract_batch error as you suggested in 999866e -- much better to get an error rather than the wrong answer.

It seems to me that there is no fundamental reason DataFusion couldn't support sliding windows with user defined aggregate functions without retract_batch -- for example it could make a new Accumulator instance for each row and feed in the entire widow.

However, given it is likely to be much more efficient to implement retract_batch, that is what I think we should encourage people to do. If someone really needs to use a UDAF without retract_batch for a sliding window, we can look into implementing the "full recompute" solution at that time

I thought about this some more and have made queries with unbounded windows without retract_batch error as you suggested in 999866e -- much better to get an error rather than the wrong answer.

It seems to me that there is no fundamental reason DataFusion couldn't support sliding windows with user defined aggregate functions without retract_batch -- for example it could make a new Accumulator instance for each row and feed in the entire widow.

However, given it is likely to be much more efficient to implement retract_batch, that is what I think we should encourage people to do. If someone really needs to use a UDAF without retract_batch for a sliding window, we can look into implementing the "full recompute" solution at that time

Agreed. Also, users can define their UDWF (User defined window function). When UDWF's uses_window_frame flag is true. evaluate method of PartitionEvaluator is like "full recompute" approach.

alamb · 2023-06-14T21:17:05Z

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

-            order_by,
-            window_frame,
-        )),
+        WindowFunction::AggregateUDF(fun) => {


The major difference with this PR and what @stuartcarnie did is that this PR still allows AggregateUDFs to run as window functions even if they don't implement retract_batch

Update: it turns out it is not correct to run AggregateUDFs as sliding window function if they don't implement retract_batch which I have made clear in comments and the tests

alamb · 2023-06-14T21:19:54Z

datafusion/expr/src/accumulator.rs

        ))
    }

+    /// Does the accumulator support incrementally updating its value


Here is a new API addition, but it should be backwards compatible

stuartcarnie

Looks great, @alamb. I updated my IOx code to point to this branch, and my existing UDAFs work great 👍🏻

mustafasrepo · 2023-06-15T06:43:43Z

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

+            let aggregate =
+                udaf::create_aggregate_expr(fun.as_ref(), args, input_schema, name)?;
+
+            if !unbounded_window && aggregate.retractable()? {


I think this line should be

Suggested change

if !unbounded_window && aggregate.retractable()? {

if !unbounded_window {

If aggregate doesn't implement retract we should get an error. There is no way to run aggregator whose start is different than (UNBOUNDED PRECEDING) without retract_batch method.

I will make this change -- thank you @mustafasrepo -- I didn't appreciate that there is no correct way to run the AggregateUDF correctly without retract_batch. Thus I think the correct solution is to return an error when such a combination is attempted. I will update the PR to reflect this

in 999866e

mustafasrepo · 2023-06-15T07:21:46Z

Also I think with the introduction of supports_retract_batch method, supports_bounded_execution method is no longer necessary for the AggregateExpr trait. I filed a small PR on top this PR to show required changes for removing supports_bounded_execution method from AggregateExpr trait.

More detailed explanation and PR can be found in #PR13

Other than the points in my review. This PR is LGTM!. With these changes users can write accumulators with retract_batch method, then can use it in aggregate and/or window queries.

…_batch is used as a sliding accumulator

alamb marked this pull request as draft June 14, 2023 19:56

github-actions bot added core Core DataFusion crate logical-expr Logical plan and expressions labels Jun 14, 2023

alamb force-pushed the alamb/sliding_udf_real branch from 6282c54 to 5198ae8 Compare June 14, 2023 19:59

alamb changed the title ~~Allow AggregateUDF to define retractable batch @alamb~~ Allow AggregateUDF to define retractable batch Jun 14, 2023

This was referenced Jun 14, 2023

Support sliding windows in AggregateUDF / User Defined Window Functions #6611

Closed

Minor: Add tests for User Defined Aggregate functions #6669

Merged

feat: support sliding window accumulators

d53d6c4

Rationale: The default implementation of the `Accumulator` trait returns an error for the `retract_batch` API.

alamb force-pushed the alamb/sliding_udf_real branch from 5198ae8 to 0d5adc4 Compare June 14, 2023 21:14

Allow AggregateUDF to define retractable batch

579b4d9

alamb force-pushed the alamb/sliding_udf_real branch from 0d5adc4 to 579b4d9 Compare June 14, 2023 21:19

alamb commented Jun 14, 2023

View reviewed changes

alamb changed the title ~~Allow AggregateUDF to define retractable batch~~ Allow AggregateUDF to define retractable batch , efficiently implement window functions Jun 14, 2023

alamb marked this pull request as ready for review June 14, 2023 21:20

stuartcarnie approved these changes Jun 14, 2023

View reviewed changes

mustafasrepo reviewed Jun 15, 2023

View reviewed changes

alamb added 2 commits June 15, 2023 10:42

Merge remote-tracking branch 'apache/main' into alamb/sliding_udf_real

e833914

Return error rather than wrong results when aggregate without retract…

999866e

…_batch is used as a sliding accumulator

alamb changed the title ~~Allow AggregateUDF to define retractable batch , efficiently implement window functions~~ Allow AggregateUDF to define retractable batch , implement sliding window functions Jun 15, 2023

alamb mentioned this pull request Jun 15, 2023

Replace supports_bounded_execution with supports_retract_batch alamb/datafusion#13

Closed

Merge remote-tracking branch 'apache/main' into alamb/sliding_udf_real

f1949a7

alamb merged commit 84e4977 into apache:main Jun 15, 2023

mustafasrepo mentioned this pull request Jun 16, 2023

Replace supports_bounded_execution with supports_retract_batch #6695

Merged

	if !unbounded_window && aggregate.retractable()? {
	if !unbounded_window {

Allow AggregateUDF to define retractable batch , implement sliding window functions #6671

Allow AggregateUDF to define retractable batch , implement sliding window functions #6671

Uh oh!

Conversation

alamb commented Jun 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mustafasrepo Jun 15, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stuartcarnie left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mustafasrepo commented Jun 15, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Allow `AggregateUDF` to define retractable batch , implement sliding window functions #6671

Allow `AggregateUDF` to define retractable batch , implement sliding window functions #6671

alamb commented Jun 14, 2023 •

edited

Loading

mustafasrepo Jun 15, 2023 •

edited

Loading

mustafasrepo commented Jun 15, 2023 •

edited

Loading