diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 468216034..294ebfb3a 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -116,6 +116,12 @@ one of them, and the failure is a bare `TypeError` from a `call1`. So: `python/datafusion/user_defined.py`, where the `Protocol` type hints for these methods live. +Changing what a codec puts *on the wire* is equally breaking, and easier to +miss because no signature moves and nothing fails to compile. Serialized plans +outlive the process that wrote them, so the same checklist applies: upgrade +guide, `api change` label, and a statement of exactly which sessions produce +different bytes. + ## Rule 6 — a session keeps one `Arc` for life `FFI_TaskContextProvider` holds its provider **weakly**, and every codec handed @@ -181,8 +187,11 @@ pins that; changing it should be deliberate. - `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. - `docs/source/user-guide/upgrade-guides.md` — every past migration. +- `crates/core/src/codec.rs` — the codec chain: the envelope, identity dispatch, + and the two unframed cases from Rule 8. - `examples/datafusion-ffi-example/src/` — provider, catalog, function, codec - getters, all in current form. + getters, all in current form. `name_only_codec.rs` is the codec that encodes + nothing. - `examples/datafusion-ffi-query-planner-example/src/planner.rs` — planner getter. - `examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` diff --git a/AGENTS.md b/AGENTS.md index 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. diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 94942a2d2..fd32fae7f 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -29,16 +29,18 @@ //! //! [`PythonLogicalCodec`] is the [`LogicalExtensionCodec`] that //! datafusion-python parks on every `SessionContext`. It wraps a -//! user-supplied (or default) inner codec and adds Python-aware -//! in-band encoding on top: when the encoder sees a Python-defined -//! UDF, the codec cloudpickles the callable + signature into the -//! `fun_definition` proto field; when the decoder sees a payload it -//! produced, it reconstructs the UDF from the bytes alone — no -//! pre-registration on the receiver. UDFs the codec does not -//! recognise are delegated to `inner`, which is typically -//! `DefaultLogicalExtensionCodec` but may be a downstream-supplied -//! FFI codec installed via -//! `SessionContext.with_logical_extension_codec(...)`. +//! chain of composable codecs and adds Python-aware in-band encoding +//! on top: when the encoder sees a Python-defined UDF, the codec +//! cloudpickles the callable + signature into the `fun_definition` +//! proto field; when the decoder sees a payload it produced, it +//! reconstructs the UDF from the bytes alone — no pre-registration on +//! the receiver. Everything the codec does not recognise is delegated +//! to the chain: each downstream FFI codec installed via +//! `SessionContext.with_logical_extension_codec(...)` is appended, and +//! encoding consults them in install order with +//! `DefaultLogicalExtensionCodec` as the terminal fallback. Decoding +//! does not walk the chain at all — a payload names the codec that +//! wrote it. See [`PythonLogicalCodec`]. //! //! [`PythonPhysicalCodec`] is the symmetric wrapper around //! [`PhysicalExtensionCodec`]. Logical and physical layers each have @@ -58,7 +60,7 @@ //! actionable error instead of an opaque `marshal` failure on load //! (cloudpickle payloads are not portable across Python minor //! versions). Dispatch precedence on decode: **family match + -//! supported version + matching Python version → `inner` codec → +//! supported version + matching Python version → codec chain → //! caller's `FunctionRegistry` fallback.** //! //! ## Wire-format family registry @@ -81,10 +83,15 @@ //! for an older shape. //! //! Downstream FFI codecs should pick non-colliding family prefixes -//! (use a `DF` namespace plus a crate-specific suffix). The codec -//! implementations in this module currently delegate every method to -//! `inner`; the encoder/decoder hooks for each kind are added as the -//! corresponding Python-side type becomes serializable. +//! (use a `DF` namespace plus a crate-specific suffix) and return an +//! error for *objects* they do not own — on encode, that error is the +//! chain's "not mine" signal, letting the next codec take a turn. A +//! codec that answers `Ok` for objects outside its family claims them +//! ahead of every codec installed after it. +//! +//! Rejecting foreign *payloads* is not asked of a codec, because a +//! codec cannot reliably do it: a payload is only ever handed to the +//! codec whose id it carries. See [`PythonLogicalCodec`]. use std::sync::Arc; @@ -167,7 +174,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 +230,337 @@ fn strip_wire_header<'a>( Ok(Some(&buf[py_minor_idx + 1..])) } +/// Family prefix for the envelope wrapping a chained codec's payload. +/// +/// A distinct magic is what makes "is this framed?" a definite test +/// rather than a speculative decode. Probing by attempting to parse the +/// envelope would reintroduce exactly the protobuf ambiguity this +/// framing exists to remove: prost skips unknown fields and defaults +/// missing ones, so a foreign payload can parse cleanly as an envelope. +pub(crate) const CHAINED_PAYLOAD_FAMILY: &[u8] = b"DFPYCHN"; + +/// Wire-format version for the chained-payload envelope. Independent of +/// [`WIRE_VERSION_CURRENT`], which versions the cloudpickle framing. +pub(crate) const CHAIN_WIRE_VERSION_CURRENT: u8 = 1; + +/// Oldest chained-payload envelope version this build decodes. +pub(crate) const CHAIN_WIRE_VERSION_MIN_SUPPORTED: u8 = 1; + +/// Prefix for the synthetic id given to a codec installed from a bare +/// PyCapsule, which exposes nothing stable to derive an identity from. +/// The rest of the id is random per install, so no other session can +/// mint it: a payload carrying one decodes within the installing +/// session's lineage, which clones the id along with the chain, and +/// fails with a pointed error anywhere else rather than resolving to a +/// different codec. A counter or a chain position would not do — every +/// session numbers from the same end, so the first bare capsule +/// installed anywhere would answer for every other session's first. +pub(crate) const ANONYMOUS_CODEC_ID_PREFIX: &str = "anon:"; + +/// Prefix for the id a `SessionContext` reports when its own codec stack +/// is installed as an extension codec on another session. +/// +/// Deriving that id from the class, as an ordinary codec object's is, +/// would name every session at once: they all share one class. Two +/// sessions installed as codecs on one target would collide, and a +/// payload written by one would resolve to the other on decode. The +/// session id is per session and already stable, so it is what the +/// remainder of this id carries. +pub(crate) const SESSION_CODEC_ID_PREFIX: &str = "session:"; + +/// One installed codec plus the identity its payloads are tagged with. +/// +/// The id is what makes dispatch order-independent. Keying on position +/// in the chain — as `ComposedPhysicalExtensionCodec` does upstream — +/// is sound only when both ends assemble the same list in the same +/// order. That holds for the consumers upstream was written for, whose +/// lists structurally cannot disagree: Ballista's is a compile-time +/// constant, and `datafusion-distributed` pins its own codec at index 0 +/// and appends user codecs rebuilt from the same startup code on every +/// node. It does not hold here. A chain is assembled by user Python, +/// and `Expr.to_bytes(ctx1)` / `Expr.from_bytes(ctx2)` puts two +/// independently configured sessions on either end of one payload, so +/// an index would name a different codec in the decoder as soon as +/// install order differed. +struct ChainEntry { + 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. +/// +/// Do not replace this with a walk that hands `buf` to each codec until +/// one returns `Ok`. That looks simpler and removes the envelope, and it +/// is unsound: protobuf carries no type identity, so a `prost` message +/// decodes cleanly from an unrelated message's bytes whenever their +/// leading field numbers and wire types line up, and an all-defaults +/// message encodes to zero bytes that decode as anything. Asking codecs +/// to check a byte prefix does not fix it either, because the natural +/// implementation is `MyMessage::decode(buf)`, which has no prefix to +/// check and cannot decline. Upstream shipped that design and reverted +/// it after a Parquet payload decoded as CSV — apache/datafusion#16980, +/// fixed in #16986. +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 entry in chain { + match f(entry.codec.as_ref()) { + Ok(value) => return Ok(value), + Err(err) => errors.push(err), + } + } + match f(terminal.as_ref()) { + Ok(value) => Ok(value), + Err(err) => { + errors.push(err); + Err(aggregate_chain_errors(what, errors)) + } + } +} + +/// Collapse per-codec failures into one error. A single failure is +/// returned as-is so a session with no extension codecs behaves exactly +/// like a build without codec chaining. +fn aggregate_chain_errors( + what: &str, + mut errors: Vec, +) -> datafusion::error::DataFusionError { + 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!( + "No installed extension codec handled {what}: {joined}" + )) + } + } +} + +/// Encode through the chain, tagging the payload with its author. +/// +/// Entries are consulted in install order and the first one to write +/// bytes wins, so installing a codec can only claim objects no +/// earlier codec claimed. Adding a library therefore never changes how +/// an already-installed library's objects encode. +/// +/// Each codec encodes into a scratch buffer so a failed attempt cannot +/// leave partial bytes behind. `Ok` with an empty buffer is "no +/// opinion" rather than a claim, so the walk continues; if nothing +/// writes bytes the result is `Ok` with nothing written, which is +/// DataFusion's encode-by-name signal. Framing that empty result would +/// set `fun_definition` and permanently skip the registry lookup the +/// decoder does first. +/// +/// The terminal codec writes unframed, so a session with no extension +/// codecs is byte-compatible with a build predating the chain. +fn chain_encode( + 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 entry in chain { + let mut scratch = Vec::new(); + match f(entry.codec.as_ref(), &mut scratch) { + Ok(()) if !scratch.is_empty() => { + write_chained_payload(buf, &entry.id, &scratch); + return Ok(()); + } + Ok(()) => saw_empty_ok = true, + Err(err) => errors.push(err), + } + } + let mut scratch = Vec::new(); + match f(terminal.as_ref(), &mut scratch) { + Ok(()) if !scratch.is_empty() => { + buf.extend_from_slice(&scratch); + return Ok(()); + } + Ok(()) => saw_empty_ok = true, + Err(err) => errors.push(err), + } + if saw_empty_ok { + return Ok(()); + } + Err(aggregate_chain_errors(what, errors)) +} + /// `LogicalExtensionCodec` parked on every `SessionContext`. Holds /// the Python-aware encoding hooks for logical-layer types /// (`LogicalPlan`, `Expr`) and delegates everything it does not -/// handle to the composable `inner` codec — typically -/// `DefaultLogicalExtensionCodec`, or a downstream FFI codec -/// installed via `SessionContext.with_logical_extension_codec(...)`. +/// handle to a chain of composable codecs. Each downstream FFI codec +/// installed via `SessionContext.with_logical_extension_codec(...)` is +/// appended to the chain, and `terminal` — normally +/// `DefaultLogicalExtensionCodec` — handles whatever no installed codec +/// claims. +/// +/// Every payload an installed codec writes is wrapped in an envelope +/// naming that codec (see [`write_chained_payload`]), so decoding +/// consults exactly the codec that encoded it. Dispatch does not depend +/// on a codec recognizing and rejecting foreign payloads, which is not +/// something a codec can reliably do: a prost message decodes cleanly +/// from another message's bytes whenever their leading field numbers and +/// wire types line up. /// /// Sitting at the top of the session's logical codec stack means /// every serializer that reads `session.logical_codec()` automatically @@ -241,22 +573,71 @@ 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>, + 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 { - inner, + chain: Vec::new(), + terminal: inner, python_udf_inlining: true, } } - pub fn inner(&self) -> &Arc { - &self.inner + /// Return a copy of this codec with `codec` appended to the chain + /// under `id`, preserving the Python-UDF-inlining setting. + /// + /// Appending rather than prepending keeps the operation additive: + /// the new codec is consulted for encoding only after every codec + /// already installed, so it can claim objects nothing else claimed + /// but cannot take over an existing library's objects. + pub fn with_additional_codec( + &self, + id: impl Into>, + 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 @@ -297,11 +678,23 @@ impl LogicalExtensionCodec for PythonLogicalCodec { inputs: &[LogicalPlan], ctx: &TaskContext, ) -> Result { - self.inner.try_decode(buf, inputs, ctx) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an extension logical plan node", + |codec, buf| codec.try_decode(buf, inputs, ctx), + ) } fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { - self.inner.try_encode(node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "an extension logical plan node", + |codec, buf| codec.try_encode(node, buf), + ) } fn try_decode_table_provider( @@ -311,8 +704,13 @@ impl LogicalExtensionCodec for PythonLogicalCodec { schema: SchemaRef, ctx: &TaskContext, ) -> Result> { - self.inner - .try_decode_table_provider(buf, table_ref, schema, ctx) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a table provider", + |codec, buf| codec.try_decode_table_provider(buf, table_ref, Arc::clone(&schema), ctx), + ) } fn try_encode_table_provider( @@ -321,7 +719,13 @@ impl LogicalExtensionCodec for PythonLogicalCodec { node: Arc, buf: &mut Vec, ) -> Result<()> { - self.inner.try_encode_table_provider(table_ref, node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a table provider", + |codec, buf| codec.try_encode_table_provider(table_ref, Arc::clone(&node), buf), + ) } fn try_decode_file_format( @@ -329,7 +733,13 @@ impl LogicalExtensionCodec for PythonLogicalCodec { buf: &[u8], ctx: &TaskContext, ) -> Result> { - self.inner.try_decode_file_format(buf, ctx) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a file format", + |codec, buf| codec.try_decode_file_format(buf, ctx), + ) } fn try_encode_file_format( @@ -337,14 +747,26 @@ impl LogicalExtensionCodec for PythonLogicalCodec { buf: &mut Vec, node: Arc, ) -> Result<()> { - self.inner.try_encode_file_format(buf, node) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a file format", + |codec, buf| codec.try_encode_file_format(buf, Arc::clone(&node)), + ) } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } - self.inner.try_encode_udf(node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_encode_udf(node, buf), + ) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { @@ -355,14 +777,26 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } - self.inner.try_decode_udf(name, buf) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_decode_udf(name, buf), + ) } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } - self.inner.try_encode_udaf(node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_encode_udaf(node, buf), + ) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { @@ -373,14 +807,26 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } - self.inner.try_decode_udaf(name, buf) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_decode_udaf(name, buf), + ) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } - self.inner.try_encode_udwf(node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_encode_udwf(node, buf), + ) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { @@ -391,13 +837,19 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } - self.inner.try_decode_udwf(name, buf) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_decode_udwf(name, buf), + ) } } /// Strict-mode gate: if `buf` is a well-framed inline payload for /// `family`, return the strict-refusal error; otherwise return -/// `Ok(())` so the caller can delegate to its `inner` codec. +/// `Ok(())` so the caller can delegate to its codec chain. /// /// Routing through [`read_framed_payload`] (rather than a bare /// `starts_with` probe) means malformed inline bytes — wrong @@ -442,7 +894,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 +907,59 @@ 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>, + 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 { - inner, + chain: Vec::new(), + terminal: inner, python_udf_inlining: true, } } - pub fn inner(&self) -> &Arc { - &self.inner + /// Return a copy of this codec with `codec` appended to the chain + /// under `id`. See [`PythonLogicalCodec::with_additional_codec`]. + pub fn with_additional_codec( + &self, + id: impl Into>, + 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. @@ -500,7 +990,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - self.inner.try_decode(buf, inputs, ctx, proto_converter) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an execution plan", + |codec, buf| codec.try_decode(buf, inputs, ctx, proto_converter), + ) } fn try_encode( @@ -509,14 +1005,26 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - self.inner.try_encode(node, buf, proto_converter) + chain_encode( + &self.chain, + &self.terminal, + buf, + "an execution plan", + |codec, buf| codec.try_encode(Arc::clone(&node), buf, proto_converter), + ) } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } - self.inner.try_encode_udf(node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_encode_udf(node, buf), + ) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { @@ -527,7 +1035,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } - self.inner.try_decode_udf(name, buf) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a scalar UDF", + |codec, buf| codec.try_decode_udf(name, buf), + ) } fn try_encode_expr( @@ -536,7 +1050,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { buf: &mut Vec, ctx: &PhysicalExprEncodeCtx<'_>, ) -> Result<()> { - self.inner.try_encode_expr(node, buf, ctx) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a physical expression", + |codec, buf| codec.try_encode_expr(node, buf, ctx), + ) } fn try_decode_expr( @@ -545,14 +1065,26 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { inputs: &[Arc], ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - self.inner.try_decode_expr(buf, inputs, ctx) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a physical expression", + |codec, buf| codec.try_decode_expr(buf, inputs, ctx), + ) } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } - self.inner.try_encode_udaf(node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_encode_udaf(node, buf), + ) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { @@ -563,14 +1095,26 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } - self.inner.try_decode_udaf(name, buf) + chain_decode( + &self.chain, + &self.terminal, + buf, + "an aggregate UDF", + |codec, buf| codec.try_decode_udaf(name, buf), + ) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } - self.inner.try_encode_udwf(node, buf) + chain_encode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_encode_udwf(node, buf), + ) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { @@ -581,7 +1125,13 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } - self.inner.try_decode_udwf(name, buf) + chain_decode( + &self.chain, + &self.terminal, + buf, + "a window UDF", + |codec, buf| codec.try_decode_udwf(name, buf), + ) } } @@ -598,7 +1148,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 +1163,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) { @@ -1039,137 +1589,3 @@ fn decode_python_udaf(py: Python<'_>, payload: &[u8]) -> PyResult String { + format!("{SESSION_CODEC_ID_PREFIX}{}", self.ctx.session_id()) + } + /// `session` exists so this matches the protocol an extension library /// implements, where the argument is how the library reaches the session /// it is being installed on. A session already is one, so it is ignored. @@ -1457,22 +1477,28 @@ 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(); - // Carry the receiver's inlining setting over. `PythonLogicalCodec::new` - // defaults it to on, so building the replacement without this would - // silently re-enable inline Python UDF encoding on a context that had - // opted out with `with_python_udf_inlining(enabled=False)`. - let logical_codec = Arc::new( - PythonLogicalCodec::new(inner) - .with_python_udf_inlining(this.logical_codec.python_udf_inlining()), - ); + // Append rather than replace: previously installed codecs stay active, + // and every payload this one writes is tagged with `id` so decoding + // reaches it directly rather than by trying codecs in turn. Appending + // also carries the receiver's inlining setting over, which a fresh + // `PythonLogicalCodec::new` would not — it defaults inlining to on, and + // would silently re-enable inline Python UDF encoding on a context that + // had opted out with `with_python_udf_inlining(enabled=False)`. + let logical_codec = Arc::new(this.logical_codec.with_additional_codec(id, inner)); let derived = Self { ctx: Arc::clone(&this.ctx), logical_codec, @@ -1484,6 +1510,27 @@ 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>( @@ -1495,19 +1542,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 the flag is carried over. - let physical_codec = Arc::new( - PythonPhysicalCodec::new(inner) - .with_python_udf_inlining(this.physical_codec.python_udf_inlining()), - ); + // See `with_logical_extension_codec` for why this appends rather than + // replaces, and why that is also what carries the inlining flag over. + let physical_codec = Arc::new(this.physical_codec.with_additional_codec(id, inner)); let derived = Self { ctx: Arc::clone(&this.ctx), logical_codec: Arc::clone(&this.logical_codec), @@ -1525,7 +1576,7 @@ 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 +1588,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 { @@ -1705,6 +1760,80 @@ 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. +/// `SessionContext` pins its own through this arm — see +/// [`PySessionContext::__datafusion_codec_id__`] — because the class-derived +/// id below would name every session at once. +/// 3. The exporting object's `module.QualName`, which is the library's own +/// import path and therefore already stable across processes. This is the +/// common case and asks nothing of existing extension libraries. +/// 4. For a bare `PyCapsule` there is nothing stable to read — every capsule +/// reports the same type — so mint a fresh random id. Payloads tagged this +/// way decode correctly within the session lineage that installed the +/// codec, because the chain is cloned along with the id, and fail with a +/// pointed error everywhere else. Randomness is the point: an id drawn from +/// a namespace another session can mint the same value from — a counter, a +/// chain position — would let an unrelated codec answer for these bytes. +/// +/// An id already in use is rejected rather than shadowed. Two codecs sharing an +/// id are indistinguishable on decode, and the API cannot tell whether two +/// instances of one class write the same wire format — so the ambiguity is +/// surfaced at install time, where the caller can resolve it, instead of at +/// decode time, where it would pick whichever entry came first. +fn resolve_codec_id( + codec: &Bound<'_, PyAny>, + explicit: Option, + existing: &[&str], +) -> PyResult { + let id = derive_codec_id(codec, explicit)?; + if existing.contains(&id.as_str()) { + return Err(PyValueError::new_err(format!( + "An extension codec with id '{id}' is already installed on this session. Two \ + codecs cannot share an id, because a payload names its codec by id when it is \ + decoded. Pass `codec_id=` to give this one a distinct identity." + ))); + } + Ok(id) +} + +fn derive_codec_id(codec: &Bound<'_, PyAny>, explicit: Option) -> 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}{}", + Uuid::new_v4() + .simple() + .encode_lower(&mut Uuid::encode_buffer()) + )); + } + let ty = codec.get_type(); + let module = ty + .getattr("__module__") + .and_then(|m| m.extract::()) + .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/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 31cd9391f..f27251f5d 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -248,10 +248,115 @@ 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` 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. + `SessionContext` declares one itself, carrying its session id: installing one + context's codec stack on another session is the case where the class-derived id + below would name every session at once. +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 random identity minted at install time, 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. + +The randomness there is deliberate, not laziness. An identity another session can +mint the same value from — a counter, a position in the chain, a class every +candidate shares — reintroduces positional dispatch through the back door: every +session numbers from the same end, so one session's first bare capsule would answer +for every other session's first. + +Composing whole sessions is worth one caution beyond identity. Install the context, +not the capsule it exports: `ctx.with_logical_extension_codec(other_ctx)` carries +`other_ctx`'s session id as the identity, whereas +`ctx.with_logical_extension_codec(other_ctx.__datafusion_logical_extension_codec__())` +hands over a bare capsule and gets a random one that no other session can decode. +Either way the imported codecs resolve their task context against `other_ctx` and +stop working when it is dropped — see +[One session, one `Arc`](#one-session-one-arcsessioncontext) — so +this composes sessions, it does not copy codecs out of one. + +`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) + +# 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())) +``` + +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 @@ -347,7 +452,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 +466,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/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 0fa10d7f3..4dc6c2a6a 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -35,7 +35,13 @@ 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 appends the codec to the session's codec chain, and DataFusion's default codec handles whatever no installed codec claims. Every payload a codec writes is wrapped in an envelope naming that codec, and decoding consults exactly the codec that encoded it — so several independent plugin libraries can install codecs on the same session without any of them having to recognise or reject the others' payloads. The codecs here are written as ordinary `LogicalExtensionCodec` / `PhysicalExtensionCodec` implementations; the envelope is applied and stripped by `datafusion-python` and never reaches them. + +`MyLogicalExtensionCodec` takes an optional `provider_prefix` argument (`MyLogicalExtensionCodec(provider_prefix="TOKENAAA")`) that overrides the byte prefix it stamps on encoded table providers. It exists so the tests can install two instances that own disjoint slices of the wire format, which is what makes install ordering observable from Python. Two instances of one class 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: ```python ctx = ctx.with_logical_extension_codec(provider_logical_codec) @@ -45,4 +51,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..3147bcc51 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py +++ b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py @@ -17,8 +17,46 @@ from __future__ import annotations -from datafusion import LogicalPlan, SessionContext -from datafusion_ffi_example import MyLogicalExtensionCodec +import pyarrow as pa +import pytest +from datafusion import Expr, LogicalPlan, SessionContext, col, udf +from datafusion_ffi_example import ( + MyLogicalExtensionCodec, + MyTableProvider, + NameOnlyFunction, + NameOnlyUdfCodec, +) + + +def _double_udf(): + return udf( + lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]), + [pa.int64()], + pa.int64(), + volatility="immutable", + name="double", + ) + + +def _encode_provider_plan(token: str) -> tuple[bytes, MyLogicalExtensionCodec]: + """Serialize a plan over this library's table provider using a codec + that stamps `token` on the encoded provider. + + Returns the blob and the codec, so callers can assert on its call + counters. The token is chosen per test so a second codec installed + later is provably unable to claim these bytes. + + The codec is installed under ``token`` as its id as well, so a + caller can reinstall the same instance elsewhere and have the tag on + these bytes resolve. Identity would otherwise be derived from the + class, which every instance shares. + """ + codec = MyLogicalExtensionCodec(provider_prefix=token) + ctx = SessionContext().with_logical_extension_codec(codec, codec_id=token) + ctx.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) + assert token.encode() in blob + return blob, codec def _setup_session_with_codec() -> tuple[SessionContext, MyLogicalExtensionCodec]: @@ -33,8 +71,8 @@ def _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 @@ -80,3 +118,430 @@ 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 appends it to the + session's codec chain instead of replacing the first. The second + codec here (a default-backed codec exported from a fresh session) + cannot encode this library's table provider, so the first codec + still claims it. Under replace semantics this test fails with + `LogicalExtensionCodec is not provided`.""" + 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() + + +def test_first_installed_codec_encodes(): + """Encoding walks the chain in install order, so the earliest + installed codec that can claim an object gets it. + + Both orders run in one test on purpose. Asserting a single order + would also pass under replace semantics, where the second install + simply discards the first codec; swapping the order and getting the + other token proves the losing codec was still installed and merely + lost the race. + + The two instances need explicit ids: identity is otherwise derived + from the class, and these are two instances of one class owning + disjoint slices of the wire format. + """ + for winner, loser in (("TOKENAAA", "TOKENBBB"), ("TOKENBBB", "TOKENAAA")): + winner_codec = MyLogicalExtensionCodec(provider_prefix=winner) + loser_codec = MyLogicalExtensionCodec(provider_prefix=loser) + ctx = SessionContext().with_logical_extension_codec( + winner_codec, codec_id=winner + ) + ctx = ctx.with_logical_extension_codec(loser_codec, codec_id=loser) + + ctx.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) + + assert winner.encode() in blob + assert loser.encode() not in blob + assert winner_codec.table_provider_encode_calls() == 1 + assert loser_codec.table_provider_encode_calls() == 0 + + +def test_installing_a_codec_cannot_hijack_an_earlier_codecs_objects(): + """Appending is additive: a later install can claim objects nothing + else claimed, but never takes over an object an earlier codec was + already encoding. + + This is why install order is append rather than prepend. Under + prepend, adding an unrelated library would silently change how an + existing library's objects encode -- and, once payloads are tagged, + would renumber ids that older payloads already reference. + """ + first = MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + ctx = SessionContext().with_logical_extension_codec(first, codec_id="TOKENAAA") + ctx.register_table("numbers", MyTableProvider(1, 4, 1)) + before = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) + + later = MyLogicalExtensionCodec(provider_prefix="TOKENBBB") + ctx = ctx.with_logical_extension_codec(later, codec_id="TOKENBBB") + after = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx) + + # Provider tokens are minted per encode, so the payloads differ in the + # token id. What must not change is which codec claimed the provider. + assert b"TOKENAAA" in after + assert b"TOKENBBB" not in after + assert later.table_provider_encode_calls() == 0 + assert len(before) == len(after) + + +def test_decode_dispatches_to_the_codec_that_encoded(): + """A payload names the codec that wrote it, so decoding consults + exactly that codec and never offers the bytes to any other. + + The blob here is written by ``first`` in its own session. Installing + ``second`` alongside it must not put ``second`` anywhere near those + bytes -- under trial-and-error dispatch it would be asked first, and + a codec that decodes structurally similar protobuf would answer. + """ + blob, first = _encode_provider_plan("TOKENAAA") + + second = MyLogicalExtensionCodec(provider_prefix="TOKENBBB") + ctx = SessionContext().with_logical_extension_codec(first, codec_id="TOKENAAA") + ctx = ctx.with_logical_extension_codec(second, codec_id="TOKENBBB") + + restored = LogicalPlan.from_bytes(ctx, blob) + assert ctx.create_dataframe_from_logical_plan(restored).collect() + + assert first.table_provider_decode_calls() == 1 + assert second.table_provider_decode_calls() == 0 + + +def test_decode_survives_a_different_install_order(): + """Dispatch keys off codec identity, not chain position, so the + decoding session may install the same codecs in any order. + + This is the case positional dispatch cannot handle: the encoding + session has the owning codec at index 0 and the decoding session has + it at index 1. Keying on position would hand the payload to whatever + sits at index 0 in the decoder -- silently, and with a plausible + result. + """ + owner = MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + other = MyLogicalExtensionCodec(provider_prefix="TOKENBBB") + + encoder = SessionContext().with_logical_extension_codec(owner, codec_id="TOKENAAA") + encoder = encoder.with_logical_extension_codec(other, codec_id="TOKENBBB") + encoder.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder) + + # Same codecs, opposite order. + decoder = SessionContext().with_logical_extension_codec(other, codec_id="TOKENBBB") + decoder = decoder.with_logical_extension_codec(owner, codec_id="TOKENAAA") + + restored = LogicalPlan.from_bytes(decoder, blob) + assert decoder.create_dataframe_from_logical_plan(restored).collect() + assert other.table_provider_decode_calls() == 0 + + +def test_decode_names_the_codec_that_is_not_installed(): + """When the owning codec is absent the error names it and lists what + is installed, instead of reporting DataFusion's generic "not + provided" from whichever codec was tried last.""" + blob, _owner = _encode_provider_plan("TOKENBBB") + + ctx = SessionContext().with_logical_extension_codec( + MyLogicalExtensionCodec(provider_prefix="TOKENCCC"), codec_id="lib_c.Codec" + ) + + with pytest.raises(Exception, match="TOKENBBB") as excinfo: + LogicalPlan.from_bytes(ctx, blob) + + message = str(excinfo.value) + # Names the codec the payload belongs to, and what is actually here. + assert "not installed on this session" in message + assert "lib_c.Codec" in message + + +def test_a_codec_id_defaults_to_its_class_module_and_qualname(): + """A codec that declares no identity is named by its class. + + That string goes on the wire in front of every payload the codec + writes, so it is a compatibility surface: renaming the class or moving + the module changes it, and plans stored by an earlier release stop + decoding. A library that needs to rename declares + ``__datafusion_codec_id__`` instead, which the next test covers. + + Pinned twice on purpose -- once against the literal, so a rename has to + come here and be acknowledged, and once against the rule, so the + literal cannot drift into something the code no longer derives. + """ + ctx = SessionContext().with_logical_extension_codec(MyLogicalExtensionCodec()) + + assert ctx.logical_extension_codec_ids() == [ + "datafusion_ffi_example.MyLogicalExtensionCodec" + ] + assert ctx.logical_extension_codec_ids() == [ + f"{MyLogicalExtensionCodec.__module__}.{MyLogicalExtensionCodec.__qualname__}" + ] + + +class _CodecUnderItsOldName: + """A library codec that pins its identity, so the class can be renamed. + + Delegates the capsule getter to a real FFI codec. ``session`` is + forwarded, because that argument is how the underlying library reaches + the session the codec is being installed on. + """ + + __datafusion_codec_id__ = "pinned.example.Codec" + + def __init__(self, inner: MyLogicalExtensionCodec) -> None: + self._inner = inner + + def __datafusion_logical_extension_codec__(self, session: object = None) -> object: + return self._inner.__datafusion_logical_extension_codec__(session) + + +class _CodecUnderItsNewName(_CodecUnderItsOldName): + """The same codec after a rename. Same pinned id, different class.""" + + +def test_a_pinned_codec_id_survives_a_class_rename(): + """``__datafusion_codec_id__`` decouples identity from the class name, + which is the reason to declare one. + + A plan encoded by the codec under its old name decodes on a session + that only knows the new name. Under the class-derived default the two + would be different ids and the payload would be undecodable. + """ + assert _CodecUnderItsOldName.__qualname__ != _CodecUnderItsNewName.__qualname__ + + old = MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + encoder = SessionContext().with_logical_extension_codec(_CodecUnderItsOldName(old)) + assert encoder.logical_extension_codec_ids() == ["pinned.example.Codec"] + + encoder.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder) + + # The pinned id, not the class name, is what the payload carries. + assert b"pinned.example.Codec" in blob + assert b"_CodecUnderItsOldName" not in blob + + new = MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + decoder = SessionContext().with_logical_extension_codec(_CodecUnderItsNewName(new)) + assert decoder.logical_extension_codec_ids() == ["pinned.example.Codec"] + + restored = LogicalPlan.from_bytes(decoder, blob) + assert decoder.create_dataframe_from_logical_plan(restored).collect() + assert new.table_provider_decode_calls() == 1 + + +def test_installing_two_codecs_under_one_id_is_rejected(): + """Identity is derived from the class, so installing two instances of + one class collides. Rejecting at install time is the point: two + codecs sharing an id are indistinguishable when a payload is decoded, + and only the caller knows whether they write the same wire format.""" + ctx = SessionContext().with_logical_extension_codec(MyLogicalExtensionCodec()) + + with pytest.raises(ValueError, match="already installed"): + ctx.with_logical_extension_codec(MyLogicalExtensionCodec()) + + +def _encode_through_a_bare_capsule(token: str) -> tuple[bytes, list[str]]: + """Serialize a provider plan through a codec installed as a bare + capsule, which is the case with no derivable identity. + + Returns the blob and the encoding session's codec ids, so a caller + can compare them against another session's. + """ + exporter = SessionContext().with_logical_extension_codec( + MyLogicalExtensionCodec(provider_prefix=token) + ) + encoder = SessionContext().with_logical_extension_codec( + exporter.__datafusion_logical_extension_codec__() + ) + encoder.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder) + return blob, encoder.logical_extension_codec_ids() + + +def test_bare_capsule_codec_is_session_local(): + """A bare PyCapsule exposes nothing stable to derive an identity + from -- every capsule reports the same type -- so it is tagged with a + session-local id. A plan it encodes fails on an unrelated session + with an error naming the fix, rather than being decoded by whichever + codec happens to sit at the same position.""" + blob, _ids = _encode_through_a_bare_capsule("TOKENAAA") + + with pytest.raises(Exception, match="bare PyCapsule") as excinfo: + LogicalPlan.from_bytes(SessionContext(), blob) + assert "codec_id" in str(excinfo.value) + + +def test_a_bare_capsule_codec_id_is_not_re_mintable_by_another_session(): + """The id given to a capsule-installed codec must be one no other + session can arrive at. + + Both sessions here install exactly one bare capsule, so any identity + drawn from a namespace both sessions number the same way -- a + counter, a position in the chain -- collides, and the payload is + handed to the other library's codec. That failure is quiet: a codec + offered bytes it does not recognise falls through to its own inner + default codec, so the error names neither codec and no counter moves. + Asserting on the message is what separates the two schemes. + + The empty-chain case in the test above passes under either scheme, + because a lookup in an empty chain misses whatever the id is. + """ + blob, encoder_ids = _encode_through_a_bare_capsule("TOKENAAA") + + # An unrelated session, also holding exactly one bare capsule. + other = SessionContext().with_logical_extension_codec( + MyLogicalExtensionCodec(provider_prefix="TOKENBBB") + ) + decoder = SessionContext().with_logical_extension_codec( + other.__datafusion_logical_extension_codec__() + ) + decoder_ids = decoder.logical_extension_codec_ids() + assert len(encoder_ids) == len(decoder_ids) == 1 + assert encoder_ids != decoder_ids + + with pytest.raises(Exception, match="bare PyCapsule") as excinfo: + LogicalPlan.from_bytes(decoder, blob) + assert "codec_id" in str(excinfo.value) + + +def test_installing_a_session_as_a_codec_uses_a_per_session_id(): + """A context installed as a codec is identified by its session, not by + its class. + + Every ``SessionContext`` shares one class, so a class-derived id would + name all of them: two contexts could not coexist on one target, and a + payload written through one would resolve to the other on decode. The + sources are held in locals because the imported codecs resolve their + task context against them. + """ + src_a = SessionContext().with_logical_extension_codec( + MyLogicalExtensionCodec(provider_prefix="TOKENAAA") + ) + src_b = SessionContext().with_logical_extension_codec( + MyLogicalExtensionCodec(provider_prefix="TOKENBBB") + ) + assert src_a.__datafusion_codec_id__ != src_b.__datafusion_codec_id__ + + # Both sources compose onto one session, which a shared id would refuse. + both = SessionContext().with_logical_extension_codec(src_a) + both = both.with_logical_extension_codec(src_b) + assert both.logical_extension_codec_ids() == [ + src_a.__datafusion_codec_id__, + src_b.__datafusion_codec_id__, + ] + + # A payload written through one source does not resolve to the other. + encoder = SessionContext().with_logical_extension_codec(src_a) + encoder.register_table("numbers", MyTableProvider(1, 4, 1)) + blob = encoder.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(encoder) + + decoder = SessionContext().with_logical_extension_codec(src_b) + with pytest.raises(Exception, match="not installed on this session") as excinfo: + LogicalPlan.from_bytes(decoder, blob) + assert src_a.__datafusion_codec_id__ in str(excinfo.value) + + +def test_a_derived_handle_reports_its_session_id(): + """Handles derived from one session share its id, so only one of them + can be installed on a given target. Their payloads would be + indistinguishable on decode, so refusing is the right answer.""" + base = SessionContext() + derived = base.with_python_udf_inlining(enabled=False) + assert derived.__datafusion_codec_id__ == base.__datafusion_codec_id__ + + target = SessionContext().with_logical_extension_codec(base) + with pytest.raises(ValueError, match="already installed"): + target.with_logical_extension_codec(derived) + + +def test_name_only_codec_round_trips_without_a_payload(): + """A codec may own functions that need no payload: the name is the + whole encoding. ``try_encode_udf`` writes nothing, and the decoder + rebuilds the function from the name with no registry entry. + + DataFusion supports this directly -- an empty ``fun_definition`` + sends the decoder to the registry first and the codec second. This + test pins that arm from the Python side, because it is the one path + where a payload is still offered to every installed codec: there are + no bytes, so there is no identity to dispatch on. + + It is also the guard against a plausible "improvement". Wrapping + every chained encode in the identity envelope would make this + payload non-empty, which sets ``fun_definition`` and permanently + skips the registry lookup -- breaking both this codec and ordinary + by-name round trips, with nothing else in the suite noticing. + """ + codec = NameOnlyUdfCodec() + name = NameOnlyUdfCodec.function_name() + + # FROM-less, so serialization never reaches try_encode_table_provider -- + # this codec owns functions, not providers. + encoder = SessionContext().with_logical_extension_codec(codec) + encoder.register_udf(udf(NameOnlyFunction())) + blob = encoder.sql(f"SELECT {name}(1) AS x").logical_plan().to_bytes(encoder) + + # The name is the entire encoding, so the codec contributed no bytes + # and the payload carries no identity envelope for it. + assert codec.encode_udf_calls() > 0 + assert b"DFPYCHN" not in blob + + # A fresh session that never registered the function: only the codec + # can supply it, and only from the name. + decoder = SessionContext().with_logical_extension_codec(codec) + restored = LogicalPlan.from_bytes(decoder, blob) + + assert codec.decode_udf_calls() > 0 + assert decoder.create_dataframe_from_logical_plan(restored).collect() + + +def test_default_only_session_writes_no_envelope(): + """A session with no extension codecs installed produces the same + bytes as a build without codec chaining: the terminal codec writes + unframed, so the envelope only appears once a codec is installed. + + Keeps the wire break scoped to sessions that actually compose.""" + ctx = SessionContext() + blob = ctx.sql("SELECT abs(-1) AS x").logical_plan().to_bytes(ctx) + assert b"DFPYCHN" not in blob + + +def test_udf_inlining_setting_survives_codec_install(): + """Installing an extension codec must not silently re-enable inline + Python UDF encoding on a session that opted out. Regression guard in + both directions: the encoder still emits the by-name form, and the + decoder still refuses an inline payload. + + The codec installed here delegates UDF encoding to DataFusion's + default codec. A codec exported from another `SessionContext` would + not work as a probe: that export is itself a Python-aware codec with + inlining enabled, so the strict outer codec would delegate to it and + the inline payload would reappear. + """ + strict = SessionContext().with_python_udf_inlining(enabled=False) + extended = strict.with_logical_extension_codec( + MyLogicalExtensionCodec(provider_prefix="TOKENFFF") + ) + + e = _double_udf()(col("a")) + assert b"DFPYUDF" not in e.to_bytes(extended) + + inline_blob = e.to_bytes(SessionContext()) + assert b"DFPYUDF" in inline_blob + with pytest.raises(Exception, match="inlining is disabled"): + Expr.from_bytes(inline_blob, ctx=extended) diff --git a/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py index 28eaaf2d9..c7a6ede7b 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 appends to the chain instead + of replacing the first codec. The second codec here (default-backed + export from a fresh session) encodes UDFs by name without writing + bytes, which the chain treats as "no opinion" — so the user codec + installed first is still consulted. Under replace semantics its + counter stays at zero.""" + ctx, codec = _setup_session_with_codec() + ctx = ctx.with_physical_extension_codec( + SessionContext().__datafusion_physical_extension_codec__() + ) + + df = ctx.sql("SELECT abs(a) AS x FROM t") + original = df.execution_plan() + + before = codec.encode_udf_calls() + blob = original.to_bytes(ctx) + assert codec.encode_udf_calls() > before + + restored = ExecutionPlan.from_bytes(ctx, blob) + assert str(original) == str(restored) diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 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/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/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 72f96bb8e..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 pair is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against it, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). +The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and 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 — 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..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 @@ -667,3 +686,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..9cea74555 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 @@ -2235,6 +2236,32 @@ def __datafusion_task_context_provider__(self) -> Any: """Access the PyCapsule FFI_TaskContextProvider.""" return self.ctx.__datafusion_task_context_provider__() + @property + def __datafusion_codec_id__(self) -> str: + """Identity this context carries when installed as an extension codec. + + Installing one context's codec stack on another session tags the + payloads it writes with this string. It is derived from the session id + rather than from the class, because every context shares one class: + a class-derived identity would name them all, so two contexts installed + on one session would collide and a payload written through one would + resolve to the other when decoded. + + Contexts derived from the same session — including the ones returned by + :py:meth:`with_logical_extension_codec` and + :py:meth:`with_python_udf_inlining` — report the same id, so only one of + them can be installed on a given session. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx.__datafusion_codec_id__.startswith("session:") + True + >>> ctx.__datafusion_codec_id__ == SessionContext().__datafusion_codec_id__ + False + """ + return self.ctx.__datafusion_codec_id__ + def __datafusion_logical_extension_codec__(self, session: Any = None) -> Any: """Access the PyCapsule FFI_LogicalExtensionCodec. @@ -2253,26 +2280,86 @@ 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 specified codec. + """Create a new session context with an additional logical codec. Only FFI codecs are supported. Pass any object implementing ``__datafusion_logical_extension_codec__`` (see :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`). + Codecs compose: each call appends the codec 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 - 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. + + 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 + 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. @@ -2280,23 +2367,55 @@ 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 + self, + codec: PhysicalExtensionCodecExportable | _PyCapsule, + codec_id: str | None = None, ) -> SessionContext: - """Create a new session context with the specified physical codec. + """Create a new session context with an additional physical codec. Only FFI codecs are supported. Pass any object implementing ``__datafusion_physical_extension_codec__`` (see :py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`). + Codecs compose the same way as in + :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 - 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. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx = ctx.with_physical_extension_codec( + ... my_library.PhysicalCodec() + ... ) # doctest: +SKIP + + >>> ctx = ctx.with_physical_extension_codec( + ... capsule, codec_id="my_library.PhysicalCodec" + ... ) # doctest: +SKIP """ - new_internal = self.ctx.with_physical_extension_codec(codec) + new_internal = self.ctx.with_physical_extension_codec(codec, codec_id) new = SessionContext.__new__(SessionContext) new.ctx = new_internal return new diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index 43b53e469..70dd4958b 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -120,6 +120,18 @@ class LogicalExtensionCodecExportable(Protocol): is being installed on. Take the task context provider from it rather than building a session of your own, so the decode callbacks resolve names against the session that runs the query. + + Implement the codec itself exactly as you would for a session that installs + only yours. A session may hold several codecs, but each payload records the + codec that wrote it and is only ever handed back to that codec, so there is + no need to recognise or reject another library's payloads. + + An optional ``__datafusion_codec_id__`` attribute pins the identity a + payload records. It is not part of this protocol and is rarely needed: + identity otherwise comes from the class's module and qualified name, which + is already stable across processes. Declare it when a class rename must not + invalidate previously encoded plans, or when one library installs two + instances that own disjoint slices of the wire format. """ def __datafusion_logical_extension_codec__( # noqa: D105 @@ -130,7 +142,9 @@ def __datafusion_logical_extension_codec__( # noqa: D105 class PhysicalExtensionCodecExportable(Protocol): """Type hint for objects exposing ``__datafusion_physical_extension_codec__``. - See :py:class:`LogicalExtensionCodecExportable` for ``session``. + See :py:class:`LogicalExtensionCodecExportable` for ``session``, for why a + codec need not recognise other libraries' payloads, and for + ``__datafusion_codec_id__``. """ def __datafusion_physical_extension_codec__( # noqa: D105 diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py index 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