fix: validate arg_types when requesting return_type

Author: Sl1mb0

host/src/lib.rs

@@ -438,6 +438,35 @@ impl WasmScalarUdf {
     pub fn as_async_udf(self) -> AsyncScalarUDF {
         AsyncScalarUDF::new(Arc::new(self))
     }
+
+    /// Check that the provided argument types match the UDF signature.
+    fn check_arg_types(&self, arg_types: &[DataType]) -> DataFusionResult<()> {
+        if let TypeSignature::Exact(expected_types) = &self.signature.type_signature {
+            if arg_types.len() != expected_types.len() {
+                return Err(DataFusionError::Plan(format!(
+                    "`{}` expects {} parameters but got {}",
+                    self.name,
+                    expected_types.len(),
+                    arg_types.len()
+                )));
+            }
+
+            for (i, (provided, expected)) in arg_types.iter().zip(expected_types.iter()).enumerate()
+            {
+                if provided != expected {
+                    return Err(DataFusionError::Plan(format!(
+                        "argument {} of `{}` should be {:?}, got {:?}",
+                        i + 1,
+                        self.name,
+                        expected,
+                        provided
+                    )));
+                }
+            }
+        };
+
+        Ok(())
+    }
 }
 
 impl std::fmt::Debug for WasmScalarUdf {
@@ -476,6 +505,8 @@ impl ScalarUDFImpl for WasmScalarUdf {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        self.check_arg_types(arg_types)?;
+
         if let Some(return_type) = &self.return_type {
             return Ok(return_type.clone());
         }

host/tests/integration_tests/python/examples.rs

@@ -32,9 +32,10 @@ def add_one(x: int) -> int:
         udf.return_type(&[DataType::Int64]).unwrap(),
         DataType::Int64,
     );
-    // Signature is exact, so it would have been cached;
-    // thus - no error even with a wrong type of argument.
-    assert_eq!(udf.return_type(&[]).unwrap(), DataType::Int64,);
+    assert_eq!(
+        udf.return_type(&[]).unwrap_err().to_string(),
+        "Error during planning: `add_one` expects 1 parameters but got 0",
+    );
 
     // call with array
     let array = udf

host/tests/integration_tests/python/runtime/errors.rs

@@ -21,11 +21,19 @@ def foo(x: int) -> int:
 
     let udf = python_scalar_udf(CODE).await.unwrap();
 
-    // Signature is exact, so it would have been cached;
-    // thus - no error even with a wrong type of argument.
     insta::assert_snapshot!(
-        udf.return_type(&[]).unwrap(),
-        @"Int64",
+        udf.return_type(&[]).unwrap_err(),
+        @"Error during planning: `foo` expects 1 parameters but got 0",
+    );
+
+    insta::assert_snapshot!(udf.return_type(
+        &[DataType::Int64, DataType::Int64]).unwrap_err(),
+        @"Error during planning: `foo` expects 1 parameters but got 2",
+    );
+
+    insta::assert_snapshot!(
+        udf.return_type(&[DataType::Float64]).unwrap_err(),
+        @"Error during planning: argument 1 of `foo` should be Int64, got Float64",
     );
 }