Support ORDER BY in AggregateUDF #9249

jayzhan211 · 2024-02-16T14:41:13Z

Which issue does this PR close?

Closes #8984.

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

alamb · 2024-02-20T07:46:58Z

Thank you @jayzhan211 -- I plan to review this tomorrow

alamb

Thank you so much for working on this @jayzhan211 . I had some comments about the API. Sorry for the delay in reviews

alamb · 2024-02-21T08:03:32Z

datafusion/physical-plan/src/windows/mod.rs

@@ -90,8 +90,15 @@ pub fn create_window_expr(
            ))
        }
        WindowFunctionDefinition::AggregateUDF(fun) => {
-            let aggregate =
-                udaf::create_aggregate_expr(fun.as_ref(), args, input_schema, name)?;
+            // TODO: Ordering not supported for Window UDFs yet


We should probably file a ticket for this as well -- I can do so as part of this PR's review

datafusion/core/src/physical_planner.rs

datafusion/expr/src/expr_fn.rs

datafusion/expr/src/function.rs

alamb

Thank you @jayzhan211 -- this pr is looking close. I still think the API needs a little tweaking as I mentioned in the comments

I apologize for the delay in review and feedback and I feel back about the back and forth. I am still catching up from last week. If you don't have a chance I hope to find some time to help push this PR through myself later this week

alamb · 2024-02-27T11:33:30Z

datafusion/core/tests/user_defined/user_defined_aggregates.rs

+        &self,
+        _arg: &DataType,
+        _sort_exprs: Vec<Expr>,
+        _schmea: Option<Schema>,


Suggested change

_schmea: Option<Schema>,

_schema: Option<Schema>,

datafusion/expr/src/expr_fn.rs

alamb · 2024-02-27T11:39:40Z

datafusion/expr/src/udaf.rs

@@ -255,15 +261,20 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {

    /// Return a new [`Accumulator`] that aggregates values for a specific
    /// group during query execution.
-    fn accumulator(&self, arg: &DataType) -> Result<Box<dyn Accumulator>>;
+    fn accumulator(


I think this API:

Doesn't need an owned Schema (passing in a Schema requires cloning which I think we can avoid

I am not sure why it needs an Option (why is it not always passed in)

Should have the documentation updated to explain what sort_exprs are and what schema is

Option is because Schema is only used for ordering, so if we don't care about ordering we can just pass None.

Or maybe we should return the schema in case we need it in the future not only for ordering? 🤔

Unless there is some reason not to always pass in schema I think we should always pass it in to make for an easier to use API. For example, if it is an Option the user would have to call unwrap() or check / error when they got sort exprs

alamb

Thanks @jayzhan211 -- I had a play around with the APIs and the code in this PR and I came up with this: jayzhan211#1

I see the challenge now that the create_accumulator function is invoked by the physical plan when there are no more Exprs / Schemas.

I think the next thing to do is figure out how ARRAY_AGG or some other built in aggregate function handles ordering and see if we can emulate the same thing

alamb · 2024-02-27T16:05:10Z

datafusion/expr/src/udaf.rs

@@ -255,15 +261,20 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {

    /// Return a new [`Accumulator`] that aggregates values for a specific
    /// group during query execution.
-    fn accumulator(&self, arg: &DataType) -> Result<Box<dyn Accumulator>>;
+    fn accumulator(


Unless there is some reason not to always pass in schema I think we should always pass it in to make for an easier to use API. For example, if it is an Option the user would have to call unwrap() or check / error when they got sort exprs

alamb · 2024-02-27T16:09:46Z

datafusion/expr/src/udaf.rs

+    /// group during query execution. sort_exprs is a list of ordering expressions,
+    /// and schema is used while ordering.


Suggested change

/// group during query execution. sort_exprs is a list of ordering expressions,

/// and schema is used while ordering.

/// group during query execution.

///

/// `arg`: the type of the argument to this accumulator

///

/// `sort_exprs`: contains a list of `Expr::SortExpr`s if the

/// aggregate is called with an explicit `ORDER BY`. For example,

/// `ARRAY_AGG(x ORDER BY y ASC)`. In this case, `sort_exprs` would contain `[y ASC]`

///

/// `schema` is the input schema to the aggregate

alamb · 2024-02-27T16:14:49Z

datafusion/expr/src/udaf.rs

@@ -276,6 +288,16 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
    fn create_groups_accumulator(&self) -> Result<Box<dyn GroupsAccumulator>> {
        not_impl_err!("GroupsAccumulator hasn't been implemented for {self:?} yet")
    }
+
+    /// Return the ordering expressions for the accumulator
+    fn sort_exprs(&self) -> Vec<Expr> {


I don't understand how the sort_exprs could be supplied by the UDAF. I would expect them to be provided to the UDAF based on what was in the query

jayzhan211 · 2024-03-01T06:35:14Z

Edit: I think we can just pass the logical ordering expr and not move the conversion downstream.

I think the problem here is that we need logical expr (ordering and schema) for accumulator. But the arguments pass into create_aggregate_expr are already physical expr. The reason why we need logical expr is that when we defining udaf, we only know logical expr but not physical expr.

https://github.com/apache/arrow-datafusion/blob/10d5f2df1c8f81a0d64a8644cfa429331f1e8ac3/datafusion/physical-plan/src/udaf.rs#L40-L45

I'm thinking about changing the function here to let us done the converting of logical to physical lately, so we can keep the logical expr for accumulator.

Other builtin accumulator don't have the problem because they are all defined in physical-expr unlike udaf.

@alamb Does the move of physical expression conversion make senses?

datafusion/core/tests/user_defined/user_defined_aggregates.rs

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 · 2024-03-25T12:43:34Z

datafusion/core/tests/user_defined/user_defined_aggregates.rs

+            data_type,
+            &ordering_dtypes,
+            ordering_req,
+            false,


ignore_nulls or other arguments are given in the accumulator provided by the user.

jayzhan211 · 2024-03-25T12:52:27Z

datafusion/proto/src/physical_plan/mod.rs

@@ -517,7 +517,10 @@ impl AsExecutionPlan for PhysicalPlanNode {
                                        }
                                        AggregateFunction::UserDefinedAggrFunction(udaf_name) => {
                                            let agg_udf = registry.udaf(udaf_name)?;
-                                            udaf::create_aggregate_expr(agg_udf.as_ref(), &input_phy_expr, &physical_schema, name)
+                                            // TODO: `order by` is not supported for UDAF yet
+                                            let sort_exprs = &[];


We need to convert Vec<PhysicalSortExprNode> to Vec<Expr>, not sure if we should include it in this PR or not.

jayzhan211 · 2024-03-26T00:08:37Z

@alamb I think it is ready for review

alamb · 2024-03-26T02:13:16Z

Thanks @jayzhan211 -- I will check it out, hopefully tomorrow

jayzhan211 · 2024-03-30T13:52:55Z

I plan to change first_value to udaf-based directly, the current implementation doesn't seem correct.

github-actions bot added sql logical-expr Logical plan and expressions physical-expr Physical Expressions core Core datafusion crate labels Feb 16, 2024

jayzhan211 marked this pull request as ready for review February 18, 2024 11:59

alamb reviewed Feb 21, 2024

View reviewed changes

jayzhan211 marked this pull request as draft February 21, 2024 13:14

github-actions bot added optimizer Optimizer rules substrait labels Feb 21, 2024

jayzhan211 commented Feb 21, 2024

View reviewed changes

datafusion/expr/src/function.rs Show resolved Hide resolved

jayzhan211 marked this pull request as ready for review February 21, 2024 15:17

jayzhan211 requested a review from alamb February 22, 2024 00:20

alamb mentioned this pull request Feb 26, 2024

DataFusion weekly project plan (Andrew Lamb) - Feb 26, 2024 #9345

Closed

9 tasks

alamb reviewed Feb 27, 2024

View reviewed changes

jayzhan211 marked this pull request as draft February 27, 2024 15:03

jayzhan211 marked this pull request as ready for review February 27, 2024 15:16

alamb mentioned this pull request Feb 27, 2024

Proposed API change for AggregateUDF to support ORDER BY jayzhan211/arrow-datafusion#1

Closed

alamb reviewed Feb 27, 2024

View reviewed changes

jayzhan211 marked this pull request as draft February 28, 2024 00:55

jayzhan211 force-pushed the udf-order-2 branch 2 times, most recently from 0dbfbad to 80eae25 Compare March 1, 2024 12:44

jayzhan211 commented Mar 1, 2024

View reviewed changes

datafusion/core/tests/user_defined/user_defined_aggregates.rs Show resolved Hide resolved

jayzhan211 marked this pull request as ready for review March 1, 2024 13:56

jayzhan211 requested a review from alamb March 1, 2024 13:56

jayzhan211 added 4 commits March 7, 2024 21:37

first draft

b94f70f

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

clippy fix

c743d13

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

cleanup

3a7e965

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

use one vector for ordering req

4917f56

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 added 8 commits March 7, 2024 21:37

change to ref

6fcdaac

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fix typo

c3512a6

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fix doc

092d46e

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fmt

8592e6b

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

move schema and logical ordering exprs

0f8fc24

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

remove redudant info

3185f9f

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

rename

3ecc772

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

cleanup

faadc63

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 force-pushed the udf-order-2 branch from 8b6ff2d to faadc63 Compare March 7, 2024 13:46

jayzhan211 marked this pull request as draft March 7, 2024 14:37

add ignore nulls

7e33910

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 marked this pull request as ready for review March 8, 2024 00:22

This was referenced Mar 8, 2024

DataFusion weekly project plan (Andrew Lamb) - March 4, 2024 #9453

Closed

DataFusion weekly project plan (Andrew Lamb) - March 11, 2024 #9555

Closed

This comment was marked as outdated.

Sign in to view

jayzhan211 marked this pull request as draft March 15, 2024 05:51

jayzhan211 added 2 commits March 25, 2024 19:59

Merge remote-tracking branch 'upstream/main' into udf-order-2

cfffcbf

fix conflict

6aaa15c

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 force-pushed the udf-order-2 branch from 27d5fdc to 6aaa15c Compare March 25, 2024 12:39

jayzhan211 commented Mar 25, 2024

View reviewed changes

jayzhan211 marked this pull request as ready for review March 25, 2024 13:30

jayzhan211 mentioned this pull request Mar 26, 2024

[Epic] Unify AggregateFunction Interface (remove built in list of AggregateFunction s), improve the system #8708

Open

alamb mentioned this pull request Mar 26, 2024

DataFusion weekly project plan (Andrew Lamb) - March 25, 2024 #9796

Closed

6 tasks

jayzhan211 marked this pull request as draft March 30, 2024 13:53

jayzhan211 mentioned this pull request Mar 31, 2024

Make FirstValue an UDAF, Change AggregateUDFImpl::accumulator signature, support ORDER BY for UDAFs #9874

Merged

jayzhan211 closed this in #9874 Apr 3, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support ORDER BY in AggregateUDF #9249

Support ORDER BY in AggregateUDF #9249

jayzhan211 commented Feb 16, 2024

alamb commented Feb 20, 2024

alamb left a comment

alamb Feb 21, 2024

alamb left a comment

alamb Feb 27, 2024

alamb Feb 27, 2024

jayzhan211 Feb 27, 2024

jayzhan211 Feb 27, 2024

alamb Feb 27, 2024

alamb left a comment

alamb Feb 27, 2024

alamb Feb 27, 2024

alamb Feb 27, 2024

jayzhan211 commented Mar 1, 2024 •

edited

Loading

This comment was marked as outdated.

jayzhan211 Mar 25, 2024 •

edited

Loading

jayzhan211 Mar 25, 2024 •

edited

Loading

jayzhan211 commented Mar 26, 2024

alamb commented Mar 26, 2024

jayzhan211 commented Mar 30, 2024 •

edited

Loading

		/// group during query execution. sort_exprs is a list of ordering expressions,
		/// and schema is used while ordering.

-    /// group during query execution. sort_exprs is a list of ordering expressions,
-    /// and schema is used while ordering.
+    /// group during query execution.
+    ///
+    /// `arg`: the type of the argument to this accumulator
+    ///
+    /// `sort_exprs`: contains a list of `Expr::SortExpr`s if the
+    /// aggregate is called with an explicit `ORDER BY`. For example,
+    /// `ARRAY_AGG(x ORDER BY y ASC)`. In this case, `sort_exprs` would contain `[y ASC]`
+    ///
+    /// `schema` is the input schema to the aggregate

Support ORDER BY in AggregateUDF #9249

Support ORDER BY in AggregateUDF #9249

Conversation

jayzhan211 commented Feb 16, 2024

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?

alamb commented Feb 20, 2024

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

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

jayzhan211 commented Mar 1, 2024 • edited Loading

This comment was marked as outdated.

jayzhan211 Mar 25, 2024 • edited Loading

Choose a reason for hiding this comment

jayzhan211 Mar 25, 2024 • edited Loading

Choose a reason for hiding this comment

jayzhan211 commented Mar 26, 2024

alamb commented Mar 26, 2024

jayzhan211 commented Mar 30, 2024 • edited Loading

jayzhan211 commented Mar 1, 2024 •

edited

Loading

jayzhan211 Mar 25, 2024 •

edited

Loading

jayzhan211 Mar 25, 2024 •

edited

Loading

jayzhan211 commented Mar 30, 2024 •

edited

Loading