andygrove commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3713243875
########## 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: Agreed, and the current shape only works because the host remembers to write a default back into each slot it moves out of. That is exactly the kind of thing that stays correct until someone edits the loader. Not doing it here: making the taker NULL the source `release` changes the ABI contract itself, so it wants to land deliberately rather than as a review fix. Filed as #5250, which also covers the mid-import question you raised further down. ########## 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: Applied in ceb93655d, and extended to the other three entry points that reach `private_data`: `c_factory_new_impl`, `c_kernel_init` and `c_kernel_execute`. A released struct has `release: None`, so checking both catches a call made *after* release rather than only an uninitialized one, which is the more likely mistake. ########## 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: Possible, and worth doing. `catch_panic_infallible` now prints to stderr with a context string naming the callback that panicked (ceb93655d). Not a logging facade, though: `arrow` is deliberately the SDK's only dependency, and pulling in `log` for this would put a facade with no subscriber inside every user cdylib. stderr is the honest option at this layer. The cases it covers are all cleanup paths where the ABI gives no way to return an error, so previously a panicking user `Drop` left a leaked allocation and no trace at all. ########## 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: Added a `# Stability` section to the module docs in ceb93655d, saying plainly that these structs are specific to one Comet version, that they are internal under Comet's versioning policy and may change in any release including a patch, and that the practical consequence is rebuilding the cdylib per Comet release. Your second point is right and I have written it down too: neither the arrow nor the datafusion version has to match, because only `FFI_ArrowArray` and `FFI_ArrowSchema` cross the boundary and those are `#[repr(C)]` renderings of the C Data Interface. The binding constraint is what `comet-udf-sdk` itself compiles against, since the SDK is built into the user's cdylib and Cargo has to unify its `arrow` requirement with theirs. That is the relaxation you raise on `Cargo.toml`, tracked as #5253. ########## 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: Filed as #5253 rather than done here. Two reasons: the crate is `publish = false` and consumed by git, so there is no external consumer to unblock yet, and a relaxed range on one workspace member needs an override rather than `workspace = true`, plus a CI job actually building against the floor. Otherwise the range is a claim nothing tests, which is worse than the current pin. Noted in the issue that this becomes the blocking constraint the moment the SDK is published. ########## 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: Not a security choice, no. `canonicalize()` falls back to the raw path, so a bare name reaches `Library::new` and resolves through the platform loader search path exactly as you describe. The documentation was simply wrong to imply otherwise. Fixed the docs rather than the code in ceb93655d: an absolute path is still the sensible thing on a cluster, but the page now says what actually happens, and says explicitly that it is not a trust boundary, since Comet does not restrict which paths may be loaded either way. ########## 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: Probably, and it is worth doing together with removing the `Mutex` rather than separately. That lock serializes every batch for a given UDF through one mutex per process, which is already on the follow-on list as a performance item, and the reason it is there at all is that the FFI `Drop` is not `Sync`-safe. Both come down to the same question about how the kernel's lifetime is managed. Filed as #5252, noting the two invariants any replacement has to keep: `release` runs exactly once, and the kernel never outlives the `LoadedLibrary` that dlopened it. ########## 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: Documented in ceb93655d, in three places: a `# Only immutable functions are supported` section on the `CometCScalarUdf` trait, a limitation in the user guide, and a comment at the `Volatility::Immutable` line itself explaining why it is not derived from anything. Writing it up turned out to matter more than expected. The guide previously told readers to pass `deterministic = false` for an impure function, and that flag was being silently dropped. See my reply on the proto field. -- 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]
