apache / apache/datafusion-comet

Rust UDF adapter rebuilds the kernel impl and re-resolves the return type on every batch

Open
#5,296 0 comments 0 reactions 0 assignees View on GitHub
enhancement performance
Dominant language
Scala
Stars
1.3k
Forks
373
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_r3730399898)), raised by @mbutrovich.

`ImportedCScalarUdf::invoke_with_args` does the following on **every batch**:

1. builds a fresh `CometCScalarKernelImpl` via `new_impl`,
2. re-encodes every argument `Field` to an `FFI_ArrowSchema`,
3. calls the kernel's `init`, which re-invokes its `return_field`,
4. decodes the resulting `FFI_ArrowSchema` back into a `Field`,
5. drops the impl, running its `release` callback.

Argument types do not change across batches for a given `ScalarFunctionExpr`, and `return_type()` already did the same work once at planning time, so steps 2 and 4 are pure repetition.

Two separable pieces of work:

**Drop the return-type decode (small).** DataFusion 54's `ScalarFunctionArgs` carries `return_field`, so `Field::try_from(&out_schema)` in `invoke_with_args` is unnecessary. Careful, though: today the output array is imported with the type the kernel reported for *this* batch, and switching to the plan-time type means `from_ffi_and_data_type` receives a type that can differ in nested nullability from what the kernel actually produced — the same asymmetry the declared-vs-actual check in `planner.rs` deliberately erases. Needs the complex-type matrix in `CometRustUdfSuite` re-run rather than a spot check.

**Cache the impl across batches (larger, ABI decision).** The per-batch impl lifecycle is currently documented in `comet-udf-sdk`, so a kernel author may assume a fresh impl per call. Reusing one across batches changes that contract and should be decided deliberately, alongside the thread-safety question in the Mutex issue.

Contributor guide

Open the contributing guide

Research direction

Start by reading ImportedCScalarUdf::invoke_with_args and return_type(), then inspect planner.rs and the per-batch lifecycle documentation in comet-udf-sdk. Run the complex-type matrix in CometRustUdfSuite; done means the redundant work is addressed deliberately, output types remain correct, and the impl-reuse and thread-safety contract is resolved.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
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.