Migrate arrow_cast to a UDF
#9610
There are no files selected for viewing
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
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,63 +15,123 @@ | |
| // specific language governing permissions and limitations | ||
| // under the License. | ||
|
|
||
| //! [`ArrowCastFunc`]: Implementation of the `arrow_cast` | ||
|
|
||
| use std::any::Any; | ||
| use std::{fmt::Display, iter::Peekable, str::Chars, sync::Arc}; | ||
|
|
||
| use arrow_schema::{DataType, Field, IntervalUnit, TimeUnit}; | ||
| use datafusion_common::{ | ||
| internal_err, plan_datafusion_err, plan_err, DataFusionError, ExprSchema, Result, | ||
| ScalarValue, | ||
| }; | ||
|
|
||
| use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo}; | ||
| use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility}; | ||
|
|
||
| /// Implements casting to arbitrary arrow types (rather than SQL types) | ||
| /// | ||
| /// Note that the `arrow_cast` function is somewhat special in that its | ||
| /// return depends only on the *value* of its second argument (not its type) | ||
| /// | ||
| /// It is implemented by calling the same underlying arrow `cast` kernel as | ||
| /// normal SQL casts. | ||
| /// | ||
| /// For example to cast to `int` using SQL (which is then mapped to the arrow | ||
| /// type `Int32`) | ||
| /// | ||
| /// ```sql | ||
| /// select cast(column_x as int) ... | ||
| /// ``` | ||
| /// | ||
| /// You can use the `arrow_cast` functiont to cast to a specific arrow type | ||
| /// | ||
| /// For example | ||
| /// ```sql | ||
| /// select arrow_cast(column_x, 'Float64') | ||
| /// ``` | ||
| #[derive(Debug)] | ||
| pub(super) struct ArrowCastFunc { | ||
| signature: Signature, | ||
| } | ||
|
|
||
| impl ArrowCastFunc { | ||
| pub fn new() -> Self { | ||
| Self { | ||
| signature: Signature::any(2, Volatility::Immutable), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl ScalarUDFImpl for ArrowCastFunc { | ||
| fn as_any(&self) -> &dyn Any { | ||
| self | ||
| } | ||
|
|
||
| fn name(&self) -> &str { | ||
| "arrow_cast" | ||
| } | ||
|
|
||
| fn signature(&self) -> &Signature { | ||
| &self.signature | ||
| } | ||
|
|
||
| fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { | ||
| parse_data_type(&arg_types[1].to_string()) | ||
|
||
| } | ||
|
|
||
| fn return_type_from_exprs( | ||
| &self, | ||
| args: &[Expr], | ||
| _schema: &dyn ExprSchema, | ||
| _arg_types: &[DataType], | ||
| ) -> Result<DataType> { | ||
| data_type_from_args(args) | ||
| } | ||
|
|
||
| fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> { | ||
| internal_err!("arrow_cast should have been simplified to cast") | ||
| } | ||
|
|
||
| fn simplify( | ||
| &self, | ||
| mut args: Vec<Expr>, | ||
| info: &dyn SimplifyInfo, | ||
| ) -> Result<ExprSimplifyResult> { | ||
| // convert this into a real cast | ||
| let target_type = data_type_from_args(&args)?; | ||
|
Contributor
Author
There was a problem hiding this comment. This simplify logic mirrors the previous behavior in that arrow_cast is replaced with a normal cast |
||
| // remove second (type) argument | ||
| args.pop().unwrap(); | ||
| let arg = args.pop().unwrap(); | ||
|
|
||
| let source_type = info.get_data_type(&arg)?; | ||
| let new_expr = if source_type == target_type { | ||
| // the argument's data type is already the correct type | ||
| arg | ||
| } else { | ||
| // Use an actual cast to get the correct type | ||
| Expr::Cast(datafusion_expr::Cast { | ||
| expr: Box::new(arg), | ||
| data_type: target_type, | ||
| }) | ||
| }; | ||
| // return the newly written argument to DataFusion | ||
| Ok(ExprSimplifyResult::Simplified(new_expr)) | ||
| } | ||
| } | ||
|
|
||
| /// Returns the requested type from the arguments | ||
| fn data_type_from_args(args: &[Expr]) -> Result<DataType> { | ||
| if args.len() != 2 { | ||
| return plan_err!("arrow_cast needs 2 arguments, {} provided", args.len()); | ||
| } | ||
| let Expr::Literal(ScalarValue::Utf8(Some(val))) = &args[1] else { | ||
| return plan_err!( | ||
| "arrow_cast requires its second argument to be a constant string, got {:?}", | ||
| &args[1] | ||
| ); | ||
| }; | ||
| parse_data_type(val) | ||
| } | ||
|
|
||
| /// Parses `str` into a `DataType`. | ||
|
|
@@ -80,22 +140,8 @@ pub fn create_arrow_cast(mut args: Vec<Expr>, schema: &DFSchema) -> Result<Expr> | |
| /// impl, and maintains the invariant that | ||
| /// `parse_data_type(data_type.to_string()) == data_type` | ||
| /// | ||
| /// Remove if added to arrow: <https://github.com/apache/arrow-rs/issues/3821> | ||
| fn parse_data_type(val: &str) -> Result<DataType> { | ||
| Parser::new(val).parse() | ||
| } | ||
|
|
||
|
|
@@ -844,7 +890,6 @@ mod test { | |
| assert!(message.contains("Must be a supported arrow type name such as 'Int32' or 'Timestamp(Nanosecond, None)'")); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
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
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These differences are due to the fact that
arrow_castis just a normal function now rather than a special case in the parser. Thus the naming reflects normal function namingThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Interesting I got stuck implementing the simpliy function because I thought it should convert
arrow_cast(t.values,Utf8(\"Utf8\"))tot.valuesand other similar cases as well.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah -- this is pretty tricky.
arrow_castwas quite special in the parser, so now that it is handled like a normal function it has the same (somewhat strange) function effect of column naming