Introduce expr builder for aggregate function by jayzhan211 · Pull Request #10560 · apache/datafusion

jayzhan211 · 2024-05-17T13:54:24Z

Which issue does this PR close?

Rationale for this change

After this PR, there are two kinds of expression API.
One is the normal one like count(expr), convenient for expression that has few arguments.
Another is builder mode, useful for expressions that expects multiple kind of arguments, like first_value_builder(expr).order_by().build()

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

jayzhan211 · 2024-05-17T13:56:21Z

+
 create_func!($UDAF, $AGGREGATE_UDF_FN);
 };
- ($UDAF:ty, $EXPR_FN:ident, $($arg:ident)*, $distinct:ident, $DOC:expr, $AGGREGATE_UDF_FN:ident) => {


I plan to introduce another macro for distinct, so we have count_distinct, count_distinct_builder expr::fn.
I change to order_by because first_value needs it. Also, I change the expression for first_value to single expression, since it does not expect variadic args.

Instead of count_distinct_builder what do you think about adding a distinct type method instead?
Like
let agg = count_builder().args(col("a")).distinct().build()?
🤔

jayzhan211 · 2024-05-17T13:57:31Z

I plan to add distinct macro in the follow up PR for count.

alamb · 2024-05-17T20:06:05Z

This looks super cool @jayzhan211

alamb · 2024-05-17T20:07:17Z

 use datafusion::execution::FunctionRegistry;
-use datafusion::functions_aggregate::covariance::{covar_pop, covar_samp};
-use datafusion::functions_aggregate::expr_fn::first_value;
+use datafusion::functions_aggregate::expr_fn::{


Could we add some documentation with examples somewhere? That would both serve to help document and test this feature.
But thank you @jayzhan211 -- this is looking very cool

jayzhan211 · 2024-05-18T00:14:05Z

 covar_samp(lit(1.5), lit(2.2)),
+ covar_samp_builder(lit(1.5), lit(2.3)).build(),
 covar_pop(lit(1.5), lit(2.2)),
+ covar_pop_builder(lit(1.5), lit(2.3)).build(),


I think roundtrip tests also cover the testing for these expression APIs.
I expect we add all the functions here for test, and only little functions in datafusion-example.

jayzhan211 · 2024-05-19T06:13:23Z

+
+| Syntax | Equivalent to |
+| ------------------------------------------------------ | ----------------------------------- |
+| first_value_builder(expr).order_by(vec![expr]).build() | first_value(expr, Some(vec![expr])) |


I think we can add non-trivial examples in doc, but not all of them.

Yes I think adding examples in doc comments is likely the best for this particular API as people will be looking at the code already if they want to programatically create Exprs

alamb · 2024-05-22T00:43:57Z

THanks @jayzhan211 -- I will plan to review this tomorrow.

alamb

Thank you @jayzhan211 -- this is looking really cool. I have some feeback on the API design

One major thing I think that might be worth considering is now to get one of these builders. Right now they are implemented as count_builder free functions, which means both adding many new functions as well as that they won't work for user defined aggregate functions

I wonder if it would be possible to add something to the AggregateUDF instead maybe make an extension trait so people could make a builder for any arbitrary aggregate expression,

So today to call a UDF you do

// Given an aggregate UDF:let agg_udf:AggregateUdf = todo!();// agg(a)let expr = agg_udf.call(col("a"))

Could we implement a builder that works like

// agg(a ORDER BY b)let expr = agg_udf.builder().args(col("a")).order_by("b").build()

Or I potentailly keep the "call" syntax:

// agg(a ORDER BY b)let expr = agg_udf.builder().order_by(col("b")).call(col("a")).build()

What do you think?

alamb · 2024-05-22T20:03:41Z

+use datafusion_expr::{expr::AggregateFunction, Expr};
+use sqlparser::ast::NullTreatment;
+
+/// Builder for creating an aggregate function expression


I suggest:
Add a docstring example here
Consider putting this in datafusion/expr (not datafusion-functions-aggregate) -- if we put it in datafusion-functions-aggregate people can't use it unless they use the built in aggregates of DataFusion 🤔

if we put it in datafusion-functions-aggregate people can't use it unless they use the built in aggregates of DataFusion
Nice! I forgot this

alamb · 2024-05-22T20:05:35Z

+
 create_func!($UDAF, $AGGREGATE_UDF_FN);
 };
- ($UDAF:ty, $EXPR_FN:ident, $($arg:ident)*, $distinct:ident, $DOC:expr, $AGGREGATE_UDF_FN:ident) => {


Instead of count_distinct_builder what do you think about adding a distinct type method instead?
Like
let agg = count_builder().args(col("a")).distinct().build()?
🤔

alamb · 2024-05-22T20:06:34Z


 [dependencies]
 arrow = { workspace = true }
+concat-idents = "1.1.5"


Would it be possible to use https://docs.rs/paste/latest/paste/ (which is already a dependency) instead? I think that is what binary expression evaluation uses 🤔

alamb · 2024-05-22T20:07:46Z

+
+| Syntax | Equivalent to |
+| ------------------------------------------------------ | ----------------------------------- |
+| first_value_builder(expr).order_by(vec![expr]).build() | first_value(expr, Some(vec![expr])) |


Yes I think adding examples in doc comments is likely the best for this particular API as people will be looking at the code already if they want to programatically create Exprs

alamb · 2024-05-22T20:08:27Z

+/// Builder for creating an aggregate function expression
+///
+/// Has the same arguments from [AggregateFunction]
+pub struct ExprBuilder {


I wonder if we should call it AggBuilder 🤔 or AggregeExprBuilder 🤔

jayzhan211 · 2024-05-23T01:06:25Z

which means both adding many new functions as well as that they won't work for user defined aggregate functions

I think they can build the function with macro in function-aggregate, so they can have their function with ExprBuilder extension. But then I need to move the macro to datafusion-expr 🤔

maybe make an extension trait

It could be a good idea

jayzhan211 · 2024-05-23T11:50:50Z

Maybe we can just extend function for AggregateFunction 🤔

pub fn call(&self, args: Vec<Expr>) -> AggregateFunction

implAggregateFunction{fnbuild() -> Expr}

alamb · 2024-06-04T16:23:21Z


 // perform the aggregation
- let df = df.aggregate(vec![], vec![geometric_mean.call(vec![col("a")])])?;
+ let df = df.aggregate(vec![], vec![geometric_mean.call(vec![col("a")]).build()])?;


Sorry for the delay in reviewing this @jayzhan211 -- I have been thinking about this PR and I am struggling
On one hand, having a builder allows for exposing the ability to provide other types of arguments to aggregate functions (e.g. ORDER BY or FILTER)
On the other hand, I feel like this API change (requiring a call to build() after geometric_mean.call makes DataFusion harder to learn / use as most calls to aggregate functions are made with just the arguments.
I was thinking about what an ideal user experience would be. What do you think about something that permits using the existing expr_fn API, but also allows easier construction of more advanced aggregate calls
What about something like this (as today, no need to call build)
// geometric_mean(a)let agg = geometric_mean.call(vec![col("a")]);let df = df.aggregate(vec![],vec![agg])?;
And then then to create an ORDER BY the user could create and use a builder like this:
// geometric_mean(a FILTER b > c)let agg = geometric_mean.call(vec![col("a")]).agg_builder()// creates a builder that wraps an Expr, no error return.filter(col("b").gt(col("c"((.build()?;// check here if everything was oklet df = df.aggregate(vec![],vec![agg])?;
We could probably implement a trait for Expr to add agg_builder (or maybe just add it directly)

It looks good to me to have both style for function call.
We can also have builder for expr_fn API
expr_fn::first_value(args).agg_builder().filter().builder()

I found that we don't even need builder() with trait impl

// geometric_mean(a FILTER b > c)let agg = geometric_mean.call(vec![col("a")]).filter(col("b").gt(col("c")))

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

alamb · 2024-06-07T10:28:27Z

I'm not sure if we really need the error check since we don't have any now

I guess what I was thinking is that the reason there is no check now is because the rust compiler will ensure the types are correct (specifically that they are AggregateUDF.

With the builder API as it is written now, you can call order_by on any arbitrary Expr (not just the appropriate Expr::AggregateFunction variant)

So I was imagining if any of the methods had been called on an invalid Expr, build() would return an error

Maybe something like

/// Traits to help build aggregate functionstraitAggregateExt{// return a builderfnorder_by(self,order_by:Vec<Expr>) -> AggregateBuilder;fnfilter(self,filter:Box<Expr>) -> AggregateBuilder;fnnull_treatment(self,null_treatment:NullTreatment) -> AggregateBuilder;fndistinct(self) -> AggregateBuilder;}pubstructAggregateBuilder{// if set to none, builder errorsudf:Option<Arc<AggregateUdf>>,order_by:Vec<Expr>,
....}implAggregateBuilder{fnorder_by(self,order_by:Vec<Expr>) -> AggregateBuilder{self.order_by = order_by;self}fnfilter(self,filter:Box<Expr>) -> AggregateBuilder{..}fnnull_treatment(self,null_treatment:NullTreatment) -> AggregateBuilder{..}fndistinct(self) -> AggregateBuilder{..}// builds up any in progress aggregatefnbuild(self) -> Result<Expr>{letSome(udf) = self.udfelse{returnplan_err!("Expr of type XXX is not an aggregate")} udf.order_by = self.order_by;
...Ok(Expr::AggregateFunction(udf))}}implAggregateExtforExpr{fnorder_by(self,order_by:Vec<Expr>) -> AggregateBuilder{matchself{Expr::AggregateFunction(mut udaf) => {AggregateBuilder{udf:Some(udaf)}}// wrong type passed -- error when build is called
_ => {AggregateBuilder{udf:None}}}}

jayzhan211 · 2024-06-07T12:23:47Z

I'm not sure if we really need the error check since we don't have any now

I guess what I was thinking is that the reason there is no check now is because the rust compiler will ensure the types are correct (specifically that they are AggregateUDF.

With the builder API as it is written now, you can call order_by on any arbitrary Expr (not just the appropriate Expr::AggregateFunction variant)

So I was imagining if any of the methods had been called on an invalid Expr, build() would return an error

Maybe something like

/// Traits to help build aggregate functionstraitAggregateExt{// return a builderfnorder_by(self,order_by:Vec<Expr>) -> AggregateBuilder;fnfilter(self,filter:Box<Expr>) -> AggregateBuilder;fnnull_treatment(self,null_treatment:NullTreatment) -> AggregateBuilder;fndistinct(self) -> AggregateBuilder;}pubstructAggregateBuilder{// if set to none, builder errorsudf:Option<Arc<AggregateUdf>>,order_by:Vec<Expr>,
....}implAggregateBuilder{fnorder_by(self,order_by:Vec<Expr>) -> AggregateBuilder{self.order_by = order_by;self}fnfilter(self,filter:Box<Expr>) -> AggregateBuilder{..}fnnull_treatment(self,null_treatment:NullTreatment) -> AggregateBuilder{..}fndistinct(self) -> AggregateBuilder{..}// builds up any in progress aggregatefnbuild(self) -> Result<Expr>{letSome(udf) = self.udfelse{returnplan_err!("Expr of type XXX is not an aggregate")} udf.order_by = self.order_by;
...Ok(Expr::AggregateFunction(udf))}}implAggregateExtforExpr{fnorder_by(self,order_by:Vec<Expr>) -> AggregateBuilder{matchself{Expr::AggregateFunction(mut udaf) => {AggregateBuilder{udf:Some(udaf)}}// wrong type passed -- error when build is called
_ => {AggregateBuilder{udf:None}}}}

I'm not sure is it worth to introduce build() just for handling non Expr::AggregateFunction cases 🤔

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

jayzhan211 · 2024-06-07T12:25:15Z

+ .call(vec![expression])
+ .order_by(order_by)
+ .build()
+ .unwrap()


I think it is fine to unwrap since udaf.call() is guaranteed to be Expr::AggregateFunction

100% agree
Suggested change
.unwrap()
// guaranteed to be `Expr::AggregateFunction`
.unwrap()

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

alamb

I think this looks great @jayzhan211 🙏

I have some ideas about additional comments / documentation that I would be happy to help add -- but I think we could do it as a follow on PR as well

alamb · 2024-06-07T12:43:51Z

+ .call(vec![expression])
+ .order_by(order_by)
+ .build()
+ .unwrap()


100% agree
Suggested change
.unwrap()
// guaranteed to be `Expr::AggregateFunction`
.unwrap()

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

jayzhan211 · 2024-06-07T13:35:48Z

I have some ideas about additional comments / documentation that I would be happy to help add

Sure!

alamb · 2024-06-08T18:33:20Z

I have some ideas about additional comments / documentation that I would be happy to help add
Sure!

Thank you for your patience @jayzhan211 -- I just pushed a bunch of docs and tests (and a small API refinement):

Consolidated the example into the expr_api.rs examples
Simplified the api for filter from filter(Box<Expr>) to just filter(Expr)
Added documentation and examples to the trait
Added tests
Checked for SortExprs in order_by

alamb

I am quite biased, but I think this API now looks really nice and hopefully is both easy to use as well as serve as a model for WindowFunctions

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

jayzhan211 · 2024-06-09T05:49:14Z

It looks pretty nice now! Thanks @alamb

* expr builder Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * fmt Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * build Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * upd user-guide Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * fix builder Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * Consolidate example in udaf_expr.rs, simplify filter API * Add doc strings and examples * Add tests and checks * Improve documentation more * fixup * rm spce Signed-off-by: jayzhan211 <jayzhan211@gmail.com> --------- Signed-off-by: jayzhan211 <jayzhan211@gmail.com> Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Upstream signatures were changed for the new new `AggregateBuilder` api [0]. This simply gets the code to work. We should better incorporate that API into `datafusion-python`. [0] apache/datafusion#10560

* chore: update datafusion deps * feat: impl ExecutionPlan::static_name() for DatasetExec This required trait method was added upstream [0] and recommends to simply forward to `static_name`. [0]: apache/datafusion#10266 * feat: update first_value and last_value wrappers. Upstream signatures were changed for the new new `AggregateBuilder` api [0]. This simply gets the code to work. We should better incorporate that API into `datafusion-python`. [0] apache/datafusion#10560 * migrate count to UDAF Builtin Count was removed upstream. TBD whether we want to re-implement `count_star` with new API. Ref: apache/datafusion#10893 * migrate approx_percentile_cont, approx_distinct, and approx_median to UDAF Ref: approx_distinct apache/datafusion#10851 Ref: approx_median apache/datafusion#10840 Ref: approx_percentile_cont and _with_weight apache/datafusion#10917 * migrate avg to UDAF Ref: apache/datafusion#10964 * migrage corr to UDAF Ref: apache/datafusion#10884 * migrate grouping to UDAF Ref: apache/datafusion#10906 * add alias `mean` for UDAF `avg` * migrate stddev to UDAF Ref: apache/datafusion#10827 * remove rust alias for stddev The python wrapper now provides stddev_samp alias. * migrage var_pop to UDAF Ref: apache/datafusion#10836 * migrate regr_* functions to UDAF Ref: apache/datafusion#10898 * migrate bitwise functions to UDAF The functions now take a single expression instead of a Vec<_>. Ref: apache/datafusion#10930 * add missing variants for ScalarValue with todo * fix typo in approx_percentile_cont * add distinct arg to count * comment out failing test `approx_percentile_cont` is now returning a DoubleArray instead of an IntArray. This may be a bug upstream; it requires further investigation. * update tests to expect lowercase `sum` in query plans This was changed upstream. Ref: apache/datafusion#10831 * update ScalarType data_type map * add docs dependency pickleshare * re-implement count_star * lint: ruff python lint * lint: rust cargo fmt * include name of window function in error for find_window_fn * refactor `find_window_fn` for debug clarity * search default aggregate functions by both name and aliases The alias list no longer includes the name of the function. Ref: apache/datafusion#10658 * fix markdown in find_window_fn docs * parameterize test_window_functions `first_value` and `last_value` are currently failing and marked as xfail. * add test ids to test_simple_select tests marked xfail * update find_window_fn to search built-ins first The behavior of `first_value` and `last_value` UDAFs currently does not match the built-in behavior. This allowed me to remove `marks=pytest.xfail` from the window tests. * improve first_call and last_call use of the builder API * remove trailing todos * fix examples/substrait.py * chore: remove explicit aliases from functions.rs Ref: #779 * remove `array_fn!` aliases * remove alias rules for `expr_fn_vec!` * remove alias rules from `expr_fn!` macro * remove unnecessary pyo3 var-arg signatures in functions.rs * remove pyo3 signatures that provided defaults for first_value and last_value * parametrize test_string_functions * test regr_ function wrappers Closes#778

jayzhan211 commented May 17, 2024
View reviewed changes

jayzhan211 changed the title ~~introduce expr builder for aggregate function~~Introduce expr builder for aggregate functionMay 17, 2024

jayzhan211 marked this pull request as ready for review May 17, 2024 14:29

alamb reviewed May 17, 2024
View reviewed changes

jayzhan211 commented May 18, 2024
View reviewed changes

jayzhan211 commented May 19, 2024
View reviewed changes

jayzhan211 requested a review from alamb May 22, 2024 00:01

alamb reviewed May 22, 2024
View reviewed changes

jayzhan211 marked this pull request as draft May 23, 2024 05:15

github-actionsBot added logical-expr Logical plan and expressions sqllogictest SQL Logic Tests (.slt) labels May 23, 2024

jayzhan211 commented May 23, 2024
View reviewed changes

Comment threaddatafusion-examples/examples/udaf_expr.rs Outdated

jayzhan211 marked this pull request as ready for review May 23, 2024 12:55

jayzhan211 mentioned this pull request May 25, 2024
[Epic] Unify AggregateFunction Interface (remove built in list of AggregateFunction s), improve the system #8708
Closed

github-actionsBot removed the sqllogictest SQL Logic Tests (.slt) label May 26, 2024

jayzhan211 marked this pull request as draft May 27, 2024 11:58

github-actionsBot added the optimizer Optimizer rules label May 27, 2024

jayzhan211 marked this pull request as ready for review May 27, 2024 13:11

jayzhan211 requested a review from alamb May 30, 2024 12:35

alamb reviewed Jun 4, 2024
View reviewed changes

jayzhan211 marked this pull request as draft June 6, 2024 00:08

jayzhan211 force-pushed the expr-builder branch from a73af0f to c689980Compare June 6, 2024 12:36

jayzhan211 added 2 commits June 6, 2024 20:39

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

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

jayzhan211 force-pushed the expr-builder branch from 3754f58 to 698ab8fCompare June 6, 2024 12:39

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

jayzhan211 commented Jun 7, 2024
View reviewed changes

upd user-guide
3c10e71
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

alamb reviewed Jun 7, 2024
View reviewed changes

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

alamb added 5 commits June 8, 2024 08:20

Merge remote-tracking branch 'apache/main' into expr-builder
722e6dc

Consolidate example in udaf_expr.rs, simplify filter API
f2fc8d5

Add doc strings and examples
a471385

Add tests and checks
fc5a6a5

Improve documentation more
d547853

github-actionsBot added the core Core DataFusion crate label Jun 8, 2024

alamb approved these changes Jun 8, 2024
View reviewed changes

This was referenced Jun 8, 2024
Make it easier to create WindowFunctions with the Expr API #6747
Closed
AggregateUDF expression API design #10545
Closed

alamband others added 2 commits June 8, 2024 14:51

fixup
980c30c

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

jayzhan211 merged commit 24a0846 into apache:mainJun 9, 2024

jayzhan211 deleted the expr-builder branch June 9, 2024 05:49

alamb mentioned this pull request Jun 11, 2024
DataFusion weekly project plan (Andrew Lamb) - June 10, 2024 #10869
Closed
7 tasks

timsaucer mentioned this pull request Jul 19, 2024
Consistent API to set parameters of aggregate and window functions (AggregateExt --> ExprFunctionExt) #11550
Merged

	.unwrap()
	// guaranteed to be `Expr::AggregateFunction`
	.unwrap()

Uh oh!

Conversation

jayzhan211 commented May 17, 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?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayzhan211 commented May 17, 2024

Uh oh!

alamb commented May 17, 2024

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!

alamb commented May 22, 2024

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

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!

jayzhan211 commented May 23, 2024• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jayzhan211 commented May 23, 2024

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayzhan211Jun 4, 2024• 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!

jayzhan211Jun 6, 2024• 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!

jayzhan211Jun 6, 2024• 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!

alamb commented Jun 7, 2024

Uh oh!

jayzhan211 commented Jun 7, 2024• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayzhan211 commented Jun 7, 2024

Uh oh!

alamb commented Jun 8, 2024

jayzhan211 commented May 23, 2024•
edited
Loading

jayzhan211Jun 4, 2024•
edited
Loading

jayzhan211Jun 6, 2024•
edited
Loading

jayzhan211Jun 6, 2024•
edited
Loading

jayzhan211 commented Jun 7, 2024•
edited
Loading