Skip to content

Make extension codecs composable - #1678

Open
timsaucer wants to merge 3 commits into
feat/ffi-query-planner-corefrom
feat/ffi-composable-codecs
Open

Make extension codecs composable#1678
timsaucer wants to merge 3 commits into
feat/ffi-query-planner-corefrom
feat/ffi-composable-codecs

Conversation

@timsaucer

@timsaucer timsaucer commented Aug 7, 2026

Copy link
Copy Markdown
Member

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_codec now 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.
  • 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 in the chain fails, the errors are aggregated so the owning codec's diagnostic is not masked by the default codec's generic error.
  • Fixed a latent bug where installing a codec silently reset python_udf_inlining back to enabled.
  • docs/source/contributor-guide/ffi.md gains 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 in crates/core/src/codec.rs. Those tests never ran: no workflow invokes cargo test, and the only Rust checks are cargo fmt --check and cargo clippy --no-deps --all-targets. --all-targets compiles test code, so the tests could not rot into a non-compiling state, but a behavioral regression would not have failed the build. Adding a cargo test job is also not a one-line change, because crates/core/Cargo.toml enables pyo3/extension-module unconditionally and the test binary therefore fails to link against Py_* on Linux. The coverage was moved to pytest instead, matching this repository's practice of treating the user-facing Python surface as the primary focus.

  • Both Rust test modules are removed. Five of the eighteen cases were already covered by existing pytest cases, and one (strip_errors_on_too_old_version) asserted nothing because WIRE_VERSION_MIN_SUPPORTED equals WIRE_VERSION_CURRENT.
  • python/tests/test_pickle_expr.py gains coverage for the wire-header diagnostics — an unsupported wire-format version and a Python major-version mismatch — by patching the header in place inside the encoded protobuf. The patches preserve length so the outer message stays parseable and the bytes reach the codec.
  • The FFI example suite gains coverage for the chain itself, exercised across a real FFI boundary: encode ordering, decode fall-through to an earlier-installed codec, the aggregated multi-codec error, the verbatim single-codec error, and preservation of python_udf_inlining across a codec install.
  • Three cases were dropped rather than ported. They truncate the wire header, which changes the payload length and breaks the enclosing protobuf framing, so they cannot reach the header check from Python.
  • The inlining test asserts encode and decode behavior rather than the python_udf_inlining() accessor the Rust test checked. This is a stronger assertion: a codec exported from another SessionContext is 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.md records the Python-first testing preference and the fact that CI does not run Rust tests, so the tradeoff does not have to be rediscovered.

Are there any user-facing changes?

Behavior change: installing an extension codec 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 datafusion package beyond the above. MyLogicalExtensionCodec in examples/datafusion-ffi-example gains 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.

timsaucer and others added 3 commits August 8, 2026 14:09
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
timsaucer force-pushed the feat/ffi-composable-codecs branch from 8242ceb to b32667e Compare August 8, 2026 18:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant