ARROW-9836: [Rust][DataFusion] Improve API for usage of UDFs

[jorgecarleitao]

See associated issue and document for details.

The gist is that currently, users call UDFs through

df.select(scalar_functions(“my_sqrt”, vec![col(“a”)], DataType::Float64))

and this PR proposes a change to

let functions = df.registry()?;

df.select(functions.udf(“my_sqrt”, vec![col(“a”)])?)

so that they do not have to remember the UDFs return type when using it (and a whole lot other things for us internally).

[github-actions]

https://issues.apache.org/jira/browse/ARROW-9836

[andygrove]

For built-in functions like sqrt I would expect DataFusion to provide convenience functions to create an expression, like we do with col and the aggregate functions. I assume we could also do that with the design proposed here?

For example, I would like to be able to write:

df.select(vec![col("foo"), sqrt(col("bar"))])?

[jorgecarleitao]

For built-in functions like sqrt I would expect DataFusion to provide convenience functions to create an expression, like we do with col and the aggregate functions. I assume we could also do that with the design proposed here?

For example, I would like to be able to write:

df.select(vec![col("foo"), sqrt(col("bar"))])?

This PR does not support this, as it threats every function (built-in or not) equally. To include that case, IMO this PR needs to add a new enum in the logical Expr:

• Expr::ScalarFunction { name: String/Enum, args: Vec<Expr> } that we logically know its return type based on name (e.g. "sqrt"), exactly like Expr::BinaryExpr. This is mapped to a physical expression during planning. These can be build without access to the registry, as we hard-code the return type on the logical plan to be consistent with the physical one, like we do for our aggregates, binary expressions, etc.
• Expr:ScalarUDF { fun: ScalarFunction, args: Vec<Expr> }, whose return type is only known after going to the registry to check what the user set its return type to be (as this PR currently does).

sqrt would return Expr::ScalarFunction, that knows its own return type, and the planner converts it to ScalarFunction via a hard-coded map, while Expr:UDF's physical planning is just planning args and pass them to ScalarFunction like this PR already does.

I.e. at the physical level, built-in and UDFs are indistinguishable, but at the logical plan, one only knows its name (built-in), the other also knows its physical representation ScalarUDF.

[alamb]

For anyone else reading along, the associated document I think is https://docs.google.com/document/d/1Kzz642ScizeKXmVE1bBlbLvR663BKQaGqVIyy9cAscY/edit?usp=sharing

[jorgecarleitao]

@alamb thank you very much for your comments, I will now work on addressing them now. I still learning the Arc/Box/Ref, so thank you a lot for also teaching me.

@andygrove , I agree with you that built-in functions should not require access to the registry. Unfortunately, doing so required some re-work, which is the reason I retracted #7967 back to draft to focus on this one first.

I pushed a new commit to this PR to address this point. Specifically, that commit adds:

• a new enum with all built-in functions
• functionally gluing the logical plan with the physical plan so that the function's return types are invariant.
• made type coercion on built-in functions to be on the physical plane, to preserve schema invariance during planning.

I am pretty happy with this PR, as IMO has the flexibility we need to expand DataFusion's pool of built-in functions to multiple input and return types. The main features of this PR:

• users no longer have to pass the return type of the UDF when calling them (the proposal)
• planning built-in functions continue to not need access to the registry (@andygrove 's point)
• built-in functions now support multiple input types (e.g. sqrt(f32), sqrt(f64))
• built-in functions now support multiple return types (e.g. sqrt(f32) -> f32, sqrt(f64) -> f64)
• coercion rules are no longer applied in the sql planning or physical planning to built-in functions, to avoid breaking schema invariance during planning

I have not completed the valid return types of built-in math functions as this PR was already too long.

Overall, I think that this has not been a pleasant experience for you @andygrove and @alamb, as I constantly open and close PRs around functions/UDFs, and for that I am really sorry. I've been hitting some design challenge after another, which requires me to go back and forth.

I am still in pursuit of my original quests:

• built-in aggregate functions whose logical types are known from the physical expressions
• type coercion on aggregate functions
• built-in aggregate functions whose return types (e.g. min(f32) -> f32, min(f64) -> f64) are directly derived from the physical plan (there is an old fixme/todo in the code around that)
• aggregate udfs
• udfs with multiple incoming and return types, to bring them to the same level of functionality of built-ins
• planning a udf without registering it (a-la spark) in the DF's API.

I have code for some of this, I... just... need... to... finish... the... scalar... stuff... first... 😃

[alamb]

Good luck -- at this stage of a project (when architecture is changing a bunch) know it is hard to make small / easy to review PRs. I hope the comments are helpful and I am sorry I don't have more time to devote to reviews.

…

On Tue, Aug 25, 2020 at 4:10 PM Jorge Leitao ***@***.***> wrote: @alamb <https://github.com/alamb> thank you very much for your comments, I will now work on addressing them now. I still learning the Arc/Box/Ref, so thank you a lot for also teaching me. @andygrove <https://github.com/andygrove> , I agree with you that built-in functions should not require access to the registry. Unfortunately, doing so required some re-work, which is the reason I retracted #7967 <#7967> back to draft to focus on this one first. I pushed a new commit to this PR to address this point. Specifically, that commit adds: - a new enum with all built-in functions - functionally gluing the logical plan with the physical plan so that the function's return types are invariant. - made type coercion on built-in functions to be on the physical plane, to preserve schema invariance during planning. I am pretty happy with this PR, as IMO has the flexibility we need to expand DataFusion's pool of built-in functions to multiple input and return types. The main features of this PR: - users no longer have to pass the return type of the UDF when calling them (the proposal) - planning built-in functions continue to not need access to the registry ***@***.*** <https://github.com/andygrove> 's point) - built-in functions now support multiple input types (e.g. sqrt(f32), sqrt(f64)) - built-in functions now support multiple return types (e.g. sqrt(f32) -> f32, sqrt(f64) -> f64) - coercion rules are no longer applied in the sql planning or physical planning to built-in functions, to avoid breaking schema invariance during planning I have not completed the valid return types of built-in math functions as this PR was already too long. Overall, I think that this has not been a pleasant experience for you @andygrove <https://github.com/andygrove> and @alamb <https://github.com/alamb>, as I constantly open and close PRs around functions/UDFs, and for that I am really sorry. I've been hitting some design challenge after another, which requires me to go back and forth. I am still in pursuit of my original quests: - built-in aggregate functions whose logical types are known from the physical expressions - type coercion on aggregate functions - built-in aggregate functions whose return types (e.g. min(f32) -> f32, min(f64) -> f64) are directly derived from the physical plan (there is an old fixme/todo in the code around that) - aggregate udfs - udfs with multiple incoming and return types, to bring them to the same level of functionality of built-ins - planning a udf without registering it (a-la spark) in the DF's API. I have code for some of this, I... just... need... to... finish... the... scalar... stuff... first... 😃 — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#8032 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AADXZMLPUU3TGXRLTCKWWKLSCQLEFANCNFSM4QIW6KXQ> .

[jorgecarleitao]

FYI, this is what a PR to add support for f32 to mathematical expressions (keeping the return type f64) looks like:
https://github.com/jorgecarleitao/arrow/pull/1

I.e. IMO with this PR we can support almost any built-in function: fixed type, variable return type, multiple input types, etc on Expr::ScalarFunction, which gives a lot of flexibility to add a new function, as we do not need to fiddle with Expr, only with input types, return types, etc. on the physical plane.

I split built-ins from the UDFs because built-ins type is known without access to the registry, which is currently required to allow users to use them outside the registry.

[jorgecarleitao]

FYI, this is what a PR would look like for the concatenate function: https://github.com/jorgecarleitao/arrow/pull/2/files using this API.

I am not advocating that we follow this design (option 3 in this comment) would also be fine.

My point is that regardless of which option we pick, we will need to have the functionality in this PR:

1. function signatures
2. methods to check and coerce types based on those signatures
3. methods to return the return type of the function

[jorgecarleitao]

And finally, this is how we would add the array function, that receives an arbitrary but uniformly-typed number of arguments: https://github.com/jorgecarleitao/arrow/pull/3/files

To summarize:

• math functions:
  • accept 1 argument of a fixed type (f32 or f64)
  • return a fixed type (f64)
• concat:
  • accepts an arbitrary number of arguments of fixed type (utf8 atm, but large is trivial to add)
  • returns a fixed type (utf8)
• array:
  • accepts an arbitrary number of arguments of variable but uniform types (i.e. all arguments must be of equal type)
  • returns FixedSizeList(input_types[0], input.len())

The API that I am proposing here addresses all these cases out of the box. The 3 PRs in my repo,

• all math accepts f32: https://github.com/jorgecarleitao/arrow/pull/1/files
• concatenate for utf8: https://github.com/jorgecarleitao/arrow/pull/2/files
• array for utf8: https://github.com/jorgecarleitao/arrow/pull/3/files

add support to each of them at the physical and logical level.

This PR also includes all the required coercion rules for this to work. E.g. array(f32, i32, u32) would be coerced to array(f32, f32, f32) (first argument dominates).

[jorgecarleitao]

After a long digression through the realm of built-ins, this has now been simplified and rebased against master.

@andygrove and @alamb , ready for a re-review.

Again, the core goal here is to allow users to use UDFs without having to worry about their return type.

I've incorporated all points from @alamb and @andygrove so far:

• Registry is now a trait
• Arc is used consistently
• built-ins continue to be available without access to the registry

[alamb]

I really like this code -- I think the UDF usage / definition is clearer now and the code looks better. Really nice work @jorgecarleitao

[alamb]

I want to verify my understanding -- this is a code cleanup that is not directly required for UDFs, right?

[jorgecarleitao]

Indirectly, it is related: expr.contains(&Expr::Wildcard) requires PartialEq, which was dropped in this line due to the addition of an Arc<ScalarFunction> in this line.

However, since if expr.contains(&Expr::Wildcard) thankfully was entirely optional, I dropped it ^_^

[andygrove]

Is this registry specific to scalar functions or will it also be used for aggregate functions? Perhaps we should name the method either function_registry or scalar_function_registry?

[andygrove]

LGTM and this is a nice improvement! My only comment is that maybe some of the naming could be made more consistent, or documentation made clearer .. the things that weren't immediately clear to me (although I didn't spend much time on the review) were:

• Is the registry just for user-defined scalar functions, or does it include the built-in scalar functions as well?
• The registry has methods like udf but should we consider naming scalar_udf? I wasn't sure what the plan was for user-defined aggregate functions (UDAFs)

[jorgecarleitao]

I agree with you, @andygrove .

• Built-in functions are not part of the registry as they are exposed directly from the prelude. So, maybe udf_registry could be more explicit about this.

• I thought about FunctionRegistry being used for both udfs and UDAFs (user-defined aggregate function, à la spark), so that a user does not have to call df.registry() or df.aggregate_registry(), they just need to remember about one.

• I thought about FunctionRegistry::udf and FunctionRegistry::udaf, just because short names make it easier to read statements. But, as always, it is a trade-off between short vs understanding.

I have no strong opinions about naming nor UX here: I will implement whatever you agree upon :-). My main concern was to fix the data type thing :P

[andygrove]

@jorgecarleitao @alamb I'm catching up on the PRs today. It looks like this one is ready to merge?

[jorgecarleitao]

I think so, @andygrove . There is probably some renaming once we have UDAFs. For now, I think it is fine.