fix(udf): fix wrong number of rows (risingwavelabs#9003)

Signed-off-by: Runji Wang <wangrunji0408@163.com>
Little-Wallace · Apr 5, 2023 · ce25cf5 · ce25cf5
1 parent db072cb
commit ce25cf5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git 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
@@ -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.