diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 294ebfb3a..1efc93ec5 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -62,6 +62,33 @@ library reaches things only the session has. session satisfies the protocol too — `ctx.__datafusion_query_planner__()` and `ctx.__datafusion_query_planner__(ctx)` are both valid. +`__datafusion_session_planner__(ctx, fallback)` is the exception to the shape +above: it takes a second argument, the planner assembled so far. A session has +one planner slot, so planners compose by nesting rather than by chaining, and +the host hands each bundle the previous layer instead of letting it capture one. +Wrap `fallback` and delegate to it; returning a planner that ignores it discards +every layer beneath, including one the session already had. It runs after every +bundle's codecs are installed, so `ctx` carries the final chains. + +That is also the only hook where it does. `__datafusion_session_extension__` +runs before anything is installed, so its `ctx` still carries the chains the +receiver had — the same session, and the same task-context provider, but not +this call's codecs, not even your own. Read the host's codec chains in the +planner hook, never in the extension hook. + +A *codec* must always be handed over as an object implementing its getter, never +as the bare capsule the getter returns; `with_extensions` refuses a capsule. +A codec's wire id — the string a payload names on decode, which has to mean the +same thing in whichever process decodes — is read off the object it arrives as, +and a capsule has no type to read one from. Deriving the id from the bundle that +contributed the capsule is not the fix: the bundle is whatever object the caller +passed, so an application packaging your library inside a bundle of its own +would re-tag your payloads and they would stop decoding where they are read. If +the object's class name is not the identity you want on the wire, declare +`__datafusion_codec_id__` on it. `BundledLogicalCodec` in +`examples/datafusion-ffi-query-planner-example/src/extension.rs` is the shape. +This applies only to codecs — a query planner carries no wire id. + ## Rule 3 — never construct a `SessionContext` in an extension library The FFI constructors ask for things a library does not have: @@ -154,8 +181,19 @@ guards this. Its `WHERE` clause is load-bearing: filter pushdown upgrades the weak handle during logical optimization, before plan serialization could fail first for an unrelated reason. +`SessionContext.with_extensions` is where this rule is easiest to get wrong, +because "bind the components to the context you are about to return" reads like +an instruction to derive one first. It is not: the factories are handed the +receiver, and the returned handle shares its allocation. There is nothing to +keep alive separately and nothing to garbage-collect out from under a provider. + `SessionContext.enable_url_table` is the one method that mints a second -allocation for a session. Its result must not outlive the receiver. +allocation for a session. Its result must not outlive the receiver, and it also +forks the session's `SessionState` while keeping its id, so two handles report +one `session_id()` with divergent configuration. That is a bug rather than a +design — tracked in +[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708) +— so do not cite it as precedent for deriving a replacement context. ## Rule 7 — installing a planner mutates the session, and says so diff --git a/Cargo.lock b/Cargo.lock index c7632732a..6a7f68438 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1209,6 +1209,7 @@ dependencies = [ "datafusion-catalog", "datafusion-common", "datafusion-ffi", + "datafusion-proto", "datafusion-python-util", "datafusion-session", "pyo3", diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 84182ff19..c711a62dc 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -66,7 +66,7 @@ use datafusion_python_util::{ }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; -use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; +use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; use url::Url; @@ -424,10 +424,12 @@ impl PySessionContext { pub fn enable_url_table(&self) -> PyResult { // Pre-existing caveat, unrelated to query planners: this is the one - // method that mints a second `Arc` for a session. Any - // weak `FFI_TaskContextProvider` handed out by the receiver stays bound - // to the receiver, so the returned context must not outlive it. See + // method that mints a second `Arc` for a session, and + // it also forks the session's state while keeping its id. Any weak + // `FFI_TaskContextProvider` handed out by the receiver stays bound to + // the receiver, so the returned context must not outlive it. See // `set_session_query_planner` for why everything else mutates in place. + // Tracked as a bug in . Ok(PySessionContext { ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()), logical_codec: Arc::clone(&self.logical_codec), @@ -1433,10 +1435,13 @@ impl PySessionContext { /// decode. See [`SESSION_CODEC_ID_PREFIX`]. /// /// Handles derived from one session — `with_python_udf_inlining`, - /// `with_logical_extension_codec` — report the same id even though their - /// codec chains differ, so installing two of them on one target is - /// refused. That is the intended answer: their payloads would be - /// indistinguishable on decode. + /// `with_logical_extension_codec`, [`Self::_install_extension_codecs`] — + /// report the same id even though their codec chains differ, so installing + /// two of them on one target is refused. That is the intended answer: they + /// share a `state_ref`, so their payloads would resolve against the same + /// session and are indistinguishable on decode. Every derivation shares the + /// session for exactly this reason; `enable_url_table` is the one that does + /// not, and it is tracked as a bug. #[getter] pub fn __datafusion_codec_id__(&self) -> String { format!("{SESSION_CODEC_ID_PREFIX}{}", self.ctx.session_id()) @@ -1608,6 +1613,118 @@ impl PySessionContext { derived.set_session_query_planner(None); derived } + + /// Build the codec chains for a `with_extensions` call. + /// + /// Private support method for `SessionContext.with_extensions`, and the + /// first of the two phases that method runs. `self` is the context the + /// extensions bound their components against, and is also the + /// `Arc` every FFI task-context provider they created + /// targets, so the returned handle shares it rather than deriving a new + /// one. + /// + /// **Writes nothing.** The codec chains belong to the returned handle + /// rather than to `SessionState`, so this phase is transactional for free: + /// a codec that fails to import, or that collides with an installed id, + /// leaves the caller's context exactly as it was. Binding the planner is + /// the only step that touches the session, and it is deferred to + /// [`Self::_install_extension_planner`] so the planner hooks can run + /// against the final chains. + /// + /// Codecs must arrive as objects exposing the capsule getter, never as + /// bare capsules — see [`resolve_bundle_codec_id`]. + pub fn _install_extension_codecs<'py>( + slf: &Bound<'py, Self>, + logical_codecs: Vec>, + physical_codecs: Vec>, + ) -> PyDataFusionResult { + // Chains are built as local values, so a codec that fails to import -- + // or that collides with an id already installed -- leaves the session + // untouched. Nothing is borrowed across a call back into Python. + let (mut logical_codec, mut physical_codec) = { + let this = slf.borrow(); + ( + this.logical_codec.as_ref().clone(), + this.physical_codec.as_ref().clone(), + ) + }; + + for codec in logical_codecs { + let id = resolve_bundle_codec_id( + &codec, + "__datafusion_logical_extension_codec__", + &logical_codec.codec_ids(), + )?; + let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; + let inner: Arc = (&inner_ffi).into(); + logical_codec = logical_codec.with_additional_codec(id, inner); + } + let logical_codec = Arc::new(logical_codec); + + for codec in physical_codecs { + let id = resolve_bundle_codec_id( + &codec, + "__datafusion_physical_extension_codec__", + &physical_codec.codec_ids(), + )?; + let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; + let inner: Arc = (&inner_ffi).into(); + physical_codec = physical_codec.with_additional_codec(id, inner); + } + let physical_codec = Arc::new(physical_codec); + + Ok(Self { + ctx: Arc::clone(&slf.borrow().ctx), + logical_codec, + physical_codec, + }) + } + + /// Re-export a planner a `__datafusion_session_planner__` hook returned as + /// a capsule, so the next hook in the chain receives one either way. + /// + /// A hook may hand back an object exposing `__datafusion_query_planner__` + /// or a raw capsule; the next hook wraps whatever it is given and should + /// not have to branch on which. Importing here also surfaces a malformed + /// planner at the hook that produced it rather than at the final install. + /// Writes nothing. + pub fn _export_query_planner<'py>( + slf: &Bound<'py, Self>, + planner: Bound<'py, PyAny>, + ) -> PyDataFusionResult> { + let ffi = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?; + Ok(create_query_planner_capsule(slf.py(), &ffi)?) + } + + /// Commit the query planner for a `with_extensions` call. + /// + /// The second phase, run once every codec is installed and every planner + /// hook has returned, so the planner is bound against the final chains. + /// This is the one call in `with_extensions` that writes to the session, + /// and it goes through this context's own `state_ref()`, so providers + /// bound to it stay valid. + /// + /// `None` means no bundle supplied a planner. That still rebuilds + /// whichever planner the session already holds against the new chains, + /// exactly as `with_logical_extension_codec` does, and writes nothing at + /// all if the session has no FFI planner to rebuild. + /// + /// The caller skips this step entirely when the call installed no codec + /// and no planner, the same way [`Self::with_python_udf_inlining`] returns + /// early for a no-op toggle: there is nothing to rebind against, and the + /// rebuild would drag a planner sitting on another handle's codecs onto + /// this one's. + #[pyo3(signature = (planner=None))] + pub fn _install_extension_planner<'py>( + slf: &Bound<'py, Self>, + planner: Option>, + ) -> PyDataFusionResult<()> { + let planner = planner + .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))) + .transpose()?; + slf.borrow().set_session_query_planner(planner); + Ok(()) + } } impl PySessionContext { @@ -1783,6 +1900,10 @@ impl PySessionContext { /// 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. +/// Reachable only from `with_logical_extension_codec` and +/// `with_physical_extension_codec`, where `codec_id=` is the way out; +/// `with_extensions` refuses bare capsules outright rather than naming them +/// after something that is not the codec. See [`resolve_bundle_codec_id`]. /// /// 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 @@ -1799,12 +1920,53 @@ fn resolve_codec_id( 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." + decoded. Give this one a distinct identity: declare \ + `__datafusion_codec_id__` on the object being installed, or pass `codec_id=` \ + if you are calling `with_logical_extension_codec` or \ + `with_physical_extension_codec` directly." ))); } Ok(id) } +/// Resolve the wire id for a codec contributed through `with_extensions`, +/// requiring an object that can name itself. +/// +/// `with_extensions` takes no `codec_id=`, so the only naming channels are the +/// ones [`derive_codec_id`] reads off the handed-over object: a declared +/// `__datafusion_codec_id__`, or its class's `module.QualName`. A bare capsule +/// has neither. Naming it after the bundle that contributed it looks like an +/// answer and is not one: the bundle is whatever object the caller passed to +/// `with_extensions`, so a bundle that wraps another library's bundle — the +/// natural way for an application to package several libraries as one — would +/// stamp its own identity onto the inner library's codecs and silently change +/// the wire format. The inner library cannot defend against that no matter what +/// it declares, and the mismatch does not surface until a plan fails to decode +/// in another process. +/// +/// So the capsule is refused here, where the author can fix it by wrapping it +/// in an object. Wrapping also decouples the codec's wire identity from the +/// bundle's Python class name, which is the whole point of +/// `__datafusion_codec_id__`. +fn resolve_bundle_codec_id( + codec: &Bound<'_, PyAny>, + getter: &str, + existing: &[&str], +) -> PyResult { + if codec.is_instance_of::() { + return Err(PyTypeError::new_err(format!( + "A codec contributed through `with_extensions` must be an object exposing \ + `{getter}`, not a bare PyCapsule. A capsule carries no type of its own, so \ + there is nothing to name the codec by, and a payload names its codec by id \ + when it is decoded — an id that has to mean the same thing in whichever \ + process decodes. Wrap the capsule in an object that exposes `{getter}` and, \ + if the class name is not the identity you want on the wire, declares \ + `__datafusion_codec_id__`." + ))); + } + resolve_codec_id(codec, None, existing) +} + fn derive_codec_id(codec: &Bound<'_, PyAny>, explicit: Option) -> PyResult { if let Some(id) = explicit { return Ok(id); diff --git a/docs/source/conf.py b/docs/source/conf.py index 22bace809..8b3dbd0dd 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -199,8 +199,10 @@ def setup(sphinx) -> None: "**": ["sidebar-globaltoc.html"], } -# tell myst_parser to auto-generate anchor links for headers h1, h2, h3 -myst_heading_anchors = 3 +# tell myst_parser to auto-generate anchor links for headers h1 through h4. +# h4 is included because the FFI guide cross-references its own `####` +# subsections; without an anchor those links render but resolve nowhere. +myst_heading_anchors = 4 # MyST extensions: # - tasklist: GitHub-style `- [x]` checkboxes diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index d86858a83..244505c2c 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -277,10 +277,15 @@ three cases: - **Two instances of one class.** Both get the same id, so the second install raises `ValueError`. Pass `codec_id=` to tell them apart. -- **A bare `PyCapsule`.** A capsule has no class to take a name from, so it gets an - id private to the session that installed it. Plans it encodes fail with a clear - error on any other session, rather than being decoded by the wrong codec. Pass +- **A bare `PyCapsule`.** A capsule has no class to take a name from, so installing + one through `with_logical_extension_codec` or `with_physical_extension_codec` gives + it an id private to the session that installed it; plans it encodes fail with a + clear error on any other session rather than being decoded by the wrong codec. Pass `codec_id=` if those plans have to cross sessions. + + `with_extensions` takes no `codec_id=`, so it refuses a bare capsule outright and + tells you to wrap it. See + [Extension bundles: `with_extensions`](#extension-bundles-with_extensions). - **A class you intend to rename.** The id follows the class name, so renaming stops older plans from decoding. Declare `__datafusion_codec_id__` on the exporting object to pin an id that survives the rename. @@ -343,6 +348,228 @@ The current FFI logical codec supports providers and UDFs but not arbitrary cust `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and local build commands. +### Extension bundles: `with_extensions` + +The chaining above works, but it makes the caller responsible for ordering: the codecs +have to be installed before the planner, because a planner is built against whatever +codec chains exist when it is installed, and a codec added afterwards rebinds it. Get +that wrong and the planner encodes through a chain that is missing a library. + +`SessionContext.with_extensions` removes the ordering question. An extension library +exposes a bundle object implementing one or both of two hooks: + +```python +class MyEngineExtension: + def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: + # Phase one. Create fresh components bound to `ctx` on every call. + return SessionExtensionComponents( + logical_extension_codecs=(self._make_logical_codec(ctx),), + physical_extension_codecs=(self._make_physical_codec(ctx),), + ) + + def __datafusion_session_planner__(self, ctx: SessionContext, fallback): + # Phase two. `ctx` now carries every bundle's codecs, and `fallback` is + # the planner built so far. Wrapping it is what makes this library + # compose with the other planners in the call. + return self._make_planner(ctx, fallback=fallback) +``` + +Implement whichever apply: a codec-only library defines the first, a library that ships +only an optimizing planner defines the second. + +```python +ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) +ctx.register_table("t", lib_a.TableProvider()) +ctx.register_udf(udf(lib_b.SomeUDF())) +``` + +#### Two phases, because codecs and planners compose differently + +A session chains **many** codecs and dispatches between them by id. Codecs therefore +just accumulate: order affects encoding only, and decoding always routes to the codec +that wrote the payload. A session holds exactly **one** query planner, so planners +cannot accumulate — they compose by *nesting*, each wrapping the one before it and +delegating to it for work it does not handle. + +So `with_extensions` runs every `__datafusion_session_extension__` and installs all the +codecs, and only then runs each `__datafusion_session_planner__`, in argument order, +handing each the planner built so far. Two consequences worth holding onto: + +- **Bundle order matters differently for each.** For planners it sets the nesting: the + last extension listed ends up outermost and is consulted first. For codecs it never + affects decoding, and affects encoding only when two codecs would claim the same node + — see [When codec order does matter](#when-codec-order-does-matter). +- **A planner is always built against the complete codec set**, including codecs from + bundles listed after it. This is what the low-level chaining cannot give you, and it + matters most for a nested planner: the rebuild that follows a later codec install + reaches only the outermost layer (see [Rebinding a planner's codecs is one level + deep](#rebinding-a-planners-codecs-is-one-level-deep)), so a fallback captured before + the codecs were complete would stay stale forever. + +The two hooks therefore see the same session through different chains. Both receive a +handle on the one session, so the task-context provider taken off either is the same and +stays valid — but the `ctx` in phase one still carries the chains the receiver had, since +nothing is installed yet, while the `ctx` in phase two carries every bundle's codecs. A +bundle that reads the host's codec chains — `MyPlannerExtension` does, to give its +planner the host's codecs rather than minting its own — must do that in the planner hook. +Reading them in phase one gets the chains from before the call, missing even the bundle's +own codecs. + +An extension that ignores `fallback` and returns an unrelated planner replaces every +layer beneath it, including any planner the session already had. That is legal — a +library that must be the only planner does it deliberately — but it is not composable, +and nothing detects it. Returning `None` contributes no planner and leaves `fallback` +in place. + +`None` is the no-op, not `fallback`. The capsule handed to the first bundle wraps the +session's planner for export, so returning it unchanged installs that planner as a +foreign one and every plan built afterwards crosses an FFI boundary it did not before. +A bundle that decides at runtime it has nothing to contribute returns `None`. + +Three libraries that each ship a planner therefore install like this, with the +outermost last: + +```python +ctx = SessionContext(config).with_extensions( + tables.Extension(), # codecs only + functions.Extension(), # codecs only + optimizer.Extension(), # planner, wrapping the session default + distributed.Extension(), # planner, wrapping the optimizer +) +``` + +#### When codec order does matter + +Decoding is never order-dependent: a payload names its codec by id and the chain +dispatches straight to it. Encoding walks the chain in install order and stops at the +first codec that claims the node. Most of the time that is invisible, because libraries +claim disjoint things — one owns its table providers, another its UDFs, a third its own +execution plan nodes. + +It stops being invisible when a codec claims *broadly*. A node that came from another +library arrives as an opaque `ForeignExecutionPlan`, and a codec that claims any of +those will take nodes it does not own from any library installed after it. The query +still succeeds. What changes is which library wrote the bytes — so a plan that has to +decode in another process now needs whichever library happened to win, not the one whose +node it is. `MyPhysicalExtensionCodec` in the provider example claims this way, and +`test_a_greedy_codec_installed_first_claims_another_librarys_node` pins the consequence. + +Two rules of thumb: + +- **Writing a codec, claim narrowly.** Downcast to your own types. Claiming a broad + category makes your library order-sensitive for everyone downstream of it. +- **Shipping plans out of the process, verify.** Do not assume your node reached your + codec just because both are installed. Round-trip a plan through + `ExecutionPlan.to_bytes` / `from_bytes` in a test and assert your codec did the work. + +#### When the two orders conflict + +Because codec position and planner position both come from one argument list, a library +can in principle need to be early for one and late for the other: its codec must precede +a broad claimer, while its planner must nest outside that library's planner. + +Do not try to satisfy both by reordering — contribute each half at its own position. The +two hooks are independent, so a three-line adapter each is enough: + +```python +class CodecsOf: + """Contribute only the codec half of a bundle, at this position.""" + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_extension__(self, ctx): + return self.inner.__datafusion_session_extension__(ctx) + + +class PlannerOf: + """Contribute only the planner half of a bundle, at this position.""" + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_planner__(self, ctx, fallback): + return self.inner.__datafusion_session_planner__(ctx, fallback) + + +ctx = SessionContext(config).with_extensions( + CodecsOf(engine), CodecsOf(tables), # engine's codec first + PlannerOf(tables), PlannerOf(engine), # engine's planner outermost +) +``` + +This keeps everything `with_extensions` guarantees: one transaction, codecs complete +before any planner is built, codec ids untouched — an id is read off the codec object, +not off the extension that contributed it, so splitting a bundle cannot re-tag its +payloads. A library that expects to be composed this way should expose the halves itself +rather than make callers write the adapters. + +Falling back to the low-level `with_logical_extension_codec` / +`with_physical_extension_codec` / `set_query_planner` sequence also works, and it is the +right answer when the pieces do not come as bundles at all. But it is a real downgrade, +not just a more verbose spelling: you take back responsibility for installing every +codec before every planner, and a planner you layer by hand keeps the codecs it captured +— the [one-level rebind](#rebinding-a-planners-codecs-is-one-level-deep) does not reach +inside it. Reach for it last. + +There is no attempt here to make every permutation expressible from one call. Two +positions per bundle covers the cases that arise; anything stranger is a sign the +libraries disagree about what they own, which is better fixed there. + +#### Codecs are objects, not capsules + +`with_extensions` requires each codec to be an object exposing the capsule getter, and +refuses a bare `PyCapsule`. A codec's id is read off the object it is handed over as, +and a capsule has no type to read one from; since this method takes no `codec_id=`, +there would be nothing left to name it by. A library holding a raw capsule — which is +what a Rust implementation has — wraps it: + +```python +class MyLogicalCodec: + # Optional. Without it the id is this class's import path, which is already + # stable; declare it if you may rename the class and need old plans to decode. + __datafusion_codec_id__ = "my_library.logical.v1" + + def __init__(self, capsule): + self._capsule = capsule + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule +``` + +Wrapping is not just bookkeeping. It ties the id to the codec rather than to the bundle +that contributed it, and that difference is load-bearing: an application commonly +presents several libraries as one bundle of its own, and the id has to survive that. +Were the id taken from the contributing bundle, wrapping `my_engine.Extension` inside +`my_app.Extension` would silently re-tag the engine's payloads, and a scheduler that +installs the engine's codec by its documented id would fail to decode plans from +composed clients while succeeding for direct ones. The wrapper travels with the codec; +the bundle does not. + +The query planner is exempt — it carries no wire id, so it may be an object or a +capsule. + +Nothing is written to the session until every factory has returned and every capsule +has been validated, so a factory that raises leaves the session exactly as it was. A +factory that mutates the context it is handed — registering a table, say — is not +rolled back, which is why bundle objects must be configuration-only: create fresh +components on each call, never cache bound components, and do not retain the context +passed in. + +Like every other derivation, the returned context is a handle on the *same* session as +the receiver — see [What a derived context shares](#what-a-derived-context-shares). +Only the Python-side codec chains belong to the returned handle; the planner is +installed on the shared session and takes effect even if that handle is discarded. + +The session owns every installed component's task-context provider, and dependent +objects do not extend its lifetime. A `DataFrame`, logical plan, or capsule can outlive +every context on the session, but any operation that reaches an FFI codec after the +last one is collected fails with `TaskContextProvider went out of scope over FFI +boundary`. Keep a context alive for as long as objects derived from it are in use. + +`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust +implementation of the protocol, including taking the task-context provider off the +supplied context, wrapping its codecs in `BundledLogicalCodec` / `BundledPhysicalCodec` +so they carry declared ids, and constructing a Python `SessionExtensionComponents`. + ### Capsule getters receive the session they are installed on `__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and @@ -375,6 +602,14 @@ the session, and never touches a provider directly. That also matches what insta does anyway: `set_query_planner` builds the planner against the codecs of the session that will run the query. +Duck-type `session`, and do not check its type. It is the PyO3 context the binding +installs through, not the `datafusion.context.SessionContext` wrapper, so it carries +every capsule getter and `__datafusion_codec_id__` — everything the protocol asks of it +— but `isinstance(session, SessionContext)` is `False` in Python even though its `repr` +reads `datafusion.SessionContext`. The two bundle hooks +`__datafusion_session_extension__` and `__datafusion_session_planner__` are the +exception: `with_extensions` dispatches them from Python and hands them the wrapper. + `SessionContext` accepts the argument on all three getters and ignores it, so a session satisfies the same protocol an extension library implements. When you export the current planner to wrap it, `ctx.__datafusion_query_planner__()` and @@ -419,15 +654,24 @@ registered straight back into that same session, which would close the cycle `SessionContext.enable_url_table` is the one exception. It clones the underlying `SessionContext`, so the returned context has an allocation of its own and must not -outlive the receiver. +outlive the receiver. It also forks the session's state while keeping its id, so two +handles report one `session_id()` with divergent configuration. That is a bug rather +than a design, tracked in +[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708); +do not copy the pattern. ### What a derived context shares -`with_logical_extension_codec`, `with_physical_extension_codec`, and -`with_python_udf_inlining` return a new `SessionContext` wrapping the *same* underlying -session. Only the Python-side codec settings differ; catalogs, tables, registered -functions, and configuration are the one shared session, so a registration on either -side is visible to both. +`with_logical_extension_codec`, `with_physical_extension_codec`, +`with_python_udf_inlining`, and `with_extensions` return a new `SessionContext` wrapping +the *same* underlying session. Only the Python-side codec settings differ; catalogs, +tables, registered functions, and configuration are the one shared session, so a +registration on either side is visible to both. + +There is one `Arc` per session, which is what makes the weak +`FFI_TaskContextProvider` scheme work: a component bound through any handle stays valid +while *any* handle on that session is alive, so there is no way to bind a component to +an intermediate handle and have it dangle when that handle is dropped. `set_query_planner` does not return anything. The query planner lives in `SessionState`, so it is a property of the session rather than of a handle on it, and installing one is @@ -516,6 +760,12 @@ the original handle rebinds the session's planner back to the original handle's instead, which is the trap `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` pins. +`with_extensions` sidesteps this entirely, and for nested planners too: every codec from +every bundle is installed before the first planner hook runs, so no layer — outer or +fallback — is ever captured against a partial chain. There is no "afterwards" within a +call. Prefer it over hand-layering whenever the planners you are composing all ship as +bundles. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 257749c3a..4bd42c192 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -83,6 +83,16 @@ way `add_physical_optimizer_rule` does and returns nothing — the query planner lives in `SessionState`, so it belongs to the session rather than to a particular handle on it. See the {ref}`ffi` guide for the full protocol. +If a library ships codecs *and* a planner, prefer +`SessionContext.with_extensions(bundle)` over installing each piece by hand. It +installs every codec before it binds any planner, so a planner cannot end up +carrying a chain that a later `with_logical_extension_codec` call has grown. +The library exposes a bundle object implementing +`__datafusion_session_extension__` for its codecs and +`__datafusion_session_planner__` for its planner — the latter is handed the +planner installed so far, so several libraries that each ship one nest instead +of displacing each other. See the {ref}`ffi` guide. + ### Mismatched extension libraries now fail loudly Objects imported through the capsule protocol are checked against the major diff --git a/examples/datafusion-ffi-query-planner-example/Cargo.toml b/examples/datafusion-ffi-query-planner-example/Cargo.toml index 4d02c69f1..263f034b8 100644 --- a/examples/datafusion-ffi-query-planner-example/Cargo.toml +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -31,6 +31,7 @@ datafusion = { workspace = true } datafusion-catalog = { workspace = true, default-features = false } datafusion-common = { workspace = true, default-features = false } datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } datafusion-session = { workspace = true } async-trait = { workspace = true } datafusion-python-util.workspace = true diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 66bc45196..4b9140b53 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -41,7 +41,30 @@ uv run pytest \ examples/datafusion-ffi-query-planner-example/python/tests/_test*.py ``` -The integration test follows this setup: +The preferred setup uses `SessionContext.with_extensions` with extension bundles: + +```python +config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) +ctx = SessionContext(config).with_extensions(provider_bundle, MyPlannerExtension()) +ctx.register_table("numbers", provider) +ctx.register_udf(provider_udf) +``` + +`MyPlannerExtension` implements both extension hooks. `__datafusion_session_extension__` +receives the session it is being installed on, binds fresh codecs to that session's +task-context provider, and returns them as `SessionExtensionComponents`. +`__datafusion_session_planner__` then runs in the host's second phase, after every +bundle's codecs are installed, and builds a planner that delegates to the `fallback` it +is handed — so several libraries that each ship a planner nest instead of displacing one +another, and no planner is left carrying a chain that has since grown. + +Its codecs are handed over as `BundledLogicalCodec` and `BundledPhysicalCodec` rather +than as bare capsules. `with_extensions` requires an object, because a codec's wire id +is read off the object it arrives as and a capsule has no type to read one from. Each +wrapper declares `__datafusion_codec_id__`, so the id belongs to this library and does +not change when the bundle is nested inside an application's own bundle. + +The integration tests also cover the low-level chaining setup: ```python config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) 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 c6ef2072a..80ba7c724 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 @@ -17,11 +17,24 @@ from __future__ import annotations +import doctest import gc +import inspect +import io +import sys +import types import pyarrow as pa import pytest -from datafusion import Expr, SessionConfig, SessionContext, col, udf +from datafusion import ( + Expr, + SessionConfig, + SessionContext, + SessionExtensionComponents, + col, + udf, +) +from datafusion.plan import ExecutionPlan from datafusion_ffi_example import ( IsNullUDF, MyCatalogProvider, @@ -30,7 +43,11 @@ MyPhysicalOptimizerRule, MyTableProvider, ) -from datafusion_ffi_query_planner_example import MyPlannerConfig, MyQueryPlanner +from datafusion_ffi_query_planner_example import ( + MyPlannerConfig, + MyPlannerExtension, + MyQueryPlanner, +) def configured_context(max_rows: int): @@ -688,6 +705,614 @@ def test_query_planner_rejects_invalid_config(max_rows: str): ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() +# Ids `MyPlannerExtension`'s codec wrappers declare, mirroring +# `LOGICAL_CODEC_ID` / `PHYSICAL_CODEC_ID` in the crate's `extension.rs`. A +# scheduler decoding this library's plans installs codecs under these names, so +# they are part of its wire format rather than an implementation detail. +LOGICAL_CODEC_ID = "datafusion_ffi_query_planner_example.logical.v1" +PHYSICAL_CODEC_ID = "datafusion_ffi_query_planner_example.physical.v1" + + +class ProviderCodecsExtension: + """Bundles the provider library's codecs for ``with_extensions``. + + These codecs keep their own private task-context provider, so they only + need to be created once; the bundle can hand out the same exporters on + every call. + """ + + def __init__(self) -> None: + self.logical_codec = MyLogicalExtensionCodec() + self.physical_codec = MyPhysicalExtensionCodec() + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(self.logical_codec,), + physical_extension_codecs=(self.physical_codec,), + ) + + +class _NamedCodec: + """Forwards a codec's capsule getters under a declared id. + + ``with_extensions`` takes no ``codec_id=``, so an extension that ships a + codec class another extension also ships declares + ``__datafusion_codec_id__`` on the object it hands over. Both getters are + forwarded because one wrapper stands in for whichever kind it wraps. + """ + + def __init__(self, codec: object, codec_id: str) -> None: + self._codec = codec + self.__datafusion_codec_id__ = codec_id + + def __datafusion_logical_extension_codec__(self, session: object = None) -> object: + return self._codec.__datafusion_logical_extension_codec__(session) + + def __datafusion_physical_extension_codec__(self, session: object = None) -> object: + return self._codec.__datafusion_physical_extension_codec__(session) + + +class IdentifiedProviderCodecsExtension(ProviderCodecsExtension): + """``ProviderCodecsExtension`` whose codecs carry ids of their own.""" + + def __init__(self, prefix: str) -> None: + super().__init__() + self.logical = _NamedCodec(self.logical_codec, f"{prefix}.logical") + self.physical = _NamedCodec(self.physical_codec, f"{prefix}.physical") + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(self.logical,), + physical_extension_codecs=(self.physical,), + ) + + +def test_with_extensions_three_library_query(): + """One with_extensions call installs provider codecs and a planner bundle, + and a real non-empty plan flows across the three libraries.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + provider_ext = ProviderCodecsExtension() + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(provider_ext, planner_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert batches[0].column(1).to_pylist() == [False, False, False] + assert planner_ext.plan_calls() >= 1 + assert planner_ext.last_max_rows() == 3 + assert planner_ext.foreign_session_observed() + assert planner_ext.foreign_provider_observed() + assert planner_ext.foreign_plan_observed() + assert provider_ext.logical_codec.table_provider_encode_calls() > 0 + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_encode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +def test_with_extensions_names_a_rust_bundles_codecs_by_their_declared_id(): + """A Rust bundle wraps each codec in an object that declares its own id. + + This is the identity that has to survive leaving the process: a plan a + distributed engine writes here is decoded by its scheduler, which installs + a codec under the same id. A session-private random id — what a bare + capsule would get if `with_extensions` accepted one — would make the plan + undecodable there, which is why bare capsules are refused. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + + assert LOGICAL_CODEC_ID in ctx.logical_extension_codec_ids() + assert PHYSICAL_CODEC_ID in ctx.physical_extension_codec_ids() + + # The provider bundle's codecs declare no id, so they fall back to their + # own class names — never to the bundle's. + assert ( + "datafusion_ffi_example.MyLogicalExtensionCodec" + in ctx.logical_extension_codec_ids() + ) + assert not any( + codec_id.startswith("anon:") for codec_id in ctx.logical_extension_codec_ids() + ) + + +def test_with_extensions_rejects_a_rust_bundles_bare_capsule(): + """A bundle handing over a raw capsule is refused, with the fix named. + + This is the shape a Rust library reaches for first — `MyPlannerExtension` + wraps its capsules precisely to avoid it. + """ + + class BareCapsuleExtension: + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=( + ctx.__datafusion_logical_extension_codec__(), + ), + ) + + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + with pytest.raises(TypeError, match="must be an object exposing"): + SessionContext(config).with_extensions(BareCapsuleExtension()) + + +def test_bundle_codec_carries_its_own_planners_node(): + """The two halves of a bundle meet: its codec serializes its planner's node. + + ``DistributedQueryPlanner`` emits a ``DistributedExec``, a type private to + this library. No other codec in the session knows it, so the planner is + only useful alongside the codec that carries it — which is why the two ship + as one bundle, and why ``with_extensions`` installs every codec before it + binds any planner. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(bundle, ProviderCodecsExtension()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert bundle.distributed_exec_encode_calls() > 0 + assert bundle.distributed_exec_decode_calls() > 0 + # Reaching the session config from inside those decode callbacks is what + # shows the provider bound at installation resolves against this session. + assert bundle.decode_max_rows_seen() == [2] * len(bundle.decode_max_rows_seen()) + assert bundle.decode_max_rows_seen() + + +def test_bundle_planners_node_survives_a_plan_round_trip(): + """A plan carrying the node serializes and comes back intact. + + This is the path a distributed engine takes to ship a plan to a remote + executor, and the reason its node's id has to mean the same thing there. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(bundle, ProviderCodecsExtension()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + plan = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').execution_plan() + assert "DistributedExec" in plan.display() + + before = bundle.distributed_exec_encode_calls() + restored = ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx)) + + assert bundle.distributed_exec_encode_calls() > before + assert "DistributedExec" in restored.display() + + +def test_a_greedy_codec_installed_first_claims_another_librarys_node(): + """Encode order decides *which* library serializes a node. + + Decoding routes by id, so codec order never affects it. Encoding walks the + chain in install order and stops at the first codec that claims the node — + and a codec may claim broadly. ``MyPhysicalExtensionCodec`` claims any + ``ForeignExecutionPlan``, which is what a node from another library looks + like once it crosses the boundary, so installing it ahead of this bundle + takes the bundle's own node away from it. + + Nothing detects this. A library whose plans must decode elsewhere should + not assume its node reached its own codec just because both are installed. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + # Provider codecs first, so their broad claim wins. + ctx = SessionContext(config).with_extensions(ProviderCodecsExtension(), bundle) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + # The query still runs — the node was carried, just not by its own library. + assert bundle.distributed_exec_encode_calls() == 0 + + +class CodecsOf: + """Contribute only the codec half of a bundle, at this position.""" + + def __init__(self, inner: object) -> None: + self.inner = inner + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return self.inner.__datafusion_session_extension__(ctx) + + +class PlannerOf: + """Contribute only the planner half of a bundle, at this position.""" + + def __init__(self, inner: object) -> None: + self.inner = inner + + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + return self.inner.__datafusion_session_planner__(ctx, fallback) + + +def test_splitting_a_bundle_resolves_conflicting_orders(): + """A bundle can take one position for its codec and another for its planner. + + Codec position and planner position both come from one argument list, so a + library can need to be early for one and late for the other: here the + bundle's codec must precede the provider's broad claim, while its planner + must stay outermost. Splitting the halves satisfies both without giving up + what ``with_extensions`` guarantees. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + provider = ProviderCodecsExtension() + ctx = SessionContext(config).with_extensions( + CodecsOf(bundle), + CodecsOf(provider), + PlannerOf(bundle), + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + # Its codec ran ahead of the provider's broad claim, so the bundle kept its + # own node... + assert bundle.distributed_exec_encode_calls() > 0 + # ...and splitting did not re-tag anything: an id is read off the codec + # object, never off the extension that contributed it. + assert PHYSICAL_CODEC_ID in ctx.physical_extension_codec_ids() + assert ( + "datafusion_ffi_example.MyPhysicalExtensionCodec" + in ctx.physical_extension_codec_ids() + ) + + +class PlannerOnlyExtension: + """A library that ships a planner and no codecs. + + Implements only the planner hook — there is nothing to contribute in phase + one, and the protocol should not make it say so. + """ + + def __init__(self) -> None: + self.planner = None + + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + self.planner = MyQueryPlanner(fallback=fallback) + return self.planner + + +def test_with_extensions_nests_planners_in_argument_order(): + """Two planner-shipping libraries compose instead of displacing each other. + + This is the four-library case: A and C contribute codecs, B an optimizing + planner, D a distributed one that should sit outside B. Both planners run + for one query, which is only possible if D delegates to B rather than + replacing it — a session holds exactly one planner, so the nesting is the + only way both are reachable. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + codecs = ProviderCodecsExtension() + inner = PlannerOnlyExtension() + outer = PlannerOnlyExtension() + ctx = SessionContext(config).with_extensions(codecs, inner, outer) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + + assert inner.planner.plan_calls() >= 1 + assert outer.planner.plan_calls() >= 1 + # The last extension listed is outermost, so it is the one that had to + # delegate. The inner planner is the fallback, and reaches the session's + # original planner through its own. + assert outer.planner.used_fallback() + + +def test_with_extensions_planner_sees_every_bundles_codecs(): + """Phase two runs after phase one, for every bundle. + + A planner contributed by an early argument is still built against codecs a + later argument installed — the ordering trap that chaining the low-level + methods by hand leaves to the caller. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + planner_first = PlannerOnlyExtension() + codecs_last = ProviderCodecsExtension() + ctx = SessionContext(config).with_extensions(planner_first, codecs_last) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + # The provider codecs were installed after the planner was listed, and the + # query still round-trips its table provider through them. + assert codecs_last.logical_codec.table_provider_decode_calls() > 0 + + +def test_with_extensions_codec_ids_survive_bundle_composition(): + """Nesting a bundle inside another does not re-tag its codecs. + + An application that presents several libraries as one bundle is the + natural shape, and it must not change what the inner libraries write on + the wire — a scheduler installing `MyPlannerExtension`'s codec by id has + no idea which application wrapper the client used. Reading the id off the + handed-over object rather than off the contributing bundle is what makes + that hold. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + direct = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + composed = SessionContext(config).with_extensions( + _DocstringExampleExtension("scheduler:50050") + ) + + assert composed.logical_extension_codec_ids() == ( + direct.logical_extension_codec_ids() + ) + assert composed.physical_extension_codec_ids() == ( + direct.physical_extension_codec_ids() + ) + assert LOGICAL_CODEC_ID in composed.logical_extension_codec_ids() + + +def test_with_extensions_shares_the_session_with_the_source(): + """``with_extensions`` returns a handle on the source's session, and the + bundle's task-context provider resolves against that one session. + + There is one ``Arc`` per session, so a component bound + during installation cannot be left pointing at a handle that is dropped + later. A `SET` issued through the *source* after installation is therefore + visible to the provider the bundle bound, which is what a codec's decode + callback resolves through. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + source = SessionContext(config) + planner_ext = MyPlannerExtension() + result = source.with_extensions(ProviderCodecsExtension(), planner_ext) + + assert result.session_id() == source.session_id() + + # Registrations and config changes go through the source handle only. + source.register_table("numbers", MyTableProvider(1, 6, 1)) + source.sql("SET ffi_query_planner.max_rows = 2").collect() + + batches = result.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + assert planner_ext.max_rows_through_provider() == 2 + + # Symmetrically, the codec chains installed on the shared session are in + # force for the source handle too. + batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_with_extensions_survives_dropping_source_and_bundles(): + """The returned handle alone keeps the installed components alive. + + The context ``with_extensions`` was called on is a temporary here, and the + bundle objects are dropped with it. Both share their allocation with the + returned handle, so the components' task-context provider stays valid. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + gc.collect() + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_with_extensions_sees_state_changes_after_install(): + """Tables, UDFs, and config changes made after installation are visible + to the planner and to provider callbacks.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=4)) + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(ProviderCodecsExtension(), planner_ext) + + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + ctx.sql("SET ffi_query_planner.max_rows = 2").collect() + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + assert planner_ext.max_rows_through_provider() == 2 + + +def test_with_extensions_bundle_is_reusable(): + """Installing the same bundle into two contexts binds fresh components to + each destination.""" + planner_ext = MyPlannerExtension() + + config_a = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx_a = SessionContext(config_a).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_a.register_table("numbers", MyTableProvider(1, 6, 1)) + + config_b = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + ctx_b = SessionContext(config_b).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_b.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx_a.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + + batches = ctx_b.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert planner_ext.last_max_rows() == 3 + + +def test_with_extensions_failure_leaves_source_usable(): + """A failing factory after a successful one leaves the source context + fully functional.""" + + class BoomExtension: + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + msg = "boom" + raise RuntimeError(msg) + + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + + with pytest.raises(RuntimeError, match="boom"): + source.with_extensions(MyPlannerExtension(), BoomExtension()) + + # No planner was installed, so the default planner runs unrestricted. + batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] + + +def test_with_extensions_rebinds_existing_planner(): + """Codec-only bundles installed on a context that already has an FFI + planner rebind that planner to the new codec chains.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + planner = MyQueryPlanner() + ctx = SessionContext(config) + ctx.set_query_planner(planner) + provider_ext = ProviderCodecsExtension() + ctx = ctx.with_extensions(provider_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.last_max_rows() == 2 + # The planner only sees these codecs if it was rebound to the chains + # built during with_extensions. + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +class NoOpExtension: + """A bundle that turns out to contribute nothing. + + A plugin that finds no work to do -- an engine pointed at no scheduler, a + codec pack for a feature the session did not enable -- still gets listed, + and both its hooks answer empty rather than the caller having to filter it + out. + """ + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents() + + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + return None + + +@pytest.mark.parametrize("bundles", [(), (NoOpExtension(),)], ids=["empty", "no_op"]) +def test_with_extensions_installing_nothing_leaves_the_planner_alone(bundles): + """A call that installs nothing must not rebind the session's planner. + + The sibling of ``test_an_unchanged_inlining_setting_leaves_the_planner_alone``, + and the same hazard: committing the planner rebuilds ``SessionState`` to + rebind an existing FFI planner to *this handle's* chains. When no codec was + installed there is nothing to rebind against, so the rebuild buys nothing + and can only do harm. + + Observable only once the planner holds some *other* handle's codec, which + is what the discarded ``with_logical_extension_codec`` below arranges. + Without the guard the no-op call drags the planner back onto ``ctx``'s + codecs -- and ``ctx`` has no logical codec, so the planner is left with an + empty chain and the query fails outright instead of quietly using the wrong + codec. + """ + ctx, _physical_codec = physical_only_context() + ctx.set_query_planner(MyQueryPlanner()) + + planner_codec = MyLogicalExtensionCodec() + ctx.with_logical_extension_codec(planner_codec) # discarded; planner keeps it + gc.collect() + + ctx.with_extensions(*bundles) + gc.collect() + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner_codec.table_provider_encode_calls() > 0 + + +def test_with_extensions_rejects_two_bundles_of_the_same_codec_class(): + """Two bundles contributing the same codec class collide on id. + + Ids are derived from the exporting class, so two instances of one class + claim the same id. A payload names its codec by id when it is decoded, so + the ambiguity is refused at install time rather than resolved by position. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + with pytest.raises(ValueError, match="is already installed on this session"): + SessionContext(config).with_extensions( + ProviderCodecsExtension(), ProviderCodecsExtension(), MyPlannerExtension() + ) + + +def test_with_extensions_accepts_distinct_codec_ids(): + """Declaring ``__datafusion_codec_id__`` resolves the collision above. + + Both codec pairs then install, and the query still runs end to end: only + the codec that wrote a payload is asked to decode it, so the second pair + is simply never consulted. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ext_a = ProviderCodecsExtension() + ext_b = IdentifiedProviderCodecsExtension("second") + ctx = SessionContext(config).with_extensions(ext_a, ext_b, MyPlannerExtension()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + ids = ctx.logical_extension_codec_ids() + assert "datafusion_ffi_example.MyLogicalExtensionCodec" in ids + assert "second.logical" in ids + + # The first pair wrote the payloads, so decoding routes back to it alone. + assert ext_a.logical_codec.table_provider_encode_calls() > 0 + assert ext_a.logical_codec.table_provider_decode_calls() > 0 + assert ext_b.logical_codec.table_provider_decode_calls() == 0 + + +def test_dataframe_outliving_context_fails_cleanly(): + """A DataFrame does not keep its SessionContext alive. FFI components + resolve the task context through a weak reference, so using the + DataFrame after dropping the context raises a clean error instead of + crashing. This locks in the documented ownership contract: the context + must outlive DataFrames that depend on FFI codecs.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + df = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"') + del ctx + gc.collect() + + with pytest.raises(Exception, match="went out of scope"): + df.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. @@ -714,3 +1339,87 @@ def test_composed_codecs_with_query_planner(): assert logical_codec.table_provider_encode_calls() > 0 assert logical_codec.table_provider_decode_calls() > 0 assert physical_codec.execution_plan_decode_calls() > 0 + + +class _DocstringExampleExtension: + """Stand-in for the ``my_extension`` bundle named in the docstring. + + The docstring shows a single engine bundle taking a scheduler address, + which is what a real distributed engine ships: one object contributing a + planner *and* the codecs that carry its plans. Here that is assembled from + this repository's two example libraries. The address is accepted and + ignored; everything else the example touches is the real API. + """ + + def __init__(self, endpoint: str) -> None: + self.endpoint = endpoint + self._codecs = ProviderCodecsExtension() + self._planner = MyPlannerExtension() + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + codecs = self._codecs.__datafusion_session_extension__(ctx) + planner = self._planner.__datafusion_session_extension__(ctx) + return SessionExtensionComponents( + logical_extension_codecs=( + *codecs.logical_extension_codecs, + *planner.logical_extension_codecs, + ), + physical_extension_codecs=( + *codecs.physical_extension_codecs, + *planner.physical_extension_codecs, + ), + ) + + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + return self._planner.__datafusion_session_planner__(ctx, fallback) + + +def test_with_extensions_docstring_example_still_runs(): + """Run the ``with_extensions`` docstring example verbatim. + + The example is marked ``+SKIP`` because the main suite has no built FFI + extension to import, which is exactly how such an example rots. Here the + statements are parsed out of the live docstring, the skip is dropped, and + each one is executed and its output compared. + + Only names are redirected: ``my_extension`` resolves to the bundle above, + and ``SessionContext`` supplies the config this library's planner reads. + A renamed method, a changed signature, or a wrong expected output in the + docstring fails here. + """ + examples = doctest.DocTestParser().get_examples( + inspect.getdoc(SessionContext.with_extensions) + ) + assert examples, "with_extensions docstring has no examples to check" + for example in examples: + example.options.pop(doctest.SKIP, None) + + module = types.ModuleType("my_extension") + module.DistributedEngineExtension = _DocstringExampleExtension + + def make_context(config: SessionConfig | None = None) -> SessionContext: + # Accept a config so the example is free to pass one. Supplying it + # positionally the way the real constructor does keeps a docstring + # edit failing as a doctest diff rather than as a TypeError in here. + config = SessionConfig() if config is None else config + return SessionContext(config.with_extension(MyPlannerConfig(max_rows=3))) + + test = doctest.DocTest( + examples, + {"SessionContext": make_context}, + "SessionContext.with_extensions", + None, + None, + None, + ) + output = io.StringIO() + sys.modules["my_extension"] = module + try: + results = doctest.DocTestRunner().run(test, out=output.write) + finally: + del sys.modules["my_extension"] + assert results.failed == 0, output.getvalue() diff --git a/examples/datafusion-ffi-query-planner-example/src/distributed_exec.rs b/examples/datafusion-ffi-query-planner-example/src/distributed_exec.rs new file mode 100644 index 000000000..a0a347246 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/distributed_exec.rs @@ -0,0 +1,108 @@ +// 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 custom execution plan node owned by this library. +//! +//! This is the half of an extension bundle that makes the other half +//! necessary. A query planner that only rearranges stock DataFusion nodes needs +//! no codec of its own; one that emits a node *it* defines does, because +//! nothing else in the process knows how to serialize it. Shipping the planner +//! and the codec that carries its nodes as one bundle is the normal case, and +//! it is why `with_extensions` installs every codec before it binds any +//! planner. +//! +//! The node itself is deliberately trivial — it passes its child's stream +//! through untouched. A real distributed engine would ship the child plan to a +//! remote executor here; what matters for the example is that the node exists, +//! that this library's planner produces it, and that only this library's codec +//! can encode and decode it. + +use std::fmt; +use std::sync::Arc; + +use datafusion::common::Result; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; + +/// Marks a subtree this library claims for remote execution. +#[derive(Debug)] +pub(crate) struct DistributedExec { + input: Arc, + properties: Arc, +} + +impl DistributedExec { + pub(crate) fn new(input: Arc) -> Self { + // The node is pass-through, so it inherits its child's properties + // rather than describing anything of its own. + let properties = Arc::clone(input.properties()); + Self { input, properties } + } +} + +impl DisplayAs for DistributedExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "DistributedExec") + } +} + +impl ExecutionPlan for DistributedExec { + fn name(&self) -> &str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Owns no physical expressions of its own; the child holds them all. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return datafusion::common::internal_err!( + "DistributedExec expects exactly one child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new(children.swap_remove(0)))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs new file mode 100644 index 000000000..c9a42dd6e --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -0,0 +1,413 @@ +// 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. + +use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use datafusion::common::{Result, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_ffi::execution::FFI_TaskContextProvider; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, +}; +use datafusion_python_util::{ + create_logical_extension_capsule, create_physical_extension_capsule, + create_query_planner_capsule, ffi_logical_codec_from_pycapsule, + ffi_physical_codec_from_pycapsule, ffi_query_planner_from_pycapsule, + ffi_task_context_provider_from_pycapsule, get_tokio_runtime, +}; +use datafusion_session::QueryPlanner; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyDict}; + +use crate::distributed_exec::DistributedExec; +use crate::planner::{DistributedQueryPlanner, PlannerObservations, planner_config_from_options}; + +/// Values of `ffi_query_planner.max_rows` observed through the task-context +/// provider bound at installation time. +/// +/// Recorded on every decode call the chain routes to this bundle's physical +/// codec, including ones it declines. Reaching the session config from inside +/// a decode callback is what proves the provider bound at installation +/// resolves against the session running the query. +type ObservedMaxRows = Arc>>; + +/// The task-context provider handed to this bundle's components, if it has been +/// installed. `FFI_TaskContextProvider` holds its session weakly, so keeping one +/// here does not keep that session alive. +type BoundProvider = Arc>>; + +fn record_task_ctx(observed: &ObservedMaxRows, ctx: &TaskContext) { + if let Ok(config) = planner_config_from_options(ctx.session_config().options()) + && let Ok(mut observed) = observed.lock() + { + observed.push(config.max_rows); + } +} + +/// Carries this library's own [`DistributedExec`] nodes. +/// +/// This is the codec half of the bundle, and the reason the bundle ships both. +/// `DistributedQueryPlanner` emits a `DistributedExec`; no other codec in the +/// session knows the type, so without this one the plans that planner produces +/// cannot be serialized at all. Anything else is declined by delegating to the +/// default codec, so the host's chain falls through to whichever library owns +/// the node. +/// +/// The payload is a marker rather than a serialized node. `DistributedExec` is +/// pass-through and its child arrives already decoded in `inputs`, so there is +/// nothing else to write down; a node with state of its own would encode that +/// state here. +struct ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec, + observed: ObservedMaxRows, + claims: Arc, +} + +/// How often this bundle's codec claimed one of its own nodes. +/// +/// Distinct from [`ObservedMaxRows`], which counts every call the chain made, +/// including ones this codec declined. +#[derive(Default, Debug)] +pub(crate) struct DistributedExecClaims { + encoded: AtomicUsize, + decoded: AtomicUsize, +} + +/// Payload written for a [`DistributedExec`]. See +/// [`ObservingPhysicalExtensionCodec`]. +const DISTRIBUTED_EXEC_MARKER: &[u8] = b"datafusion_ffi_query_planner_example:DistributedExec"; + +impl fmt::Debug for ObservingPhysicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ObservingPhysicalExtensionCodec") + .finish_non_exhaustive() + } +} + +impl PhysicalExtensionCodec for ObservingPhysicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + // Reading the config through `ctx` is what proves the task-context + // provider bound at installation resolves against the session running + // the query. It happens on the real decode path now, not a synthetic + // one. + record_task_ctx(&self.observed, ctx); + if buf == DISTRIBUTED_EXEC_MARKER { + let [input] = inputs else { + return internal_err!( + "DistributedExec expects exactly one input, got {}", + inputs.len() + ); + }; + self.claims.decoded.fetch_add(1, Ordering::SeqCst); + return Ok(Arc::new(DistributedExec::new(Arc::clone(input)))); + } + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + if node.is::() { + self.claims.encoded.fetch_add(1, Ordering::SeqCst); + buf.extend_from_slice(DISTRIBUTED_EXEC_MARKER); + return Ok(()); + } + self.inner.try_encode(node, buf, proto_converter) + } +} + +/// Wire id this library's logical codec claims, pinned so that renaming the +/// Rust or Python types does not invalidate plans already encoded. +const LOGICAL_CODEC_ID: &str = "datafusion_ffi_query_planner_example.logical.v1"; + +/// Physical companion to [`LOGICAL_CODEC_ID`]. +const PHYSICAL_CODEC_ID: &str = "datafusion_ffi_query_planner_example.physical.v1"; + +/// Carries this bundle's logical codec as an object rather than a bare capsule. +/// +/// `with_extensions` requires an object: a codec's wire id is read off the +/// thing it is handed over as, and a capsule has no type to read one from. +/// Wrapping is also what keeps the id *this library's*. An id derived from the +/// contributing bundle would follow whichever object the caller passed to +/// `with_extensions`, so an application that packages this library inside a +/// bundle of its own would silently re-tag these payloads and they would stop +/// decoding in the process that reads them. The wrapper travels with the codec; +/// the bundle does not. +/// +/// Declaring `__datafusion_codec_id__` is optional — the class's +/// `module.QualName` would serve — but a library whose plans leave the process +/// should pin the id rather than let a refactor move it. +#[pyclass( + name = "BundledLogicalCodec", + module = "datafusion_ffi_query_planner_example" +)] +pub(crate) struct BundledLogicalCodec { + codec: FFI_LogicalExtensionCodec, +} + +#[pymethods] +impl BundledLogicalCodec { + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + LOGICAL_CODEC_ID + } + + /// `session` is unused: the codec was bound to its task-context provider + /// when the bundle was installed, which is the whole reason the bundle + /// receives the context. + #[pyo3(signature = (session=None))] + fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + create_logical_extension_capsule(py, &self.codec) + } +} + +/// Physical companion to [`BundledLogicalCodec`]. +#[pyclass( + name = "BundledPhysicalCodec", + module = "datafusion_ffi_query_planner_example" +)] +pub(crate) struct BundledPhysicalCodec { + codec: FFI_PhysicalExtensionCodec, +} + +#[pymethods] +impl BundledPhysicalCodec { + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + PHYSICAL_CODEC_ID + } + + /// See [`BundledLogicalCodec::__datafusion_logical_extension_codec__`]. + #[pyo3(signature = (session=None))] + fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + create_physical_extension_capsule(py, &self.codec) + } +} + +/// Extension bundle for `SessionContext.with_extensions`. +/// +/// Mirrors how a distributed engine such as Ballista packages its session +/// extensions: the object itself is reusable configuration, and every +/// `__datafusion_session_extension__` call creates fresh codec and planner +/// components bound to the task-context provider of the context it receives. +#[pyclass( + from_py_object, + name = "MyPlannerExtension", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Default, Clone)] +pub(crate) struct MyPlannerExtension { + observations: Arc, + observed_max_rows: ObservedMaxRows, + claims: Arc, + bound_provider: BoundProvider, +} + +impl fmt::Debug for MyPlannerExtension { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("MyPlannerExtension") + .field("observations", &self.observations) + .finish_non_exhaustive() + } +} + +#[pymethods] +impl MyPlannerExtension { + #[new] + fn new() -> Self { + Self::default() + } + + fn plan_calls(&self) -> usize { + self.observations.plan_calls.load(Ordering::SeqCst) + } + + fn last_max_rows(&self) -> usize { + self.observations.last_max_rows.load(Ordering::SeqCst) + } + + fn foreign_session_observed(&self) -> bool { + self.observations.foreign_session.load(Ordering::SeqCst) + } + + fn foreign_provider_observed(&self) -> bool { + self.observations.foreign_provider.load(Ordering::SeqCst) + } + + fn foreign_plan_observed(&self) -> bool { + self.observations.foreign_plan.load(Ordering::SeqCst) + } + + /// `ffi_query_planner.max_rows` values seen through the bound + /// task-context provider during codec decode calls. + /// + /// One entry per decode call the chain routed to this bundle, so a query + /// whose plan carries a `DistributedExec` leaves several. + fn decode_max_rows_seen(&self) -> Vec { + self.observed_max_rows + .lock() + .map(|observed| observed.clone()) + .unwrap_or_default() + } + + /// How often this bundle's physical codec encoded one of the + /// `DistributedExec` nodes its own planner produced. + /// + /// Non-zero only when the plan was actually serialized — running a query + /// does not do that, because an FFI planner hands its result back as an + /// opaque plan handle. A distributed engine shipping the plan to a remote + /// executor does, which is the case the pairing exists for; in this + /// repository `ExecutionPlan.to_bytes` stands in for it. + fn distributed_exec_encode_calls(&self) -> usize { + self.claims.encoded.load(Ordering::SeqCst) + } + + /// Companion to [`Self::distributed_exec_encode_calls`], counting the + /// nodes rebuilt on the way back in. + fn distributed_exec_decode_calls(&self) -> usize { + self.claims.decoded.load(Ordering::SeqCst) + } + + /// `ffi_query_planner.max_rows` read through the task-context provider + /// this bundle was last bound to. + /// + /// Resolving the provider is what a codec's decode callback does, so this + /// answers which session those callbacks would resolve against -- the + /// context `with_extensions` returned, not the one it was called on. + /// Returns ``None`` if the bundle was never installed, or if the context it + /// was bound to has been dropped: the provider holds it weakly. + fn max_rows_through_provider(&self) -> Option { + let provider = self.bound_provider.lock().ok()?.clone()?; + let task_ctx = Arc::::try_from(&provider).ok()?; + planner_config_from_options(task_ctx.session_config().options()) + .ok() + .map(|config| config.max_rows) + } + + fn __datafusion_session_extension__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + // Bind every component to the context supplied by the host, which is + // the session the components will run on. Components must not be + // cached across calls: each installation may target a different + // session. + // + // The task-context provider comes off that context rather than from a + // `SessionContext` built here, so the codecs' decode callbacks resolve + // names against the session that will actually run the query. + let provider = ffi_task_context_provider_from_pycapsule(&ctx)?; + if let Ok(mut bound) = self.bound_provider.lock() { + *bound = Some(provider.clone()); + } + let runtime = get_tokio_runtime().handle().clone(); + + // Plain default: this library defines no logical extension node, so + // there is nothing for a logical codec of its own to claim. It is still + // contributed so the bundle carries both codec kinds under ids it + // declares -- see `BundledLogicalCodec`. + let logical: Arc = Arc::new(DefaultLogicalExtensionCodec {}); + let ffi_logical = + FFI_LogicalExtensionCodec::new(logical, Some(runtime.clone()), provider.clone()); + // Handed over as an object, not a capsule, so the codec carries an id + // of its own. See `BundledLogicalCodec`. + let logical_codec = Py::new(py, BundledLogicalCodec { codec: ffi_logical })?; + + let physical: Arc = + Arc::new(ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + claims: Arc::clone(&self.claims), + }); + let ffi_physical = + FFI_PhysicalExtensionCodec::new(physical, Some(runtime), provider.clone()); + let physical_codec = Py::new( + py, + BundledPhysicalCodec { + codec: ffi_physical, + }, + )?; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("logical_extension_codecs", (logical_codec,))?; + kwargs.set_item("physical_extension_codecs", (physical_codec,))?; + components.call((), Some(&kwargs)) + } + + /// Contribute this library's planner, nesting it on whatever came before. + /// + /// Runs in the host's second phase, after every bundle's codecs are + /// installed, so `ctx` carries the final chains and the planner this + /// builds is not left encoding through a partial set. `fallback` is the + /// planner assembled so far — the session's existing one for the first + /// bundle, the previous bundle's for the rest — and delegating to it is + /// what makes several planner-shipping libraries composable. Returning a + /// planner that ignored it would discard every layer beneath. + fn __datafusion_session_planner__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + fallback: Bound<'py, PyAny>, + ) -> PyResult> { + let fallback = ffi_query_planner_from_pycapsule(&fallback, Some(&ctx))?; + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + fallback: Some((&fallback).into()), + }); + // The planner takes the host's codecs, not ones built here. By now + // those are the final chains, and this library has no business minting + // a provider of its own. + let host_logical = ffi_logical_codec_from_pycapsule(ctx.clone(), None)?; + let host_physical = ffi_physical_codec_from_pycapsule(ctx, None)?; + let ffi_planner = + FFI_QueryPlanner::new_with_ffi_codecs(planner, host_logical, host_physical); + create_query_planner_capsule(py, &ffi_planner) + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs index c505c1ce7..30a2d5e4f 100644 --- a/examples/datafusion-ffi-query-planner-example/src/lib.rs +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -18,15 +18,21 @@ use pyo3::prelude::*; use crate::config::MyPlannerConfig; +use crate::extension::{BundledLogicalCodec, BundledPhysicalCodec, MyPlannerExtension}; use crate::planner::MyQueryPlanner; mod config; +mod distributed_exec; +mod extension; mod planner; #[pymodule] fn datafusion_ffi_query_planner_example(m: &Bound<'_, PyModule>) -> PyResult<()> { pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 67262e39c..744b67458 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -40,6 +40,7 @@ use pyo3::prelude::*; use pyo3::types::PyCapsule; use crate::config::MyPlannerConfig; +use crate::distributed_exec::DistributedExec; /// What the planner saw, accumulated across every call rather than reset each /// time. @@ -52,13 +53,15 @@ use crate::config::MyPlannerConfig; /// most recent plan would be answering a different question than the one its /// accessor name asks. #[derive(Default)] -struct PlannerObservations { - plan_calls: AtomicUsize, - last_max_rows: AtomicUsize, - foreign_session: AtomicBool, - foreign_provider: AtomicBool, - foreign_plan: AtomicBool, - /// Only ever set to `true`, so it is already cumulative. +pub(crate) struct PlannerObservations { + pub(crate) plan_calls: AtomicUsize, + pub(crate) last_max_rows: AtomicUsize, + pub(crate) foreign_session: AtomicBool, + pub(crate) foreign_provider: AtomicBool, + pub(crate) foreign_plan: AtomicBool, + /// Only ever set to `true`, so it is already cumulative. Read only through + /// `MyQueryPlanner::used_fallback` in this module, so unlike its + /// neighbours it needs no wider visibility. used_fallback: AtomicBool, } @@ -104,8 +107,12 @@ const MAX_ROWS_KEY: &str = "ffi_query_planner.max_rows"; const FFI_MAX_ROWS_KEY: &str = "datafusion_ffi.ffi_query_planner.max_rows"; fn planner_config(session: &dyn Session) -> datafusion::common::Result { - let options = session.config_options(); + planner_config_from_options(session.config_options()) +} +pub(crate) fn planner_config_from_options( + options: &datafusion::common::config::ConfigOptions, +) -> datafusion::common::Result { // Prefer the raw entry. `local_or_ffi_extension` discards a value it cannot // parse and hands back `MyPlannerConfig::default()`, which would quietly turn // a typo into a different row limit instead of reporting it. @@ -143,8 +150,8 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result, +pub(crate) struct DistributedQueryPlanner { + pub(crate) observations: Arc, /// Planner to hand the work to instead of planning here. /// /// This is how a real planner layers on top of an existing one. The capsule @@ -156,7 +163,7 @@ struct DistributedQueryPlanner { /// Note that `Session::create_physical_plan` cannot be used for this. It /// dispatches through the session's installed query planner, so calling it /// from inside that planner recurses until the stack overflows. - fallback: Option>, + pub(crate) fallback: Option>, } #[async_trait] @@ -201,11 +208,14 @@ impl QueryPlanner for DistributedQueryPlanner { .foreign_plan .fetch_or(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst); - Ok(Arc::new(GlobalLimitExec::new( - plan, - 0, - Some(config.max_rows), - ))) + // Wrap the result in a node this library owns. Nothing else in the + // process can serialize a `DistributedExec`, so a session that installs + // this planner without the matching physical codec cannot round-trip + // the plans it produces -- which is the reason the two ship as one + // bundle. See `ObservingPhysicalExtensionCodec`. + Ok(Arc::new(DistributedExec::new(Arc::new( + GlobalLimitExec::new(plan, 0, Some(config.max_rows)), + )))) } } diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9c55f446c..fd6f27b82 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -92,6 +92,12 @@ ) from .dataframe_formatter import configure_formatter from .expr import Expr, WindowFrame +from .extensions import ( + QueryPlannerExportable, + SessionExtensionComponents, + SessionExtensionExportable, + SessionPlannerExportable, +) from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions from .plan import ExecutionPlan, LogicalPlan, Metric, MetricsSet @@ -127,6 +133,7 @@ "MetricsSet", "ParquetColumnOptions", "ParquetWriterOptions", + "QueryPlannerExportable", "RecordBatch", "RecordBatchStream", "RuntimeEnvBuilder", @@ -134,6 +141,9 @@ "ScalarUDF", "SessionConfig", "SessionContext", + "SessionExtensionComponents", + "SessionExtensionExportable", + "SessionPlannerExportable", "Table", "TableFunction", "TableProviderFactory", @@ -146,6 +156,7 @@ "common", "configure_formatter", "expr", + "extensions", "functions", "ipc", "lit", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 644c7b445..69339a33f 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -69,6 +69,12 @@ ) from datafusion.dataframe import DataFrame from datafusion.expr import sort_list_to_raw_sort_list +from datafusion.extensions import ( + QueryPlannerExportable, + SessionExtensionComponents, + SessionExtensionExportable, + SessionPlannerExportable, +) from datafusion.options import ( DEFAULT_MAX_INFER_SCHEMA, CsvReadOptions, @@ -145,18 +151,6 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 -class QueryPlannerExportable(Protocol): - """Type hint for object that has a __datafusion_query_planner__ PyCapsule. - - The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically - produced by a separate compiled extension. ``session`` is the - :py:class:`SessionContext` the planner is being installed on; take the - extension codecs from it rather than building your own. - """ - - def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 - - class SessionConfig: """Session configuration options.""" @@ -1798,7 +1792,7 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non Args: planner: Object exposing ``__datafusion_query_planner__`` (see - :class:`QueryPlannerExportable`) or a raw + :py:class:`~datafusion.extensions.QueryPlannerExportable`) or a raw ``datafusion_query_planner`` PyCapsule. Examples: @@ -1817,6 +1811,198 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non """ self.ctx.set_query_planner(planner) + def with_extensions( + self, *extensions: SessionExtensionExportable | SessionPlannerExportable + ) -> SessionContext: + """Create a new session context with the given extension bundles. + + This is the preferred way to install FFI extensions that need a + task-context provider (extension codecs and query planners). It avoids + the pitfalls of chaining :py:meth:`with_logical_extension_codec`, + :py:meth:`with_physical_extension_codec`, and + :py:meth:`set_query_planner` by hand, where the codecs a planner was + built against can end up stale. + + Installation runs in two phases, because codecs and planners compose + differently: + + 1. Every extension's ``__datafusion_session_extension__`` is called + with this context and its codecs are collected, then all of them are + installed at once. A session chains many codecs and dispatches + between them by id, so they merely accumulate; order affects + encoding only. + 2. Every extension's ``__datafusion_session_planner__`` is then called, + **in argument order**, each receiving the planner built so far. A + session holds exactly one planner, so planners compose by *nesting*: + each wraps the previous one and delegates to it. The last extension + listed ends up outermost and is consulted first. + + An extension implements either hook or both. Phase two runs after every + codec is installed and receives a context carrying the final chains, so + a nested planner is never left encoding through a chain a later + extension has grown. + + If no extension supplies a planner but codecs were installed, an + existing FFI planner is rebound to the final chains; if the call + installed nothing at all, the session's planner is not touched. An + extension that ignores the ``fallback`` it is handed replaces the + planners before it instead of nesting on them, including any the + session already had. + + Codec order never affects decoding, which routes by codec id. It + affects encoding only when two codecs would claim the same node: the + chain stops at the first that does, so a codec claiming a broad + category can take nodes belonging to a library installed after it. The + query still succeeds, but the plan is written by the wrong library and + may not decode elsewhere. If an extension needs to be early for its + codec and late for its planner, contribute each half at its own + position rather than reordering — the two hooks are independent, so a + small adapter implementing one of them and delegating is enough. The + FFI extensions guide shows the pattern. + + Codecs must be handed over as objects exposing the capsule getter, not + as bare ``PyCapsule`` objects, and are named after their exporting + class as :py:meth:`with_logical_extension_codec` describes. Declare + ``__datafusion_codec_id__`` on the object to pin an id that survives a + later class rename. A capsule carries no type of its own, so there + would be nothing to name the codec by, and this method takes no + ``codec_id=``; wrap it in an object instead. That also keeps a codec's + wire identity independent of the extension that ships it, so an + extension composed inside another one still writes the same ids. + + Planners are exempt — a planner carries no wire id, so a hook may + return an object or a capsule. + + Like the individual ``with_*`` methods, the returned context shares its + session with this one: catalogs, tables, registered functions, and + configuration are the one session, so a registration on either side is + visible to both, and the planner is installed on that shared session + even if the returned context is discarded. Only the Python-side codec + chains are specific to the returned handle. + + No state is written until every extension has run and every capsule has + been validated, so an extension that raises or returns something + invalid — in either phase — leaves the session as it was. Codec chains + belong to the returned handle, and the single session write happens + after the last planner hook returns. The exception is an extension that + mutates the context it is handed — registering a table, say — which is + not rolled back. Extension factories should treat that context as + configuration-only. + + The session owns the installed components' task-context providers, and + dependent objects do not extend its lifetime. Keep a context on the + session alive for as long as DataFrames or plans derived from it are in + use; FFI operations after the last one is collected raise an error. + + Args: + extensions: Extension bundles to install. Order is irrelevant for + codecs and significant for planners, which nest in this order + with the last one outermost. Passing none installs nothing and + returns a handle on this session, so a caller assembling the + list at runtime need not special-case it being empty. + + Returns: + A new context with all extension components installed. + + Raises: + TypeError: If an argument implements neither hook, if + ``__datafusion_session_extension__`` returns something other + than a + :py:class:`~datafusion.extensions.SessionExtensionComponents`, + or if an extension contributes a codec as a bare ``PyCapsule``. + ValueError: If two codecs claim the same id. An extension that + contributes two instances of one codec class must declare + ``__datafusion_codec_id__`` on at least one of them; the + collision is refused rather than resolved by position, because + a positional id would break stored plans the first time the + extension reordered what it returns. Also if a capsule getter + returns a capsule of the wrong kind — a physical codec handed + over under ``__datafusion_logical_extension_codec__``, say — + which is reported against the name the getter should have + produced. + + Examples: + The example is skipped here because it needs a built FFI + extension library, which this package does not ship. It is run + verbatim against a real one by + ``test_with_extensions_docstring_example_still_runs`` in + ``examples/datafusion-ffi-query-planner-example``, so it cannot + drift from the API. + + >>> from my_extension import DistributedEngineExtension # doctest: +SKIP + >>> ctx = SessionContext().with_extensions( + ... DistributedEngineExtension("scheduler:50050") + ... ) # doctest: +SKIP + >>> batches = ctx.sql("SELECT 1 AS n").collect() # doctest: +SKIP + >>> batches[0].column(0).to_pylist() # doctest: +SKIP + [1] + """ + for extension in extensions: + if not isinstance( + extension, (SessionExtensionExportable, SessionPlannerExportable) + ): + msg = ( + "Extension implements neither " + "__datafusion_session_extension__ nor " + f"__datafusion_session_planner__: {extension!r}" + ) + raise TypeError(msg) + + # Phase one: collect every bundle's codecs. Components are bound + # against this context, not a context derived from it. There is one + # `Arc` per session, so a component bound here holds a + # task-context provider that the returned handle keeps alive. + logical_codecs: list[LogicalExtensionCodecExportable] = [] + physical_codecs: list[PhysicalExtensionCodecExportable] = [] + for extension in extensions: + if not isinstance(extension, SessionExtensionExportable): + continue + components = extension.__datafusion_session_extension__(self) + if not isinstance(components, SessionExtensionComponents): + msg = ( + "__datafusion_session_extension__ must return " + "SessionExtensionComponents, got " + f"{type(components).__name__} from {extension!r}" + ) + raise TypeError(msg) + logical_codecs.extend(components.logical_extension_codecs) + physical_codecs.extend(components.physical_extension_codecs) + + # Writes nothing: the chains belong to the new handle, so a failure + # above or below leaves this context as it was. + new = SessionContext.__new__(SessionContext) + new.ctx = self.ctx._install_extension_codecs(logical_codecs, physical_codecs) + + # Phase two: nest the planners, outermost last. Each hook runs against + # `new`, which carries the final chains, so a planner captured here + # never sees a partial codec set. `planner` stays None when no bundle + # supplies one, which leaves an already-installed planner in place + # rather than wrapping the session's default in an FFI hop. + planner: _PyCapsule | None = None + for extension in extensions: + if not isinstance(extension, SessionPlannerExportable): + continue + fallback = ( + planner + if planner is not None + else new.ctx.__datafusion_query_planner__() + ) + supplied = extension.__datafusion_session_planner__(new, fallback) + if supplied is None: + continue + planner = new.ctx._export_query_planner(supplied) + + # Rebinding the session's planner is a side effect on state shared with + # every other handle, so do not pay it for a call that installs nothing + # -- the same guard `with_python_udf_inlining` carries. With no codec + # installed the chains the planner would be rebuilt against are the ones + # it already holds, so the rebuild is unobservable except in the one case + # where it does harm: a planner sitting on some *other* handle's codecs + # gets dragged onto this handle's, silently undoing that install. + if planner is not None or logical_codecs or physical_codecs: + new.ctx._install_extension_planner(planner) + return new + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. @@ -2246,9 +2432,11 @@ def __datafusion_codec_id__(self) -> str: written through one will not be decoded by the other. 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. + :py:meth:`with_logical_extension_codec`, + :py:meth:`with_python_udf_inlining`, and :py:meth:`with_extensions` — + report the same id, so only one of them can be installed on a given + session. That is the intended answer: they are one session, so their + payloads would be indistinguishable on decode. Examples: >>> from datafusion import SessionContext diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py new file mode 100644 index 000000000..631fe239b --- /dev/null +++ b/python/datafusion/extensions.py @@ -0,0 +1,332 @@ +# 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. + +"""Protocols and value types for installing extensions on a session context. + +An *extension* is a reusable configuration object — typically shipped by a +separate compiled library — that contributes components to a +:py:class:`~datafusion.context.SessionContext`. It implements +:py:class:`SessionExtensionExportable` by returning a +:py:class:`SessionExtensionComponents` describing what it contributes, and is +installed with :py:meth:`~datafusion.context.SessionContext.with_extensions`:: + + ctx = SessionContext().with_extensions(MyLibraryExtension()) + +Installing through ``with_extensions`` rather than by chaining the individual +``with_*`` methods matters for components that hold a task-context provider: +the extension is handed the session its components will run on, and every +codec is installed before any query planner is bound against them, so no +planner is left carrying a codec chain that has since grown. See the FFI +extensions guide in the contributor documentation for the full rationale. + +Codecs and planners install in two phases, because they compose differently. A +session's codec chain holds many codecs and dispatches between them by id, so +codecs merely accumulate and their order does not affect decoding. A session +holds exactly *one* query planner, so planners compose by nesting: each wraps +the one before it. Phase one collects the codecs of every bundle implementing +:py:class:`SessionExtensionExportable` and installs them; phase two runs +:py:class:`SessionPlannerExportable` once for each bundle that implements it, +in argument order, handing each the planner built so far. A bundle implements +either hook or both, and one it does not implement is simply not called. + +That split is what lets several libraries that each ship a planner coexist. It +also means bundle order is significant for planners and irrelevant for codecs. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from _typeshed import CapsuleType as _PyCapsule + + from datafusion.context import SessionContext + from datafusion.user_defined import ( + LogicalExtensionCodecExportable, + PhysicalExtensionCodecExportable, + ) + +__all__ = [ + "QueryPlannerExportable", + "SessionExtensionComponents", + "SessionExtensionExportable", + "SessionPlannerExportable", +] + + +class QueryPlannerExportable(Protocol): + """Type hint for object that has a __datafusion_query_planner__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically + produced by a separate compiled extension. ``session`` is a handle on the + session the planner is being installed on; take the extension codecs from + it rather than building your own. + + Duck-type that handle rather than checking its type. It is the PyO3 + context from ``datafusion._internal``, not the + :py:class:`~datafusion.context.SessionContext` wrapper, so it exposes every + capsule getter and ``__datafusion_codec_id__`` — which is all the protocol + asks of it — but ``isinstance(session, SessionContext)`` is ``False`` even + though its ``repr`` reads ``datafusion.SessionContext``. The same is true + of the codec getters in :py:mod:`datafusion.user_defined`. The two bundle + hooks are the exception: :py:class:`SessionExtensionExportable` and + :py:class:`SessionPlannerExportable` are dispatched from Python and receive + the wrapper. + """ + + def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 + + +def _not_a_codec_iterable(field: str, value: object) -> str: + """Message for a codec field that cannot be read as a collection.""" + return ( + f"{field} must be an iterable of codec objects, not a single " + f"{type(value).__name__}. A lone codec is written as a one-element " + f"tuple — {field}=(codec,) — and the trailing comma is what makes it one." + ) + + +@dataclass(frozen=True) +class SessionExtensionComponents: + """Components an extension contributes to a session context. + + Returned by :py:meth:`SessionExtensionExportable.__datafusion_session_extension__` + and consumed by + :py:meth:`~datafusion.context.SessionContext.with_extensions`. Every + component must be created against the context passed to that method; + components bound to a different session hold a task-context provider for + that other session and cannot be rebound. + + Query planners are not listed here. They install in a second phase so each + can wrap the one before it — see :py:class:`SessionPlannerExportable`. + + Codecs must be objects exposing the capsule getters, never bare + ``PyCapsule`` objects: a codec's id is read off the object it is handed + over as, and a capsule has no type to read. A library holding a raw capsule + wraps it in an object, which is also what gives the codec an identity of + its own — one that does not change when the codec is contributed through a + different extension. + + Examples: + A bundle that contributes no codecs is valid — a planner-only library + returns this, or omits the hook entirely: + + >>> from datafusion import SessionExtensionComponents + >>> components = SessionExtensionComponents() + >>> components.logical_extension_codecs + () + >>> components.physical_extension_codecs + () + + A bundle that contributes one kind of component names it, leaving the + rest empty. Here the codec is a capsule wrapped in an object that + declares the id its payloads will carry: + + >>> from datafusion import SessionContext + >>> class NamedCodec: + ... __datafusion_codec_id__ = "my_library.v1" + ... + ... def __init__(self, capsule): + ... self._capsule = capsule + ... + ... def __datafusion_logical_extension_codec__(self, session=None): + ... return self._capsule + + The context stays in scope for as long as the codec does. An + ``FFI_LogicalExtensionCodec`` holds its task-context provider *weakly*, + so a capsule taken off a throwaway ``SessionContext()`` names a session + that is already gone and fails on first use with ``TaskContextProvider + went out of scope over FFI boundary``: + + >>> ctx = SessionContext() + >>> capsule = ctx.__datafusion_logical_extension_codec__() + >>> components = SessionExtensionComponents( + ... logical_extension_codecs=(NamedCodec(capsule),) + ... ) + >>> components.logical_extension_codecs[0].__datafusion_codec_id__ + 'my_library.v1' + >>> components.physical_extension_codecs + () + + Any iterable is accepted and stored as a tuple, so a bundle that builds + its codecs with a list comprehension does not have to convert: + + >>> components = SessionExtensionComponents( + ... logical_extension_codecs=[NamedCodec(capsule)] + ... ) + >>> type(components.logical_extension_codecs).__name__ + 'tuple' + + A single codec is not an iterable of codecs, and forgetting the + trailing comma is the easy way to write one by accident: + + >>> SessionExtensionComponents(logical_extension_codecs=NamedCodec(capsule)) + Traceback (most recent call last): + ... + TypeError: logical_extension_codecs must be an iterable of codec objects... + """ + + logical_extension_codecs: tuple[LogicalExtensionCodecExportable, ...] = () + """Logical codecs to add to the session's codec chain, in declaration order.""" + + physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = () + """Physical codecs to add to the session's codec chain, in declaration order.""" + + def __post_init__(self) -> None: + """Normalize each field to a tuple, rejecting what cannot become one. + + A bundle that writes ``logical_extension_codecs=codec`` instead of + ``(codec,)`` is contributing one codec, not an iterable of them. + Without this, the mistake surfaces inside + :py:meth:`~datafusion.context.SessionContext.with_extensions` as + ``'MyCodec' object is not iterable``, which names neither the field + nor the hook that built it. Checking here puts the error in the + extension library's own frame. + + Normalizing is worth doing on its own: the declared type is a tuple + and the class is frozen, so a list left in place would be a mutable + member of an immutable value, and a generator would be exhausted by + the first read. + + Driven off :py:func:`dataclasses.fields` rather than a written-out + list, so a codec field added later is normalized without anyone + remembering to name it here. The ``_codecs`` suffix is what marks a + field as one of them, leaving room for a future field that is not a + codec collection and must not be turned into a tuple. + """ + for field in fields(self): + name = field.name + if not name.endswith("_codecs"): + continue + value = getattr(self, name) + # A str is iterable, so it would otherwise normalize into a tuple + # of characters and fail much later as that many bogus codecs. + if isinstance(value, (str, bytes)): + raise TypeError(_not_a_codec_iterable(name, value)) + try: + codecs = tuple(value) + except TypeError: + raise TypeError(_not_a_codec_iterable(name, value)) from None + object.__setattr__(self, name, codecs) + + +@runtime_checkable +class SessionExtensionExportable(Protocol): + """Type hint for extension bundles installable via ``with_extensions``. + + Runtime-checkable, so ``isinstance`` answers whether an object implements + the protocol. Only the presence of the method is checked, which is the same + question :py:meth:`~datafusion.context.SessionContext.with_extensions` asks + before calling it. + + Implementations are reusable configuration objects: they must create fresh + components on every call using the context supplied by + :py:meth:`~datafusion.context.SessionContext.with_extensions`, and must not + retain that context or cache the components they bound to it, since the + next call may install onto a different session. They should also avoid + mutating the context they are handed — a registration made during binding + is not rolled back if a later extension fails. + + ``ctx`` is the right session but not yet the final codec chains: this hook + runs before anything is installed, so ``ctx`` still carries whatever chains + the receiver had. Take the task-context provider off it — that is bound to + the session and is what the components need — but do not read its codec + chains expecting to find this call's codecs, including your own. + :py:class:`SessionPlannerExportable` is the hook that sees the completed + chains, which is why a planner that wraps the host's codecs builds them + there rather than here. + + A bundle that also contributes a query planner implements + :py:class:`SessionPlannerExportable` alongside this protocol. Planners are + installed in a second phase, so they are not part of the components + returned here. + + Examples: + >>> from datafusion import ( + ... SessionExtensionComponents, + ... SessionExtensionExportable, + ... ) + >>> class MyLibraryExtension: + ... def __datafusion_session_extension__(self, ctx): + ... return SessionExtensionComponents() + >>> isinstance(MyLibraryExtension(), SessionExtensionExportable) + True + >>> isinstance(object(), SessionExtensionExportable) + False + """ + + def __datafusion_session_extension__( # noqa: D105 + self, ctx: SessionContext + ) -> SessionExtensionComponents: ... + + +@runtime_checkable +class SessionPlannerExportable(Protocol): + """Type hint for extension bundles that contribute a query planner. + + A session holds exactly one query planner, so planners compose by nesting + rather than by chaining: each wraps the one before it and delegates to it + for the work it does not handle. + :py:meth:`~datafusion.context.SessionContext.with_extensions` runs this + hook once per bundle that implements it, **in argument order**, handing + each the planner built so far. Returning a planner that wraps ``fallback`` + puts this bundle *outside* the previous one, so the last bundle listed ends + up outermost and is consulted first. + + The hook runs after every codec from every bundle is installed, and ``ctx`` + is the context carrying those final chains. That ordering is the point: a + planner captured here sees the complete codec set, so a nested planner is + not left encoding through a chain that a later bundle has grown. + + Return ``None`` to contribute no planner and leave ``fallback`` in place. + That is the no-op, and it is not the same as returning ``fallback``: the + capsule the first bundle receives wraps the session's planner for export, so + handing it back installs it as a foreign planner and every later plan crosses + an FFI boundary that was not there before. A bundle with nothing to + contribute returns ``None``. + + Ignoring ``fallback`` and returning a planner that does not delegate to it + is legal and means "replace" — but it discards every planner listed before + this one, including any the session already had. + + Args: + ctx: The session the planner will run on, carrying the final codec + chains. + fallback: The planner built so far, as a ``PyCapsule``. For the first + bundle this is the session's existing planner, which is the + DataFusion default unless one was installed earlier. + + Examples: + >>> from datafusion import SessionPlannerExportable + >>> class MyEngineExtension: + ... def __datafusion_session_planner__(self, ctx, fallback): + ... # A real library returns its own planner wrapping + ... # `fallback`, e.g. ``my_library.Planner(fallback=fallback)``. + ... # Handing it straight back is the degenerate wrap: legal, + ... # but it still installs `fallback` as a foreign planner. + ... # Return None instead to contribute nothing. + ... return fallback + >>> isinstance(MyEngineExtension(), SessionPlannerExportable) + True + >>> isinstance(object(), SessionPlannerExportable) + False + """ + + def __datafusion_session_planner__( # noqa: D105 + self, ctx: SessionContext, fallback: _PyCapsule + ) -> QueryPlannerExportable | _PyCapsule | None: ... diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 3c95835af..c1aa19a4e 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -16,6 +16,7 @@ # under the License. import ctypes import datetime as dt +import gc import gzip import pathlib import shutil @@ -29,6 +30,7 @@ RuntimeEnvBuilder, SessionConfig, SessionContext, + SessionExtensionComponents, SQLOptions, Table, column, @@ -879,6 +881,477 @@ def test_contexts_sharing_a_session_share_the_planner(ctx): assert sibling.session_id() == ctx.session_id() +class _NamedCodec: + """Wraps a codec capsule in an object that can name itself. + + ``with_extensions`` requires objects rather than bare capsules, because a + codec's wire id is read off the object it is handed over as. This is the + shape a library holding a raw capsule hands over. + """ + + def __init__(self, capsule, codec_id): + self._capsule = capsule + self.__datafusion_codec_id__ = codec_id + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule + + def __datafusion_physical_extension_codec__(self, session=None): + return self._capsule + + +class _CodecOnlyExtension: + """Contributes decline-all codecs exported from an unrelated session. + + Retaining ``ctx`` is what the protocol tells real extensions not to do — + a bundle is reusable, so a cached context belongs to whichever session it + was last installed on. It is kept here only so a test can assert *which* + context the factory was handed. + """ + + def __init__(self, prefix="my_library"): + self.exporter = SessionContext() + self.prefix = prefix + self.bound_ctx = None + + def __datafusion_session_extension__(self, ctx): + self.bound_ctx = ctx + return SessionExtensionComponents( + logical_extension_codecs=( + _NamedCodec( + self.exporter.__datafusion_logical_extension_codec__(), + f"{self.prefix}.logical", + ), + ), + physical_extension_codecs=( + _NamedCodec( + self.exporter.__datafusion_physical_extension_codec__(), + f"{self.prefix}.physical", + ), + ), + ) + + +class _PlannerExtension: + """Contributes a planner, recording the fallback it was handed. + + Passing ``fallback`` straight back through is the degenerate wrap: it plans + the same queries to the same plans, which is what lets a pure-Python test + assert the threading without a real layering planner. It is not a no-op — + the capsule gets installed, so the session ends up planning through a + foreign planner — but nothing here depends on that either way. + """ + + def __init__(self, calls=None): + self.fallbacks = [] + self.planner_ctx = None + # Shared list the hooks append themselves to, so a test can assert the + # order they ran in rather than only that each ran. + self.calls = [] if calls is None else calls + + def __datafusion_session_planner__(self, ctx, fallback): + self.calls.append(self) + self.planner_ctx = ctx + self.fallbacks.append(fallback) + return fallback + + +def test_with_extensions_accepts_no_extensions(ctx): + """No extensions installs nothing and returns a handle on this session. + + A caller assembling the list at runtime — from a plugin registry, say — + should not have to special-case it being empty, and every sibling varargs + method on ``DataFrame`` accepts zero arguments the same way. + """ + ctx.register_record_batches( + "empty_extensions_test", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + result = ctx.with_extensions() + + assert result.session_id() == ctx.session_id() + assert result.table_exist("empty_extensions_test") + assert result.logical_extension_codec_ids() == ctx.logical_extension_codec_ids() + + +def test_with_extensions_no_extensions_keeps_an_installed_planner(ctx): + """The empty case must not disturb a planner the session already has. + + A call that installs no codec and no planner skips the planner commit + entirely, so the installed planner keeps the chains it was bound to and the + session still plans through it. + """ + extension = _CodecOnlyExtension() + installed = ctx.with_extensions(extension, _PlannerExtension()) + + result = installed.with_extensions() + + assert result.logical_extension_codec_ids() == ( + installed.logical_extension_codec_ids() + ) + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_rejects_non_extension(ctx): + with pytest.raises(TypeError, match="__datafusion_session_planner__"): + ctx.with_extensions(object()) + + +def test_with_extensions_rejects_bad_components(ctx): + class BadExtension: + def __datafusion_session_extension__(self, ctx): + return 42 + + with pytest.raises(TypeError, match="SessionExtensionComponents"): + ctx.with_extensions(BadExtension()) + + +@pytest.mark.parametrize( + "field", ["logical_extension_codecs", "physical_extension_codecs"] +) +def test_session_extension_components_rejects_a_single_codec(field): + """A lone codec is not an iterable of codecs. + + Dropping the trailing comma is the easy way to write one by accident. The + check lives on the value type so the error lands in the extension + library's own frame, naming the field it got wrong, rather than surfacing + later inside ``with_extensions`` as ``'_NamedCodec' object is not + iterable``. + """ + codec = _NamedCodec( + SessionContext().__datafusion_logical_extension_codec__(), + "my_library.logical", + ) + + with pytest.raises(TypeError, match=r"must be an iterable of codec objects"): + SessionExtensionComponents(**{field: codec}) + + +def test_session_extension_components_rejects_a_string(): + """A str is iterable, so it needs refusing on its own. + + Left alone it would normalize into a tuple of characters and fail much + later as that many bogus codecs. + """ + with pytest.raises(TypeError, match=r"not a single str"): + SessionExtensionComponents(logical_extension_codecs="my_library.logical") + + +def test_with_extensions_accepts_a_planner_only_extension(ctx): + """An extension may implement the planner hook alone. + + A library that ships an optimizing planner and no codecs — nothing to + contribute in phase one — should not have to return empty components. + """ + extension = _PlannerExtension() + result = ctx.with_extensions(extension) + + assert len(extension.fallbacks) == 1 + assert result.session_id() == ctx.session_id() + + +def test_with_extensions_threads_the_planner_through_in_order(ctx): + """Each planner hook receives what the previous one returned. + + Planners nest rather than chain, so the host hands each bundle the planner + built so far. Argument order is nesting order, last one outermost. + """ + calls = [] + first, second = _PlannerExtension(calls), _PlannerExtension(calls) + ctx.with_extensions(first, second) + + # Argument order, once each. Nothing else pins the order: both hooks + # return capsules, and a host that ran them backwards would still leave + # each with one fallback recorded. + assert calls == [first, second] + + # `first` returned its fallback unchanged, but the host re-exports every + # hook's return value before handing it on, so `second` receives a capsule + # of its own rather than the object `first` was handed. + assert second.fallbacks[0] is not first.fallbacks[0] + + # That is as far as pure Python reaches: a capsule is opaque, so this + # cannot tell a re-export of `first`'s planner from a fresh read of the + # session's. `test_with_extensions_nests_planners_in_argument_order` in + # examples/datafusion-ffi-query-planner-example is what pins the nesting, + # by asserting the outer planner delegated to the inner one. + + +def test_with_extensions_planner_hook_sees_the_new_handle(ctx): + """Phase two runs against the handle carrying the final codec chains. + + A planner captured against the pre-install handle would encode through a + chain missing every codec this call installed. + """ + codecs = _CodecOnlyExtension() + planner = _PlannerExtension() + result = ctx.with_extensions(codecs, planner) + + assert planner.planner_ctx.logical_extension_codec_ids() == ["my_library.logical"] + assert result.logical_extension_codec_ids() == ["my_library.logical"] + + +def test_with_extensions_skips_a_planner_hook_returning_none(ctx): + """Returning ``None`` contributes no planner and keeps the fallback. + + The skip is what lets the call succeed at all: a host that treated the + ``None`` as a contribution would hand it to the export step and fail with + ``'None' is not an instance of 'PyCapsule'`` before ``downstream`` ran. + """ + + class NoPlanner: + def __init__(self): + self.fallbacks = [] + + def __datafusion_session_planner__(self, ctx, fallback): + self.fallbacks.append(fallback) + # Spelled out rather than left to fall off the end: `None` is the + # protocol's "contribute no planner", which is what this test is + # about, and an implicit one would read as an oversight. + return None # noqa: RET501, PLR1711 + + skipped = NoPlanner() + downstream = _PlannerExtension() + result = ctx.with_extensions(skipped, downstream) + + # Both hooks ran, and `downstream` was handed a planner rather than the + # `None` in front of it. It is a fresh read of the session's planner, not + # the object `skipped` was given, so a host that fell back by reusing the + # previous hook's *input* is ruled out too. + assert len(skipped.fallbacks) == 1 + assert len(downstream.fallbacks) == 1 + assert downstream.fallbacks[0] is not skipped.fallbacks[0] + + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_rejects_bad_codec_capsule(ctx): + """A correctly shaped object still has to return the right capsule.""" + + class BadCodecExtension: + def __datafusion_session_extension__(self, ctx): + wrong_capsule = ctx.__datafusion_task_context_provider__() + return SessionExtensionComponents( + logical_extension_codecs=( + _NamedCodec(wrong_capsule, "my_library.logical"), + ), + ) + + with pytest.raises( + ValueError, match="Expected name 'datafusion_logical_extension_codec'" + ): + ctx.with_extensions(BadCodecExtension()) + + +def test_with_extensions_rejects_a_bare_capsule_codec(ctx): + """A codec must be an object that can name itself, not a bare capsule. + + An id is read off the object a codec is handed over as, and a capsule has + no type to read one from. ``with_extensions`` takes no ``codec_id=``, so + the capsule is refused here rather than given an id derived from something + that is not the codec. + """ + + class BareCapsuleExtension: + def __init__(self): + self.exporter = SessionContext() + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=( + self.exporter.__datafusion_logical_extension_codec__(), + ), + ) + + with pytest.raises( + TypeError, + match="must be an object exposing `__datafusion_logical_extension_codec__`", + ): + ctx.with_extensions(BareCapsuleExtension()) + + +def test_with_extensions_rejects_a_bare_physical_capsule_codec(ctx): + """The physical getter is named in its own diagnostic.""" + + class BareCapsuleExtension: + def __init__(self): + self.exporter = SessionContext() + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + physical_extension_codecs=( + self.exporter.__datafusion_physical_extension_codec__(), + ), + ) + + with pytest.raises( + TypeError, + match="must be an object exposing `__datafusion_physical_extension_codec__`", + ): + ctx.with_extensions(BareCapsuleExtension()) + + +def test_with_extensions_codec_ids_survive_composition(ctx): + """A codec keeps its id when its extension is nested inside another one. + + Wire ids have to mean the same thing in whichever process decodes, so + packaging one extension inside another — the natural way for an + application to present several libraries as one — must not re-tag the + inner library's payloads. Reading the id off the handed-over object rather + than off the contributing extension is what guarantees that. + """ + + class ComposedExtension: + """Presents another extension's components as its own.""" + + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_extension__(self, ctx): + return self.inner.__datafusion_session_extension__(ctx) + + direct = ctx.with_extensions(_CodecOnlyExtension()) + wrapped = SessionContext().with_extensions(ComposedExtension(_CodecOnlyExtension())) + + assert direct.logical_extension_codec_ids() == ["my_library.logical"] + assert wrapped.logical_extension_codec_ids() == ["my_library.logical"] + assert direct.physical_extension_codec_ids() == ["my_library.physical"] + assert wrapped.physical_extension_codec_ids() == ["my_library.physical"] + + +def test_with_extensions_uses_ids_declared_on_the_codec(ctx): + """``__datafusion_codec_id__`` on the handed-over object names the codec. + + This is how an extension contributing more than one codec of a kind tells + them apart. + """ + + class TwoNamedCodecs: + def __init__(self): + self.exporter = SessionContext() + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=( + _NamedCodec( + self.exporter.__datafusion_logical_extension_codec__(), + "my_library.first", + ), + _NamedCodec( + self.exporter.__datafusion_logical_extension_codec__(), + "my_library.second", + ), + ), + ) + + result = ctx.with_extensions(TwoNamedCodecs()) + assert result.logical_extension_codec_ids() == [ + "my_library.first", + "my_library.second", + ] + + +def test_with_extensions_rejects_two_codecs_of_one_class(ctx): + """Two wrappers of one class claim one class-derived id, so they collide. + + Numbering them by position would be an id another library can mint the + same value from, and would break stored plans the first time the extension + reordered what it returns, so the ambiguity is refused instead. + """ + + class UnnamedCodec: + def __init__(self, capsule): + self._capsule = capsule + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule + + class TwoUnnamedCodecs: + def __init__(self): + self.exporter = SessionContext() + + def __datafusion_session_extension__(self, ctx): + capsule = self.exporter.__datafusion_logical_extension_codec__ + return SessionExtensionComponents( + logical_extension_codecs=( + UnnamedCodec(capsule()), + UnnamedCodec(capsule()), + ), + ) + + with pytest.raises(ValueError, match="__datafusion_codec_id__"): + ctx.with_extensions(TwoUnnamedCodecs()) + + +def test_with_extensions_leaves_an_exporting_object_its_own_id(ctx): + """A codec handed over as an object keeps the identity it declares.""" + exporter = SessionContext() + + class ObjectCodecExtension: + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents(logical_extension_codecs=(exporter,)) + + result = ctx.with_extensions(ObjectCodecExtension()) + assert result.logical_extension_codec_ids() == [exporter.__datafusion_codec_id__] + + +def test_with_extensions_installs_codecs_and_planner(ctx): + ctx.register_record_batches( + "extensions_test", + [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]], + ) + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension, _PlannerExtension()) + + assert result.table_exist("extensions_test") + # In-memory tables need a real extension codec to round-trip through the + # FFI planner, so query plans that don't serialize a table provider. + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_binds_to_the_receiving_session(ctx): + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension) + + # Factories are handed the receiver itself, so a component bound during + # installation targets the session the returned handle also wraps. There + # is no intermediate context that could be collected out from under it. + assert extension.bound_ctx is ctx + assert result.session_id() == ctx.session_id() + + # One session: a registration through either handle is visible to both. + ctx.register_record_batches( + "bound_test", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + assert result.table_exist("bound_test") + + +def test_with_extensions_survives_source_collection(): + extension = _CodecOnlyExtension() + result = SessionContext().with_extensions(extension, _PlannerExtension()) + gc.collect() + + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_failure_leaves_source_usable(ctx): + class BoomExtension: + def __datafusion_session_extension__(self, ctx): + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions(_CodecOnlyExtension(), BoomExtension()) + + batches = ctx.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py index fea4cc91f..9764e2973 100644 --- a/python/tests/test_imports.py +++ b/python/tests/test_imports.py @@ -94,6 +94,24 @@ def test_datafusion_python_version(): assert datafusion.__version__ is not None +def test_extension_protocols_are_exported_together(): + """The extension protocol family is reachable from the package root. + + ``QueryPlannerExportable`` types the planner a + ``__datafusion_session_planner__`` hook returns, so a bundle author needs + it exactly as much as the other three; leaving it in the submodule made + one member of one family import differently from the rest. + """ + for name in [ + "QueryPlannerExportable", + "SessionExtensionComponents", + "SessionExtensionExportable", + "SessionPlannerExportable", + ]: + assert name in datafusion.__all__, f"{name} missing from datafusion.__all__" + assert getattr(datafusion, name) is getattr(datafusion.extensions, name) + + def test_class_module_is_datafusion(): # context for klass in [ diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index cf6719ecf..6927632b9 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -28,6 +28,21 @@ from enum import EnumMeta as EnumType +# Internal methods a wrapper calls but does not re-export. Add to this only +# when the method exists to serve a public wrapper, never to silence a genuine +# gap in coverage. +PRIVATE_SUPPORT_METHODS = frozenset( + { + # The three steps of SessionContext.with_extensions: install the + # codecs, re-export each planner hook's return value as a capsule, + # commit the planner. + "_install_extension_codecs", + "_export_query_planner", + "_install_extension_planner", + } +) + + def _check_enum_exports(internal_obj, wrapped_obj) -> None: """Check that all enum values are present in wrapped object.""" expected_values = [v for v in dir(internal_obj) if not v.startswith("__")] @@ -67,6 +82,14 @@ def missing_exports(internal_obj, wrapped_obj) -> None: pytest.fail(f"Missing __repr__: {internal_obj.__name__}") for internal_attr_name in dir(internal_obj): + # Private support methods that exist only for a wrapper to call, so + # they are not part of the public surface and need no wrapper of their + # own. Listed rather than matched by leading underscore, which would + # also excuse names like `_repr_html_` that a wrapper does have to + # provide. + if internal_attr_name in PRIVATE_SUPPORT_METHODS: + continue + wrapped_attr_name = internal_attr_name.removeprefix("Raw") assert wrapped_attr_name in dir(wrapped_obj)