apache / apache/datafusion-comet

Scope the Rust UDF registry to the session that registered the UDF

Open
#5,294 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Scala
Stars
1.3k
Forks
375
Avg merge
2d 4h
Merged PRs (30d)
198

Description

Follow-up from review of #4459 ([thread](https://github.com/apache/datafusion-comet/pull/4459#discussion_r3730388880)).

`CometRustUdfRegistry` is a process-wide `ConcurrentHashMap[String, RustUdfMetadata]` keyed by bare function name, with no session scoping and no unregister path. As @mbutrovich points out, the [contributor guide's "Global singletons" section](https://datafusion.apache.org/comet/contributor-guide/development.html#global-singletons) lists this shape as one to avoid: the metadata it holds (library path, declared input and return types) is per-registration configuration that varies by caller, which is the guide's "state depends on configuration that can vary between jobs or queries" case.

Concretely, in any JVM hosting more than one session — a Spark Connect server, a notebook, a test JVM running several suites — one session registering a Rust UDF named `transform` claims that name for every other session for the life of the process, and the last registration of a name wins for all of them.

Registration should be scoped to the session (or `SparkContext`) that performed it, so that a name registered in one session is not visible to another, and entries go away with the session rather than living for the life of the JVM.

This is a driver-side concern only: executors never consult the registry, since the library path travels with the plan in the `RustUdfCall` proto.

Related to the name-matching half of the same problem (an ordinary Scala UDF registered under a Rust UDF's name is answered out of the Rust library). A fix that gives each registration an identity Comet can recognize may address both.

`CometRustUdfRegistry.instance` currently carries a comment explaining the lifetime and bounds, per the guide's requirement for a singleton that is kept anyway; that comment should come out with the fix.

Contributor guide

Open the contributing guide

Research direction

Start at CometRustUdfRegistry.instance and trace registration and lookup through the driver-side Rust UDF path; use the contributor guide's Global singletons section as the lifetime constraint. Check how RustUdfCall carries library metadata, then verify that registrations are isolated between sessions and removed with their owner, including conflicting names across sessions.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust, scala
Domain
backend, data-engineering
Issue type
Refactor
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.