Merge pull request #217 from influxdata/crepererum/check-array-size

feat: check return row count of UDF

Author: crepererum

guests/evil/src/lib.rs

@@ -11,6 +11,7 @@ mod common;
 mod env;
 mod fs;
 mod net;
+mod return_data;
 mod root;
 mod runtime;
 mod spin;
@@ -46,6 +47,10 @@ impl Evil {
                 root: Box::new(common::root_empty),
                 udfs: Box::new(net::udfs),
             },
+            "return_data" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(return_data::udfs),
+            },
             "root::invalid_entry" => Self {
                 root: Box::new(root::invalid_entry::root),
                 udfs: Box::new(common::udfs_empty),

guests/evil/src/return_data.rs

@@ -0,0 +1,52 @@
+//! Payload that returns invalid data.
+
+use std::sync::Arc;
+
+use arrow::{array::StringArray, datatypes::DataType};
+use datafusion_common::error::Result as DataFusionResult;
+use datafusion_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+};
+
+/// UDF that return the wrong number of rows.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct WrongNumberOfRows;
+
+impl ScalarUDFImpl for WrongNumberOfRows {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "wrong-number-of-rows"
+    }
+
+    fn signature(&self) -> &Signature {
+        static S: Signature = Signature {
+            type_signature: TypeSignature::Uniform(0, vec![]),
+            volatility: Volatility::Immutable,
+        };
+
+        &S
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        Ok(DataType::Utf8)
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
+        Ok(ColumnarValue::Array(Arc::new(
+            (0..=args.number_rows)
+                .map(|idx| Some(idx.to_string()))
+                .collect::<StringArray>(),
+        )))
+    }
+}
+
+/// Returns our evil UDFs.
+///
+/// The passed `source` is ignored.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
+    Ok(vec![Arc::new(WrongNumberOfRows)])
+}

host/src/lib.rs

@@ -797,6 +797,20 @@ impl AsyncScalarUDFImpl for WasmScalarUdf {
 
         drop(store_guard);
 
-        return_type.checked_into_root(&self.trusted_data_limits)
+        match return_type.checked_into_root(&self.trusted_data_limits) {
+            Ok(ColumnarValue::Scalar(scalar)) => Ok(ColumnarValue::Scalar(scalar)),
+            Ok(ColumnarValue::Array(array)) if array.len() as u64 != args.number_rows => {
+                Err(DataFusionError::External(
+                    format!(
+                        "UDF returned array of length {} but should produce {} rows",
+                        array.len(),
+                        args.number_rows
+                    )
+                    .into(),
+                ))
+            }
+            Ok(ColumnarValue::Array(array)) => Ok(ColumnarValue::Array(array)),
+            Err(e) => Err(e),
+        }
     }
 }

host/tests/integration_tests/evil/mod.rs

@@ -1,6 +1,7 @@
 mod env;
 mod fs;
 mod net;
+mod return_data;
 mod root;
 mod runtime;
 mod spin;

host/tests/integration_tests/evil/return_data.rs

@@ -0,0 +1,32 @@
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field};
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::{ScalarFunctionArgs, async_udf::AsyncScalarUDFImpl};
+
+use crate::integration_tests::evil::test_utils::try_scalar_udfs;
+
+#[tokio::test]
+async fn test_wrong_number_of_rows() {
+    let [udf] = try_scalar_udfs("return_data")
+        .await
+        .unwrap()
+        .try_into()
+        .unwrap();
+
+    let err = udf
+        .invoke_async_with_args(ScalarFunctionArgs {
+            args: vec![],
+            arg_fields: vec![],
+            number_rows: 42,
+            return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
+            config_options: Arc::new(ConfigOptions::default()),
+        })
+        .await
+        .unwrap_err();
+
+    insta::assert_snapshot!(
+        err,
+        @"External error: UDF returned array of length 43 but should produce 42 rows",
+    );
+}