andygrove commented on code in PR #4459: URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r4188399417
########## native/core/src/execution/c_udf/cache.rs: ########## @@ -0,0 +1,90 @@ +// 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)); Review Comment: Fixed in 0043f3826. The canonical lookup is now a statement of its own, so the read guard is dropped before the write lock is taken. The new test looks a loaded library up through a fresh symlink on a separate thread and fails after 30 seconds instead of hanging the run. With the old `if let` it times out. This is most likely also why five of the `c_udf` tests hung for me under plain multi-threaded `cargo test` from a worktree under `/tmp`, which is a symlink on macOS: two threads racing on the first load can land in the same branch. nextest runs each test in its own process, which hid it in CI. ########## native/core/src/execution/planner.rs: ########## @@ -971,6 +971,81 @@ impl PhysicalPlanner { self.class_loader.clone(), ))) } + ExprStruct::NativeScalarUdf(call) => { + let arg_exprs: Vec<Arc<dyn PhysicalExpr>> = call + .args + .iter() + .map(|e| self.create_expr(e, Arc::clone(&input_schema))) + .collect::<Result<Vec<_>, _>>()?; + + let lib = crate::execution::c_udf::cache::get_or_load(&call.library_path).map_err( + |e| GeneralError(format!("native UDF load '{}': {e}", call.library_path)), + )?; + + let loaded = lib + .udfs + .iter() + .find(|u| u.name == call.name) + .ok_or_else(|| { + GeneralError(format!( + "native UDF '{}' not found in '{}'", + call.name, call.library_path + )) + })?; + + let udf = Arc::new(ScalarUDF::new_from_shared_impl(Arc::clone( + &loaded.udf_impl, + ))); + + let return_type = + to_arrow_datatype(call.return_type.as_ref().ok_or_else(|| { + GeneralError("NativeScalarUdf missing return_type".into()) + })?); + + // The declared return type comes from the JVM-side `CometNativeUDF.register` call + // and is what Spark planned against; the kernel's own `return_field` is what will + // actually be produced. If they disagree, fail here with both types named rather + // than letting it surface later as a bare type assertion mid-execution. + let arg_types = arg_exprs + .iter() + .map(|e| e.data_type(input_schema.as_ref())) + .collect::<Result<Vec<_>, _>>()?; + let kernel_return_type = loaded.udf_impl.return_type(&arg_types)?; + if !crate::execution::c_udf::return_types_compatible( + &return_type, + &kernel_return_type, + ) { + return Err(GeneralError(format!( + "native UDF '{}' was registered as returning {return_type} but its \ + return_field reports {kernel_return_type} for argument types {arg_types:?}. \ + Make the type passed to CometNativeUDF.register match what the UDF returns. \ + Note that a timestamp's timezone, a decimal's precision and scale, and \ + struct field names all have to match exactly; Spark's TimestampType is \ + Timestamp(Microsecond, Some(\"UTC\")) and TimestampNTZType is \ + Timestamp(Microsecond, None).", + call.name + ))); + } + + // Promise DataFusion the kernel's own type rather than the declared one, since the + // two can differ in nested nullability and the kernel's is what arrives. List and + // map child fields are renamed to Comet's canonical names, which is what every + // other expression producing that type uses, and the adapter relabels each result + // to match. + let return_field = Arc::new(Field::new( + &call.name, + crate::execution::c_udf::canonicalize_child_names(&kernel_return_type), Review Comment: Fixed in dc962d5a7. The planner now promises the kernel's type with every nested field nullable, except a map's entries and keys, which Arrow requires to be non-null, and the adapter conforms each result with a cast that only renames children and widens nullability. `make_struct_c` in the test library reports a non-nullable field, and new planner and adapter tests pin the widened type. Your `if` query is in the suite in both branch orders. Since #6458 landed on main, native IF reconciles its branches' nested nullability by itself, so that query passes even without this change. The planner change still matters for every other consumer of the UDF's type. ########## spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala: ########## @@ -55,8 +57,74 @@ 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] = { + // A registered native UDF is emitted as NativeScalarUdf and dispatched to the loaded shared + // library rather than to the JVM codegen dispatcher. + // + // The match is on the name alone, which is not enough to identify one: Spark sets `udfName` for + // every `spark.udf.register` call, and the registry is process-wide and keyed by bare name, so + // an ordinary Scala UDF sharing the name is currently answered out of the native library. The + // registration would have to be identified some other way to fix that, since the closure Spark + // holds for the catalog stub is one `functions.udf` wrapped rather than the one Comet passed + // in. See https://github.com/apache/datafusion-comet/issues/5295. + expr.udfName.flatMap(CometNativeUdfRegistry.get) match { + case Some(meta) => + emitNativeScalarUdf(expr, meta, inputs, binding) + case None => + emitJvmCodegenDispatch(expr, inputs, binding) + } + } + + private def emitNativeScalarUdf( + expr: ScalaUDF, + meta: NativeUdfMetadata, + inputs: Seq[Attribute], + binding: Boolean): Option[Expr] = { + val name = expr.udfName.get + checkArgumentTypes(name, expr, meta) + val argProtos = expr.children.map(c => exprToProtoInternal(c, inputs, binding)) + if (argProtos.exists(_.isEmpty)) { + withFallbackReason(expr, "one or more native UDF arguments are not supported") + return None + } + val returnTypeProto = serializeDataType(meta.returnType).getOrElse { + withFallbackReason(expr, s"return type ${meta.returnType} not serializable") + return None + } + val callBuilder = ExprOuterClass.NativeScalarUdf + .newBuilder() + .setName(name) + .setLibraryPath(meta.libraryPath) + .setReturnType(returnTypeProto) + .setDeterministic(expr.deterministic) + argProtos.foreach(a => callBuilder.addArgs(a.get)) + Some(ExprOuterClass.Expr.newBuilder().setNativeScalarUdf(callBuilder.build()).build()) + } + + /** + * Refuse a call whose argument types differ from the ones the UDF was registered with. + * + * The catalog stub Comet installs is untyped, so Spark inserts no casts for it and a call + * reaches this point with whatever types its arguments happen to have. Converting them here + * would be a semantic choice Spark never made, so the call is refused instead, naming both + * signatures. Nullability is disregarded because it does not change the values a UDF receives. + * + * This throws rather than falling back: the stub cannot evaluate the UDF on the JVM, so a + * fallback would only fail later with a less useful message. + */ + private def checkArgumentTypes(name: String, expr: ScalaUDF, meta: NativeUdfMetadata): Unit = { + val actual = expr.children.map(_.dataType) + val matches = actual.length == meta.inputTypes.length && + actual.zip(meta.inputTypes).forall { case (a, d) => deepNullable(a) == deepNullable(d) } Review Comment: Fixed in 74d7560d6. The check now uses `DataTypeSupport.equalsIgnoreNullability`, which re-derives Spark's `DataType.equalsIgnoreNullability` because that one is `private[sql]` on 3.4. It compares struct field names but neither nullability nor metadata. The new test passes a struct whose field carries a comment, built with `struct(col.as(name, metadata))`. On the old comparison it fails with `struct<a:int>` on both sides of the message. ########## spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala: ########## @@ -0,0 +1,168 @@ +/* + * 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 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 + +/** + * Entry point for registering scalar UDFs that run as native code inside Comet. + * + * A UDF lives in a shared library that exports the Comet UDF C ABI. That ABI is parameterized + * only by the Arrow C Data Interface and mentions no DataFusion or Rust types, so a compiled UDF + * is not tied to Comet's DataFusion version, and the host does not care what language produced + * the library. Only the Rust SDK (`comet-udf-sdk`) is supported and tested today; a C or C++ + * library that implemented the same ABI would load, but nothing ships a header for it. Hence the + * neutral name here, rather than one that promises Rust specifically. + * + * 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 CometNativeUDF { + + /** + * Register a single native UDF with an explicit signature. + * + * Validates the library on the driver (loads it, confirms a UDF named `name` exists). On + * success the driver-side registry is updated and a stub Spark catalog UDF is installed, in + * that order, so SQL/DataFrame name resolution succeeds only once the plan can be serialized as + * a `NativeScalarUdf`. + * + * Executors do not consult the driver's registry: the library path travels with the plan in the + * `NativeScalarUdf` proto, and each executor loads the library itself on first use. The path + * must therefore be valid on every executor, not just the driver. + * + * `inputTypes` is the signature every call must match. Comet does not convert arguments to + * these types: a call whose argument types differ, other than in nullability, is refused at + * planning time with both signatures named, so cast the arguments in the query instead. + * + * `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"native UDF '$name': deterministic = false is not supported yet. Comet plans native " + + "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";) + } + validateLibrary(libraryPath, name) + val meta = NativeUdfMetadata(libraryPath, inputTypes, returnType, deterministic) + CometNativeUdfRegistry.register(name, meta) + // Last, because this is the step that makes the name resolvable to Spark's analyzer. A query + // planned against a resolvable name that has no registry entry yet would route the call to the + // JVM codegen dispatcher and hit the stub's "not evaluated" exception, so the registry entry + // has to be in place first. + installCatalogStub(spark, name, inputTypes, returnType, deterministic) + } + + /** + * Load the library on the driver and confirm it exposes a UDF named `name`, translating the + * native failure into a typed exception. + */ + private def validateLibrary(libraryPath: String, name: String): Unit = { + try { + CometNativeUdfBridge.validateLibrary(libraryPath, name) + } catch { + case t: Throwable => throw classifyNativeError(libraryPath, t) + } + } + + /** + * Map a native loader failure onto a typed exception. + * + * The native side reports these as plain messages, so the mapping keys on the wording produced + * by `LoaderError`'s `Display` impl (`native/core/src/execution/c_udf/loader.rs`) and by + * `comet_native_udf_bridge.rs`. Each phrase below is matched in full rather than by a fragment + * like "ABI", because every one of those messages interpolates the library path: a library + * under a directory named `ABI` would otherwise have its "failed to open" reported as an ABI + * mismatch. `CometNativeUdfSuite` pins each failure mode to the type it produces here, so a + * reworded message on the native side fails a test rather than silently changing the exception + * a caller sees. + */ + private def classifyNativeError(libraryPath: String, t: Throwable): RuntimeException = { + val m = Option(t.getMessage).getOrElse("") + if (m.contains("missing required symbol") || m.contains("reports ABI v") || + m.contains("does not export")) { + new CometNativeUdfAbiException(m) + } else if (m.contains("' not found in ")) { + new java.util.NoSuchElementException(m) + } else { + new CometNativeUdfLoadException(s"failed to load $libraryPath: $m", t) + } + } + + /** + * Install a Spark catalog UDF under `name` so that SQL and DataFrame name resolution succeed. + * + * The stub only ever throws: a native UDF that reaches the JVM means Comet did not replace the + * expression with a native call. Note that the closure Spark keeps is not the one passed here, + * because `functions.udf` wraps the `UDFn` it is handed, which is why the serde cannot + * recognize this registration by identity and has to match on the name alone. See + * [[https://github.com/apache/datafusion-comet/issues/5295]]. + */ + private def installCatalogStub( + spark: SparkSession, + name: String, + inputTypes: Seq[DataType], + returnType: DataType, + deterministic: Boolean): Unit = { + val u: UserDefinedFunction = inputTypes.size match { + case 0 => + udf(() => throw new CometNativeUdfNotEvaluatedException(name), returnType) + case 1 => + udf((_: Any) => throw new CometNativeUdfNotEvaluatedException(name), returnType) + case 2 => + udf((_: Any, _: Any) => throw new CometNativeUdfNotEvaluatedException(name), returnType) + case 3 => + udf( + (_: Any, _: Any, _: Any) => throw new CometNativeUdfNotEvaluatedException(name), + returnType) + case 4 => + udf( + (_: Any, _: Any, _: Any, _: Any) => throw new CometNativeUdfNotEvaluatedException(name), + returnType) + case n => + throw new IllegalArgumentException( + s"native UDF '$name' arity $n not supported by stub. Reduce arity " + + "or open a feature request to extend stub coverage.") + } + val finalUdf = if (deterministic) u else u.asNondeterministic() + spark.udf.register(name, finalUdf) Review Comment: Fixed in f23556e1c with `val _ =` at both sites. `./mvnw test-compile -Pspark-3.5 -Pstrict-warnings` passes now, and the four errors in the untouched shuffle files are gone too. They only showed up because the fatal error stopped scalac early. -- 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]
