ARROW-9756: [Rust] [DataFusion] Added support for scalar UDFs of arbitrary return types#7974
Closed
jorgecarleitao wants to merge 13 commits intoapache:masterfrom
jorgecarleitao:scalar
Closed
ARROW-9756: [Rust] [DataFusion] Added support for scalar UDFs of arbitrary return types#7974jorgecarleitao wants to merge 13 commits intoapache:masterfrom jorgecarleitao:scalar
jorgecarleitao wants to merge 13 commits intoapache:masterfrom
jorgecarleitao:scalar
Conversation
jorgecarleitao
changed the title
Member
Author
|
FYI @andygrove @alamb @houqp : this is a draft because it depends on other PRs being reviewed and accepted. |
jorgecarleitao
commented
Aug 16, 2020
rust/datafusion/src/logicalplan.rs
Outdated
Member
Author
There was a problem hiding this comment.
I had to drop partialEq for expressions, and therefore also dropped this since it was strictly not needed.
jorgecarleitao
commented
Aug 16, 2020
jorgecarleitao
commented
Aug 16, 2020
Currently a UDF's argument can only be a single type. This PR adds support for multiple types per argument, thus allowing users to register UDFs that can handle multiple types at once.
This operation is already by the optimizer and is more complete.
This allows to declare UDFs that support multiple types. The existing UDFs (math) now support float32 and float64.
jorgecarleitao
commented
Aug 17, 2020
| #[derive(Debug)] | ||
| struct SumAccumulator { | ||
| sum: Option<ScalarValue>, | ||
| pub sum: Option<ScalarValue>, |
jorgecarleitao
commented
Aug 17, 2020
|
|
||
| pub mod dataframe; | ||
| pub mod datasource; | ||
| mod datatyped; |
Member
Author
There was a problem hiding this comment.
pub so that others can implement a different DataTyped?
jorgecarleitao
commented
Aug 17, 2020
| args: Vec<Expr>, | ||
| /// The `DataType` the expression will yield | ||
| return_type: DataType, | ||
| return_type: ReturnType, |
Member
Author
There was a problem hiding this comment.
THIS IS A MAJOR CHANGE: it makes Expr hold an Arc<&dyn>, which make them unserializable. I was unable to find another way, unfortunately.
jorgecarleitao
commented
Aug 17, 2020
| Expr::AggregateFunction { return_type, .. } => Ok(return_type.clone()), | ||
| Expr::ScalarFunction { | ||
| args, return_type, .. | ||
| } => return_type(&args.iter().map(|x| x.as_datatyped()).collect(), schema), |
Member
Author
There was a problem hiding this comment.
not very beautiful, but was unable to find a better pattern for this :(
jorgecarleitao
commented
Aug 17, 2020
| match self.schema_provider.get_function_meta(&name) { | ||
| Some(fm) => { | ||
| let args = if name == "count" { | ||
| // optimization to avoid computing expressions |
Member
Author
There was a problem hiding this comment.
This should probably be moved to another place (an optimizer?). For now I left it where it here as it was.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR is done on top of #7971 ,
Its current runtime consequence is that all our math functions now return float32 or float64 depending on their incoming column (and use float32 for other numeric types). Its API consequences is that it allows to register UDFs of variable return types.
This PR hits both physical plans and logical plans, and abstracts a bit the typing of some of our structs. In particular, it introduces a new trait that only cares about the return datatype of an object (PhysicalExpr and LogicalExpr).