Make extension codecs composable - #1678
Open
timsaucer wants to merge 3 commits into
Open
Conversation
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>
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>
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>
timsaucer
force-pushed
the
feat/ffi-composable-codecs
branch
from
August 8, 2026 18:09
8242ceb to
b32667e
Compare
ntjohnson1
approved these changes
Aug 14, 2026
ntjohnson1
left a comment
ntjohnson1
left a comment
Contributor
There was a problem hiding this comment.
Hmm is it useful to be able to see for a given call which codec triggers? I see you can check the counts before/after with something like codec.table_provider_encode_calls() but I wonder if this does/should show up in an explain plan or something.
| } | ||
| } | ||
|
|
||
| pub fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> { |
Contributor
There was a problem hiding this comment.
This was pub before. Do people care about inspecting the codecs directly?
| } | ||
| } | ||
|
|
||
| pub fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> { |
Contributor
There was a problem hiding this comment.
Similar note about directly referencing codec chain
|
|
||
| # 1. 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()) |
Contributor
There was a problem hiding this comment.
So is the resulting stage at this point [default, lib_a, lib_b]? I could read it or ask claude if some assumed based codec for standard datafusion python gets installed always as a fallback. Mostly I'm curious if it's a reasonable workflow for someone to set things up to only get [lib_a, lib_b] and if anything unsupported by their custom codecs barfs.
| The example codecs do not inspect the callback `TaskContext`. A production codec that depends on session configuration or registered functions must ensure its exported FFI codec is bound to, and retains, the appropriate host `TaskContextProvider`. | ||
|
|
||
| The current Python API installs one external logical codec and one external physical codec. It does not yet compose codecs from several independent plugin owners. This example therefore makes the provider library the sole external codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host. | ||
| Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call prepends the codec to the session's codec chain, with the most recently installed codec consulted first and DataFusion's default codec as the terminal fallback. A codec signals "not mine" by returning an error, so several independent plugin libraries can install codecs on the same session as long as each only answers for payloads it owns (frame them with a distinct byte prefix). In this example the provider library is the only codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host. |
Contributor
There was a problem hiding this comment.
Ok this line says DF default codec is the terminal fallback but I think the constructor takes a codec in rust at least so it could potentially be updated there but maybe not in the exposed python surface.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Part 2 of 3 in the split of #1672. These are enabled as a github stack so you should be able to swab between the 3 PRs in github interface (above, next to the "Open" oval).
Rationale for this change
Supporting a foreign planner surfaced a codec problem: a query can involve three independent native libraries (datafusion-python, a provider library, and a planner library), and each library needs its extension codecs active on the session at the same time. Previously, installing a logical or physical extension codec replaced the prior codec, so the second library's install silently discarded the first — plans then failed later with a confusing decode error.
What changes are included in this PR?
with_logical_extension_codec/with_physical_extension_codecnow prepend to a codec chain instead of replacing the prior codec. The most recently installed codec is consulted first, falling through codec by codec to DataFusion's default codec. A codec signals "not mine" by returning an error.python_udf_inliningback to enabled.docs/source/contributor-guide/ffi.mdgains sections on composable codecs: family-prefix discipline, and that registration order between libraries no longer matters.Test coverage
The chain behavior was originally covered by two
#[cfg(test)]modules incrates/core/src/codec.rs. Those tests never ran: no workflow invokescargo test, and the only Rust checks arecargo fmt --checkandcargo clippy --no-deps --all-targets.--all-targetscompiles test code, so the tests could not rot into a non-compiling state, but a behavioral regression would not have failed the build. Adding acargo testjob is also not a one-line change, becausecrates/core/Cargo.tomlenablespyo3/extension-moduleunconditionally and the test binary therefore fails to link againstPy_*on Linux. The coverage was moved to pytest instead, matching this repository's practice of treating the user-facing Python surface as the primary focus.strip_errors_on_too_old_version) asserted nothing becauseWIRE_VERSION_MIN_SUPPORTEDequalsWIRE_VERSION_CURRENT.python/tests/test_pickle_expr.pygains coverage for the wire-header diagnostics — an unsupported wire-format version and a Python major-version mismatch — by patching the header in place inside the encoded protobuf. The patches preserve length so the outer message stays parseable and the bytes reach the codec.python_udf_inliningacross a codec install.python_udf_inlining()accessor the Rust test checked. This is a stronger assertion: a codec exported from anotherSessionContextis itself a Python-aware codec with inlining enabled, so the strict outer codec delegates to it and the inline payload reappears even though the accessor still reports strict. The test uses an extension codec that delegates UDF encoding to DataFusion's default codec, which is the realistic case.AGENTS.mdrecords the Python-first testing preference and the fact that CI does not run Rust tests, so the tradeoff does not have to be rediscovered.Are there any user-facing changes?
Behavior change: installing an extension codec now composes with previously installed codecs instead of replacing them. Code that relied on replacement semantics (installing a codec to remove a prior one) is affected; all other usage keeps working and no longer loses earlier codecs.
No change to the published
datafusionpackage beyond the above.MyLogicalExtensionCodecinexamples/datafusion-ffi-examplegains an optional token argument that overrides the byte prefix it stamps on encoded table providers, which is what lets the tests install two instances owning disjoint slices of the wire format. It defaults to the previous constant, so existing call sites are unchanged, and the example README notes that real plugin libraries should hard-code a prefix unique to the library rather than accept one from the caller.