Simplify AsyncScalarUdfImpl so it extends ScalarUdfImpl (#16523)

* Simplify AsyncScalarUdfImpl so it extends ScalarUdfImpl

* Update one example

* Update one example

* prettier

Author: alamb

datafusion-examples/examples/async_udf.rs

@@ -20,15 +20,14 @@ use arrow::compute::kernels::cmp::eq;
 use arrow_schema::{DataType, Field, Schema};
 use async_trait::async_trait;
 use datafusion::common::error::Result;
-use datafusion::common::internal_err;
 use datafusion::common::types::{logical_int64, logical_string};
 use datafusion::common::utils::take_function_args;
+use datafusion::common::{internal_err, not_impl_err};
 use datafusion::config::ConfigOptions;
-use datafusion::logical_expr::async_udf::{
-    AsyncScalarFunctionArgs, AsyncScalarUDF, AsyncScalarUDFImpl,
-};
+use datafusion::logical_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
 use datafusion::logical_expr::{
-    ColumnarValue, Signature, TypeSignature, TypeSignatureClass, Volatility,
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
+    TypeSignatureClass, Volatility,
 };
 use datafusion::logical_expr_common::signature::Coercion;
 use datafusion::physical_expr_common::datum::apply_cmp;
@@ -153,7 +152,7 @@ impl AsyncUpper {
 }
 
 #[async_trait]
-impl AsyncScalarUDFImpl for AsyncUpper {
+impl ScalarUDFImpl for AsyncUpper {
     fn as_any(&self) -> &dyn Any {
         self
     }
@@ -170,13 +169,20 @@ impl AsyncScalarUDFImpl for AsyncUpper {
         Ok(DataType::Utf8)
     }
 
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        not_impl_err!("AsyncUpper can only be called from async contexts")
+    }
+}
+
+#[async_trait]
+impl AsyncScalarUDFImpl for AsyncUpper {
     fn ideal_batch_size(&self) -> Option<usize> {
         Some(10)
     }
 
     async fn invoke_async_with_args(
         &self,
-        args: AsyncScalarFunctionArgs,
+        args: ScalarFunctionArgs,
         _option: &ConfigOptions,
     ) -> Result<ArrayRef> {
         trace!("Invoking async_upper with args: {:?}", args);
@@ -226,7 +232,7 @@ impl AsyncEqual {
 }
 
 #[async_trait]
-impl AsyncScalarUDFImpl for AsyncEqual {
+impl ScalarUDFImpl for AsyncEqual {
     fn as_any(&self) -> &dyn Any {
         self
     }
@@ -243,9 +249,16 @@ impl AsyncScalarUDFImpl for AsyncEqual {
         Ok(DataType::Boolean)
     }
 
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        not_impl_err!("AsyncEqual can only be called from async contexts")
+    }
+}
+
+#[async_trait]
+impl AsyncScalarUDFImpl for AsyncEqual {
     async fn invoke_async_with_args(
         &self,
-        args: AsyncScalarFunctionArgs,
+        args: ScalarFunctionArgs,
         _option: &ConfigOptions,
     ) -> Result<ArrayRef> {
         let [arg1, arg2] = take_function_args(self.name(), &args.args)?;

datafusion/core/src/physical_planner.rs

@@ -779,9 +779,11 @@ impl DefaultPhysicalPlanner {
                 let runtime_expr =
                     self.create_physical_expr(predicate, input_dfschema, session_state)?;
 
+                let input_schema = input.schema();
                 let filter = match self.try_plan_async_exprs(
-                    input.schema().fields().len(),
+                    input_schema.fields().len(),
                     PlannedExprResult::Expr(vec![runtime_expr]),
+                    input_schema.as_arrow(),
                 )? {
                     PlanAsyncExpr::Sync(PlannedExprResult::Expr(runtime_expr)) => {
                         FilterExec::try_new(Arc::clone(&runtime_expr[0]), physical_input)?
@@ -2082,6 +2084,7 @@ impl DefaultPhysicalPlanner {
         match self.try_plan_async_exprs(
             num_input_columns,
             PlannedExprResult::ExprWithName(physical_exprs),
+            input_physical_schema.as_ref(),
         )? {
             PlanAsyncExpr::Sync(PlannedExprResult::ExprWithName(physical_exprs)) => Ok(
                 Arc::new(ProjectionExec::try_new(physical_exprs, input_exec)?),
@@ -2104,18 +2107,19 @@ impl DefaultPhysicalPlanner {
         &self,
         num_input_columns: usize,
         physical_expr: PlannedExprResult,
+        schema: &Schema,
     ) -> Result<PlanAsyncExpr> {
         let mut async_map = AsyncMapper::new(num_input_columns);
         match &physical_expr {
             PlannedExprResult::ExprWithName(exprs) => {
                 exprs
                     .iter()
-                    .try_for_each(|(expr, _)| async_map.find_references(expr))?;
+                    .try_for_each(|(expr, _)| async_map.find_references(expr, schema))?;
             }
             PlannedExprResult::Expr(exprs) => {
                 exprs
                     .iter()
-                    .try_for_each(|expr| async_map.find_references(expr))?;
+                    .try_for_each(|expr| async_map.find_references(expr, schema))?;
             }
         }
 

datafusion/expr/src/async_udf.rs

@@ -17,7 +17,7 @@
 
 use crate::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl};
 use arrow::array::ArrayRef;
-use arrow::datatypes::{DataType, Field, FieldRef, SchemaRef};
+use arrow::datatypes::{DataType, FieldRef};
 use async_trait::async_trait;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
@@ -35,34 +35,7 @@ use std::sync::Arc;
 ///
 /// The name is chosen to mirror ScalarUDFImpl
 #[async_trait]
-pub trait AsyncScalarUDFImpl: Debug + Send + Sync {
-    /// the function cast as any
-    fn as_any(&self) -> &dyn Any;
-
-    /// The name of the function
-    fn name(&self) -> &str;
-
-    /// The signature of the function
-    fn signature(&self) -> &Signature;
-
-    /// The return type of the function
-    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType>;
-
-    /// What type will be returned by this function, given the arguments?
-    ///
-    /// By default, this function calls [`Self::return_type`] with the
-    /// types of each argument.
-    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
-        let data_types = args
-            .arg_fields
-            .iter()
-            .map(|f| f.data_type())
-            .cloned()
-            .collect::<Vec<_>>();
-        let return_type = self.return_type(&data_types)?;
-        Ok(Arc::new(Field::new(self.name(), return_type, true)))
-    }
-
+pub trait AsyncScalarUDFImpl: ScalarUDFImpl {
     /// The ideal batch size for this function.
     ///
     /// This is used to determine what size of data to be evaluated at once.
@@ -74,7 +47,7 @@ pub trait AsyncScalarUDFImpl: Debug + Send + Sync {
     /// Invoke the function asynchronously with the async arguments
     async fn invoke_async_with_args(
         &self,
-        args: AsyncScalarFunctionArgs,
+        args: ScalarFunctionArgs,
         option: &ConfigOptions,
     ) -> Result<ArrayRef>;
 }
@@ -107,7 +80,7 @@ impl AsyncScalarUDF {
     /// Invoke the function asynchronously with the async arguments
     pub async fn invoke_async_with_args(
         &self,
-        args: AsyncScalarFunctionArgs,
+        args: ScalarFunctionArgs,
         option: &ConfigOptions,
     ) -> Result<ArrayRef> {
         self.inner.invoke_async_with_args(args, option).await
@@ -145,10 +118,3 @@ impl Display for AsyncScalarUDF {
         write!(f, "AsyncScalarUDF: {}", self.inner.name())
     }
 }
-
-#[derive(Debug)]
-pub struct AsyncScalarFunctionArgs {
-    pub args: Vec<ColumnarValue>,
-    pub number_rows: usize,
-    pub schema: SchemaRef,
-}

datafusion/physical-expr/src/async_scalar_function.rs

@@ -17,11 +17,12 @@
 
 use crate::ScalarFunctionExpr;
 use arrow::array::{make_array, MutableArrayData, RecordBatch};
-use arrow::datatypes::{DataType, Field, Schema};
+use arrow::datatypes::{DataType, Field, FieldRef, Schema};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::Result;
 use datafusion_common::{internal_err, not_impl_err};
-use datafusion_expr::async_udf::{AsyncScalarFunctionArgs, AsyncScalarUDF};
+use datafusion_expr::async_udf::AsyncScalarUDF;
+use datafusion_expr::ScalarFunctionArgs;
 use datafusion_expr_common::columnar_value::ColumnarValue;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use std::any::Any;
@@ -36,6 +37,8 @@ pub struct AsyncFuncExpr {
     pub name: String,
     /// The actual function (always `ScalarFunctionExpr`)
     pub func: Arc<dyn PhysicalExpr>,
+    /// The field that this function will return
+    return_field: FieldRef,
 }
 
 impl Display for AsyncFuncExpr {
@@ -59,17 +62,23 @@ impl Hash for AsyncFuncExpr {
 
 impl AsyncFuncExpr {
     /// create a new AsyncFuncExpr
-    pub fn try_new(name: impl Into<String>, func: Arc<dyn PhysicalExpr>) -> Result<Self> {
+    pub fn try_new(
+        name: impl Into<String>,
+        func: Arc<dyn PhysicalExpr>,
+        schema: &Schema,
+    ) -> Result<Self> {
         let Some(_) = func.as_any().downcast_ref::<ScalarFunctionExpr>() else {
             return internal_err!(
                 "unexpected function type, expected ScalarFunctionExpr, got: {:?}",
                 func
             );
         };
 
+        let return_field = func.return_field(schema)?;
         Ok(Self {
             name: name.into(),
             func,
+            return_field,
         })
     }
 
@@ -128,6 +137,12 @@ impl AsyncFuncExpr {
             );
         };
 
+        let arg_fields = scalar_function_expr
+            .args()
+            .iter()
+            .map(|e| e.return_field(batch.schema_ref()))
+            .collect::<Result<Vec<_>>>()?;
+
         let mut result_batches = vec![];
         if let Some(ideal_batch_size) = self.ideal_batch_size()? {
             let mut remainder = batch.clone();
@@ -148,10 +163,11 @@ impl AsyncFuncExpr {
                 result_batches.push(
                     async_udf
                         .invoke_async_with_args(
-                            AsyncScalarFunctionArgs {
-                                args: args.to_vec(),
+                            ScalarFunctionArgs {
+                                args,
+                                arg_fields: arg_fields.clone(),
                                 number_rows: current_batch.num_rows(),
-                                schema: current_batch.schema(),
+                                return_field: Arc::clone(&self.return_field),
                             },
                             option,
                         )
@@ -168,10 +184,11 @@ impl AsyncFuncExpr {
             result_batches.push(
                 async_udf
                     .invoke_async_with_args(
-                        AsyncScalarFunctionArgs {
+                        ScalarFunctionArgs {
                             args: args.to_vec(),
+                            arg_fields,
                             number_rows: batch.num_rows(),
-                            schema: batch.schema(),
+                            return_field: Arc::clone(&self.return_field),
                         },
                         option,
                     )
@@ -223,6 +240,7 @@ impl PhysicalExpr for AsyncFuncExpr {
         Ok(Arc::new(AsyncFuncExpr {
             name: self.name.clone(),
             func: new_func,
+            return_field: Arc::clone(&self.return_field),
         }))
     }
 

datafusion/physical-plan/src/async_func.rs

@@ -245,6 +245,7 @@ impl AsyncMapper {
     pub fn find_references(
         &mut self,
         physical_expr: &Arc<dyn PhysicalExpr>,
+        schema: &Schema,
     ) -> Result<()> {
         // recursively look for references to async functions
         physical_expr.apply(|expr| {
@@ -256,6 +257,7 @@ impl AsyncMapper {
                     self.async_exprs.push(Arc::new(AsyncFuncExpr::try_new(
                         next_name,
                         Arc::clone(expr),
+                        schema,
                     )?));
                 }
             }