Repository navigation
Remove the per-UDF Mutex serializing Rust UDF batches #5252
Description
Follow-up from review of #4459 (thread).
ImportedCScalarUdf holds its kernel as Mutex<Box<CometCScalarKernel>>. Every return_type and invoke_with_args takes that lock, so all concurrent batches for a given UDF in a process serialize on one mutex, even though the kernel is logically immutable after load and the per-execution state already lives in a separate CometCScalarKernelImpl built fresh per call.
The lock is defensive rather than required: the comment notes DataFusion serializes invocations of a given ScalarUDFImpl anyway, and the FFI Drop is what is not Sync-safe.
@paleolimbot suggested Arc<CometCScalarKernel> would maintain the reference counts and release the instance correctly when the last reference goes. Worth evaluating, along with simply holding Box<CometCScalarKernel> directly given the struct is already Send + Sync and the kernel is only read after load.
Whatever shape this takes, it should preserve the property that the kernel's release runs exactly once, and it should not outlive the LoadedLibrary that dlopened it (see the field ordering in LoadedLibrary).
Activity
review: address mbutrovich's first pass on the Rust UDF path
andygrove commentedon Aug 7, 2026
Correcting the rationale in the description, per @mbutrovich in review of #4459.
The description says the lock is "defensive rather than required", on the grounds that DataFusion serializes invocations of a given
ScalarUDFImplanyway. That is wrong. DataFusion does not serialize anything here —ScalarUDFImpl: Send + Syncexists so it does not have to — andplanner.rsbuilds every task'sScalarFunctionExprwithScalarUDF::new_from_shared_implover theArcfrom the process-wide cache incache.rs. Concurrent Spark tasks in an executor are separate threads sharing oneImportedCScalarUdf, so concurrentinvoke_with_argson it is the normal case, not an unusual one.Simply removing the lock would introduce a data race.
What the lock actually stands in for is an ABI guarantee that does not exist yet. Rust's aliasing rules would not require it, since the callbacks reached through it (
function_name,new_impl) take*const CometCScalarKerneland per-batch mutable state lives in theCometCScalarKernelImpleach call builds for itself. But ABI v1 does not require a kernel'snew_implto be callable concurrently, and the kernel is arbitrary user code.So this needs one of:
Arc<CometCScalarKernel>as @paleolimbot suggested, or plainBoxgiven the struct is alreadySend + Sync.The consequence today, which is also worth stating in this issue: all batches of a given Rust UDF serialize through one mutex per executor process, so a Rust UDF over a wide scan does not scale with task parallelism.
The comment on the
kernelfield inimported_c.rshas been corrected in #4459 and now describes the above.feat: custom Rust UDFs via arrow-ffi [experimental] (apache#4459)