`ComposedPhysicalExtensionCodec::encode_protobuf` claims by-name encodes and breaks UDF round-trips
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
**Describe the bug**
`encode_protobuf` (`datafusion/proto/src/physical_plan/mod.rs:1935`) has three related problems that combine to break UDF serialization whenever codecs are composed.
**1. It breaks on the first `Ok`, whether or not bytes were written.**
```rust
for (position, codec) in self.codecs.iter().enumerate() {
match encode(codec.as_ref(), &mut data) {
Ok(_) => { encoder_position = Some(position as u32); break; }
Err(err) => last_err = Some(err),
}
}
```
`PhysicalExtensionCodec::try_encode_udf` and `try_encode_udaf` default to `Ok(())` writing nothing — that is the encode-by-name signal. So the first codec in the list that does not override `try_encode_udf` claims *every* scalar UDF at position 0, and the codec that actually owns them is never asked.
**2. An empty result is framed anyway.**
`DataEncoderTuple { encoder_position, blob }` is emitted unconditionally, so `buf` comes back non-empty even when `blob` is empty. DataFusion reads an empty `fun_definition` as "resolve by name" — `(!buf.is_empty()).then_some(buf)` in `ConverterPlanEncoder::encode_udf` — so framing it sets the field and permanently skips the registry-first decode path in `from_proto.rs`:
```rust
None => ctx.udf(fun_name.as_str())
.or_else(|_| codec.try_decode_udf(fun_name, &[]))?,
```
**3. `data` is not cleared between attempts.** A codec that writes bytes and *then* errors leaves them in the buffer, and whichever codec succeeds next commits them along with its own.
**To Reproduce**
Compose `[CodecA, CodecB]` where `CodecA` does not override `try_encode_udf` and `CodecB` owns the UDFs. Serialize a plan referencing one of `CodecB`'s functions:
- encoding stops at `CodecA`, writes no bytes, stamps `encoder_position: 0`
- the tuple is written anyway, so `fun_definition` is `Some`
- decoding dispatches to `CodecA::try_decode_udf(name, &[])`, whose default is `not_impl_err!`
- the registry is never consulted, because the payload looked non-empty
The same sequence breaks a plain by-name UDF that would round-trip fine through any single codec.
**Expected behavior**
- The search continues past a codec that returns `Ok` without writing bytes, so a later codec can claim the object.
- When no codec writes bytes, `buf` is left empty, preserving the encode-by-name signal and the registry-first decode path.
- Each attempt encodes into a fresh buffer, so a failed attempt cannot contribute bytes.
**Additional context**
Found while evaluating `ComposedPhysicalExtensionCodec` for `datafusion-python` (https://github.com/apache/datafusion-python/pull/1678), where we hit exactly this and ended up with the three behaviours above in our own chain implementation. Happy to port the fix upstream.
Related: #24829 — the same type implements only half the trait.
Contributor guide
Research direction
Start in datafusion/proto/src/physical_plan/mod.rs:1935 and inspect ComposedPhysicalExtensionCodec::encode_protobuf together with the encode-by-name behavior of try_encode_udf and try_encode_udaf. Reproduce the composed [CodecA, CodecB] case described in the issue, then verify that empty successes do not claim encoding, failed attempts do not leak bytes, and UDF round-trips preserve the registry-first decode path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- data-engineering
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 75/100