andygrove commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3736239699
########## 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: Fair. The concrete one I can act on is the banner comments — `// -------- internals --------` here and `// ---------- type coverage ----------` in the suite. Neither style appears in any other `.scala` or `.java` file in the repo, so they're gone. If there's more you can point at, I'll take it — I'd rather match the surrounding code than argue about taste. ########## 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're right to double-check it, and the claim is false. Corrected the comment. The planner does `ScalarUDF::new_from_shared_impl(Arc::clone(&loaded.udf_impl))` against the process-wide cache in `cache.rs`, so every task in the executor shares one `ImportedCScalarUdf`, and concurrent Spark tasks are separate threads. Concurrent `invoke_with_args` on this instance is the normal case, not an unusual one. DataFusion does not serialize anything here — `ScalarUDFImpl: Send + Sync` exists precisely so it doesn't have to. So the lock is load-bearing, and the comment now says what it actually stands in for: Rust's aliasing rules would not require it (the callbacks reached from here take `*const CometCScalarKernel`, and per-batch mutable state lives in the `CometCScalarKernelImpl` each call builds), but ABI v1 does not require a kernel's `new_impl` to be callable concurrently, and the kernel is arbitrary user code. The cost is that all batches of a UDF serialize through one mutex per process, which is worth knowing before anyone runs a Rust UDF over a wide scan. That also changes #5252: its body repeats the "DataFusion serializes invocations anyway" rationale, so removing the lock on that basis would introduce a real data race. Removing it needs either a thread-safety requirement in the ABI or a kernel per task. I'll fix the issue body. -- 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]
