Merge pull request #161 from influxdata/tm/return-type-async

feat: cache return type if function signature is exact

Author: Sl1mb0

host/src/lib.rs

@@ -8,7 +8,7 @@ use ::http::HeaderName;
 use arrow::datatypes::DataType;
 use datafusion_common::{DataFusionError, Result as DataFusionResult, config::ConfigOptions};
 use datafusion_expr::{
-    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
     async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl},
 };
 use tokio::{runtime::Handle, sync::Mutex};
@@ -313,21 +313,31 @@ pub struct WasmScalarUdf {
 
     /// Name of the UDF.
     ///
-    /// This was pre-fetched during UDF generation because [`ScalarUDFImpl::name`] is sync and requires us to return a
-    /// reference.
+    /// This was pre-fetched during UDF generation because
+    /// [`ScalarUDFImpl::name`] is sync and requires us to return a reference.
     name: String,
 
     /// Signature of the UDF.
     ///
-    /// This was pre-fetched during UDF generation because [`ScalarUDFImpl::signature`] is sync and requires us to return a
+    /// This was pre-fetched during UDF generation because
+    /// [`ScalarUDFImpl::signature`] is sync and requires us to return a
     /// reference.
     signature: Signature,
+
+    /// Return type of the UDF.
+    ///
+    /// This was pre-fetched during UDF generation because
+    /// [`ScalarUDFImpl::return_type`] is sync and requires us to return a
+    /// reference. We can only compute the return type if the underlying
+    /// [TypeSignature] is [Exact](TypeSignature::Exact).
+    return_type: Option<DataType>,
 }
 
 impl WasmScalarUdf {
     /// Create multiple UDFs from a single WASM VM.
     ///
-    /// UDFs bound to the same VM share state, however calling this method multiple times will yield independent WASM VMs.
+    /// UDFs bound to the same VM share state, however calling this method
+    /// multiple times will yield independent WASM VMs.
     pub async fn new(
         component: &WasmComponentPrecompiled,
         permissions: &WasmPermissions,
@@ -404,7 +414,7 @@ impl WasmScalarUdf {
                 )?;
 
             let store2: &mut Store<WasmStateImpl> = &mut store_guard;
-            let signature = bindings
+            let signature: Signature = bindings
                 .datafusion_udf_wasm_udf_types()
                 .scalar_udf()
                 .call_signature(store2, resource)
@@ -415,12 +425,36 @@ impl WasmScalarUdf {
                 )?
                 .try_into()?;
 
+            let return_type = match &signature.type_signature {
+                TypeSignature::Exact(t) => {
+                    let store2: &mut Store<WasmStateImpl> = &mut store_guard;
+                    let r = bindings
+                        .datafusion_udf_wasm_udf_types()
+                        .scalar_udf()
+                        .call_return_type(
+                            store2,
+                            resource,
+                            &t.iter()
+                                .map(|dt| wit_types::DataType::from(dt.clone()))
+                                .collect::<Vec<_>>(),
+                        )
+                        .await
+                        .context(
+                            "call ScalarUdf::return_type",
+                            Some(&store_guard.data().stderr.contents()),
+                        )??;
+                    Some(r.try_into()?)
+                }
+                _ => None,
+            };
+
             udfs.push(Self {
                 store: Arc::clone(&store),
                 bindings: Arc::clone(&bindings),
                 resource,
                 name,
                 signature,
+                return_type,
             });
         }
 
@@ -431,6 +465,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 {
@@ -441,6 +504,7 @@ impl std::fmt::Debug for WasmScalarUdf {
             resource,
             name,
             signature,
+            return_type,
         } = self;
 
         f.debug_struct("WasmScalarUdf")
@@ -449,6 +513,7 @@ impl std::fmt::Debug for WasmScalarUdf {
             .field("resource", resource)
             .field("name", name)
             .field("signature", signature)
+            .field("return_type", return_type)
             .finish()
     }
 }
@@ -467,6 +532,12 @@ 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());
+        }
+
         async_in_sync_context(async {
             let arg_types = arg_types
                 .iter()

host/tests/integration_tests/python/argument_forms.rs

@@ -15,7 +15,7 @@ use datafusion_expr::{
 
 use crate::integration_tests::python::test_utils::python_scalar_udf;
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_positional_or_keyword() {
     const CODE: &str = "
 def foo(x: int) -> int:
@@ -55,7 +55,7 @@ def foo(x: int) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_positional_or_keyword_default() {
     const CODE: &str = "
 def foo(x: int = 1) -> int:
@@ -77,7 +77,7 @@ def foo(x: int = 1) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_positional_only() {
     const CODE: &str = "
 def foo(x: int, /) -> int:
@@ -117,7 +117,7 @@ def foo(x: int, /) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_positional_only_default() {
     const CODE: &str = "
 def foo(x: int = 1, /) -> int:
@@ -139,7 +139,7 @@ def foo(x: int = 1, /) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_positional_or_keyword_and_positional_only() {
     const CODE: &str = "
 def foo(x: int, /, y: int) -> int:
@@ -182,7 +182,7 @@ def foo(x: int, /, y: int) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_var_positional() {
     const CODE: &str = "
 def foo(*x: int) -> int:
@@ -204,7 +204,7 @@ def foo(*x: int) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_keyword_only() {
     const CODE: &str = "
 def foo(*, x: int) -> int:
@@ -226,7 +226,7 @@ def foo(*, x: int) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_var_keyword() {
     const CODE: &str = "
 def foo(**x: int) -> int:

host/tests/integration_tests/python/examples.rs

@@ -13,7 +13,7 @@ use datafusion_expr::{
 
 use crate::integration_tests::python::test_utils::python_scalar_udf;
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_add_one() {
     const CODE: &str = "
 def add_one(x: int) -> int:

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

@@ -1,7 +1,7 @@
 use crate::integration_tests::python::test_utils::python_scalar_udfs;
 use datafusion_common::DataFusionError;
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_invalid_syntax() {
     const CODE: &str = ")";
 
@@ -18,7 +18,7 @@ async fn test_invalid_syntax() {
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_missing_return_type() {
     const CODE: &str = "
 def add_one(x: int):
@@ -43,7 +43,7 @@ def add_one(x: int):
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_missing_arg_type() {
     const CODE: &str = "
 def add_one(x) -> int:
@@ -68,7 +68,7 @@ def add_one(x) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_union_type_2() {
     const CODE: &str = "
 def add_one(x: int | str) -> int:
@@ -93,7 +93,7 @@ def add_one(x: int | str) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_union_type_2_and_none() {
     const CODE: &str = "
 def add_one(x: int | str | None) -> int:
@@ -118,7 +118,7 @@ def add_one(x: int | str | None) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_union_type_2_identical() {
     const CODE: &str = "
 def add_one(x: int | str | int) -> int:
@@ -143,7 +143,7 @@ def add_one(x: int | str | int) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_union_type_2_identical_and_none() {
     const CODE: &str = "
 def add_one(x: int | None | str | int) -> int:
@@ -168,7 +168,7 @@ def add_one(x: int | None | str | int) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_union_type_3() {
     const CODE: &str = "
 def add_one(x: int | str | float) -> int:
@@ -193,7 +193,7 @@ def add_one(x: int | str | float) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_type_annotation_is_not_a_type() {
     const CODE: &str = "
 def add_one(x: 1337) -> int:
@@ -218,7 +218,7 @@ def add_one(x: 1337) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_unsupported_type() {
     const CODE: &str = "
 def add_one(x: list[int]) -> int:
@@ -243,7 +243,7 @@ def add_one(x: list[int]) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_custom_type() {
     const CODE: &str = "
 class C:
@@ -271,7 +271,7 @@ def add_one(x: C) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_exception() {
     const CODE: &str = "
 raise Exception('foo')

host/tests/integration_tests/python/inspection/filter.rs

@@ -1,7 +1,7 @@
 use crate::integration_tests::python::test_utils::python_scalar_udfs;
 use datafusion_expr::ScalarUDFImpl;
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_underscore() {
     const CODE: &str = "
 def foo(x: int) -> int:
@@ -14,7 +14,7 @@ def _bar(x: int) -> int:
     assert_eq!(found_udfs(CODE).await, ["foo".to_owned()]);
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_non_callalbes() {
     const CODE: &str = "
 variable = 1
@@ -26,7 +26,7 @@ def foo(x: int) -> int:
     assert_eq!(found_udfs(CODE).await, ["foo".to_owned()]);
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_imports() {
     const CODE: &str = "
 from sys import exit
@@ -38,7 +38,7 @@ def foo(x: int) -> int:
     assert_eq!(found_udfs(CODE).await, ["foo".to_owned()]);
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_classes() {
     const CODE: &str = "
 class C:

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

@@ -9,7 +9,7 @@ use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, async_udf::AsyncScalarU
 
 use crate::integration_tests::python::test_utils::python_scalar_udf;
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn call_other_function() {
     const CODE: &str = "
 def _sub1(x: int) -> int:
@@ -45,7 +45,7 @@ def foo(x: int) -> int:
     );
 }
 
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn functools_cache() {
     const CODE: &str = "
 from functools import cache