fix: udtf single-row expectation

Author: KKould

src/query/expression/src/utils/udf_client.rs

@@ -302,15 +302,14 @@ impl UDFFlightClient {
         &mut self,
         name: &str,
         func_name: &str,
-        num_rows: usize,
+        num_rows: Option<usize>,
         args: Vec<BlockEntry>,
         return_type: &DataType,
     ) -> Result<DataBlock> {
         let instant = Instant::now();
 
         Profile::record_usize_profile(ProfileStatisticsName::ExternalServerRequestCount, 1);
         record_running_requests_external_start(name, 1);
-        record_request_external_batch_rows(func_name, num_rows);
 
         let args = args
             .into_iter()
@@ -333,7 +332,9 @@ impl UDFFlightClient {
             .collect::<Vec<_>>();
         let data_schema = DataSchema::new(fields);
 
-        let input_batch = DataBlock::new(args, num_rows)
+        // at least 1 for `UDFFlightClient::batch_rows`
+        let input_num_rows = args.first().map(|entry| entry.len()).unwrap_or(1);
+        let input_batch = DataBlock::new(args, input_num_rows)
             .to_record_batch_with_dataschema(&data_schema)
             .map_err(|err| ErrorCode::from_string(format!("{err}")))?;
 
@@ -359,13 +360,17 @@ impl UDFFlightClient {
             ));
         }
 
-        if result_block.num_rows() != num_rows {
-            return Err(ErrorCode::UDFDataError(format!(
-                "UDF server should return {} rows, but it returned {} rows",
-                num_rows,
-                result_block.num_rows()
-            )));
+        if let Some(expected_rows) = num_rows {
+            if result_block.num_rows() != expected_rows {
+                return Err(ErrorCode::UDFDataError(format!(
+                    "UDF server should return {} rows, but it returned {} rows",
+                    expected_rows,
+                    result_block.num_rows()
+                )));
+            }
         }
+        record_request_external_batch_rows(func_name, result_block.num_rows());
+
         if return_type.remove_nullable().is_tuple() && result_fields.len() > 1 {
             if let DataType::Tuple(tys) = return_type.remove_nullable() {
                 if tys

src/query/service/src/pipelines/builders/builder_udtf.rs

@@ -123,7 +123,7 @@ impl UdtfServerSource {
 
         debug_assert!(func.return_ty.is_tuple());
         let result = client
-            .do_exchange(&func.name, &func.func_name, 1, args, &func.return_ty)
+            .do_exchange(&func.name, &func.func_name, None, args, &func.return_ty)
             .await?;
 
         drop(permit);

src/query/service/src/pipelines/processors/transforms/transform_udf_server.rs

@@ -140,7 +140,7 @@ impl TransformUdfServer {
             .do_exchange(
                 &func.name,
                 &func.func_name,
-                num_rows,
+                Some(num_rows),
                 block_entries,
                 &func.data_type,
             )

src/query/service/src/table_functions/others/udf.rs

@@ -163,7 +163,7 @@ impl Table for UdfEchoTable {
         let return_type = DataType::Nullable(Box::new(DataType::String));
 
         let result = client
-            .do_exchange(name, name, num_rows, block_entries, &return_type)
+            .do_exchange(name, name, Some(num_rows), block_entries, &return_type)
             .await?;
 
         let scalar = unsafe { result.get_by_offset(0).index_unchecked(0) };

src/query/sql/src/planner/semantic/type_check.rs

@@ -5347,7 +5347,7 @@ impl<'a> TypeChecker<'a> {
             .do_exchange(
                 name,
                 &udf_definition.handler,
-                num_rows,
+                Some(num_rows),
                 block_entries,
                 &udf_definition.return_type,
             )

tests/sqllogictests/suites/udf_server/udf_server_test.test

@@ -717,8 +717,10 @@ query I
 SELECT * from stage_summary_udtf(@s3_stage/output/2024, 21)
 ----
 s3_stage External test output/2024 21 s3_stage:test:output/2024:21
+s3_stage External test output/2024 22 s3_stage:test:output/2024:22
 
 query I
 SELECT * from multi_stage_process_udtf(@s3_stage/input/2024/, @s3_stage/output/2024, 21)
 ----
 s3_stage s3_stage test test input/2024/ output/2024 29
+s3_stage s3_stage test test input/2024/ output/2024 30

tests/udf/udf_server.py

@@ -454,17 +454,23 @@ def stage_summary_udtf(data_stage: StageLocation, value: int):
     assert data_stage.stage_type.lower() == "external"
     assert data_stage.storage
     bucket = _stage_bucket(data_stage)
-    summary = f"{data_stage.stage_name}:{bucket}:{data_stage.relative_path}:{value}"
-    return [
-        {
-            "stage_name": data_stage.stage_name or "",
-            "stage_type": data_stage.stage_type or "",
-            "bucket": bucket,
-            "relative_path": data_stage.relative_path or "",
-            "value": value,
-            "summary": summary,
-        }
-    ]
+    rows = []
+    for offset in (0, 1):
+        current_value = value + offset
+        summary = (
+            f"{data_stage.stage_name}:{bucket}:{data_stage.relative_path}:{current_value}"
+        )
+        rows.append(
+            {
+                "stage_name": data_stage.stage_name or "",
+                "stage_type": data_stage.stage_type or "",
+                "bucket": bucket,
+                "relative_path": data_stage.relative_path or "",
+                "value": current_value,
+                "summary": summary,
+            }
+        )
+    return rows
 
 
 @udf(
@@ -507,18 +513,22 @@ def multi_stage_process_udtf(
     assert output_stage.stage_type.lower() == "external"
     input_bucket = _stage_bucket(input_stage)
     output_bucket = _stage_bucket(output_stage)
-    result = value + len(input_bucket) + len(output_bucket)
-    return [
-        {
-            "input_stage": input_stage.stage_name or "",
-            "output_stage": output_stage.stage_name or "",
-            "input_bucket": input_bucket,
-            "output_bucket": output_bucket,
-            "input_relative_path": input_stage.relative_path or "",
-            "output_relative_path": output_stage.relative_path or "",
-            "result": result,
-        }
-    ]
+    rows = []
+    for offset in (0, 1):
+        current_value = value + offset
+        result = current_value + len(input_bucket) + len(output_bucket)
+        rows.append(
+            {
+                "input_stage": input_stage.stage_name or "",
+                "output_stage": output_stage.stage_name or "",
+                "input_bucket": input_bucket,
+                "output_bucket": output_bucket,
+                "input_relative_path": input_stage.relative_path or "",
+                "output_relative_path": output_stage.relative_path or "",
+                "result": result,
+            }
+        )
+    return rows
 
 
 @udf(