refactor: extend scalar-udfs WIT interface

guests/python/src/lib.rs

@@ -77,10 +77,10 @@ fn root() -> Option<Vec<u8>> {
     Some(ROOT_TAR.to_vec())
 }
 
-fn udfs() -> Vec<Arc<dyn ScalarUDFImpl>> {
+fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
     pyo3::prepare_freethreaded_python();
 
-    vec![Arc::new(Test::new())]
+    Ok(vec![Arc::new(Test::new())])
 }
 
 export! {

guests/rust/examples/add_one.rs

@@ -90,8 +90,8 @@ fn root() -> Option<Vec<u8>> {
     None
 }
 
-fn udfs() -> Vec<Arc<dyn ScalarUDFImpl>> {
-    vec![Arc::new(AddOne::new())]
+fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
+    Ok(vec![Arc::new(AddOne::new())])
 }
 
 export! {

guests/rust/src/lib.rs

@@ -19,13 +19,21 @@ macro_rules! export {
                 $root_fs_tar()
             }
 
-            fn scalar_udfs() -> Vec<$crate::bindings::exports::datafusion_udf_wasm::udf::types::ScalarUdf> {
-                $scalar_udfs()
-                    .into_iter()
+            fn scalar_udfs(
+                source: String,
+            ) -> Result<
+                Vec<$crate::bindings::exports::datafusion_udf_wasm::udf::types::ScalarUdf>,
+                $crate::bindings::exports::datafusion_udf_wasm::udf::types::DataFusionError,
+            > {
+                let udfs = $scalar_udfs(source)?;
+
+                Ok(
+                    udfs.into_iter()
                     .map(|udf| $crate::bindings::exports::datafusion_udf_wasm::udf::types::ScalarUdf::new(
                         $crate::wrapper::ScalarUdfWrapper::new(udf)
                     ))
                     .collect()
+                )
             }
         }
 

host/src/error.rs

@@ -29,3 +29,20 @@ impl<T> WasmToDataFusionResultExt for Result<T, wasmtime::Error> {
         self.map_err(|err| WasmToDataFusionErrorExt::context(err, msg, stderr))
     }
 }
+
+pub(crate) trait DataFusionResultExt {
+    type T;
+
+    fn context(self, description: impl Into<String>) -> Result<Self::T, DataFusionError>;
+}
+
+impl<T, E> DataFusionResultExt for Result<T, E>
+where
+    E: Into<DataFusionError>,
+{
+    type T = T;
+
+    fn context(self, description: impl Into<String>) -> Result<Self::T, DataFusionError> {
+        self.map_err(|e| e.into().context(description))
+    }
+}

host/src/lib.rs

@@ -15,7 +15,8 @@ use wasmtime_wasi::{
 
 use crate::{
     bindings::exports::datafusion_udf_wasm::udf::types as wit_types,
-    compilation_cache::CompilationCache, tokio_helpers::async_in_sync_context,
+    compilation_cache::CompilationCache, error::DataFusionResultExt,
+    tokio_helpers::async_in_sync_context,
 };
 use crate::{error::WasmToDataFusionResultExt, tokio_helpers::blocking_io};
 
@@ -67,7 +68,7 @@ pub struct WasmScalarUdf {
 }
 
 impl WasmScalarUdf {
-    pub async fn new(wasm_binary: &[u8]) -> DataFusionResult<Vec<Self>> {
+    pub async fn new(wasm_binary: &[u8], source: String) -> DataFusionResult<Vec<Self>> {
         // TODO: we need an in-mem file system for this, see
         //       - https://github.com/bytecodealliance/wasmtime/issues/8963
         //       - https://github.com/Timmmm/wasmtime_fs_demo
@@ -130,12 +131,13 @@ impl WasmScalarUdf {
 
         let udf_resources = bindings
             .datafusion_udf_wasm_udf_types()
-            .call_scalar_udfs(&mut store)
+            .call_scalar_udfs(&mut store, &source)
             .await
             .context(
-                "call scalar_udfs() method",
+                "calling scalar_udfs() method failed",
                 Some(&store.data().stderr.contents()),
-            )?;
+            )?
+            .context("scalar_udfs")?;
 
         let store = Arc::new(Mutex::new(store));
 

host/tests/integration_tests/python.rs

@@ -14,7 +14,7 @@ async fn test() {
     .await
     .unwrap();
 
-    let mut udfs = WasmScalarUdf::new(&data).await.unwrap();
+    let mut udfs = WasmScalarUdf::new(&data, "".to_owned()).await.unwrap();
     assert_eq!(udfs.len(), 1);
     let udf = udfs.pop().unwrap();
 

host/tests/integration_tests/rust.rs

@@ -17,7 +17,7 @@ async fn test_add_one() {
     .await
     .unwrap();
 
-    let mut udfs = WasmScalarUdf::new(&data).await.unwrap();
+    let mut udfs = WasmScalarUdf::new(&data, "".to_owned()).await.unwrap();
     assert_eq!(udfs.len(), 1);
     let udf = udfs.pop().unwrap();
 

wit/world.wit

@@ -87,7 +87,7 @@ interface types {
     }
 
     root-fs-tar: func() -> option<list<u8>>;
-    scalar-udfs: func() -> list<scalar-udf>;
+    scalar-udfs: func(source: string) -> result<list<scalar-udf>, data-fusion-error>;
 }
 
 world datafusion {