Properly return a scalar value from a UDF when only input is scalar

Fixes #699
ArroyoSystems · Oct 10, 2024 · 2f05a4b · 2f05a4b
1 parent 655d388
commit 2f05a4b
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/crates/arroyo-udf/arroyo-udf-host/src/lib.rs b/crates/arroyo-udf/arroyo-udf-host/src/lib.rs
@@ -2,7 +2,7 @@
 mod test;
 
 use anyhow::{anyhow, bail};
-use arrow::array::{make_array, Array, ArrayData, ArrayRef, UInt64Array};
+use arrow::array::{make_array, Array, ArrayData, ArrayRef, Datum, UInt64Array};
 use arrow::datatypes::DataType;
 use arrow::ffi::from_ffi;
 use arroyo_udf_common::async_udf::{DrainResult, SendableFfiAsyncUdfHandle};
@@ -389,6 +389,9 @@ impl ScalarUDFImpl for SyncUdfDylib {
             })
             .max()
             .unwrap();
+
+        let all_scalars = args.iter()
+            .all(|c| matches!(c, ColumnarValue::Scalar(_)));
 
         let args = args
             .iter()
@@ -402,7 +405,13 @@ impl ScalarUDFImpl for SyncUdfDylib {
         match result {
             RunResult::Ok(FfiArraySchema(array, schema)) => {
                 let result_array = unsafe { from_ffi(array, &schema).unwrap() };
-                Ok(ColumnarValue::Array(make_array(result_array)))
+
+                let array = make_array(result_array);
+                if all_scalars {
+                    Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(&array, 0)?))
+                } else {
+                    Ok(ColumnarValue::Array(array))                    
+                }
             }
             RunResult::Err => {
                 panic!("panic in UDF {}", self.name);