mbutrovich commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3730388880
########## spark/src/main/scala/org/apache/comet/udf/CometRustUdfRegistry.scala: ########## @@ -0,0 +1,53 @@ +/* + * 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 org.apache.comet.udf + +import java.util.concurrent.ConcurrentHashMap + +import org.apache.spark.sql.types.DataType + +/** Metadata for a registered Rust UDF. */ +case class RustUdfMetadata( + libraryPath: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean) + +/** + * Driver-side registry of Rust UDFs. Looked up by `QueryPlanSerde` to recognize names that should + * be emitted as `RustUdfCall` instead of attempted as JVM-evaluated `ScalaUDF`s. + */ +class CometRustUdfRegistry { + private val byName = new ConcurrentHashMap[String, RustUdfMetadata]() + + /** Register or replace metadata for a name. */ + def register(name: String, meta: RustUdfMetadata): Unit = + byName.put(name, meta) + + /** Return metadata for a name, if registered. */ + def get(name: String): Option[RustUdfMetadata] = + Option(byName.get(name)) +} + +object CometRustUdfRegistry { + + /** Process-wide singleton. */ + lazy val instance: CometRustUdfRegistry = new CometRustUdfRegistry Review Comment: The [contributor guide's "Global singletons" section](https://datafusion.apache.org/comet/contributor-guide/development.html#global-singletons) lists exactly this shape as a case to avoid: state that "depends on configuration that can vary between jobs or queries," with the risk being "cross-job contamination... a singleton initialized by the first job silently serves wrong state to subsequent jobs." `byName` is a process-wide `ConcurrentHashMap` keyed only by bare function name, with no session scoping and no unregister path, and the metadata it holds (library path, types) is exactly per-registration config that varies by caller. In a Spark Connect server, a notebook environment, or a test JVM running multiple suites, one session registering a Rust UDF called `"transform"` claims that name for every other session sharing the JVM, for the life of the process. The guide also asks that a singleton which is kept anyway carry a comment explaining why `static`/`object` is the right lif etime and whether the state is bounded; this one has neither. Compare with the native-side cache in `cache.rs`, which is also a process-wide singleton but is closer to the guide's "acceptable" column (effectively immutable per key once a path is loaded, and the file explains the lifetime choice in a comment). Should `CometRustUdfRegistry` be scoped to the session (or `SparkContext`) that performed the registration instead of being a bare JVM-wide singleton? ########## native/core/src/execution/rust_udf/imported_c.rs: ########## @@ -0,0 +1,311 @@ +// 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.) Review Comment: The comment justifying the `Mutex`'s scope says "DataFusion serializes invocations of a given `ScalarUDFImpl` per-batch through `invoke_with_args` anyway." Is that actually true for Comet's execution model, where concurrent Spark tasks run as separate threads in one executor process and can all reach the same cached UDF through the process-wide cache in `cache.rs`? If concurrent tasks can call the same registered Rust UDF at once, this claim seems worth double-checking before #5252 removes the lock, since it's the stated rationale for why the lock is "defensive" rather than load bearing. ########## native/core/src/execution/rust_udf/loader.rs: ########## @@ -0,0 +1,280 @@ +// 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, + /// One entry per UDF, with name and ScalarUDFImpl already built. + /// + /// Declared before `library` on purpose. Struct fields drop in + /// declaration order, and each UDF's drop calls a `release` callback + /// that lives in the library's text: unloading first would call + /// through a dangling pointer. + pub udfs: Vec<LoadedUdf>, + /// 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>, +} + +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() }) +} Review Comment: `read_abi_version` calls `sym()` directly with no `catch_unwind`, unlike every other user-code entry point in this PR (the `catch_panic`/`catch_panic_infallible` wrappers in `c_abi.rs`, and the `catch_unwind` in the discovery macro). Should the ABI-version probe get the same containment, given a panic here would unwind across the `extern "C"` boundary? ########## spark/src/test/scala/org/apache/comet/CometRustUdfSuite.scala: ########## Review Comment: Given the registry question above, a test that registers a Rust UDF and then an ordinary Scala UDF under the same name (or vice versa) would settle what currently happens. A concurrent-registration test (two threads racing `register()` for the same or different names) isn't present. Worth adding given the registry has no locking of its own beyond the underlying `ConcurrentHashMap`. `classifyNativeError` (`CometRustUDF.scala:108-118`) picks an exception type by substring matching on the native error message. Is there a test pinning each failure mode (missing file / bad ABI / name not found) to its expected exception type? A wording change on the native side could silently misclassify without anything catching it. ########## spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala: ########## @@ -53,8 +54,45 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen */ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { - override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = - emitJvmCodegenDispatch(expr, inputs, binding) + override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + // First check if this udfName is a registered Rust UDF -- those get emitted as RustUdfCall + // and dispatched to the loaded cdylib rather than the JVM codegen dispatcher. + expr.udfName.flatMap(CometRustUdfRegistry.instance.get) match { Review Comment: Following from the above: this checks every `ScalaUDF`'s `udfName` against the global registry, and Spark sets `udfName` for any ordinary `spark.udf.register(...)` call, not just ones that went through `CometRustUDF.register`. If two sessions in the same driver JVM register different functions under the same name (one Rust, one an ordinary Scala UDF), does the second registration silently reroute through the first's `RustUdfCall`? The catalog stub's self-guard only covers the UDF that originally claimed the name; the other one has no guard at all. Might be worth a test that registers an ordinary Scala UDF under a name a Rust UDF already claimed, to see what actually happens. ########## spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala: ########## @@ -0,0 +1,150 @@ +/* + * 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 org.apache.comet.udf + +import scala.util.Try + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.expressions.UserDefinedFunction +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.types.DataType + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +/** + * Entry point for registering Rust scalar UDFs with Comet. + * + * The UDF cdylib is built against the `comet-udf-sdk` crate and exposes its functions through an + * ABI built only on the Arrow C Data Interface, so a compiled UDF is not tied to Comet's + * DataFusion version. + * + * This is an experimental API. It is deliberately not annotated + * `org.apache.comet.annotation.Public`, so it sits outside the enumerated public API in Comet's + * [[https://datafusion.apache.org/comet/about/versioning_policy.html versioning policy]] and + * carries no compatibility guarantee: it may change or be removed in any release, including a + * patch release, with no deprecation cycle. + */ +object CometRustUDF { + + private val mapper: ObjectMapper = new ObjectMapper() + + /** + * Register a single Rust UDF with an explicit signature. + * + * Validates the library on the driver (loads it, confirms a UDF named `name` exists). On + * success a stub Spark catalog UDF is installed (so SQL/DataFrame name resolution succeeds) and + * the driver-side registry is updated. + * + * Executors do not consult the driver's registry: the library path travels with the plan in the + * `RustUdfCall` proto, and each executor loads the library itself on first use. The path must + * therefore be valid on every executor, not just the driver. + * + * `deterministic` must be `true`. Comet plans every imported kernel as immutable, so a + * nondeterministic UDF cannot yet be expressed; passing `false` fails here rather than silently + * planning the function as pure. + */ + def register( + spark: SparkSession, + name: String, + libraryPath: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean = true): Unit = { + if (!deterministic) { + // The native signature is built once per library load with + // Volatility::Immutable, while determinism is declared per registration, so the + // flag cannot be honored without reworking how kernels are cached. Until then a + // `false` here would let DataFusion constant-fold or CSE a call the user told us + // was not safe to reuse. + throw new IllegalArgumentException( + s"Rust UDF '$name': deterministic = false is not supported yet. Comet plans Rust UDFs " + + "as immutable, so a nondeterministic function may be constant-folded or eliminated " + + "as a common subexpression. See https://github.com/apache/datafusion-comet/issues/5249";) + } + val described = describeOne(libraryPath, name) + require(described.name == name, s"unexpected name from native: ${described.name}") + installCatalogStub(spark, name, inputTypes, returnType, deterministic) + val meta = RustUdfMetadata(libraryPath, inputTypes, returnType, deterministic) + CometRustUdfRegistry.instance.register(name, meta) + } + + // -------- internals -------- Review Comment: There are weird, very Claude-y comment styles that I don't really see in the codebase otherwise. ########## spark/src/main/scala/org/apache/comet/udf/CometRustUDF.scala: ########## @@ -0,0 +1,150 @@ +/* + * 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 org.apache.comet.udf + +import scala.util.Try + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.expressions.UserDefinedFunction +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.types.DataType + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode + +/** + * Entry point for registering Rust scalar UDFs with Comet. + * + * The UDF cdylib is built against the `comet-udf-sdk` crate and exposes its functions through an + * ABI built only on the Arrow C Data Interface, so a compiled UDF is not tied to Comet's + * DataFusion version. + * + * This is an experimental API. It is deliberately not annotated + * `org.apache.comet.annotation.Public`, so it sits outside the enumerated public API in Comet's + * [[https://datafusion.apache.org/comet/about/versioning_policy.html versioning policy]] and + * carries no compatibility guarantee: it may change or be removed in any release, including a + * patch release, with no deprecation cycle. + */ +object CometRustUDF { + + private val mapper: ObjectMapper = new ObjectMapper() + + /** + * Register a single Rust UDF with an explicit signature. + * + * Validates the library on the driver (loads it, confirms a UDF named `name` exists). On + * success a stub Spark catalog UDF is installed (so SQL/DataFrame name resolution succeeds) and + * the driver-side registry is updated. + * + * Executors do not consult the driver's registry: the library path travels with the plan in the + * `RustUdfCall` proto, and each executor loads the library itself on first use. The path must + * therefore be valid on every executor, not just the driver. + * + * `deterministic` must be `true`. Comet plans every imported kernel as immutable, so a + * nondeterministic UDF cannot yet be expressed; passing `false` fails here rather than silently + * planning the function as pure. + */ + def register( + spark: SparkSession, + name: String, + libraryPath: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean = true): Unit = { + if (!deterministic) { + // The native signature is built once per library load with + // Volatility::Immutable, while determinism is declared per registration, so the + // flag cannot be honored without reworking how kernels are cached. Until then a + // `false` here would let DataFusion constant-fold or CSE a call the user told us + // was not safe to reuse. + throw new IllegalArgumentException( + s"Rust UDF '$name': deterministic = false is not supported yet. Comet plans Rust UDFs " + + "as immutable, so a nondeterministic function may be constant-folded or eliminated " + + "as a common subexpression. See https://github.com/apache/datafusion-comet/issues/5249";) + } + val described = describeOne(libraryPath, name) + require(described.name == name, s"unexpected name from native: ${described.name}") + installCatalogStub(spark, name, inputTypes, returnType, deterministic) + val meta = RustUdfMetadata(libraryPath, inputTypes, returnType, deterministic) + CometRustUdfRegistry.instance.register(name, meta) Review Comment: `installCatalogStub` runs before `CometRustUdfRegistry.instance.register(...)`. Doesn't that open a window where the name is resolvable by Spark's analyzer but not yet present in the registry a concurrent query would need to plan it as `RustUdfCall`? Reordering these two calls looks like it would close it. ########## native/core/src/execution/rust_udf/imported_c.rs: ########## @@ -0,0 +1,311 @@ +// 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. + // + // Volatility is always Immutable. The signature is built once per + // library load, while determinism is declared per registration, so + // the two do not line up: `CometRustUDF.register` rejects + // `deterministic = false` rather than let a volatile function be + // planned as if it were pure. + 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(), Review Comment: Both `init_cb` calls pass `std::ptr::null()` for `scalar_args`, and any `ColumnarValue::Scalar` is unconditionally expanded to a full array before this point. The ABI docs describe `scalar_args` as letting a kernel specialize on bound literals. Since nothing in the host ever populates it, is this parameter reachable by any kernel today, or is it effectively dead on the host side until something wires it up? ########## native/core/src/comet_rust_udf_bridge.rs: ########## @@ -0,0 +1,89 @@ +// 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. + +//! JNI entry points for driver-side validation of Rust UDF cdylibs. +//! Used by `org.apache.comet.udf.CometRustUdfBridge` on the driver. + +use crate::errors::{try_unwrap_or_throw, CometError}; +use crate::execution::rust_udf::cache::get_or_load; +use crate::execution::rust_udf::loader::LoadedUdf; +use jni::objects::{JClass, JString}; +use jni::sys::jobject; +use jni::EnvUnowned; + +/// Best-effort serialization of a single discovered UDF as JSON. +/// +/// Reporting `args`/`return_type` would require calling the kernel's +/// `init` to discover a return type; that is deferred, so this returns +/// only what the Scala registry needs today (`name`). +fn udf_to_json(udf: &LoadedUdf) -> serde_json::Value { + serde_json::json!({ + "name": udf.name, + }) +} Review Comment: `validateLibrary` already filters `lib.udfs` by `u.name == name` and errors if it's absent, so `udf_to_json` can only ever return `{"name": "<the name that was passed in>"}`. `describeOne` / `parseDescribed` then round-trips that through Jackson just to check `described.name == name`, which can't fail given the native-side filter already ran. Could this JSON construct/parse cycle be dropped, since `validateLibrary` succeeding is already the confirmation? ########## native/core/src/execution/rust_udf/imported_c.rs: ########## @@ -0,0 +1,311 @@ +// 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. + // + // Volatility is always Immutable. The signature is built once per + // library load, while determinism is declared per registration, so + // the two do not line up: `CometRustUDF.register` rejects + // `deterministic = false` rather than let a volatile function be + // planned as if it were pure. + 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: `invoke_with_args` builds a fresh `CometCScalarKernelImpl` and calls `init` (re-encoding every arg `Field` to `FFI_ArrowSchema` and re-invoking the kernel's `return_field`) on every batch. `return_type()` already does the same work once at planning time, and argument types don't change across batches for a given `ScalarFunctionExpr`. Could the resolved return `Field` (or the whole `impl_state`) be cached on the adapter and reused across batches instead of recomputed each time? ########## 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(); + + if let Some(lib) = cache().read().unwrap().get(&raw).cloned() { + return Ok(lib); + } + + let canonical = raw.canonicalize().unwrap_or_else(|_| raw.clone()); + if canonical != raw { + if let Some(lib) = cache().read().unwrap().get(&canonical).cloned() { + cache().write().unwrap().insert(raw, Arc::clone(&lib)); + return Ok(lib); + } + } + + let mut w = cache().write().unwrap(); + if let Some(lib) = w.get(&canonical).cloned() { + if canonical != raw { + w.insert(raw, Arc::clone(&lib)); + } + return Ok(lib); + } + let loaded = Arc::new(load(&canonical)?); + w.insert(canonical.clone(), Arc::clone(&loaded)); + if canonical != raw { + w.insert(raw, Arc::clone(&loaded)); + } Review Comment: The write lock is held for the duration of `load(&canonical)`, which runs the cdylib's static initializers and discovery routine. Since the cache is one process-wide `RwLock` rather than sharded per path, would a slow (or, per the point above, panicking) load of one library block `get_or_load` for every unrelated library path in the process for as long as that load takes? ########## native/core/src/comet_rust_udf_bridge.rs: ########## @@ -0,0 +1,89 @@ +// 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. + +//! JNI entry points for driver-side validation of Rust UDF cdylibs. +//! Used by `org.apache.comet.udf.CometRustUdfBridge` on the driver. + +use crate::errors::{try_unwrap_or_throw, CometError}; +use crate::execution::rust_udf::cache::get_or_load; +use crate::execution::rust_udf::loader::LoadedUdf; +use jni::objects::{JClass, JString}; +use jni::sys::jobject; +use jni::EnvUnowned; + +/// Best-effort serialization of a single discovered UDF as JSON. +/// +/// Reporting `args`/`return_type` would require calling the kernel's +/// `init` to discover a return type; that is deferred, so this returns +/// only what the Scala registry needs today (`name`). +fn udf_to_json(udf: &LoadedUdf) -> serde_json::Value { + serde_json::json!({ + "name": udf.name, + }) +} + +/// Validate that `library_path` loads, exposes a UDF named +/// `expected_name`, and return a JSON description of that UDF. Throws +/// on any error. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_udf_CometRustUdfBridge_validateLibrary( + e: EnvUnowned, + _class: JClass, + library_path: JString, + expected_name: JString, +) -> jobject { + try_unwrap_or_throw(&e, |env| { + let path: String = library_path + .try_to_string(env) + .map_err(|e| CometError::Internal(e.to_string()))?; + let name: String = expected_name + .try_to_string(env) + .map_err(|e| CometError::Internal(e.to_string()))?; + let lib = get_or_load(&path).map_err(|e| CometError::Internal(e.to_string()))?; + let udf = lib + .udfs + .iter() + .find(|u| u.name == name) + .ok_or_else(|| CometError::Internal(format!("UDF '{name}' not found in {path}")))?; + let json = udf_to_json(udf).to_string(); + let jstr = env + .new_string(json) + .map_err(|e| CometError::Internal(e.to_string()))?; + Ok(jstr.into_raw()) + }) +} + +/// Return a JSON array describing every UDF exposed by `library_path`. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_udf_CometRustUdfBridge_listUdfs( + e: EnvUnowned, + _class: JClass, + library_path: JString, +) -> jobject { + try_unwrap_or_throw(&e, |env| { + let path: String = library_path + .try_to_string(env) + .map_err(|e| CometError::Internal(e.to_string()))?; + let lib = get_or_load(&path).map_err(|e| CometError::Internal(e.to_string()))?; + let entries: Vec<serde_json::Value> = lib.udfs.iter().map(udf_to_json).collect(); + let json = serde_json::Value::Array(entries).to_string(); + let jstr = env + .new_string(json) + .map_err(|e| CometError::Internal(e.to_string()))?; + Ok(jstr.into_raw()) + }) +} Review Comment: `listUdfs` (and its Java declaration at `CometRustUdfBridge.java:40`) doesn't appear to have a caller anywhere in the diff. Is this meant to be wired up in a follow-up, or is it dead code that should come out? -- 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]
