From d28615db6b92b0600f0eb47fa694499c02127d2f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 10:34:45 -0400 Subject: [PATCH 01/13] 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 --- 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 +- .../tests/_test_logical_extension_codec.py | 27 +- .../tests/_test_physical_extension_codec.py | 23 + .../README.md | 4 +- .../_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 94942a2d2..0f7e20c95 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, 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(chain: &[Arc], what: &str, f: impl Fn(&C) -> Result) -> Result { + let mut errors: Vec = 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 { + 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::>() + .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( + chain: &[Arc], + buf: &mut Vec, + what: &str, + f: impl Fn(&C, &mut Vec) -> Result<()>, +) -> Result<()> { + let mut saw_empty_ok = false; + let mut errors: Vec = 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`; see /// `PySessionContext::set_session_query_planner`. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct PythonLogicalCodec { - inner: Arc, + chain: Vec>, python_udf_inlining: bool, } impl PythonLogicalCodec { pub fn new(inner: Arc) -> Self { Self { - inner, + chain: vec![inner], python_udf_inlining: true, } } - pub fn inner(&self) -> &Arc { - &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) -> 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 { - 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) -> 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> { - 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, buf: &mut Vec, ) -> 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> { - 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, node: Arc, ) -> 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) -> 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> { @@ -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) -> 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> { @@ -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) -> 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> { @@ -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, + chain: Vec>, python_udf_inlining: bool, } impl PythonPhysicalCodec { pub fn new(inner: Arc) -> Self { Self { - inner, + chain: vec![inner], python_udf_inlining: true, } } - pub fn inner(&self) -> &Arc { - &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) -> 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> { - 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, 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) -> 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> { @@ -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, 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], ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - 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) -> 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> { @@ -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) -> 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> { @@ -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) -> Result { let Some(py_udf) = node.inner().downcast_ref::() else { return Ok(false); @@ -613,7 +767,7 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec) /// 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>> { 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 { + 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 { + 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 { + exec_err!("TokenCodec does not decode extension nodes") + } + + fn try_encode(&self, _node: &Extension, _buf: &mut Vec) -> 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> { + 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, + buf: &mut Vec, + ) -> 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 { + 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 75bfed601..4f44ee841 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1465,14 +1465,13 @@ impl PySessionContext { let inner: Arc = (&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 = (&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 31cd9391f..67a1e7385 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 0fa10d7f3..4934b1a91 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 example keeps the provider library as the sole external codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host. + +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), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on. -For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. +For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. diff --git a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py index cd0c5a61a..0fd1d7431 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 28eaaf2d9..82116bef7 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 72f96bb8e..f188ec740 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 — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). +The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call 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 nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). -For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. +For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index d046f67a6..c6f9189d4 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 2f2cc6119..6a4fbdb18 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1793,7 +1793,8 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non 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 @@ def __datafusion_query_planner__(self, session: Any = None) -> Any: 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 @@ def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: 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. From 62c3d38ccbed22806f8633901a7f4a459cc8189a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 8 Aug 2026 11:56:18 -0400 Subject: [PATCH 02/13] test: port codec Rust tests to pytest MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- crates/core/src/codec.rs | 312 ------------------ examples/datafusion-ffi-example/README.md | 2 + .../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 0f7e20c95..0e990dec6 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 Arc { - 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 { - 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 { - exec_err!("TokenCodec does not decode extension nodes") - } - - fn try_encode(&self, _node: &Extension, _buf: &mut Vec) -> 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> { - 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, - buf: &mut Vec, - ) -> 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 { - 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 4934b1a91..cef39815d 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 example keeps the provider library as the sole external codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host. +`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 0fd1d7431..90864a19a 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 1fcaaef4c..5660489d4 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, + /// 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> { 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, required_udf: Option, + 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) -> Self { + #[pyo3(signature = (require_udf_on_decode=None, provider_prefix=None))] + fn new(require_udf_on_decode: Option, 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 451f5d215..dc55a0767 100644 --- a/python/tests/test_pickle_expr.py +++ b/python/tests/test_pickle_expr.py @@ -323,6 +323,49 @@ def test_cross_version_error_message(self): ): 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 From c1ffd81ca02ec969e84a349e7fef1d8fcf9d854f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 8 Aug 2026 11:58:07 -0400 Subject: [PATCH 03/13] 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) --- AGENTS.md | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index 327ebd643..659094ec0 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. From 2a458cc49548dd4410b9eae9e052ee60659edef0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 31 Aug 2026 16:23:06 -0400 Subject: [PATCH 04/13] refactor: dispatch chained codecs by identity, not by trial Decoding walked the codec chain and took the first codec that returned Ok. That is unsound, and upstream has already been bitten by it: apache/datafusion#16980 records ComposedPhysicalExtensionCodec decoding an encoded parquet file as csv because the payload happened to parse with an earlier codec, and #16986 fixed it by recording the encoder. Protobuf carries no type identity, so two codecs whose leading field numbers and wire types line up decode each other's payloads cleanly. A byte-prefix convention does not help: the natural implementation is `Message::decode(buf)`, which has no prefix to check and cannot decline. Payloads written by a chained codec now carry an envelope naming the codec that wrote them, and decoding consults exactly that codec. The envelope is applied and stripped inside PythonLogicalCodec / PythonPhysicalCodec, so third-party codecs are unmodified and never see it -- they receive the bare bytes they wrote. Keyed on a stable identity rather than chain position. Position is sound for Ballista and datafusion-distributed because their codec lists are pinned -- a compile-time constant in one, a fixed entry plus appended user codecs rebuilt from the same startup code in the other. A datafusion-python chain is assembled by user Python across sessions that share no struct, and a library shipping a codec cannot know its own index, so position would silently name the wrong codec whenever two sessions install in different orders. Identity is derived and asks nothing of existing libraries: an explicit codec_id, else __datafusion_codec_id__, else the exporting class's module and qualified name. A bare PyCapsule exposes nothing stable -- every capsule reports the same type -- so it gets a session-local id and a pointed error if its payloads reach an unrelated session. Installing two codecs under one id is rejected at install time rather than resolving to whichever entry came first. Codecs now append rather than prepend, so encoding is claimed by the first codec installed that wants the object. Installing a library can only claim objects nothing else claimed; it can never take over an existing library's objects, and it never renumbers ids that older payloads reference. Two payloads are deliberately left unframed. The terminal codec writes bare, so a session with no extension codecs is byte-identical to a build without chaining. And an encode that writes nothing stays empty: an empty fun_definition is DataFusion's encode-by-name signal, and framing it would set the field, permanently skipping the registry lookup the decoder does first and breaking codecs that reconstruct a function from its name alone (from_proto.rs, the `None => ctx.udf(..).or_else(..)` arm). That arm is also why an empty buffer still consults every codec: there are no bytes to tag. It is not the hazard tagging removes -- the question asked is "do you own the function named x", which is name-scoped, and a disagreement needs two libraries claiming one name. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/codec.rs | 594 ++++++++++++++---- crates/core/src/context.rs | 95 ++- .../tests/_test_logical_extension_codec.py | 179 ++++-- .../_test_three_library_query_planner.py | 55 +- python/datafusion/context.py | 70 ++- 5 files changed, 793 insertions(+), 200 deletions(-) diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 0e990dec6..c58b81aff 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -224,29 +224,215 @@ fn strip_wire_header<'a>( Ok(Some(&buf[py_minor_idx + 1..])) } -/// Run `f` against each codec in `chain`, returning the first `Ok`. +/// Family prefix for the envelope wrapping a chained codec's payload. /// -/// 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(chain: &[Arc], what: &str, f: impl Fn(&C) -> Result) -> Result { +/// A distinct magic is what makes "is this framed?" a definite test +/// rather than a speculative decode. Probing by attempting to parse the +/// envelope would reintroduce exactly the protobuf ambiguity this +/// framing exists to remove: prost skips unknown fields and defaults +/// missing ones, so a foreign payload can parse cleanly as an envelope. +pub(crate) const CHAINED_PAYLOAD_FAMILY: &[u8] = b"DFPYCHN"; + +/// Wire-format version for the chained-payload envelope. Independent of +/// [`WIRE_VERSION_CURRENT`], which versions the cloudpickle framing. +pub(crate) const CHAIN_WIRE_VERSION_CURRENT: u8 = 1; + +/// Oldest chained-payload envelope version this build decodes. +pub(crate) const CHAIN_WIRE_VERSION_MIN_SUPPORTED: u8 = 1; + +/// Prefix for the synthetic id given to a codec installed from a bare +/// PyCapsule, which exposes nothing stable to derive an identity from. +/// Such an id is unique to the installing session, so a payload +/// carrying one decodes within that session's lineage and fails with a +/// pointed error anywhere else rather than resolving to a different +/// codec that happens to sit at the same position. +pub(crate) const ANONYMOUS_CODEC_ID_PREFIX: &str = "anon:"; + +/// One installed codec plus the identity its payloads are tagged with. +/// +/// The id is what makes dispatch order-independent. Keying on position +/// in the chain — as `ComposedPhysicalExtensionCodec` does upstream — +/// is sound only when both ends assemble the same list in the same +/// order, which holds for a compile-time codec list but not for a +/// chain built by user Python across two independently configured +/// sessions. +struct ChainEntry { + id: Arc, + codec: Arc, +} + +impl Clone for ChainEntry { + fn clone(&self) -> Self { + Self { + id: Arc::clone(&self.id), + codec: Arc::clone(&self.codec), + } + } +} + +impl std::fmt::Debug for ChainEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChainEntry") + .field("id", &self.id) + .field("codec", &self.codec) + .finish() + } +} + +/// Wrap `blob` in the envelope identifying the codec that produced it. +/// +/// Layout: `DFPYCHN | version: u8 | id_len: u32 (LE) | id | blob`. +fn write_chained_payload(buf: &mut Vec, codec_id: &str, blob: &[u8]) { + buf.extend_from_slice(CHAINED_PAYLOAD_FAMILY); + buf.push(CHAIN_WIRE_VERSION_CURRENT); + buf.extend_from_slice(&(codec_id.len() as u32).to_le_bytes()); + buf.extend_from_slice(codec_id.as_bytes()); + buf.extend_from_slice(blob); +} + +/// Inspect the chained-payload envelope on `buf`. +/// +/// * `Ok(None)` — no envelope. The payload came from the terminal +/// codec, which writes unframed so a session with no extension +/// codecs installed produces bytes identical to a build without +/// codec chaining. +/// * `Ok(Some((codec_id, blob)))` — the owning codec's id and its +/// original bytes, byte-for-byte as it wrote them. +fn read_chained_payload(buf: &[u8]) -> Result> { + if !buf.starts_with(CHAINED_PAYLOAD_FAMILY) { + return Ok(None); + } + let mut idx = CHAINED_PAYLOAD_FAMILY.len(); + let Some(&version) = buf.get(idx) else { + return Err(datafusion::error::DataFusionError::Execution( + "Truncated extension codec payload: missing envelope version byte".to_string(), + )); + }; + if !(CHAIN_WIRE_VERSION_MIN_SUPPORTED..=CHAIN_WIRE_VERSION_CURRENT).contains(&version) { + return Err(datafusion::error::DataFusionError::Execution(format!( + "Extension codec payload envelope version v{version}; this build supports \ + v{CHAIN_WIRE_VERSION_MIN_SUPPORTED}..=v{CHAIN_WIRE_VERSION_CURRENT}. \ + Align datafusion-python versions on sender and receiver." + ))); + } + idx += 1; + let Some(len_bytes) = buf.get(idx..idx + 4) else { + return Err(datafusion::error::DataFusionError::Execution( + "Truncated extension codec payload: missing codec id length".to_string(), + )); + }; + let id_len = u32::from_le_bytes(len_bytes.try_into().expect("4 bytes")) as usize; + idx += 4; + let Some(id_bytes) = buf.get(idx..idx + id_len) else { + return Err(datafusion::error::DataFusionError::Execution( + "Truncated extension codec payload: codec id shorter than its declared length" + .to_string(), + )); + }; + let codec_id = std::str::from_utf8(id_bytes).map_err(|err| { + datafusion::error::DataFusionError::Execution(format!( + "Extension codec payload carries a non-UTF-8 codec id: {err}" + )) + })?; + Ok(Some((codec_id, &buf[idx + id_len..]))) +} + +/// Decode `buf` with the single codec that encoded it. +/// +/// Three cases: +/// +/// * **Empty `buf`** — nothing was encoded, so there is no tag to +/// dispatch on and every codec is offered the empty buffer in install +/// order. See [`chain_resolve_by_name`] for why that is sound here and +/// why the case exists at all. +/// * **Framed `buf`** — the envelope names its author, so exactly one +/// codec is consulted and its error surfaces verbatim. +/// * **Unframed non-empty `buf`** — the terminal codec wrote it. +/// +/// Outside the empty case nothing is ever offered to a codec that did +/// not write it, which is what stops a structurally similar prost +/// message from decoding in the wrong library. +fn chain_decode( + chain: &[ChainEntry], + terminal: &Arc, + buf: &[u8], + what: &str, + f: impl Fn(&C, &[u8]) -> Result, +) -> Result { + if buf.is_empty() { + return chain_resolve_by_name(chain, terminal, what, |codec| f(codec, buf)); + } + let Some((codec_id, blob)) = read_chained_payload(buf)? else { + return f(terminal.as_ref(), buf); + }; + let Some(entry) = chain.iter().find(|entry| &*entry.id == codec_id) else { + let installed = if chain.is_empty() { + "no extension codecs are installed on this session".to_string() + } else { + format!( + "installed: {}", + chain + .iter() + .map(|entry| entry.id.as_ref()) + .collect::>() + .join(", ") + ) + }; + let hint = if codec_id.starts_with(ANONYMOUS_CODEC_ID_PREFIX) { + ". This payload was written by a codec installed from a bare PyCapsule, which \ + carries no portable identity. Pass `codec_id=` when installing it if plans must \ + cross sessions." + } else { + "" + }; + return Err(datafusion::error::DataFusionError::Execution(format!( + "{what} was encoded by extension codec '{codec_id}', which is not installed on \ + this session ({installed}){hint}" + ))); + }; + f(entry.codec.as_ref(), blob) +} + +/// Resolve an object carrying no payload, by consulting each codec. +/// +/// Used only where DataFusion encodes by name: `try_encode_udf` and its +/// aggregate/window siblings return `Ok` writing nothing, and the +/// decoder then tries the `FunctionRegistry` first and the codec second +/// (`from_proto.rs`, the `None => ctx.udf(..).or_else(..)` arm). A codec +/// whose functions are reconstructible from the name alone is reached +/// through that arm and must still be offered the empty buffer. +/// +/// This is the one place dispatch cannot be tagged — there are no bytes +/// to tag. It is not the hazard that tagging exists to remove: the +/// question asked here is "do you own the function named `x`", which is +/// name-scoped and answerable, not "do these bytes happen to parse as +/// your message type". Two codecs disagreeing requires them to claim +/// the same function name, which already collides in the registry. +fn chain_resolve_by_name( + chain: &[ChainEntry], + terminal: &Arc, + what: &str, + f: impl Fn(&C) -> Result, +) -> Result { let mut errors: Vec = Vec::new(); - for codec in chain { - match f(codec) { + for entry in chain { + match f(entry.codec.as_ref()) { Ok(value) => return Ok(value), Err(err) => errors.push(err), } } - Err(aggregate_chain_errors(what, errors)) + match f(terminal.as_ref()) { + Ok(value) => Ok(value), + Err(err) => { + errors.push(err); + Err(aggregate_chain_errors(what, errors)) + } + } } /// Collapse per-codec failures into one error. A single failure is -/// returned as-is so the one-codec (default-only) chain behaves -/// exactly like the pre-chain implementation. +/// returned as-is so a session with no extension codecs behaves exactly +/// like a build without codec chaining. fn aggregate_chain_errors( what: &str, mut errors: Vec, @@ -263,43 +449,58 @@ fn aggregate_chain_errors( .collect::>() .join("; "); datafusion::error::DataFusionError::Execution(format!( - "None of the {} composed extension codecs handled {what}: {joined}", - errors.len() + "No installed extension codec handled {what}: {joined}" )) } } } -/// Encode variant of [`chain_try`] for methods that write into a -/// caller-provided buffer. +/// Encode through the chain, tagging the payload with its author. +/// +/// Entries are consulted in install order and the first one to write +/// bytes wins, so installing a codec can only claim objects no +/// earlier codec claimed. Adding a library therefore never changes how +/// an already-installed library's objects encode. /// /// Each codec encodes into a scratch buffer so a failed attempt cannot -/// leave partial bytes behind. `Ok` with 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). +/// leave partial bytes behind. `Ok` with an empty buffer is "no +/// opinion" rather than a claim, so the walk continues; if nothing +/// writes bytes the result is `Ok` with nothing written, which is +/// DataFusion's encode-by-name signal. Framing that empty result would +/// set `fun_definition` and permanently skip the registry lookup the +/// decoder does first. +/// +/// The terminal codec writes unframed, so a session with no extension +/// codecs is byte-compatible with a build predating the chain. fn chain_encode( - chain: &[Arc], + chain: &[ChainEntry], + terminal: &Arc, buf: &mut Vec, what: &str, f: impl Fn(&C, &mut Vec) -> Result<()>, ) -> Result<()> { let mut saw_empty_ok = false; let mut errors: Vec = Vec::new(); - for codec in chain { + for entry in chain { let mut scratch = Vec::new(); - match f(codec, &mut scratch) { + match f(entry.codec.as_ref(), &mut scratch) { Ok(()) if !scratch.is_empty() => { - buf.extend_from_slice(&scratch); + write_chained_payload(buf, &entry.id, &scratch); return Ok(()); } Ok(()) => saw_empty_ok = true, Err(err) => errors.push(err), } } + let mut scratch = Vec::new(); + match f(terminal.as_ref(), &mut scratch) { + Ok(()) if !scratch.is_empty() => { + buf.extend_from_slice(&scratch); + return Ok(()); + } + Ok(()) => saw_empty_ok = true, + Err(err) => errors.push(err), + } if saw_empty_ok { return Ok(()); } @@ -309,15 +510,19 @@ fn chain_encode( /// `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 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. +/// handle to a chain of composable codecs. Each downstream FFI codec +/// installed via `SessionContext.with_logical_extension_codec(...)` is +/// appended to the chain, and `terminal` — normally +/// `DefaultLogicalExtensionCodec` — handles whatever no installed codec +/// claims. /// -/// 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. +/// Every payload an installed codec writes is wrapped in an envelope +/// naming that codec (see [`write_chained_payload`]), so decoding +/// consults exactly the codec that encoded it. Dispatch does not depend +/// on a codec recognizing and rejecting foreign payloads, which is not +/// something a codec can reliably do: a prost message decodes cleanly +/// from another message's bytes whenever their leading field numbers and +/// wire types line up. /// /// Sitting at the top of the session's logical codec stack means /// every serializer that reads `session.logical_codec()` automatically @@ -332,31 +537,71 @@ fn chain_encode( /// `PySessionContext::set_session_query_planner`. #[derive(Debug, Clone)] pub struct PythonLogicalCodec { - chain: Vec>, + chain: Vec>, + terminal: Arc, python_udf_inlining: bool, } impl PythonLogicalCodec { + /// Build a codec with no installed extension codecs and `inner` as + /// the terminal fallback. `inner` is not part of the keyed chain and + /// its payloads are written unframed, so a context built this way + /// serializes byte-identically to one with no chaining at all. pub fn new(inner: Arc) -> Self { Self { - chain: vec![inner], + chain: Vec::new(), + terminal: inner, python_udf_inlining: true, } } - /// 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) -> Self { - let mut chain = Vec::with_capacity(self.chain.len() + 1); - chain.push(codec); - chain.extend(self.chain.iter().map(Arc::clone)); + /// Return a copy of this codec with `codec` appended to the chain + /// under `id`, preserving the Python-UDF-inlining setting. + /// + /// Appending rather than prepending keeps the operation additive: + /// the new codec is consulted for encoding only after every codec + /// already installed, so it can claim objects nothing else claimed + /// but cannot take over an existing library's objects. + pub fn with_additional_codec( + &self, + id: impl Into>, + codec: Arc, + ) -> Self { + let mut chain = self.chain.clone(); + chain.push(ChainEntry { + id: id.into(), + codec, + }); Self { chain, + terminal: Arc::clone(&self.terminal), python_udf_inlining: self.python_udf_inlining, } } + /// Ids of the installed extension codecs, in install order. + /// + /// The terminal codec is not listed: it is not addressable by id + /// because its payloads are written unframed. + pub fn codec_ids(&self) -> Vec<&str> { + self.chain.iter().map(|entry| entry.id.as_ref()).collect() + } + + /// Installed extension codecs paired with their ids, in install + /// order. Restores the inspection that the removed `inner()` + /// accessor provided, and exposes the id dispatch keys along with it. + pub fn codecs(&self) -> Vec<(&str, &Arc)> { + self.chain + .iter() + .map(|entry| (entry.id.as_ref(), &entry.codec)) + .collect() + } + + /// Terminal codec consulted when no installed codec claims an object. + pub fn terminal(&self) -> &Arc { + &self.terminal + } + /// Toggle inline encoding of Python UDFs. See /// `SessionContext.with_python_udf_inlining` (Python) for full /// behavior and use cases. @@ -395,14 +640,19 @@ impl LogicalExtensionCodec for PythonLogicalCodec { inputs: &[LogicalPlan], ctx: &TaskContext, ) -> Result { - chain_try(&self.chain, "an extension logical plan node", |codec| { - codec.try_decode(buf, inputs, ctx) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an extension logical plan node", + |codec, buf| codec.try_decode(buf, inputs, ctx), + ) } fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { chain_encode( &self.chain, + &self.terminal, buf, "an extension logical plan node", |codec, buf| codec.try_encode(node, buf), @@ -416,9 +666,13 @@ impl LogicalExtensionCodec for PythonLogicalCodec { schema: SchemaRef, ctx: &TaskContext, ) -> Result> { - chain_try(&self.chain, "a table provider", |codec| { - codec.try_decode_table_provider(buf, table_ref, Arc::clone(&schema), ctx) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a table provider", + |codec, buf| codec.try_decode_table_provider(buf, table_ref, Arc::clone(&schema), ctx), + ) } fn try_encode_table_provider( @@ -427,9 +681,13 @@ impl LogicalExtensionCodec for PythonLogicalCodec { node: Arc, buf: &mut Vec, ) -> Result<()> { - chain_encode(&self.chain, buf, "a table provider", |codec, buf| { - codec.try_encode_table_provider(table_ref, Arc::clone(&node), buf) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a table provider", + |codec, buf| codec.try_encode_table_provider(table_ref, Arc::clone(&node), buf), + ) } fn try_decode_file_format( @@ -437,9 +695,13 @@ impl LogicalExtensionCodec for PythonLogicalCodec { buf: &[u8], ctx: &TaskContext, ) -> Result> { - chain_try(&self.chain, "a file format", |codec| { - codec.try_decode_file_format(buf, ctx) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a file format", + |codec, buf| codec.try_decode_file_format(buf, ctx), + ) } fn try_encode_file_format( @@ -447,18 +709,26 @@ impl LogicalExtensionCodec for PythonLogicalCodec { buf: &mut Vec, node: Arc, ) -> Result<()> { - chain_encode(&self.chain, buf, "a file format", |codec, buf| { - codec.try_encode_file_format(buf, Arc::clone(&node)) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a file format", + |codec, buf| codec.try_encode_file_format(buf, Arc::clone(&node)), + ) } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } - chain_encode(&self.chain, buf, "a scalar UDF", |codec, buf| { - codec.try_encode_udf(node, buf) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_encode_udf(node, buf), + ) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { @@ -469,18 +739,26 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } - chain_try(&self.chain, "a scalar UDF", |codec| { - codec.try_decode_udf(name, buf) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_decode_udf(name, buf), + ) } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } - chain_encode(&self.chain, buf, "an aggregate UDF", |codec, buf| { - codec.try_encode_udaf(node, buf) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_encode_udaf(node, buf), + ) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { @@ -491,18 +769,26 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } - chain_try(&self.chain, "an aggregate UDF", |codec| { - codec.try_decode_udaf(name, buf) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_decode_udaf(name, buf), + ) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } - chain_encode(&self.chain, buf, "a window UDF", |codec, buf| { - codec.try_encode_udwf(node, buf) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_encode_udwf(node, buf), + ) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { @@ -513,9 +799,13 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } - chain_try(&self.chain, "a window UDF", |codec| { - codec.try_decode_udwf(name, buf) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_decode_udwf(name, buf), + ) } } @@ -581,31 +871,59 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio /// from; see that type for why. #[derive(Debug, Clone)] pub struct PythonPhysicalCodec { - chain: Vec>, + chain: Vec>, + terminal: Arc, python_udf_inlining: bool, } impl PythonPhysicalCodec { + /// See [`PythonLogicalCodec::new`]; `inner` is the terminal codec + /// rather than a chain entry. pub fn new(inner: Arc) -> Self { Self { - chain: vec![inner], + chain: Vec::new(), + terminal: inner, python_udf_inlining: true, } } - /// 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) -> Self { - let mut chain = Vec::with_capacity(self.chain.len() + 1); - chain.push(codec); - chain.extend(self.chain.iter().map(Arc::clone)); + /// Return a copy of this codec with `codec` appended to the chain + /// under `id`. See [`PythonLogicalCodec::with_additional_codec`]. + pub fn with_additional_codec( + &self, + id: impl Into>, + codec: Arc, + ) -> Self { + let mut chain = self.chain.clone(); + chain.push(ChainEntry { + id: id.into(), + codec, + }); Self { chain, + terminal: Arc::clone(&self.terminal), python_udf_inlining: self.python_udf_inlining, } } + /// Ids of the installed extension codecs, in install order. + pub fn codec_ids(&self) -> Vec<&str> { + self.chain.iter().map(|entry| entry.id.as_ref()).collect() + } + + /// Installed extension codecs paired with their ids, in install order. + pub fn codecs(&self) -> Vec<(&str, &Arc)> { + self.chain + .iter() + .map(|entry| (entry.id.as_ref(), &entry.codec)) + .collect() + } + + /// Terminal codec consulted when no installed codec claims an object. + pub fn terminal(&self) -> &Arc { + &self.terminal + } + /// Toggle inline encoding of Python UDFs on this physical codec. /// /// Mirrors [`PythonLogicalCodec::with_python_udf_inlining`]; see @@ -634,9 +952,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - chain_try(&self.chain, "an execution plan", |codec| { - codec.try_decode(buf, inputs, ctx, proto_converter) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an execution plan", + |codec, buf| codec.try_decode(buf, inputs, ctx, proto_converter), + ) } fn try_encode( @@ -645,18 +967,26 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - chain_encode(&self.chain, buf, "an execution plan", |codec, buf| { - codec.try_encode(Arc::clone(&node), buf, proto_converter) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "an execution plan", + |codec, buf| codec.try_encode(Arc::clone(&node), buf, proto_converter), + ) } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } - chain_encode(&self.chain, buf, "a scalar UDF", |codec, buf| { - codec.try_encode_udf(node, buf) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_encode_udf(node, buf), + ) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { @@ -667,9 +997,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } - chain_try(&self.chain, "a scalar UDF", |codec| { - codec.try_decode_udf(name, buf) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_decode_udf(name, buf), + ) } fn try_encode_expr( @@ -678,9 +1012,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { buf: &mut Vec, ctx: &PhysicalExprEncodeCtx<'_>, ) -> Result<()> { - chain_encode(&self.chain, buf, "a physical expression", |codec, buf| { - codec.try_encode_expr(node, buf, ctx) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a physical expression", + |codec, buf| codec.try_encode_expr(node, buf, ctx), + ) } fn try_decode_expr( @@ -689,18 +1027,26 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { inputs: &[Arc], ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - chain_try(&self.chain, "a physical expression", |codec| { - codec.try_decode_expr(buf, inputs, ctx) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a physical expression", + |codec, buf| codec.try_decode_expr(buf, inputs, ctx), + ) } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } - chain_encode(&self.chain, buf, "an aggregate UDF", |codec, buf| { - codec.try_encode_udaf(node, buf) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_encode_udaf(node, buf), + ) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { @@ -711,18 +1057,26 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } - chain_try(&self.chain, "an aggregate UDF", |codec| { - codec.try_decode_udaf(name, buf) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_decode_udaf(name, buf), + ) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } - chain_encode(&self.chain, buf, "a window UDF", |codec, buf| { - codec.try_encode_udwf(node, buf) - }) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_encode_udwf(node, buf), + ) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { @@ -733,9 +1087,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } - chain_try(&self.chain, "a window UDF", |codec| { - codec.try_decode_udwf(name, buf) - }) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_decode_udwf(name, buf), + ) } } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 4f44ee841..8a87e36ac 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -75,7 +75,7 @@ use uuid::Uuid; use crate::catalog::{ PyCatalog, PyCatalogList, RustWrappedPyCatalogProvider, RustWrappedPyCatalogProviderList, }; -use crate::codec::{PythonLogicalCodec, PythonPhysicalCodec}; +use crate::codec::{ANONYMOUS_CODEC_ID_PREFIX, PythonLogicalCodec, PythonPhysicalCodec}; use crate::common::data_type::PyScalarValue; use crate::common::df_schema::PyDFSchema; use crate::dataframe::PyDataFrame; @@ -1457,21 +1457,28 @@ impl PySessionContext { create_query_planner_capsule(py, &ffi) } + #[pyo3(signature = (codec, codec_id=None))] pub fn with_logical_extension_codec<'py>( slf: &Bound<'py, Self>, codec: Bound<'py, PyAny>, + codec_id: Option, ) -> PyDataFusionResult { + let id = { + let this = slf.borrow(); + resolve_codec_id(&codec, codec_id, &this.logical_codec.codec_ids())? + }; let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); let this = slf.borrow(); - // 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 + // Append rather than replace: previously installed codecs stay active, + // and every payload this one writes is tagged with `id` so decoding + // reaches it directly rather than by trying codecs in turn. Appending + // also carries the receiver's inlining setting over, which a fresh // `PythonLogicalCodec::new` would not — it defaults inlining to on, and // would silently re-enable inline Python UDF encoding on a context that // had opted out with `with_python_udf_inlining(enabled=False)`. - let logical_codec = Arc::new(this.logical_codec.with_additional_codec(inner)); + let logical_codec = Arc::new(this.logical_codec.with_additional_codec(id, inner)); let derived = Self { ctx: Arc::clone(&this.ctx), logical_codec, @@ -1494,17 +1501,23 @@ impl PySessionContext { create_physical_extension_capsule(py, self.ffi_physical_codec().as_ref()) } + #[pyo3(signature = (codec, codec_id=None))] pub fn with_physical_extension_codec<'py>( slf: &Bound<'py, Self>, codec: Bound<'py, PyAny>, + codec_id: Option, ) -> PyDataFusionResult { + let id = { + let this = slf.borrow(); + resolve_codec_id(&codec, codec_id, &this.physical_codec.codec_ids())? + }; let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); let this = slf.borrow(); - // See `with_logical_extension_codec` for why this prepends rather than + // See `with_logical_extension_codec` for why this appends rather than // replaces, and why that is also what carries the inlining flag over. - let physical_codec = Arc::new(this.physical_codec.with_additional_codec(inner)); + let physical_codec = Arc::new(this.physical_codec.with_additional_codec(id, inner)); let derived = Self { ctx: Arc::clone(&this.ctx), logical_codec: Arc::clone(&this.logical_codec), @@ -1706,6 +1719,74 @@ impl PySessionContext { } } +/// Determine the wire identity to tag an installed codec's payloads with. +/// +/// Every payload a chained codec writes carries this string, and decoding +/// dispatches on it, so it has to name the same codec in the process that +/// decodes as it did in the process that encoded. Resolution order: +/// +/// 1. An explicit `codec_id` argument. +/// 2. `codec.__datafusion_codec_id__`, letting a library pin its own identity +/// so a class rename does not invalidate previously encoded plans, and so +/// two instances of one class can own disjoint slices of the wire format. +/// 3. The exporting object's `module.QualName`, which is the library's own +/// import path and therefore already stable across processes. This is the +/// common case and asks nothing of existing extension libraries. +/// 4. For a bare `PyCapsule` there is nothing stable to read — every capsule +/// reports the same type — so fall back to a session-local id. Payloads +/// tagged this way decode correctly within the session lineage that +/// installed the codec and fail with a pointed error elsewhere, rather +/// than resolving to whichever codec happens to sit at the same position. +/// +/// An id already in use is rejected rather than shadowed. Two codecs sharing an +/// id are indistinguishable on decode, and the API cannot tell whether two +/// instances of one class write the same wire format — so the ambiguity is +/// surfaced at install time, where the caller can resolve it, instead of at +/// decode time, where it would pick whichever entry came first. +fn resolve_codec_id( + codec: &Bound<'_, PyAny>, + explicit: Option, + existing: &[&str], +) -> PyResult { + let id = derive_codec_id(codec, explicit, existing.len())?; + if existing.contains(&id.as_str()) { + return Err(PyValueError::new_err(format!( + "An extension codec with id '{id}' is already installed on this session. Two \ + codecs cannot share an id, because a payload names its codec by id when it is \ + decoded. Pass `codec_id=` to give this one a distinct identity." + ))); + } + Ok(id) +} + +fn derive_codec_id( + codec: &Bound<'_, PyAny>, + explicit: Option, + installed: usize, +) -> PyResult { + if let Some(id) = explicit { + return Ok(id); + } + if let Ok(declared) = codec.getattr("__datafusion_codec_id__") + && !declared.is_none() + { + return declared.extract::(); + } + if codec.is_instance_of::() { + return Ok(format!("{ANONYMOUS_CODEC_ID_PREFIX}{installed}")); + } + let ty = codec.get_type(); + let module = ty + .getattr("__module__") + .and_then(|m| m.extract::()) + .unwrap_or_else(|_| "".to_string()); + let qualname = ty + .getattr("__qualname__") + .and_then(|q| q.extract::()) + .or_else(|_| ty.name().and_then(|n| n.extract::()))?; + Ok(format!("{module}.{qualname}")) +} + pub fn parse_file_compression_type( file_compression_type: Option, ) -> Result { 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 90864a19a..16588e6f3 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 @@ -40,9 +40,14 @@ def _encode_provider_plan(token: str) -> tuple[bytes, MyLogicalExtensionCodec]: Returns the blob and the codec, so callers can assert on its call counters. The token is chosen per test so a second codec installed later is provably unable to claim these bytes. + + The codec is installed under ``token`` as its id as well, so a + caller can reinstall the same instance elsewhere and have the tag on + these bytes resolve. Identity would otherwise be derived from the + class, which every instance shares. """ codec = MyLogicalExtensionCodec(provider_prefix=token) - ctx = SessionContext().with_logical_extension_codec(codec) + ctx = SessionContext().with_logical_extension_codec(codec, codec_id=token) ctx.register_table("numbers", MyTableProvider(1, 4, 1)) blob = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) assert token.encode() in blob @@ -135,22 +140,27 @@ def test_ffi_logical_codec_composes_with_later_install(): 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. +def test_first_installed_codec_encodes(): + """Encoding walks the chain in install order, so the earliest + installed codec that can claim an object gets it. + + Both orders run in one test on purpose. Asserting a single order + would also pass under replace semantics, where the second install + simply discards the first codec; swapping the order and getting the + other token proves the losing codec was still installed and merely + lost the race. - 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. + The two instances need explicit ids: identity is otherwise derived + from the class, and these are two instances of one class owning + disjoint slices of the wire format. """ for winner, loser in (("TOKENAAA", "TOKENBBB"), ("TOKENBBB", "TOKENAAA")): - 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) + loser_codec = MyLogicalExtensionCodec(provider_prefix=loser) + ctx = SessionContext().with_logical_extension_codec( + winner_codec, codec_id=winner + ) + ctx = ctx.with_logical_extension_codec(loser_codec, codec_id=loser) ctx.register_table("numbers", MyTableProvider(1, 4, 1)) blob = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) @@ -161,16 +171,47 @@ def test_most_recently_installed_codec_encodes_first(): 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.""" +def test_installing_a_codec_cannot_hijack_an_earlier_codecs_objects(): + """Appending is additive: a later install can claim objects nothing + else claimed, but never takes over an object an earlier codec was + already encoding. + + This is why install order is append rather than prepend. Under + prepend, adding an unrelated library would silently change how an + existing library's objects encode -- and, once payloads are tagged, + would renumber ids that older payloads already reference. + """ + first = MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + ctx = SessionContext().with_logical_extension_codec(first, codec_id="TOKENAAA") + ctx.register_table("numbers", MyTableProvider(1, 4, 1)) + before = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) + + later = MyLogicalExtensionCodec(provider_prefix="TOKENBBB") + ctx = ctx.with_logical_extension_codec(later, codec_id="TOKENBBB") + after = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) + + # Provider tokens are minted per encode, so the payloads differ in the + # token id. What must not change is which codec claimed the provider. + assert b"TOKENAAA" in after + assert b"TOKENBBB" not in after + assert later.table_provider_encode_calls() == 0 + assert len(before) == len(after) + + +def test_decode_dispatches_to_the_codec_that_encoded(): + """A payload names the codec that wrote it, so decoding consults + exactly that codec and never offers the bytes to any other. + + The blob here is written by ``first`` in its own session. Installing + ``second`` alongside it must not put ``second`` anywhere near those + bytes -- under trial-and-error dispatch it would be asked first, and + a codec that decodes structurally similar protobuf would answer. + """ blob, first = _encode_provider_plan("TOKENAAA") second = MyLogicalExtensionCodec(provider_prefix="TOKENBBB") - ctx = SessionContext().with_logical_extension_codec(first) - ctx = ctx.with_logical_extension_codec(second) + ctx = SessionContext().with_logical_extension_codec(first, codec_id="TOKENAAA") + ctx = ctx.with_logical_extension_codec(second, codec_id="TOKENBBB") restored = LogicalPlan.from_bytes(ctx, blob) assert ctx.create_dataframe_from_logical_plan(restored).collect() @@ -179,40 +220,92 @@ def test_decode_falls_through_to_earlier_installed_codec(): 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.""" +def test_decode_survives_a_different_install_order(): + """Dispatch keys off codec identity, not chain position, so the + decoding session may install the same codecs in any order. + + This is the case positional dispatch cannot handle: the encoding + session has the owning codec at index 0 and the decoding session has + it at index 1. Keying on position would hand the payload to whatever + sits at index 0 in the decoder -- silently, and with a plausible + result. + """ + owner = MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + other = MyLogicalExtensionCodec(provider_prefix="TOKENBBB") + + encoder = SessionContext().with_logical_extension_codec(owner, codec_id="TOKENAAA") + encoder = encoder.with_logical_extension_codec(other, codec_id="TOKENBBB") + encoder.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder) + + # Same codecs, opposite order. + decoder = SessionContext().with_logical_extension_codec(other, codec_id="TOKENBBB") + decoder = decoder.with_logical_extension_codec(owner, codec_id="TOKENAAA") + + restored = LogicalPlan.from_bytes(decoder, blob) + assert decoder.create_dataframe_from_logical_plan(restored).collect() + assert other.table_provider_decode_calls() == 0 + + +def test_decode_names_the_codec_that_is_not_installed(): + """When the owning codec is absent the error names it and lists what + is installed, instead of reporting DataFusion's generic "not + provided" from whichever codec was tried last.""" blob, _owner = _encode_provider_plan("TOKENBBB") - # 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") + MyLogicalExtensionCodec(provider_prefix="TOKENCCC"), codec_id="lib_c.Codec" ) - with pytest.raises(Exception, match="None of the 3 composed extension codecs"): + with pytest.raises(Exception, match="TOKENBBB") as excinfo: LogicalPlan.from_bytes(ctx, blob) + message = str(excinfo.value) + # Names the codec the payload belongs to, and what is actually here. + assert "not installed on this session" in message + assert "lib_c.Codec" in message + -def test_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") +def test_installing_two_codecs_under_one_id_is_rejected(): + """Identity is derived from the class, so installing two instances of + one class collides. Rejecting at install time is the point: two + codecs sharing an id are indistinguishable when a payload is decoded, + and only the caller knows whether they write the same wire format.""" + ctx = SessionContext().with_logical_extension_codec(MyLogicalExtensionCodec()) - # 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: + with pytest.raises(ValueError, match="already installed"): + ctx.with_logical_extension_codec(MyLogicalExtensionCodec()) + + +def test_bare_capsule_codec_is_session_local(): + """A bare PyCapsule exposes nothing stable to derive an identity + from -- every capsule reports the same type -- so it is tagged with a + session-local id. A plan it encodes fails on an unrelated session + with an error naming the fix, rather than being decoded by whichever + codec happens to sit at the same position.""" + exporter = SessionContext().with_logical_extension_codec( + MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + ) + encoder = SessionContext().with_logical_extension_codec( + exporter.__datafusion_logical_extension_codec__() + ) + encoder.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder) + + with pytest.raises(Exception, match="bare PyCapsule") as excinfo: LogicalPlan.from_bytes(SessionContext(), blob) + assert "codec_id" in str(excinfo.value) + + +def test_default_only_session_writes_no_envelope(): + """A session with no extension codecs installed produces the same + bytes as a build without codec chaining: the terminal codec writes + unframed, so the envelope only appears once a codec is installed. - assert "composed extension codecs" not in str(excinfo.value) + Keeps the wire break scoped to sessions that actually compose.""" + ctx = SessionContext() + blob = ctx.sql("SELECT abs(-1) AS x").logical_plan().to_bytes(ctx) + assert b"DFPYCHN" not in blob def test_udf_inlining_setting_survives_codec_install(): 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 c6f9189d4..97453b279 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -222,6 +222,23 @@ def codec_context(max_rows: int = 3): return ctx, logical_codec, physical_codec +def physical_only_context(max_rows: int = 3): + """Context with the physical codec installed but no logical codec. + + Encoding consults chained codecs in install order and the first to claim an + object wins, so a codec installed later cannot be observed while an earlier + one is already claiming table providers. Leaving the logical slot empty lets + a test install exactly one logical codec -- through a handle it then throws + away -- and watch its counters. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows)) + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + return ctx, physical_codec + + def test_installing_a_planner_keeps_the_session_id(): """A session and its decode callbacks must agree on the session id. @@ -521,9 +538,11 @@ def test_a_discarded_derived_context_still_rebinds_the_planner(): A fresh codec instance is what makes it observable -- it carries its own counters, and the planner encodes the outbound logical plan with whichever - codec it is holding. + codec it is holding. The base context deliberately installs no logical + codec, so this one is the only candidate; an already-installed codec would + claim the provider first and hide the rebind. """ - ctx, _logical_codec, _physical_codec = configured_context(max_rows=3) + ctx, _physical_codec = physical_only_context(max_rows=3) ctx.set_query_planner(MyQueryPlanner()) later = MyLogicalExtensionCodec() @@ -548,14 +567,17 @@ def test_the_planner_and_the_handle_can_hold_different_codecs(): Chaining ``ctx = ctx.with_...(...)`` keeps the two in step; this pins what happens when they are allowed to diverge. + ``ctx`` installs no logical codec of its own, so its chain is empty and the + planner's holds exactly one entry. That asymmetry is what makes the split + visible: an entry on both chains would be claimed by the same codec either + way, since encoding stops at the first codec to claim an object. + Inlining has to be off for the assertion to say anything: with it on, a Python UDF is encoded inline by ``PythonLogicalCodec`` and never reaches the installed codec's ``try_encode_udf``. """ config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) - handle_codec = MyLogicalExtensionCodec() ctx = SessionContext(config).with_python_udf_inlining(enabled=False) - ctx = ctx.with_logical_extension_codec(handle_codec) ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) ctx.register_table("numbers", MyTableProvider(1, 6, 1)) ctx.set_query_planner(MyQueryPlanner()) @@ -575,15 +597,14 @@ def test_the_planner_and_the_handle_can_hold_different_codecs(): ctx.register_udf(identity) Expr.to_bytes(identity(col("A")), ctx) - # Serializing through `ctx` uses `ctx`'s own codec field. - assert handle_codec.encode_udf_calls() > 0 + # Serializing through `ctx` uses `ctx`'s own codec field, which is empty -- + # the UDF goes out by name through the terminal codec. assert planner_codec.encode_udf_calls() == 0 ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() # Planning through the same `ctx` uses the codec the planner was rebound to. assert planner_codec.table_provider_encode_calls() > 0 - assert handle_codec.table_provider_encode_calls() == 0 def test_an_unchanged_inlining_setting_leaves_the_planner_alone(): @@ -595,11 +616,11 @@ def test_an_unchanged_inlining_setting_leaves_the_planner_alone(): Observable only once the planner is holding some *other* handle's codec: without the guard, a defensive no-op toggle on `ctx` drags the planner back - onto `ctx`'s codec and silently undoes the install below. The rebuilt - codecs otherwise wrap the same inner codec, so nothing else distinguishes - the two paths. + onto `ctx`'s codecs and silently undoes the install below. `ctx` installs no + logical codec, so being dragged back leaves the planner with an empty chain + and the query fails outright rather than quietly using the wrong codec. """ - ctx, handle_codec, _physical_codec = codec_context() + ctx, _physical_codec = physical_only_context() ctx.set_query_planner(MyQueryPlanner()) planner_codec = MyLogicalExtensionCodec() @@ -612,7 +633,6 @@ def test_an_unchanged_inlining_setting_leaves_the_planner_alone(): ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() assert planner_codec.table_provider_encode_calls() > 0 - assert handle_codec.table_provider_encode_calls() == 0 def test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs(): @@ -627,7 +647,7 @@ def test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs(): So "re-install the planner after installing a codec" only repairs anything when it is done from the handle holding the new codec. """ - ctx, original_logical, _physical_codec = codec_context() + ctx, _physical_codec = physical_only_context() planner = MyQueryPlanner() ctx.set_query_planner(planner) @@ -638,15 +658,14 @@ def test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs(): ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() assert later.table_provider_encode_calls() > 0 - assert original_logical.table_provider_encode_calls() == 0 - # `ctx`'s own codec field never changed, so this rebuilds the planner - # against `original_logical` and drops `later` from the session's planner. + # `ctx`'s own codec field never changed -- it never had a logical codec -- + # so this rebuilds the planner against an empty chain and drops `later`. ctx.set_query_planner(planner) encodes_by_later = later.table_provider_encode_calls() - ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() - assert original_logical.table_provider_encode_calls() > 0 + with pytest.raises(Exception, match=r"LogicalExtensionCodec|TableProvider"): + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() assert later.table_provider_encode_calls() == encodes_by_later diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 6a4fbdb18..b13cc26bb 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2254,7 +2254,9 @@ def __datafusion_query_planner__(self, session: Any = None) -> Any: return self.ctx.__datafusion_query_planner__(session) def with_logical_extension_codec( - self, codec: LogicalExtensionCodecExportable | _PyCapsule + self, + codec: LogicalExtensionCodecExportable | _PyCapsule, + codec_id: str | None = None, ) -> SessionContext: """Create a new session context with an additional logical codec. @@ -2262,13 +2264,25 @@ def with_logical_extension_codec( ``__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. + Codecs compose: each call appends the codec to the session's codec + chain rather than replacing prior codecs. Every payload a codec writes + is tagged with that codec's identity, so decoding consults exactly the + codec that encoded it and never offers bytes to a codec that did not + write them. Install order therefore does not affect decoding at all. + + On encoding, codecs are consulted in install order and the first one to + claim an object wins, so installing a codec can only claim objects no + earlier codec claimed. Order is only observable when two codecs both + claim the same object, which is a collision worth avoiding regardless. + + ``codec_id`` sets the identity used for tagging. It is normally + unnecessary: an identity is derived from the codec's + ``__datafusion_codec_id__`` attribute if present, otherwise from its + class's module and qualified name, which is stable across processes. + Pass it explicitly when installing from a bare ``PyCapsule``, which + exposes nothing stable to derive from — such a codec is tagged with a + session-local identity, and plans it encodes will not decode on an + unrelated session. The returned context shares its session state with the original, so a later registration on either is visible to both. If a custom query @@ -2276,8 +2290,22 @@ def with_logical_extension_codec( 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. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx = ctx.with_logical_extension_codec( + ... my_library.Codec() + ... ) # doctest: +SKIP + + Installing from a bare capsule, pinning the identity so encoded + plans remain decodable on another session: + + >>> ctx = ctx.with_logical_extension_codec( + ... capsule, codec_id="my_library.Codec" + ... ) # doctest: +SKIP """ - new_internal = self.ctx.with_logical_extension_codec(codec) + new_internal = self.ctx.with_logical_extension_codec(codec, codec_id) new = SessionContext.__new__(SessionContext) new.ctx = new_internal return new @@ -2290,7 +2318,9 @@ def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: return self.ctx.__datafusion_physical_extension_codec__(session) def with_physical_extension_codec( - self, codec: PhysicalExtensionCodecExportable | _PyCapsule + self, + codec: PhysicalExtensionCodecExportable | _PyCapsule, + codec_id: str | None = None, ) -> SessionContext: """Create a new session context with an additional physical codec. @@ -2299,9 +2329,10 @@ def with_physical_extension_codec( :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. + :py:meth:`with_logical_extension_codec`: each call appends to the + session's codec chain, payloads are tagged with the identity of the + codec that wrote them, and ``codec_id`` overrides that identity. See + that method for the full description. The returned context shares its session state with the original, so a later registration on either is visible to both. If a custom query @@ -2309,8 +2340,19 @@ def with_physical_extension_codec( 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. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx = ctx.with_physical_extension_codec( + ... my_library.PhysicalCodec() + ... ) # doctest: +SKIP + + >>> ctx = ctx.with_physical_extension_codec( + ... capsule, codec_id="my_library.PhysicalCodec" + ... ) # doctest: +SKIP """ - new_internal = self.ctx.with_physical_extension_codec(codec) + new_internal = self.ctx.with_physical_extension_codec(codec, codec_id) new = SessionContext.__new__(SessionContext) new.ctx = new_internal return new From d8b5ad7109c0d1bf3b1424ab2bb5bfa6f1b50752 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 31 Aug 2026 16:42:15 -0400 Subject: [PATCH 05/13] test: pin the by-name decode path, and document identity dispatch Adds NameOnlyUdfCodec to the FFI example: a codec owning functions that are fully described by their names, so try_encode_udf writes nothing and try_decode_udf rebuilds the function from the name with no registry entry. DataFusion supports that shape directly -- an empty fun_definition sends the decoder to the FunctionRegistry first and the codec second (from_proto.rs, the `None => ctx.udf(..).or_else(..)` arm). Nothing covered that arm before, and it is the path most at risk from a plausible change: wrapping every chained encode in the identity envelope would make an empty payload non-empty, set fun_definition, and skip the registry lookup permanently. That breaks ordinary by-name round trips as well as codecs like this one, and no other test would notice. The decoding session here deliberately never registers the function, so only the codec can supply it. Exposes logical_extension_codec_ids() and physical_extension_codec_ids() on SessionContext. The ids are the dispatch keys a payload names, so listing them answers "which library owns this plan" and "can this session decode it" -- and they are what a decode failure reports. Documentation rewritten around identity rather than ordering: - ffi.md states that codecs need no changes, why rejecting foreign payloads is not something a codec can reliably do, the identity resolution ladder, and the two payloads left unframed with the reason for each. - Both example READMEs drop the "most recently installed is consulted first" language, which described dispatch that no longer exists. - upgrade-guides.md gains a section for the behaviour change: codecs compose rather than replace, no codec-side change is required, and the wire format changes only for sessions that install one. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 21 ++ docs/source/contributor-guide/ffi.md | 81 +++++- docs/source/user-guide/upgrade-guides.md | 47 +++ examples/datafusion-ffi-example/README.md | 6 +- .../tests/_test_logical_extension_codec.py | 59 +++- examples/datafusion-ffi-example/src/lib.rs | 4 + .../src/name_only_codec.rs | 268 ++++++++++++++++++ .../README.md | 2 +- python/datafusion/context.py | 37 +++ 9 files changed, 501 insertions(+), 24 deletions(-) create mode 100644 examples/datafusion-ffi-example/src/name_only_codec.rs diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 8a87e36ac..6b072b7e1 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1490,6 +1490,27 @@ impl PySessionContext { Ok(derived) } + /// Ids of the logical extension codecs installed on this session, in + /// install order — the same order encoding consults them in, and the keys + /// a payload names when it is decoded. + pub fn logical_extension_codec_ids(&self) -> Vec { + self.logical_codec + .codec_ids() + .into_iter() + .map(str::to_string) + .collect() + } + + /// Ids of the physical extension codecs installed on this session. + /// See [`Self::logical_extension_codec_ids`]. + pub fn physical_extension_codec_ids(&self) -> Vec { + self.physical_codec + .codec_ids() + .into_iter() + .map(str::to_string) + .collect() + } + /// See [`Self::__datafusion_logical_extension_codec__`] for `session`. #[pyo3(signature = (session=None))] pub fn __datafusion_physical_extension_codec__<'py>( diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 67a1e7385..edc47ef74 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -251,20 +251,52 @@ durable metadata instead. ### 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: +`with_physical_extension_codec` appends the codec to the session's codec chain +rather than replacing prior codecs. + +**Nothing is asked of the codec itself.** Implement `LogicalExtensionCodec` or +`PhysicalExtensionCodec` exactly as you would for a session that installs only +yours. `Python{Logical,Physical}Codec` sits between DataFusion and every installed +codec, and it wraps each payload in an envelope naming the codec that produced it. +Your codec receives, byte for byte, the payload it wrote, and never sees the +envelope. + +Decoding reads that name and consults exactly one codec. A codec is never offered +bytes it did not write, so it does not have to recognise and reject foreign +payloads — which is not something a codec can reliably do anyway. Protobuf carries +no type identity: a `prost` message decodes cleanly from an unrelated message's +bytes whenever their leading field numbers and wire types line up, and the natural +implementation, `MyMessage::decode(buf)`, has no prefix to check and cannot +decline. Trying codecs in turn until one succeeds is how upstream's +`ComposedPhysicalExtensionCodec` came to decode a Parquet payload as CSV +([apache/datafusion#16980](https://github.com/apache/datafusion/issues/16980)). + +Codec identity is derived automatically and is stable across processes: + +1. An explicit `codec_id=` argument to the install call. +2. `__datafusion_codec_id__` on the exporting object, if it defines one. Declare it + when a class rename must not invalidate previously encoded plans, or when one + library installs two instances owning disjoint slices of the wire format. +3. Otherwise the exporting class's `module.QualName`, which is the library's own + import path. + +Two codecs cannot share an identity — installing a second under an id already in +use raises rather than shadowing the first, because a payload naming that id would +otherwise resolve to whichever entry came first. A codec installed from a bare +`PyCapsule` is the one case with nothing stable to derive from, since every capsule +reports the same type; it gets a session-local identity, and plans it encodes fail +with a pointed error on an unrelated session instead of being decoded by the wrong +codec. Pass `codec_id=` for those. + +`SessionContext.logical_extension_codec_ids()` and its physical counterpart list +what is installed, which is also what a decode failure names. + +Because decoding keys off identity rather than install position, registration order +between independent libraries does not affect decoding at all. It is visible only +on encoding, where codecs are consulted in install order and the first to claim an +object wins — so installing a library can claim objects nothing else claimed, but +never takes over an object an earlier codec was already encoding. Two libraries +that each own tables, functions, and a planner register like this: ```python ctx = SessionContext(config) @@ -286,6 +318,27 @@ ctx.register_table("t", lib_a.TableProvider()) ctx.register_udf(udf(lib_b.SomeUDF())) ``` +Two payloads are deliberately left unframed, and both matter if you are changing +this code. + +The terminal codec — `Default{Logical,Physical}ExtensionCodec` unless a Rust caller +supplied another to `Python{Logical,Physical}Codec::new` — handles whatever no +installed codec claims and writes bare. A session with no extension codecs +installed therefore serializes byte-identically to a build without codec chaining. + +An encode that writes nothing also stays empty. `try_encode_udf` returning `Ok` +with an empty buffer is DataFusion's encode-by-name signal: it leaves +`fun_definition` unset, and the decoder then tries the `FunctionRegistry` first and +the codec second. Framing an empty payload would set the field and skip that +registry lookup permanently, breaking both ordinary by-name round trips and codecs +whose functions are reconstructible from a name alone — the case +`NameOnlyUdfCodec` in the FFI example covers. That empty-buffer decode is also the +one path where every installed codec is still consulted in turn, because there are +no bytes to carry an identity. It is not the hazard the envelope removes: the +question asked is "do you own the function named `x`", which is name-scoped, and +two codecs disagreeing requires them to claim the same function name — already a +collision in the function registry. + 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 local build commands. diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 29085bc3d..a6ce1348a 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -100,6 +100,53 @@ clear message rather than undefined behaviour on first use. `FFI_TaskContextProvider`, `FFI_TableProviderFactory`, and `FFI_ExtensionOptions` carry no version field, so objects of those types cannot be checked. +### Extension codecs compose instead of replacing + +`SessionContext.with_logical_extension_codec` and +`with_physical_extension_codec` previously replaced whichever codec was already +installed, so a session could only ever have one. Installing a second codec +silently discarded the first, and plans failed later with a confusing decode +error. Both methods now append to a chain, and a session can carry codecs from +several independent libraries at once. + +**No change is required in an extension codec.** Keep implementing +`LogicalExtensionCodec` or `PhysicalExtensionCodec` exactly as before. Payloads +are wrapped in an envelope naming their author by `datafusion-python`, which +strips it again before your codec sees the bytes. + +Callers relying on replacement semantics — installing a codec in order to remove +a previous one — are affected. There is no way to remove an installed codec. + +Two behaviours are worth knowing: + +- Installing two codecs under the same identity raises a `ValueError`. Identity + is derived from the exporting class's module and qualified name, so this comes + up when installing two instances of one class. Pass `codec_id=` to distinguish + them. +- A codec installed from a bare `PyCapsule` has no portable identity, because + every capsule reports the same type. It is tagged with a session-local + identity and works normally within that session, but a plan it encodes cannot + be decoded on an unrelated session. Pass `codec_id=` if plans must cross + sessions. + +```python +ctx = ctx.with_logical_extension_codec(lib_a.codec()) +ctx = ctx.with_logical_extension_codec(lib_b.codec()) # no longer discards lib_a + +# Two instances of one class need distinct identities. +ctx = ctx.with_logical_extension_codec(lib_a.Codec(), codec_id="lib_a.reader") +ctx = ctx.with_logical_extension_codec(lib_a.Codec(), codec_id="lib_a.writer") + +ctx.logical_extension_codec_ids() +``` + +Serialized plans change shape once an extension codec is installed: payloads +written by a chained codec now carry an identity envelope. A session with no +extension codecs installed is unaffected and produces the same bytes as before, +as do functions encoded by name. Plans serialized by an earlier release and +stored for later use should be regenerated if they were produced by a session +with an extension codec installed. + ### Changes to the `datafusion-python-util` crate Extension libraries written in Rust usually depend on the diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index cef39815d..4dc6c2a6a 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -35,9 +35,11 @@ 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. -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 example keeps the provider library as 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 appends the codec to the session's codec chain, and DataFusion's default codec handles whatever no installed codec claims. Every payload a codec writes is wrapped in an envelope naming that codec, and decoding consults exactly the codec that encoded it — so several independent plugin libraries can install codecs on the same session without any of them having to recognise or reject the others' payloads. The codecs here are written as ordinary `LogicalExtensionCodec` / `PhysicalExtensionCodec` implementations; the envelope is applied and stripped by `datafusion-python` and never reaches them. -`MyLogicalExtensionCodec` takes an optional `provider_prefix` argument (`MyLogicalExtensionCodec(provider_prefix="TOKENAAA")`) that overrides the byte prefix it stamps on encoded table providers. It exists so the tests can install two instances that own disjoint slices of the wire format, which is what makes 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. +`MyLogicalExtensionCodec` takes an optional `provider_prefix` argument (`MyLogicalExtensionCodec(provider_prefix="TOKENAAA")`) that overrides the byte prefix it stamps on encoded table providers. It exists so the tests can install two instances that own disjoint slices of the wire format, which is what makes install ordering observable from Python. Two instances of one class share a derived identity, so those tests also pass `codec_id=` to tell them apart. Real plugin libraries should hard-code a prefix unique to the library rather than accept one from the caller. + +`NameOnlyUdfCodec` is the opposite shape: it owns functions that are fully described by their names, so it encodes no bytes at all and rebuilds each function from the name on decode. It exists to pin the by-name path, which is the one place a payload carries no identity to dispatch on. Register both provider codecs before installing the planner: 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 16588e6f3..89f4053a9 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 @@ -20,7 +20,12 @@ import pyarrow as pa import pytest from datafusion import Expr, LogicalPlan, SessionContext, col, udf -from datafusion_ffi_example import MyLogicalExtensionCodec, MyTableProvider +from datafusion_ffi_example import ( + MyLogicalExtensionCodec, + MyTableProvider, + NameOnlyFunction, + NameOnlyUdfCodec, +) def _double_udf(): @@ -66,8 +71,8 @@ def _setup_session_with_codec() -> tuple[SessionContext, MyLogicalExtensionCodec def test_ffi_logical_codec_install_and_export(): - """Installing a user FFI codec replaces the session's logical - codec; the capsule getter on the session re-exports it.""" + """Installing a user FFI codec adds it to the session's logical + codec chain; the capsule getter on the session re-exports it.""" ctx, _codec = _setup_session_with_codec() capsule = ctx.__datafusion_logical_extension_codec__() assert capsule is not None @@ -116,12 +121,12 @@ def test_ffi_logical_codec_roundtrip(): def test_ffi_logical_codec_composes_with_later_install(): - """Codecs compose: installing a second codec prepends it to the + """Codecs compose: installing a second codec appends it to the session's codec chain instead of replacing the first. The second codec here (a default-backed codec exported from a fresh session) - cannot encode this library's table provider, so encoding falls - through to the user codec installed first. Under replace semantics - this test fails with `LogicalExtensionCodec is not provided`.""" + cannot encode this library's table provider, so the first codec + still claims it. Under replace semantics this test fails with + `LogicalExtensionCodec is not provided`.""" ctx, codec = _setup_session_with_codec() ctx = ctx.with_logical_extension_codec( SessionContext().__datafusion_logical_extension_codec__() @@ -297,6 +302,46 @@ def test_bare_capsule_codec_is_session_local(): assert "codec_id" in str(excinfo.value) +def test_name_only_codec_round_trips_without_a_payload(): + """A codec may own functions that need no payload: the name is the + whole encoding. ``try_encode_udf`` writes nothing, and the decoder + rebuilds the function from the name with no registry entry. + + DataFusion supports this directly -- an empty ``fun_definition`` + sends the decoder to the registry first and the codec second. This + test pins that arm from the Python side, because it is the one path + where a payload is still offered to every installed codec: there are + no bytes, so there is no identity to dispatch on. + + It is also the guard against a plausible "improvement". Wrapping + every chained encode in the identity envelope would make this + payload non-empty, which sets ``fun_definition`` and permanently + skips the registry lookup -- breaking both this codec and ordinary + by-name round trips, with nothing else in the suite noticing. + """ + codec = NameOnlyUdfCodec() + name = NameOnlyUdfCodec.function_name() + + # FROM-less, so serialization never reaches try_encode_table_provider -- + # this codec owns functions, not providers. + encoder = SessionContext().with_logical_extension_codec(codec) + encoder.register_udf(udf(NameOnlyFunction())) + blob = encoder.sql(f"SELECT {name}(1) AS x").logical_plan().to_bytes(encoder) + + # The name is the entire encoding, so the codec contributed no bytes + # and the payload carries no identity envelope for it. + assert codec.encode_udf_calls() > 0 + assert b"DFPYCHN" not in blob + + # A fresh session that never registered the function: only the codec + # can supply it, and only from the name. + decoder = SessionContext().with_logical_extension_codec(codec) + restored = LogicalPlan.from_bytes(decoder, blob) + + assert codec.decode_udf_calls() > 0 + assert decoder.create_dataframe_from_logical_plan(restored).collect() + + def test_default_only_session_writes_no_envelope(): """A session with no extension codecs installed produces the same bytes as a build without codec chaining: the terminal codec writes diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 3d00fdb3e..92fccb1e2 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -21,6 +21,7 @@ use crate::aggregate_udf::MySumUDF; use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList}; use crate::config::MyConfig; use crate::logical_extension_codec::MyLogicalExtensionCodec; +use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec}; use crate::physical_extension_codec::MyPhysicalExtensionCodec; use crate::physical_optimizer::MyPhysicalOptimizerRule; use crate::scalar_udf::IsNullUDF; @@ -33,6 +34,7 @@ pub(crate) mod aggregate_udf; pub(crate) mod catalog_provider; pub(crate) mod config; pub(crate) mod logical_extension_codec; +pub(crate) mod name_only_codec; pub(crate) mod physical_extension_codec; pub(crate) mod physical_optimizer; pub(crate) mod required_udf; @@ -57,6 +59,8 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/examples/datafusion-ffi-example/src/name_only_codec.rs b/examples/datafusion-ffi-example/src/name_only_codec.rs new file mode 100644 index 000000000..9b82c4bd1 --- /dev/null +++ b/examples/datafusion-ffi-example/src/name_only_codec.rs @@ -0,0 +1,268 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A codec whose functions need no payload at all. +//! +//! Most extension codecs answer with bytes. This one owns a fixed catalog of +//! functions that are fully described by their names, so `try_encode_udf` +//! writes nothing and `try_decode_udf` rebuilds the function from `name` +//! alone. DataFusion supports that shape directly: an encoder that writes no +//! bytes leaves `fun_definition` unset, and the decoder then tries the +//! `FunctionRegistry` first and the codec second — see the +//! `None => ctx.udf(..).or_else(|_| codec.try_decode_udf(name, &[]))` arm in +//! `datafusion-proto`'s `from_proto.rs`. +//! +//! It exists here to pin that arm. Because there are no bytes, there is +//! nothing to tag with the codec's identity, so this is the one path where +//! `PythonLogicalCodec` still offers a payload to every installed codec in +//! turn. A change that wrapped empty encodings in an envelope would set +//! `fun_definition`, skip the registry lookup permanently, and break both this +//! codec and plain by-name round trips — with no other test noticing. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow_schema::DataType; +use datafusion::common::error::Result; +use datafusion::common::not_impl_err; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, + Volatility, +}; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio_runtime}; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +/// Prefix marking the functions this library owns. A name is the entire +/// encoding, so the prefix is the whole ownership test. +const NAME_PREFIX: &str = "name_only_"; + +/// Scalar function reconstructed purely from its name. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct NameOnlyUdf { + name: String, + signature: Signature, +} + +impl NameOnlyUdf { + fn new(name: impl Into) -> Self { + Self { + name: name.into(), + signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for NameOnlyUdf { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + Ok(args.args[0].clone()) + } +} + +#[derive(Default)] +struct Counters { + encode_udf: AtomicUsize, + decode_udf: AtomicUsize, +} + +impl fmt::Debug for Counters { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.debug_struct("Counters").finish_non_exhaustive() + } +} + +struct NameOnlyLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec, + counters: Arc, +} + +impl fmt::Debug for NameOnlyLogicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("NameOnlyLogicalExtensionCodec") + .finish_non_exhaustive() + } +} + +impl LogicalExtensionCodec for NameOnlyLogicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[datafusion::logical_expr::LogicalPlan], + ctx: &datafusion::execution::TaskContext, + ) -> Result { + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode( + &self, + node: &datafusion::logical_expr::Extension, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &datafusion::common::TableReference, + schema: arrow_schema::SchemaRef, + ctx: &datafusion::execution::TaskContext, + ) -> Result> { + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &datafusion::common::TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode_table_provider(table_ref, node, buf) + } + + /// Writes nothing on purpose. The name is the whole encoding, so there is + /// no payload to emit, and returning `Ok` with an empty buffer is how a + /// codec says "encoded by name" to DataFusion. + fn try_encode_udf(&self, node: &ScalarUDF, _buf: &mut Vec) -> Result<()> { + if node.name().starts_with(NAME_PREFIX) { + self.counters.encode_udf.fetch_add(1, Ordering::SeqCst); + } + Ok(()) + } + + /// Rebuilds the function from `name`, with no registry entry and no bytes. + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + if !name.starts_with(NAME_PREFIX) { + return not_impl_err!("Not a name-only function: {name}"); + } + if !buf.is_empty() { + return not_impl_err!( + "name-only functions carry no payload, but {} bytes were supplied for {name}", + buf.len() + ); + } + self.counters.decode_udf.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(ScalarUDF::from(NameOnlyUdf::new(name)))) + } +} + +/// The function [`NameOnlyUdfCodec`] owns, exported so a session can register +/// it and build a plan that references it. +/// +/// Only the *encoding* session needs it registered. The decoding session +/// deliberately does not, which is what forces the codec's name-only decode +/// path to run. +#[pyclass( + from_py_object, + name = "NameOnlyFunction", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone)] +pub(crate) struct NameOnlyFunction; + +#[pymethods] +impl NameOnlyFunction { + #[new] + fn new() -> Self { + Self + } + + fn __datafusion_scalar_udf__<'py>(&self, py: Python<'py>) -> PyResult> { + let func = Arc::new(ScalarUDF::from(NameOnlyUdf::new(format!( + "{NAME_PREFIX}identity" + )))); + PyCapsule::new_with_value( + py, + datafusion_ffi::udf::FFI_ScalarUDF::from(func), + cr"datafusion_scalar_udf", + ) + } +} + +/// Codec owning functions that are reconstructible from their names alone. +/// +/// A real library shaped like this would be one shipping a fixed catalog of +/// built-ins: nothing about a call site varies, so there is nothing to encode. +#[pyclass( + from_py_object, + name = "NameOnlyUdfCodec", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Clone)] +pub(crate) struct NameOnlyUdfCodec { + counters: Arc, +} + +#[pymethods] +impl NameOnlyUdfCodec { + #[new] + fn new() -> Self { + Self { + counters: Arc::new(Counters::default()), + } + } + + /// Name of the function this codec can rebuild, for use in a query. + #[staticmethod] + fn function_name() -> String { + format!("{NAME_PREFIX}identity") + } + + fn encode_udf_calls(&self) -> usize { + self.counters.encode_udf.load(Ordering::SeqCst) + } + + fn decode_udf_calls(&self) -> usize { + self.counters.decode_udf.load(Ordering::SeqCst) + } + + fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let inner: Arc = Arc::new(NameOnlyLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec {}, + counters: Arc::clone(&self.counters), + }); + + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), ctx_provider); + + PyCapsule::new_with_value(py, ffi, cr"datafusion_logical_extension_codec") + } +} diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index f188ec740..dd829cc46 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 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 nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). +The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and payloads are tagged with the identity of the codec that wrote them, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). For the limits behind that choice — 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/python/datafusion/context.py b/python/datafusion/context.py index b13cc26bb..a6470e015 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2310,6 +2310,30 @@ def with_logical_extension_codec( new.ctx = new_internal return new + def logical_extension_codec_ids(self) -> list[str]: + """List the logical extension codecs installed on this session. + + Returns the identity of each installed codec, in install order. Those + identities are what encoding stamps onto a payload and what decoding + dispatches on, so this is how to check which library owns a plan and + whether a session is able to decode one. + + The terminal codec is not listed. It handles whatever no installed + codec claims and writes unframed, so it is not addressable by id. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx.logical_extension_codec_ids() + [] + >>> ctx = ctx.with_logical_extension_codec( + ... my_library.Codec() + ... ) # doctest: +SKIP + >>> ctx.logical_extension_codec_ids() # doctest: +SKIP + ['my_library.Codec'] + """ + return self.ctx.logical_extension_codec_ids() + def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: """Access the PyCapsule FFI_PhysicalExtensionCodec. @@ -2317,6 +2341,19 @@ def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: """ return self.ctx.__datafusion_physical_extension_codec__(session) + def physical_extension_codec_ids(self) -> list[str]: + """List the physical extension codecs installed on this session. + + See :py:meth:`logical_extension_codec_ids`. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx.physical_extension_codec_ids() + [] + """ + return self.ctx.physical_extension_codec_ids() + def with_physical_extension_codec( self, codec: PhysicalExtensionCodecExportable | _PyCapsule, From ef80b64c3ee2bc8faf909090826fd833ddebb952 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 31 Aug 2026 16:48:52 -0400 Subject: [PATCH 06/13] docs: record codec chain dispatch in the FFI capsule protocol skill The skill already triggers on "any FFI_* export that asks for ... an extension codec", but had nothing to say about how a chain decides which codec handles a payload -- the part of this area most likely to be got wrong, and the part with a rationale that is not visible from the code. Rule 8 records it: dispatch by identity rather than by trying codecs until one returns Ok, with the upstream incident that settles the question (apache/datafusion#16980, a Parquet payload decoded as CSV) and the protobuf reason a byte-prefix convention cannot fix it. It also records why we key on identity where Ballista, datafusion-distributed, and upstream all key on chain position. Their lists cannot disagree -- one is a compile-time constant, the other a pinned entry plus user codecs rebuilt from the same startup code. Ours is assembled by user Python, and to_bytes/from_bytes puts two independently built sessions on either end of one payload. Copying the upstream design here reintroduces the same silent mis-decode from the other direction, which is exactly the kind of thing an agent reading only upstream would do. And it records the two payloads that are never framed, because framing the empty one is a plausible tidy-up that breaks by-name decoding silently. Points at NameOnlyUdfCodec as the guard. Extends the trigger to cover codec.rs dispatch, which the old wording did not obviously reach, and documents the optional __datafusion_codec_id__ hook on the two codec Protocols -- where Rule 5 says such things belong. Reviewed the other four skills: check-upstream, make-pythonic, audit-skill-md, and the user-facing datafusion_python skill are all unaffected, none of them touching codec dispatch. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 73 +++++++++++++++++++++++- python/datafusion/user_defined.py | 16 +++++- 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 468216034..936381d66 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -1,7 +1,7 @@ --- name: ffi-capsule-protocol -description: "TRIGGER — read before adding, changing, or reviewing any __datafusion_*__ capsule getter, any FFI_* export that asks for a TaskContextProvider or an extension codec, or any code that calls FFI_QueryPlanner::new / FFI_TableProvider::new / FFI_{Logical,Physical}ExtensionCodec::new. These methods are one protocol with a settled convention. Do not design it fresh; do not construct a SessionContext inside an extension library." -argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", or omit to review the whole family)" +description: "TRIGGER — read before adding, changing, or reviewing any __datafusion_*__ capsule getter, any FFI_* export that asks for a TaskContextProvider or an extension codec, any code that calls FFI_QueryPlanner::new / FFI_TableProvider::new / FFI_{Logical,Physical}ExtensionCodec::new, or anything in crates/core/src/codec.rs that decides which extension codec handles a payload. These methods are one protocol with a settled convention. Do not design it fresh; do not construct a SessionContext inside an extension library; do not dispatch a codec chain by trying codecs until one succeeds." +argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", \"codec chain\", or omit to review the whole family)" ---