Convert first, last aggregate function to UDAF #10648

mustafasrepo · 2024-05-24T11:32:05Z

Which issue does this PR close?

Closes #9957
Part of #8708

Rationale for this change

This work is a continuation of the work done by @jayzhan211 for UDAF support.
This PR converts FIRST_VALUE_AGG, LAST_VALUE_AGG Aggregate functions to the UDAF.

As part of this PR, I have also updated refactored OptimizeAggregateOrder rule such that it no longer depends on the
.as_any().downcast_ref for to work (e.g it is not overfitted to the FirstValue, LastValue aggregate expressions). File name is changed from convert_first_last.rs to update_aggr_exprs.rs also.

What changes are included in this PR?

Are these changes tested?

Yes

Are there any user-facing changes?

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

jayzhan211

This PR is really impressive!

jayzhan211 · 2024-05-24T11:53:52Z

datafusion/physical-expr-common/src/aggregate/mod.rs

+        .map(Some)
+    }
+
+    fn reverse_expr(&self) -> Option<Arc<dyn AggregateExpr>> {


datafusion/functions-aggregate/src/first_last.rs

ozankabak

LGTM as well

alamb

I took a look at this PR and I agree, great work @mustafasrepo .

I had some comments on the public API design but nothing that would block the PR merging, in my opinon.

alamb · 2024-05-25T09:35:42Z

datafusion/expr/src/udaf.rs

@@ -360,6 +391,27 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
        &[]
    }

+    /// Sets the flag specifying whether the requirement of the UDF is satisfied.


I found this name / documenation to be a little confusing - specifically it isn't clear from just this documentation (the public API of AggregageUDFImpl what type of requirement is being satisfied.

The documentation on AggregateExpr makes it clearer -- maybe we could find a way to bring some of that description to this location, or link to there for more details

Also I think we could improve the name of this function

with_... is often used for builder style APIs in rust, which this is not

It uses "requirement" which is general

Maybe it would be clearer it we called it has_beneficial_ordering or input_ordered_correctly 🤔

datafusion/physical-expr-common/src/utils.rs

jayzhan211 · 2024-05-27T12:34:55Z

datafusion/optimizer/src/replace_distinct_aggregate.rs

-                        ))
-                    })
-                    .collect::<Vec<Expr>>();
+                    .map(|e| first_value(vec![e], false, None, sort_expr.clone(), None));


I just noticed that we import functions-aggregate in optimizer for first_value, should we move this rule into functions-aggregate?

As far as I know, All LogicalPlan rules are under optimizer/src/ folder. Also, I don't see any downside in having this rule here. Hence, I think it is better to keep rule in its current place. I wonder what others think though?

Perhaps what @jayzhan211 is noticing is that by directly relying on datafusion-aggregate-functions in the optimizer it means that now the optimizer is still treating "built in" aggregates specially (and for example, user defined functions couldn't take advantage of this)

So if we moved this rule (or somehow made the dependency via an interface) it would be more general and decouple the code and make datafusion more extensible.

One way is to move the entire rule into datafusion-aggregate-functions but maybe we could instead have replace_distinct_aggregate depend on FunctionRegistry rather than a direct call to first_value for example

Maybe we can file a TODO ticket to track it

Ref apache/datafusion#10648

* deps: update datafusion to 39.0.0, pyo3 to 0.21, and object_store to 0.10.1 `datafusion-common` also depends on `pyo3`, so they need to be upgraded together. * feat: remove GetIndexField datafusion replaced Expr::GetIndexField with a FieldAccessor trait. Ref apache/datafusion#10568 Ref apache/datafusion#10769 * feat: update ScalarFunction The field `func_name` was changed to `func` as part of removing `ScalarFunctionDefinition` upstream. Ref apache/datafusion#10325 * feat: incorporate upstream array_slice fixes Fixes #670 * update ExectionPlan::children impl for DatasetExec Ref apache/datafusion#10543 * update value_interval_daytime Ref apache/arrow-rs#5769 * update regexp_replace and regexp_match Fixes #677 * add gil-refs feature to pyo3 This silences pyo3's deprecation warnings for its new Bounds api. It's the 1st step of the migration, and should be removed before merge. Ref https://pyo3.rs/v0.21.0/migration#from-020-to-021 * fix signature for octet_length Ref apache/datafusion#10726 * update signature for covar_samp AggregateUDF expressions now have a builder API design, which removes arguments like filter and order_by Ref apache/datafusion#10545 Ref apache/datafusion#10492 * convert covar_pop to expr_fn api Ref: https://github.com/apache/datafusion/pull/10418/files * convert median to expr_fn api Ref apache/datafusion#10644 * convert variance sample to UDF Ref apache/datafusion#10667 * convert first_value and last_value to UDFs Ref apache/datafusion#10648 * checkpointing with a few todos to fix remaining compile errors * impl PyExpr::python_value for IntervalDayTime and IntervalMonthDayNano * convert sum aggregate function to UDF * remove unnecessary clone on double reference * apply cargo fmt * remove duplicate allow-dead-code annotation * update tpch examples for new pyarrow interval Fixes #665 * marked q11 tpch example as expected fail Ref #730 * add default stride of None back to array_slice

* move out the ordering ruel Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * introduce rule Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * revert test result Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * pass mulit order test Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * cleanup Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * with new childes Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * revert slt Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * revert back Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * rm rewrite in new child Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * backup Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * only move conversion to optimizer Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * find test that do reverse Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * add test for first and last Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * pass all test Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * upd test Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * upd test Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * cleanup Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * cleanup Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * cleanup Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * add aggregate test Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * cleanup Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * final draft Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * cleanup again Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * pull out finer ordering code and reuse Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * clippy Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * remove finer in optimize rule Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * add comments and clenaup Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * rename fun Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * rename fun Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * fmt Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * avoid unnecessary recursion and rename Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * Minor changes * Add new API for aggregate optimization * Minor changes * Minor changes * Remove old code * Minor changes * Minor changes * Minor changes * Minor changes * Minor changes * Minor changes * Update comments * Minor changes * Minor changes * Review Part 1 * TMP * Update display of aggregate fun exprs * TMP * TMP * Update tests * TMP buggy * modify name in place * Minor changes * Tmp * Tmp * Tmp * TMP * Simplifications * Tmp * Tmp * Compiles * Resolve linter errors * Resolve linter errors * Minor changes * Simplifications * Minor chagnes * Move cast to common * Minor changes * Fix test * Minor changes * Simplifications * Review * Address reviews * Address reviews * Update documentation, rename method * Minor changes --------- Signed-off-by: jayzhan211 <jayzhan211@gmail.com> Co-authored-by: jayzhan211 <jayzhan211@gmail.com> Co-authored-by: berkaysynnada <berkay.sahin@synnada.ai> Co-authored-by: Mehmet Ozan Kabak <ozankabak@gmail.com>

jayzhan211 added 30 commits April 10, 2024 14:55

move out the ordering ruel

e506662

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

introduce rule

2812886

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

revert test result

d20dde3

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

pass mulit order test

c671033

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

cleanup

3f43dc1

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

with new childes

c8b46a4

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

revert slt

fcb297b

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

revert back

ab10168

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

rm rewrite in new child

fffb54a

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

backup

fec3eaf

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

only move conversion to optimizer

5b68b48

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

find test that do reverse

2ab74c9

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

add test for first and last

03abf19

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

pass all test

83c721a

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

upd test

34b8f01

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

upd test

b17e8f7

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

cleanup

a313af0

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

cleanup

a37fb3e

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

cleanup

8de501c

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

add aggregate test

107f4c8

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

cleanup

345d4c6

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

final draft

030d7b9

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

cleanup again

cf3f33e

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

pull out finer ordering code and reuse

6d9ee9f

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

clippy

abd40c6

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

remove finer in optimize rule

3e2c186

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

add comments and clenaup

d7c2590

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

rename fun

fed356a

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

rename fun

0d3d461

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

fmt

d3daa4e

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

mustafasrepo and others added 7 commits May 23, 2024 13:33

Minor chagnes

5fa6b5a

Move cast to common

2aa4244

Minor changes

7393158

Fix test

92fec67

Minor changes

809dc7f

Simplifications

e091afc

Review

89a601b

github-actions bot added logical-expr Logical plan and expressions physical-expr Physical Expressions optimizer Optimizer rules core Core DataFusion crate labels May 24, 2024

jayzhan211 approved these changes May 24, 2024

View reviewed changes

Address reviews

9e50da6

jayzhan211 mentioned this pull request May 24, 2024

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

Open

ozankabak approved these changes May 24, 2024

View reviewed changes

alamb approved these changes May 25, 2024

View reviewed changes

jayzhan211 mentioned this pull request May 27, 2024

Remove duplicate function name in its aliases list #10661

Merged

mustafasrepo added 5 commits May 27, 2024 09:35

Address reviews

c9bcc89

Merge branch 'apache_main' into feature/10061_subsequent_pr

13eed8b

Update documentation, rename method

93cccdb

Minor changes

4c430d3

Merge branch 'apache_main' into feature/10061_subsequent_pr

cbd450c

mustafasrepo merged commit 549cf84 into apache:main May 27, 2024
25 checks passed

jayzhan211 reviewed May 27, 2024

View reviewed changes

alamb mentioned this pull request May 28, 2024

DataFusion weekly project plan (Andrew Lamb) - May 27, 2024 #10699

Closed

9 tasks

jayzhan211 mentioned this pull request May 29, 2024

Decouple builtin aggregate function first_value from optimize rule replace_distinct_aggregate #10703

Closed

appletreeisyellow mentioned this pull request Jun 3, 2024

Regression in first_value and last_value coercion #10781

Closed

Michael-J-Ward added a commit to Michael-J-Ward/datafusion-python that referenced this pull request Jun 13, 2024

convert first_value and last_value to UDFs

5f08588

Ref apache/datafusion#10648

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Convert first, last aggregate function to UDAF #10648

Convert first, last aggregate function to UDAF #10648

mustafasrepo commented May 24, 2024 •

edited by jayzhan211

Loading

jayzhan211 left a comment

jayzhan211 May 24, 2024

ozankabak left a comment

alamb left a comment

alamb May 25, 2024

jayzhan211 May 27, 2024

mustafasrepo May 27, 2024

alamb May 28, 2024

Convert first, last aggregate function to UDAF #10648

Convert first, last aggregate function to UDAF #10648

Conversation

mustafasrepo commented May 24, 2024 • edited by jayzhan211 Loading

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?

jayzhan211 left a comment

Choose a reason for hiding this comment

jayzhan211 May 24, 2024

Choose a reason for hiding this comment

ozankabak left a comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

alamb May 25, 2024

Choose a reason for hiding this comment

jayzhan211 May 27, 2024

Choose a reason for hiding this comment

mustafasrepo May 27, 2024

Choose a reason for hiding this comment

alamb May 28, 2024

Choose a reason for hiding this comment

mustafasrepo commented May 24, 2024 •

edited by jayzhan211

Loading