A Rust UDF silently answers calls to an ordinary Scala UDF registered under the same name #5295

Follow-up from review of #4459 (thread), found by @mbutrovich.

Problem

CometScalaUDF.convert recognizes a Rust UDF by name alone:

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 ignored 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:

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 ignored 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.