ARROW-9756: [Rust] [DataFusion] Added support for scalar UDFs of arbitrary return types #7974
Conversation
This file contains 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
jorgecarleitao
changed the title
|
FYI @andygrove @alamb @houqp : this is a draft because it depends on other PRs being reviewed and accepted. |
jorgecarleitao
commented
Aug 16, 2020
| @@ -1087,7 +1101,7 @@ impl LogicalPlanBuilder { | |||
| /// Apply a projection | |||
| pub fn project(&self, expr: Vec<Expr>) -> Result<Self> { | |||
| let input_schema = self.plan.schema(); | |||
| let projected_expr = if expr.contains(&Expr::Wildcard) { | |||
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
| @@ -87,7 +95,7 @@ impl Debug for ScalarFunctionExpr { | |||
| .field("fun", &"<FUNC>") | |||
| .field("name", &self.name) | |||
| .field("args", &self.args) | |||
| .field("return_type", &self.return_type) | |||
| //.field("return_type", &self.return_type) | |||
jorgecarleitao
commented
Aug 16, 2020
| @@ -50,7 +58,7 @@ impl Debug for ScalarFunction { | |||
| f.debug_struct("ScalarFunction") | |||
| .field("name", &self.name) | |||
| .field("args", &self.args) | |||
| .field("return_type", &self.return_type) | |||
| //.field("return_type", &self.return_type) | |||
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
| @@ -156,7 +168,7 @@ macro_rules! sum_accumulate { | |||
|
|
|||
| #[derive(Debug)] | |||
| struct SumAccumulator { | |||
| sum: Option<ScalarValue>, | |||
| pub sum: Option<ScalarValue>, | |||
jorgecarleitao
commented
Aug 17, 2020
| @@ -31,6 +31,7 @@ extern crate sqlparser; | |||
|
|
|||
| pub mod dataframe; | |||
| pub mod datasource; | |||
| mod datatyped; | |||
There was a problem hiding this comment.
pub so that others can implement a different DataTyped?
jorgecarleitao
commented
Aug 17, 2020
| @@ -355,7 +355,7 @@ pub enum Expr { | |||
| /// List of expressions to feed to the functions as arguments | |||
| args: Vec<Expr>, | |||
| /// The `DataType` the expression will yield | |||
| return_type: DataType, | |||
| return_type: ReturnType, | |||
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), |
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 |
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.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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).