Make extension codecs composable

[timsaucer]

Which issue does this PR close?

Part 2 of 3 in the split of #1672. Part 1 (#1677) is merged; this branch is rebased on it.

Rationale for this change

Supporting a foreign planner surfaced a codec problem: a query can involve three independent native libraries (datafusion-python, a provider library, and a planner library), and each library needs its extension codecs active on the session at the same time. Previously, installing a logical or physical extension codec replaced the prior codec, so the second library's install silently discarded the first — plans then failed later with a confusing decode error.

What changes are included in this PR?

Codecs compose instead of replacing

with_logical_extension_codec / with_physical_extension_codec append to a codec chain rather than replacing the prior codec. Encoding consults the chain in install order and the first codec to write bytes claims the object, so installing a library can claim objects nothing else claimed but never takes over an object an earlier codec was already encoding. Each codec encodes into a scratch buffer, so a failed attempt leaves no partial bytes. Ok with an empty buffer is treated as "no opinion" rather than a claim, so later codecs still get a turn.

with_python_udf_inlining clones the codec chain rather than rebuilding around a single inner codec, so toggling inlining no longer collapses a multi-codec chain to its first entry.

Decoding dispatches on identity, not on position or trial

This part was redesigned in response to review from @milenkovicm and @gabotechs — the earlier version of this PR tried each codec in turn until one returned Ok, which is unsound. Every payload an installed codec writes is now wrapped in an envelope naming the codec that produced it, laid out as DFPYCHN | version: u8 | id_len: u32 (LE) | id | blob. Decoding reads the id and consults exactly one codec, so a codec is never offered bytes it did not write, and its error surfaces verbatim rather than being masked by whichever codec was tried last.

Nothing is asked of the codec itself. An extension library implements LogicalExtensionCodec / PhysicalExtensionCodec exactly as it would for a session that installs only its own, and receives its payload byte for byte; the envelope is applied and stripped by datafusion-python and never reaches it. In particular a codec does not have to recognise and reject foreign payloads, which is not something a codec can reliably do: protobuf carries no type identity, so a prost message decodes cleanly from an unrelated message's bytes whenever their leading field numbers and wire types line up, and the natural implementation, MyMessage::decode(buf), has no prefix to check and cannot decline.

datafusion_proto::physical_plan::ComposedPhysicalExtensionCodec was suggested and is not reused, for three reasons:

1. It keys dispatch on encoder_position, which is sound only when both ends assemble the same codec list in the same order. That holds for Ballista (a compile-time constant) and datafusion-distributed (own codec pinned at index 0, user codecs rebuilt from the same startup code on every node). It does not hold here: a chain is assembled by user Python, and Expr.to_bytes(ctx1) / Expr.from_bytes(ctx2) puts two independently configured sessions on either end of one payload, so an index names a different codec in the decoder as soon as install order differs.
2. It exists only for the physical layer. There is no ComposedLogicalExtensionCodec, and this change needs both layers.
3. encode_protobuf always wraps the result in DataEncoderTuple, so a codec that writes zero bytes still produces a non-empty payload. That sets fun_definition and permanently skips the FunctionRegistry lookup the decoder does first, breaking DataFusion's encode-by-name path.

Codec identity

An identity is derived automatically, in this order:

1. An explicit codec_id= argument to the install call.
2. __datafusion_codec_id__ on the exporting object. SessionContext declares one carrying its session id, because the class-derived fallback below would name every session at once.
3. Otherwise the exporting class's module.QualName, which is the library's own import path and already stable across processes. This is the common case and asks nothing of existing extension libraries.
4. For a bare PyCapsule there is nothing stable to read, since every capsule reports the same type. It gets a random id minted at install time, so a payload it writes decodes within the installing session's lineage and fails with a pointed error anywhere else. Pass codec_id= when such plans must cross sessions.

Installing two codecs under one id raises ValueError rather than shadowing the first, because a payload naming that id would otherwise resolve to whichever entry came first and only the caller knows whether the two write the same wire format. SessionContext.logical_extension_codec_ids() and physical_extension_codec_ids() list what is installed, which is also what a decode failure names.

Two payloads deliberately stay unframed

The terminal codec — Default{Logical,Physical}ExtensionCodec unless a Rust caller supplied another — writes bare, so a session with no extension codecs installed serializes byte-identically to a build without codec chaining.

An encode that writes nothing also stays empty, because try_encode_udf returning Ok with an empty buffer is DataFusion's encode-by-name signal. Framing it would set fun_definition and skip the registry lookup permanently. That empty-buffer path is the one place every installed codec is still consulted in turn, since there are no bytes to carry an identity; it is not the hazard the envelope removes, because the question asked is "do you own the function named x", which is name-scoped, and two codecs disagreeing requires them to claim the same function name — already a collision in the function registry.

Documentation

docs/source/contributor-guide/ffi.md gains a "Composable codecs" section, scoped to what an extension author or a caller wiring several libraries together can act on: that nothing is asked of the codec itself, that defensive prefix checks against another library's payloads are unnecessary, how identity is derived, when to pass codec_id=, that registration order between libraries does not affect decoding, that a codec may encode a function by name alone, the caution that composing whole sessions couples the imported codecs to the source session's lifetime, and a two-library registration recipe. It replaces the paragraph stating that installing a codec replaces the prior one. The design rationale behind those rules — the envelope layout, why trial decoding is unsound, why the bare-capsule identity is random, and the two payloads that stay unframed — lives in crates/core/src/codec.rs, which the section points to, rather than in the published guide. docs/source/user-guide/upgrade-guides.md gains a migration section, and both example READMEs are updated.

Test coverage

The chain behavior was originally covered by two #[cfg(test)] modules in crates/core/src/codec.rs. Those tests never ran: no workflow invokes cargo test, and the only Rust checks are cargo fmt --check and cargo clippy --no-deps --all-targets. --all-targets compiles test code, so the tests could not rot into a non-compiling state, but a behavioral regression would not have failed the build. Adding a cargo test job is also not a one-line change, because crates/core/Cargo.toml enables pyo3/extension-module unconditionally and the test binary therefore fails to link against Py_* on Linux. The coverage was moved to pytest instead, matching this repository's practice of treating the user-facing Python surface as the primary focus.

• Both Rust test modules are removed. Five of the eighteen cases were already covered by existing pytest cases, and one (strip_errors_on_too_old_version) asserted nothing because WIRE_VERSION_MIN_SUPPORTED equals WIRE_VERSION_CURRENT.
• python/tests/test_pickle_expr.py gains coverage for the wire-header diagnostics — an unsupported wire-format version and a Python major-version mismatch — by patching the header in place inside the encoded protobuf. The patches preserve length so the outer message stays parseable and the bytes reach the codec.
• Three cases were dropped rather than ported. They truncate the wire header, which changes the payload length and breaks the enclosing protobuf framing, so they cannot reach the header check from Python.
• The FFI example suite gains coverage for the chain itself, exercised across a real FFI boundary: encode ordering, that a later install cannot hijack an earlier codec's objects, that decoding reaches the codec that encoded and survives a different install order in the decoding session, the error naming a codec that is not installed, rejection of a duplicate id, the per-session and per-install identities, that a class-derived id is what goes on the wire, and that a pinned __datafusion_codec_id__ survives a class rename.
• NameOnlyUdfCodec in the FFI example pins the by-name path: it owns functions fully described by their names, encodes no bytes, and rebuilds each function from the name on decode. It is the guard against framing empty payloads, which would break by-name round trips with nothing else in the suite noticing.
• One inlining test is retained in the FFI example suite even though Add FFI query planner support #1677 added equivalent coverage in test_pickle_expr.py. That version installs a codec exported from another SessionContext, which is itself a Python-aware codec with inlining enabled; this one installs a codec that delegates UDF encoding to DataFusion's default codec, which is the realistic extension-library case and the only one that pins the strict outer codec's own behavior.
• AGENTS.md records the Python-first testing preference and the fact that CI does not run Rust tests, so the tradeoff does not have to be rediscovered.

Are there any user-facing changes?

Behavior change. Installing an extension codec composes with previously installed codecs instead of replacing them. Code that relied on replacement semantics (installing a codec to remove a prior one) is affected; all other usage keeps working and no longer loses earlier codecs. There is no way to remove an installed codec.

Wire-format change. Once an extension codec is installed, payloads it writes carry the identity envelope. A session with no extension codecs installed produces the same bytes as before, as do functions encoded by name. Plans serialized by an earlier release and stored for later use should be regenerated if they were produced by a session with an extension codec installed.

New API surface, all additive:

• codec_id= on with_logical_extension_codec / with_physical_extension_codec.
• SessionContext.logical_extension_codec_ids() and physical_extension_codec_ids().
• SessionContext.__datafusion_codec_id__, and recognition of __datafusion_codec_id__ on any object being installed as a codec.
• Installing two codecs under one id now raises ValueError.

Carries the api change label, with a before/after section added to docs/source/user-guide/upgrade-guides.md.

MyLogicalExtensionCodec in examples/datafusion-ffi-example gains an optional provider_prefix argument that overrides the byte prefix it stamps on encoded table providers, which is what lets the tests install two instances owning disjoint slices of the wire format. It defaults to the previous constant, so existing call sites are unchanged, and the example README notes that real plugin libraries should hard-code a prefix unique to the library rather than accept one from the caller. The example also gains NameOnlyUdfCodec / NameOnlyFunction for the by-name path.

[ntjohnson1]

Hmm is it useful to be able to see for a given call which codec triggers? I see you can check the counts before/after with something like codec.table_provider_encode_calls() but I wonder if this does/should show up in an explain plan or something.

[ntjohnson1]

This was pub before. Do people care about inspecting the codecs directly?

[ntjohnson1]

Similar note about directly referencing codec chain

[ntjohnson1]

So is the resulting stage at this point [default, lib_a, lib_b]? I could read it or ask claude if some assumed based codec for standard datafusion python gets installed always as a fallback. Mostly I'm curious if it's a reasonable workflow for someone to set things up to only get [lib_a, lib_b] and if anything unsupported by their custom codecs barfs.

[ntjohnson1]

Ok this line says DF default codec is the terminal fallback but I think the constructor takes a codec in rust at least so it could potentially be updated there but maybe not in the exposed python surface.

[milenkovicm]

Please have a look at datafusion_proto::physical_plan::ComposedPhysicalExtensionCodec which should handle this case. try decode can produce wrong results in some cases and its hard to debug when it does

[milenkovicm]

this is not safe, you can get things decoded by accident. please dont ask how i know and how much time i spent debuging it

[gabotechs]

☝️ this is important, some more context in apache/datafusion#16980 and apache/datafusion#16986.

Because of how protobuf decoding works, it's very easy to be in situations where the same protobuf payload can decode to multiple prost structs, specially in cases where the structs are specially simple.

[timsaucer]

You were both right, and I've reworked the dispatch rather than keeping trial decoding. Thanks for the #16980 / #16986 pointers — the CSV-decoded-as-Parquet case is exactly the failure mode, and it's not one a codec can defend against from inside, since MyMessage::decode(buf) has no prefix to check and cannot decline.

Every payload an installed codec writes is now wrapped in an envelope naming its author: DFPYCHN | version: u8 | id_len: u32 (LE) | id | blob. Decoding reads the id and consults exactly one codec. Nothing is ever offered to a codec that did not write it, so structurally similar prost messages can't cross libraries. Codecs themselves are unchanged — they receive their payload byte for byte and never see the envelope.

On reusing ComposedPhysicalExtensionCodec: I looked at it closely and it doesn't fit here, for three reasons.

1. It keys on encoder_position, which is sound when both ends build the same list in the same order. That's true for the consumers it was written for — Ballista's list is a compile-time constant, and datafusion-distributed pins its own codec at index 0 and appends user codecs rebuilt from the same startup code on every node. It isn't true here: the chain is assembled by user Python, and Expr.to_bytes(ctx1) / Expr.from_bytes(ctx2) puts two independently configured sessions on either end of one payload. An index names a different codec in the decoder as soon as install order differs. test_decode_survives_a_different_install_order is the case.
2. It only exists for the physical layer. There's no ComposedLogicalExtensionCodec, and this needs both — table providers and UDFs go through the logical codec.
3. encode_protobuf always wraps in DataEncoderTuple, so a codec that writes zero bytes still emits a non-empty payload. That sets fun_definition and permanently skips the FunctionRegistry lookup the decoder does first, which breaks DataFusion's encode-by-name path. I need that path to keep working — NameOnlyUdfCodec in the FFI example owns functions that are fully described by their names and encodes nothing at all.

The reasoning is recorded in the code at crates/core/src/codec.rs (see ChainEntry and chain_decode, which cite both upstream issues) so the next person doesn't re-derive it, and the "Composable codecs" section of docs/source/contributor-guide/ffi.md covers it for extension authors. The empty-payload case is the one place codecs are still consulted in turn, because there are no bytes to carry an identity; that question is "do you own the function named x", which is name-scoped, and two codecs disagreeing means they've already collided in the function registry.

[milenkovicm]

this might work, but still can be tricky

[gabotechs]

I think it's worth following @milenkovicm's suggestion and reuse the ComposedPhysicalExtensionCodec.

[gabotechs]

☝️ this is important, some more context in apache/datafusion#16980 and apache/datafusion#16986.

Because of how protobuf decoding works, it's very easy to be in situations where the same protobuf payload can decode to multiple prost structs, specially in cases where the structs are specially simple.

[timsaucer]

@ntjohnson1 on your question about seeing which codec triggers for a given call — partly answered now, though not in EXPLAIN.

SessionContext.logical_extension_codec_ids() and physical_extension_codec_ids() list what's installed, in install order, and those ids are exactly what a payload carries — so a decode failure names the codec that wrote the bytes and lists what the session actually has, e.g. a table provider was encoded by extension codec 'lib_a.Codec', which is not installed on this session (installed: lib_c.Codec).

What's still missing is your actual question: per-call attribution, "which codec handled this node in this plan". Encoding decides that by walking the chain until one claims the object, and nothing records the winner. Surfacing it in EXPLAIN would mean threading it through plan formatting, which I'd rather not fold into this PR. Happy to file it as a follow-up if you think it's worth having.

Your other three comments are addressed as well: the inner() accessor is replaced by codecs() / codec_ids() / terminal() on the Rust side and logical_extension_codec_ids() / physical_extension_codec_ids() from Python, and both the contributor guide and the example README now state that the terminal fallback is Default{Logical,Physical}ExtensionCodec unless a Rust caller supplies another to Python{Logical,Physical}Codec::new, which is not reachable from the Python surface — so a chain of [lib_a, lib_b] with no default behind it isn't something you can set up from Python.