Properly return a scalar value from a UDF when only input is scalar (#…

…753) Fixes #699
ArroyoSystems · Oct 10, 2024 · 1d95e13 · 1d95e13
1 parent 655d388
commit 1d95e13
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/crates/arroyo-udf/arroyo-udf-host/src/lib.rs b/crates/arroyo-udf/arroyo-udf-host/src/lib.rs
@@ -390,6 +390,8 @@ impl ScalarUDFImpl for SyncUdfDylib {
             .max()
             .unwrap();
 
+        let all_scalars = args.iter().all(|c| matches!(c, ColumnarValue::Scalar(_)));
+
         let args = args
             .iter()
             .map(|arg| arg.clone().into_array(num_rows).unwrap().to_data())
@@ -402,7 +404,15 @@ 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);