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>diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md
index 4682160..294ebfb 100644
--- a/.ai/skills/ffi-capsule-protocol/SKILL.md
+++ b/.ai/skills/ffi-capsule-protocol/SKILL.md
@@ -116,6 +116,12 @@
`python/datafusion/user_defined.py`, where the `Protocol` type hints for
these methods live.
+Changing what a codec puts *on the wire* is equally breaking, and easier to
+miss because no signature moves and nothing fails to compile. Serialized plans
+outlive the process that wrote them, so the same checklist applies: upgrade
+guide, `api change` label, and a statement of exactly which sessions produce
+different bytes.
+
## Rule 6 — a session keeps one `Arc<SessionContext>` for life
`FFI_TaskContextProvider` holds its provider **weakly**, and every codec handed
@@ -181,8 +187,11 @@
- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat.
- `docs/source/user-guide/upgrade-guides.md` — every past migration.
+- `crates/core/src/codec.rs` — the codec chain: the envelope, identity dispatch,
+ and the two unframed cases from Rule 8.
- `examples/datafusion-ffi-example/src/` — provider, catalog, function, codec
- getters, all in current form.
+ getters, all in current form. `name_only_codec.rs` is the codec that encodes
+ nothing.
- `examples/datafusion-ffi-query-planner-example/src/planner.rs` — planner
getter.
- `examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`
diff --git a/AGENTS.md b/AGENTS.md
index 327ebd6..659094e 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -82,6 +82,40 @@
Fix any failures before committing.
+## Test Coverage
+
+Always prefer Python coverage — a doctest example in a docstring, or a pytest
+case. The user-facing Python surface is the first line of defense and the
+primary focus, so behavior should be pinned where users actually meet it.
+
+**CI does not run Rust tests.** No workflow invokes `cargo test`; the only
+Rust checks are `cargo fmt --check` and
+`cargo clippy --no-deps --all-targets`. `--all-targets` compiles
+`#[cfg(test)]` code, so a Rust test cannot rot into a non-compiling state, but
+it is never executed and a behavioral regression will not fail the build. A
+Rust test added today is dead weight.
+
+Adding a `cargo test` job is not a one-line change: `crates/core/Cargo.toml`
+enables `pyo3/extension-module` unconditionally, so the test binary fails to
+link against `Py_*` symbols on Linux. The feature would have to be gated first.
+
+Write a Rust test only when the behavior is genuinely unreachable from Python,
+and wire up CI in the same change so it actually runs. Before concluding it is
+unreachable, check the suites that already exist:
+
+- `python/tests/` — the main suite. Run `pytest python/`, **not**
+ `pytest python/tests/`: `--doctest-modules` is on by default and the
+ narrower path skips the doctests in `python/datafusion/`.
+- `examples/datafusion-ffi-example/python/tests/` and
+ `examples/datafusion-ffi-query-planner-example/python/tests/` — integration
+ coverage across a real FFI boundary, for anything involving extension
+ codecs, table providers, query planners, or capsule export. These need the
+ example crates built (`maturin build`, then install the wheel).
+- `examples/tpch/` — end-to-end query coverage.
+
+Prefer asserting observable behavior over internal accessors. A test that
+checks a getter can pass while the path a user actually takes is broken.
+
## Python Function Docstrings
Every Python function must include a docstring with usage examples.
diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs
index 94942a2..8f43bdb 100644
--- a/crates/core/src/codec.rs
+++ b/crates/core/src/codec.rs
@@ -29,16 +29,18 @@
//!
//! [`PythonLogicalCodec`] is the [`LogicalExtensionCodec`] that
//! datafusion-python parks on every `SessionContext`. It wraps a
-//! user-supplied (or default) inner codec and adds Python-aware
-//! in-band encoding on top: when the encoder sees a Python-defined
-//! UDF, the codec cloudpickles the callable + signature into the
-//! `fun_definition` proto field; when the decoder sees a payload it
-//! produced, it reconstructs the UDF from the bytes alone — no
-//! pre-registration on the receiver. UDFs the codec does not
-//! recognise are delegated to `inner`, which is typically
-//! `DefaultLogicalExtensionCodec` but may be a downstream-supplied
-//! FFI codec installed via
-//! `SessionContext.with_logical_extension_codec(...)`.
+//! chain of composable codecs and adds Python-aware in-band encoding
+//! on top: when the encoder sees a Python-defined UDF, the codec
+//! cloudpickles the callable + signature into the `fun_definition`
+//! proto field; when the decoder sees a payload it produced, it
+//! reconstructs the UDF from the bytes alone — no pre-registration on
+//! the receiver. Everything the codec does not recognise is delegated
+//! to the chain: each downstream FFI codec installed via
+//! `SessionContext.with_logical_extension_codec(...)` is appended, and
+//! encoding consults them in install order with
+//! `DefaultLogicalExtensionCodec` as the terminal fallback. Decoding
+//! does not walk the chain at all — a payload names the codec that
+//! wrote it. See [`PythonLogicalCodec`].
//!
//! [`PythonPhysicalCodec`] is the symmetric wrapper around
//! [`PhysicalExtensionCodec`]. Logical and physical layers each have
@@ -58,7 +60,7 @@
//! actionable error instead of an opaque `marshal` failure on load
//! (cloudpickle payloads are not portable across Python minor
//! versions). Dispatch precedence on decode: **family match +
-//! supported version + matching Python version → `inner` codec →
+//! supported version + matching Python version → codec chain →
//! caller's `FunctionRegistry` fallback.**
//!
//! ## Wire-format family registry
@@ -81,10 +83,18 @@
//! for an older shape.
//!
//! Downstream FFI codecs should pick non-colliding family prefixes
-//! (use a `DF` namespace plus a crate-specific suffix). The codec
-//! implementations in this module currently delegate every method to
-//! `inner`; the encoder/decoder hooks for each kind are added as the
-//! corresponding Python-side type becomes serializable.
+//! (use a `DF` namespace plus a crate-specific suffix) and return an
+//! error for *objects* they do not own — on encode, that error is the
+//! chain's "not mine" signal, letting the next codec take a turn. A
+//! codec that answers `Ok` for objects outside its family claims them
+//! ahead of every codec installed after it.
+//!
+//! Rejecting foreign *payloads* is not asked of a codec, because a
+//! payload is only ever handed to the codec whose id it carries. A
+//! codec is free to check for a marker of its own anyway — that costs
+//! nothing here and is worth keeping for hosts that dispatch by
+//! position or by trial — but dispatch never depends on it. See
+//! [`PythonLogicalCodec`].
use std::sync::Arc;
@@ -167,7 +177,7 @@
/// Inspect the framing on `buf`.
///
/// * `Ok(None)` — `buf` does not carry `family`. The caller should
-/// delegate to its `inner` codec.
+/// delegate to its codec chain.
/// * `Ok(Some(payload))` — `buf` carries `family` at a version this
/// build accepts and a Python `(major, minor)` matching
/// `expected_py`; `payload` is the cloudpickle blob.
@@ -223,12 +233,345 @@
Ok(Some(&buf[py_minor_idx + 1..]))
}
+/// Family prefix for the envelope wrapping a chained codec's payload.
+///
+/// A distinct magic is what makes "is this framed?" a definite test
+/// rather than a speculative decode. Probing by attempting to parse the
+/// envelope would reintroduce exactly the protobuf ambiguity this
+/// framing exists to remove: prost skips unknown fields and defaults
+/// missing ones, so a foreign payload can parse cleanly as an envelope.
+pub(crate) const CHAINED_PAYLOAD_FAMILY: &[u8] = b"DFPYCHN";
+
+/// Wire-format version for the chained-payload envelope. Independent of
+/// [`WIRE_VERSION_CURRENT`], which versions the cloudpickle framing.
+pub(crate) const CHAIN_WIRE_VERSION_CURRENT: u8 = 1;
+
+/// Oldest chained-payload envelope version this build decodes.
+pub(crate) const CHAIN_WIRE_VERSION_MIN_SUPPORTED: u8 = 1;
+
+/// Prefix for the id given to a codec installed from a bare
+/// `PyCapsule`.
+///
+/// Every capsule reports the same type, so there is nothing on it to
+/// derive an id from. The rest of the id is a fresh UUID, minted when
+/// the codec is installed. Plans it encodes decode on the session that
+/// installed it, and on sessions cloned from that one, because cloning
+/// copies the chain along with its ids. On any other session the id is
+/// simply missing, and decoding says so.
+///
+/// Numbering the capsules instead — `anon:0`, `anon:1` — would be
+/// worse. Every session starts counting at zero, so one session's
+/// `anon:0` would be accepted by another session and decoded with
+/// whatever codec happened to be its own first capsule.
+pub(crate) const ANONYMOUS_CODEC_ID_PREFIX: &str = "anon:";
+
+/// Prefix for the id a `SessionContext` reports when its own codec stack
+/// is installed as an extension codec on another session.
+///
+/// An ordinary codec object takes its id from its class, which is the
+/// library's import path. That does not work for a session: every
+/// session is an instance of the same class, so they would all report
+/// the same id. Installing two sessions as codecs on one target would
+/// then collide, and a plan encoded by one would be decoded by the
+/// other. A session id is unique per session and stable, so the rest
+/// of the id carries that.
+pub(crate) const SESSION_CODEC_ID_PREFIX: &str = "session:";
+
+/// One installed codec plus the identity its payloads are tagged with.
+///
+/// The id is what makes dispatch order-independent. Keying on position
+/// in the chain — as `ComposedPhysicalExtensionCodec` does upstream —
+/// is sound only when both ends assemble the same list in the same
+/// order. That holds for the consumers upstream was written for, whose
+/// lists structurally cannot disagree: Ballista's 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 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 would name a different codec in the decoder as soon as
+/// install order differed.
+struct ChainEntry<C: ?Sized> {
+ id: Arc<str>,
+ codec: Arc<C>,
+}
+
+impl<C: ?Sized> Clone for ChainEntry<C> {
+ fn clone(&self) -> Self {
+ Self {
+ id: Arc::clone(&self.id),
+ codec: Arc::clone(&self.codec),
+ }
+ }
+}
+
+impl<C: ?Sized + std::fmt::Debug> std::fmt::Debug for ChainEntry<C> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("ChainEntry")
+ .field("id", &self.id)
+ .field("codec", &self.codec)
+ .finish()
+ }
+}
+
+/// Wrap `blob` in the envelope identifying the codec that produced it.
+///
+/// Layout: `DFPYCHN | version: u8 | id_len: u32 (LE) | id | blob`.
+fn write_chained_payload(buf: &mut Vec<u8>, codec_id: &str, blob: &[u8]) {
+ buf.extend_from_slice(CHAINED_PAYLOAD_FAMILY);
+ buf.push(CHAIN_WIRE_VERSION_CURRENT);
+ buf.extend_from_slice(&(codec_id.len() as u32).to_le_bytes());
+ buf.extend_from_slice(codec_id.as_bytes());
+ buf.extend_from_slice(blob);
+}
+
+/// Inspect the chained-payload envelope on `buf`.
+///
+/// * `Ok(None)` — no envelope. The payload came from the terminal
+/// codec, which writes unframed so a session with no extension
+/// codecs installed produces bytes identical to a build without
+/// codec chaining.
+/// * `Ok(Some((codec_id, blob)))` — the owning codec's id and its
+/// original bytes, byte-for-byte as it wrote them.
+fn read_chained_payload(buf: &[u8]) -> Result<Option<(&str, &[u8])>> {
+ if !buf.starts_with(CHAINED_PAYLOAD_FAMILY) {
+ return Ok(None);
+ }
+ let mut idx = CHAINED_PAYLOAD_FAMILY.len();
+ let Some(&version) = buf.get(idx) else {
+ return Err(datafusion::error::DataFusionError::Execution(
+ "Truncated extension codec payload: missing envelope version byte".to_string(),
+ ));
+ };
+ if !(CHAIN_WIRE_VERSION_MIN_SUPPORTED..=CHAIN_WIRE_VERSION_CURRENT).contains(&version) {
+ return Err(datafusion::error::DataFusionError::Execution(format!(
+ "Extension codec payload envelope version v{version}; this build supports \
+ v{CHAIN_WIRE_VERSION_MIN_SUPPORTED}..=v{CHAIN_WIRE_VERSION_CURRENT}. \
+ Align datafusion-python versions on sender and receiver."
+ )));
+ }
+ idx += 1;
+ let Some(len_bytes) = buf.get(idx..idx + 4) else {
+ return Err(datafusion::error::DataFusionError::Execution(
+ "Truncated extension codec payload: missing codec id length".to_string(),
+ ));
+ };
+ let id_len = u32::from_le_bytes(len_bytes.try_into().expect("4 bytes")) as usize;
+ idx += 4;
+ let Some(id_bytes) = buf.get(idx..idx + id_len) else {
+ return Err(datafusion::error::DataFusionError::Execution(
+ "Truncated extension codec payload: codec id shorter than its declared length"
+ .to_string(),
+ ));
+ };
+ let codec_id = std::str::from_utf8(id_bytes).map_err(|err| {
+ datafusion::error::DataFusionError::Execution(format!(
+ "Extension codec payload carries a non-UTF-8 codec id: {err}"
+ ))
+ })?;
+ Ok(Some((codec_id, &buf[idx + id_len..])))
+}
+
+/// Decode `buf` with the single codec that encoded it.
+///
+/// Three cases:
+///
+/// * **Empty `buf`** — nothing was encoded, so there is no tag to
+/// dispatch on and every codec is offered the empty buffer in install
+/// order. See [`chain_resolve_by_name`] for why that is sound here and
+/// why the case exists at all.
+/// * **Framed `buf`** — the envelope names its author, so exactly one
+/// codec is consulted and its error surfaces verbatim.
+/// * **Unframed non-empty `buf`** — the terminal codec wrote it.
+///
+/// Outside the empty case nothing is ever offered to a codec that did
+/// not write it, which is what stops a structurally similar prost
+/// message from decoding in the wrong library.
+///
+/// Do not replace this with a walk that hands `buf` to each codec until
+/// one returns `Ok`. That looks simpler and removes the envelope, and it
+/// is unsound: 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 an all-defaults
+/// message encodes to zero bytes that decode as anything. Requiring
+/// each codec to check a marker of its own does not fix it either: the
+/// codecs come from libraries this crate does not control, the natural
+/// implementation is `MyMessage::decode(buf)`, which has no marker to
+/// check and cannot decline, and one library skipping the convention is
+/// enough to lose someone else's payload. Upstream shipped that design
+/// and reverted it after a Parquet payload decoded as CSV —
+/// apache/datafusion#16980, fixed in #16986.
+fn chain_decode<C: ?Sized, R>(
+ chain: &[ChainEntry<C>],
+ terminal: &Arc<C>,
+ buf: &[u8],
+ what: &str,
+ f: impl Fn(&C, &[u8]) -> Result<R>,
+) -> Result<R> {
+ if buf.is_empty() {
+ return chain_resolve_by_name(chain, terminal, what, |codec| f(codec, buf));
+ }
+ let Some((codec_id, blob)) = read_chained_payload(buf)? else {
+ return f(terminal.as_ref(), buf);
+ };
+ let Some(entry) = chain.iter().find(|entry| &*entry.id == codec_id) else {
+ let installed = if chain.is_empty() {
+ "no extension codecs are installed on this session".to_string()
+ } else {
+ format!(
+ "installed: {}",
+ chain
+ .iter()
+ .map(|entry| entry.id.as_ref())
+ .collect::<Vec<_>>()
+ .join(", ")
+ )
+ };
+ let hint = if codec_id.starts_with(ANONYMOUS_CODEC_ID_PREFIX) {
+ ". This payload was written by a codec installed from a bare PyCapsule, which \
+ carries no portable identity. Pass `codec_id=` when installing it if plans must \
+ cross sessions."
+ } else {
+ ""
+ };
+ return Err(datafusion::error::DataFusionError::Execution(format!(
+ "{what} was encoded by extension codec '{codec_id}', which is not installed on \
+ this session ({installed}){hint}"
+ )));
+ };
+ f(entry.codec.as_ref(), blob)
+}
+
+/// Resolve an object carrying no payload, by consulting each codec.
+///
+/// Used only where DataFusion encodes by name: `try_encode_udf` and its
+/// aggregate/window siblings return `Ok` writing nothing, and the
+/// decoder then tries the `FunctionRegistry` first and the codec second
+/// (`from_proto.rs`, the `None => ctx.udf(..).or_else(..)` arm). A codec
+/// whose functions are reconstructible from the name alone is reached
+/// through that arm and must still be offered the empty buffer.
+///
+/// This is the one place dispatch cannot be tagged — there are no bytes
+/// to tag. It is not the hazard that tagging exists to remove: the
+/// question asked here is "do you own the function named `x`", which is
+/// name-scoped and answerable, not "do these bytes happen to parse as
+/// your message type". Two codecs disagreeing requires them to claim
+/// the same function name, which already collides in the registry.
+fn chain_resolve_by_name<C: ?Sized, R>(
+ chain: &[ChainEntry<C>],
+ terminal: &Arc<C>,
+ what: &str,
+ f: impl Fn(&C) -> Result<R>,
+) -> Result<R> {
+ let mut errors: Vec<datafusion::error::DataFusionError> = Vec::new();
+ for entry in chain {
+ match f(entry.codec.as_ref()) {
+ Ok(value) => return Ok(value),
+ Err(err) => errors.push(err),
+ }
+ }
+ match f(terminal.as_ref()) {
+ Ok(value) => Ok(value),
+ Err(err) => {
+ errors.push(err);
+ Err(aggregate_chain_errors(what, errors))
+ }
+ }
+}
+
+/// Collapse per-codec failures into one error. A single failure is
+/// returned as-is so a session with no extension codecs behaves exactly
+/// like a build without codec chaining.
+fn aggregate_chain_errors(
+ what: &str,
+ mut errors: Vec<datafusion::error::DataFusionError>,
+) -> datafusion::error::DataFusionError {
+ match errors.len() {
+ 0 => datafusion::error::DataFusionError::Internal(format!(
+ "Empty extension codec chain while handling {what}"
+ )),
+ 1 => errors.swap_remove(0),
+ _ => {
+ let joined = errors
+ .iter()
+ .map(|err| err.to_string())
+ .collect::<Vec<_>>()
+ .join("; ");
+ datafusion::error::DataFusionError::Execution(format!(
+ "No installed extension codec handled {what}: {joined}"
+ ))
+ }
+ }
+}
+
+/// Encode through the chain, tagging the payload with its author.
+///
+/// Entries are consulted in install order and the first one to write
+/// bytes wins, so installing a codec can only claim objects no
+/// earlier codec claimed. Adding a library therefore never changes how
+/// an already-installed library's objects encode.
+///
+/// Each codec encodes into a scratch buffer so a failed attempt cannot
+/// leave partial bytes behind. `Ok` with an empty buffer is "no
+/// opinion" rather than a claim, so the walk continues; if nothing
+/// writes bytes the result is `Ok` with nothing written, which is
+/// DataFusion's encode-by-name signal. Framing that empty result would
+/// set `fun_definition` and permanently skip the registry lookup the
+/// decoder does first.
+///
+/// The terminal codec writes unframed, so a session with no extension
+/// codecs is byte-compatible with a build predating the chain.
+fn chain_encode<C: ?Sized>(
+ chain: &[ChainEntry<C>],
+ terminal: &Arc<C>,
+ buf: &mut Vec<u8>,
+ what: &str,
+ f: impl Fn(&C, &mut Vec<u8>) -> Result<()>,
+) -> Result<()> {
+ let mut saw_empty_ok = false;
+ let mut errors: Vec<datafusion::error::DataFusionError> = Vec::new();
+ for entry in chain {
+ let mut scratch = Vec::new();
+ match f(entry.codec.as_ref(), &mut scratch) {
+ Ok(()) if !scratch.is_empty() => {
+ write_chained_payload(buf, &entry.id, &scratch);
+ return Ok(());
+ }
+ Ok(()) => saw_empty_ok = true,
+ Err(err) => errors.push(err),
+ }
+ }
+ let mut scratch = Vec::new();
+ match f(terminal.as_ref(), &mut scratch) {
+ Ok(()) if !scratch.is_empty() => {
+ buf.extend_from_slice(&scratch);
+ return Ok(());
+ }
+ Ok(()) => saw_empty_ok = true,
+ Err(err) => errors.push(err),
+ }
+ if saw_empty_ok {
+ return Ok(());
+ }
+ Err(aggregate_chain_errors(what, errors))
+}
+
/// `LogicalExtensionCodec` parked on every `SessionContext`. Holds
/// the Python-aware encoding hooks for logical-layer types
/// (`LogicalPlan`, `Expr`) and delegates everything it does not
-/// handle to the composable `inner` codec — typically
-/// `DefaultLogicalExtensionCodec`, or a downstream FFI codec
-/// installed via `SessionContext.with_logical_extension_codec(...)`.
+/// handle to a chain of composable codecs. Each downstream FFI codec
+/// installed via `SessionContext.with_logical_extension_codec(...)` is
+/// appended to the chain, and `terminal` — normally
+/// `DefaultLogicalExtensionCodec` — handles whatever no installed codec
+/// claims.
+///
+/// Every payload an installed codec writes is wrapped in an envelope
+/// naming that codec (see [`write_chained_payload`]), so decoding
+/// consults exactly the codec that encoded it. Dispatch does not depend
+/// on a codec recognizing and rejecting foreign payloads, which is not
+/// something a codec can reliably do: a prost message decodes cleanly
+/// from another message's bytes whenever their leading field numbers and
+/// wire types line up.
///
/// Sitting at the top of the session's logical codec stack means
/// every serializer that reads `session.logical_codec()` automatically
@@ -241,22 +584,71 @@
/// the weak `FFI_TaskContextProvider` valid is instead a matter of never
/// replacing the session's `Arc<SessionContext>`; see
/// `PySessionContext::set_session_query_planner`.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
pub struct PythonLogicalCodec {
- inner: Arc<dyn LogicalExtensionCodec>,
+ chain: Vec<ChainEntry<dyn LogicalExtensionCodec>>,
+ terminal: Arc<dyn LogicalExtensionCodec>,
python_udf_inlining: bool,
}
impl PythonLogicalCodec {
+ /// Build a codec with no installed extension codecs and `inner` as
+ /// the terminal fallback. `inner` is not part of the keyed chain and
+ /// its payloads are written unframed, so a context built this way
+ /// serializes byte-identically to one with no chaining at all.
pub fn new(inner: Arc<dyn LogicalExtensionCodec>) -> Self {
Self {
- inner,
+ chain: Vec::new(),
+ terminal: inner,
python_udf_inlining: true,
}
}
- pub fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> {
- &self.inner
+ /// Return a copy of this codec with `codec` appended to the chain
+ /// under `id`, preserving the Python-UDF-inlining setting.
+ ///
+ /// Appending rather than prepending keeps the operation additive:
+ /// the new codec is consulted for encoding only after every codec
+ /// already installed, so it can claim objects nothing else claimed
+ /// but cannot take over an existing library's objects.
+ pub fn with_additional_codec(
+ &self,
+ id: impl Into<Arc<str>>,
+ codec: Arc<dyn LogicalExtensionCodec>,
+ ) -> Self {
+ let mut chain = self.chain.clone();
+ chain.push(ChainEntry {
+ id: id.into(),
+ codec,
+ });
+ Self {
+ chain,
+ terminal: Arc::clone(&self.terminal),
+ python_udf_inlining: self.python_udf_inlining,
+ }
+ }
+
+ /// Ids of the installed extension codecs, in install order.
+ ///
+ /// The terminal codec is not listed: it is not addressable by id
+ /// because its payloads are written unframed.
+ pub fn codec_ids(&self) -> Vec<&str> {
+ self.chain.iter().map(|entry| entry.id.as_ref()).collect()
+ }
+
+ /// Installed extension codecs paired with their ids, in install
+ /// order. Restores the inspection that the removed `inner()`
+ /// accessor provided, and exposes the id dispatch keys along with it.
+ pub fn codecs(&self) -> Vec<(&str, &Arc<dyn LogicalExtensionCodec>)> {
+ self.chain
+ .iter()
+ .map(|entry| (entry.id.as_ref(), &entry.codec))
+ .collect()
+ }
+
+ /// Terminal codec consulted when no installed codec claims an object.
+ pub fn terminal(&self) -> &Arc<dyn LogicalExtensionCodec> {
+ &self.terminal
}
/// Toggle inline encoding of Python UDFs. See
@@ -297,11 +689,23 @@
inputs: &[LogicalPlan],
ctx: &TaskContext,
) -> Result<Extension> {
- self.inner.try_decode(buf, inputs, ctx)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an extension logical plan node",
+ |codec, buf| codec.try_decode(buf, inputs, ctx),
+ )
}
fn try_encode(&self, node: &Extension, buf: &mut Vec<u8>) -> Result<()> {
- self.inner.try_encode(node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an extension logical plan node",
+ |codec, buf| codec.try_encode(node, buf),
+ )
}
fn try_decode_table_provider(
@@ -311,8 +715,13 @@
schema: SchemaRef,
ctx: &TaskContext,
) -> Result<Arc<dyn TableProvider>> {
- self.inner
- .try_decode_table_provider(buf, table_ref, schema, ctx)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a table provider",
+ |codec, buf| codec.try_decode_table_provider(buf, table_ref, Arc::clone(&schema), ctx),
+ )
}
fn try_encode_table_provider(
@@ -321,7 +730,13 @@
node: Arc<dyn TableProvider>,
buf: &mut Vec<u8>,
) -> Result<()> {
- self.inner.try_encode_table_provider(table_ref, node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a table provider",
+ |codec, buf| codec.try_encode_table_provider(table_ref, Arc::clone(&node), buf),
+ )
}
fn try_decode_file_format(
@@ -329,7 +744,13 @@
buf: &[u8],
ctx: &TaskContext,
) -> Result<Arc<dyn FileFormatFactory>> {
- self.inner.try_decode_file_format(buf, ctx)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a file format",
+ |codec, buf| codec.try_decode_file_format(buf, ctx),
+ )
}
fn try_encode_file_format(
@@ -337,14 +758,26 @@
buf: &mut Vec<u8>,
node: Arc<dyn FileFormatFactory>,
) -> Result<()> {
- self.inner.try_encode_file_format(buf, node)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a file format",
+ |codec, buf| codec.try_encode_file_format(buf, Arc::clone(&node)),
+ )
}
fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? {
return Ok(());
}
- self.inner.try_encode_udf(node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a scalar UDF",
+ |codec, buf| codec.try_encode_udf(node, buf),
+ )
}
fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
@@ -355,14 +788,26 @@
} else {
refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
}
- self.inner.try_decode_udf(name, buf)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a scalar UDF",
+ |codec, buf| codec.try_decode_udf(name, buf),
+ )
}
fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {
if self.python_udf_inlining && try_encode_python_udaf(node, buf)? {
return Ok(());
}
- self.inner.try_encode_udaf(node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an aggregate UDF",
+ |codec, buf| codec.try_encode_udaf(node, buf),
+ )
}
fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> {
@@ -373,14 +818,26 @@
} else {
refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
}
- self.inner.try_decode_udaf(name, buf)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an aggregate UDF",
+ |codec, buf| codec.try_decode_udaf(name, buf),
+ )
}
fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> {
if self.python_udf_inlining && try_encode_python_udwf(node, buf)? {
return Ok(());
}
- self.inner.try_encode_udwf(node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a window UDF",
+ |codec, buf| codec.try_encode_udwf(node, buf),
+ )
}
fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> {
@@ -391,13 +848,19 @@
} else {
refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
}
- self.inner.try_decode_udwf(name, buf)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a window UDF",
+ |codec, buf| codec.try_decode_udwf(name, buf),
+ )
}
}
/// Strict-mode gate: if `buf` is a well-framed inline payload for
/// `family`, return the strict-refusal error; otherwise return
-/// `Ok(())` so the caller can delegate to its `inner` codec.
+/// `Ok(())` so the caller can delegate to its codec chain.
///
/// Routing through [`read_framed_payload`] (rather than a bare
/// `starts_with` probe) means malformed inline bytes — wrong
@@ -442,7 +905,8 @@
/// `PhysicalExtensionCodec` mirror of [`PythonLogicalCodec`] parked
/// on the same `SessionContext`. Carries the Python-aware encoding
/// hooks for physical-layer types (`ExecutionPlan`, `PhysicalExpr`)
-/// and delegates the rest to `inner`.
+/// and delegates the rest to the composable codec chain (see
+/// [`PythonLogicalCodec`] for chain ordering and dispatch rules).
///
/// The `PhysicalExtensionCodec` trait has its own `try_encode_udf`
/// / `try_decode_udf` pair distinct from the logical one, so a
@@ -454,22 +918,59 @@
///
/// Like [`PythonLogicalCodec`], this does not retain the session it was built
/// from; see that type for why.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
pub struct PythonPhysicalCodec {
- inner: Arc<dyn PhysicalExtensionCodec>,
+ chain: Vec<ChainEntry<dyn PhysicalExtensionCodec>>,
+ terminal: Arc<dyn PhysicalExtensionCodec>,
python_udf_inlining: bool,
}
impl PythonPhysicalCodec {
+ /// See [`PythonLogicalCodec::new`]; `inner` is the terminal codec
+ /// rather than a chain entry.
pub fn new(inner: Arc<dyn PhysicalExtensionCodec>) -> Self {
Self {
- inner,
+ chain: Vec::new(),
+ terminal: inner,
python_udf_inlining: true,
}
}
- pub fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> {
- &self.inner
+ /// Return a copy of this codec with `codec` appended to the chain
+ /// under `id`. See [`PythonLogicalCodec::with_additional_codec`].
+ pub fn with_additional_codec(
+ &self,
+ id: impl Into<Arc<str>>,
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ ) -> Self {
+ let mut chain = self.chain.clone();
+ chain.push(ChainEntry {
+ id: id.into(),
+ codec,
+ });
+ Self {
+ chain,
+ terminal: Arc::clone(&self.terminal),
+ python_udf_inlining: self.python_udf_inlining,
+ }
+ }
+
+ /// Ids of the installed extension codecs, in install order.
+ pub fn codec_ids(&self) -> Vec<&str> {
+ self.chain.iter().map(|entry| entry.id.as_ref()).collect()
+ }
+
+ /// Installed extension codecs paired with their ids, in install order.
+ pub fn codecs(&self) -> Vec<(&str, &Arc<dyn PhysicalExtensionCodec>)> {
+ self.chain
+ .iter()
+ .map(|entry| (entry.id.as_ref(), &entry.codec))
+ .collect()
+ }
+
+ /// Terminal codec consulted when no installed codec claims an object.
+ pub fn terminal(&self) -> &Arc<dyn PhysicalExtensionCodec> {
+ &self.terminal
}
/// Toggle inline encoding of Python UDFs on this physical codec.
@@ -500,7 +1001,13 @@
ctx: &TaskContext,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
- self.inner.try_decode(buf, inputs, ctx, proto_converter)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an execution plan",
+ |codec, buf| codec.try_decode(buf, inputs, ctx, proto_converter),
+ )
}
fn try_encode(
@@ -509,14 +1016,26 @@
buf: &mut Vec<u8>,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
- self.inner.try_encode(node, buf, proto_converter)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an execution plan",
+ |codec, buf| codec.try_encode(Arc::clone(&node), buf, proto_converter),
+ )
}
fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? {
return Ok(());
}
- self.inner.try_encode_udf(node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a scalar UDF",
+ |codec, buf| codec.try_encode_udf(node, buf),
+ )
}
fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
@@ -527,7 +1046,13 @@
} else {
refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
}
- self.inner.try_decode_udf(name, buf)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a scalar UDF",
+ |codec, buf| codec.try_decode_udf(name, buf),
+ )
}
fn try_encode_expr(
@@ -536,7 +1061,13 @@
buf: &mut Vec<u8>,
ctx: &PhysicalExprEncodeCtx<'_>,
) -> Result<()> {
- self.inner.try_encode_expr(node, buf, ctx)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a physical expression",
+ |codec, buf| codec.try_encode_expr(node, buf, ctx),
+ )
}
fn try_decode_expr(
@@ -545,14 +1076,26 @@
inputs: &[Arc<dyn PhysicalExpr>],
ctx: &PhysicalExprDecodeCtx<'_>,
) -> Result<Arc<dyn PhysicalExpr>> {
- self.inner.try_decode_expr(buf, inputs, ctx)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a physical expression",
+ |codec, buf| codec.try_decode_expr(buf, inputs, ctx),
+ )
}
fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()> {
if self.python_udf_inlining && try_encode_python_udaf(node, buf)? {
return Ok(());
}
- self.inner.try_encode_udaf(node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an aggregate UDF",
+ |codec, buf| codec.try_encode_udaf(node, buf),
+ )
}
fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> {
@@ -563,14 +1106,26 @@
} else {
refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
}
- self.inner.try_decode_udaf(name, buf)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "an aggregate UDF",
+ |codec, buf| codec.try_decode_udaf(name, buf),
+ )
}
fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()> {
if self.python_udf_inlining && try_encode_python_udwf(node, buf)? {
return Ok(());
}
- self.inner.try_encode_udwf(node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a window UDF",
+ |codec, buf| codec.try_encode_udwf(node, buf),
+ )
}
fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> {
@@ -581,7 +1136,13 @@
} else {
refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
}
- self.inner.try_decode_udwf(name, buf)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a window UDF",
+ |codec, buf| codec.try_decode_udwf(name, buf),
+ )
}
}
@@ -598,7 +1159,7 @@
/// `Ok(true)` when the payload (`DFPYUDF` family prefix, version byte,
/// cloudpickled tuple) was written and the caller should skip its
/// inner codec. Returns `Ok(false)` for any non-Python UDF, signalling
-/// the caller to delegate to its `inner`.
+/// the caller to delegate to its codec chain.
pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<bool> {
let Some(py_udf) = node.inner().downcast_ref::<PythonFunctionScalarUDF>() else {
return Ok(false);
@@ -613,7 +1174,7 @@
/// Decode an inline Python scalar UDF payload. Returns `Ok(None)`
/// when `buf` does not carry the `DFPYUDF` family prefix, signalling
-/// the caller to delegate to its `inner` codec (and eventually the
+/// the caller to delegate to its codec chain (and eventually the
/// `FunctionRegistry`).
pub(crate) fn try_decode_python_scalar_udf(buf: &[u8]) -> Result<Option<Arc<ScalarUDF>>> {
if !buf.starts_with(PY_SCALAR_UDF_FAMILY) {
@@ -1039,137 +1600,3 @@
volatility,
))
}
-
-#[cfg(test)]
-mod wire_header_tests {
- use super::*;
-
- const TEST_PY: (u8, u8) = (3, 12);
-
- #[test]
- fn strip_returns_none_when_family_absent() {
- let buf = b"OTHER_PAYLOAD";
- assert!(matches!(
- strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY),
- Ok(None)
- ));
- }
-
- #[test]
- fn strip_errors_on_truncated_version_byte() {
- let buf = PY_SCALAR_UDF_FAMILY;
- let err = strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
- assert!(format!("{err}").contains("missing wire-format version byte"));
- }
-
- #[test]
- fn strip_errors_on_too_new_version() {
- let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
- buf.push(WIRE_VERSION_CURRENT.saturating_add(1));
- buf.push(TEST_PY.0);
- buf.push(TEST_PY.1);
- buf.extend_from_slice(b"payload");
- let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
- let msg = format!("{err}");
- assert!(msg.contains("wire-format version v"));
- assert!(msg.contains("supports"));
- assert!(msg.contains("Align datafusion-python versions"));
- }
-
- #[test]
- fn strip_errors_on_too_old_version() {
- if WIRE_VERSION_MIN_SUPPORTED == 0 {
- return;
- }
- let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
- buf.push(WIRE_VERSION_MIN_SUPPORTED - 1);
- buf.push(TEST_PY.0);
- buf.push(TEST_PY.1);
- buf.extend_from_slice(b"payload");
- assert!(strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).is_err());
- }
-
- #[test]
- fn strip_errors_on_truncated_py_major() {
- let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
- buf.push(WIRE_VERSION_CURRENT);
- let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
- assert!(format!("{err}").contains("missing Python major version byte"));
- }
-
- #[test]
- fn strip_errors_on_truncated_py_minor() {
- let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
- buf.push(WIRE_VERSION_CURRENT);
- buf.push(TEST_PY.0);
- let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
- assert!(format!("{err}").contains("missing Python minor version byte"));
- }
-
- #[test]
- fn strip_errors_on_py_minor_mismatch() {
- let mut buf = Vec::new();
- write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, (3, 11));
- buf.extend_from_slice(b"payload");
- let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", (3, 12)).unwrap_err();
- let msg = format!("{err}");
- assert!(msg.contains("Python 3.11"));
- assert!(msg.contains("Python 3.12"));
- assert!(msg.contains("not portable across Python minor versions"));
- }
-
- #[test]
- fn strip_errors_on_py_major_mismatch() {
- let mut buf = Vec::new();
- write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, (3, 12));
- buf.extend_from_slice(b"payload");
- assert!(strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", (4, 0)).is_err());
- }
-
- #[test]
- fn write_then_strip_round_trips_scalar_payload() {
- let mut buf = Vec::new();
- write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, TEST_PY);
- buf.extend_from_slice(b"scalar-payload");
-
- let payload = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY)
- .unwrap()
- .unwrap();
- assert_eq!(payload, b"scalar-payload");
- }
-
- #[test]
- fn write_then_strip_round_trips_agg_payload() {
- let mut buf = Vec::new();
- write_wire_header(&mut buf, PY_AGG_UDF_FAMILY, TEST_PY);
- buf.extend_from_slice(b"agg-payload");
-
- let payload = strip_wire_header(&buf, PY_AGG_UDF_FAMILY, "aggregate UDF", TEST_PY)
- .unwrap()
- .unwrap();
- assert_eq!(payload, b"agg-payload");
- }
-
- #[test]
- fn write_then_strip_round_trips_window_payload() {
- let mut buf = Vec::new();
- write_wire_header(&mut buf, PY_WINDOW_UDF_FAMILY, TEST_PY);
- buf.extend_from_slice(b"window-payload");
-
- let payload = strip_wire_header(&buf, PY_WINDOW_UDF_FAMILY, "window UDF", TEST_PY)
- .unwrap()
- .unwrap();
- assert_eq!(payload, b"window-payload");
- }
-
- #[test]
- fn strip_does_not_match_a_different_family() {
- let mut buf = Vec::new();
- write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, TEST_PY);
- buf.extend_from_slice(b"payload");
- assert!(matches!(
- strip_wire_header(&buf, PY_WINDOW_UDF_FAMILY, "window UDF", TEST_PY),
- Ok(None)
- ));
- }
-}
diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs
index 75bfed6..84182ff 100644
--- a/crates/core/src/context.rs
+++ b/crates/core/src/context.rs
@@ -75,7 +75,9 @@
use crate::catalog::{
PyCatalog, PyCatalogList, RustWrappedPyCatalogProvider, RustWrappedPyCatalogProviderList,
};
-use crate::codec::{PythonLogicalCodec, PythonPhysicalCodec};
+use crate::codec::{
+ ANONYMOUS_CODEC_ID_PREFIX, PythonLogicalCodec, PythonPhysicalCodec, SESSION_CODEC_ID_PREFIX,
+};
use crate::common::data_type::PyScalarValue;
use crate::common::df_schema::PyDFSchema;
use crate::dataframe::PyDataFrame;
@@ -1422,6 +1424,24 @@
PyCapsule::new_with_value(py, ffi_ctx_provider, cr"datafusion_task_context_provider")
}
+ /// Identity this session's codec stack carries when it is itself installed
+ /// as an extension codec on another session.
+ ///
+ /// Without it the id would be derived from the class, which every session
+ /// shares: two sessions installed on one target would collide at install
+ /// time, and a payload written through one would resolve to the other on
+ /// decode. See [`SESSION_CODEC_ID_PREFIX`].
+ ///
+ /// Handles derived from one session — `with_python_udf_inlining`,
+ /// `with_logical_extension_codec` — report the same id even though their
+ /// codec chains differ, so installing two of them on one target is
+ /// refused. That is the intended answer: their payloads would be
+ /// indistinguishable on decode.
+ #[getter]
+ pub fn __datafusion_codec_id__(&self) -> String {
+ format!("{SESSION_CODEC_ID_PREFIX}{}", self.ctx.session_id())
+ }
+
/// `session` exists so this matches the protocol an extension library
/// implements, where the argument is how the library reaches the session
/// it is being installed on. A session already is one, so it is ignored.
@@ -1457,22 +1477,28 @@
create_query_planner_capsule(py, &ffi)
}
+ #[pyo3(signature = (codec, codec_id=None))]
pub fn with_logical_extension_codec<'py>(
slf: &Bound<'py, Self>,
codec: Bound<'py, PyAny>,
+ codec_id: Option<String>,
) -> PyDataFusionResult<Self> {
+ let id = {
+ let this = slf.borrow();
+ resolve_codec_id(&codec, codec_id, &this.logical_codec.codec_ids())?
+ };
let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
let this = slf.borrow();
- // Carry the receiver's inlining setting over. `PythonLogicalCodec::new`
- // defaults it to on, so building the replacement without this would
- // silently re-enable inline Python UDF encoding on a context that had
- // opted out with `with_python_udf_inlining(enabled=False)`.
- let logical_codec = Arc::new(
- PythonLogicalCodec::new(inner)
- .with_python_udf_inlining(this.logical_codec.python_udf_inlining()),
- );
+ // Append rather than replace: previously installed codecs stay active,
+ // and every payload this one writes is tagged with `id` so decoding
+ // reaches it directly rather than by trying codecs in turn. Appending
+ // also carries the receiver's inlining setting over, which a fresh
+ // `PythonLogicalCodec::new` would not — it defaults inlining to on, and
+ // would silently re-enable inline Python UDF encoding on a context that
+ // had opted out with `with_python_udf_inlining(enabled=False)`.
+ let logical_codec = Arc::new(this.logical_codec.with_additional_codec(id, inner));
let derived = Self {
ctx: Arc::clone(&this.ctx),
logical_codec,
@@ -1484,6 +1510,27 @@
Ok(derived)
}
+ /// Ids of the logical extension codecs installed on this session, in
+ /// install order — the same order encoding consults them in, and the keys
+ /// a payload names when it is decoded.
+ pub fn logical_extension_codec_ids(&self) -> Vec<String> {
+ self.logical_codec
+ .codec_ids()
+ .into_iter()
+ .map(str::to_string)
+ .collect()
+ }
+
+ /// Ids of the physical extension codecs installed on this session.
+ /// See [`Self::logical_extension_codec_ids`].
+ pub fn physical_extension_codec_ids(&self) -> Vec<String> {
+ self.physical_codec
+ .codec_ids()
+ .into_iter()
+ .map(str::to_string)
+ .collect()
+ }
+
/// See [`Self::__datafusion_logical_extension_codec__`] for `session`.
#[pyo3(signature = (session=None))]
pub fn __datafusion_physical_extension_codec__<'py>(
@@ -1495,19 +1542,23 @@
create_physical_extension_capsule(py, self.ffi_physical_codec().as_ref())
}
+ #[pyo3(signature = (codec, codec_id=None))]
pub fn with_physical_extension_codec<'py>(
slf: &Bound<'py, Self>,
codec: Bound<'py, PyAny>,
+ codec_id: Option<String>,
) -> PyDataFusionResult<Self> {
+ let id = {
+ let this = slf.borrow();
+ resolve_codec_id(&codec, codec_id, &this.physical_codec.codec_ids())?
+ };
let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
let inner: Arc<dyn PhysicalExtensionCodec> = (&inner_ffi).into();
let this = slf.borrow();
- // See `with_logical_extension_codec` for why the flag is carried over.
- let physical_codec = Arc::new(
- PythonPhysicalCodec::new(inner)
- .with_python_udf_inlining(this.physical_codec.python_udf_inlining()),
- );
+ // See `with_logical_extension_codec` for why this appends rather than
+ // replaces, and why that is also what carries the inlining flag over.
+ let physical_codec = Arc::new(this.physical_codec.with_additional_codec(id, inner));
let derived = Self {
ctx: Arc::clone(&this.ctx),
logical_codec: Arc::clone(&this.logical_codec),
@@ -1525,7 +1576,7 @@
// already inlines would otherwise rebind the session's planner to this
// handle's codecs, and callers routinely discard the result. Returning
// the codecs as-is is observationally equivalent to the rebuild below,
- // which wraps the same inner codec in a fresh `Python*Codec`.
+ // which clones the same codec chain and only flips the flag.
if self.logical_codec.python_udf_inlining() == enabled
&& self.physical_codec.python_udf_inlining() == enabled
{
@@ -1537,11 +1588,15 @@
}
let logical_codec = Arc::new(
- PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner()))
+ self.logical_codec
+ .as_ref()
+ .clone()
.with_python_udf_inlining(enabled),
);
let physical_codec = Arc::new(
- PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner()))
+ self.physical_codec
+ .as_ref()
+ .clone()
.with_python_udf_inlining(enabled),
);
let derived = Self {
@@ -1705,6 +1760,80 @@
}
}
+/// Determine the wire identity to tag an installed codec's payloads with.
+///
+/// Every payload a chained codec writes carries this string, and decoding
+/// dispatches on it, so it has to name the same codec in the process that
+/// decodes as it did in the process that encoded. Resolution order:
+///
+/// 1. An explicit `codec_id` argument.
+/// 2. `codec.__datafusion_codec_id__`, letting a library pin its own identity
+/// so a class rename does not invalidate previously encoded plans, and so
+/// two instances of one class can own disjoint slices of the wire format.
+/// `SessionContext` pins its own through this arm — see
+/// [`PySessionContext::__datafusion_codec_id__`] — because the class-derived
+/// id below would name every session at once.
+/// 3. The exporting object's `module.QualName`, which is the library's own
+/// import path and therefore 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 — every capsule
+/// reports the same type — so mint a fresh random id. Payloads tagged this
+/// way decode correctly within the session lineage that installed the
+/// codec, because the chain is cloned along with the id, and fail with a
+/// pointed error everywhere else. Randomness is the point: an id drawn from
+/// a namespace another session can mint the same value from — a counter, a
+/// chain position — would let an unrelated codec answer for these bytes.
+///
+/// An id already in use is rejected rather than shadowed. Two codecs sharing an
+/// id are indistinguishable on decode, and the API cannot tell whether two
+/// instances of one class write the same wire format — so the ambiguity is
+/// surfaced at install time, where the caller can resolve it, instead of at
+/// decode time, where it would pick whichever entry came first.
+fn resolve_codec_id(
+ codec: &Bound<'_, PyAny>,
+ explicit: Option<String>,
+ existing: &[&str],
+) -> PyResult<String> {
+ let id = derive_codec_id(codec, explicit)?;
+ if existing.contains(&id.as_str()) {
+ return Err(PyValueError::new_err(format!(
+ "An extension codec with id '{id}' is already installed on this session. Two \
+ codecs cannot share an id, because a payload names its codec by id when it is \
+ decoded. Pass `codec_id=` to give this one a distinct identity."
+ )));
+ }
+ Ok(id)
+}
+
+fn derive_codec_id(codec: &Bound<'_, PyAny>, explicit: Option<String>) -> PyResult<String> {
+ if let Some(id) = explicit {
+ return Ok(id);
+ }
+ if let Ok(declared) = codec.getattr("__datafusion_codec_id__")
+ && !declared.is_none()
+ {
+ return declared.extract::<String>();
+ }
+ if codec.is_instance_of::<PyCapsule>() {
+ return Ok(format!(
+ "{ANONYMOUS_CODEC_ID_PREFIX}{}",
+ Uuid::new_v4()
+ .simple()
+ .encode_lower(&mut Uuid::encode_buffer())
+ ));
+ }
+ let ty = codec.get_type();
+ let module = ty
+ .getattr("__module__")
+ .and_then(|m| m.extract::<String>())
+ .unwrap_or_else(|_| "<unknown>".to_string());
+ let qualname = ty
+ .getattr("__qualname__")
+ .and_then(|q| q.extract::<String>())
+ .or_else(|_| ty.name().and_then(|n| n.extract::<String>()))?;
+ Ok(format!("{module}.{qualname}"))
+}
+
pub fn parse_file_compression_type(
file_compression_type: Option<String>,
) -> Result<FileCompressionType, PyErr> {
diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md
index 31cd939..d86858a 100644
--- a/docs/source/contributor-guide/ffi.md
+++ b/docs/source/contributor-guide/ffi.md
@@ -248,10 +248,96 @@
process-local tokens to demonstrate ownership; production codecs should serialize
durable metadata instead.
-The current Python API has one external logical codec and one external physical codec.
-Installing another codec replaces the prior codec rather than composing a registry.
-The example therefore has one external codec owner, and the planner uses built-in
-physical nodes. Install the provider codecs before the planner where possible.
+### Composable codecs
+
+Extension codecs compose. Each call to `with_logical_extension_codec` or
+`with_physical_extension_codec` appends the codec to the session's codec chain
+rather than replacing prior codecs.
+
+**Nothing is asked of the codec itself.** Implement `LogicalExtensionCodec` or
+`PhysicalExtensionCodec` exactly as you would for a session that installs only
+yours. When your codec writes bytes into a serialized plan, datafusion-python
+records which codec wrote them, and strips that record off again before handing the
+bytes back. So your codec receives, byte for byte, the payload it wrote, and is
+never offered a payload another codec wrote.
+
+A codec that also ships to hosts which dispatch differently may still want its own
+guard against foreign payloads. Keeping one is fine; it is simply not needed for the
+datafusion-python path.
+
+That record is the codec's **id**: a short string stored inside the plan, naming the
+codec that wrote each payload. Because plans are decoded in another process — or
+another program — the id has to name the same codec there as it did where the plan
+was written.
+
+Ids are assigned for you. A codec's id is normally its exporting class's import
+path, such as `my_library.Codec`, which is what you will see in
+`logical_extension_codec_ids()` and in decode errors. You choose one yourself in
+three cases:
+
+- **Two instances of one class.** Both get the same id, so the second install
+ raises `ValueError`. Pass `codec_id=` to tell them apart.
+- **A bare `PyCapsule`.** A capsule has no class to take a name from, so it gets an
+ id private to the session that installed it. Plans it encodes fail with a clear
+ error on any other session, rather than being decoded by the wrong codec. Pass
+ `codec_id=` if those plans have to cross sessions.
+- **A class you intend to rename.** The id follows the class name, so renaming stops
+ older plans from decoding. Declare `__datafusion_codec_id__` on the exporting
+ object to pin an id that survives the rename.
+
+`SessionContext.logical_extension_codec_ids()` and its physical counterpart list the
+ids installed on a session, which is also what a decode failure names.
+
+Installing one context's codec stack on another session composes the two sessions
+rather than copying codecs out of one: the imported codecs resolve their task context
+against the original and stop working when it is dropped — see
+[One session, one `Arc<SessionContext>`](#one-session-one-arcsessioncontext). Pass
+the context itself rather than the capsule it exports, so its codecs get an id that
+other sessions can decode.
+
+Because decoding keys off the id rather than install position, registration order
+between independent libraries does not affect decoding at all. It is visible only
+on encoding, where codecs are consulted in install order and the first to claim an
+object wins — so installing a library can claim objects nothing else claimed, but
+never takes over an object an earlier codec was already encoding. Two libraries
+that each own tables, functions, and a planner register like this:
+
+```python
+ctx = SessionContext(config)
+
+# Codecs from both libraries. Order between libraries does not matter.
+ctx = ctx.with_logical_extension_codec(lib_a.codec())
+ctx = ctx.with_logical_extension_codec(lib_b.codec())
+ctx = ctx.with_physical_extension_codec(lib_a.physical_codec())
+ctx = ctx.with_physical_extension_codec(lib_b.physical_codec())
+
+# A session holds one planner, so layering is explicit delegation. Install the
+# codecs first: the fallback captured here keeps the codecs it was exported
+# with. See "Rebinding a planner's codecs is one level deep" below.
+ctx.set_query_planner(lib_a.Planner())
+ctx.set_query_planner(lib_b.Planner(fallback=ctx.__datafusion_query_planner__()))
+
+# Tables and functions — any time before the first query.
+ctx.register_table("t", lib_a.TableProvider())
+ctx.register_udf(udf(lib_b.SomeUDF()))
+```
+
+A codec may own functions that need no payload at all, where the name is the whole
+encoding: `try_encode_udf` writes nothing and `try_decode_udf` rebuilds the function
+from `name`. That is supported and needs no id, because an `Ok` with an empty
+buffer is read as "no opinion" and passes the object to the next codec.
+`NameOnlyUdfCodec` in the FFI example is the worked case. Anything no installed
+codec claims falls through to `Default{Logical,Physical}ExtensionCodec`.
+
+This is the one case where your decoder is consulted about something you may not
+own, because an empty payload has no id to route on. `try_decode_udf` and
+its aggregate and window siblings can therefore be called with an empty `buf` and a
+`name` belonging to another library. Decide from `name` and return an error if it is
+not yours; do not assume `buf` is non-empty.
+
+The framing itself — how an id is stored alongside a payload and routed back, and the two cases
+that stay unframed — is internal to datafusion-python and documented in
+`crates/core/src/codec.rs` for anyone changing it.
The current FFI logical codec supports providers and UDFs but not arbitrary custom
`LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and
@@ -347,7 +433,7 @@
so it is a property of the session rather than of a handle on it, and installing one is
visible to every context sharing that session — including ones a `with_*` call returned
earlier. Installing a codec on a session that already has a foreign planner rebuilds
-that planner against the new codec for the same reason: there is one planner, and it has
+that planner against the new chain for the same reason: there is one planner, and it has
to carry the codecs currently in force. This happens on the shared session, so it takes
effect even if the returned context is discarded — `ctx.with_python_udf_inlining(...)`
whose result is thrown away still leaves the session's planner carrying the codecs of
@@ -361,15 +447,17 @@
> installed one. Every other path — `Expr.to_bytes(ctx)`, `ExecutionPlan.to_bytes(ctx)`,
> registering a provider — uses the codecs of the handle you call it on.
-Those can be different handles, and then one session has two codecs in effect at once:
+Those can be different handles, and then one session has two codec chains in effect at
+once:
```python
ctx = ctx.with_logical_extension_codec(codec_a)
ctx.set_query_planner(planner)
ctx.with_logical_extension_codec(codec_b) # discarded
-Expr.to_bytes(expr, ctx) # encodes with codec_a -- ctx's own field
-ctx.sql(...).collect() # plans with codec_b -- installed via the discarded handle
+Expr.to_bytes(expr, ctx) # encodes with [codec_a, default] -- ctx's own field
+ctx.sql(...).collect() # plans with [codec_b, codec_a, default] -- the discarded
+ # handle's chain, installed on the shared session
```
Chaining `ctx = ctx.with_...(...)`, as the example below does, keeps the two in step.
diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md
index 29085bc..257749c 100644
--- a/docs/source/user-guide/upgrade-guides.md
+++ b/docs/source/user-guide/upgrade-guides.md
@@ -100,6 +100,50 @@
`FFI_TaskContextProvider`, `FFI_TableProviderFactory`, and `FFI_ExtensionOptions`
carry no version field, so objects of those types cannot be checked.
+### Extension codecs compose instead of replacing
+
+`SessionContext.with_logical_extension_codec` and
+`with_physical_extension_codec` previously replaced whichever codec was already
+installed, so a session could only ever have one. Installing a second codec
+silently discarded the first, and plans failed later with a confusing decode
+error. Both methods now append to a chain, and a session can carry codecs from
+several independent libraries at once.
+
+**No change is required in an extension codec.** Keep implementing
+`LogicalExtensionCodec` or `PhysicalExtensionCodec` exactly as before. Your codec
+is still handed back exactly the bytes it wrote, and is never handed a payload
+another library's codec wrote.
+
+Callers relying on replacement semantics — installing a codec in order to remove
+a previous one — are affected. There is no way to remove an installed codec.
+
+A serialized plan now records which codec wrote each payload, as a short id taken
+from the codec's class. Two behaviours follow from that:
+
+- Installing two instances of one class raises a `ValueError`, because both would
+ claim the same id. Pass `codec_id=` to tell them apart.
+- A codec installed from a bare `PyCapsule` has no class to take an id from, so it
+ gets one private to the session that installed it. It works normally on that
+ session, but a plan it encodes cannot be decoded on an unrelated one. Pass
+ `codec_id=` if those plans have to cross sessions.
+
+```python
+ctx = ctx.with_logical_extension_codec(lib_a.codec())
+ctx = ctx.with_logical_extension_codec(lib_b.codec()) # no longer discards lib_a
+
+# Two instances of one class need distinct ids.
+ctx = ctx.with_logical_extension_codec(lib_a.Codec(), codec_id="lib_a.reader")
+ctx = ctx.with_logical_extension_codec(lib_a.Codec(), codec_id="lib_a.writer")
+
+ctx.logical_extension_codec_ids()
+```
+
+Serialized plans change shape once an extension codec is installed, because each
+payload now records which codec wrote it. A session with no extension codecs
+installed produces the same bytes as before, as do functions encoded by name.
+Regenerate any plan you serialized with an earlier release and stored for later
+use, if it was produced by a session with an extension codec installed.
+
### Changes to the `datafusion-python-util` crate
Extension libraries written in Rust usually depend on the
diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md
index 0fa10d7..401eb61 100644
--- a/examples/datafusion-ffi-example/README.md
+++ b/examples/datafusion-ffi-example/README.md
@@ -35,7 +35,13 @@
Both codec getters take the `SessionContext` they are being installed on and pull the `TaskContextProvider` off it, so decode callbacks resolve session configuration and registered functions against the session that is running the query. Passing `require_udf_on_decode` to either constructor makes every decode call resolve a named scalar function out of that context, which is how the tests check where the registry came from.
-This example makes the provider library the sole external codec owner. Register both provider codecs before installing the planner:
+Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends the codec to the session's codec chain, and DataFusion's default codec handles whatever no installed codec claims. Every payload a codec writes is wrapped in an envelope naming that codec, and decoding consults exactly the codec that encoded it — so several independent plugin libraries can install codecs on the same session without any of them having to recognise or reject the others' payloads. The codecs here are written as ordinary `LogicalExtensionCodec` / `PhysicalExtensionCodec` implementations; the envelope is applied and stripped by `datafusion-python` and never reaches them.
+
+`MyLogicalExtensionCodec` takes an optional `provider_prefix` argument (`MyLogicalExtensionCodec(provider_prefix="TOKENAAA")`) that overrides the byte prefix it stamps on encoded table providers. It exists so the tests can install two instances that own disjoint slices of the wire format, which is what makes install ordering observable from Python. Two instances of one class derive the same id, so those tests also pass `codec_id=` to tell them apart. Real plugin libraries should hard-code a prefix unique to the library rather than accept one from the caller.
+
+`NameOnlyUdfCodec` is the opposite shape: it owns functions that are fully described by their names, so it encodes no bytes at all and rebuilds each function from the name on decode. It exists to pin the by-name path, which is the one place a payload has no id to dispatch on.
+
+Register both provider codecs before installing the planner:
```python
ctx = ctx.with_logical_extension_codec(provider_logical_codec)
@@ -45,4 +51,4 @@
Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. The exception is a planner that wraps a fallback: the rebuild reaches the installed planner only, not the fallback inside it, so codecs-first is a requirement there. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on.
-For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
+For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
diff --git a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py
index cd0c5a6..ac4c69d 100644
--- a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py
+++ b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py
@@ -17,8 +17,46 @@
from __future__ import annotations
-from datafusion import LogicalPlan, SessionContext
-from datafusion_ffi_example import MyLogicalExtensionCodec
+import pyarrow as pa
+import pytest
+from datafusion import Expr, LogicalPlan, SessionContext, col, udf
+from datafusion_ffi_example import (
+ MyLogicalExtensionCodec,
+ MyTableProvider,
+ NameOnlyFunction,
+ NameOnlyUdfCodec,
+)
+
+
+def _double_udf():
+ return udf(
+ lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),
+ [pa.int64()],
+ pa.int64(),
+ volatility="immutable",
+ name="double",
+ )
+
+
+def _encode_provider_plan(token: str) -> tuple[bytes, MyLogicalExtensionCodec]:
+ """Serialize a plan over this library's table provider using a codec
+ that stamps `token` on the encoded provider.
+
+ Returns the blob and the codec, so callers can assert on its call
+ counters. The token is chosen per test so a second codec installed
+ later is provably unable to claim these bytes.
+
+ The codec is installed under ``token`` as its id as well, so a
+ caller can reinstall the same instance elsewhere and have the tag on
+ these bytes resolve. Identity would otherwise be derived from the
+ class, which every instance shares.
+ """
+ codec = MyLogicalExtensionCodec(provider_prefix=token)
+ ctx = SessionContext().with_logical_extension_codec(codec, codec_id=token)
+ ctx.register_table("numbers", MyTableProvider(1, 4, 1))
+ blob = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx)
+ assert token.encode() in blob
+ return blob, codec
def _setup_session_with_codec() -> tuple[SessionContext, MyLogicalExtensionCodec]:
@@ -33,8 +71,8 @@
def test_ffi_logical_codec_install_and_export():
- """Installing a user FFI codec replaces the session's logical
- codec; the capsule getter on the session re-exports it."""
+ """Installing a user FFI codec adds it to the session's logical
+ codec chain; the capsule getter on the session re-exports it."""
ctx, _codec = _setup_session_with_codec()
capsule = ctx.__datafusion_logical_extension_codec__()
assert capsule is not None
@@ -80,3 +118,438 @@
restored = LogicalPlan.from_bytes(ctx, blob)
df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
assert df.collect() == df_round_trip.collect()
+
+
+def test_ffi_logical_codec_composes_with_later_install():
+ """Codecs compose: installing a second codec appends it to the
+ session's codec chain instead of replacing the first. The second
+ codec here (a default-backed codec exported from a fresh session)
+ cannot encode this library's table provider, so the first codec
+ still claims it. Under replace semantics this test fails with
+ `LogicalExtensionCodec is not provided`.
+
+ Both directions are counted. Asserting only the round-trip result
+ would pass if decoding resolved the table some other way; the
+ decode count pins that the bytes went back to the codec that wrote
+ them."""
+ ctx, codec = _setup_session_with_codec()
+ ctx = ctx.with_logical_extension_codec(
+ SessionContext().__datafusion_logical_extension_codec__()
+ )
+
+ ctx.register_table("numbers", MyTableProvider(1, 4, 1))
+ df = ctx.sql('SELECT "A" FROM numbers')
+ plan = df.logical_plan()
+
+ encode_before = codec.table_provider_encode_calls()
+ decode_before = codec.table_provider_decode_calls()
+ blob = plan.to_bytes(ctx)
+ assert codec.table_provider_encode_calls() > encode_before
+
+ restored = LogicalPlan.from_bytes(ctx, blob)
+ assert codec.table_provider_decode_calls() > decode_before
+
+ df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
+ assert df.collect() == df_round_trip.collect()
+
+
+def test_first_installed_codec_encodes():
+ """Encoding walks the chain in install order, so the earliest
+ installed codec that can claim an object gets it.
+
+ Both orders run in one test on purpose. Asserting a single order
+ would also pass under replace semantics, where the second install
+ simply discards the first codec; swapping the order and getting the
+ other token proves the losing codec was still installed and merely
+ lost the race.
+
+ The two instances need explicit ids: identity is otherwise derived
+ from the class, and these are two instances of one class owning
+ disjoint slices of the wire format.
+ """
+ for winner, loser in (("TOKENAAA", "TOKENBBB"), ("TOKENBBB", "TOKENAAA")):
+ winner_codec = MyLogicalExtensionCodec(provider_prefix=winner)
+ loser_codec = MyLogicalExtensionCodec(provider_prefix=loser)
+ ctx = SessionContext().with_logical_extension_codec(
+ winner_codec, codec_id=winner
+ )
+ ctx = ctx.with_logical_extension_codec(loser_codec, codec_id=loser)
+
+ ctx.register_table("numbers", MyTableProvider(1, 4, 1))
+ blob = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx)
+
+ assert winner.encode() in blob
+ assert loser.encode() not in blob
+ assert winner_codec.table_provider_encode_calls() == 1
+ assert loser_codec.table_provider_encode_calls() == 0
+
+
+def test_installing_a_codec_cannot_hijack_an_earlier_codecs_objects():
+ """Appending is additive: a later install can claim objects nothing
+ else claimed, but never takes over an object an earlier codec was
+ already encoding.
+
+ This is why install order is append rather than prepend. Under
+ prepend, adding an unrelated library would silently change how an
+ existing library's objects encode -- and, once payloads are tagged,
+ would renumber ids that older payloads already reference.
+ """
+ first = MyLogicalExtensionCodec(provider_prefix="TOKENAAA")
+ ctx = SessionContext().with_logical_extension_codec(first, codec_id="TOKENAAA")
+ ctx.register_table("numbers", MyTableProvider(1, 4, 1))
+ before = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx)
+
+ later = MyLogicalExtensionCodec(provider_prefix="TOKENBBB")
+ ctx = ctx.with_logical_extension_codec(later, codec_id="TOKENBBB")
+ after = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx)
+
+ # Provider tokens are minted per encode, so the payloads differ in the
+ # token id. What must not change is which codec claimed the provider.
+ assert b"TOKENAAA" in after
+ assert b"TOKENBBB" not in after
+ assert later.table_provider_encode_calls() == 0
+ assert len(before) == len(after)
+
+
+def test_decode_dispatches_to_the_codec_that_encoded():
+ """A payload names the codec that wrote it, so decoding consults
+ exactly that codec and never offers the bytes to any other.
+
+ The blob here is written by ``first`` in its own session. Installing
+ ``second`` alongside it must not put ``second`` anywhere near those
+ bytes -- under trial-and-error dispatch it would be asked first, and
+ a codec that decodes structurally similar protobuf would answer.
+ """
+ blob, first = _encode_provider_plan("TOKENAAA")
+
+ second = MyLogicalExtensionCodec(provider_prefix="TOKENBBB")
+ ctx = SessionContext().with_logical_extension_codec(first, codec_id="TOKENAAA")
+ ctx = ctx.with_logical_extension_codec(second, codec_id="TOKENBBB")
+
+ restored = LogicalPlan.from_bytes(ctx, blob)
+ assert ctx.create_dataframe_from_logical_plan(restored).collect()
+
+ assert first.table_provider_decode_calls() == 1
+ assert second.table_provider_decode_calls() == 0
+
+
+def test_decode_survives_a_different_install_order():
+ """Dispatch keys off codec identity, not chain position, so the
+ decoding session may install the same codecs in any order.
+
+ This is the case positional dispatch cannot handle: the encoding
+ session has the owning codec at index 0 and the decoding session has
+ it at index 1. Keying on position would hand the payload to whatever
+ sits at index 0 in the decoder -- silently, and with a plausible
+ result.
+ """
+ owner = MyLogicalExtensionCodec(provider_prefix="TOKENAAA")
+ other = MyLogicalExtensionCodec(provider_prefix="TOKENBBB")
+
+ encoder = SessionContext().with_logical_extension_codec(owner, codec_id="TOKENAAA")
+ encoder = encoder.with_logical_extension_codec(other, codec_id="TOKENBBB")
+ encoder.register_table("numbers", MyTableProvider(1, 4, 1))
+ blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder)
+
+ # Same codecs, opposite order.
+ decoder = SessionContext().with_logical_extension_codec(other, codec_id="TOKENBBB")
+ decoder = decoder.with_logical_extension_codec(owner, codec_id="TOKENAAA")
+
+ restored = LogicalPlan.from_bytes(decoder, blob)
+ assert decoder.create_dataframe_from_logical_plan(restored).collect()
+ assert other.table_provider_decode_calls() == 0
+
+
+def test_decode_names_the_codec_that_is_not_installed():
+ """When the owning codec is absent the error names it and lists what
+ is installed, instead of reporting DataFusion's generic "not
+ provided" from whichever codec was tried last."""
+ blob, _owner = _encode_provider_plan("TOKENBBB")
+
+ ctx = SessionContext().with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix="TOKENCCC"), codec_id="lib_c.Codec"
+ )
+
+ with pytest.raises(Exception, match="TOKENBBB") as excinfo:
+ LogicalPlan.from_bytes(ctx, blob)
+
+ message = str(excinfo.value)
+ # Names the codec the payload belongs to, and what is actually here.
+ assert "not installed on this session" in message
+ assert "lib_c.Codec" in message
+
+
+def test_a_codec_id_defaults_to_its_class_module_and_qualname():
+ """A codec that declares no identity is named by its class.
+
+ That string goes on the wire in front of every payload the codec
+ writes, so it is a compatibility surface: renaming the class or moving
+ the module changes it, and plans stored by an earlier release stop
+ decoding. A library that needs to rename declares
+ ``__datafusion_codec_id__`` instead, which the next test covers.
+
+ Pinned twice on purpose -- once against the literal, so a rename has to
+ come here and be acknowledged, and once against the rule, so the
+ literal cannot drift into something the code no longer derives.
+ """
+ ctx = SessionContext().with_logical_extension_codec(MyLogicalExtensionCodec())
+
+ assert ctx.logical_extension_codec_ids() == [
+ "datafusion_ffi_example.MyLogicalExtensionCodec"
+ ]
+ assert ctx.logical_extension_codec_ids() == [
+ f"{MyLogicalExtensionCodec.__module__}.{MyLogicalExtensionCodec.__qualname__}"
+ ]
+
+
+class _CodecUnderItsOldName:
+ """A library codec that pins its identity, so the class can be renamed.
+
+ Delegates the capsule getter to a real FFI codec. ``session`` is
+ forwarded, because that argument is how the underlying library reaches
+ the session the codec is being installed on.
+ """
+
+ __datafusion_codec_id__ = "pinned.example.Codec"
+
+ def __init__(self, inner: MyLogicalExtensionCodec) -> None:
+ self._inner = inner
+
+ def __datafusion_logical_extension_codec__(self, session: object = None) -> object:
+ return self._inner.__datafusion_logical_extension_codec__(session)
+
+
+class _CodecUnderItsNewName(_CodecUnderItsOldName):
+ """The same codec after a rename. Same pinned id, different class."""
+
+
+def test_a_pinned_codec_id_survives_a_class_rename():
+ """``__datafusion_codec_id__`` decouples identity from the class name,
+ which is the reason to declare one.
+
+ A plan encoded by the codec under its old name decodes on a session
+ that only knows the new name. Under the class-derived default the two
+ would be different ids and the payload would be undecodable.
+ """
+ assert _CodecUnderItsOldName.__qualname__ != _CodecUnderItsNewName.__qualname__
+
+ old = MyLogicalExtensionCodec(provider_prefix="TOKENAAA")
+ encoder = SessionContext().with_logical_extension_codec(_CodecUnderItsOldName(old))
+ assert encoder.logical_extension_codec_ids() == ["pinned.example.Codec"]
+
+ encoder.register_table("numbers", MyTableProvider(1, 4, 1))
+ blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder)
+
+ # The pinned id, not the class name, is what the payload carries.
+ assert b"pinned.example.Codec" in blob
+ assert b"_CodecUnderItsOldName" not in blob
+
+ new = MyLogicalExtensionCodec(provider_prefix="TOKENAAA")
+ decoder = SessionContext().with_logical_extension_codec(_CodecUnderItsNewName(new))
+ assert decoder.logical_extension_codec_ids() == ["pinned.example.Codec"]
+
+ restored = LogicalPlan.from_bytes(decoder, blob)
+ assert decoder.create_dataframe_from_logical_plan(restored).collect()
+ assert new.table_provider_decode_calls() == 1
+
+
+def test_installing_two_codecs_under_one_id_is_rejected():
+ """Identity is derived from the class, so installing two instances of
+ one class collides. Rejecting at install time is the point: two
+ codecs sharing an id are indistinguishable when a payload is decoded,
+ and only the caller knows whether they write the same wire format."""
+ ctx = SessionContext().with_logical_extension_codec(MyLogicalExtensionCodec())
+
+ with pytest.raises(ValueError, match="already installed"):
+ ctx.with_logical_extension_codec(MyLogicalExtensionCodec())
+
+
+def _encode_through_a_bare_capsule(token: str) -> tuple[bytes, list[str]]:
+ """Serialize a provider plan through a codec installed as a bare
+ capsule, which is the case with no derivable identity.
+
+ Returns the blob and the encoding session's codec ids, so a caller
+ can compare them against another session's.
+ """
+ exporter = SessionContext().with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix=token)
+ )
+ encoder = SessionContext().with_logical_extension_codec(
+ exporter.__datafusion_logical_extension_codec__()
+ )
+ encoder.register_table("numbers", MyTableProvider(1, 4, 1))
+ blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder)
+ return blob, encoder.logical_extension_codec_ids()
+
+
+def test_bare_capsule_codec_is_session_local():
+ """A bare PyCapsule exposes nothing stable to derive an identity
+ from -- every capsule reports the same type -- so it is tagged with a
+ session-local id. A plan it encodes fails on an unrelated session
+ with an error naming the fix, rather than being decoded by whichever
+ codec happens to sit at the same position."""
+ blob, _ids = _encode_through_a_bare_capsule("TOKENAAA")
+
+ with pytest.raises(Exception, match="bare PyCapsule") as excinfo:
+ LogicalPlan.from_bytes(SessionContext(), blob)
+ assert "codec_id" in str(excinfo.value)
+
+
+def test_a_bare_capsule_codec_id_is_not_re_mintable_by_another_session():
+ """The id given to a capsule-installed codec must be one no other
+ session can arrive at.
+
+ Both sessions here install exactly one bare capsule, so any identity
+ drawn from a namespace both sessions number the same way -- a
+ counter, a position in the chain -- collides, and the payload is
+ handed to the other library's codec. That failure is quiet: a codec
+ offered bytes it does not recognise falls through to its own inner
+ default codec, so the error names neither codec and no counter moves.
+ Asserting on the message is what separates the two schemes.
+
+ The empty-chain case in the test above passes under either scheme,
+ because a lookup in an empty chain misses whatever the id is.
+ """
+ blob, encoder_ids = _encode_through_a_bare_capsule("TOKENAAA")
+
+ # An unrelated session, also holding exactly one bare capsule.
+ other = SessionContext().with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix="TOKENBBB")
+ )
+ decoder = SessionContext().with_logical_extension_codec(
+ other.__datafusion_logical_extension_codec__()
+ )
+ decoder_ids = decoder.logical_extension_codec_ids()
+ assert len(encoder_ids) == len(decoder_ids) == 1
+ assert encoder_ids != decoder_ids
+
+ with pytest.raises(Exception, match="bare PyCapsule") as excinfo:
+ LogicalPlan.from_bytes(decoder, blob)
+ assert "codec_id" in str(excinfo.value)
+
+
+def test_installing_a_session_as_a_codec_uses_a_per_session_id():
+ """A context installed as a codec is identified by its session, not by
+ its class.
+
+ Every ``SessionContext`` shares one class, so a class-derived id would
+ name all of them: two contexts could not coexist on one target, and a
+ payload written through one would resolve to the other on decode. The
+ sources are held in locals because the imported codecs resolve their
+ task context against them.
+ """
+ src_a = SessionContext().with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix="TOKENAAA")
+ )
+ src_b = SessionContext().with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix="TOKENBBB")
+ )
+ assert src_a.__datafusion_codec_id__ != src_b.__datafusion_codec_id__
+
+ # Both sources compose onto one session, which a shared id would refuse.
+ both = SessionContext().with_logical_extension_codec(src_a)
+ both = both.with_logical_extension_codec(src_b)
+ assert both.logical_extension_codec_ids() == [
+ src_a.__datafusion_codec_id__,
+ src_b.__datafusion_codec_id__,
+ ]
+
+ # A payload written through one source does not resolve to the other.
+ encoder = SessionContext().with_logical_extension_codec(src_a)
+ encoder.register_table("numbers", MyTableProvider(1, 4, 1))
+ blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder)
+
+ decoder = SessionContext().with_logical_extension_codec(src_b)
+ with pytest.raises(Exception, match="not installed on this session") as excinfo:
+ LogicalPlan.from_bytes(decoder, blob)
+ assert src_a.__datafusion_codec_id__ in str(excinfo.value)
+
+
+def test_a_derived_handle_reports_its_session_id():
+ """Handles derived from one session share its id, so only one of them
+ can be installed on a given target. Their payloads would be
+ indistinguishable on decode, so refusing is the right answer."""
+ base = SessionContext()
+ derived = base.with_python_udf_inlining(enabled=False)
+ assert derived.__datafusion_codec_id__ == base.__datafusion_codec_id__
+
+ target = SessionContext().with_logical_extension_codec(base)
+ with pytest.raises(ValueError, match="already installed"):
+ target.with_logical_extension_codec(derived)
+
+
+def test_name_only_codec_round_trips_without_a_payload():
+ """A codec may own functions that need no payload: the name is the
+ whole encoding. ``try_encode_udf`` writes nothing, and the decoder
+ rebuilds the function from the name with no registry entry.
+
+ DataFusion supports this directly -- an empty ``fun_definition``
+ sends the decoder to the registry first and the codec second. This
+ test pins that arm from the Python side, because it is the one path
+ where a payload is still offered to every installed codec: there are
+ no bytes, so there is no identity to dispatch on.
+
+ It is also the guard against a plausible "improvement". Wrapping
+ every chained encode in the identity envelope would make this
+ payload non-empty, which sets ``fun_definition`` and permanently
+ skips the registry lookup -- breaking both this codec and ordinary
+ by-name round trips, with nothing else in the suite noticing.
+ """
+ codec = NameOnlyUdfCodec()
+ name = NameOnlyUdfCodec.function_name()
+
+ # FROM-less, so serialization never reaches try_encode_table_provider --
+ # this codec owns functions, not providers.
+ encoder = SessionContext().with_logical_extension_codec(codec)
+ encoder.register_udf(udf(NameOnlyFunction()))
+ blob = encoder.sql(f"SELECT {name}(1) AS x").logical_plan().to_bytes(encoder)
+
+ # The name is the entire encoding, so the codec contributed no bytes
+ # and the payload carries no identity envelope for it.
+ assert codec.encode_udf_calls() > 0
+ assert b"DFPYCHN" not in blob
+
+ # A fresh session that never registered the function: only the codec
+ # can supply it, and only from the name.
+ decoder = SessionContext().with_logical_extension_codec(codec)
+ restored = LogicalPlan.from_bytes(decoder, blob)
+
+ assert codec.decode_udf_calls() > 0
+ assert decoder.create_dataframe_from_logical_plan(restored).collect()
+
+
+def test_default_only_session_writes_no_envelope():
+ """A session with no extension codecs installed produces the same
+ bytes as a build without codec chaining: the terminal codec writes
+ unframed, so the envelope only appears once a codec is installed.
+
+ Keeps the wire break scoped to sessions that actually compose."""
+ ctx = SessionContext()
+ blob = ctx.sql("SELECT abs(-1) AS x").logical_plan().to_bytes(ctx)
+ assert b"DFPYCHN" not in blob
+
+
+def test_udf_inlining_setting_survives_codec_install():
+ """Installing an extension codec must not silently re-enable inline
+ Python UDF encoding on a session that opted out. Regression guard in
+ both directions: the encoder still emits the by-name form, and the
+ decoder still refuses an inline payload.
+
+ The codec installed here delegates UDF encoding to DataFusion's
+ default codec. A codec exported from another `SessionContext` would
+ not work as a probe: that export is itself a Python-aware codec with
+ inlining enabled, so the strict outer codec would delegate to it and
+ the inline payload would reappear.
+ """
+ strict = SessionContext().with_python_udf_inlining(enabled=False)
+ extended = strict.with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix="TOKENFFF")
+ )
+
+ e = _double_udf()(col("a"))
+ assert b"DFPYUDF" not in e.to_bytes(extended)
+
+ inline_blob = e.to_bytes(SessionContext())
+ assert b"DFPYUDF" in inline_blob
+ with pytest.raises(Exception, match="inlining is disabled"):
+ Expr.from_bytes(inline_blob, ctx=extended)
diff --git a/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py
index 28eaaf2..c7a6ede 100644
--- a/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py
+++ b/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py
@@ -76,3 +76,26 @@
restored = ExecutionPlan.from_bytes(ctx, blob)
assert str(original) == str(restored)
+
+
+def test_ffi_physical_codec_composes_with_later_install():
+ """Codecs compose: a second install appends to the chain instead
+ of replacing the first codec. The second codec here (default-backed
+ export from a fresh session) encodes UDFs by name without writing
+ bytes, which the chain treats as "no opinion" — so the user codec
+ installed first is still consulted. Under replace semantics its
+ counter stays at zero."""
+ ctx, codec = _setup_session_with_codec()
+ ctx = ctx.with_physical_extension_codec(
+ SessionContext().__datafusion_physical_extension_codec__()
+ )
+
+ df = ctx.sql("SELECT abs(a) AS x FROM t")
+ original = df.execution_plan()
+
+ before = codec.encode_udf_calls()
+ blob = original.to_bytes(ctx)
+ assert codec.encode_udf_calls() > before
+
+ restored = ExecutionPlan.from_bytes(ctx, blob)
+ assert str(original) == str(restored)
diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs
index 3d00fdb..92fccb1 100644
--- a/examples/datafusion-ffi-example/src/lib.rs
+++ b/examples/datafusion-ffi-example/src/lib.rs
@@ -21,6 +21,7 @@
use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList};
use crate::config::MyConfig;
use crate::logical_extension_codec::MyLogicalExtensionCodec;
+use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec};
use crate::physical_extension_codec::MyPhysicalExtensionCodec;
use crate::physical_optimizer::MyPhysicalOptimizerRule;
use crate::scalar_udf::IsNullUDF;
@@ -33,6 +34,7 @@
pub(crate) mod catalog_provider;
pub(crate) mod config;
pub(crate) mod logical_extension_codec;
+pub(crate) mod name_only_codec;
pub(crate) mod physical_extension_codec;
pub(crate) mod physical_optimizer;
pub(crate) mod required_udf;
@@ -57,6 +59,8 @@
m.add_class::<MyRankUDF>()?;
m.add_class::<MyConfig>()?;
m.add_class::<MyLogicalExtensionCodec>()?;
+ m.add_class::<NameOnlyUdfCodec>()?;
+ m.add_class::<NameOnlyFunction>()?;
m.add_class::<MyPhysicalExtensionCodec>()?;
m.add_class::<MyPhysicalOptimizerRule>()?;
Ok(())
diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs
index 1fcaaef..5660489 100644
--- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs
+++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs
@@ -90,6 +90,9 @@
/// Scalar function every table-provider decode must resolve from the
/// `TaskContext` it is handed. See [`crate::required_udf`].
required_udf: Option<String>,
+ /// Byte prefix identifying providers this codec owns. Distinct tokens let a
+ /// test install several instances and observe which one the chain picks.
+ token: Arc<[u8]>,
}
impl fmt::Debug for CountingLogicalExtensionCodec {
@@ -124,7 +127,7 @@
ctx: &TaskContext,
) -> Result<Arc<dyn TableProvider>> {
resolve_required_udf(self.required_udf.as_deref(), ctx, &self.counters.task_ctx)?;
- if let Some(id) = token_id(buf, TABLE_PROVIDER_TOKEN) {
+ if let Some(id) = token_id(buf, &self.token) {
self.counters
.decode_table_provider
.fetch_add(1, Ordering::SeqCst);
@@ -157,7 +160,7 @@
.lock()
.map_err(|err| DataFusionError::Internal(err.to_string()))?
.insert(id, node);
- buf.extend_from_slice(TABLE_PROVIDER_TOKEN);
+ buf.extend_from_slice(&self.token);
buf.extend_from_slice(&id.to_le_bytes());
return Ok(());
}
@@ -185,6 +188,7 @@
pub(crate) struct MyLogicalExtensionCodec {
counters: Arc<CallCounters>,
required_udf: Option<String>,
+ token: Arc<[u8]>,
}
#[pymethods]
@@ -195,12 +199,22 @@
/// provider decode must find in the `TaskContext` it is handed. Leave it
/// unset for the ordinary behaviour; set it to observe *which* session's
/// registry the FFI decode callback actually receives.
+ ///
+ /// `provider_prefix` overrides [`TABLE_PROVIDER_TOKEN`], the byte prefix
+ /// stamped on encoded table providers. Two instances built with different
+ /// prefixes each own a disjoint slice of the wire format, which is what
+ /// lets a test install both and tell from the decoded bytes which one the
+ /// session's codec chain consulted.
#[new]
- #[pyo3(signature = (require_udf_on_decode=None))]
- fn new(require_udf_on_decode: Option<String>) -> Self {
+ #[pyo3(signature = (require_udf_on_decode=None, provider_prefix=None))]
+ fn new(require_udf_on_decode: Option<String>, provider_prefix: Option<&str>) -> Self {
Self {
counters: Arc::new(CallCounters::default()),
required_udf: require_udf_on_decode,
+ token: provider_prefix.map_or_else(
+ || Arc::from(TABLE_PROVIDER_TOKEN),
+ |prefix| Arc::from(prefix.as_bytes()),
+ ),
}
}
@@ -245,6 +259,7 @@
inner: DefaultLogicalExtensionCodec {},
counters: Arc::clone(&self.counters),
required_udf: self.required_udf.clone(),
+ token: Arc::clone(&self.token),
});
let runtime = get_tokio_runtime().handle().clone();
diff --git a/examples/datafusion-ffi-example/src/name_only_codec.rs b/examples/datafusion-ffi-example/src/name_only_codec.rs
new file mode 100644
index 0000000..9b82c4b
--- /dev/null
+++ b/examples/datafusion-ffi-example/src/name_only_codec.rs
@@ -0,0 +1,268 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A codec whose functions need no payload at all.
+//!
+//! Most extension codecs answer with bytes. This one owns a fixed catalog of
+//! functions that are fully described by their names, so `try_encode_udf`
+//! writes nothing and `try_decode_udf` rebuilds the function from `name`
+//! alone. DataFusion supports that shape directly: an encoder that writes no
+//! bytes leaves `fun_definition` unset, and the decoder then tries the
+//! `FunctionRegistry` first and the codec second — see the
+//! `None => ctx.udf(..).or_else(|_| codec.try_decode_udf(name, &[]))` arm in
+//! `datafusion-proto`'s `from_proto.rs`.
+//!
+//! It exists here to pin that arm. Because there are no bytes, there is
+//! nothing to tag with the codec's identity, so this is the one path where
+//! `PythonLogicalCodec` still offers a payload to every installed codec in
+//! turn. A change that wrapped empty encodings in an envelope would set
+//! `fun_definition`, skip the registry lookup permanently, and break both this
+//! codec and plain by-name round trips — with no other test noticing.
+
+use std::fmt;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use arrow_schema::DataType;
+use datafusion::common::error::Result;
+use datafusion::common::not_impl_err;
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
+ Volatility,
+};
+use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec};
+use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio_runtime};
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+/// Prefix marking the functions this library owns. A name is the entire
+/// encoding, so the prefix is the whole ownership test.
+const NAME_PREFIX: &str = "name_only_";
+
+/// Scalar function reconstructed purely from its name.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct NameOnlyUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl NameOnlyUdf {
+ fn new(name: impl Into<String>) -> Self {
+ Self {
+ name: name.into(),
+ signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for NameOnlyUdf {
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ Ok(args.args[0].clone())
+ }
+}
+
+#[derive(Default)]
+struct Counters {
+ encode_udf: AtomicUsize,
+ decode_udf: AtomicUsize,
+}
+
+impl fmt::Debug for Counters {
+ fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+ formatter.debug_struct("Counters").finish_non_exhaustive()
+ }
+}
+
+struct NameOnlyLogicalExtensionCodec {
+ inner: DefaultLogicalExtensionCodec,
+ counters: Arc<Counters>,
+}
+
+impl fmt::Debug for NameOnlyLogicalExtensionCodec {
+ fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+ formatter
+ .debug_struct("NameOnlyLogicalExtensionCodec")
+ .finish_non_exhaustive()
+ }
+}
+
+impl LogicalExtensionCodec for NameOnlyLogicalExtensionCodec {
+ fn try_decode(
+ &self,
+ buf: &[u8],
+ inputs: &[datafusion::logical_expr::LogicalPlan],
+ ctx: &datafusion::execution::TaskContext,
+ ) -> Result<datafusion::logical_expr::Extension> {
+ self.inner.try_decode(buf, inputs, ctx)
+ }
+
+ fn try_encode(
+ &self,
+ node: &datafusion::logical_expr::Extension,
+ buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ self.inner.try_encode(node, buf)
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ buf: &[u8],
+ table_ref: &datafusion::common::TableReference,
+ schema: arrow_schema::SchemaRef,
+ ctx: &datafusion::execution::TaskContext,
+ ) -> Result<Arc<dyn datafusion::datasource::TableProvider>> {
+ self.inner
+ .try_decode_table_provider(buf, table_ref, schema, ctx)
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ table_ref: &datafusion::common::TableReference,
+ node: Arc<dyn datafusion::datasource::TableProvider>,
+ buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ self.inner.try_encode_table_provider(table_ref, node, buf)
+ }
+
+ /// Writes nothing on purpose. The name is the whole encoding, so there is
+ /// no payload to emit, and returning `Ok` with an empty buffer is how a
+ /// codec says "encoded by name" to DataFusion.
+ fn try_encode_udf(&self, node: &ScalarUDF, _buf: &mut Vec<u8>) -> Result<()> {
+ if node.name().starts_with(NAME_PREFIX) {
+ self.counters.encode_udf.fetch_add(1, Ordering::SeqCst);
+ }
+ Ok(())
+ }
+
+ /// Rebuilds the function from `name`, with no registry entry and no bytes.
+ fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+ if !name.starts_with(NAME_PREFIX) {
+ return not_impl_err!("Not a name-only function: {name}");
+ }
+ if !buf.is_empty() {
+ return not_impl_err!(
+ "name-only functions carry no payload, but {} bytes were supplied for {name}",
+ buf.len()
+ );
+ }
+ self.counters.decode_udf.fetch_add(1, Ordering::SeqCst);
+ Ok(Arc::new(ScalarUDF::from(NameOnlyUdf::new(name))))
+ }
+}
+
+/// The function [`NameOnlyUdfCodec`] owns, exported so a session can register
+/// it and build a plan that references it.
+///
+/// Only the *encoding* session needs it registered. The decoding session
+/// deliberately does not, which is what forces the codec's name-only decode
+/// path to run.
+#[pyclass(
+ from_py_object,
+ name = "NameOnlyFunction",
+ module = "datafusion_ffi_example",
+ subclass
+)]
+#[derive(Debug, Clone)]
+pub(crate) struct NameOnlyFunction;
+
+#[pymethods]
+impl NameOnlyFunction {
+ #[new]
+ fn new() -> Self {
+ Self
+ }
+
+ fn __datafusion_scalar_udf__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyCapsule>> {
+ let func = Arc::new(ScalarUDF::from(NameOnlyUdf::new(format!(
+ "{NAME_PREFIX}identity"
+ ))));
+ PyCapsule::new_with_value(
+ py,
+ datafusion_ffi::udf::FFI_ScalarUDF::from(func),
+ cr"datafusion_scalar_udf",
+ )
+ }
+}
+
+/// Codec owning functions that are reconstructible from their names alone.
+///
+/// A real library shaped like this would be one shipping a fixed catalog of
+/// built-ins: nothing about a call site varies, so there is nothing to encode.
+#[pyclass(
+ from_py_object,
+ name = "NameOnlyUdfCodec",
+ module = "datafusion_ffi_example",
+ subclass
+)]
+#[derive(Clone)]
+pub(crate) struct NameOnlyUdfCodec {
+ counters: Arc<Counters>,
+}
+
+#[pymethods]
+impl NameOnlyUdfCodec {
+ #[new]
+ fn new() -> Self {
+ Self {
+ counters: Arc::new(Counters::default()),
+ }
+ }
+
+ /// Name of the function this codec can rebuild, for use in a query.
+ #[staticmethod]
+ fn function_name() -> String {
+ format!("{NAME_PREFIX}identity")
+ }
+
+ fn encode_udf_calls(&self) -> usize {
+ self.counters.encode_udf.load(Ordering::SeqCst)
+ }
+
+ fn decode_udf_calls(&self) -> usize {
+ self.counters.decode_udf.load(Ordering::SeqCst)
+ }
+
+ fn __datafusion_logical_extension_codec__<'py>(
+ &self,
+ py: Python<'py>,
+ session: Bound<'py, PyAny>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ let inner: Arc<dyn LogicalExtensionCodec> = Arc::new(NameOnlyLogicalExtensionCodec {
+ inner: DefaultLogicalExtensionCodec {},
+ counters: Arc::clone(&self.counters),
+ });
+
+ let runtime = get_tokio_runtime().handle().clone();
+ let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?;
+ let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), ctx_provider);
+
+ PyCapsule::new_with_value(py, ffi, cr"datafusion_logical_extension_codec")
+ }
+}
diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md
index 72f96bb..66bc451 100644
--- a/examples/datafusion-ffi-query-planner-example/README.md
+++ b/examples/datafusion-ffi-query-planner-example/README.md
@@ -55,6 +55,6 @@
`MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit.
-The provider's codec pair is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against it, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep).
+The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep).
-For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
+For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py
index d046f67..c6ef207 100644
--- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py
+++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py
@@ -222,6 +222,23 @@
return ctx, logical_codec, physical_codec
+def physical_only_context(max_rows: int = 3):
+ """Context with the physical codec installed but no logical codec.
+
+ Encoding consults chained codecs in install order and the first to claim an
+ object wins, so a codec installed later cannot be observed while an earlier
+ one is already claiming table providers. Leaving the logical slot empty lets
+ a test install exactly one logical codec -- through a handle it then throws
+ away -- and watch its counters.
+ """
+ config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows))
+ physical_codec = MyPhysicalExtensionCodec()
+ ctx = SessionContext(config)
+ ctx = ctx.with_physical_extension_codec(physical_codec)
+ ctx.register_table("numbers", MyTableProvider(1, 6, 1))
+ return ctx, physical_codec
+
+
def test_installing_a_planner_keeps_the_session_id():
"""A session and its decode callbacks must agree on the session id.
@@ -521,9 +538,11 @@
A fresh codec instance is what makes it observable -- it carries its own
counters, and the planner encodes the outbound logical plan with whichever
- codec it is holding.
+ codec it is holding. The base context deliberately installs no logical
+ codec, so this one is the only candidate; an already-installed codec would
+ claim the provider first and hide the rebind.
"""
- ctx, _logical_codec, _physical_codec = configured_context(max_rows=3)
+ ctx, _physical_codec = physical_only_context(max_rows=3)
ctx.set_query_planner(MyQueryPlanner())
later = MyLogicalExtensionCodec()
@@ -548,14 +567,17 @@
Chaining ``ctx = ctx.with_...(...)`` keeps the two in step; this pins what
happens when they are allowed to diverge.
+ ``ctx`` installs no logical codec of its own, so its chain is empty and the
+ planner's holds exactly one entry. That asymmetry is what makes the split
+ visible: an entry on both chains would be claimed by the same codec either
+ way, since encoding stops at the first codec to claim an object.
+
Inlining has to be off for the assertion to say anything: with it on, a
Python UDF is encoded inline by ``PythonLogicalCodec`` and never reaches
the installed codec's ``try_encode_udf``.
"""
config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
- handle_codec = MyLogicalExtensionCodec()
ctx = SessionContext(config).with_python_udf_inlining(enabled=False)
- ctx = ctx.with_logical_extension_codec(handle_codec)
ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
ctx.register_table("numbers", MyTableProvider(1, 6, 1))
ctx.set_query_planner(MyQueryPlanner())
@@ -575,15 +597,14 @@
ctx.register_udf(identity)
Expr.to_bytes(identity(col("A")), ctx)
- # Serializing through `ctx` uses `ctx`'s own codec field.
- assert handle_codec.encode_udf_calls() > 0
+ # Serializing through `ctx` uses `ctx`'s own codec field, which is empty --
+ # the UDF goes out by name through the terminal codec.
assert planner_codec.encode_udf_calls() == 0
ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
# Planning through the same `ctx` uses the codec the planner was rebound to.
assert planner_codec.table_provider_encode_calls() > 0
- assert handle_codec.table_provider_encode_calls() == 0
def test_an_unchanged_inlining_setting_leaves_the_planner_alone():
@@ -595,11 +616,11 @@
Observable only once the planner is holding some *other* handle's codec:
without the guard, a defensive no-op toggle on `ctx` drags the planner back
- onto `ctx`'s codec and silently undoes the install below. The rebuilt
- codecs otherwise wrap the same inner codec, so nothing else distinguishes
- the two paths.
+ onto `ctx`'s codecs and silently undoes the install below. `ctx` installs no
+ logical codec, so being dragged back leaves the planner with an empty chain
+ and the query fails outright rather than quietly using the wrong codec.
"""
- ctx, handle_codec, _physical_codec = codec_context()
+ ctx, _physical_codec = physical_only_context()
ctx.set_query_planner(MyQueryPlanner())
planner_codec = MyLogicalExtensionCodec()
@@ -612,7 +633,6 @@
ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert planner_codec.table_provider_encode_calls() > 0
- assert handle_codec.table_provider_encode_calls() == 0
def test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs():
@@ -627,7 +647,7 @@
So "re-install the planner after installing a codec" only repairs anything
when it is done from the handle holding the new codec.
"""
- ctx, original_logical, _physical_codec = codec_context()
+ ctx, _physical_codec = physical_only_context()
planner = MyQueryPlanner()
ctx.set_query_planner(planner)
@@ -638,15 +658,14 @@
ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert later.table_provider_encode_calls() > 0
- assert original_logical.table_provider_encode_calls() == 0
- # `ctx`'s own codec field never changed, so this rebuilds the planner
- # against `original_logical` and drops `later` from the session's planner.
+ # `ctx`'s own codec field never changed -- it never had a logical codec --
+ # so this rebuilds the planner against an empty chain and drops `later`.
ctx.set_query_planner(planner)
encodes_by_later = later.table_provider_encode_calls()
- ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
- assert original_logical.table_provider_encode_calls() > 0
+ with pytest.raises(Exception, match=r"LogicalExtensionCodec|TableProvider"):
+ ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert later.table_provider_encode_calls() == encodes_by_later
@@ -667,3 +686,31 @@
with pytest.raises(Exception, match=r"max_rows|Invalid value"):
ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect()
+
+
+def test_composed_codecs_with_query_planner():
+ """A second pair of codecs installed on top of the provider codecs
+ composes with them instead of replacing them.
+
+ The provider codecs are installed first, so encoding consults them
+ first and they claim this library's tables and plans before the
+ extra codecs (default-backed exports from a fresh session) get a
+ turn; decoding goes straight to whichever codec wrote the bytes.
+ The extra pair therefore changes nothing observable here, which is
+ the assertion: under replace semantics it would have discarded the
+ provider codecs and the planner-driven round trip would fail."""
+ ctx, logical_codec, physical_codec = configured_context(max_rows=2)
+ other = SessionContext()
+ ctx = ctx.with_logical_extension_codec(
+ other.__datafusion_logical_extension_codec__()
+ )
+ ctx = ctx.with_physical_extension_codec(
+ other.__datafusion_physical_extension_codec__()
+ )
+ ctx.set_query_planner(MyQueryPlanner())
+
+ batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+ assert batches[0].column(0).to_pylist() == [0, 1]
+ assert logical_codec.table_provider_encode_calls() > 0
+ assert logical_codec.table_provider_decode_calls() > 0
+ assert physical_codec.execution_plan_decode_calls() > 0
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index 2f2cc61..644c7b4 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -1793,7 +1793,8 @@
fallback inside it, which keeps the codecs it was imported with. Note
also that the planner is built against the codecs of the context this
method is called on, so installing the same planner again on a different
- handle rebinds the session's planner to *that* handle's codecs.
+ handle rebinds the session's planner to *that* handle's codecs. See the
+ FFI extensions guide for the full multi-library registration recipe.
Args:
planner: Object exposing ``__datafusion_query_planner__`` (see
@@ -2235,6 +2236,30 @@
"""Access the PyCapsule FFI_TaskContextProvider."""
return self.ctx.__datafusion_task_context_provider__()
+ @property
+ def __datafusion_codec_id__(self) -> str:
+ """Identity this context carries when installed as an extension codec.
+
+ A context can be installed on another session as an extension codec,
+ which tags the payloads it writes with this string. It is unique per
+ session, so two contexts can be installed on one session and a plan
+ written through one will not be decoded by the other.
+
+ Contexts derived from the same session — including the ones returned by
+ :py:meth:`with_logical_extension_codec` and
+ :py:meth:`with_python_udf_inlining` — report the same id, so only one of
+ them can be installed on a given session.
+
+ Examples:
+ >>> from datafusion import SessionContext
+ >>> ctx = SessionContext()
+ >>> ctx.__datafusion_codec_id__.startswith("session:")
+ True
+ >>> ctx.__datafusion_codec_id__ == SessionContext().__datafusion_codec_id__
+ False
+ """
+ return self.ctx.__datafusion_codec_id__
+
def __datafusion_logical_extension_codec__(self, session: Any = None) -> Any:
"""Access the PyCapsule FFI_LogicalExtensionCodec.
@@ -2253,26 +2278,79 @@
return self.ctx.__datafusion_query_planner__(session)
def with_logical_extension_codec(
- self, codec: LogicalExtensionCodecExportable | _PyCapsule
+ self,
+ codec: LogicalExtensionCodecExportable | _PyCapsule,
+ codec_id: str | None = None,
) -> SessionContext:
- """Create a new session context with specified codec.
+ """Create a new session context with an additional logical codec.
Only FFI codecs are supported. Pass any object implementing
``__datafusion_logical_extension_codec__`` (see
:py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`).
+ Codecs compose: each call appends the codec rather than replacing
+ codecs installed earlier, so one session can carry codecs from several
+ independent libraries and the order they are installed in does not
+ affect decoding.
+
+ A serialized plan records which codec wrote each payload, as a short id
+ taken from the codec's class. ``codec_id`` overrides that id and is
+ normally unnecessary. Pass it when installing from a bare ``PyCapsule``,
+ which has no class to take an id from, or when installing two instances
+ of one class, which otherwise claim the same id and raise ``ValueError``.
+
The returned context shares its session state with the original, so a
- later registration on either is visible to both. If a custom query
- planner is installed, it is rebuilt against the new codec on the shared
- session, so the original context plans with the new codec too. This
- happens on the shared session, so it takes effect even if the returned
- context is discarded.
+ later registration on either is visible to both, and an installed query
+ planner is rebound on the shared session even if the returned context is
+ discarded.
+
+ See :ref:`ffi` in the online documentation for how ids are assigned,
+ what an extension codec has to implement, and a worked multi-library
+ registration recipe.
+
+ Examples:
+ >>> from datafusion import SessionContext
+ >>> ctx = SessionContext()
+ >>> ctx = ctx.with_logical_extension_codec(
+ ... my_library.Codec()
+ ... ) # doctest: +SKIP
+
+ Installing from a bare capsule, pinning the id so encoded
+ plans remain decodable on another session:
+
+ >>> ctx = ctx.with_logical_extension_codec(
+ ... capsule, codec_id="my_library.Codec"
+ ... ) # doctest: +SKIP
"""
- new_internal = self.ctx.with_logical_extension_codec(codec)
+ new_internal = self.ctx.with_logical_extension_codec(codec, codec_id)
new = SessionContext.__new__(SessionContext)
new.ctx = new_internal
return new
+ def logical_extension_codec_ids(self) -> list[str]:
+ """List the logical extension codecs installed on this session.
+
+ Returns the identity of each installed codec, in install order. Those
+ identities are what encoding stamps onto a payload and what decoding
+ dispatches on, so this is how to check which library owns a plan and
+ whether a session is able to decode one.
+
+ DataFusion's own default codec is not listed. It handles whatever no
+ installed codec claims, and it carries no identity to list.
+
+ Examples:
+ >>> from datafusion import SessionContext
+ >>> ctx = SessionContext()
+ >>> ctx.logical_extension_codec_ids()
+ []
+ >>> ctx = ctx.with_logical_extension_codec(
+ ... my_library.Codec()
+ ... ) # doctest: +SKIP
+ >>> ctx.logical_extension_codec_ids() # doctest: +SKIP
+ ['my_library.Codec']
+ """
+ return self.ctx.logical_extension_codec_ids()
+
def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any:
"""Access the PyCapsule FFI_PhysicalExtensionCodec.
@@ -2280,23 +2358,46 @@
"""
return self.ctx.__datafusion_physical_extension_codec__(session)
+ def physical_extension_codec_ids(self) -> list[str]:
+ """List the physical extension codecs installed on this session.
+
+ See :py:meth:`logical_extension_codec_ids`.
+
+ Examples:
+ >>> from datafusion import SessionContext
+ >>> ctx = SessionContext()
+ >>> ctx.physical_extension_codec_ids()
+ []
+ """
+ return self.ctx.physical_extension_codec_ids()
+
def with_physical_extension_codec(
- self, codec: PhysicalExtensionCodecExportable | _PyCapsule
+ self,
+ codec: PhysicalExtensionCodecExportable | _PyCapsule,
+ codec_id: str | None = None,
) -> SessionContext:
- """Create a new session context with the specified physical codec.
+ """Create a new session context with an additional physical codec.
Only FFI codecs are supported. Pass any object implementing
``__datafusion_physical_extension_codec__`` (see
:py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`).
- The returned context shares its session state with the original, so a
- later registration on either is visible to both. If a custom query
- planner is installed, it is rebuilt against the new codec on the shared
- session, so the original context plans with the new codec too. This
- happens on the shared session, so it takes effect even if the returned
- context is discarded.
+ Composes and assigns an id exactly as
+ :py:meth:`with_logical_extension_codec` does, including when to pass
+ ``codec_id`` and what the returned context shares. See that method.
+
+ Examples:
+ >>> from datafusion import SessionContext
+ >>> ctx = SessionContext()
+ >>> ctx = ctx.with_physical_extension_codec(
+ ... my_library.PhysicalCodec()
+ ... ) # doctest: +SKIP
+
+ >>> ctx = ctx.with_physical_extension_codec(
+ ... capsule, codec_id="my_library.PhysicalCodec"
+ ... ) # doctest: +SKIP
"""
- new_internal = self.ctx.with_physical_extension_codec(codec)
+ new_internal = self.ctx.with_physical_extension_codec(codec, codec_id)
new = SessionContext.__new__(SessionContext)
new.ctx = new_internal
return new
diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py
index 43b53e4..eafcefd 100644
--- a/python/datafusion/user_defined.py
+++ b/python/datafusion/user_defined.py
@@ -120,6 +120,18 @@
is being installed on. Take the task context provider from it rather than
building a session of your own, so the decode callbacks resolve names
against the session that runs the query.
+
+ Implement the codec itself exactly as you would for a session that installs
+ only yours. A session may hold several codecs, but each payload records the
+ codec that wrote it and is only ever handed back to that codec, so there is
+ no need to recognise or reject another library's payloads.
+
+ A serialized plan records which codec wrote each payload, as a short id
+ taken from your class's module and qualified name. An optional
+ ``__datafusion_codec_id__`` attribute pins that id instead. It is not part
+ of this protocol and is rarely needed: declare it when renaming your class
+ must not stop older plans from decoding, or when one library installs two
+ instances that own disjoint slices of the wire format.
"""
def __datafusion_logical_extension_codec__( # noqa: D105
@@ -130,7 +142,9 @@
class PhysicalExtensionCodecExportable(Protocol):
"""Type hint for objects exposing ``__datafusion_physical_extension_codec__``.
- See :py:class:`LogicalExtensionCodecExportable` for ``session``.
+ See :py:class:`LogicalExtensionCodecExportable` for ``session``, for why a
+ codec need not recognise other libraries' payloads, and for
+ ``__datafusion_codec_id__``.
"""
def __datafusion_physical_extension_codec__( # noqa: D105
diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py
index 451f5d2..dc55a07 100644
--- a/python/tests/test_pickle_expr.py
+++ b/python/tests/test_pickle_expr.py
@@ -323,6 +323,49 @@
):
Expr.from_bytes(bytes(tampered))
+ def test_unsupported_wire_version_error_message(self):
+ """A payload stamped with a wire-format version newer than this
+ build supports names both versions and points at the fix, rather
+ than failing deep inside cloudpickle with an opaque tuple-unpack
+ error.
+
+ Patches the version byte at offset 7 of the frame described in
+ :meth:`test_cross_version_error_message`. The patch is
+ length-preserving, so the enclosing protobuf stays parseable and
+ the bytes reach the codec.
+ """
+ e = _double_udf()(col("a"))
+ blob = e.to_bytes()
+
+ idx = blob.find(b"DFPYUDF")
+ assert idx >= 0, "DFPYUDF frame not found in payload"
+
+ tampered = bytearray(blob)
+ tampered[idx + 7] = 2 # WIRE_VERSION_CURRENT is 1
+
+ with pytest.raises(Exception, match="wire-format version v2") as excinfo:
+ Expr.from_bytes(bytes(tampered))
+ assert "Align datafusion-python versions" in str(excinfo.value)
+
+ def test_cross_major_version_error_message(self):
+ """Same diagnostic as the minor-version mismatch, driven from the
+ major byte at offset 8. Guards against a check that compares only
+ the minor component."""
+ import sys
+
+ e = _double_udf()(col("a"))
+ blob = e.to_bytes()
+
+ idx = blob.find(b"DFPYUDF")
+ assert idx >= 0, "DFPYUDF frame not found in payload"
+
+ tampered = bytearray(blob)
+ tampered[idx + 8] = (sys.version_info.major + 1) % 256
+
+ with pytest.raises(Exception, match="not portable") as excinfo:
+ Expr.from_bytes(bytes(tampered))
+ assert f"Python {sys.version_info.major + 1}." in str(excinfo.value)
+
class TestPythonUdfInliningToggle:
"""`SessionContext.with_python_udf_inlining(enabled=False)` opts out of
