Skip to content

ComposedPhysicalExtensionCodec::encode_protobuf claims by-name encodes and breaks UDF round-trips #24830

Description

@timsaucer

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.

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:

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 (apache/datafusion-python#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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions