Change AsyncScalarUDFImpl::invoke_async_with_args return type to ColumnarValue

Author: geetanshjuneja

datafusion-examples/examples/async_udf.rs

@@ -199,7 +199,7 @@ impl AsyncScalarUDFImpl for AskLLM {
         &self,
         args: ScalarFunctionArgs,
         _option: &ConfigOptions,
-    ) -> Result<ArrayRef> {
+    ) -> Result<ColumnarValue> {
         // in a real UDF you would likely want to special case constant
         // arguments to improve performance, but this example converts the
         // arguments to arrays for simplicity.
@@ -234,6 +234,6 @@ impl AsyncScalarUDFImpl for AskLLM {
             })
             .collect();
 
-        Ok(Arc::new(result_array))
+        Ok(ColumnarValue::from(Arc::new(result_array) as ArrayRef))
     }
 }

datafusion/expr/src/async_udf.rs

@@ -16,7 +16,6 @@
 // under the License.
 
 use crate::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl};
-use arrow::array::ArrayRef;
 use arrow::datatypes::{DataType, FieldRef};
 use async_trait::async_trait;
 use datafusion_common::config::ConfigOptions;
@@ -50,7 +49,7 @@ pub trait AsyncScalarUDFImpl: ScalarUDFImpl {
         &self,
         args: ScalarFunctionArgs,
         option: &ConfigOptions,
-    ) -> Result<ArrayRef>;
+    ) -> Result<ColumnarValue>;
 }
 
 /// A scalar UDF that must be invoked using async methods
@@ -83,7 +82,7 @@ impl AsyncScalarUDF {
         &self,
         args: ScalarFunctionArgs,
         option: &ConfigOptions,
-    ) -> Result<ArrayRef> {
+    ) -> Result<ColumnarValue> {
         self.inner.invoke_async_with_args(args, option).await
     }
 }

datafusion/physical-expr/src/async_scalar_function.rs

@@ -196,7 +196,7 @@ impl AsyncFuncExpr {
             );
         }
 
-        let datas = result_batches
+        let datas = ColumnarValue::values_to_arrays(&result_batches)?
             .iter()
             .map(|b| b.to_data())
             .collect::<Vec<_>>();