paleolimbot commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3712337503
########## native/comet-udf-sdk/src/c_abi.rs: ########## @@ -0,0 +1,835 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The Comet UDF C ABI — sedona-style. +//! +//! The wire format is two `#[repr(C)]` structs of function pointers, +//! parameterized only by Arrow's C Data Interface +//! (`FFI_ArrowSchema` / `FFI_ArrowArray`). No DataFusion types appear in +//! the FFI surface, so the user's cdylib only needs a matching `arrow` +//! crate, not a matching `datafusion` version. Review Comment: You have this in the document above, but perhaps worth reiterating here that these are comet version specific (i.e., these are not yet ABI stable between comet versions). I believe that neither the arrow nor DataFusion version have to match here (but one may have to relax the arrow version compatibility in Cargo.toml for Rust dependents of this crate to be able to compile it). ########## native/comet-udf-sdk/Cargo.toml: ########## @@ -0,0 +1,32 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "comet-udf-sdk" +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +license = { workspace = true } +description = "SDK for writing custom Rust UDFs that run inside Apache DataFusion Comet (arrow-ffi based)" + +publish = false + +[dependencies] +arrow = { workspace = true } Review Comment: I am not sure if this will wreak havoc the build, but it may be possible to relax the version compatibility here to whenever the FFI_ArrowArray/Schema were added. ########## native/core/src/execution/rust_udf/cache.rs: ########## @@ -0,0 +1,87 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Process-wide cache of loaded UDF cdylibs. +//! +//! Same-path lookups always return the same `Arc<LoadedLibrary>` for +//! the lifetime of the process — libraries are deliberately never +//! unloaded. Calling `dlclose` while a thread is mid-call would be a +//! use-after-free, and there is no safe point to unload without +//! per-invocation refcounting we don't want on the hot path. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, OnceLock, RwLock}; + +use super::loader::{load, LoadedLibrary, LoaderError}; + +static CACHE: OnceLock<RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>>> = OnceLock::new(); + +fn cache() -> &'static RwLock<HashMap<PathBuf, Arc<LoadedLibrary>>> { + CACHE.get_or_init(|| RwLock::new(HashMap::new())) +} + +/// Get an already-loaded library, or load and cache it. +pub fn get_or_load(path: impl AsRef<Path>) -> Result<Arc<LoadedLibrary>, LoaderError> { + let raw = path.as_ref().to_path_buf(); Review Comment: In the documentation you say these have to be absolute paths, but I think relative paths to the working directory would make it through here. I think that's fine, just flagging if the absolute path consideration was a security choice. Relative to LD_LIBRARY_PATH would also be fine (maybe more portable too). ########## native/comet-udf-sdk/src/c_abi.rs: ########## @@ -0,0 +1,835 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The Comet UDF C ABI — sedona-style. +//! +//! The wire format is two `#[repr(C)]` structs of function pointers, +//! parameterized only by Arrow's C Data Interface +//! (`FFI_ArrowSchema` / `FFI_ArrowArray`). No DataFusion types appear in +//! the FFI surface, so the user's cdylib only needs a matching `arrow` +//! crate, not a matching `datafusion` version. +//! +//! # Authoring a UDF +//! +//! Implement [`CometCScalarUdf`] for a type that also implements `Default`, +//! then use the [`comet_c_udf_export!`] macro to emit the discovery entry +//! point: +//! +//! ```ignore +//! use comet_udf_sdk::c_abi::*; +//! use arrow::array::{ArrayRef, Int64Array}; +//! use arrow::datatypes::{DataType, Field}; +//! use std::sync::Arc; +//! +//! #[derive(Default)] +//! pub struct AddOne; +//! impl CometCScalarUdf for AddOne { +//! fn name(&self) -> &str { "add_one_c" } +//! fn return_field(&self, args: &[Field]) -> Result<Field, String> { +//! if args.len() != 1 || args[0].data_type() != &DataType::Int64 { +//! return Err("expected (Int64) -> Int64".into()); +//! } +//! Ok(Field::new("add_one_c", DataType::Int64, true)) +//! } +//! fn invoke(&self, args: &[ArrayRef], _n: usize) -> Result<ArrayRef, String> { +//! let a = args[0].as_any().downcast_ref::<Int64Array>().unwrap(); +//! Ok(Arc::new(a.iter().map(|v| v.map(|x| x + 1)).collect::<Int64Array>())) +//! } +//! } +//! +//! comet_udf_sdk::comet_c_udf_export!(AddOne); +//! ``` + +use std::ffi::{c_char, c_int, c_void}; + +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; + +/// Generic non-zero error code returned by `init` / `execute` to signal +/// failure. The host treats any non-zero return as an error and calls +/// `get_last_error` for the message; the specific code is informational. +const C_ABI_ERR: c_int = 1; + +// -- panic containment ----------------------------------------------------- +// +// Every `extern "C"` function in this module is an unwind boundary. A panic +// that escapes one aborts the whole process (Rust's default `extern "C"` +// unwind behavior since 1.81), which for Comet means killing the executor +// JVM and losing every task on it -- not just the query that used the UDF. +// +// User UDF code is arbitrary and panicking is idiomatic Rust (`unwrap`, +// slice indexing, integer overflow in debug), so the SDK treats a panic in +// user code as an ordinary error: catch it at the boundary, convert it to a +// message, and report it through the same `get_last_error` channel as a +// returned `Err`. The query fails; the executor survives. + +/// Render a caught panic payload as an error message. +fn panic_message(panic: Box<dyn std::any::Any + Send>) -> String { + let detail = panic + .downcast_ref::<&'static str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::<String>().cloned()) + .unwrap_or_else(|| "<non-string panic payload>".to_string()); + format!("panic in UDF code: {detail}") +} + +/// Run `f`, converting a panic into `Err(message)`. +fn catch_panic<T>(f: impl FnOnce() -> Result<T, String>) -> Result<T, String> { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(result) => result, + Err(panic) => Err(panic_message(panic)), + } +} + +/// Run an infallible `f` (typically a release/cleanup callback), swallowing +/// any panic. Used where the ABI gives us no way to report an error and +/// aborting would be a worse outcome than leaking. +fn catch_panic_infallible(f: impl FnOnce()) { + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); +} Review Comment: Can this type of error be logged in some way or is that not safe/possible here? ########## native/comet-udf-sdk/src/c_abi.rs: ########## @@ -0,0 +1,835 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The Comet UDF C ABI — sedona-style. +//! +//! The wire format is two `#[repr(C)]` structs of function pointers, +//! parameterized only by Arrow's C Data Interface +//! (`FFI_ArrowSchema` / `FFI_ArrowArray`). No DataFusion types appear in +//! the FFI surface, so the user's cdylib only needs a matching `arrow` +//! crate, not a matching `datafusion` version. +//! +//! # Authoring a UDF +//! +//! Implement [`CometCScalarUdf`] for a type that also implements `Default`, +//! then use the [`comet_c_udf_export!`] macro to emit the discovery entry +//! point: +//! +//! ```ignore +//! use comet_udf_sdk::c_abi::*; +//! use arrow::array::{ArrayRef, Int64Array}; +//! use arrow::datatypes::{DataType, Field}; +//! use std::sync::Arc; +//! +//! #[derive(Default)] +//! pub struct AddOne; +//! impl CometCScalarUdf for AddOne { +//! fn name(&self) -> &str { "add_one_c" } +//! fn return_field(&self, args: &[Field]) -> Result<Field, String> { +//! if args.len() != 1 || args[0].data_type() != &DataType::Int64 { +//! return Err("expected (Int64) -> Int64".into()); +//! } +//! Ok(Field::new("add_one_c", DataType::Int64, true)) +//! } +//! fn invoke(&self, args: &[ArrayRef], _n: usize) -> Result<ArrayRef, String> { +//! let a = args[0].as_any().downcast_ref::<Int64Array>().unwrap(); +//! Ok(Arc::new(a.iter().map(|v| v.map(|x| x + 1)).collect::<Int64Array>())) +//! } +//! } +//! +//! comet_udf_sdk::comet_c_udf_export!(AddOne); +//! ``` + +use std::ffi::{c_char, c_int, c_void}; + +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; + +/// Generic non-zero error code returned by `init` / `execute` to signal +/// failure. The host treats any non-zero return as an error and calls +/// `get_last_error` for the message; the specific code is informational. +const C_ABI_ERR: c_int = 1; + +// -- panic containment ----------------------------------------------------- +// +// Every `extern "C"` function in this module is an unwind boundary. A panic +// that escapes one aborts the whole process (Rust's default `extern "C"` +// unwind behavior since 1.81), which for Comet means killing the executor +// JVM and losing every task on it -- not just the query that used the UDF. +// +// User UDF code is arbitrary and panicking is idiomatic Rust (`unwrap`, +// slice indexing, integer overflow in debug), so the SDK treats a panic in +// user code as an ordinary error: catch it at the boundary, convert it to a +// message, and report it through the same `get_last_error` channel as a +// returned `Err`. The query fails; the executor survives. + +/// Render a caught panic payload as an error message. +fn panic_message(panic: Box<dyn std::any::Any + Send>) -> String { + let detail = panic + .downcast_ref::<&'static str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::<String>().cloned()) + .unwrap_or_else(|| "<non-string panic payload>".to_string()); + format!("panic in UDF code: {detail}") +} + +/// Run `f`, converting a panic into `Err(message)`. +fn catch_panic<T>(f: impl FnOnce() -> Result<T, String>) -> Result<T, String> { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(result) => result, + Err(panic) => Err(panic_message(panic)), + } +} + +/// Run an infallible `f` (typically a release/cleanup callback), swallowing +/// any panic. Used where the ABI gives us no way to report an error and +/// aborting would be a worse outcome than leaking. +fn catch_panic_infallible(f: impl FnOnce()) { + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); +} + +// -- factory struct -------------------------------------------------------- + +/// Factory for [`CometCScalarKernelImpl`] instances. +/// +/// Lives in a registry, may be cloned across an FFI boundary. Calls to +/// `function_name` and `new_impl` must be thread-safe (the implementation +/// is responsible for any internal synchronization). +/// +/// `#[repr(C)]` layout, matched by the host loader. Adding new fields +/// requires bumping `COMET_UDF_ABI_VERSION`. +#[repr(C)] +pub struct CometCScalarKernel { + /// Return the function name this kernel implements as a NUL-terminated + /// UTF-8 C string. The pointer must remain valid for the lifetime of + /// the [`CometCScalarKernel`]. + /// + /// May be `None`, in which case the kernel is treated as anonymous and + /// won't be discoverable by name. (Comet always sets this; field is + /// optional for parity with sedona's design.) + pub function_name: Option<unsafe extern "C" fn(*const CometCScalarKernel) -> *const c_char>, + + /// Initialize a new [`CometCScalarKernelImpl`] into `out`. Called once + /// per execution, on the thread that will then drive `init`/`execute`. + pub new_impl: + Option<unsafe extern "C" fn(*const CometCScalarKernel, out: *mut CometCScalarKernelImpl)>, + + /// Release this kernel. After release, all callbacks must be set to + /// `None`. Called when the host's `LoadedLibrary` is dropped. + pub release: Option<unsafe extern "C" fn(*mut CometCScalarKernel)>, + + /// Implementation-private data, opaque to the host. + pub private_data: *mut c_void, +} + +// SAFETY: `CometCScalarKernel` is a thin wrapper around C function +// pointers with caller-defined synchronization semantics; the trait impls +// are required so loaded kernels can be referenced from multi-threaded +// host code. Implementations of the FFI must respect thread safety as +// described in the doc comments. +unsafe impl Send for CometCScalarKernel {} +unsafe impl Sync for CometCScalarKernel {} + +impl Default for CometCScalarKernel { + fn default() -> Self { + Self { + function_name: None, + new_impl: None, + release: None, + private_data: std::ptr::null_mut(), + } + } +} + +impl Drop for CometCScalarKernel { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: release is the FFI-defined cleanup callback; + // implementations must reset `release` to None per the contract. + unsafe { release(self) }; + } + } +} + +// -- per-execution instance struct ---------------------------------------- + +/// Per-execution instance produced by [`CometCScalarKernel::new_impl`]. +/// +/// Not thread-safe; the caller must serialize access. Typically used on +/// one thread for one batch then dropped. +#[repr(C)] +pub struct CometCScalarKernelImpl { + /// Compute the return type from arg types and (optionally) bound + /// scalar arguments. + /// + /// On success, `out` is populated with the return type as an + /// `FFI_ArrowSchema` and the function returns 0. On failure, returns + /// a non-zero errno and the host calls `get_last_error` to retrieve + /// the message. + /// + /// `arg_types` points to an array of `n_args` `*const FFI_ArrowSchema`. + /// `scalar_args` may be NULL (no scalars) or point to an array of + /// `n_args` `*mut FFI_ArrowArray`, each of length 1 (or NULL when + /// the corresponding argument is not a scalar). Implementations may + /// take ownership of scalar entries by replacing them with released + /// arrays. + pub init: Option< + unsafe extern "C" fn( + *mut CometCScalarKernelImpl, + arg_types: *const *const FFI_ArrowSchema, + scalar_args: *const *mut FFI_ArrowArray, + n_args: i64, + out: *mut FFI_ArrowSchema, + ) -> c_int, + >, + + /// Execute one batch. + /// + /// `args` points to an array of `n_args` `*mut FFI_ArrowArray`. + /// Each input must have length `n_rows` or length 1 (scalar broadcast). + /// On success writes the result into `out` and returns 0. + pub execute: Option< + unsafe extern "C" fn( + *mut CometCScalarKernelImpl, + args: *const *mut FFI_ArrowArray, + n_args: i64, + n_rows: i64, + out: *mut FFI_ArrowArray, + ) -> c_int, + >, + + /// Return the last error message produced by `init` or `execute`. + /// + /// Returns NULL if there is no error. The pointer is valid until the + /// next call to any method on this instance (or `release`). + pub get_last_error: Option<unsafe extern "C" fn(*mut CometCScalarKernelImpl) -> *const c_char>, + + /// Release this instance. After release `release` must be `None`. + pub release: Option<unsafe extern "C" fn(*mut CometCScalarKernelImpl)>, + + /// Implementation-private data, opaque to the host. + pub private_data: *mut c_void, +} + +impl Default for CometCScalarKernelImpl { + fn default() -> Self { + Self { + init: None, + execute: None, + get_last_error: None, + release: None, + private_data: std::ptr::null_mut(), + } + } +} + +impl Drop for CometCScalarKernelImpl { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: per the FFI contract `release` cleans up + // private_data and resets `release` to None. + unsafe { release(self) }; + } + } +} + +// -- discovery list -------------------------------------------------------- + +/// List of kernels exposed by a cdylib via `comet_c_udf_list_v1`. +/// +/// Ownership of the underlying `CometCScalarKernel` array is transferred +/// to the host: the host invokes each kernel's `release` and then frees +/// the list via `release_list`. +#[repr(C)] +pub struct CometCScalarKernelList { + /// Pointer to the kernel array, or null if `len == 0`. + pub kernels: *mut CometCScalarKernel, + /// Number of kernels in `kernels`. + pub len: i64, + /// Free the array of kernels. Implementations must invoke each + /// kernel's `release` first, then release the array storage. + pub release: Option<unsafe extern "C" fn(*mut CometCScalarKernelList)>, +} + +impl Default for CometCScalarKernelList { + fn default() -> Self { + Self { + kernels: std::ptr::null_mut(), + len: 0, + release: None, + } + } +} + +impl Drop for CometCScalarKernelList { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: `release` is responsible for freeing each kernel and + // the array storage that backs `kernels`. + unsafe { release(self) }; + } + } +} + +// -- high-level Rust trait + adapter -------------------------------------- + +use arrow::array::ArrayRef; +use arrow::datatypes::Field; + +/// High-level Rust trait the user implements to author a UDF. +/// +/// Adapted to the C ABI by [`ExportedScalarKernel`]. +pub trait CometCScalarUdf: Send + Sync { + /// Stable function name. Returned via `function_name` over the FFI. + fn name(&self) -> &str; + + /// Compute the output `Field` from the input `Field`s. + /// + /// Called once per execution, before `invoke`. May reject input + /// arities or types by returning an error; the host then surfaces + /// the message to the planner. + fn return_field(&self, args: &[Field]) -> Result<Field, String>; + + /// Evaluate one batch of `n_rows` rows. + fn invoke(&self, args: &[ArrayRef], n_rows: usize) -> Result<ArrayRef, String>; +} + +/// Wraps a user `CometCScalarUdf` impl as a [`CometCScalarKernel`] +/// suitable for emission via the C ABI discovery list. +pub struct ExportedScalarKernel { + inner: std::sync::Arc<dyn CometCScalarUdf>, + /// NUL-terminated C string holding the function name. Lifetime is + /// tied to `self` so the pointer returned to the host stays valid. + name_c: std::ffi::CString, +} + +impl ExportedScalarKernel { + /// Wrap `udf` for export. + pub fn new<U: CometCScalarUdf + 'static>(udf: U) -> Self { + let name_c = std::ffi::CString::new(udf.name().to_string()) + .expect("UDF name must not contain interior NUL bytes"); + Self { + inner: std::sync::Arc::new(udf), + name_c, + } + } +} + +impl From<ExportedScalarKernel> for CometCScalarKernel { + fn from(value: ExportedScalarKernel) -> Self { + let boxed: Box<ExportedScalarKernel> = Box::new(value); + let private = Box::into_raw(boxed) as *mut c_void; + CometCScalarKernel { + function_name: Some(c_factory_function_name), + new_impl: Some(c_factory_new_impl), + release: Some(c_factory_release), + private_data: private, + } + } +} + +unsafe extern "C" fn c_factory_function_name(this: *const CometCScalarKernel) -> *const c_char { + debug_assert!(!this.is_null()); + let this = unsafe { &*this }; + debug_assert!(!this.private_data.is_null()); Review Comment: The release callback might also be good to check here (and in other debug asserts) ```suggestion debug_assert!(!this.private_data.is_null() && this.release.is_some()); ``` ########## native/comet-udf-sdk/src/c_abi.rs: ########## @@ -0,0 +1,835 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The Comet UDF C ABI — sedona-style. +//! +//! The wire format is two `#[repr(C)]` structs of function pointers, +//! parameterized only by Arrow's C Data Interface +//! (`FFI_ArrowSchema` / `FFI_ArrowArray`). No DataFusion types appear in +//! the FFI surface, so the user's cdylib only needs a matching `arrow` +//! crate, not a matching `datafusion` version. +//! +//! # Authoring a UDF +//! +//! Implement [`CometCScalarUdf`] for a type that also implements `Default`, +//! then use the [`comet_c_udf_export!`] macro to emit the discovery entry +//! point: +//! +//! ```ignore +//! use comet_udf_sdk::c_abi::*; +//! use arrow::array::{ArrayRef, Int64Array}; +//! use arrow::datatypes::{DataType, Field}; +//! use std::sync::Arc; +//! +//! #[derive(Default)] +//! pub struct AddOne; +//! impl CometCScalarUdf for AddOne { +//! fn name(&self) -> &str { "add_one_c" } +//! fn return_field(&self, args: &[Field]) -> Result<Field, String> { +//! if args.len() != 1 || args[0].data_type() != &DataType::Int64 { +//! return Err("expected (Int64) -> Int64".into()); +//! } +//! Ok(Field::new("add_one_c", DataType::Int64, true)) +//! } +//! fn invoke(&self, args: &[ArrayRef], _n: usize) -> Result<ArrayRef, String> { +//! let a = args[0].as_any().downcast_ref::<Int64Array>().unwrap(); +//! Ok(Arc::new(a.iter().map(|v| v.map(|x| x + 1)).collect::<Int64Array>())) +//! } +//! } +//! +//! comet_udf_sdk::comet_c_udf_export!(AddOne); +//! ``` + +use std::ffi::{c_char, c_int, c_void}; + +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; + +/// Generic non-zero error code returned by `init` / `execute` to signal +/// failure. The host treats any non-zero return as an error and calls +/// `get_last_error` for the message; the specific code is informational. +const C_ABI_ERR: c_int = 1; + +// -- panic containment ----------------------------------------------------- +// +// Every `extern "C"` function in this module is an unwind boundary. A panic +// that escapes one aborts the whole process (Rust's default `extern "C"` +// unwind behavior since 1.81), which for Comet means killing the executor +// JVM and losing every task on it -- not just the query that used the UDF. +// +// User UDF code is arbitrary and panicking is idiomatic Rust (`unwrap`, +// slice indexing, integer overflow in debug), so the SDK treats a panic in +// user code as an ordinary error: catch it at the boundary, convert it to a +// message, and report it through the same `get_last_error` channel as a +// returned `Err`. The query fails; the executor survives. + +/// Render a caught panic payload as an error message. +fn panic_message(panic: Box<dyn std::any::Any + Send>) -> String { + let detail = panic + .downcast_ref::<&'static str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::<String>().cloned()) + .unwrap_or_else(|| "<non-string panic payload>".to_string()); + format!("panic in UDF code: {detail}") +} + +/// Run `f`, converting a panic into `Err(message)`. +fn catch_panic<T>(f: impl FnOnce() -> Result<T, String>) -> Result<T, String> { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(result) => result, + Err(panic) => Err(panic_message(panic)), + } +} + +/// Run an infallible `f` (typically a release/cleanup callback), swallowing +/// any panic. Used where the ABI gives us no way to report an error and +/// aborting would be a worse outcome than leaking. +fn catch_panic_infallible(f: impl FnOnce()) { + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); +} + +// -- factory struct -------------------------------------------------------- + +/// Factory for [`CometCScalarKernelImpl`] instances. +/// +/// Lives in a registry, may be cloned across an FFI boundary. Calls to +/// `function_name` and `new_impl` must be thread-safe (the implementation +/// is responsible for any internal synchronization). +/// +/// `#[repr(C)]` layout, matched by the host loader. Adding new fields +/// requires bumping `COMET_UDF_ABI_VERSION`. +#[repr(C)] +pub struct CometCScalarKernel { + /// Return the function name this kernel implements as a NUL-terminated + /// UTF-8 C string. The pointer must remain valid for the lifetime of + /// the [`CometCScalarKernel`]. + /// + /// May be `None`, in which case the kernel is treated as anonymous and + /// won't be discoverable by name. (Comet always sets this; field is + /// optional for parity with sedona's design.) + pub function_name: Option<unsafe extern "C" fn(*const CometCScalarKernel) -> *const c_char>, + + /// Initialize a new [`CometCScalarKernelImpl`] into `out`. Called once + /// per execution, on the thread that will then drive `init`/`execute`. + pub new_impl: + Option<unsafe extern "C" fn(*const CometCScalarKernel, out: *mut CometCScalarKernelImpl)>, + + /// Release this kernel. After release, all callbacks must be set to + /// `None`. Called when the host's `LoadedLibrary` is dropped. + pub release: Option<unsafe extern "C" fn(*mut CometCScalarKernel)>, + + /// Implementation-private data, opaque to the host. + pub private_data: *mut c_void, +} + +// SAFETY: `CometCScalarKernel` is a thin wrapper around C function +// pointers with caller-defined synchronization semantics; the trait impls +// are required so loaded kernels can be referenced from multi-threaded +// host code. Implementations of the FFI must respect thread safety as +// described in the doc comments. +unsafe impl Send for CometCScalarKernel {} +unsafe impl Sync for CometCScalarKernel {} + +impl Default for CometCScalarKernel { + fn default() -> Self { + Self { + function_name: None, + new_impl: None, + release: None, + private_data: std::ptr::null_mut(), + } + } +} + +impl Drop for CometCScalarKernel { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: release is the FFI-defined cleanup callback; + // implementations must reset `release` to None per the contract. + unsafe { release(self) }; + } + } +} + +// -- per-execution instance struct ---------------------------------------- + +/// Per-execution instance produced by [`CometCScalarKernel::new_impl`]. +/// +/// Not thread-safe; the caller must serialize access. Typically used on +/// one thread for one batch then dropped. +#[repr(C)] +pub struct CometCScalarKernelImpl { + /// Compute the return type from arg types and (optionally) bound + /// scalar arguments. + /// + /// On success, `out` is populated with the return type as an + /// `FFI_ArrowSchema` and the function returns 0. On failure, returns + /// a non-zero errno and the host calls `get_last_error` to retrieve + /// the message. + /// + /// `arg_types` points to an array of `n_args` `*const FFI_ArrowSchema`. + /// `scalar_args` may be NULL (no scalars) or point to an array of + /// `n_args` `*mut FFI_ArrowArray`, each of length 1 (or NULL when + /// the corresponding argument is not a scalar). Implementations may + /// take ownership of scalar entries by replacing them with released + /// arrays. + pub init: Option< + unsafe extern "C" fn( + *mut CometCScalarKernelImpl, + arg_types: *const *const FFI_ArrowSchema, + scalar_args: *const *mut FFI_ArrowArray, + n_args: i64, + out: *mut FFI_ArrowSchema, + ) -> c_int, + >, + + /// Execute one batch. + /// + /// `args` points to an array of `n_args` `*mut FFI_ArrowArray`. + /// Each input must have length `n_rows` or length 1 (scalar broadcast). + /// On success writes the result into `out` and returns 0. + pub execute: Option< + unsafe extern "C" fn( + *mut CometCScalarKernelImpl, + args: *const *mut FFI_ArrowArray, + n_args: i64, + n_rows: i64, + out: *mut FFI_ArrowArray, + ) -> c_int, + >, + + /// Return the last error message produced by `init` or `execute`. + /// + /// Returns NULL if there is no error. The pointer is valid until the + /// next call to any method on this instance (or `release`). + pub get_last_error: Option<unsafe extern "C" fn(*mut CometCScalarKernelImpl) -> *const c_char>, + + /// Release this instance. After release `release` must be `None`. + pub release: Option<unsafe extern "C" fn(*mut CometCScalarKernelImpl)>, + + /// Implementation-private data, opaque to the host. + pub private_data: *mut c_void, +} + +impl Default for CometCScalarKernelImpl { + fn default() -> Self { + Self { + init: None, + execute: None, + get_last_error: None, + release: None, + private_data: std::ptr::null_mut(), + } + } +} + +impl Drop for CometCScalarKernelImpl { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + // SAFETY: per the FFI contract `release` cleans up + // private_data and resets `release` to None. + unsafe { release(self) }; + } + } +} + +// -- discovery list -------------------------------------------------------- + +/// List of kernels exposed by a cdylib via `comet_c_udf_list_v1`. +/// +/// Ownership of the underlying `CometCScalarKernel` array is transferred +/// to the host: the host invokes each kernel's `release` and then frees +/// the list via `release_list`. +#[repr(C)] +pub struct CometCScalarKernelList { + /// Pointer to the kernel array, or null if `len == 0`. + pub kernels: *mut CometCScalarKernel, + /// Number of kernels in `kernels`. + pub len: i64, + /// Free the array of kernels. Implementations must invoke each + /// kernel's `release` first, then release the array storage. + pub release: Option<unsafe extern "C" fn(*mut CometCScalarKernelList)>, Review Comment: Optional, but this seems like it would be easy to forget to do...if it works with the use of this, it would probably be less error prone to force a caller to "move" the kernel (i.e., set the release callback of the array version wrapped here to NULL and force the caller to take responsibility of the C struct. Then this release callback would drop any valid kernels that remained. ########## native/core/src/execution/rust_udf/loader.rs: ########## @@ -0,0 +1,275 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Loader: open a UDF cdylib via libloading, validate the ABI version, +//! discover UDFs via the C-ABI entry point, and produce DataFusion +//! `ScalarUDFImpl` impls for each. +//! +//! See `super::mod.rs` for an overview of the ABI. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use comet_udf_sdk::c_abi::CometCScalarKernelList; +use comet_udf_sdk::{ABI_VERSION_SYMBOL, COMET_UDF_ABI_VERSION, C_ABI_DISCOVERY_SYMBOL}; +use datafusion::logical_expr::ScalarUDFImpl; +use libloading::{Library, Symbol}; + +use super::imported_c::ImportedCScalarUdf; + +/// One loaded UDF: name plus a `ScalarUDFImpl` ready to plug into the +/// planner. +pub struct LoadedUdf { + /// UDF name as exposed by the cdylib. + pub name: String, + /// The `ScalarUDFImpl` adapter the planner will wrap. + pub udf_impl: Arc<dyn ScalarUDFImpl>, +} + +impl std::fmt::Debug for LoadedUdf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedUdf") + .field("name", &self.name) + .finish() + } +} + +/// Result of loading a UDF cdylib: the live `Library` plus per-UDF +/// adapters. +pub struct LoadedLibrary { + /// Canonicalized path the library was loaded from. + pub path: PathBuf, + /// The loaded `Library`. Held inside an `Arc` so loaded UDFs can + /// outlive lookups. Library is never unloaded for the process lifetime. + pub library: Arc<Library>, + /// One entry per UDF, with name and ScalarUDFImpl already built. + pub udfs: Vec<LoadedUdf>, +} + +impl std::fmt::Debug for LoadedLibrary { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedLibrary") + .field("path", &self.path) + .field("udfs", &self.udfs) + .finish() + } +} + +/// Errors returned by the loader. +#[derive(Debug)] +pub enum LoaderError { + /// `libloading::Library::new` failed. + Open { + /// Path that was passed to `Library::new`. + path: PathBuf, + /// Underlying error. + source: libloading::Error, + }, + /// `comet_udf_abi_version` is missing or returned an unexpected value. + AbiMismatch { + /// Path of the offending library. + path: PathBuf, + /// Version reported by the cdylib (or `None` if the symbol is missing). + found: Option<u32>, + /// Version this host expects. + expected: u32, + }, + /// Library does not expose `comet_c_udf_list_v1`. + NoDiscovery { + /// Path of the offending library. + path: PathBuf, + }, + /// The discovery function returned a non-zero rc, or a kernel entry + /// was malformed. + Discovery { + /// Path of the library. + path: PathBuf, + /// Human-readable reason. + reason: String, + }, +} + +impl std::fmt::Display for LoaderError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use LoaderError::*; + match self { + Open { path, source } => write!(f, "failed to open {}: {source}", path.display()), + AbiMismatch { + path, + found, + expected, + } => match found { + Some(v) => write!( + f, + "{} reports ABI v{v}, host expects v{expected}", + path.display() + ), + None => write!( + f, + "{} missing required symbol {ABI_VERSION_SYMBOL}", + path.display() + ), + }, + NoDiscovery { path } => write!( + f, + "{} does not export {C_ABI_DISCOVERY_SYMBOL}", + path.display() + ), + Discovery { path, reason } => write!(f, "{}: {reason}", path.display()), + } + } +} + +impl std::error::Error for LoaderError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + LoaderError::Open { source, .. } => Some(source), + _ => None, + } + } +} + +/// Open and validate a UDF cdylib. +pub fn load(path: impl AsRef<Path>) -> Result<LoadedLibrary, LoaderError> { + let path = path.as_ref().to_path_buf(); + // SAFETY: `Library::new` runs the cdylib's static initializers. We + // accept this risk because user UDF cdylibs are explicitly registered + // by an operator via `CometRustUDF.register`. + let library = unsafe { Library::new(&path) }.map_err(|source| LoaderError::Open { Review Comment: This version of loading I believe will also resolve libraries on LD_LIBRARY_PATH (the previous version canonicalizes the path first, so I think only resolves against the working directory or an absolute path) ########## native/proto/src/proto/expr.proto: ########## @@ -618,3 +619,22 @@ message JvmScalarUdf { // Whether the result column may contain nulls. bool return_nullable = 4; } + +// Call to a user-supplied Rust UDF loaded from a cdylib. +// +// The native side resolves (library_path, name) against its loaded-library +// cache, looks up the kernel by name, and invokes it through whichever ABI +// flavor (C ABI / datafusion-ffi) the cdylib registered the kernel under. +message RustUdfCall { + // Function name as registered through CometRustUDF.register on the JVM + // side; matched against names exposed by the cdylib. + string name = 1; + // Filesystem path of the cdylib. + string library_path = 2; + // Argument expressions, evaluated before invocation. + repeated Expr args = 3; + // Expected return type, declared at register time on the JVM side. + DataType return_type = 4; + // Whether the call is deterministic (mirrors Spark's deterministic flag). + bool deterministic = 5; Review Comment: Is this always `true` because RustUdfs are always immutable? ########## native/comet-udf-sdk/src/c_abi.rs: ########## @@ -0,0 +1,835 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The Comet UDF C ABI — sedona-style. +//! +//! The wire format is two `#[repr(C)]` structs of function pointers, +//! parameterized only by Arrow's C Data Interface +//! (`FFI_ArrowSchema` / `FFI_ArrowArray`). No DataFusion types appear in +//! the FFI surface, so the user's cdylib only needs a matching `arrow` +//! crate, not a matching `datafusion` version. +//! +//! # Authoring a UDF +//! +//! Implement [`CometCScalarUdf`] for a type that also implements `Default`, +//! then use the [`comet_c_udf_export!`] macro to emit the discovery entry +//! point: +//! +//! ```ignore +//! use comet_udf_sdk::c_abi::*; +//! use arrow::array::{ArrayRef, Int64Array}; +//! use arrow::datatypes::{DataType, Field}; +//! use std::sync::Arc; +//! +//! #[derive(Default)] +//! pub struct AddOne; +//! impl CometCScalarUdf for AddOne { +//! fn name(&self) -> &str { "add_one_c" } +//! fn return_field(&self, args: &[Field]) -> Result<Field, String> { +//! if args.len() != 1 || args[0].data_type() != &DataType::Int64 { +//! return Err("expected (Int64) -> Int64".into()); +//! } +//! Ok(Field::new("add_one_c", DataType::Int64, true)) +//! } +//! fn invoke(&self, args: &[ArrayRef], _n: usize) -> Result<ArrayRef, String> { +//! let a = args[0].as_any().downcast_ref::<Int64Array>().unwrap(); +//! Ok(Arc::new(a.iter().map(|v| v.map(|x| x + 1)).collect::<Int64Array>())) +//! } +//! } +//! +//! comet_udf_sdk::comet_c_udf_export!(AddOne); +//! ``` + +use std::ffi::{c_char, c_int, c_void}; + +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; + +/// Generic non-zero error code returned by `init` / `execute` to signal +/// failure. The host treats any non-zero return as an error and calls +/// `get_last_error` for the message; the specific code is informational. +const C_ABI_ERR: c_int = 1; + +// -- panic containment ----------------------------------------------------- +// +// Every `extern "C"` function in this module is an unwind boundary. A panic +// that escapes one aborts the whole process (Rust's default `extern "C"` +// unwind behavior since 1.81), which for Comet means killing the executor +// JVM and losing every task on it -- not just the query that used the UDF. +// +// User UDF code is arbitrary and panicking is idiomatic Rust (`unwrap`, +// slice indexing, integer overflow in debug), so the SDK treats a panic in +// user code as an ordinary error: catch it at the boundary, convert it to a +// message, and report it through the same `get_last_error` channel as a +// returned `Err`. The query fails; the executor survives. + +/// Render a caught panic payload as an error message. +fn panic_message(panic: Box<dyn std::any::Any + Send>) -> String { + let detail = panic + .downcast_ref::<&'static str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::<String>().cloned()) + .unwrap_or_else(|| "<non-string panic payload>".to_string()); + format!("panic in UDF code: {detail}") +} + +/// Run `f`, converting a panic into `Err(message)`. +fn catch_panic<T>(f: impl FnOnce() -> Result<T, String>) -> Result<T, String> { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(result) => result, + Err(panic) => Err(panic_message(panic)), + } +} + +/// Run an infallible `f` (typically a release/cleanup callback), swallowing +/// any panic. Used where the ABI gives us no way to report an error and +/// aborting would be a worse outcome than leaking. +fn catch_panic_infallible(f: impl FnOnce()) { + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); +} + +// -- factory struct -------------------------------------------------------- + +/// Factory for [`CometCScalarKernelImpl`] instances. +/// +/// Lives in a registry, may be cloned across an FFI boundary. Calls to +/// `function_name` and `new_impl` must be thread-safe (the implementation +/// is responsible for any internal synchronization). +/// +/// `#[repr(C)]` layout, matched by the host loader. Adding new fields +/// requires bumping `COMET_UDF_ABI_VERSION`. +#[repr(C)] +pub struct CometCScalarKernel { + /// Return the function name this kernel implements as a NUL-terminated + /// UTF-8 C string. The pointer must remain valid for the lifetime of + /// the [`CometCScalarKernel`]. + /// + /// May be `None`, in which case the kernel is treated as anonymous and + /// won't be discoverable by name. (Comet always sets this; field is + /// optional for parity with sedona's design.) + pub function_name: Option<unsafe extern "C" fn(*const CometCScalarKernel) -> *const c_char>, + Review Comment: No pressure to do this here, but you could do something like this to pass the volatility, Display, Debug, etc. over FFI in a forward-flexible way. I haven't added this for UDFs yet but I'm using it for table providers, exec plans, and expressions. The annoying part is returning a variable length string (I use a FFI_ArrowArray for this at the moment, which is possibly overkill). ```suggestion pub get_property: Option< unsafe extern "C" fn( *mut CometCScalarKernelImpl, property: *const c_char, args: *const c_char, out: *mut FFI_ArrowArray, ) -> c_int, >, ``` ########## native/core/src/execution/rust_udf/imported_c.rs: ########## @@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Adapter wrapping a C-ABI [`CometCScalarKernel`] as a DataFusion +//! [`ScalarUDFImpl`]. +//! +//! Lifecycle inside `invoke_with_args`: +//! +//! 1. Build a fresh [`CometCScalarKernelImpl`] via the kernel's `new_impl`. +//! 2. Call `init` with the input field types (and any scalar args) to get +//! the return type. +//! 3. Call `execute` once with the batch. +//! 4. Drop the impl (its `release` callback runs). + +use std::ffi::CStr; +use std::sync::Mutex; + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field}; +use arrow::ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema}; +use comet_udf_sdk::c_abi::{CometCScalarKernel, CometCScalarKernelImpl}; +use datafusion::common::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; + +/// Adapter wrapping a [`CometCScalarKernel`] as a DataFusion +/// [`ScalarUDFImpl`]. +pub struct ImportedCScalarUdf { + name: String, + /// Boxed so the kernel's address is stable; held inside a Mutex + /// because the FFI Drop is not Sync-safe under concurrent invocation. + /// The kernel itself is logically immutable post-load — the lock only + /// protects the FFI calls' aliasing rules. (DataFusion serializes + /// invocations of a given ScalarUDFImpl per-batch through + /// invoke_with_args anyway; the lock is defensive.) + kernel: Mutex<Box<CometCScalarKernel>>, Review Comment: Would `Arc<CometCScalarKernel>` sufficiently maintain reference counts to the underlying kernel and release the instance correctly when there is exactly one reference left? ########## native/core/src/execution/rust_udf/imported_c.rs: ########## @@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Adapter wrapping a C-ABI [`CometCScalarKernel`] as a DataFusion +//! [`ScalarUDFImpl`]. +//! +//! Lifecycle inside `invoke_with_args`: +//! +//! 1. Build a fresh [`CometCScalarKernelImpl`] via the kernel's `new_impl`. +//! 2. Call `init` with the input field types (and any scalar args) to get +//! the return type. +//! 3. Call `execute` once with the batch. +//! 4. Drop the impl (its `release` callback runs). + +use std::ffi::CStr; +use std::sync::Mutex; + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field}; +use arrow::ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema}; +use comet_udf_sdk::c_abi::{CometCScalarKernel, CometCScalarKernelImpl}; +use datafusion::common::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; + +/// Adapter wrapping a [`CometCScalarKernel`] as a DataFusion +/// [`ScalarUDFImpl`]. +pub struct ImportedCScalarUdf { + name: String, + /// Boxed so the kernel's address is stable; held inside a Mutex + /// because the FFI Drop is not Sync-safe under concurrent invocation. + /// The kernel itself is logically immutable post-load — the lock only + /// protects the FFI calls' aliasing rules. (DataFusion serializes + /// invocations of a given ScalarUDFImpl per-batch through + /// invoke_with_args anyway; the lock is defensive.) + kernel: Mutex<Box<CometCScalarKernel>>, + signature: Signature, +} + +impl ImportedCScalarUdf { + /// Construct from an owned C kernel. + /// + /// Reads the kernel's name via its `function_name` callback and + /// stores it for `name()` lookups; the kernel itself is held inside + /// a mutex. + pub fn try_new(kernel: Box<CometCScalarKernel>) -> Result<Self, String> { + let function_name_cb = kernel + .function_name + .ok_or_else(|| "kernel.function_name is null".to_string())?; + let _ = kernel + .new_impl + .ok_or_else(|| "kernel.new_impl is null".to_string())?; + + // SAFETY: function_name_cb is the FFI-supplied callback; + // implementations promise the returned pointer is a NUL-terminated + // UTF-8 string valid for the lifetime of the kernel. + let name_ptr = unsafe { function_name_cb(kernel.as_ref() as *const _) }; + if name_ptr.is_null() { + return Err("function_name returned null".into()); + } + let name = unsafe { CStr::from_ptr(name_ptr) } + .to_str() + .map_err(|e| format!("function_name not UTF-8: {e}"))? + .to_string(); + + // Use UserDefined signature: per-call init() is what decides + // whether the input types are acceptable. `coerce_types` is not + // implemented; user must pass exact types from the JVM register call. + let signature = Signature::new(TypeSignature::UserDefined, Volatility::Immutable); + + Ok(Self { + name, + kernel: Mutex::new(kernel), + signature, + }) + } +} + +impl std::fmt::Debug for ImportedCScalarUdf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ImportedCScalarUdf") + .field("name", &self.name) + .finish() + } +} + +impl PartialEq for ImportedCScalarUdf { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + } +} + +impl Eq for ImportedCScalarUdf {} + +impl std::hash::Hash for ImportedCScalarUdf { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.name.hash(state); + } +} + +impl ScalarUDFImpl for ImportedCScalarUdf { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, args: &[DataType]) -> datafusion::common::Result<DataType> { + // Build a fresh impl, call init, drop. Done at planning time so + // the planner can know the output type before execution. + let kernel = self.kernel.lock().unwrap(); + let mut impl_state = CometCScalarKernelImpl::default(); + let new_impl_cb = kernel + .new_impl + .ok_or_else(|| DataFusionError::Internal("new_impl is null".into()))?; + // SAFETY: new_impl_cb is the FFI-supplied factory; impl_state is a + // caller-allocated default value the cdylib writes into. + unsafe { + new_impl_cb(kernel.as_ref() as *const _, &mut impl_state); + } + + // Build input fields and FFI schemas. + let fields: Vec<Field> = args + .iter() + .map(|dt| Field::new("", dt.clone(), true)) + .collect(); + let ffi_schemas = build_ffi_schemas(&fields)?; + let ffi_schema_ptrs: Vec<*const FFI_ArrowSchema> = + ffi_schemas.iter().map(|s| s as *const _).collect(); + + let init_cb = impl_state + .init + .ok_or_else(|| DataFusionError::Internal("kernel impl missing init".into()))?; + let mut out_schema = FFI_ArrowSchema::empty(); + // SAFETY: pointers are valid for the duration of the call. + let rc = unsafe { + init_cb( + &mut impl_state, + ffi_schema_ptrs.as_ptr(), + std::ptr::null(), + fields.len() as i64, + &mut out_schema, + ) + }; + if rc != 0 { + let msg = read_last_error(&mut impl_state); + return Err(DataFusionError::Plan(format!( + "{}: init failed: {msg}", + self.name + ))); + } + let return_field = Field::try_from(&out_schema) + .map_err(|e| DataFusionError::Internal(format!("decoding return type: {e}")))?; + Ok(return_field.data_type().clone()) Review Comment: I know you don't support them one level up, but since you're calculating a field here anyway can you implement `return_field()`? ########## native/core/src/execution/rust_udf/loader.rs: ########## @@ -0,0 +1,275 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Loader: open a UDF cdylib via libloading, validate the ABI version, +//! discover UDFs via the C-ABI entry point, and produce DataFusion +//! `ScalarUDFImpl` impls for each. +//! +//! See `super::mod.rs` for an overview of the ABI. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use comet_udf_sdk::c_abi::CometCScalarKernelList; +use comet_udf_sdk::{ABI_VERSION_SYMBOL, COMET_UDF_ABI_VERSION, C_ABI_DISCOVERY_SYMBOL}; +use datafusion::logical_expr::ScalarUDFImpl; +use libloading::{Library, Symbol}; + +use super::imported_c::ImportedCScalarUdf; + +/// One loaded UDF: name plus a `ScalarUDFImpl` ready to plug into the +/// planner. +pub struct LoadedUdf { + /// UDF name as exposed by the cdylib. + pub name: String, + /// The `ScalarUDFImpl` adapter the planner will wrap. + pub udf_impl: Arc<dyn ScalarUDFImpl>, +} + +impl std::fmt::Debug for LoadedUdf { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedUdf") + .field("name", &self.name) + .finish() + } +} + +/// Result of loading a UDF cdylib: the live `Library` plus per-UDF +/// adapters. +pub struct LoadedLibrary { + /// Canonicalized path the library was loaded from. + pub path: PathBuf, + /// The loaded `Library`. Held inside an `Arc` so loaded UDFs can + /// outlive lookups. Library is never unloaded for the process lifetime. + pub library: Arc<Library>, + /// One entry per UDF, with name and ScalarUDFImpl already built. + pub udfs: Vec<LoadedUdf>, +} + +impl std::fmt::Debug for LoadedLibrary { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LoadedLibrary") + .field("path", &self.path) + .field("udfs", &self.udfs) + .finish() + } +} + +/// Errors returned by the loader. +#[derive(Debug)] +pub enum LoaderError { + /// `libloading::Library::new` failed. + Open { + /// Path that was passed to `Library::new`. + path: PathBuf, + /// Underlying error. + source: libloading::Error, + }, + /// `comet_udf_abi_version` is missing or returned an unexpected value. + AbiMismatch { + /// Path of the offending library. + path: PathBuf, + /// Version reported by the cdylib (or `None` if the symbol is missing). + found: Option<u32>, + /// Version this host expects. + expected: u32, + }, + /// Library does not expose `comet_c_udf_list_v1`. + NoDiscovery { + /// Path of the offending library. + path: PathBuf, + }, + /// The discovery function returned a non-zero rc, or a kernel entry + /// was malformed. + Discovery { + /// Path of the library. + path: PathBuf, + /// Human-readable reason. + reason: String, + }, +} + +impl std::fmt::Display for LoaderError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use LoaderError::*; + match self { + Open { path, source } => write!(f, "failed to open {}: {source}", path.display()), + AbiMismatch { + path, + found, + expected, + } => match found { + Some(v) => write!( + f, + "{} reports ABI v{v}, host expects v{expected}", + path.display() + ), + None => write!( + f, + "{} missing required symbol {ABI_VERSION_SYMBOL}", + path.display() + ), + }, + NoDiscovery { path } => write!( + f, + "{} does not export {C_ABI_DISCOVERY_SYMBOL}", + path.display() + ), + Discovery { path, reason } => write!(f, "{}: {reason}", path.display()), + } + } +} + +impl std::error::Error for LoaderError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + LoaderError::Open { source, .. } => Some(source), + _ => None, + } + } +} + +/// Open and validate a UDF cdylib. +pub fn load(path: impl AsRef<Path>) -> Result<LoadedLibrary, LoaderError> { + let path = path.as_ref().to_path_buf(); + // SAFETY: `Library::new` runs the cdylib's static initializers. We + // accept this risk because user UDF cdylibs are explicitly registered + // by an operator via `CometRustUDF.register`. + let library = unsafe { Library::new(&path) }.map_err(|source| LoaderError::Open { + path: path.clone(), + source, + })?; + + // ABI version probe. + let v = read_abi_version(&library, &path)?; + if v != COMET_UDF_ABI_VERSION { + return Err(LoaderError::AbiMismatch { + path, + found: Some(v), + expected: COMET_UDF_ABI_VERSION, + }); + } + + let udfs = match read_c_kernels(&library, &path)? { + Some(udfs) => udfs, + None => return Err(LoaderError::NoDiscovery { path }), + }; + + Ok(LoadedLibrary { + path, + library: Arc::new(library), + udfs, + }) +} + +fn read_abi_version(lib: &Library, path: &Path) -> Result<u32, LoaderError> { + let sym: Symbol<unsafe extern "C" fn() -> u32> = unsafe { + lib.get(ABI_VERSION_SYMBOL.as_bytes()) + } + .map_err(|_| LoaderError::AbiMismatch { + path: path.to_path_buf(), + found: None, + expected: COMET_UDF_ABI_VERSION, + })?; + // SAFETY: comet_udf_abi_version takes no arguments, returns u32, no side effects. + Ok(unsafe { sym() }) +} + +fn read_c_kernels(lib: &Library, path: &Path) -> Result<Option<Vec<LoadedUdf>>, LoaderError> { + let sym: Symbol<unsafe extern "C" fn(*mut CometCScalarKernelList) -> i32> = + match unsafe { lib.get(C_ABI_DISCOVERY_SYMBOL.as_bytes()) } { + Ok(s) => s, + Err(_) => return Ok(None), + }; + let mut list = CometCScalarKernelList::default(); + // SAFETY: list is caller-allocated; the cdylib writes into it via `out`. + let rc = unsafe { sym(&mut list) }; + if rc != 0 { + return Err(LoaderError::Discovery { + path: path.to_path_buf(), + reason: format!("{C_ABI_DISCOVERY_SYMBOL} returned rc={rc}"), + }); + } + let mut udfs = Vec::with_capacity(list.len.max(0) as usize); + if !list.kernels.is_null() && list.len > 0 { + // Move each kernel out of the array into a Box so it owns itself. + // We can't simply read each entry because they implement Drop; + // doing it via std::ptr::read transfers ownership cleanly. + let len = list.len as usize; + for i in 0..len { + // SAFETY: the kernel array was produced by the cdylib's + // `comet_c_udf_export!` and contains `len` valid entries. + // We move each entry out into a Box so its Drop runs when + // the host releases the loaded library. + let raw = unsafe { list.kernels.add(i) }; + let kernel = unsafe { std::ptr::read(raw) }; + // Replace the slot with a default kernel (no callbacks) so + // the array's release doesn't double-free. Review Comment: I think the array's release won't double free here as currently implemented (but should!). This probably never leaks because none of this is ever likely to fail but in theory it would if an error occurred mid-import. ########## native/core/src/execution/rust_udf/imported_c.rs: ########## @@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Adapter wrapping a C-ABI [`CometCScalarKernel`] as a DataFusion +//! [`ScalarUDFImpl`]. +//! +//! Lifecycle inside `invoke_with_args`: +//! +//! 1. Build a fresh [`CometCScalarKernelImpl`] via the kernel's `new_impl`. +//! 2. Call `init` with the input field types (and any scalar args) to get +//! the return type. +//! 3. Call `execute` once with the batch. +//! 4. Drop the impl (its `release` callback runs). + +use std::ffi::CStr; +use std::sync::Mutex; + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field}; +use arrow::ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema}; +use comet_udf_sdk::c_abi::{CometCScalarKernel, CometCScalarKernelImpl}; +use datafusion::common::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; + +/// Adapter wrapping a [`CometCScalarKernel`] as a DataFusion +/// [`ScalarUDFImpl`]. +pub struct ImportedCScalarUdf { + name: String, + /// Boxed so the kernel's address is stable; held inside a Mutex + /// because the FFI Drop is not Sync-safe under concurrent invocation. + /// The kernel itself is logically immutable post-load — the lock only + /// protects the FFI calls' aliasing rules. (DataFusion serializes + /// invocations of a given ScalarUDFImpl per-batch through + /// invoke_with_args anyway; the lock is defensive.) + kernel: Mutex<Box<CometCScalarKernel>>, + signature: Signature, +} + +impl ImportedCScalarUdf { + /// Construct from an owned C kernel. + /// + /// Reads the kernel's name via its `function_name` callback and + /// stores it for `name()` lookups; the kernel itself is held inside + /// a mutex. + pub fn try_new(kernel: Box<CometCScalarKernel>) -> Result<Self, String> { + let function_name_cb = kernel + .function_name + .ok_or_else(|| "kernel.function_name is null".to_string())?; + let _ = kernel + .new_impl + .ok_or_else(|| "kernel.new_impl is null".to_string())?; + + // SAFETY: function_name_cb is the FFI-supplied callback; + // implementations promise the returned pointer is a NUL-terminated + // UTF-8 string valid for the lifetime of the kernel. + let name_ptr = unsafe { function_name_cb(kernel.as_ref() as *const _) }; + if name_ptr.is_null() { + return Err("function_name returned null".into()); + } + let name = unsafe { CStr::from_ptr(name_ptr) } + .to_str() + .map_err(|e| format!("function_name not UTF-8: {e}"))? + .to_string(); + + // Use UserDefined signature: per-call init() is what decides + // whether the input types are acceptable. `coerce_types` is not + // implemented; user must pass exact types from the JVM register call. + let signature = Signature::new(TypeSignature::UserDefined, Volatility::Immutable); Review Comment: It may be worth documenting in the C ABI limitations section that only immutable functions are supported -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
