sunchao commented on code in PR #4459:
URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r4104678388
##########
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:
[P2] Ignore incidental struct-field metadata when checking argument types.
`deepNullable` preserves `StructField.metadata`, so equality rejects a struct
column whose field carries a comment when the registered signature has the same
field name and SQL type without that comment. Both signatures render as
`struct<a:int>`, yet the call throws `CometNativeUdfArgumentTypeException`.
This prevents otherwise supported UDF calls over annotated nested schemas. Use
Spark's `DataType.equalsIgnoreNullability` or an equivalent structural
comparison that ignores metadata.
Evidence: A Scala reproduction using Spark 4.1.3, the unchanged
`deepNullable` implementation, and this comparison tested identical
`struct<a:int>` types differing only by field metadata `{"comment":"column
documentation"}`. The PR comparison returned false, while
`DataType.equalsIgnoreNullability` returned true. Both `catalogString` values
were identical.
##########
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:
[P2] Normalize nested nullability when exposing the UDF result to the native
plan. A valid kernel can return `Struct<a: Int32 non-null>` while its Spark
registration declares nullable `a`. The compatibility check accepts this, but
this line preserves the narrower kernel type. Placing that UDF first in
`if(condition, udf(id), other_struct)` makes DataFusion attempt to narrow the
other branch and fail with `Cannot cast nullable struct field 'a' to
non-nullable field`. Spark accepts the query. Expose a consistently nullable
nested type and safely conform the result, for example through a widening cast.
Evidence: A focused harness used the unchanged SDK, adapter,
canonicalization, and Comet `IfExpr` with a kernel producing non-nullable `a`.
Type validation passed, but the UDF-first branch failed with the quoted error.
Reversing branches succeeded. Widening the UDF expression's nested nullability
also succeeded. Spark 4.1.3 evaluated `if(id > 0, nn_struct(id),
named_struct('a', cast(id as int)))` over `range(2)` successfully, returning
structs containing 0 and 7.
##########
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:
[P2] Release the read guard before acquiring the write lock. After a library
is loaded through its canonical path, requesting it through a previously unused
symlink reaches this branch. With the workspace's Rust 2021 edition, the read
guard created by the `if let` condition remains alive inside its body.
Consequently, `cache().write()` waits forever for this same thread's read lock.
Instead of reusing the cached library, driver registration or executor planning
hangs. Store the cloned lookup result in a separate, completed scope before
acquiring the write lock.
Evidence: A bounded harness using the unchanged loader/cache and freshly
built `libcomet_test_udfs.so` completed canonical-path and repeated-path
lookups, then timed out after three seconds on a symlink lookup. A separate
cache control that explicitly dropped the read guard completed all three
lookups. This occurs before `load()`, independently of the already-discussed
lock held across `dlopen`.
--
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]
