apache / apache/datafusion-comet
Rust UDF adapter rebuilds the kernel impl and re-resolves the return type on every batch
- 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
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