From ce25cf5a201e241fde8c8d6dac0cb5326c2629a3 Mon Sep 17 00:00:00 2001
From: Runji Wang <wangrunji0408@163.com>
Date: Wed, 5 Apr 2023 14:20:55 +0800
Subject: [PATCH] fix(udf): fix wrong number of rows (#9003)

Signed-off-by: Runji Wang <wangrunji0408@163.com>
---
 src/expr/src/expr/expr_udf.rs | 16 +++++++++++-----
 src/udf/src/lib.rs            |  2 --
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/expr/src/expr/expr_udf.rs b/src/expr/src/expr/expr_udf.rs
index b6c5ff7f8e23..67f1befa36f1 100644
--- a/src/expr/src/expr/expr_udf.rs
+++ b/src/expr/src/expr/expr_udf.rs
@@ -52,15 +52,21 @@ impl Expression for UdfExpression {
             columns.push(array.as_ref().into());
         }
         let opts =
-            arrow_array::RecordBatchOptions::default().with_row_count(Some(input.cardinality()));
+            arrow_array::RecordBatchOptions::default().with_row_count(Some(input.capacity()));
         let input =
             arrow_array::RecordBatch::try_new_with_options(self.arg_schema.clone(), columns, &opts)
                 .expect("failed to build record batch");
         let output = self.client.call(&self.identifier, input).await?;
-        let arrow_array = output
-            .columns()
-            .get(0)
-            .ok_or(risingwave_udf::Error::NoColumn)?;
+        if output.num_rows() != vis.len() {
+            bail!(
+                "UDF returned {} rows, but expected {}",
+                output.num_rows(),
+                vis.len(),
+            );
+        }
+        let Some(arrow_array) = output.columns().get(0) else {
+            bail!("UDF returned no columns");
+        };
         let mut array = ArrayImpl::from(arrow_array);
         array.set_bitmap(array.null_bitmap() & vis);
         Ok(Arc::new(array))
diff --git a/src/udf/src/lib.rs b/src/udf/src/lib.rs
index 0256dfe0ed6e..f35a2f832eb4 100644
--- a/src/udf/src/lib.rs
+++ b/src/udf/src/lib.rs
@@ -159,8 +159,6 @@ pub enum Error {
     },
     #[error("UDF service returned no data")]
     NoReturned,
-    #[error("UDF service returned a batch with no column")]
-    NoColumn,
 }
 
 /// Check if two list of data types match, ignoring field names.