From 1d95e13dcad1afe4debc67fd078290a4cc6d5f15 Mon Sep 17 00:00:00 2001
From: Micah Wylde <micah@arroyo.systems>
Date: Wed, 9 Oct 2024 19:16:27 -0700
Subject: [PATCH] Properly return a scalar value from a UDF when only input is
 scalar (#753)

Fixes #699
---
 crates/arroyo-udf/arroyo-udf-host/src/lib.rs | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/crates/arroyo-udf/arroyo-udf-host/src/lib.rs b/crates/arroyo-udf/arroyo-udf-host/src/lib.rs
index cc600b829..25bcbe2cb 100644
--- 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);