andygrove commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3713246159
########## 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: Can, and `return_type` is already building the whole `Field` and throwing away everything but the `DataType`, while the planner separately hardcodes the output field as nullable. So it is mostly a matter of returning what is already computed. Filed as #5251 rather than done here, because it changes what nullability Comet promises: Spark UDF results are nullable in Spark's own schema, and the declared-vs-actual check in `planner.rs` deliberately erases *nested* nullability for a related reason. I would rather that land with a test for a kernel reporting a non-nullable field than slip in as part of a docs pass. ########## 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: Correct. `cache::get_or_load` canonicalizes with `unwrap_or_else(|_| raw)`, so a path that does not resolve on the filesystem is handed to `Library::new` unchanged and gets the normal dlopen search. The docs claimed absolute paths were required; they now describe what the code does. Same thread as your `cache.rs` comment. ########## 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: Traced this and I do not think it leaks, though it took following three separate owners to convince myself. On the `?` early return: kernels `0..i` are owned by `udfs` and released through their `Arc<ImportedCScalarUdf>`; kernel `i` was moved into the `Box` passed to `try_new`, which drops it on the error path; kernels `i+1..len` are still live in the array, and `list`'s `Drop` calls `c_list_release`, which reconstructs the boxed slice and drops each one. The moved-out slots hold defaults with `release: None`, so they are no-ops rather than double frees. That said, the fact that it takes a paragraph to establish is your point on the other thread. A test forcing a failure mid-import would make it durable instead of incidental, and it needs a fixture library exporting a deliberately bad kernel, so I have folded it into #5250 alongside the move-semantics change. ########## 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: It was not always true, and this is the one comment in the review that turned out to be a live bug rather than a rough edge. `CometRustUDF.register` takes `deterministic` as a public parameter, defaulting to true, and calls `asNondeterministic()` on the Spark catalog stub when it is false. The value then travels in this proto field and is never read again: the planner ignores it, and `ImportedCScalarUdf::try_new` hardcodes `Volatility::Immutable`. So a UDF the caller explicitly declared nondeterministic was planned as pure, and DataFusion was free to fold it over constants, evaluate it once and reuse the result, or drop a repeated call as a common subexpression. Honoring the flag properly is not a one-liner, because the signature is built once per library load and cached process-wide, while determinism is declared per registration. Two registrations of the same kernel with different determinism would need different volatility out of one cached `ScalarUDFImpl`. That is #5249. For this PR, 3e999c45d makes `register` reject `deterministic = false` with an explicit not-yet-supported error, so the parameter cannot quietly lie about what Comet does with it, with a test. The field comment now records that it is always true today and why it is still carried on the wire. ########## 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: This is the right shape for it, and it would have paid for itself already: volatility (#5249) is exactly a property that today would need a new struct field and an ABI version bump, breaking every existing cdylib to add one boolean. Filed as #5254 with your sketch. Not taking it here because adding the field is itself an ABI change, so it wants to land before anyone depends on the current layout or ride along with the next bump, rather than being appended to a PR that is already large. Agreed the variable-length string return is the awkward part; noted the `FFI_ArrowArray` approach and that it may be heavier than needed. -- 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]
