Merge pull request #175 from influxdata/crepererum/df50

chore: DataFusion 50

Author: crepererum

Cargo.lock

@@ -17,12 +17,12 @@ license = "MIT OR Apache-2.0"
 
 [workspace.dependencies]
 anyhow = { version = "1.0.100", default-features = false }
-arrow = { version = "55.2.0", default-features = false, features = ["ipc"] }
+arrow = { version = "56.2.0", default-features = false, features = ["ipc"] }
 chrono = { version = "0.4.42", default-features = false }
-datafusion = { version = "49.0.1", default-features = false }
-datafusion-common = { version = "49.0.1", default-features = false }
-datafusion-expr = { version = "49.0.1", default-features = false }
-datafusion-sql = { version = "49.0.1", default-features = false }
+datafusion = { version = "50.3.0", default-features = false }
+datafusion-common = { version = "50.3.0", default-features = false }
+datafusion-expr = { version = "50.3.0", default-features = false }
+datafusion-sql = { version = "50.3.0", default-features = false }
 datafusion-udf-wasm-arrow2bytes = { path = "arrow2bytes", version = "0.1.0" }
 datafusion-udf-wasm-bundle = { path = "guests/bundle", version = "0.1.0" }
 datafusion-udf-wasm-evil = { path = "guests/evil", version = "0.1.0" }
@@ -36,13 +36,14 @@ insta = { version = "1.43.2", "default-features" = false }
 log = { version = "0.4.28", default-features = false }
 pyo3 = { version = "0.27.1", default-features = false, features = ["macros"] }
 regex = { version = "1", default-features = false }
-sqlparser = { version = "0.55.0", default-features = false, features = [
+sqlparser = { version = "0.58.0", default-features = false, features = [
   "std",
   "visitor"
 ] }
 tar = { version = "0.4.44", default-features = false }
 tempfile = { version = "3.23.0", default-features = false }
 tokio = { version = "1.48.0", default-features = false }
+uuid = { version = "1.18.1", default-features = false, features = ["v4"] }
 wasip2 = { version = "1" }
 wasmtime = { version = "38.0.4", default-features = false, features = [
   "async",
@@ -90,14 +91,6 @@ bare_urls = "deny"
 broken_intra_doc_links = "deny"
 private_intra_doc_links = "deny"
 
-[patch.crates-io]
-# use same DataFusion fork as InfluxDB
-# See https://github.com/influxdata/arrow-datafusion/pull/72
-datafusion = { git = "https://github.com/influxdata/arrow-datafusion.git", rev = "8347a71f62d4fef8d37548f22b93877170039357" }
-datafusion-common = { git = "https://github.com/influxdata/arrow-datafusion.git", rev = "8347a71f62d4fef8d37548f22b93877170039357" }
-datafusion-expr = { git = "https://github.com/influxdata/arrow-datafusion.git", rev = "8347a71f62d4fef8d37548f22b93877170039357" }
-datafusion-sql = { git = "https://github.com/influxdata/arrow-datafusion.git", rev = "8347a71f62d4fef8d37548f22b93877170039357" }
-
 # faster tests
 [profile.dev.package]
 # improve wasmtime compilation speed

Cargo.toml

@@ -1,5 +1,5 @@
 //! Evil payloads that try to break the sandbox when the UDF is called.
-use std::sync::Arc;
+use std::{hash::Hash, sync::Arc};
 
 use arrow::datatypes::DataType;
 use datafusion_common::{Result as DataFusionResult, ScalarValue};
@@ -49,6 +49,20 @@ impl std::fmt::Debug for SideEffect {
     }
 }
 
+impl PartialEq<Self> for SideEffect {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name
+    }
+}
+
+impl Eq for SideEffect {}
+
+impl Hash for SideEffect {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+    }
+}
+
 impl ScalarUDFImpl for SideEffect {
     fn as_any(&self) -> &dyn std::any::Any {
         self

guests/evil/src/runtime.rs

@@ -14,6 +14,7 @@ datafusion-common.workspace = true
 datafusion-expr.workspace = true
 datafusion-udf-wasm-guest.workspace = true
 pyo3.workspace = true
+uuid.workspace = true
 wasip2.workspace = true
 
 [build-dependencies]

guests/python/Cargo.toml

@@ -4,6 +4,7 @@
 //! [CPython]: https://www.python.org/
 //! [`pyo3`]: https://pyo3.rs/
 use std::any::Any;
+use std::hash::Hash;
 use std::ops::{ControlFlow, Range};
 use std::sync::{Arc, Once};
 
@@ -15,6 +16,7 @@ use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signatur
 use datafusion_udf_wasm_guest::export;
 use pyo3::prelude::*;
 use pyo3::types::PyTuple;
+use uuid::Uuid;
 
 use crate::error::py_err_to_string;
 use crate::inspect::inspect_python_code;
@@ -35,6 +37,9 @@ struct PythonScalarUDF {
     /// Handle of the wrapped python function.
     python_function: PythonFn,
 
+    /// We treat every python UDF as unique, but we need a proxy value to express that.
+    id: Uuid,
+
     /// Signature of the UDF.
     ///
     /// We store this here because [`ScalarUDFImpl::signature`] requires us to return a reference.
@@ -56,6 +61,7 @@ impl PythonScalarUDF {
 
         Self {
             python_function,
+            id: Uuid::new_v4(),
             signature,
         }
     }
@@ -94,6 +100,20 @@ impl PythonScalarUDF {
     }
 }
 
+impl PartialEq<Self> for PythonScalarUDF {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id
+    }
+}
+
+impl Eq for PythonScalarUDF {}
+
+impl Hash for PythonScalarUDF {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.id.hash(state);
+    }
+}
+
 impl ScalarUDFImpl for PythonScalarUDF {
     fn as_any(&self) -> &dyn Any {
         self
@@ -118,6 +138,7 @@ impl ScalarUDFImpl for PythonScalarUDF {
             arg_fields,
             number_rows,
             return_field,
+            config_options: _,
         } = args;
 
         let return_dt = self.python_function.signature.return_type.t.data_type();

guests/python/src/lib.rs

@@ -13,7 +13,7 @@ use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signatur
 use datafusion_udf_wasm_guest::export;
 
 /// UDF that implements "add one".
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, Hash)]
 struct AddOne {
     /// Signature of the UDF.
     ///
@@ -58,6 +58,7 @@ impl ScalarUDFImpl for AddOne {
             arg_fields: _,
             number_rows: _,
             return_field: _,
+            config_options: _,
         } = args;
 
         // extract inputs

guests/rust/examples/add_one.rs

@@ -1,9 +1,11 @@
 //! Conversion routes from/to [WIT types](crate::bindings).
+use std::sync::Arc;
+
 use arrow::{
     array::ArrayRef,
     datatypes::{DataType, Field, FieldRef},
 };
-use datafusion_common::{error::DataFusionError, scalar::ScalarValue};
+use datafusion_common::{config::ConfigOptions, error::DataFusionError, scalar::ScalarValue};
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
 use datafusion_udf_wasm_arrow2bytes::{array2bytes, bytes2array, bytes2datatype, datatype2bytes};
 
@@ -238,6 +240,9 @@ impl TryFrom<wit_types::ScalarFunctionArgs> for ScalarFunctionArgs {
                 .collect::<Result<_, _>>()?,
             number_rows: value.number_rows as usize,
             return_field: value.return_field.try_into()?,
+            config_options: Arc::new(ConfigOptions::from_string_hash_map(
+                &value.config_options.into_iter().collect(),
+            )?),
         })
     }
 }

guests/rust/src/conversion.rs

@@ -17,6 +17,7 @@ rand = { version = "0.9" }
 siphasher = { version = "1", default-features = false }
 tar.workspace = true
 tokio = { workspace = true, features = ["rt", "rt-multi-thread", "sync"] }
+uuid.workspace = true
 wasmtime.workspace = true
 wasmtime-wasi.workspace = true
 wasmtime-wasi-http.workspace = true

host/Cargo.toml

@@ -221,6 +221,16 @@ impl TryFrom<ScalarFunctionArgs> for wit_types::ScalarFunctionArgs {
                 .collect(),
             number_rows: value.number_rows as u64,
             return_field: value.return_field.as_ref().clone().into(),
+            config_options: value
+                .config_options
+                .entries()
+                .into_iter()
+                .filter_map(|e| {
+                    let k = e.key;
+                    let v = e.value?;
+                    Some((k, v))
+                })
+                .collect(),
         })
     }
 }

host/src/conversion.rs

@@ -2,16 +2,17 @@
 //!
 //!
 //! [DataFusion]: https://datafusion.apache.org/
-use std::{any::Any, collections::BTreeMap, ops::DerefMut, sync::Arc};
+use std::{any::Any, collections::BTreeMap, hash::Hash, ops::DerefMut, sync::Arc};
 
 use ::http::HeaderName;
 use arrow::datatypes::DataType;
-use datafusion_common::{DataFusionError, Result as DataFusionResult, config::ConfigOptions};
+use datafusion_common::{DataFusionError, Result as DataFusionResult};
 use datafusion_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
     async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl},
 };
 use tokio::{runtime::Handle, sync::Mutex};
+use uuid::Uuid;
 use wasmtime::{
     Engine, Store,
     component::{Component, ResourceAny},
@@ -322,6 +323,9 @@ pub struct WasmScalarUdf {
     /// [`ScalarUDFImpl::name`] is sync and requires us to return a reference.
     name: String,
 
+    /// We treat every UDF as unique, but we need a proxy value to express that.
+    id: Uuid,
+
     /// Signature of the UDF.
     ///
     /// This was pre-fetched during UDF generation because
@@ -458,6 +462,7 @@ impl WasmScalarUdf {
                 bindings: Arc::clone(&bindings),
                 resource,
                 name,
+                id: Uuid::new_v4(),
                 signature,
                 return_type,
             });
@@ -508,6 +513,7 @@ impl std::fmt::Debug for WasmScalarUdf {
             bindings: _,
             resource,
             name,
+            id,
             signature,
             return_type,
         } = self;
@@ -517,12 +523,27 @@ impl std::fmt::Debug for WasmScalarUdf {
             .field("bindings", &"<BINDINGS>")
             .field("resource", resource)
             .field("name", name)
+            .field("id", id)
             .field("signature", signature)
             .field("return_type", return_type)
             .finish()
     }
 }
 
+impl PartialEq<Self> for WasmScalarUdf {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id
+    }
+}
+
+impl Eq for WasmScalarUdf {}
+
+impl Hash for WasmScalarUdf {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.id.hash(state);
+    }
+}
+
 impl ScalarUDFImpl for WasmScalarUdf {
     fn as_any(&self) -> &dyn Any {
         self
@@ -579,8 +600,7 @@ impl AsyncScalarUDFImpl for WasmScalarUdf {
     async fn invoke_async_with_args(
         &self,
         args: ScalarFunctionArgs,
-        _option: &ConfigOptions,
-    ) -> DataFusionResult<arrow::array::ArrayRef> {
+    ) -> DataFusionResult<ColumnarValue> {
         let args = args.try_into()?;
         let mut store_guard = self.store.lock().await;
         let return_type = self
@@ -596,10 +616,6 @@ impl AsyncScalarUDFImpl for WasmScalarUdf {
 
         drop(store_guard);
 
-        let columnar_value: ColumnarValue = return_type.try_into()?;
-        match columnar_value {
-            ColumnarValue::Array(v) => Ok(v),
-            ColumnarValue::Scalar(v) => v.to_array_of_size(args.number_rows as usize),
-        }
+        return_type.try_into()
     }
 }