alamb commented on code in PR #9333:
URL: https://github.com/apache/arrow-datafusion/pull/9333#discussion_r1508176155
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction]
statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+
datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+
datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) =>
Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
> lets say we have a SQL UDFs which are expressed by expr, we would need to
have some kind of registry (which is not udf_registry) which can be consulted
by plan analyser.
Another potential approach would be to encapsulate the code directly into
the ScalarUDF and then use an API like `simplify()` that @jayzhan211 and I are
working on in https://github.com/apache/arrow-datafusion/pull/9304 to replace
the `ScalarUDF` with its `Expr` equivelnt...
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -619,6 +624,117 @@ async fn verify_udf_return_type() -> Result<()> {
Ok(())
}
+#[derive(Debug, Default)]
+struct MockFunctionFactory {
+ pub captured_expr: Mutex<Option<Expr>>,
+}
+
+#[async_trait::async_trait]
+impl FunctionFactory for MockFunctionFactory {
+ #[doc = r" Crates and registers a function from [CreateFunction]
statement"]
+ #[must_use]
+ #[allow(clippy::type_complexity, clippy::type_repetition_in_bounds)]
+ async fn create(
+ &self,
+ _config: &SessionConfig,
+ statement: CreateFunction,
+ ) -> datafusion::error::Result<RegisterFunction> {
+ // this function is a mock for testing
+ // `CreateFunction` should be used to derive this function
+
+ let mock_add = Arc::new(|args: &[datafusion_expr::ColumnarValue]| {
+ let args = datafusion_expr::ColumnarValue::values_to_arrays(args)?;
+ let base =
+
datafusion_common::cast::as_float64_array(&args[0]).expect("cast failed");
+ let exponent =
+
datafusion_common::cast::as_float64_array(&args[1]).expect("cast failed");
+
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| match (base, exponent) {
+ (Some(base), Some(exponent)) =>
Some(base.add_wrapping(exponent)),
+ _ => None,
+ })
+ .collect::<arrow_array::Float64Array>();
+ Ok(datafusion_expr::ColumnarValue::from(
+ Arc::new(array) as arrow_array::ArrayRef
+ ))
+ });
+
+ let args = statement.args.unwrap();
+ let mock_udf = create_udf(
+ &statement.name,
+ vec![args[0].data_type.clone(), args[1].data_type.clone()],
+ Arc::new(statement.return_type.unwrap()),
+ datafusion_expr::Volatility::Immutable,
+ mock_add,
+ );
+
+ // capture expression so we can verify
+ // it has been parsed
+ *self.captured_expr.lock() = statement.params.return_;
+
+ Ok(RegisterFunction::Scalar(Arc::new(mock_udf)))
+ }
+
+ async fn remove(
Review Comment:
> lets say we have a SQL UDFs which are expressed by expr, we would need to
have some kind of registry (which is not udf_registry) which can be consulted
by plan analyser.
Another potential approach would be to encapsulate the code directly into
the ScalarUDF and then use an API like `simplify()` that @jayzhan211 and I are
working on in https://github.com/apache/arrow-datafusion/pull/9304 to replace
the `ScalarUDF` with its `Expr` equivalent...
--
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]
