andygrove commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3896259596
########## spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.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( Review Comment: Good question, and the three split differently. **`returnType`** is the closest to derivable, and in one sense it is already validated rather than trusted: the planner calls the kernel's `return_field` with the resolved argument types and rejects a declaration that disagrees, naming both types. So a wrong value fails rather than being believed. What blocks deriving it is ordering, not information. Spark's analyzer needs a concrete `DataType` at the moment `register` installs the catalog stub, and the argument types of the eventual call sites are not known then. Deriving it would mean calling the library on the driver with the `inputTypes` the caller supplied, which removes one declaration but not both. **`inputTypes`** cannot be derived as the ABI stands. `return_field` is a predicate over argument types, not a description of them: `echo_c` accepts every type and derives its output from whatever it is handed. Getting a signature list out of the library would need a new ABI entry point, which is an ABI break and deserves its own design rather than being folded in here. **`deterministic`** is genuinely a property of the kernel rather than of the registration, so it does belong on the library side. It is stuck behind honoring it at all: the native signature is built once per library load with `Volatility::Immutable`, so `register` currently rejects `false` outright rather than planning a volatile function as pure (#5249). Filed #5597 for the `returnType` overload, with the reasoning above written up. Worth noting there that even a derived type should keep the plan-time check, because derivation happens on the driver while the check runs on the executor against the library actually loaded there, and those need not be the same file. ########## 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: Update on this: the `Mutex` half is resolved (702f952d1, #5252), but not by changing how the kernel's lifetime is managed, so `Arc<CometCScalarKernel>` is still open on its own merits. What the lock turned out to be standing in for was an ABI guarantee rather than an aliasing problem. `new_impl` and `function_name` both take `*const CometCScalarKernel` and the exporter only reads through that pointer, so concurrent calls were already sound; the ABI just did not say so. It now does, and the adapter holds a plain `Box<CometCScalarKernel>`. That leaves your original question unchanged. `Box` keeps the two invariants I noted on #5252 (`release` runs exactly once via `Drop`, and the kernel cannot outlive the `LoadedLibrary` because the loader owns both), and an `Arc` would keep them too while making the second one enforceable rather than merely true by construction. It is no longer entangled with a performance problem, which makes it a smaller and more honest change than it was. -- 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]
