Make extension codecs composable (#1678)

* Make extension codecs composable

Installing a logical or physical extension codec now prepends it to a
codec chain instead of replacing the prior codec. The most recently
installed codec is consulted first, falling through codec by codec to
the default codec. This lets multiple independent extension libraries
install codecs on the same session, and removes the codec registration
ordering requirement between libraries.

Chain dispatch treats a codec error as "not mine". Encoding runs each
codec against a scratch buffer so failed attempts leave no partial
bytes, and treats Ok-with-no-bytes (encode by name) as no opinion so
later codecs still get a chance. When every codec fails, the errors
are aggregated so the owning codec's diagnostic is not masked by the
default codec's generic error.

Also preserves the python_udf_inlining setting when installing a
codec; previously it was silently reset to enabled.

Documents the remaining planner constraint: a session holds one query
planner, layering is explicit via fallback capsules, and codecs must
be installed before exporting or chaining planners because a planner
capsule captures the codecs at export time.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* test: port codec Rust tests to pytest

The Rust tests added for the composable codec work never ran: CI invokes
`cargo fmt` and `cargo clippy --all-targets` but no `cargo test`, so the
tests compiled and were never executed. Rather than add a `cargo test`
job — which would also require feature-gating `pyo3/extension-module`,
since the test binary cannot link on Linux while it is unconditional —
move the coverage to pytest, matching this repository's practice of
treating the user-facing Python surface as the first line of defense.

Remove both `#[cfg(test)]` modules from crates/core/src/codec.rs and
replace them as follows:

- Four wire-header round-trip tests and the Python-minor-mismatch test
  were already covered by existing cases in test_pickle_expr.py.
- `strip_errors_on_too_old_version` asserted nothing: it returns early
  because WIRE_VERSION_MIN_SUPPORTED equals WIRE_VERSION_CURRENT.
- The unsupported-wire-version and Python-major-mismatch cases move to
  test_pickle_expr.py, 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.
- The three truncated-header cases are dropped. Truncation changes the
  payload length and breaks the protobuf framing, so they fail before
  reaching the header check and cannot be expressed from Python.
- The codec-chain tests move to the FFI example suite, which exercises
  the same chain through the real FFI boundary.

MyLogicalExtensionCodec gains an optional token overriding the byte
prefix it stamps on encoded table providers. Two instances with distinct
tokens own disjoint slices of the wire format, which is what makes chain
ordering and fall-through observable from Python.

The ported inlining test asserts encode and decode behavior rather than
the `python_udf_inlining()` getter the Rust test checked. This is a
stronger assertion: the getter is preserved even when a composed
Python-aware codec re-inlines a UDF that the outer strict codec declined
to inline, so the original test could not have caught that path.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: record the Python-first testing preference in AGENTS.md

New coverage should land as a doctest example or a pytest case. Agents have
been adding Rust tests that CI never executes: no workflow invokes
`cargo test`, and `cargo clippy --all-targets` only compiles the test code.
Write down that constraint, along with the reason a `cargo test` job is not a
trivial addition, so the tradeoff does not have to be rediscovered.

Also point at the FFI example suites, which are easy to overlook when judging
whether behavior is reachable from Python.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* refactor: dispatch chained codecs by identity, not by trial

Decoding walked the codec chain and took the first codec that returned
Ok. That is unsound, and upstream has already been bitten by it:
apache/datafusion#16980 records ComposedPhysicalExtensionCodec decoding
an encoded parquet file as csv because the payload happened to parse
with an earlier codec, and #16986 fixed it by recording the encoder.
Protobuf carries no type identity, so two codecs whose leading field
numbers and wire types line up decode each other's payloads cleanly.
A byte-prefix convention does not help: the natural implementation is
`Message::decode(buf)`, which has no prefix to check and cannot decline.

Payloads written by a chained codec now carry an envelope naming the
codec that wrote them, and decoding consults exactly that codec. The
envelope is applied and stripped inside PythonLogicalCodec /
PythonPhysicalCodec, so third-party codecs are unmodified and never see
it -- they receive the bare bytes they wrote.

Keyed on a stable identity rather than chain position. Position is
sound for Ballista and datafusion-distributed because their codec lists
are pinned -- a compile-time constant in one, a fixed entry plus
appended user codecs rebuilt from the same startup code in the other.
A datafusion-python chain is assembled by user Python across sessions
that share no struct, and a library shipping a codec cannot know its own
index, so position would silently name the wrong codec whenever two
sessions install in different orders.

Identity is derived and asks nothing of existing libraries: an explicit
codec_id, else __datafusion_codec_id__, else the exporting class's
module and qualified name. A bare PyCapsule exposes nothing stable --
every capsule reports the same type -- so it gets a session-local id and
a pointed error if its payloads reach an unrelated session. Installing
two codecs under one id is rejected at install time rather than
resolving to whichever entry came first.

Codecs now append rather than prepend, so encoding is claimed by the
first codec installed that wants the object. Installing a library can
only claim objects nothing else claimed; it can never take over an
existing library's objects, and it never renumbers ids that older
payloads reference.

Two payloads are deliberately left unframed. The terminal codec writes
bare, so a session with no extension codecs is byte-identical to a build
without chaining. And an encode that writes nothing stays empty: an
empty fun_definition is DataFusion's encode-by-name signal, and framing
it would set the field, permanently skipping the registry lookup the
decoder does first and breaking codecs that reconstruct a function from
its name alone (from_proto.rs, the `None => ctx.udf(..).or_else(..)`
arm). That arm is also why an empty buffer still consults every codec:
there are no bytes to tag. It is not the hazard tagging removes -- the
question asked is "do you own the function named x", which is
name-scoped, and a disagreement needs two libraries claiming one name.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* test: pin the by-name decode path, and document identity dispatch

Adds NameOnlyUdfCodec to the FFI example: a codec owning functions that
are fully described by their names, so try_encode_udf writes nothing and
try_decode_udf rebuilds the function from the name with no registry
entry. DataFusion supports that shape directly -- an empty
fun_definition sends the decoder to the FunctionRegistry first and the
codec second (from_proto.rs, the `None => ctx.udf(..).or_else(..)` arm).

Nothing covered that arm before, and it is the path most at risk from a
plausible change: wrapping every chained encode in the identity envelope
would make an empty payload non-empty, set fun_definition, and skip the
registry lookup permanently. That breaks ordinary by-name round trips as
well as codecs like this one, and no other test would notice. The
decoding session here deliberately never registers the function, so only
the codec can supply it.

Exposes logical_extension_codec_ids() and physical_extension_codec_ids()
on SessionContext. The ids are the dispatch keys a payload names, so
listing them answers "which library owns this plan" and "can this
session decode it" -- and they are what a decode failure reports.

Documentation rewritten around identity rather than ordering:

- ffi.md states that codecs need no changes, why rejecting foreign
  payloads is not something a codec can reliably do, the identity
  resolution ladder, and the two payloads left unframed with the reason
  for each.
- Both example READMEs drop the "most recently installed is consulted
  first" language, which described dispatch that no longer exists.
- upgrade-guides.md gains a section for the behaviour change: codecs
  compose rather than replace, no codec-side change is required, and the
  wire format changes only for sessions that install one.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: record codec chain dispatch in the FFI capsule protocol skill

The skill already triggers on "any FFI_* export that asks for ... an
extension codec", but had nothing to say about how a chain decides which
codec handles a payload -- the part of this area most likely to be got
wrong, and the part with a rationale that is not visible from the code.

Rule 8 records it: dispatch by identity rather than by trying codecs
until one returns Ok, with the upstream incident that settles the
question (apache/datafusion#16980, a Parquet payload decoded as CSV) and
the protobuf reason a byte-prefix convention cannot fix it.

It also records why we key on identity where Ballista,
datafusion-distributed, and upstream all key on chain position. Their
lists cannot disagree -- one is a compile-time constant, the other a
pinned entry plus user codecs rebuilt from the same startup code. Ours is
assembled by user Python, and to_bytes/from_bytes puts two independently
built sessions on either end of one payload. Copying the upstream design
here reintroduces the same silent mis-decode from the other direction,
which is exactly the kind of thing an agent reading only upstream would
do.

And it records the two payloads that are never framed, because framing
the empty one is a plausible tidy-up that breaks by-name decoding
silently. Points at NameOnlyUdfCodec as the guard.

Extends the trigger to cover codec.rs dispatch, which the old wording
did not obviously reach, and documents the optional
__datafusion_codec_id__ hook on the two codec Protocols -- where Rule 5
says such things belong.

Reviewed the other four skills: check-upstream, make-pythonic,
audit-skill-md, and the user-facing datafusion_python skill are all
unaffected, none of them touching codec dispatch.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: move codec dispatch rationale out of the skill and into the code

Rule 8 restated, less precisely, what the docstrings in codec.rs already
say: identity dispatch, why not position, the two unframed payloads,
append-and-first-claim-wins. A skill earns its place by carrying what
you cannot learn from the file you are about to edit -- Rule 1 greps a
family across files, Rule 3 governs code in an extension author's own
repo, Rule 5 is a four-file process checklist. "How the dispatch in
codec.rs works" is what codec.rs is for, and duplicating it there just
creates a second copy to keep in sync.

Two pieces were worth keeping, and neither is a rule.

The apache/datafusion#16980 citation moves to `chain_decode`, next to
the code someone would be editing when tempted to delete the envelope
and walk the chain instead. It is the concrete evidence that makes the
warning land, and it was the one thing the skill had that the code did
not.

Rule 5 gains a clause: changing what a codec puts on the wire is as
breaking as changing a getter signature, and easier to miss because
nothing fails to compile. That is genuinely cross-file process, which is
what the skill is for.

Also enriches ChainEntry's docstring with the datafusion-distributed
half of the "why not position" argument, which previously named only
Ballista and upstream. Both consumers matter: the point is that their
codec lists structurally cannot disagree and ours can.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: give a capsule-installed codec an id no other session can mint

A codec installed from a bare PyCapsule has nothing stable to derive an
identity from, so it was tagged `anon:{chain_length}`. That is a position,
not an identity: every session numbers from the same end, so two unrelated
sessions that each install one capsule both mint `anon:0`. A payload written
by the first was then handed to the second session's codec — the positional
dispatch this chain design exists to avoid, reintroduced in the one case
that has no derivable id.

The failure was quiet. A codec offered bytes it does not recognise falls
through to its own inner default codec, so the error came back as
`LogicalExtensionCodec is not provided` and named neither codec, instead of
the intended "encoded by extension codec X, which is not installed on this
session" with the `codec_id=` hint.

Mint a random id per install instead, reusing the `Uuid` idiom already in
`context.rs`. The chain clones the id along with the codec, so a payload
still decodes anywhere in the installing session's lineage; everywhere else
it now fails with the pointed error. `derive_codec_id` no longer needs the
chain length, so that parameter goes away.

Also fixes two docs that described the pre-identity-dispatch design: the
`codec.rs` module doc claimed codecs are consulted most-recently-installed
first (it is install order, and decoding does not walk the chain at all) and
asked downstream codecs to reject foreign payloads (an encode-side contract
only — a payload only ever reaches the codec whose id it carries), and a
physical codec test docstring said a second install prepends.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: identify a session installed as a codec by its session, not its class

Installing one context's codec stack on another session derived the
identity from the class, and every `SessionContext` shares one class. So
every session reported `datafusion.context.SessionContext`: two of them
could not coexist on one target, and a payload written through one
resolved to the other on decode, failing as
`LogicalExtensionCodec is not provided` from the stranger's own inner
default codec. Same shape as the positional `anon:{n}` collision, but via
an id that looks portable, so the error carried no hint either.

`SessionContext` now declares `__datafusion_codec_id__` carrying its
session id. That goes through the existing resolution arm for an object
pinning its own identity, so one implementation covers both the Python
wrapper and the internal object and `derive_codec_id` keeps its four
documented arms. Handles derived from one session report the same id, so
installing two of them on one target is refused — their payloads would be
indistinguishable on decode.

Framing is unchanged. A session-exported codec is opaque across the FFI
boundary, so the outer session cannot enumerate what is inside it and the
envelope naming its own entry is the only handle it has; nesting costs
about 50 bytes per hop and only arises when composing sessions, which a
library exporting its own codec class never does. What that route needed
was a usable identity, not fewer frames.

Documents the lifetime coupling that comes with it: imported codecs
resolve their task context against the source session and stop working
when it is dropped, so this composes sessions rather than copying codecs
out of one.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* test: pin the class-derived codec id, and the escape hatch from it

A codec that declares no identity is named `module.QualName`, and that
string goes on the wire in front of every payload it writes. Nothing
asserted it, so renaming `MyLogicalExtensionCodec` or moving its module
would have silently changed the wire format with a green suite.

Pins it twice: against the literal, so a rename has to come here and be
acknowledged, and against `__module__`/`__qualname__`, so the literal
cannot drift away from what the code actually derives.

Also covers the reason `__datafusion_codec_id__` exists, which had no
test either. A plan encoded by a codec under its old class name decodes
on a session that only knows the new one, and the payload carries the
pinned id rather than either class name. Under the class-derived default
those are two different ids and the plan is undecodable.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: aim the codec documentation at its readers

The codec documentation added in this PR drifted into narrating the
implementation. An extension author cannot act on the name of a private
Rust type, and a caller installing a library's codec does not need the
envelope layout or the argument for why trial decoding is unsound.

Removed from `ffi.md`: the paragraph naming `Python{Logical,Physical}Codec`
as the thing that wraps payloads, the protobuf-ambiguity argument with the
upstream postmortem, the note on why the bare-capsule id is random, and the
two-unframed-payloads section, which opened by addressing whoever is
"changing this code". What each reader can act on stays: implement your
codec as though yours is the only one, do not write defensive prefix checks,
how identity is derived, when to pass `codec_id=`, that install order does
not affect decoding, and that a codec may encode a function by name alone.
All of the removed rationale already lives in `crates/core/src/codec.rs`,
which the section now points to for anyone changing the framing.

Trimmed `with_logical_extension_codec` to what a caller acts on, with a
`:ref:` to the FFI guide for the rest, and added the `ValueError` on a
duplicate id, which the longer version buried. `with_physical_extension_codec`
now delegates in one sentence instead of restating the same paragraphs, which
had already begun to drift.

Reworded the internal vocabulary out of the upgrade guide and
`logical_extension_codec_ids`: "terminal codec", "writes unframed", and
"identity envelope" name internals a reader cannot see. Dropped the
explanation of why a session id rather than a class name identifies a
context installed as a codec, keeping the constraint that follows from it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: state the one case where a codec sees bytes it may not own

The FFI guide told codec authors not to guard against another library's
payloads, on the grounds that such a check "simply never fires". That was
wrong twice over. The check executes on every decode, and its reject branch
is reachable: an empty payload carries no identity to route on, so
`chain_resolve_by_name` offers it to every installed codec in turn, and a
codec asked about a function it does not own rejects — correctly.

That exception is worth documenting on its own terms, because an author who
trusts the "you only ever see your own payload" guarantee can index into an
empty `buf` and panic. `try_decode_udf` and its aggregate and window
siblings can be called with an empty buffer and another library's `name`, so
the guide now says to decide from `name` and not to assume `buf` is
non-empty. `name_only_codec.rs` in the FFI example already does exactly
that, so the guide and the worked example agree.

The advice against defensive checks is gone rather than reworded. An author
who already writes one keeps it and nothing breaks; an author who does not
needs nothing. What remains is the portability case: a codec that also ships
to hosts dispatching by position or by trial may still want its own guard,
and keeping one costs nothing here.

Two claims in codec.rs conflated a marker check with relying on decode
failure. A marker an encoder writes and a decoder checks is reliable — it is
what the envelope itself does. What cannot be relied on is
`MyMessage::decode(buf)` declining someone else's bytes, and what cannot be
required is third-party codecs all honouring such a convention. Both are now
stated as what they are.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: call a codec id an id, and say what one is

The codec documentation used "identity" as a term of art without ever
defining it, while the API it describes calls the same thing `codec_id`.
A reader met an abstract noun, had to guess what it referred to, and then
had to map it onto the parameter they actually type.

The FFI guide now defines it once, where the surrounding text has already
described the mechanism: datafusion-python records which codec wrote each
payload, and that record is the codec's id — a short string stored inside
the plan, which has to name the same codec in the process that decodes as
it did where the plan was written. That last clause is the reason ids
exist and is the part a reader needs.

Replaced the derivation-order list with the three cases where a reader has
to do something: two instances of one class collide and raise `ValueError`,
a bare `PyCapsule` gets an id private to its session, and a class you
intend to rename wants `__datafusion_codec_id__`. Each names the symptom
and the fix. The order in which an id is resolved answered a question
nobody asks, and the sub-clause about how a `SessionContext` names itself
described an id nobody types. The session-composition paragraph loses its
id explanation, keeping the lifetime coupling that is the real hazard.

Net 34 lines to 26 in that part of the guide, with a definition it did not
have before. The same substitution follows through the upgrade guide, the
`SessionContext` and protocol docstrings, and both example READMEs, since
leaving the user-facing pages on the old vocabulary would have been worse
than the original.

Untouched: "shared-library identity" earlier in the guide, which predates
this work and refers to distinct DataFusion library markers.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: plain language for the two codec id prefixes

Both doc comments stacked several clauses per sentence and stated the
rejected alternatives abstractly, which a reviewer reported as hard to
parse. Same content, plainer register: the anonymous prefix now says
the suffix is a fresh UUID, matching what derive_codec_id mints, and
spells out the anon:0 collision instead of describing it as sessions
numbering from the same end.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* test: count decodes, not just encodes, on the composed chain

test_ffi_logical_codec_composes_with_later_install asserted the encode
count and then only that the round-trip produced equal batches. The
round-trip result alone does not distinguish the codec decoding the
bytes it wrote from the table being resolved some other way, so the
decode side of the chain was unpinned. Count it too.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: describe the composed-codec test in install order

The docstring for test_composed_codecs_with_query_planner said the
extra codecs decline everything and encoding falls through to the
provider codecs, which describes the reverse of what happens. The
provider codecs are installed first and chain_encode walks in install
order, so they claim their objects before the extra codecs are
consulted; decoding dispatches by id and does not consult them at all.
Say that, and say what the test asserts: the extra pair changes
nothing, where replace semantics would have broken the round trip.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>