Start setting up wasm

crates/wasm-functions/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "wasm-functions"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+arrow = "51.0.0"
+wasm-udfs = { version = "0.1.0", path = "../wasm-udfs" }
+
+[dev-dependencies]
+wasm-bindgen-test = "0.3.43"

crates/wasm-functions/src/lib.rs

@@ -0,0 +1,84 @@
+use arrow::array::{Array, ArrayRef, Float64Array};
+use arrow::error::ArrowError;
+use std::sync::Arc;
+use wasm_udfs::*;
+
+// ```bash
+// cargo install wasm-bindgen-cli
+// ```
+
+// ```bash
+// cargo test --target wasm32-unknown-unknown
+// ```
+
+// expose function f1 as external function
+// add required bindgen, and required serialization/deserialization
+export_udf_function!(f1);
+// function should return error
+export_udf_function!(f_return_error);
+// function should panic
+// export_udf_function!(f_panic);
+// function should return arrow error
+export_udf_function!(f_return_arrow_error);
+
+/// standard datafusion udf ... kind of
+/// should return ArrayRef or ArrowError
+fn f1(args: &[ArrayRef]) -> Result<ArrayRef, ArrowError> {
+    assert_eq!(2, args.len());
+
+    let base = args[0]
+        .as_any()
+        .downcast_ref::<Float64Array>()
+        .expect("cast 0 failed");
+    let exponent = args[1]
+        .as_any()
+        .downcast_ref::<Float64Array>()
+        .expect("cast 1 failed");
+
+    assert_eq!(exponent.len(), base.len());
+
+    let array = base
+        .iter()
+        .zip(exponent.iter())
+        .map(|(base, exponent)| match (base, exponent) {
+            (Some(base), Some(exponent)) => Some(base.powf(exponent)),
+            _ => None,
+        })
+        .collect::<Float64Array>();
+
+    // TODO: do we need arc here?
+    //       only reason to stay to keep api same
+    //       like datafusion udf's
+    Ok(Arc::new(array))
+}
+/// function returns String Error
+fn f_return_error(_args: &[ArrayRef]) -> Result<ArrayRef, String> {
+    Err("wasm function returned error".to_string())
+}
+
+/// function returns error
+fn f_return_arrow_error(_args: &[ArrayRef]) -> Result<ArrayRef, ArrowError> {
+    Err(ArrowError::DivideByZero)
+}
+
+// fn f_panic(_args: &[ArrayRef]) -> Result<ArrayRef, String> {
+//     panic!("wasm function panicked")
+// }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{ArrayRef, Float64Array};
+
+    use std::sync::Arc;
+
+    #[wasm_bindgen_test::wasm_bindgen_test]
+    fn test_f1() {
+        let a: ArrayRef = Arc::new(Float64Array::from(vec![2.1, 3.1, 4.1, 5.1]));
+        let b: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]));
+        let args = vec![a, b];
+        let result = f1(&args).unwrap();
+
+        assert_eq!(4, result.len())
+    }
+}

crates/wasm-udfs/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "wasm-udfs"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+arrow = "51.0.0"
+paste = "1.0.15"
+wasmedge-bindgen = "0.4.1"
+wasmedge-bindgen-macro = "0.4.1"

crates/wasm-udfs/src/lib.rs

@@ -0,0 +1,67 @@
+use arrow::{
+    array::{Array, ArrayRef, RecordBatch},
+    datatypes::{Field, Schema, SchemaRef},
+};
+pub use paste;
+use std::sync::Arc;
+pub use wasmedge_bindgen;
+pub use wasmedge_bindgen_macro;
+
+/// packs slice of arrays to a batch
+/// with schema generated from array types
+pub fn pack_array(args: &[ArrayRef]) -> RecordBatch {
+    let fields = args
+        .iter()
+        .enumerate()
+        .map(|(i, f)| Field::new(format!("c{}", i), f.data_type().clone(), false))
+        .collect::<Vec<_>>();
+
+    let schema = Arc::new(Schema::new(fields));
+
+    RecordBatch::try_new(schema, args.to_vec()).unwrap()
+}
+
+/// packs slice of arrays to a batch
+/// with external schema
+pub fn pack_array_with_schema(args: &[ArrayRef], schema: SchemaRef) -> RecordBatch {
+    RecordBatch::try_new(schema, args.to_vec()).unwrap()
+}
+
+/// creates a arrow ipc blob
+pub fn to_ipc(schema: &Schema, batch: RecordBatch) -> Vec<u8> {
+    let blob = vec![];
+    let mut stream_writer = arrow::ipc::writer::StreamWriter::try_new(blob, schema).unwrap();
+    stream_writer.write(&batch).unwrap();
+
+    stream_writer.into_inner().unwrap()
+}
+
+/// creates arrow arrays from arrow ipc blob
+pub fn from_ipc(payload: &[u8]) -> RecordBatch {
+    let mut batch = arrow::ipc::reader::StreamReader::try_new(payload, None).unwrap();
+    batch.next().unwrap().unwrap()
+}
+
+/// exports wasm function and performs all required
+/// arrow ipc serialization/deserialization
+///
+/// macro will create new function prefixed with `__wasm_udf_`
+///
+// TODO: make this a proc macro maybe ?
+#[macro_export]
+macro_rules! export_udf_function {
+    ($name:ident) => {
+        paste::item! {
+            #[wasmedge_bindgen_macro::wasmedge_bindgen]
+            pub fn [<__wasm_udf_$name>](payload: Vec<u8>) -> Result<Vec<u8>,String> {
+                let args_batch = from_ipc(&payload);
+                let result = $name(args_batch.columns());
+                // let batch = pack_array(&vec![result]);
+                // to_ipc(&batch.schema(), batch)
+                result.map(|result| pack_array(&vec![result]))
+                    .map(|batch| to_ipc(&batch.schema(), batch))
+                    .map_err(|e| e.to_string())
+            }
+        }
+    };
+}

crates/wasmedge-factory/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "wasmedge-factory"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+arrow-udf-wasm = "0.2.2"
+async-trait = "0.1.82"
+datafusion = "38.0.0"
+log = "0.4.22"
+project-root = "0.2.2"
+thiserror = "1.0.63"
+tokio = "1.40.0"
+wasm-udfs = { version = "0.1.0", path = "../wasm-udfs" }
+# wasmedge-sdk = "0.13.2"
+weak-table = "0.3.2"