feat: custom Rust UDFs via arrow-ffi [experimental]

[andygrove]

Which issue does this PR close?

Closes #747

Rationale for this change

Comet has no way to run a user's own native Rust code as part of a native plan. Today a user-defined function either falls back to Spark or, at best, runs through the JVM codegen dispatcher, which pulls rows back across the JNI boundary.

This PR adds custom Rust scalar UDF support. A user compiles a cdylib against a published SDK crate, points Comet at it, and the function executes natively inside the DataFusion plan, operating directly on Arrow arrays.

Why a bespoke C ABI rather than datafusion-ffi

Wrapping the user's ScalarUDFImpl as datafusion_ffi::udf::FFI_ScalarUDF is the obvious alternative and hands the author a much larger surface for free: variadic signatures, type coercion, metadata-aware return types. This PR deliberately does not expose it, because it would couple every user's cdylib to Comet's DataFusion major version.

That is not a hypothetical cost. An earlier revision of this PR carried both ABIs side by side. Rebasing onto current main moved Comet from DataFusion 53 to 54, which removed as_any from ScalarUDFImpl: every user library built against the datafusion-ffi flavor would have needed a source edit and a recompile. The C ABI needed no change, because no DataFusion type appears in it.

Comet tracks DataFusion closely and upgrades often, so a UDF ABI pinned to the DataFusion version would break users on a cadence they neither control nor can opt out of. Keeping the FFI surface to Arrow's C Data Interface means a UDF compiled today keeps working across Comet upgrades, and the same ABI is implementable from C, C++, or any language that speaks the Arrow C Data Interface.

The tradeoff is a smaller surface: authors implement CometCScalarUdf and get scalar functions over Arrow arrays, not the full ScalarUDFImpl feature set. The rationale is recorded in the comet-udf-sdk crate docs so it does not have to be rediscovered.

What changes are included in this PR?

• native/comet-udf-sdk — the public SDK users compile against. Provides the CometCScalarUdf trait and the comet_c_udf_export! macro. Depends only on arrow, with no DataFusion dependency at all.
• native/comet-test-udfs — test cdylib: add_one_c plus two deliberately panicking UDFs used to verify panic containment.
• native/core/src/execution/rust_udf — loader (libloading + ABI version probe + discovery), process-wide cache, and the ImportedCScalarUdf adapter that wraps a C-ABI kernel as a ScalarUDFImpl.
• RustUdfCall proto in expr.proto + planner branch in create_expr resolving (library_path, name) against the cache.
• JNI bridge (CometRustUdfBridge / comet_rust_udf_bridge.rs) for driver-side validateLibrary / listUdfs.
• Scala API (CometRustUDF.register, CometRustUdfRegistry, typed exception classes). In QueryPlanSerde, CometScalaUDF.convert first checks the registry: if the udfName is registered it emits RustUdfCall, otherwise it falls through to the existing JVM codegen dispatcher.
• All Spark types supported, argument and return: the non-nested types (Boolean, Byte, Short, Integer, Long, Float, Double, Decimal, String, Binary, Date, Timestamp, TimestampNTZ) plus ArrayType, MapType and StructType with arbitrary nesting, nulls preserved in both directions. No ABI change was needed, since the FFI surface is type-agnostic and already carries child arrays, but nothing had demonstrated it beyond Int64.
• Plan-time validation of the declared return type against what the kernel's return_field reports, naming both types. Registering decimal(10,2) for a column Spark had widened to decimal(11,2) previously surfaced as a bare DataFusion type assertion partway through execution. The comparison erases nested nullability, because Spark carries containsNull inside the declared type while the delivered array normalizes child fields to nullable; it stays strict about everything that changes how bytes are read.
• Panic containment at every extern "C" entry point. A panic escaping one of these aborts the process, which for Comet means killing the executor JVM and losing every task on it, not just the offending query. User UDF code is arbitrary and panicking is idiomatic Rust (unwrap, slice indexing, debug overflow), so a panic is now caught at the boundary and reported through the same get_last_error channel as a returned Err. The query fails; the executor survives.
• CI coverage. CometRustUdfSuite now runs in both the Linux and macOS PR builds. The suite locates the cdylib under native/target rather than requiring a flag, and both native build jobs upload libcomet_test_udfs alongside libcomet. If the library is ever missing the suite skips locally but fails in CI, so it cannot silently stop running again.
• User guide page (docs/source/user-guide/latest/rust_udfs.md), which describes the feature as experimental and documents the current limitations, including that the library must already exist on every executor at the given path, and that loading a UDF library means running untrusted native code in the executor process.

There is no impact on existing users: every new code path is reachable only after an explicit CometRustUDF.register call. With no registered UDFs, CometScalaUDF.convert takes exactly the branch it takes today, no library is loaded, and no new configuration is introduced.

How are these changes tested?

• SDK unit tests — ABI version, C ABI adapter roundtrip (export → import via FFI), and panic containment in both return_field and invoke, asserting the error code and that the panic message reaches get_last_error.
• Native loader tests — library load, missing-path error, and discovery of every exported kernel.
• End-to-end Spark suite (CometRustUdfSuite), 44 tests — add_one_c over a range; echo_c and stringify_c over all 19 supported types with a null row, nested types included; one test per panic site asserting the query fails with the panic message rather than the executor dying; and the declared-vs-actual return type mismatch error. Gated on -Dcomet.test.udfs.lib=<path>.

The e2e tests are self-guarding on native execution: register installs a catalog stub that throws if Spark ever evaluates the UDF itself, so a silent fallback to Spark fails the test rather than passing. This caught a real fallback during development, where an unsupported expression in the test data pulled the whole projection back to Spark.

$ cargo test -p comet-udf-sdk
test tests::abi_version_is_one ... ok
test c_abi::tests::adapter_roundtrip ... ok
test c_abi::tests::panic_in_return_field_is_contained ... ok
test c_abi::tests::panic_in_invoke_is_contained ... ok
test result: ok. 4 passed; 0 failed

$ DYLD_LIBRARY_PATH=$JAVA_HOME/lib/server cargo test -p datafusion-comet --lib rust_udf
test result: ok. 5 passed; 0 failed

$ ./mvnw -q test -Dsuites="org.apache.comet.CometRustUdfSuite" -Dtest=none \
    -Dcomet.test.udfs.lib=$PWD/native/target/debug/libcomet_test_udfs.dylib
- add_one_c returns id + 1 for a range
- panic inside UDF invoke surfaces as a query error, not a crash
- panic inside UDF return_field surfaces as a query error, not a crash
- echo_c round-trips boolean including nulls
  ... one echo_c + one stringify_c test per supported type, including array/map/struct ...
- echo_c round-trips decimal(10,2) including nulls
- stringify_c reads decimal(10,2) values
- one kernel computes its return type on demand and serves many types
- a declared return type that disagrees with the UDF names both types
- echo_c rejects a call whose argument count it does not accept
Tests: succeeded 44, failed 0

Follow-on work

Known gaps, deliberately out of scope here and tracked for follow-ups:

• Remaining test coverage: empty batches, multi-argument calls, and varied array encodings.
• Arrow extension types (Variant, Geometry, Geography), which as noted in review is a broader expr.proto concern rather than something specific to this path.
• No config key gates library loading, so an operator cannot disable the feature or restrict which paths may be loaded.
• No mechanism distributes the library to executors; the path must already be valid cluster-wide.
• deterministic is carried in the proto but not yet honored natively, and the planner hardcodes the return field as nullable.
• The per-UDF Mutex in ImportedCScalarUdf serializes batches through one lock per UDF per process.
• Arity is capped at 4 by the catalog stub, and registerAll is not implemented.

[paleolimbot]

Thank you for putting this together!

[paleolimbot]

I see the motivation for having this primarily live as shared objects! I am not sure if there's an opportunity / whether it's any easier to bundle these UDFs as .jars or Python packages but that could be implemented as a future field of this struct.

[andygrove]

Agreed. The proto struct is intentionally open here so a future variant can carry a RegistrationSource oneof (SharedObject | Jar | PythonPackage) without breaking wire compat. The shared-object form is the shortest path for the Rust ecosystem; the jar / Python paths would fan out from the same registration message.

[paleolimbot]

Is a fixed return type a requirement? (DataFusion lets you compute this on demand but perhaps there's a limiation here)

[andygrove]

Not strictly. This PR pins it because the current dispatch resolves the return type at plan time (so we can allocate the output vector once per batch) and Spark expects a fixed DataType for scalar UDFs. Widening to a compute-on-demand return type would let a single UDF handle polymorphic outputs, at the cost of a per-batch call across the FFI to resolve. If that becomes valuable I can follow up with an optional return_type_fn slot.

[paleolimbot]

The other issue with the definition is that it can't handle an extension type, and some Iceberg/Spark types do map to Arrow extension types (notably: Geometry, Geography, and Variant). That is a larger issue with many of the definitions in this file and is probably better suited to a standalone change.

[andygrove]

Revisiting this now that the ABI has settled, because my earlier answer undersold what it does.

The return type is computed on demand. CometCScalarUdf::return_field(&self, args: &[Field]) is called with the actual argument types and derives the output type per call site, so a kernel has no fixed return type of its own. The test library now leans on that: a single echo_c kernel serves all 19 types in the suite, from Boolean through decimal(10,2) to array<struct<a:int>>, returning whatever it is handed.

What is fixed is the type declared to CometRustUDF.register, and that is a Spark constraint rather than an ABI one: Spark needs a concrete DataType at analysis time to plan the query and to install the catalog entry. So the declaration is per-registration, not per-kernel. Re-registering the same function under a different signature works, and the kernel computes the matching return type each time. There is a test covering exactly that.

The proto field is therefore best read as "what Spark was told", not "the only type this UDF can return". Comet now checks the two against each other at planning time and fails with both types named if they disagree, which replaced a bare DataFusion type assertion that used to fire partway through execution. That check turned out to earn its keep immediately: Spark widens decimal(10,2) + 0.25 to decimal(11,2), and registering the un-widened type is an easy mistake to make.

One wrinkle worth recording, since it also shaped the check. Spark carries containsNull and struct field nullability inside the type, so array<int> with containsNull = false converts to a List with a non-nullable child, while the array actually delivered to the UDF has that child normalized to nullable. The comparison erases nested nullability for that reason, but stays strict about everything that changes how bytes are read: decimal precision and scale, timestamp unit, and struct field names and order.

On extension types: agreed, and I have not tried to address it here. Variant and Geometry would need the type representation in expr.proto to grow an extension-type concept, which affects far more than this path, so it seems right as a standalone change. I have noted it in the PR's follow-on list so it does not get lost.

[paleolimbot]

Cool!

[paleolimbot]

Also cool! (And hard to argue with the compact implementation!)

[andygrove]

Thanks for the review @paleolimbot. The PR has changed quite a bit since you looked, so here is a summary of where it stands. It is also rebased onto current main, which turned out to be relevant to the biggest change.

The datafusion-ffi flavor is gone

This is the one I want to flag directly, because you were positive about that implementation and it is the thing I removed. Carrying two public ABIs would have meant two permanent compatibility promises, so I picked one, and the rebase is what decided it.

Merging current main moved Comet from DataFusion 53 to 54, which removed as_any from ScalarUDFImpl. Every user library built against the datafusion-ffi flavor would have needed a source edit and a recompile. The C ABI needed no change at all, because no DataFusion type appears in it. Comet tracks DataFusion closely and upgrades often, so that break would recur on a cadence UDF authors neither control nor can opt out of, and it would land on whoever upgrades Comet rather than on us.

So the surviving ABI is the arrow-only one. comet-udf-sdk now has no DataFusion dependency at all: a user's cdylib pulls in arrow and nothing else. The tradeoff we accept is the smaller authoring surface, and if that becomes the binding constraint the answer is to widen the C ABI deliberately rather than adopt a version-pinned one. The rationale is recorded in the crate docs so it does not have to be rediscovered.

Panic containment at every FFI entry point

Previously only execute was guarded. A panic escaping any of the other extern "C" functions aborts the process, which for Comet means killing the executor JVM and losing every task on it, not just the query that used the UDF. Since user UDF code is arbitrary and panicking is idiomatic Rust, panics are now caught at the boundary and reported through the same get_last_error channel as a returned Err. The query fails; the executor survives. Covered by SDK unit tests and by end-to-end tests that drive genuinely panicking UDFs through Spark.

Full type coverage, including complex types

Arguments and return values now cover every non-nested Spark type plus ArrayType, MapType and StructType with arbitrary nesting, nulls preserved in both directions. No ABI change was needed for any of it, since the FFI surface is type-agnostic and already carries child arrays, but nothing had demonstrated it beyond Int64.

Two test UDFs cover the matrix: echo_c returns its argument unchanged, which checks that each type survives the round trip with its nulls, and stringify_c renders any input as strings, which forces the UDF to actually decode the values rather than hand the array straight back.

Return type validation

Comet now checks the type declared to CometRustUDF.register against what the kernel's return_field reports, at planning time, naming both types when they disagree. This replaced a bare DataFusion type assertion that used to fire partway through execution. More on the semantics in my reply on the return_type thread.

User guide

Added docs/source/user-guide/latest/rust_udfs.md, which describes the feature as experimental and is explicit about the current limitations: the library must already exist on every executor at the given path since Comet does not distribute it, and loading a UDF library means running unsandboxed native code in the executor process.

Still outstanding

Listed in the PR description under "Follow-on work". The ones I would most want a reviewer's opinion on: the suite is not yet wired into CI, so the ABI has never been exercised against a Linux .so; there is no config key gating library loading, so an operator cannot disable the feature or restrict which paths may be loaded; and nothing distributes the library to executors.

[paleolimbot]

Thank you!

I took a read through and flagged anything I saw but they may well be me misunderstanding the current code...take or leave!

[paleolimbot]

Optional, but this seems like it would be easy to forget to do...if it works with the use of this, it would probably be less error prone to force a caller to "move" the kernel (i.e., set the release callback of the array version wrapped here to NULL and force the caller to take responsibility of the C struct. Then this release callback would drop any valid kernels that remained.

[andygrove]

Agreed, and the current shape only works because the host remembers to write a default back into each slot it moves out of. That is exactly the kind of thing that stays correct until someone edits the loader.

Not doing it here: making the taker NULL the source release changes the ABI contract itself, so it wants to land deliberately rather than as a review fix. Filed as #5250, which also covers the mid-import question you raised further down.

[paleolimbot]

The release callback might also be good to check here (and in other debug asserts)

Suggested change

	debug_assert!(!this.private_data.is_null());
	debug_assert!(!this.private_data.is_null() && this.release.is_some());

[andygrove]

Applied in ceb9365, and extended to the other three entry points that reach private_data: c_factory_new_impl, c_kernel_init and c_kernel_execute. A released struct has release: None, so checking both catches a call made after release rather than only an uninitialized one, which is the more likely mistake.

[paleolimbot]

Can this type of error be logged in some way or is that not safe/possible here?

[andygrove]

Possible, and worth doing. catch_panic_infallible now prints to stderr with a context string naming the callback that panicked (ceb9365).

Not a logging facade, though: arrow is deliberately the SDK's only dependency, and pulling in log for this would put a facade with no subscriber inside every user cdylib. stderr is the honest option at this layer. The cases it covers are all cleanup paths where the ABI gives no way to return an error, so previously a panicking user Drop left a leaked allocation and no trace at all.

[paleolimbot]

You have this in the document above, but perhaps worth reiterating here that these are comet version specific (i.e., these are not yet ABI stable between comet versions).

I believe that neither the arrow nor DataFusion version have to match here (but one may have to relax the arrow version compatibility in Cargo.toml for Rust dependents of this crate to be able to compile it).

[andygrove]

Added a # Stability section to the module docs in ceb9365, saying plainly that these structs are specific to one Comet version, that they are internal under Comet's versioning policy and may change in any release including a patch, and that the practical consequence is rebuilding the cdylib per Comet release.

Your second point is right and I have written it down too: neither the arrow nor the datafusion version has to match, because only FFI_ArrowArray and FFI_ArrowSchema cross the boundary and those are #[repr(C)] renderings of the C Data Interface. The binding constraint is what comet-udf-sdk itself compiles against, since the SDK is built into the user's cdylib and Cargo has to unify its arrow requirement with theirs. That is the relaxation you raise on Cargo.toml, tracked as #5253.

[paleolimbot]

I am not sure if this will wreak havoc the build, but it may be possible to relax the version compatibility here to whenever the FFI_ArrowArray/Schema were added.

[andygrove]

Filed as #5253 rather than done here. Two reasons: the crate is publish = false and consumed by git, so there is no external consumer to unblock yet, and a relaxed range on one workspace member needs an override rather than workspace = true, plus a CI job actually building against the floor. Otherwise the range is a claim nothing tests, which is worse than the current pin.

Noted in the issue that this becomes the blocking constraint the moment the SDK is published.

[paleolimbot]

I know you don't support them one level up, but since you're calculating a field here anyway can you implement return_field()?

[andygrove]

Can, and return_type is already building the whole Field and throwing away everything but the DataType, while the planner separately hardcodes the output field as nullable. So it is mostly a matter of returning what is already computed.

Filed as #5251 rather than done here, because it changes what nullability Comet promises: Spark UDF results are nullable in Spark's own schema, and the declared-vs-actual check in planner.rs deliberately erases nested nullability for a related reason. I would rather that land with a test for a kernel reporting a non-nullable field than slip in as part of a docs pass.

[paleolimbot]

This version of loading I believe will also resolve libraries on LD_LIBRARY_PATH (the previous version canonicalizes the path first, so I think only resolves against the working directory or an absolute path)

[andygrove]

Correct. cache::get_or_load canonicalizes with unwrap_or_else(|_| raw), so a path that does not resolve on the filesystem is handed to Library::new unchanged and gets the normal dlopen search. The docs claimed absolute paths were required; they now describe what the code does. Same thread as your cache.rs comment.

[paleolimbot]

I think the array's release won't double free here as currently implemented (but should!). This probably never leaks because none of this is ever likely to fail but in theory it would if an error occurred mid-import.

[andygrove]

Traced this and I do not think it leaks, though it took following three separate owners to convince myself.

On the ? early return: kernels 0..i are owned by udfs and released through their Arc<ImportedCScalarUdf>; kernel i was moved into the Box passed to try_new, which drops it on the error path; kernels i+1..len are still live in the array, and list's Drop calls c_list_release, which reconstructs the boxed slice and drops each one. The moved-out slots hold defaults with release: None, so they are no-ops rather than double frees.

That said, the fact that it takes a paragraph to establish is your point on the other thread. A test forcing a failure mid-import would make it durable instead of incidental, and it needs a fixture library exporting a deliberately bad kernel, so I have folded it into #5250 alongside the move-semantics change.

[paleolimbot]

Is this always true because RustUdfs are always immutable?

[andygrove]

It was not always true, and this is the one comment in the review that turned out to be a live bug rather than a rough edge.

CometRustUDF.register takes deterministic as a public parameter, defaulting to true, and calls asNondeterministic() on the Spark catalog stub when it is false. The value then travels in this proto field and is never read again: the planner ignores it, and ImportedCScalarUdf::try_new hardcodes Volatility::Immutable. So a UDF the caller explicitly declared nondeterministic was planned as pure, and DataFusion was free to fold it over constants, evaluate it once and reuse the result, or drop a repeated call as a common subexpression.

Honoring the flag properly is not a one-liner, because the signature is built once per library load and cached process-wide, while determinism is declared per registration. Two registrations of the same kernel with different determinism would need different volatility out of one cached ScalarUDFImpl. That is #5249.

For this PR, 3e999c4 makes register reject deterministic = false with an explicit not-yet-supported error, so the parameter cannot quietly lie about what Comet does with it, with a test. The field comment now records that it is always true today and why it is still carried on the wire.

[paleolimbot]

No pressure to do this here, but you could do something like this to pass the volatility, Display, Debug, etc. over FFI in a forward-flexible way. I haven't added this for UDFs yet but I'm using it for table providers, exec plans, and expressions. The annoying part is returning a variable length string (I use a FFI_ArrowArray for this at the moment, which is possibly overkill).

Suggested change

	pub get_property: Option<
	unsafe extern "C" fn(
	*mut CometCScalarKernelImpl,
	property: *const c_char,
	args: *const c_char,
	out: *mut FFI_ArrowArray,
	) -> c_int,
	>,

[andygrove]

This is the right shape for it, and it would have paid for itself already: volatility (#5249) is exactly a property that today would need a new struct field and an ABI version bump, breaking every existing cdylib to add one boolean.

Filed as #5254 with your sketch. Not taking it here because adding the field is itself an ABI change, so it wants to land before anyone depends on the current layout or ride along with the next bump, rather than being appended to a PR that is already large. Agreed the variable-length string return is the awkward part; noted the FFI_ArrowArray approach and that it may be heavier than needed.

[andygrove]

Thanks for the detailed review @paleolimbot! I'm going to address some of these items in this PR and will file follow on issues to discuss the rest.