Add a `ScalarUDFImpl::simplfy()` API, move `SimplifyInfo` et al to datafusion_expr by jayzhan211 · Pull Request #9304 · apache/datafusion

jayzhan211 · 2024-02-21T12:52:26Z

Which issue does this PR close?

Closes #9289 .

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

datafusion/physical-expr/src/planner.rs

datafusion/expr/src/udf.rs

jayzhan211 · 2024-02-23T12:00:08Z

datafusion/optimizer/src/analyzer/rewrite_expr.rs

+    }
+}
+
+fn function_simplication_analyze_internal(plan: &LogicalPlan) -> Result<LogicalPlan> {


port from operator to function, schema is removed since not used

I wonder if there is any reason calling simplify needs to be its own analyzer pass? Can't it be called as part of the existing simplification code?

logical_plan after simplify_expressions SAME TEXT AS ABOVE

That would also allow this function to be used by anyone else that used the simplification API directly: https://github.com/apache/arrow-datafusion/blob/8f3d1ef23f93cd4303745eba76c0850b39774d07/datafusion-examples/examples/expr_api.rs#L113

No special reason. I forgot this analyzer exists. Let me take a look

alamb

Thanks @jayzhan211 -- this is looking really nice and close. Thank you so much for working on these items and I again apologize for the delay in review.

I left some comments -- let me know what you think.

alamb · 2024-02-27T12:18:18Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+    }
+    // Wrap with Expr::Cast() to Int64
+    fn simplify(&self, args: &[Expr]) -> Result<Simplified> {
+        let dfs = DFSchema::new_with_metadata(


This is a great example, thank you @jayzhan211

It seems to me that the UDF can't always know what the input schema would be, and thus the schema should be passed in.

Something like this

fn simplify(&self, args: &[Expr], schema: &DFSchema) -> Result<Simplified> { ... }

What do you think?

SimpliyContext has DFSchemaRef, so I change the type to DFSchemaRef, I think it is also fine

alamb · 2024-02-27T12:18:45Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+}
+
+#[tokio::test]
+async fn test_user_defined_functions_cast_to_i64() -> Result<()> {


Thank you so much for this test / example -- it makes seeing how the API would work really clear. 👏

alamb · 2024-02-27T12:21:44Z

datafusion/expr/src/udf.rs

+    fn simplify(&self, _args: &[Expr]) -> Result<Simplified> {
+        Ok(Simplified::Original)
+    }


Is it possible to make this take self (and avoid a copy)? Something like

Suggested change

fn simplify(&self, _args: &[Expr]) -> Result<Simplified> {

Ok(Simplified::Original)

}

fn simplify(self, _args: &[Expr]) -> Result<Simplified> {

Ok(Simplified::Original(self))

}

```?

It seems quite hard.

If I +Sized for UDF trait, it can't be made into an object.

alamb · 2024-02-27T12:24:07Z

datafusion/optimizer/src/analyzer/rewrite_expr.rs

+    }
+}
+
+fn function_simplication_analyze_internal(plan: &LogicalPlan) -> Result<LogicalPlan> {


I wonder if there is any reason calling simplify needs to be its own analyzer pass? Can't it be called as part of the existing simplification code?

logical_plan after simplify_expressions SAME TEXT AS ABOVE

That would also allow this function to be used by anyone else that used the simplification API directly: https://github.com/apache/arrow-datafusion/blob/8f3d1ef23f93cd4303745eba76c0850b39774d07/datafusion-examples/examples/expr_api.rs#L113

alamb · 2024-02-28T18:12:51Z

Thanks @jayzhan211 -- I have some more thoughts on this PR and I hope to share them soon. However it may not be until tomorrow as I have a bunch of other things going on right now so I may run out of time

alamb · 2024-02-29T21:59:56Z

I ran out of time (again) today, but I will put this at the top of my list for tomorrow

alamb · 2024-03-01T21:48:07Z

datafusion/expr/src/udf.rs

+    // Do the function rewrite.
+    // 'args': The arguments of the function
+    // 'schema': The schema of the function
+    fn simplify(&self, _args: &[Expr], _schema: DFSchemaRef) -> Result<Simplified> {


I played around with this PR today.

I think the core challenge is that the Simplification code uses SimplifyContext which purposely doesn't have a reference to DFSchema. This means the simplifier can't call ScalarUDFImpl::simplify as written.

I tried several options, and here is my suggestion:

Move SimplifyContext into datafusion_expr

Change the signature to fn simplify(&self, _args: &[Expr], info: &dyn SimplifyContext) -> Result<Simplified>

Then I think calling ScalarUDFImpl::simplify will end up being straightforward to call

The one challenge I found when I tried to do this locally is that SImplifyInfo currently depends on ExecutionPropss which is in datafusion-physical-expr -- we probably have to move that structure into DataFusion common or datafusion expr as well

I guess what you are suggesting is something like

https://github.com/apache/arrow-datafusion/blob/2a490e48c924b76a2d389227c674b9f81b0404c8/datafusion/optimizer/src/simplify_expressions/utils.rs#L351-L361

alamb · 2024-03-01T22:02:04Z

I view this PR as strategically important (as it is needed to unlock the migration of functions out of the core)

I think figuring out how to get SimplifyContext out of datafusion/optimizer is the key. It is a matter of picking in which crates ExecutionProps and VarProvider go -- I can try and help tomorrow.

jayzhan211 · 2024-03-02T00:59:11Z

datafusion/optimizer/src/simplify_expressions/function_simplifier.rs

+        {
+            let schema = self
+                .info
+                .schema()


I think we can get the DFSchema from SimplifyContext here.

datafusion/physical-expr/src/lib.rs

jayzhan211 · 2024-03-03T00:31:36Z

I pushed some commits to this PR @jayzhan211 -- I hope you don't mind.

Do you think the changes make sense ?

Thank you @alamb , I don't mind that.
I'm not sure why we should avoid returning schema, others make sense to me

jayzhan211 · 2024-03-03T00:33:04Z

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

                out_expr.rewrite(self)?
            }

+            Expr::ScalarFunction(ScalarFunction {


alamb · 2024-03-03T07:08:43Z

I'm not sure why we should avoid returning schema, others make sense to me

I was thinking we should stay true to the design intent on SimplifyInfo that says it was designed to avoid having to create DFSchema objects.

If we were going to require DFSChema I think we could simply get rid of SimplifyInfo entirely 🤔

/// The information necessary to apply algebraic simplification to an
/// [Expr]. See [SimplifyContext] for one concrete implementation.
///
/// This trait exists so that other systems can plug schema
/// information in without having to create `DFSchema` objects. If you
/// have a [`DFSchemaRef`] you can use [`SimplifyContext`]
pub trait SimplifyInfo {

jayzhan211 · 2024-03-03T07:51:53Z

LGTM

alamb · 2024-03-04T11:09:45Z

I plan to resolve the conflicts with this PR later today and merge this PR unless there are any additional comments.

milenkovicm · 2024-03-04T15:17:37Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+    // Casting should be done in `simplify`, so we just return the first argument
+    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+        assert_eq!(args.len(), 1);
+        Ok(args.first().unwrap().clone())


Would it make sense to add ExprSimplifyResult::Replace(Expr) to cover this case, and eliminate UDF call when expression is simplified?

To put some context to my comment, let's say if we define function f(INT, INT) = $1 + $2 we can eliminate UDF call with Alias($1 + $2, "f(a,b)") and get UDF free plan, which would be easier to distribute across ballista cluster

I think it is an excellent idea -- I did so in fea82cb

alamb · 2024-03-04T22:48:14Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+    // Casting should be done in `simplify`, so we just return the first argument
+    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+        assert_eq!(args.len(), 1);
+        Ok(args.first().unwrap().clone())


I think it is an excellent idea -- I did so in fea82cb

... `DROP FUNCTION` will look for function name in all available registries (udf, udaf, udwf). `remove` may be necessary if UDaF and UDwF do not get `simplify` method from apache#9304.

milenkovicm · 2024-03-05T09:26:51Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+        Ok(ExprSimplifyResult::Simplified(new_expr))
+    }
+
+    // Casting should be done in `simplify`, so we just return the first argument


comment is a bit outdated

alamb · 2024-03-05T22:17:09Z

Here we go 🚀 Thank you again @jayzhan211 for all the help getting this one tee'd up and ready to merge. I think this will be a very powerful (but simple) API

* Add plugable function factory * cover `DROP FUNCTION` as well ... ... partially, as `SessionState` does not expose unregister_udf at the moment. * update documentation * fix doc test * Address PR comments (code organization) * Address PR comments (factory interface) * fix test after rebase * `remove`'s gone from the trait ... ... `DROP FUNCTION` will look for function name in all available registries (udf, udaf, udwf). `remove` may be necessary if UDaF and UDwF do not get `simplify` method from #9304. * Rename FunctionDefinition and export it ... FunctionDefinition already exists, DefinitionStatement makes more sense. * Update datafusion/expr/src/logical_plan/ddl.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> * Update datafusion/core/src/execution/context/mod.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> * Update datafusion/core/tests/user_defined/user_defined_scalar_functions.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> * Update datafusion/expr/src/logical_plan/ddl.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> * resolve part of follow up comments * Qualified functions are not supported anymore * update docs and todos * fix clippy * address additional comments * Add sqllogicteset for CREATE/DROP function * Add coverage for DROP FUNCTION IF EXISTS * fix multiline error * revert dialect back to generic in test ... ... as `create function` gets support in latest sqlparser. * fmt --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Omega359 · 2024-03-06T03:21:41Z

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

+            }) => match udf.simplify(args, info)? {
+                ExprSimplifyResult::Original(args) => {
+                    Transformed::yes(Expr::ScalarFunction(ScalarFunction {
+                        func_def: ScalarFunctionDefinition::UDF(udf),
+                        args,
+                    }))
+                }
+                ExprSimplifyResult::Simplified(expr) => Transformed::no(expr),
+            },


Sorry, I'm a bit late to the party on this one but I keep rereading this and I can't help but think the logic is inverted - Original means unchanged which should result in Transformed::no, shouldn't it? And Simplified mean changed which should result in Transformed::yes?

I did a test where I switched the yes -> no, no -> yes and cargo test passes which makes me wonder if this code is indeed getting tested properly 🤔

Or am I completely missing something (very likely)?

I think you are correct that the logic is inverted. I wil make a PR to correct it.

I did a test where I switched the yes -> no, no -> yes and cargo test passes which makes me wonder if this code is indeed getting tested properly 🤔

It is a good question about why no tests fail if Transformed::no

I think it is because the expr simplifer code doens't do anything different if the expression was transformed or not: https://github.com/apache/arrow-datafusion/blob/f3836a53122e86f2b73c25557deaa5b800e488e9/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs#L188-L189

It just looks at data:
https://github.com/apache/arrow-datafusion/blob/f3836a53122e86f2b73c25557deaa5b800e488e9/datafusion/common/src/tree_node.rs#L456

github-actions bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates core Core DataFusion crate labels Feb 21, 2024

jayzhan211 force-pushed the simply-udf branch from 3242780 to bba4940 Compare February 21, 2024 12:53

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

jayzhan211 commented Feb 21, 2024

View reviewed changes

datafusion/physical-expr/src/planner.rs Outdated Show resolved Hide resolved

jayzhan211 commented Feb 21, 2024

View reviewed changes

datafusion/expr/src/udf.rs Outdated Show resolved Hide resolved

jayzhan211 marked this pull request as draft February 23, 2024 03:16

github-actions bot added the optimizer Optimizer rules label Feb 23, 2024

jayzhan211 commented Feb 23, 2024

View reviewed changes

github-actions bot removed the physical-expr Changes to the physical-expr crates label Feb 23, 2024

jayzhan211 marked this pull request as ready for review February 23, 2024 12:02

jayzhan211 marked this pull request as draft February 23, 2024 12:38

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Feb 23, 2024

jayzhan211 marked this pull request as ready for review February 24, 2024 00:17

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:00

alamb mentioned this pull request Feb 27, 2024

feat: arrow_cast function as UDF #9298

Closed

github-actions bot removed the sqllogictest SQL Logic Tests (.slt) label Feb 28, 2024

jayzhan211 marked this pull request as ready for review February 28, 2024 12:11

jayzhan211 requested a review from alamb February 28, 2024 12:11

alamb mentioned this pull request Feb 29, 2024

Add plugable handler for CREATE FUNCTION #9333

Merged

alamb reviewed Mar 1, 2024

View reviewed changes

jayzhan211 commented Mar 2, 2024

View reviewed changes

jayzhan211 force-pushed the simply-udf branch from 9c01a16 to 5d1d500 Compare March 2, 2024 02:18

github-actions bot added the physical-expr Changes to the physical-expr crates label Mar 2, 2024

jayzhan211 commented Mar 3, 2024

View reviewed changes

datafusion/physical-expr/src/lib.rs Show resolved Hide resolved

jayzhan211 commented Mar 3, 2024

View reviewed changes

alamb added 2 commits March 3, 2024 02:09

Update datafusion/physical-expr/src/lib.rs

24adcbf

Merge remote-tracking branch 'apache/main' into simply-udf

8cab80b

alamb changed the title ~~Add a ScalarUDFImpl::simplfy() API~~ Add a ScalarUDFImpl::simplfy() API, move SimplifyInfo et al to datafusion_expr Mar 3, 2024

milenkovicm mentioned this pull request Mar 4, 2024

Add example for CREATE FUNCTION & FunctionFactory #9448

Closed

milenkovicm reviewed Mar 4, 2024

View reviewed changes

alamb added 2 commits March 4, 2024 17:41

Merge remote-tracking branch 'apache/main' into simply-udf

550cbc4

Change example simplify to always simplify its argument

fea82cb

alamb approved these changes Mar 4, 2024

View reviewed changes

alamb mentioned this pull request Mar 4, 2024

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

Closed

5 tasks

milenkovicm reviewed Mar 5, 2024

View reviewed changes

Omega359 mentioned this pull request Mar 5, 2024

Move from_unixtime, now, current_date, current_time functions to datafusion-functions #9466

Closed

alamb added 2 commits March 5, 2024 16:28

Clarify comment

4e9eb70

Merge remote-tracking branch 'apache/main' into simply-udf

fdec54c

alamb merged commit 2873fd0 into apache:main Mar 5, 2024

Omega359 reviewed Mar 6, 2024

View reviewed changes

milenkovicm mentioned this pull request Mar 6, 2024

Add example for FunctionFactory #9482

Merged

This was referenced Mar 6, 2024

Minor: fixed transformed state in UDF Simplify #9484

Merged

Add a AggregateUDFImpl::simplfy() API #9526

Closed

Add a WindowUDFImpl::simplfy() API #9527

Closed

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

Closed

jayzhan211 mentioned this pull request Mar 16, 2024

[EPIC] Stop copying LogicalPlan during OptimizerPasses #9637

Closed

31 tasks

Conversation

jayzhan211 commented Feb 21, 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!

Uh oh!

Uh oh!

jayzhan211 Feb 23, 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!

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!

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 Feb 28, 2024

Uh oh!

alamb commented Feb 29, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Mar 1, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jayzhan211 commented Mar 3, 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!

alamb commented Mar 3, 2024

Uh oh!

jayzhan211 commented Mar 3, 2024

Uh oh!

alamb commented Mar 4, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

milenkovicm Mar 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!

alamb commented Mar 5, 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

jayzhan211 Feb 23, 2024 •

edited

Loading

jayzhan211 commented Mar 3, 2024 •

edited

Loading

milenkovicm Mar 4, 2024 •

edited

Loading