fix(udf): fix missing chunks (risingwavelabs#9025)

Signed-off-by: Runji Wang <wangrunji0408@163.com>
Co-authored-by: lmatz <lmatz823@gmail.com>

src/common/src/array/arrow.rs

@@ -198,6 +198,11 @@ macro_rules! converts {
                 array.iter().collect()
             }
         }
+        impl From<&[$ArrowType]> for $ArrayType {
+            fn from(arrays: &[$ArrowType]) -> Self {
+                arrays.iter().flat_map(|a| a.iter()).collect()
+            }
+        }
     };
     // convert values using FromIntoArrow
     ($ArrayType:ty, $ArrowType:ty, @map) => {
@@ -218,6 +223,19 @@ macro_rules! converts {
                     .collect()
             }
         }
+        impl From<&[$ArrowType]> for $ArrayType {
+            fn from(arrays: &[$ArrowType]) -> Self {
+                arrays
+                    .iter()
+                    .flat_map(|a| a.iter())
+                    .map(|o| {
+                        o.map(|v| {
+                            <<$ArrayType as Array>::RefItem<'_> as FromIntoArrow>::from_arrow(v)
+                        })
+                    })
+                    .collect()
+            }
+        }
     };
 }
 converts!(BoolArray, arrow_array::BooleanArray);

src/udf/Cargo.toml

@@ -14,6 +14,7 @@ normal = ["workspace-hack"]
 arrow-array = "36"
 arrow-flight = "36"
 arrow-schema = "36"
+arrow-select = "36"
 futures-util = "0.3.25"
 thiserror = "1"
 tokio = { version = "0.2", package = "madsim-tokio", features = ["rt", "macros"] }

src/udf/src/lib.rs

@@ -72,7 +72,20 @@ impl ArrowFlightUdfClient {
     /// Call a function.
     pub async fn call(&self, id: &str, input: RecordBatch) -> Result<RecordBatch> {
         let mut output_stream = self.call_stream(id, stream::once(async { input })).await?;
-        output_stream.next().await.ok_or(Error::NoReturned)?
+        // TODO: support no output
+        let head = output_stream.next().await.ok_or(Error::NoReturned)??;
+        let mut remaining = vec![];
+        while let Some(batch) = output_stream.next().await {
+            remaining.push(batch?);
+        }
+        if remaining.is_empty() {
+            Ok(head)
+        } else {
+            Ok(arrow_select::concat::concat_batches(
+                &head.schema(),
+                std::iter::once(&head).chain(remaining.iter()),
+            )?)
+        }
     }
 
     /// Call a function with streaming input and output.
@@ -157,6 +170,8 @@ pub enum Error {
         expected: String,
         actual: String,
     },
+    #[error("arrow error: {0}")]
+    Arrow(#[from] arrow_schema::ArrowError),
     #[error("UDF service returned no data")]
     NoReturned,
 }