andygrove commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3896257283
########## 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: You were right, and the lock is gone (702f952d1). `new_impl` and `function_name` are already `*const CometCScalarKernel`, and the exporter only reads through that pointer: `c_factory_new_impl` does an `Arc::clone` of the user's `inner` and writes a fresh `CometCScalarKernelImpl` into the caller's `out`, so nothing mutable in the kernel is touched. `CometCScalarUdf: Send + Sync` was already the bound. So the property held; what was missing was that the ABI said so, which is why I had written the lock as standing in for a guarantee that did not exist. So rather than only deleting it I have made the requirement explicit, since a kernel written against a permissive host and then run without a lock is the bad outcome: - a `# Thread safety` section on `CometCScalarKernel` saying both callbacks must be safe to call concurrently on the same kernel, and that the host relies on this because it shares one kernel across every task in an executor - the same note on the `new_impl` field, since that is where an implementer looks - a `# Thread safety` section on the `CometCScalarUdf` trait explaining what the `Send + Sync` bound is buying and that per-batch scratch belongs in `invoke`, not in the struct - the same in the user guide, next to `comet_c_udf_export!` Concurrency stops at the kernel: the `CometCScalarKernelImpl` a `new_impl` call produces is still single-threaded and is used only by the thread that asked for it, which is documented as such. Pinned by `concurrent_invocations_share_one_adapter`, which drives eight threads through one shared `ImportedCScalarUdf` on distinct batches, so a kernel that leaked state between concurrent calls returns another thread's answers. The GIL comparison is the right way to put it. This also closes #5252, which is now on the PR. ########## 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 { Review Comment: Done in f06dd7230, in the shape you wrote it: `register` and `get` forward to the singleton, and `instance` is now `private` since those were the only two call sites (`CometNativeUDF.register` and `CometScalaUDF.convert`). The justification comment stays on the `lazy val` rather than moving to the object, because what needs justifying is the singleton's lifetime, not the forwarding. -- 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]
