milenkovicm commented on code in PR #9304:
URL: https://github.com/apache/arrow-datafusion/pull/9304#discussion_r1512463416
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -514,6 +519,102 @@ async fn deregister_udf() -> Result<()> {
Ok(())
}
+#[derive(Debug)]
+struct CastToI64UDF {
+ signature: Signature,
+}
+
+impl CastToI64UDF {
+ fn new() -> Self {
+ Self {
+ signature: Signature::any(1, Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for CastToI64UDF {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+ fn name(&self) -> &str {
+ "cast_to_i64"
+ }
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ // Demonstrate simplifying a UDF
+ fn simplify(
+ &self,
+ mut args: Vec<Expr>,
+ info: &dyn SimplifyInfo,
+ ) -> Result<ExprSimplifyResult> {
+ // DataFusion should have ensured the function is called with just a
+ // single argument
+ assert_eq!(args.len(), 1);
+ let arg = args.pop().unwrap();
+
+ // Note that Expr::cast_to requires an ExprSchema but simplify gets a
+ // SimplifyInfo so we have to replicate some of the casting logic here.
+
+ let source_type = info.get_data_type(&arg)?;
+ let new_expr = if source_type == DataType::Int64 {
+ // the argument's data type is already the correct type
+ arg
+ } else {
+ // need to use an actual cast to get the correct type
+ Expr::Cast(datafusion_expr::Cast {
+ expr: Box::new(arg),
+ data_type: DataType::Int64,
+ })
+ };
+ // return the newly written argument to DataFusion
+ Ok(ExprSimplifyResult::Simplified(new_expr))
+ }
+
+ // Casting should be done in `simplify`, so we just return the first
argument
Review Comment:
comment is a bit outdated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
