andygrove commented on code in PR #4459:
URL: https://github.com/apache/datafusion-comet/pull/4459#discussion_r3736237663
##########
spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala:
##########
@@ -53,8 +54,45 @@ 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] = {
+ // First check if this udfName is a registered Rust UDF -- those get
emitted as RustUdfCall
+ // and dispatched to the loaded cdylib rather than the JVM codegen
dispatcher.
+ expr.udfName.flatMap(CometRustUdfRegistry.instance.get) match {
Review Comment:
You were right to ask, and the answer is worse than a namespace clash: the
second registration is silently answered by the first's library. I wrote the
test you suggested and it fails today —
```
- an ordinary Scala UDF is not answered by a Rust UDF of the same name ***
FAILED ***
ArraySeq(0, 1, 2) did not equal List(0, 10, 20) the Scala UDF's call was
answered by the Rust UDF
```
`echo_c` registered as a Rust UDF, then `spark.udf.register("echo_c", (x:
Long) => x * 10)`, then `SELECT echo_c(id)` returns `id` rather than `id * 10`.
The serde matches `udfName` against the registry, gets a hit, emits
`RustUdfCall`, and the plan-time return-type check passes because the declared
type happens to agree. Nothing warns.
The test is in the suite as `ignore`d, with the observed failure recorded in
a comment, so it turns into a passing test with the fix rather than being
rediscovered.
On the fix: the obvious one does not work, which is worth writing down
because I tried it. I made the registry hold the closure from the catalog stub
and compared it against `ScalaUDF.function` by identity. It fails to compile on
Spark 4 and would not have worked anyway:
- `functions.udf(f: UDF1[_, _], returnType: DataType)` wraps the `UDFn` it
is handed, so the object that ends up in `ScalaUDF.function` is a closure Spark
created, not the one Comet passed in.
- The `udf(f: AnyRef, dataType: DataType)` overload that would have
preserved identity is gone in Spark 4 — with a pre-bound Scala function the
only applicable overloads are the Java `UDF0..UDF4` ones, hence `found: Any =>
Nothing, required: org.apache.spark.sql.api.java.UDF1[_, _]`.
What I think the fix is: skip `spark.udf.register` for the stub and register
a builder directly in the session's `FunctionRegistry` that emits a `ScalaUDF`
whose `function` Comet owns. `ScalaUDF`'s constructor is byte-for-byte the same
shape in 3.4.3 and 4.1.1, so that part is portable, and it would also lift the
arity-4 stub cap for free. It needs an arity check of its own and a run across
all the profiles, so I'd rather do it as a follow-up than bolt it on here.
Filing one.
##########
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)
Review Comment:
Fixed — the registry entry is published before the stub now. You had the
failure mode right: the stub is what makes the name resolvable to the analyzer,
so in the old order a query planned in that window found no registry entry,
fell through to the JVM codegen dispatcher, and hit the stub's "not evaluated"
exception.
--
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]
