apache / apache/datafusion-comet

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

Open
#5,295 0 comments 0 reactions 0 assignees View on GitHub
area:expressions bug priority:critical
Dominant language
Scala
Stars
1.3k
Forks
373
Avg merge
2d 6h
Merged PRs (30d)
190

Description

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.

Contributor guide

Open the contributing guide

Research direction

Start with CometScalaUDF.convert, CometRustUdfRegistry, installCatalogStub, and the ignored case in CometRustUdfSuite. Trace how a session FunctionRegistry builder creates ScalaUDF and how arity checking behaves across the supported Spark profiles. Enable the ignored test; done means an ordinary Scala UDF call is not answered by the Rust UDF and the existing arity behavior remains covered.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
backend, data
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.