You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ti...@apache.org on 2026/08/31 11:05:55 UTC
(datafusion-python) branch feat/ffi-composable-codecs updated (b32667eb -> c1ffd81c)
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a change to branch feat/ffi-composable-codecs
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
omit b32667eb docs: record the Python-first testing preference in AGENTS.md
omit feed4f05 test: port codec Rust tests to pytest
omit 8fd8444b Make extension codecs composable
omit 7982ee5b fix: install FFI test wheels from nested artifact paths
omit 27147f3b add rat
omit 05ab5da6 Update FFI query planner integration
omit 47aea03d Add three-library FFI planner example
omit 9ba147c1 Add FFI query planner support
add c577c0f4 [PyLimit] adding skip and fetch with datafusion expr support (#1674)
add 38422e44 Bump to 55.0.0 rc2 (#1685)
add 2bd80eca chore: Bump release candidate version to 55.0.0rc3 (#1686)
add 2f528d2d chore: use DataFusion 55.0.0 from crates.io (#1695)
add e6c75870 Add FFI query planner support (#1677)
new d28615db Make extension codecs composable
new 62c3d38c test: port codec Rust tests to pytest
new c1ffd81c docs: record the Python-first testing preference in AGENTS.md
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (b32667eb)
\
N -- N -- N refs/heads/feat/ffi-composable-codecs (c1ffd81c)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.ai/skills/audit-skill-md/SKILL.md | 12 +-
.ai/skills/check-upstream/SKILL.md | 12 +-
.ai/skills/ffi-capsule-protocol/SKILL.md | 190 +++++++
.ai/skills/make-pythonic/SKILL.md | 12 +-
AGENTS.md | 24 +-
Cargo.lock | 170 +++---
Cargo.toml | 53 +-
crates/core/Cargo.toml | 1 -
crates/core/src/catalog.rs | 21 +-
crates/core/src/codec.rs | 11 +
crates/core/src/context.rs | 336 ++++++------
crates/core/src/dataset_exec.rs | 12 +-
crates/core/src/expr/limit.rs | 29 +-
crates/core/src/udaf.rs | 8 +-
crates/core/src/udf.rs | 8 +-
crates/core/src/udtf.rs | 18 +-
crates/core/src/udwf.rs | 9 +-
crates/util/Cargo.toml | 1 -
crates/util/src/lib.rs | 286 ++++++++--
docs/source/contributor-guide/ffi.md | 214 ++++++--
docs/source/user-guide/io/table_provider.md | 13 +-
docs/source/user-guide/upgrade-guides.md | 122 +++++
examples/datafusion-ffi-example/README.md | 12 +-
examples/datafusion-ffi-example/pyproject.toml | 3 +-
.../python/tests/_test_logical_extension_codec.py | 18 +-
.../datafusion-ffi-example/src/catalog_provider.rs | 6 +-
examples/datafusion-ffi-example/src/lib.rs | 1 +
.../src/logical_extension_codec.rs | 91 +++-
.../src/physical_extension_codec.rs | 59 ++-
.../datafusion-ffi-example/src/required_udf.rs | 102 ++++
.../datafusion-ffi-example/src/table_function.rs | 2 +-
.../datafusion-ffi-example/src/table_provider.rs | 2 +-
.../src/table_provider_factory.rs | 2 +-
.../Cargo.toml | 1 -
.../datafusion-ffi-query-planner-example/README.md | 10 +-
.../tests/_test_three_library_query_planner.py | 573 ++++++++++++++++++++-
.../python/tests/conftest.py | 0
.../src/config.rs | 16 +-
.../src/lib.rs | 4 +-
.../src/planner.rs | 207 ++++++--
python/datafusion/context.py | 125 +++--
python/datafusion/user_defined.py | 21 +-
python/tests/test_context.py | 202 +++++++-
python/tests/test_expr.py | 4 +
python/tests/test_pickle_expr.py | 39 ++
python/tests/test_plans.py | 43 ++
python/tests/test_spark_functions.py | 4 -
python/tests/test_udtf.py | 30 ++
48 files changed, 2554 insertions(+), 585 deletions(-)
create mode 100644 .ai/skills/ffi-capsule-protocol/SKILL.md
create mode 100644 examples/datafusion-ffi-example/src/required_udf.rs
copy examples/{datafusion-ffi-example => datafusion-ffi-query-planner-example}/python/tests/conftest.py (100%)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 03/03: docs: record the Python-first testing preference in AGENTS.md
Posted by ti...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a commit to branch feat/ffi-composable-codecs
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit c1ffd81ca02ec969e84a349e7fef1d8fcf9d854f
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Sat Aug 8 11:58:07 2026 -0400
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) <no...@anthropic.com>
---
AGENTS.md | 34 ++++++++++++++++++++++++++++++++++
1 file changed, 34 insertions(+)
diff --git a/AGENTS.md b/AGENTS.md
index 327ebd64..659094ec 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -82,6 +82,40 @@ pre-commit run --all-files
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.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 01/03: Make extension codecs composable
Posted by ti...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a commit to branch feat/ffi-composable-codecs
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit d28615db6b92b0600f0eb47fa694499c02127d2f
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 7 10:34:45 2026 -0400
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 <no...@anthropic.com>
---
crates/core/src/codec.rs | 444 ++++++++++++++++++---
crates/core/src/context.rs | 33 +-
docs/source/contributor-guide/ffi.md | 51 ++-
examples/datafusion-ffi-example/README.md | 6 +-
.../python/tests/_test_logical_extension_codec.py | 27 +-
.../python/tests/_test_physical_extension_codec.py | 23 ++
.../datafusion-ffi-query-planner-example/README.md | 4 +-
.../tests/_test_three_library_query_planner.py | 23 ++
python/datafusion/context.py | 24 +-
9 files changed, 545 insertions(+), 90 deletions(-)
diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs
index 94942a2d..0f7e20c9 100644
--- a/crates/core/src/codec.rs
+++ b/crates/core/src/codec.rs
@@ -29,16 +29,16 @@
//!
//! [`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 consulted in
+//! most-recently-installed-first order, with
+//! `DefaultLogicalExtensionCodec` as the terminal fallback.
//!
//! [`PythonPhysicalCodec`] is the symmetric wrapper around
//! [`PhysicalExtensionCodec`]. Logical and physical layers each have
@@ -58,7 +58,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 +81,11 @@
//! 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 payloads and objects they do not own — 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 shadows
+//! every codec installed before it.
use std::sync::Arc;
@@ -167,7 +168,7 @@ fn write_wire_header(buf: &mut Vec<u8>, family: &[u8], py_version: (u8, u8)) {
/// 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 +224,100 @@ fn strip_wire_header<'a>(
Ok(Some(&buf[py_minor_idx + 1..]))
}
+/// Run `f` against each codec in `chain`, returning the first `Ok`.
+///
+/// A codec signals "not mine" by returning an error, so the chain
+/// keeps trying until a codec succeeds. When every codec fails and the
+/// chain has more than one entry, the errors are aggregated into a
+/// single message — returning only the last error would surface the
+/// terminal `Default*ExtensionCodec` "not provided" message and mask
+/// the more specific diagnostic from an installed codec (e.g. a
+/// corrupt-token error from the codec that owns the payload family).
+fn chain_try<C: ?Sized, R>(chain: &[Arc<C>], what: &str, f: impl Fn(&C) -> Result<R>) -> Result<R> {
+ let mut errors: Vec<datafusion::error::DataFusionError> = Vec::new();
+ for codec in chain {
+ match f(codec) {
+ Ok(value) => return 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 the one-codec (default-only) chain behaves
+/// exactly like the pre-chain implementation.
+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!(
+ "None of the {} composed extension codecs handled {what}: {joined}",
+ errors.len()
+ ))
+ }
+ }
+}
+
+/// Encode variant of [`chain_try`] for methods that write into a
+/// caller-provided buffer.
+///
+/// Each codec encodes into a scratch buffer so a failed attempt cannot
+/// leave partial bytes behind. `Ok` with bytes written commits those
+/// bytes and ends the chain. `Ok` with an empty buffer is treated as
+/// "no opinion" — the standard `Default*ExtensionCodec` behavior of
+/// encoding a UDF by name writes nothing — so later codecs still get a
+/// chance to emit a richer payload. If no codec writes bytes but at
+/// least one returned `Ok`, the overall result is `Ok` with nothing
+/// written (encode by name).
+fn chain_encode<C: ?Sized>(
+ chain: &[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 codec in chain {
+ let mut scratch = Vec::new();
+ match f(codec, &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. The chain starts as just
+/// `DefaultLogicalExtensionCodec`; each downstream FFI codec installed
+/// via `SessionContext.with_logical_extension_codec(...)` is prepended,
+/// so the most recently installed codec is consulted first and the
+/// default codec always runs last.
+///
+/// Chain dispatch relies on each codec recognizing its own payloads
+/// (distinct family prefixes — see the module docs) and returning an
+/// error for everything else so the next codec gets a chance.
///
/// Sitting at the top of the session's logical codec stack means
/// every serializer that reads `session.logical_codec()` automatically
@@ -241,22 +330,31 @@ fn strip_wire_header<'a>(
/// 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<Arc<dyn LogicalExtensionCodec>>,
python_udf_inlining: bool,
}
impl PythonLogicalCodec {
pub fn new(inner: Arc<dyn LogicalExtensionCodec>) -> Self {
Self {
- inner,
+ chain: vec![inner],
python_udf_inlining: true,
}
}
- pub fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> {
- &self.inner
+ /// Return a copy of this codec with `codec` prepended to the
+ /// chain, preserving the Python-UDF-inlining setting. The new
+ /// codec is consulted before every previously installed codec.
+ pub fn with_additional_codec(&self, codec: Arc<dyn LogicalExtensionCodec>) -> Self {
+ let mut chain = Vec::with_capacity(self.chain.len() + 1);
+ chain.push(codec);
+ chain.extend(self.chain.iter().map(Arc::clone));
+ Self {
+ chain,
+ python_udf_inlining: self.python_udf_inlining,
+ }
}
/// Toggle inline encoding of Python UDFs. See
@@ -297,11 +395,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
inputs: &[LogicalPlan],
ctx: &TaskContext,
) -> Result<Extension> {
- self.inner.try_decode(buf, inputs, ctx)
+ chain_try(&self.chain, "an extension logical plan node", |codec| {
+ 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,
+ buf,
+ "an extension logical plan node",
+ |codec, buf| codec.try_encode(node, buf),
+ )
}
fn try_decode_table_provider(
@@ -311,8 +416,9 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
schema: SchemaRef,
ctx: &TaskContext,
) -> Result<Arc<dyn TableProvider>> {
- self.inner
- .try_decode_table_provider(buf, table_ref, schema, ctx)
+ chain_try(&self.chain, "a table provider", |codec| {
+ codec.try_decode_table_provider(buf, table_ref, Arc::clone(&schema), ctx)
+ })
}
fn try_encode_table_provider(
@@ -321,7 +427,9 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
node: Arc<dyn TableProvider>,
buf: &mut Vec<u8>,
) -> Result<()> {
- self.inner.try_encode_table_provider(table_ref, node, buf)
+ chain_encode(&self.chain, buf, "a table provider", |codec, buf| {
+ codec.try_encode_table_provider(table_ref, Arc::clone(&node), buf)
+ })
}
fn try_decode_file_format(
@@ -329,7 +437,9 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
buf: &[u8],
ctx: &TaskContext,
) -> Result<Arc<dyn FileFormatFactory>> {
- self.inner.try_decode_file_format(buf, ctx)
+ chain_try(&self.chain, "a file format", |codec| {
+ codec.try_decode_file_format(buf, ctx)
+ })
}
fn try_encode_file_format(
@@ -337,14 +447,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
buf: &mut Vec<u8>,
node: Arc<dyn FileFormatFactory>,
) -> Result<()> {
- self.inner.try_encode_file_format(buf, node)
+ chain_encode(&self.chain, 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, 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 +469,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
} else {
refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
}
- self.inner.try_decode_udf(name, buf)
+ chain_try(&self.chain, "a scalar UDF", |codec| {
+ 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, 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 +491,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
} else {
refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
}
- self.inner.try_decode_udaf(name, buf)
+ chain_try(&self.chain, "an aggregate UDF", |codec| {
+ 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, 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 +513,15 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
} else {
refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
}
- self.inner.try_decode_udwf(name, buf)
+ chain_try(&self.chain, "a window UDF", |codec| {
+ 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 +566,8 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio
/// `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 +579,31 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio
///
/// 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<Arc<dyn PhysicalExtensionCodec>>,
python_udf_inlining: bool,
}
impl PythonPhysicalCodec {
pub fn new(inner: Arc<dyn PhysicalExtensionCodec>) -> Self {
Self {
- inner,
+ chain: vec![inner],
python_udf_inlining: true,
}
}
- pub fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> {
- &self.inner
+ /// Return a copy of this codec with `codec` prepended to the
+ /// chain, preserving the Python-UDF-inlining setting. The new
+ /// codec is consulted before every previously installed codec.
+ pub fn with_additional_codec(&self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self {
+ let mut chain = Vec::with_capacity(self.chain.len() + 1);
+ chain.push(codec);
+ chain.extend(self.chain.iter().map(Arc::clone));
+ Self {
+ chain,
+ python_udf_inlining: self.python_udf_inlining,
+ }
}
/// Toggle inline encoding of Python UDFs on this physical codec.
@@ -500,7 +634,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
ctx: &TaskContext,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
- self.inner.try_decode(buf, inputs, ctx, proto_converter)
+ chain_try(&self.chain, "an execution plan", |codec| {
+ codec.try_decode(buf, inputs, ctx, proto_converter)
+ })
}
fn try_encode(
@@ -509,14 +645,18 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
buf: &mut Vec<u8>,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<()> {
- self.inner.try_encode(node, buf, proto_converter)
+ chain_encode(&self.chain, 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, 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 +667,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
} else {
refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?;
}
- self.inner.try_decode_udf(name, buf)
+ chain_try(&self.chain, "a scalar UDF", |codec| {
+ codec.try_decode_udf(name, buf)
+ })
}
fn try_encode_expr(
@@ -536,7 +678,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
buf: &mut Vec<u8>,
ctx: &PhysicalExprEncodeCtx<'_>,
) -> Result<()> {
- self.inner.try_encode_expr(node, buf, ctx)
+ chain_encode(&self.chain, buf, "a physical expression", |codec, buf| {
+ codec.try_encode_expr(node, buf, ctx)
+ })
}
fn try_decode_expr(
@@ -545,14 +689,18 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
inputs: &[Arc<dyn PhysicalExpr>],
ctx: &PhysicalExprDecodeCtx<'_>,
) -> Result<Arc<dyn PhysicalExpr>> {
- self.inner.try_decode_expr(buf, inputs, ctx)
+ chain_try(&self.chain, "a physical expression", |codec| {
+ 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, 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 +711,18 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
} else {
refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?;
}
- self.inner.try_decode_udaf(name, buf)
+ chain_try(&self.chain, "an aggregate UDF", |codec| {
+ 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, 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 +733,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
} else {
refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?;
}
- self.inner.try_decode_udwf(name, buf)
+ chain_try(&self.chain, "a window UDF", |codec| {
+ codec.try_decode_udwf(name, buf)
+ })
}
}
@@ -598,7 +752,7 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec {
/// `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 +767,7 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec<u8>)
/// 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) {
@@ -1173,3 +1327,181 @@ mod wire_header_tests {
));
}
}
+
+#[cfg(test)]
+mod codec_chain_tests {
+ use std::sync::atomic::{AtomicUsize, Ordering};
+
+ use datafusion::catalog::MemTable;
+ use datafusion::common::exec_err;
+
+ use super::*;
+
+ /// Codec that owns a single byte token for table providers and
+ /// errors on everything else, mirroring the family-prefix
+ /// discipline expected of downstream FFI codecs.
+ #[derive(Debug)]
+ struct TokenCodec {
+ token: &'static [u8],
+ /// Return `Ok` from `try_encode_table_provider` without
+ /// writing bytes, imitating a "no opinion" codec.
+ encode_by_name: bool,
+ decode_hits: AtomicUsize,
+ encode_hits: AtomicUsize,
+ }
+
+ impl TokenCodec {
+ fn new(token: &'static [u8]) -> Arc<Self> {
+ Arc::new(Self {
+ token,
+ encode_by_name: false,
+ decode_hits: AtomicUsize::new(0),
+ encode_hits: AtomicUsize::new(0),
+ })
+ }
+
+ fn new_by_name(token: &'static [u8]) -> Arc<Self> {
+ Arc::new(Self {
+ token,
+ encode_by_name: true,
+ decode_hits: AtomicUsize::new(0),
+ encode_hits: AtomicUsize::new(0),
+ })
+ }
+ }
+
+ impl LogicalExtensionCodec for TokenCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[LogicalPlan],
+ _ctx: &TaskContext,
+ ) -> Result<Extension> {
+ exec_err!("TokenCodec does not decode extension nodes")
+ }
+
+ fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
+ exec_err!("TokenCodec does not encode extension nodes")
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ buf: &[u8],
+ _table_ref: &TableReference,
+ schema: SchemaRef,
+ _ctx: &TaskContext,
+ ) -> Result<Arc<dyn TableProvider>> {
+ if buf != self.token {
+ return exec_err!("Unknown table provider token for TokenCodec");
+ }
+ self.decode_hits.fetch_add(1, Ordering::SeqCst);
+ Ok(Arc::new(MemTable::try_new(schema, vec![vec![]])?))
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ _table_ref: &TableReference,
+ _node: Arc<dyn TableProvider>,
+ buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ self.encode_hits.fetch_add(1, Ordering::SeqCst);
+ if !self.encode_by_name {
+ buf.extend_from_slice(self.token);
+ }
+ Ok(())
+ }
+ }
+
+ fn mem_table() -> Arc<dyn TableProvider> {
+ Arc::new(MemTable::try_new(Arc::new(Schema::empty()), vec![vec![]]).unwrap())
+ }
+
+ fn table_ref() -> TableReference {
+ TableReference::bare("t")
+ }
+
+ #[test]
+ fn decode_falls_through_to_earlier_installed_codec() {
+ let first = TokenCodec::new(b"AAAA");
+ let second = TokenCodec::new(b"BBBB");
+ let codec = PythonLogicalCodec::default()
+ .with_additional_codec(first.clone())
+ .with_additional_codec(second.clone());
+
+ let ctx = TaskContext::default();
+ codec
+ .try_decode_table_provider(b"AAAA", &table_ref(), Arc::new(Schema::empty()), &ctx)
+ .unwrap();
+
+ assert_eq!(first.decode_hits.load(Ordering::SeqCst), 1);
+ assert_eq!(second.decode_hits.load(Ordering::SeqCst), 0);
+ }
+
+ #[test]
+ fn most_recently_installed_codec_encodes_first() {
+ let first = TokenCodec::new(b"AAAA");
+ let second = TokenCodec::new(b"BBBB");
+ let codec = PythonLogicalCodec::default()
+ .with_additional_codec(first.clone())
+ .with_additional_codec(second.clone());
+
+ let mut buf = Vec::new();
+ codec
+ .try_encode_table_provider(&table_ref(), mem_table(), &mut buf)
+ .unwrap();
+
+ assert_eq!(buf, b"BBBB");
+ assert_eq!(first.encode_hits.load(Ordering::SeqCst), 0);
+ }
+
+ #[test]
+ fn empty_ok_encode_lets_later_codec_write_payload() {
+ let writer = TokenCodec::new(b"AAAA");
+ let by_name = TokenCodec::new_by_name(b"BBBB");
+ let codec = PythonLogicalCodec::default()
+ .with_additional_codec(writer.clone())
+ .with_additional_codec(by_name.clone());
+
+ let mut buf = Vec::new();
+ codec
+ .try_encode_table_provider(&table_ref(), mem_table(), &mut buf)
+ .unwrap();
+
+ assert_eq!(buf, b"AAAA");
+ assert_eq!(by_name.encode_hits.load(Ordering::SeqCst), 1);
+ assert_eq!(writer.encode_hits.load(Ordering::SeqCst), 1);
+ }
+
+ #[test]
+ fn decode_failure_aggregates_every_codec_error() {
+ let codec = PythonLogicalCodec::default()
+ .with_additional_codec(TokenCodec::new(b"AAAA"))
+ .with_additional_codec(TokenCodec::new(b"BBBB"));
+
+ let ctx = TaskContext::default();
+ let err = codec
+ .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx)
+ .unwrap_err();
+
+ let msg = err.to_string();
+ assert!(msg.contains("None of the 3 composed extension codecs"));
+ assert!(msg.contains("Unknown table provider token"));
+ }
+
+ #[test]
+ fn single_codec_chain_error_is_returned_verbatim() {
+ let codec = PythonLogicalCodec::default();
+ let ctx = TaskContext::default();
+ let err = codec
+ .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx)
+ .unwrap_err();
+ assert!(!err.to_string().contains("composed extension codecs"));
+ }
+
+ #[test]
+ fn with_additional_codec_preserves_udf_inlining_setting() {
+ let strict = PythonLogicalCodec::default().with_python_udf_inlining(false);
+ let extended = strict.with_additional_codec(TokenCodec::new(b"AAAA"));
+ assert!(!extended.python_udf_inlining());
+ }
+}
diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs
index 75bfed60..4f44ee84 100644
--- a/crates/core/src/context.rs
+++ b/crates/core/src/context.rs
@@ -1465,14 +1465,13 @@ impl PySessionContext {
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()),
- );
+ // Prepend rather than replace: previously installed codecs stay active,
+ // with the most recently installed one consulted first. Prepending 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(inner));
let derived = Self {
ctx: Arc::clone(&this.ctx),
logical_codec,
@@ -1503,11 +1502,9 @@ impl PySessionContext {
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 prepends 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(inner));
let derived = Self {
ctx: Arc::clone(&this.ctx),
logical_codec: Arc::clone(&this.logical_codec),
@@ -1525,7 +1522,7 @@ impl PySessionContext {
// 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 +1534,15 @@ impl PySessionContext {
}
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 {
diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md
index 31cd9391..67a1e738 100644
--- a/docs/source/contributor-guide/ffi.md
+++ b/docs/source/contributor-guide/ffi.md
@@ -248,10 +248,43 @@ foreign planner. This lets the planner decode provider-owned objects and lets
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` adds the codec to the front of the session's codec
+chain rather than replacing prior codecs. During encoding and decoding, 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, which
+sends the chain on to the next codec. Two conventions keep this dispatch sound:
+
+- Frame your payloads with a distinct byte prefix (pick a `DF` namespace plus a
+ crate-specific suffix) and only decode payloads carrying your prefix.
+- Return an error for objects and payloads you do not own. A codec that answers
+ success for objects outside its family shadows every codec installed before it.
+
+Because dispatch keys off payload prefixes rather than install position, codec
+registration order between independent libraries does not matter. 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()))
+```
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 +380,7 @@ side is visible to both.
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 +394,17 @@ that surprises people:
> 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/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md
index 0fa10d7f..4934b1a9 100644
--- a/examples/datafusion-ffi-example/README.md
+++ b/examples/datafusion-ffi-example/README.md
@@ -35,7 +35,9 @@ Separate shared libraries guarantee distinct DataFusion library markers. This ca
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 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). This exam [...]
+
+Register both provider codecs before installing the planner:
```python
ctx = ctx.with_logical_extension_codec(provider_logical_codec)
@@ -45,4 +47,4 @@ ctx.set_query_planner(planner)
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) [...]
-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 cd0c5a61..0fd1d743 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
@@ -18,7 +18,7 @@
from __future__ import annotations
from datafusion import LogicalPlan, SessionContext
-from datafusion_ffi_example import MyLogicalExtensionCodec
+from datafusion_ffi_example import MyLogicalExtensionCodec, MyTableProvider
def _setup_session_with_codec() -> tuple[SessionContext, MyLogicalExtensionCodec]:
@@ -80,3 +80,28 @@ def test_ffi_logical_codec_roundtrip():
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 prepends 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 encoding falls
+ through to the user codec installed first. Under replace semantics
+ this test fails with `LogicalExtensionCodec is not provided`."""
+ 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()
+
+ before = codec.table_provider_encode_calls()
+ blob = plan.to_bytes(ctx)
+ assert codec.table_provider_encode_calls() > before
+
+ restored = LogicalPlan.from_bytes(ctx, blob)
+ df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
+ assert df.collect() == df_round_trip.collect()
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 28eaaf2d..82116bef 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 @@ def test_ffi_physical_codec_roundtrip():
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 prepends 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-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md
index 72f96bb8..f188ec74 100644
--- a/examples/datafusion-ffi-query-planner-example/README.md
+++ b/examples/datafusion-ffi-query-planner-example/README.md
@@ -55,6 +55,6 @@ ctx.set_query_planner(MyQueryPlanner())
`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 [...]
+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 prepends to the session's codec chain, so several libraries can install codecs on the same session and the order between them does not matter. This planner owns no serializable types of its own and deliberately uses only built-in physical [...]
-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 d046f67a..c6f9189d 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
@@ -667,3 +667,26 @@ def test_query_planner_rejects_invalid_config(max_rows: str):
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 extra codecs
+ (default-backed exports from a fresh session) decline everything,
+ so planner-driven encode/decode falls through to the provider
+ codecs and the query still succeeds end to end."""
+ 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 2f2cc611..6a4fbdb1 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -1793,7 +1793,8 @@ class SessionContext:
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
@@ -2255,15 +2256,23 @@ class SessionContext:
def with_logical_extension_codec(
self, codec: LogicalExtensionCodecExportable | _PyCapsule
) -> 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 adds the codec to the front of the
+ session's codec chain rather than replacing prior codecs. During
+ encoding and decoding, the most recently installed codec is
+ consulted first, falling through codec by codec to DataFusion's
+ default codec. Codecs signal "not mine" by returning an error, so
+ extension codecs should only answer for payloads they own —
+ typically identified by a distinct byte prefix.
+
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
+ planner is installed, it is rebuilt against the new chain 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.
@@ -2283,15 +2292,20 @@ class SessionContext:
def with_physical_extension_codec(
self, codec: PhysicalExtensionCodecExportable | _PyCapsule
) -> 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`).
+ Codecs compose the same way as in
+ :py:meth:`with_logical_extension_codec`: each call prepends to the
+ session's codec chain, and the most recently installed codec is
+ consulted first.
+
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
+ planner is installed, it is rebuilt against the new chain 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.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 02/03: test: port codec Rust tests to pytest
Posted by ti...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a commit to branch feat/ffi-composable-codecs
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit 62c3d38ccbed22806f8633901a7f4a459cc8189a
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Sat Aug 8 11:56:18 2026 -0400
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) <no...@anthropic.com>
---
crates/core/src/codec.rs | 312 ---------------------
examples/datafusion-ffi-example/README.md | 2 +
.../python/tests/_test_logical_extension_codec.py | 136 ++++++++-
.../src/logical_extension_codec.rs | 23 +-
python/tests/test_pickle_expr.py | 43 +++
5 files changed, 199 insertions(+), 317 deletions(-)
diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs
index 0f7e20c9..0e990dec 100644
--- a/crates/core/src/codec.rs
+++ b/crates/core/src/codec.rs
@@ -1193,315 +1193,3 @@ fn decode_python_udaf(py: Python<'_>, payload: &[u8]) -> PyResult<PythonFunction
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)
- ));
- }
-}
-
-#[cfg(test)]
-mod codec_chain_tests {
- use std::sync::atomic::{AtomicUsize, Ordering};
-
- use datafusion::catalog::MemTable;
- use datafusion::common::exec_err;
-
- use super::*;
-
- /// Codec that owns a single byte token for table providers and
- /// errors on everything else, mirroring the family-prefix
- /// discipline expected of downstream FFI codecs.
- #[derive(Debug)]
- struct TokenCodec {
- token: &'static [u8],
- /// Return `Ok` from `try_encode_table_provider` without
- /// writing bytes, imitating a "no opinion" codec.
- encode_by_name: bool,
- decode_hits: AtomicUsize,
- encode_hits: AtomicUsize,
- }
-
- impl TokenCodec {
- fn new(token: &'static [u8]) -> Arc<Self> {
- Arc::new(Self {
- token,
- encode_by_name: false,
- decode_hits: AtomicUsize::new(0),
- encode_hits: AtomicUsize::new(0),
- })
- }
-
- fn new_by_name(token: &'static [u8]) -> Arc<Self> {
- Arc::new(Self {
- token,
- encode_by_name: true,
- decode_hits: AtomicUsize::new(0),
- encode_hits: AtomicUsize::new(0),
- })
- }
- }
-
- impl LogicalExtensionCodec for TokenCodec {
- fn try_decode(
- &self,
- _buf: &[u8],
- _inputs: &[LogicalPlan],
- _ctx: &TaskContext,
- ) -> Result<Extension> {
- exec_err!("TokenCodec does not decode extension nodes")
- }
-
- fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
- exec_err!("TokenCodec does not encode extension nodes")
- }
-
- fn try_decode_table_provider(
- &self,
- buf: &[u8],
- _table_ref: &TableReference,
- schema: SchemaRef,
- _ctx: &TaskContext,
- ) -> Result<Arc<dyn TableProvider>> {
- if buf != self.token {
- return exec_err!("Unknown table provider token for TokenCodec");
- }
- self.decode_hits.fetch_add(1, Ordering::SeqCst);
- Ok(Arc::new(MemTable::try_new(schema, vec![vec![]])?))
- }
-
- fn try_encode_table_provider(
- &self,
- _table_ref: &TableReference,
- _node: Arc<dyn TableProvider>,
- buf: &mut Vec<u8>,
- ) -> Result<()> {
- self.encode_hits.fetch_add(1, Ordering::SeqCst);
- if !self.encode_by_name {
- buf.extend_from_slice(self.token);
- }
- Ok(())
- }
- }
-
- fn mem_table() -> Arc<dyn TableProvider> {
- Arc::new(MemTable::try_new(Arc::new(Schema::empty()), vec![vec![]]).unwrap())
- }
-
- fn table_ref() -> TableReference {
- TableReference::bare("t")
- }
-
- #[test]
- fn decode_falls_through_to_earlier_installed_codec() {
- let first = TokenCodec::new(b"AAAA");
- let second = TokenCodec::new(b"BBBB");
- let codec = PythonLogicalCodec::default()
- .with_additional_codec(first.clone())
- .with_additional_codec(second.clone());
-
- let ctx = TaskContext::default();
- codec
- .try_decode_table_provider(b"AAAA", &table_ref(), Arc::new(Schema::empty()), &ctx)
- .unwrap();
-
- assert_eq!(first.decode_hits.load(Ordering::SeqCst), 1);
- assert_eq!(second.decode_hits.load(Ordering::SeqCst), 0);
- }
-
- #[test]
- fn most_recently_installed_codec_encodes_first() {
- let first = TokenCodec::new(b"AAAA");
- let second = TokenCodec::new(b"BBBB");
- let codec = PythonLogicalCodec::default()
- .with_additional_codec(first.clone())
- .with_additional_codec(second.clone());
-
- let mut buf = Vec::new();
- codec
- .try_encode_table_provider(&table_ref(), mem_table(), &mut buf)
- .unwrap();
-
- assert_eq!(buf, b"BBBB");
- assert_eq!(first.encode_hits.load(Ordering::SeqCst), 0);
- }
-
- #[test]
- fn empty_ok_encode_lets_later_codec_write_payload() {
- let writer = TokenCodec::new(b"AAAA");
- let by_name = TokenCodec::new_by_name(b"BBBB");
- let codec = PythonLogicalCodec::default()
- .with_additional_codec(writer.clone())
- .with_additional_codec(by_name.clone());
-
- let mut buf = Vec::new();
- codec
- .try_encode_table_provider(&table_ref(), mem_table(), &mut buf)
- .unwrap();
-
- assert_eq!(buf, b"AAAA");
- assert_eq!(by_name.encode_hits.load(Ordering::SeqCst), 1);
- assert_eq!(writer.encode_hits.load(Ordering::SeqCst), 1);
- }
-
- #[test]
- fn decode_failure_aggregates_every_codec_error() {
- let codec = PythonLogicalCodec::default()
- .with_additional_codec(TokenCodec::new(b"AAAA"))
- .with_additional_codec(TokenCodec::new(b"BBBB"));
-
- let ctx = TaskContext::default();
- let err = codec
- .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx)
- .unwrap_err();
-
- let msg = err.to_string();
- assert!(msg.contains("None of the 3 composed extension codecs"));
- assert!(msg.contains("Unknown table provider token"));
- }
-
- #[test]
- fn single_codec_chain_error_is_returned_verbatim() {
- let codec = PythonLogicalCodec::default();
- let ctx = TaskContext::default();
- let err = codec
- .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx)
- .unwrap_err();
- assert!(!err.to_string().contains("composed extension codecs"));
- }
-
- #[test]
- fn with_additional_codec_preserves_udf_inlining_setting() {
- let strict = PythonLogicalCodec::default().with_python_udf_inlining(false);
- let extended = strict.with_additional_codec(TokenCodec::new(b"AAAA"));
- assert!(!extended.python_udf_inlining());
- }
-}
diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md
index 4934b1a9..cef39815 100644
--- a/examples/datafusion-ffi-example/README.md
+++ b/examples/datafusion-ffi-example/README.md
@@ -37,6 +37,8 @@ Both codec getters take the `SessionContext` they are being installed on and pul
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). This exam [...]
+`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 chain ordering and fall-through observable from Python. Real plugin libraries should hard-code a prefix unique to the library rather than accept one from the caller.
+
Register both provider codecs before installing the planner:
```python
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 0fd1d743..90864a19 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,10 +17,38 @@
from __future__ import annotations
-from datafusion import LogicalPlan, SessionContext
+import pyarrow as pa
+import pytest
+from datafusion import Expr, LogicalPlan, SessionContext, col, udf
from datafusion_ffi_example import MyLogicalExtensionCodec, MyTableProvider
+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.
+ """
+ codec = MyLogicalExtensionCodec(provider_prefix=token)
+ ctx = SessionContext().with_logical_extension_codec(codec)
+ 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]:
"""Build a session with the user-supplied logical extension codec
installed. Tests use a FROM-less query so plan serialization does
@@ -105,3 +133,109 @@ def test_ffi_logical_codec_composes_with_later_install():
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_most_recently_installed_codec_encodes_first():
+ """Encoding walks the chain front to back, and the front is the most
+ recently installed codec. Both codecs here can encode the provider,
+ so the winner is decided purely by install order.
+
+ Both orders are exercised 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.
+ """
+ for winner, loser in (("TOKENAAA", "TOKENBBB"), ("TOKENBBB", "TOKENAAA")):
+ loser_codec = MyLogicalExtensionCodec(provider_prefix=loser)
+ winner_codec = MyLogicalExtensionCodec(provider_prefix=winner)
+ ctx = SessionContext().with_logical_extension_codec(loser_codec)
+ ctx = ctx.with_logical_extension_codec(winner_codec)
+
+ 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_decode_falls_through_to_earlier_installed_codec():
+ """A codec that does not own the payload signals "not mine" by
+ erroring, and the chain keeps walking. The bytes here are stamped
+ with the first codec's token, so the more recently installed second
+ codec must decline and let the first one decode."""
+ blob, first = _encode_provider_plan("TOKENAAA")
+
+ second = MyLogicalExtensionCodec(provider_prefix="TOKENBBB")
+ ctx = SessionContext().with_logical_extension_codec(first)
+ ctx = ctx.with_logical_extension_codec(second)
+
+ 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_failure_aggregates_every_codec_error():
+ """When no codec in the chain claims the payload, the error names
+ the number of codecs tried and carries each one's message, so an
+ operator can see which library was expected to own the bytes."""
+ blob, _owner = _encode_provider_plan("TOKENBBB")
+
+ # Neither installed codec owns TOKENBBB, so the chain is exhausted:
+ # two example codecs plus DataFusion's default codec.
+ ctx = SessionContext().with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix="TOKENCCC")
+ )
+ ctx = ctx.with_logical_extension_codec(
+ MyLogicalExtensionCodec(provider_prefix="TOKENDDD")
+ )
+
+ with pytest.raises(Exception, match="None of the 3 composed extension codecs"):
+ LogicalPlan.from_bytes(ctx, blob)
+
+
+def test_single_codec_chain_error_is_returned_verbatim():
+ """A session with no extra codec has a one-entry chain, so a decode
+ failure surfaces DataFusion's own error rather than the aggregated
+ wrapper. Keeps error messages unchanged for the common case where
+ nobody composed anything."""
+ blob, _owner = _encode_provider_plan("TOKENEEE")
+
+ # DataFusion's own wording for "no codec claimed this", surfaced
+ # unwrapped because the chain has a single entry.
+ with pytest.raises(
+ Exception, match="LogicalExtensionCodec is not provided"
+ ) as excinfo:
+ LogicalPlan.from_bytes(SessionContext(), blob)
+
+ assert "composed extension codecs" not in str(excinfo.value)
+
+
+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/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs
index 1fcaaef4..5660489d 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 @@ struct CountingLogicalExtensionCodec {
/// 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 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec {
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 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec {
.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 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec {
pub(crate) struct MyLogicalExtensionCodec {
counters: Arc<CallCounters>,
required_udf: Option<String>,
+ token: Arc<[u8]>,
}
#[pymethods]
@@ -195,12 +199,22 @@ impl MyLogicalExtensionCodec {
/// 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 @@ impl MyLogicalExtensionCodec {
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/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py
index 451f5d21..dc55a076 100644
--- a/python/tests/test_pickle_expr.py
+++ b/python/tests/test_pickle_expr.py
@@ -323,6 +323,49 @@ class TestErrorPaths:
):
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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org