Add `simplify` method to aggregate function by milenkovicm · Pull Request #10354 · apache/datafusion

milenkovicm · 2024-05-02T17:27:35Z

Which issue does this PR close?

Closes#9526.

Rationale for this change

What changes are included in this PR?

Are these changes tested?

yes, added new test and example

Are there any user-facing changes?

no, changes are backward compatible

milenkovicm · 2024-05-02T17:44:32Z

the change is quite similar to #9906 there are quite a lot optional/unused parameters in the method call, but i personally find it easier to understand as it is equivalent to AggregateFunction signature.

I wonder if ExprSimplifyResult::<T>::Original(T) where T can be something else than args (tuples included) would be a backward compatible change. This way we could send and return all required parameters without copy ... but this is discussion for different time

milenkovicm · 2024-05-03T12:17:51Z

jsut thinking aloud, maybe if we change:

pubenumExprSimplifyResult{/// The function call was simplified to an entirely new ExprSimplified(Expr),/// the function call could not be simplified, and the arguments/// are return unmodified.Original(Vec<Expr>),}

to:

pubenumExprSimplifyResult<T>{/// The function call was simplified to an entirely new ExprSimplified(Expr),/// the function call could not be simplified, and the arguments/// are return unmodified.Original(T),}

simplify method would change from:

pubfn simplify(&self,args:Vec<Expr>,distinct:&bool,filter:&Option<Box<Expr>>,order_by:&Option<Vec<Expr>>,null_treatment:&Option<NullTreatment>,info:&dynSimplifyInfo,) -> Result<ExprSimplifyResult>{

to:

pubfnsimplify(&self,args:Vec<Expr>,distinct:bool,filter:Option<Box<Expr>>,order_by:Option<Vec<Expr>>,null_treatment:Option<NullTreatment>,info:&dynSimplifyInfo,) -> Result<ExprSimplifyResult<(Vec<Expr>,bool,Option<Box<Expr>> ...)>>{}

so in both cases, when we create replacement function or we re-assemble original function we would not need to clone parameters.

maybe hand full of changes, from -> Result<ExprSimplifyResult> to -> Result<ExprSimplifyResult<Vec<Expr>>>

Downside is that we would have to re-create expression if no simplification, plus change would be backward incompatible

jayzhan211 · 2024-05-04T03:38:03Z

+ /// See [`AggregateUDFImpl::simplify`] for more details.
+ pub fn simplify(
+ &self,
+ args: Vec<Expr>,


I think we should reuse args in AggregateFunction, maybe create a new struct with AggregateArgs?
datafusion/datafusion/expr/src/expr.rs
Lines 568 to 582 in 8190cb9
/// Aggregate function
#[derive(Clone,PartialEq,Eq,Hash,Debug)]
pubstructAggregateFunction{
/// Name of the function
pubfunc_def:AggregateFunctionDefinition,
/// List of expressions to feed to the functions as arguments
pubargs:Vec<Expr>,
/// Whether this is a DISTINCT aggregation or not
pubdistinct:bool,
/// Optional filter
pubfilter:Option<Box<Expr>>,
/// Optional ordering
puborder_by:Option<Vec<Expr>>,
pubnull_treatment:Option<NullTreatment>,
}

I don't know, mixed feelings about it. IMHO I see no benefits, would not simplify anything and another structure would be introduced.

One benefit is that we could potentially add new fields without causing an API change 🤔
Another benefit could be that we could add a default method that recreated the expression. For example
implAggegateArgs{// recreate the original aggregate fucntionsfninto_expr(self) -> Expr{ ...} ... }
Another benefit is, as @jayzhan211 says, that it would be more consistent with the other APIs

Do we change signature of pub struct AggregateFunction moving args, distinct, filter ... to AggegateArgs or we keep AggegateArgs as auxiliary structure?

jayzhan211 · 2024-05-04T03:45:14Z

pubenumExprSimplifyResult<T>{/// The function call was simplified to an entirely new ExprSimplified(Expr),/// the function call could not be simplified, and the arguments/// are return unmodified.Original(T),}

I think it is a good idea.

jayzhan211 · 2024-05-04T03:46:20Z

+ fn simplify(
+ &self,
+ args: Vec<Expr>,
+ distinct: &bool,


should we also pass others args with owned value for rewrite like args: Vec<Expr>?

I agree, if we make ExprSimplifyResult generic, as I pointed out in the comment. If we change it without change of ExprSimplifyResult it would lead to cloning of arguments even if simplify does nothing

Given the work we have been doing on #9637 let's not add an API that forces cloning
Another potential thing is to use Transformed directlt

alamb · 2024-05-06T19:06:21Z

jsut thinking aloud, maybe if we change:

I think if we are going to change the signature, I think we should use Transformed directly (and it is already parameterized)

alamb

Thanks @milenkovicm and @jayzhan211 -- I agree this API is a little akward. It would be great to figure out how to avoid cloning when unecessary

alamb · 2024-05-06T18:58:51Z

+ /// See [`AggregateUDFImpl::simplify`] for more details.
+ pub fn simplify(
+ &self,
+ args: Vec<Expr>,


One benefit is that we could potentially add new fields without causing an API change 🤔
Another benefit could be that we could add a default method that recreated the expression. For example
implAggegateArgs{// recreate the original aggregate fucntionsfninto_expr(self) -> Expr{ ...} ... }
Another benefit is, as @jayzhan211 says, that it would be more consistent with the other APIs

alamb · 2024-05-06T19:04:52Z

 ExprSimplifyResult::Simplified(expr) => Transformed::yes(expr),
 },

+ Expr::AggregateFunction(AggregateFunction {


It seems like in order to be able to make this API avoid copying, it would need to pass the distinct, filter, etc fields in and then get them back
Another thing we could do is to use Transformed directly
What if we made something like
structAggregateFunctionsArgs{args:Vec<Expr>,distinct:bool,filter:Option<Expr>, order_by: ... }
And then the call to simplify to be
traitAggreagateUDFImpl{ ... fnsimplify(&self,agg_args:AggregateFunctionsArgs,) -> Result<Transformed<AggregateFunctionsArgs>{Ok(Transformed::no(agg_args))}

I agree with AggregateFunctionsArgs, consistency part got me on board 😀
I'm not sure about Transformed part, reason being is that ExprSimplifyResult returns two types, Expr and currently Vec<Expr>
in that case simplify signature should be changed to:
traitAggreagateUDFImpl{ ... fnsimplify(&self,agg_args:AggregateFunctionsArgs,) -> Result<Transformed<Expr>{let original :Expr = agg_args.into();Ok(Transformed::no(original))}
Thus simplify function should return 're-composed' original function or new function ...
Anyway, I'll have a look what can be done, we can discuss

alamb · 2024-05-06T19:06:47Z

+ fn simplify(
+ &self,
+ args: Vec<Expr>,
+ distinct: &bool,


Given the work we have been doing on #9637 let's not add an API that forces cloning
Another potential thing is to use Transformed directlt

milenkovicm · 2024-05-07T14:21:39Z

Another option could be to make a simplify return optional closure

fnsimplify(&self) -> Option<Box<dynFn(AggregateFunction) -> datafusion_common::tree_node::Transformed<Expr>>>{None}

in this case expr_simplifier part would become:

Expr::AggregateFunction(AggregateFunction{func_def:AggregateFunctionDefinition::UDF(ref udaf), ..}) => {match(udaf.simplify(), expr){(Some(simplify),Expr::AggregateFunction(a)) => simplify(a),(_, e) => Transformed::no(e),}}

this solution:

avoids cloning,
avoids decomposition/composition of function if not needed
avoids adding new struct
but, it has dynamic dispatch simplify function call (in case simplify is defined)
and a bit odd interface

other issues:

SimplifyInfo is missing from signature, can be added
Transformed::no does not make sense in this context as original expression has not been propagated (should we use actual expression instead of AggregateFunction or remove Transformed from return type
closure should return result with or without Transformed

jayzhan211 · 2024-05-08T00:37:29Z

Another option is to introduce has_simplify so we can simplify the code to

Expr::AggregateFunction(AggregateFunction{func_def:AggregateFunctionDefinition::UDF(udaf),args:AggregateFunctionArgs,})if udaf.has_simplify() => {Transformed::yes(udaf.simplify()?)}
fn simplify() -> Result<Expr>

The downside is that the user needs to both define simplify and set has_simplify to true, but I think it is much simpler than the optional closure

milenkovicm · 2024-05-08T07:56:22Z

The downside is that the user needs to both define simplify and set has_simplify to true, but I think it is much simpler than the optional closure

I personally try to avoid situations like this when creating API, this should be documented that both functions have to be implemented plus it introduces more corner cases.

Anyway, closure was a brain dump from yesterday, after a good sleep I don't think we need it, will try to elaborate in next comment

milenkovicm · 2024-05-08T08:18:53Z

Option 1 - Original expression as a parameter

fnsimplify(&self,expr:Expr,) -> Result<datafusion_common::tree_node::Transformed<Expr>>{// we know it'll always be AggregateFunction but we have to match on it // which is not terribly hard but it makes api a bit odd ifletExpr::AggregateFunction(agg) = expr {
...Ok(datafusion_common::tree_node::Transformed::yes(...))}else{// no re-creation of expression if not used Ok(datafusion_common::tree_node::Transformed::no(expr))}}

default implementation is just:

fnsimplify(&self,expr:Expr,) -> Result<datafusion_common::tree_node::Transformed<Expr>>{// no re-creation of expression if not used Ok(datafusion_common::tree_node::Transformed::no(expr))}

Option 2 - AggregationFunction as a parameter

fnsimplify(&self,agg:AggregateFunction,) -> Result<datafusion_common::tree_node::Transformed<Expr>>{// we know its aggregate function, no need to do extra matching, but we need to re-create expression // issue user can return `no` with new expression in it Ok(datafusion_common::tree_node::Transformed::no(Expr::AggregateFunction(agg)))}

Option 3 - AggregationFunction as a parameter with parametrised `ExprSimplifyResult`

fnsimplify(&self,agg:AggregateFunction,) -> Result<ExprSimplifyResult<AggregateFunction>>{// user cant return `no` with new expression // simplifier will re-crete expression in case original has been returned Ok(ExprSimplifyResult::Original(agg))}

IMHO, I prefer option number 3, which would involve parametrizing ExprSimplifyResult, second best is option number 2.
wdyt @jayzhan211, @alamb

jayzhan211 · 2024-05-08T09:12:44Z

Option 1 - Original expression as a parameter

fnsimplify(&self,expr:Expr,) -> Result<datafusion_common::tree_node::Transformed<Expr>>{// we know it'll always be AggregateFunction but we have to match on it // which is not terribly hard but it makes api a bit odd ifletExpr::AggregateFunction(agg) = expr {
...Ok(datafusion_common::tree_node::Transformed::yes(...))}else{// no re-creation of expression if not used Ok(datafusion_common::tree_node::Transformed::no(expr))}}

default implementation is just:

fnsimplify(&self,expr:Expr,) -> Result<datafusion_common::tree_node::Transformed<Expr>>{// no re-creation of expression if not used Ok(datafusion_common::tree_node::Transformed::no(expr))}

Option 2 - AggregationFunction as a parameter

fnsimplify(&self,agg:AggregateFunction,) -> Result<datafusion_common::tree_node::Transformed<Expr>>{// we know its aggregate function, no need to do extra matching, but we need to re-create expression // issue user can return `no` with new expression in it Ok(datafusion_common::tree_node::Transformed::no(Expr::AggregateFunction(agg)))}

Option 3 - AggregationFunction as a parameter with parametrised `ExprSimplifyResult`

fnsimplify(&self,agg:AggregateFunction,) -> Result<ExprSimplifyResult<AggregateFunction>>{// user cant return `no` with new expression // simplifier will re-crete expression in case original has been returned Ok(ExprSimplifyResult::Original(agg))}

IMHO, I prefer option number 3, which would involve parametrizing ExprSimplifyResult, second best is option number 2. wdyt @jayzhan211, @alamb

Given these 3, I prefer 2, because I think re-create expression is not a high cost that we should avoid. And we can do even further,

If user does not implement simplify, it is equivalent to Transformed::no or ExprSimplifyResult::Original, so we can just have Result<Expr>

fnsimplify(&self,agg:AggregateFunction,) -> Result<Expr>{Ok(Expr::AggregateFunction(agg))}

milenkovicm · 2024-05-08T12:46:25Z

I apologise for making noise, but it looks like all of the non closure options have issue with borrowing:

error[E0505]: cannot move out of `expr.0` because it is borrowed
--> datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs:1391:48
|
1388 | func_def: AggregateFunctionDefinition::UDF(ref udaf),
| -------- borrow of `expr.0.func_def.0` occurs here
...
1391 | if let Expr::AggregateFunction(aggregate_function) = expr {
| ^^^^^^^^^^^^^^^^^^ move out of `expr.0` occurs here
1392 | udaf.simplify(aggregate_function, info)?
| ---- borrow later used here

It looks like AggregateArg might be back on the table, which would involve cloning udf

milenkovicm · 2024-05-09T09:58:31Z

@jayzhan211 and @alamb whenever you get chance please have a look, IMHO I find proposal with closure to tick all the boxes, but this one is definitely simpler.

jayzhan211 · 2024-05-09T12:51:31Z

@jayzhan211 and @alamb whenever you get chance please have a look, IMHO I find proposal with closure to tick all the boxes, but this one is definitely simpler.

After playing around, I feel like maybe optional closure is our answer 🤔

fnsimplify(&self,_info:&dynSimplifyInfo,) -> Option<Box<dynFn(AggregateArgs) -> Expr>>{}

milenkovicm · 2024-05-09T12:57:09Z

After playing around, I feel like maybe optional closure is our answer 🤔

Damn! I should have created a branch with that code 😀, if we have consensus on that direction I'll redo it

jayzhan211 · 2024-05-09T13:28:55Z

After playing around, I feel like maybe optional closure is our answer 🤔
Damn! I should have created a branch with that code 😀, if we have consensus on that direction I'll redo it

My final answer for today is

// UDAFfn simplify(&self,args:AggregateArgs,_info:&dynSimplifyInfo,) -> Option<Box<dynFn(AggregateArgs) -> Result<Expr>>>{

// UDF, not necessary but for consistencyfn simplify(&self,args:Vec<Expr>,_info:&dynSimplifyInfo,) -> Option<Box<dynFn(Vec<Expr>) -> Result<Expr>>>{

Go with optional closure!

milenkovicm · 2024-05-09T13:32:46Z

My final answer for today is

fn simplify(&self,args:AggregateArgs,_info:&dynSimplifyInfo,) -> Option<Box<dynFn(AggregateArgs) -> Result<Expr>>>{

Go with optional closure!

you mean

// UDFfn simplify(&self,) -> Option<Box<dynFn(AggregateFunction,&dynSimplifyInfo) -> Result<Expr>>>{

or

// UDFfn simplify(&self,info:&dynSimplifyInfo) -> Option<Box<dynFn(AggregateFunction) -> Result<Expr>>>{

jayzhan211 · 2024-05-09T13:40:10Z

It should be this.

// UDFfn simplify(&self,) -> Option<Box<dynFn(AggregateFunction,&dynSimplifyInfo) -> Result<Expr>>>{

The reason for optional closure is that I assume we need to return Expr for simplified cases. If we expect AggregateArgs, then we might not need closure.

milenkovicm · 2024-05-09T13:45:13Z

We're on the same page @jayzhan211

jayzhan211 · 2024-05-09T13:49:26Z

We're on the same page @jayzhan211

If we don't need Expr for simplified UDAF, than we can have

pubfnsimplify(&self,args:AggregateArgs,info:&dynSimplifyInfo,) -> Result<Transformed<AggregateArgs>>{}

I think the next problem is whether we need Expr or just AggregateArgs. The former can let us return the UDAF with a different function. But for this kind of function rewrite, we can handle it in FuncitonWrite. You can take ArrayFunctionRewriter for example.

@milenkovicm

milenkovicm · 2024-05-09T13:52:13Z

We're on the same page @jayzhan211
If we don't need Expr for simplified UDAF, than we can have
pubfnsimplify(&self,args:AggregateArgs,info:&dynSimplifyInfo,) -> Result<Transformed<AggregateArgs>>{}
I think the next problem is whether we need Expr or just AggregateArgs. The former can let us return the UDAF with a different function. But for this kind of function rewrite, we can handle it in FuncitonWrite. You can take ArrayFunctionRewriter for example.
@milenkovicm

I think we need Expr, this way we can replace function with something else

jayzhan211 · 2024-05-09T14:24:16Z

We're on the same page @jayzhan211
If we don't need Expr for simplified UDAF, than we can have
pubfnsimplify(&self,args:AggregateArgs,info:&dynSimplifyInfo,) -> Result<Transformed<AggregateArgs>>{}
I think the next problem is whether we need Expr or just AggregateArgs. The former can let us return the UDAF with a different function. But for this kind of function rewrite, we can handle it in FuncitonWrite. You can take ArrayFunctionRewriter for example.
@milenkovicm
I think we need Expr, this way we can replace function with something else

I agree.

jayzhan211

👍

jayzhan211 · 2024-05-09T23:48:12Z

You can define closure in datafusion/expr/src/function.rs

alamb

Thanks @milenkovicm and @jayzhan211 -- I think it looks good and the API is clean. Let's give it a try 🚀

Looks like there is a small conflict to resolve and then we'll be good to go

alamb · 2024-05-13T11:54:03Z

Thanks again @jayzhan211 and @milenkovicm 🙏

* add simplify method for aggregate function * simplify returns closure

github-actionsBot added logical-expr Logical plan and expressions optimizer Optimizer rules labels May 2, 2024

milenkovicm force-pushed the simplify_aggregations_9526 branch from 7325587 to 0e50d42Compare May 2, 2024 17:35

milenkovicm force-pushed the simplify_aggregations_9526 branch 2 times, most recently from 42000ac to f96a6d8Compare May 2, 2024 19:37

alamb mentioned this pull request May 2, 2024
DataFusion weekly project plan (Andrew Lamb) - April 29, 2024 #10283
Closed
8 tasks

jayzhan211 reviewed May 4, 2024
View reviewed changes

alamb mentioned this pull request May 6, 2024
DataFusion weekly project plan (Andrew Lamb) - May 6, 2024 #10395
Closed
7 tasks

alamb reviewed May 6, 2024
View reviewed changes

milenkovicm force-pushed the simplify_aggregations_9526 branch from f96a6d8 to 62d381eCompare May 7, 2024 08:34

milenkovicm force-pushed the simplify_aggregations_9526 branch from 62d381e to cd68b7fCompare May 8, 2024 10:03

milenkovicm marked this pull request as draft May 8, 2024 12:52

milenkovicm force-pushed the simplify_aggregations_9526 branch 3 times, most recently from d0e6842 to de51434Compare May 9, 2024 09:52

milenkovicm marked this pull request as ready for review May 9, 2024 09:54

jayzhan211 approved these changes May 9, 2024
View reviewed changes

milenkovicm force-pushed the simplify_aggregations_9526 branch 2 times, most recently from 47108f1 to a678e6dCompare May 10, 2024 14:28

alamb approved these changes May 12, 2024
View reviewed changes

milenkovicm added 2 commits May 12, 2024 19:57

add simplify method for aggregate function
d13a57d

simplify returns closure
c40690c

milenkovicm force-pushed the simplify_aggregations_9526 branch from a678e6d to c40690cCompare May 12, 2024 19:00

alamb changed the title ~~Add simplify method to aggregate function~~Add `simplify1 method to aggregate functionMay 13, 2024

alamb changed the title ~~Add `simplify1 method to aggregate function~~Add simplify method to aggregate functionMay 13, 2024

alamb merged commit 230c68c into apache:mainMay 13, 2024

findepi pushed a commit to findepi/datafusion that referenced this pull request Jul 16, 2024
Addsimplifymethod to aggregate function (apache#10354)
e264b8b
* add simplify method for aggregate function * simplify returns closure

	/// Aggregate function
	#[derive(Clone,PartialEq,Eq,Hash,Debug)]
	pubstructAggregateFunction{
	/// Name of the function
	pubfunc_def:AggregateFunctionDefinition,
	/// List of expressions to feed to the functions as arguments
	pubargs:Vec<Expr>,
	/// Whether this is a DISTINCT aggregation or not
	pubdistinct:bool,
	/// Optional filter
	pubfilter:Option<Box<Expr>>,
	/// Optional ordering
	puborder_by:Option<Vec<Expr>>,
	pubnull_treatment:Option<NullTreatment>,
	}

Uh oh!

Conversation

milenkovicm commented May 2, 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!

milenkovicm commented May 2, 2024

Uh oh!

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

Uh oh!

Uh oh!

jayzhan211May 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!

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 4, 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!

alamb commented May 6, 2024

Uh oh!

alamb 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!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

milenkovicm commented May 8, 2024

Uh oh!

milenkovicm commented May 8, 2024

Option 1 - Original expression as a parameter

Option 2 - AggregationFunction as a parameter

Option 3 - AggregationFunction as a parameter with parametrised ExprSimplifyResult

Uh oh!

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

Uh oh!

Option 1 - Original expression as a parameter

Option 2 - AggregationFunction as a parameter

Option 3 - AggregationFunction as a parameter with parametrised ExprSimplifyResult

Uh oh!

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

Uh oh!

Uh oh!

milenkovicm commented May 9, 2024

Uh oh!

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

Uh oh!

Uh oh!

milenkovicm commented May 9, 2024

Uh oh!

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

Uh oh!

Uh oh!

milenkovicm commented May 9, 2024

Uh oh!

jayzhan211 commented May 9, 2024

milenkovicm commented May 3, 2024•
edited
Loading

jayzhan211May 4, 2024•
edited
Loading

milenkovicm commented May 7, 2024•
edited
Loading

jayzhan211 commented May 8, 2024•
edited
Loading

Option 3 - AggregationFunction as a parameter with parametrised `ExprSimplifyResult`

jayzhan211 commented May 8, 2024•
edited
Loading

Option 3 - AggregationFunction as a parameter with parametrised `ExprSimplifyResult`

milenkovicm commented May 8, 2024•
edited
Loading

jayzhan211 commented May 9, 2024•
edited
Loading

jayzhan211 commented May 9, 2024•
edited
Loading

jayzhan211 commented May 9, 2024•
edited
Loading