andygrove opened a new issue, #5295: URL: https://github.com/apache/datafusion-comet/issues/5295
Follow-up from review of #4459 ([thread](https://github.com/apache/datafusion-comet/pull/4459#discussion_r3730390223)), found by @mbutrovich. ## Problem `CometScalaUDF.convert` recognizes a Rust UDF by name alone: ```scala expr.udfName.flatMap(CometRustUdfRegistry.instance.get) match { case Some(meta) => emitRustUdfCall(...) ``` Spark sets `udfName` for every `spark.udf.register` call, not just those that went through `CometRustUDF.register`, and the registry is process-wide and keyed by bare name. So an ordinary Scala UDF registered under a name a Rust UDF already claimed is answered by the Rust library, silently returning the wrong values. The catalog stub's self-guard does not help: it only covers the UDF that originally claimed the name, and the other one is never evaluated on the JVM at all. ## Reproduction `CometRustUdfSuite` contains this as an `ignore`d test ("an ordinary Scala UDF is not answered by a Rust UDF of the same name"). Enabling it fails: ``` - 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 ``` That is `echo_c` registered as a Rust UDF, then `spark.udf.register("echo_c", (x: Long) => x * 10)`, then `SELECT echo_c(id)` returning `id` instead of `id * 10`. The plan-time return type check does not catch it, because the declared type happens to agree. ## Why the obvious fix does not work Having the registry hold the closure from the catalog stub and comparing it against `ScalaUDF.function` by identity was tried and does not work: - `functions.udf(f: UDF1[_, _], returnType: DataType)` wraps the `UDFn` it is handed, so the object in `ScalaUDF.function` is a closure Spark created rather than 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, which fails to compile on the 4.1 profile: ``` found : Any => Nothing required: org.apache.spark.sql.api.java.UDF1[_, _] ``` ## Candidate fix Skip `spark.udf.register` for the stub and register a builder directly in the session's `FunctionRegistry`, producing a `ScalaUDF` whose `function` Comet owns and can then recognize by identity: ```scala def builder(children: Seq[Expression]): Expression = ScalaUDF(ourStub, returnType, children, Seq.fill(children.size)(None), udfName = Some(name), ...) spark.sessionState.functionRegistry.createOrReplaceTempFunction(name, builder, "scala_udf") ``` `ScalaUDF`'s constructor has the same shape in Spark 3.4.3 and 4.1.1 (verified with `javap`), so it is portable across the supported profiles. Notes for whoever picks this up: - The builder owns arity checking, which `spark.udf.register` currently provides for free — the existing test "echo_c rejects a call whose argument count it does not accept" depends on it. - This also removes the arity-4 cap in `installCatalogStub`, since one builder handles any arity. - Enable the `ignore`d test with the fix, and consider the session-scoping issue at the same time: both are consequences of the registry being keyed and matched by bare name. -- 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]
