feat: create UdfCodeFormatter trait

This commit does a few things. First, it splits off  into it's own  module. It also shuffles around some test utils as part of that. Most importantly though, it defines and implements a UdfCodeFormatter trait for formatting various programming languages when parsing & registering UDFs.

Author: Sl1mb0

Cargo.lock

@@ -6,6 +6,7 @@ members = [
   "guests/python",
   "guests/rust",
   "host",
+  "query",
 ]
 
 [workspace.package]
@@ -25,6 +26,7 @@ datafusion-udf-wasm-arrow2bytes = { path = "arrow2bytes", version = "0.1.0" }
 datafusion-udf-wasm-bundle = { path = "guests/bundle", version = "0.1.0" }
 datafusion-udf-wasm-guest = { path = "guests/rust", version = "0.1.0" }
 datafusion-udf-wasm-python = { path = "guests/python", version = "0.1.0" }
+datafusion-udf-wasm-query = { path = "query", version = "0.1.0" }
 http = { version = "1.3.1", default-features = false }
 hyper = { version = "1.7", default-features = false }
 pyo3 = { version = "0.27.1", default-features = false, features = ["macros"] }

Cargo.toml

@@ -7,16 +7,17 @@ license.workspace = true
 [dependencies]
 anyhow.workspace = true
 arrow.workspace = true
-datafusion.workspace = true
 datafusion-common.workspace = true
 datafusion-expr.workspace = true
-datafusion-sql.workspace = true
 datafusion-udf-wasm-arrow2bytes.workspace = true
+datafusion-udf-wasm-bundle = { workspace = true, features = [
+  "example",
+  "python"
+] }
 http.workspace = true
 hyper.workspace = true
 rand = { version = "0.9" }
 siphasher = { version = "1", default-features = false }
-sqlparser.workspace = true
 tar.workspace = true
 tokio = { workspace = true, features = ["rt", "rt-multi-thread", "sync"] }
 wasmtime.workspace = true
@@ -25,10 +26,6 @@ wasmtime-wasi-http.workspace = true
 wasmtime-wasi-io.workspace = true
 
 [dev-dependencies]
-datafusion-udf-wasm-bundle = { workspace = true, features = [
-  "example",
-  "python"
-] }
 insta = "1.43.2"
 tokio = { workspace = true, features = ["fs", "macros"] }
 wiremock = "0.6.5"

host/Cargo.toml

@@ -42,8 +42,9 @@ mod conversion;
 mod error;
 pub mod http;
 mod linker;
+/// Utilities for testing.
+pub mod test_utils;
 mod tokio_helpers;
-pub mod udf_query;
 mod vfs;
 
 /// State of the WASM payload.

host/src/lib.rs

@@ -0,0 +1,2 @@
+/// Utilities for testing.
+pub mod python;

host/src/test_utils/mod.rs

@@ -1,10 +1,12 @@
+use crate::{WasmComponentPrecompiled, WasmScalarUdf};
 use datafusion_common::DataFusionError;
-use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmScalarUdf};
 use tokio::sync::OnceCell;
 
+/// Static precompiled Python WASM component for tests
 static COMPONENT: OnceCell<WasmComponentPrecompiled> = OnceCell::const_new();
 
-pub(crate) async fn python_component() -> &'static WasmComponentPrecompiled {
+/// Returns a static reference to the precompiled Python WASM component.
+pub async fn python_component() -> &'static WasmComponentPrecompiled {
     COMPONENT
         .get_or_init(async || {
             WasmComponentPrecompiled::new(datafusion_udf_wasm_bundle::BIN_PYTHON.into())
@@ -14,13 +16,15 @@ pub(crate) async fn python_component() -> &'static WasmComponentPrecompiled {
         .await
 }
 
-pub(crate) async fn python_scalar_udfs(code: &str) -> Result<Vec<WasmScalarUdf>, DataFusionError> {
+/// Compiles the provided Python UDF code into a list of WasmScalarUdf instances.
+pub async fn python_scalar_udfs(code: &str) -> Result<Vec<WasmScalarUdf>, DataFusionError> {
     let component = python_component().await;
 
     WasmScalarUdf::new(component, &Default::default(), code.to_owned()).await
 }
 
-pub(crate) async fn python_scalar_udf(code: &str) -> Result<WasmScalarUdf, DataFusionError> {
+/// Compiles the provided Python UDF code into a single WasmScalarUdf instance.
+pub async fn python_scalar_udf(code: &str) -> Result<WasmScalarUdf, DataFusionError> {
     let udfs = python_scalar_udfs(code).await?;
     assert_eq!(udfs.len(), 1);
     Ok(udfs.into_iter().next().expect("just checked len"))

host/tests/integration_tests/python/test_utils.rs renamed to host/src/test_utils/python.rs

@@ -1,4 +1,3 @@
 mod python;
 mod rust;
 mod test_utils;
-mod udf_query;

host/tests/integration_tests/mod.rs

@@ -8,10 +8,9 @@ use arrow::{
     datatypes::{DataType, Field},
 };
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+use datafusion_udf_wasm_host::test_utils::python::python_scalar_udf;
 
-use crate::integration_tests::{
-    python::test_utils::python_scalar_udf, test_utils::ColumnarValueExt,
-};
+use crate::integration_tests::test_utils::ColumnarValueExt;
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_positional_or_keyword() {

host/tests/integration_tests/python/argument_forms.rs

@@ -6,10 +6,9 @@ use arrow::{
 };
 use datafusion_common::ScalarValue;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+use datafusion_udf_wasm_host::test_utils::python::python_scalar_udf;
 
-use crate::integration_tests::{
-    python::test_utils::python_scalar_udf, test_utils::ColumnarValueExt,
-};
+use crate::integration_tests::test_utils::ColumnarValueExt;
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_add_one() {

host/tests/integration_tests/python/examples.rs

@@ -1,6 +1,5 @@
 use datafusion_common::DataFusionError;
-
-use crate::integration_tests::python::test_utils::python_scalar_udfs;
+use datafusion_udf_wasm_host::test_utils::python::python_scalar_udfs;
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_invalid_syntax() {