diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 294ebfb3a..37be0bc9d 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_components__` +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 @@ -185,7 +223,12 @@ pins that; changing it should be deliberate. ## Where the truth is -- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. +- `docs/source/extension-guide/` — the protocol, for the library author. + `capsule-protocol.md` has the hook convention and what the getter argument + actually is; `codecs.md`, `bundles.md`, and `query-planners.md` have the + per-component rules; `index.md` lists all 18 hooks. +- `docs/source/contributor-guide/ffi-internals.md` — why the framing is shaped + this way, including the weak-`Arc` scheme and the one-level rebind. - `docs/source/user-guide/upgrade-guides.md` — every past migration. - `crates/core/src/codec.rs` — the codec chain: the envelope, identity dispatch, and the two unframed cases from Rule 8. diff --git a/AGENTS.md b/AGENTS.md index 659094ec0..f61b05636 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -133,6 +133,47 @@ Every Python function must include a docstring with usage examples. `array_sort`) only need a one-line description and a `See Also` reference to the primary function. They do not need their own examples. +### One canonical home per claim + +A claim lives where the reader already is when they need it — **exactly once**. + +- A property of one callable's arguments, return value, or errors → that + callable's docstring. +- A property of a *type* that several callables share → the class docstring. + (Session sharing and context lifetime live on `SessionContext`, not on each + of the five `with_*` methods.) +- **Why** the API is shaped this way, a multi-library recipe, Rust-side code, a + trade-off, or a limitation with an upstream issue → the narrative guide under + `docs/source/`. + +Everywhere else: one sentence plus one Sphinx role. A docstring may *state* a +claim the guide also makes; it must not *argue* it — no "because", no "the +reason is", no counter-argument. + +**The test:** if a paragraph would survive being moved into the guide unchanged, +move it. `python/tests/test_docstrings.py` enforces a length ceiling, which is +the symptom this rule treats. + +When you point at the guide, use a real `:ref:` to a specific label. Prose +saying "see the extensions guide" with no role is a dead end in the rendered +HTML. + +### Examples that need a compiled extension + +Some APIs cannot be demonstrated without a built FFI extension library, which +this package does not ship. The convention is: + +1. A **runnable** example block first, using only the wheel. A `SessionContext` + satisfies the capsule-getter protocols, so its own exported capsule stands in + for a real library's in a doctest. +2. Then a `# doctest: +SKIP` block showing real usage, at most a few lines, with + one line of prose naming the test that runs it for real. +3. Every `+SKIP` block needs that mirror. See + `test_with_extensions_docstring_example_still_runs` in + `examples/datafusion-ffi-query-planner-example`, which parses the live + docstring, drops the skip, and executes it — so a renamed method or a wrong + expected output fails there. + ## Aggregate and Window Function Documentation When adding or updating an aggregate or window function, ensure the corresponding 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/codec.rs b/crates/core/src/codec.rs index 8f43bdb5e..cca45147d 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -660,7 +660,7 @@ impl PythonLogicalCodec { /// `cloudpickle.loads` on the inline `DFPY*` payload. It does /// **not** make `pickle.loads(untrusted_bytes)` safe; treat every /// `pickle.loads` on untrusted input as unsafe regardless of this - /// setting. See `docs/source/user-guide/io/distributing_work.rst` + /// setting. See `docs/source/user-guide/distributing-work/expressions.md` /// (Security section) for the full threat model, and Python's /// [pickle module security warning][1] for why `pickle.loads` is /// unsafe in general. 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..c9f845f12 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -54,8 +54,21 @@ # raises an extension conflict. "myst_nb", "autoapi.extension", + # Emits a meta-refresh stub at each old docname listed in `redirects` + # below, so inbound links to pages that have moved keep working. + "sphinx_reredirects", ] +# Old page URLs that have moved. The site is single-version (each release +# overwrites asf-site wholesale), so these exist purely for inbound external +# links — issue comments, release blog posts, and README links that pinned a +# heading anchor. Keys are docnames without a suffix; the source file must be +# gone, or Sphinx builds the real page and the stub is never written. +redirects = { + "user-guide/distributing-work": "distributing-work/index.html", + "contributor-guide/ffi": "../extension-guide/index.html", +} + # NOTE: .rst stays alongside .md because sphinx-autoapi generates RST # under autoapi/ and Sphinx needs the suffix to parse it. The human- # authored docs are all MyST .md now. ".md" is routed through myst-nb so @@ -90,6 +103,11 @@ # This pattern also affects html_static_path and html_extra_path. exclude_patterns = [] +# `autoapi_options` is deliberately left at the default. In particular +# `special-members` is load-bearing: it is what publishes the +# `__datafusion_*__` capsule getters, which are the public protocol an +# extension library implements against. Dropping it would delete that +# reference surface along with `DataFrame.__init__` and the Arrow dunders. autoapi_dirs = ["../../python"] autoapi_ignore = ["*tests*"] autoapi_member_order = "groupwise" @@ -103,6 +121,10 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool: # noqa # Re-exports ("class", "datafusion.DataFrame"), ("class", "datafusion.SessionContext"), + ("class", "datafusion.QueryPlannerExportable"), + ("class", "datafusion.SessionExtensionComponents"), + ("class", "datafusion.SessionComponentsExportable"), + ("class", "datafusion.SessionPlannerExportable"), ("module", "datafusion.common"), # Duplicate modules (skip module-level docs to avoid duplication) ("module", "datafusion.col"), @@ -199,8 +221,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-internals.md b/docs/source/contributor-guide/ffi-internals.md new file mode 100644 index 000000000..75c32444c --- /dev/null +++ b/docs/source/contributor-guide/ffi-internals.md @@ -0,0 +1,137 @@ + + +(ffi_internals)= + +# FFI framing internals + +Read this before changing how datafusion-python frames FFI components. You do +**not** need it to write an extension library — that is the +{ref}`Extension Guide `. + +The invariants here are the reasons the extension-facing rules are shaped the +way they are. Each one has an "obvious" simplification that does not work, and +the point of this page is to record why, so the next person does not spend a +week rediscovering it. + +For the wire format itself — how a codec id is stored alongside a payload, +routed back on decode, and which two cases stay unframed — see the module +documentation in `crates/core/src/codec.rs`. It is the authority; this page +does not restate it. + +(ffi_internals_one_arc)= + +## One session, one `Arc` + +Every codec handed to a foreign object carries an `FFI_TaskContextProvider`, +and that type holds its provider **weakly**. A registered catalog provider +upgrades the handle on every `supports_filters_pushdown` and every `scan`. +Those handles are bound to one particular `Arc` allocation, not +to the logical session, so anything that replaces the allocation orphans all of +them and the next query fails with `TaskContextProvider went out of scope over +FFI boundary`. + +So a `PySessionContext` keeps the `Arc` it was created with for +its whole life. Installing a query planner writes the new `SessionState` back +through `state_ref()`, exactly as `add_physical_optimizer_rule` does, rather +than deriving a replacement context. The session id is carried across that +rewrite — `SessionStateBuilder` mints a fresh one otherwise — so `session_id()` +and every `TaskContext` the session hands out keep agreeing. + +Repairing the damage instead of avoiding it does not work in general. A context +can rebuild the codecs it holds in its own fields, but a codec already embedded +in a registered `FFI_CatalogProvider` — and in every `FFI_SchemaProvider` and +`FFI_TableProvider` minted from it — is not reachable from Python at all. Nor +can the codec simply retain the session that built it: a codec handed to a +provider is routinely registered straight back into that same session, which +would close the cycle +`SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and +leak it. + +`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. 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. + +(ffi_internals_rebinding)= + +## Why planner codec rebinding is one level deep + +Installing a codec on a session that already has a foreign planner rebuilds +that planner against the new chain. The rebuild swaps the codecs on the +installed `ForeignQueryPlanner` handle, and only that handle. A planner that +wraps a fallback resolved that fallback when *it* was installed, and holds the +result inside its own library's private data — behind a `create_physical_plan` +function pointer, with no Python-side handle. A codec installed afterwards +therefore reaches the outer planner and not the fallback, which keeps whichever +codecs were in force when it was imported. + +Neither side can repair that: + +- **The host cannot reach it.** `FFI_QueryPlanner::new_with_ffi_codecs` unwraps + exactly one `ForeignQueryPlanner` layer. There is no deeper handle to unwrap + — the same situation as a codec embedded in a registered + `FFI_CatalogProvider`. +- **The planner library cannot re-derive it.** `FFI_QueryPlanner` holds its + codecs by value, and `Session` exposes no accessor for the ones the host + currently has, so `create_physical_plan` cannot pick them up from the session + it is handed. The rebuild has to be eager, and an eager rebuild only sees the + top layer. + +A fix has to come from upstream, and is tracked in +[apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762). + +The stale codecs stay usable rather than dangling — they hold weak handles to +the one `Arc` that the previous section keeps alive — so the +effect is a fallback hop serializing with an older codec, not a failure. It is +also invisible to the examples in this repository, which use one fallback in +the same cdylib as its wrapper; `datafusion-ffi` short-circuits a same-library +hop rather than serializing, so no codec runs. A fallback in a *different* +library would serialize, and would do it with the codecs it was imported with. + +The extension-facing consequence — install codecs before a layered planner, and +prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`. + +## Two argument kinds for one convention + +`CapsuleGetterArg` in `crates/util/src/lib.rs` distinguishes three cases: no +argument, the session, and the host's logical extension codec as a bare +capsule. Provider and catalog registration passes the codec; the extension +codecs, the query planner, and table functions get the session. + +The distinction exists so that a `TypeError` from a getter that refused its +argument can be rewritten into an `ImportError` naming what the getter *should* +accept. Getting that message wrong sends an extension author to fix the wrong +signature, which is why every capsule getter routes through +`call_capsule_getter` rather than calling `getattr` directly. Three importers +previously each had their own copy of that logic and each missed later +corrections to it. + +When adding a hook, decide which arm it needs by what its FFI constructor +requires: a task-context provider means it must have the session, since you +cannot get a provider off a capsule. A codec alone means either works, and +passing the codec keeps the host from handing out session handles it does not +need to. + +The extension-facing statement of this is {ref}`extension_getter_argument`, +which deliberately describes the argument by capability — "something you can +read the host's logical codec off" — rather than by type. diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md deleted file mode 100644 index d86858a83..000000000 --- a/docs/source/contributor-guide/ffi.md +++ /dev/null @@ -1,545 +0,0 @@ - - -(ffi)= - -# Python Extensions - -The DataFusion in Python project is designed to allow users to extend its functionality in a few core -areas. Ideally many users would like to package their extensions as a Python package and easily -integrate that package with this project. This page serves to describe some of the challenges we face -when doing these integrations and the approach our project uses. - -## The Primary Issue - -Suppose you wish to use DataFusion and you have a custom data source that can produce tables that -can then be queried against, similar to how you can register a {ref}`CSV ` or -{ref}`Parquet ` file. In DataFusion terminology, you likely want to implement a -{ref}`Custom Table Provider `. In an effort to make your data source -as performant as possible and to utilize the features of DataFusion, you may decide to write -your source in Rust and then expose it through [PyO3](https://pyo3.rs) as a Python library. - -At first glance, it may appear the best way to do this is to add the `datafusion-python` -crate as a dependency, provide a `PyTable`, and then to register it with the -`SessionContext`. Unfortunately, this will not work. - -When you produce your code as a Python library and it needs to interact with the DataFusion -library, at the lowest level they communicate through an Application Binary Interface (ABI). -The acronym sounds similar to API (Application Programming Interface), but it is distinctly -different. - -The ABI sets the standard for how these libraries can share data and functions between each -other. One of the key differences between Rust and other programming languages is that Rust -does not have a stable ABI. What this means in practice is that if you compile a Rust library -with one version of the `rustc` compiler and I compile another library to interface with it -but I use a different version of the compiler, there is no guarantee the interface will be -the same. - -In practice, this means that a Python library built with `datafusion-python` as a Rust -dependency will generally **not** be compatible with the DataFusion Python package, even -if they reference the same version of `datafusion-python`. If you attempt to do this, it may -work on your local computer if you have built both packages with the same optimizations. -This can sometimes lead to a false expectation that the code will work, but it frequently -breaks the moment you try to use your package against the released packages. - -You can find more information about the Rust ABI in their -[online documentation](https://doc.rust-lang.org/reference/abi.html). - -## The FFI Approach - -Rust supports interacting with other programming languages through it's Foreign Function -Interface (FFI). The advantage of using the FFI is that it enables you to write data structures -and functions that have a stable ABI. The allows you to use Rust code with C, Python, and -other languages. In fact, the [PyO3](https://pyo3.rs) library uses the FFI to share data -and functions between Python and Rust. - -The approach we are taking in the DataFusion in Python project is to incrementally expose -more portions of the DataFusion project via FFI interfaces. This allows users to write Rust -code that does **not** require the `datafusion-python` crate as a dependency, expose their -code in Python via PyO3, and have it interact with the DataFusion Python package. - -Early adopters of this approach include [delta-rs](https://delta-io.github.io/delta-rs/) -who has adapted their Table Provider for use in `` `datafusion-python` `` with only a few lines -of code. Also, the DataFusion Python project uses the existing definitions from -[Apache Arrow CStream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html) -to support importing **and** exporting tables. Any Python package that supports reading -the Arrow C Stream interface can work with DataFusion Python out of the box! You can read -more about working with Arrow sources in the {ref}`Data Sources ` -page. - -To learn more about the Foreign Function Interface in Rust, the -[Rustonomicon](https://doc.rust-lang.org/nomicon/ffi.html) is a good resource. - -## Inspiration from Arrow - -DataFusion is built upon [Apache Arrow](https://arrow.apache.org/). The canonical Python -Arrow implementation, [pyarrow](https://arrow.apache.org/docs/python/index.html) provides -an excellent way to share Arrow data between Python projects without performing any copy -operations on the data. They do this by using a well defined set of interfaces. You can -find the details about their stream interface -[here](https://arrow.apache.org/docs/format/CStreamInterface.html). The -[Rust Arrow Implementation](https://github.com/apache/arrow-rs) also supports these -`C` style definitions via the Foreign Function Interface. - -In addition to using these interfaces to transfer Arrow data between libraries, `pyarrow` -goes one step further to make sharing the interfaces easier in Python. They do this -by exposing PyCapsules that contain the expected functionality. - -You can learn more about PyCapsules from the official -[Python online documentation](https://docs.python.org/3/c-api/capsule.html). PyCapsules -have excellent support in PyO3 already. The -[PyO3 online documentation](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule) is a good source -for more details on using PyCapsules in Rust. - -Two lessons we leverage from the Arrow project in DataFusion Python are: - -- We reuse the existing Arrow FFI functionality wherever possible. -- We expose PyCapsules that contain a FFI stable struct. - -## Implementation Details - -The bulk of the code necessary to perform our FFI operations is in the upstream -[DataFusion](https://datafusion.apache.org/) core repository. You can review the code and -documentation in the [datafusion-ffi] crate. - -Our FFI implementation is narrowly focused at sharing data and functions with Rust backed -libraries. This allows us to use the [abi_stable crate](https://crates.io/crates/abi_stable). -This is an excellent crate that allows for easy conversion between Rust native types -and FFI-safe alternatives. For example, if you needed to pass a `Vec` via FFI, -you can simply convert it to a `RVec` in an intuitive manner. It also supports -features like `RResult` and `ROption` that do not have an obvious translation to a -C equivalent. - -The [datafusion-ffi] crate has been designed to make it easy to convert from DataFusion -traits into their FFI counterparts. For example, if you have defined a custom -[TableProvider](https://docs.rs/datafusion/45.0.0/datafusion/catalog/trait.TableProvider.html) -and you want to create a sharable FFI counterpart, you could write: - -```rust -let my_provider = MyTableProvider::default(); -let ffi_provider = FFI_TableProvider::new(Arc::new(my_provider), false, None); -``` - -(ffi_pyclass_mutability)= - -## PyO3 class mutability guidelines - -PyO3 bindings should present immutable wrappers whenever a struct stores shared or -interior-mutable state. In practice this means that any `#[pyclass]` containing an -`Arc>` or similar synchronized primitive must opt into `#[pyclass(frozen)]` -unless there is a compelling reason not to. - -The execution context illustrates the preferred pattern. `PySessionContext` in -{file}`src/context.rs` stays frozen even though it shares mutable state internally via -`SessionContext`. This ensures PyO3 tracks borrows correctly while Python-facing APIs -clone the inner `SessionContext` or return new wrappers instead of mutating the -existing instance in place: - -```rust -#[pyclass(from_py_object, frozen, name = "SessionContext", module = "datafusion", subclass)] -#[derive(Clone)] -pub struct PySessionContext { - pub ctx: SessionContext, -} -``` - -Occasionally a type must remain mutable—for example when PyO3 attribute setters need to -update fields directly. In these rare cases add an inline justification so reviewers and -future contributors understand why `frozen` is unsafe to enable. `DataTypeMap` in -{file}`src/common/data_type.rs` includes such a comment because PyO3 still needs to track -field updates: - -```rust -// TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now -#[derive(Debug, Clone)] -#[pyclass(from_py_object, name = "DataTypeMap", module = "datafusion.common", subclass)] -pub struct DataTypeMap { - #[pyo3(get, set)] - pub arrow_type: PyDataType, - #[pyo3(get, set)] - pub python_type: PythonType, - #[pyo3(get, set)] - pub sql_type: SqlType, -} -``` - -When reviewers encounter a mutable `#[pyclass]` without a comment, they should request -an explanation or ask that `frozen` be added. Keeping these wrappers frozen by default -helps avoid subtle bugs stemming from PyO3's interior mutability tracking. - -If you were interfacing with a library that provided the above `FFI_TableProvider` and -you needed to turn it back into an `TableProvider`, you can turn it into a -`ForeignTableProvider` with implements the `TableProvider` trait. - -```rust -let foreign_provider: ForeignTableProvider = ffi_provider.into(); -``` - -If you review the code in [datafusion-ffi] you will find that each of the traits we share -across the boundary has two portions, one with a `FFI_` prefix and one with a `Foreign` -prefix. This is used to distinguish which side of the FFI boundary that struct is -designed to be used on. The structures with the `FFI_` prefix are to be used on the -**provider** of the structure. In the example we're showing, this means the code that has -written the underlying `TableProvider` implementation to access your custom data source. -The structures with the `Foreign` prefix are to be used by the receiver. In this case, -it is the `datafusion-python` library. - -In order to share these FFI structures, we need to wrap them in some kind of Python object -that can be used to interface from one package to another. As described in the above -section on our inspiration from Arrow, we use `PyCapsule`. We can create a `PyCapsule` -for our provider thusly: - -```rust -let name = CString::new("datafusion_table_provider")?; -let my_capsule = PyCapsule::new_bound(py, provider, Some(name))?; -``` - -On the receiving side, turn this pycapsule object into the `FFI_TableProvider`, which -can then be turned into a `ForeignTableProvider` the associated code is: - -```rust -let capsule = capsule.cast::()?; -let data: NonNull = capsule - .pointer_checked(Some(name))? - .cast(); -let codec = unsafe { data.as_ref() }; -``` - -By convention the `datafusion-python` library expects a Python object that has a -`TableProvider` PyCapsule to have this capsule accessible by calling a function named -`__datafusion_table_provider__`. You can see a complete working example of how to -share a `TableProvider` from one python library to DataFusion Python in the -[repository examples folder](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example). - -This section has been written using `TableProvider` as an example. It is the first -extension that has been written using this approach and the most thoroughly implemented. -As we continue to expose more of the DataFusion features, we intend to follow this same -design pattern. - -## Query Planners Across Multiple Libraries - -A query can involve three independent native libraries: `datafusion-python`, a library -that owns table providers or functions, and a library that owns the query planner. The -examples use two separate extension crates so each role has a distinct shared-library -identity: - -- [`datafusion-ffi-example`] owns providers, functions, and their codecs. -- [`datafusion-ffi-query-planner-example`] owns the planner and its configuration. - -The `SessionContext` owns the codecs used for the exchange and supplies them to the -foreign planner. This lets the planner decode provider-owned objects and lets -`datafusion-python` decode the physical plan returned by the planner. The examples use -process-local tokens to demonstrate ownership; production codecs should serialize -durable metadata instead. - -### Composable codecs - -Extension codecs compose. Each call to `with_logical_extension_codec` or -`with_physical_extension_codec` appends the codec to the session's codec chain -rather than replacing prior codecs. - -**Nothing is asked of the codec itself.** Implement `LogicalExtensionCodec` or -`PhysicalExtensionCodec` exactly as you would for a session that installs only -yours. When your codec writes bytes into a serialized plan, datafusion-python -records which codec wrote them, and strips that record off again before handing the -bytes back. So your codec receives, byte for byte, the payload it wrote, and is -never offered a payload another codec wrote. - -A codec that also ships to hosts which dispatch differently may still want its own -guard against foreign payloads. Keeping one is fine; it is simply not needed for the -datafusion-python path. - -That record is the codec's **id**: a short string stored inside the plan, naming the -codec that wrote each payload. Because plans are decoded in another process — or -another program — the id has to name the same codec there as it did where the plan -was written. - -Ids are assigned for you. A codec's id is normally its exporting class's import -path, such as `my_library.Codec`, which is what you will see in -`logical_extension_codec_ids()` and in decode errors. You choose one yourself in -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 - `codec_id=` if those plans have to cross sessions. -- **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. - -`SessionContext.logical_extension_codec_ids()` and its physical counterpart list the -ids installed on a session, which is also what a decode failure names. - -Installing one context's codec stack on another session composes the two sessions -rather than copying codecs out of one: the imported codecs resolve their task context -against the original and stop working when it is dropped — see -[One session, one `Arc`](#one-session-one-arcsessioncontext). Pass -the context itself rather than the capsule it exports, so its codecs get an id that -other sessions can decode. - -Because decoding keys off the id rather than install position, registration order -between independent libraries does not affect decoding at all. It is visible only -on encoding, where codecs are consulted in install order and the first to claim an -object wins — so installing a library can claim objects nothing else claimed, but -never takes over an object an earlier codec was already encoding. Two libraries -that each own tables, functions, and a planner register like this: - -```python -ctx = SessionContext(config) - -# Codecs from both libraries. Order between libraries does not matter. -ctx = ctx.with_logical_extension_codec(lib_a.codec()) -ctx = ctx.with_logical_extension_codec(lib_b.codec()) -ctx = ctx.with_physical_extension_codec(lib_a.physical_codec()) -ctx = ctx.with_physical_extension_codec(lib_b.physical_codec()) - -# A session holds one planner, so layering is explicit delegation. Install the -# codecs first: the fallback captured here keeps the codecs it was exported -# with. See "Rebinding a planner's codecs is one level deep" below. -ctx.set_query_planner(lib_a.Planner()) -ctx.set_query_planner(lib_b.Planner(fallback=ctx.__datafusion_query_planner__())) - -# Tables and functions — any time before the first query. -ctx.register_table("t", lib_a.TableProvider()) -ctx.register_udf(udf(lib_b.SomeUDF())) -``` - -A codec may own functions that need no payload at all, where the name is the whole -encoding: `try_encode_udf` writes nothing and `try_decode_udf` rebuilds the function -from `name`. That is supported and needs no id, because an `Ok` with an empty -buffer is read as "no opinion" and passes the object to the next codec. -`NameOnlyUdfCodec` in the FFI example is the worked case. Anything no installed -codec claims falls through to `Default{Logical,Physical}ExtensionCodec`. - -This is the one case where your decoder is consulted about something you may not -own, because an empty payload has no id to route on. `try_decode_udf` and -its aggregate and window siblings can therefore be called with an empty `buf` and a -`name` belonging to another library. Decide from `name` and return an error if it is -not yours; do not assume `buf` is non-empty. - -The framing itself — how an id is stored alongside a payload and routed back, and the two cases -that stay unframed — is internal to datafusion-python and documented in -`crates/core/src/codec.rs` for anyone changing it. - -The current FFI logical codec supports providers and UDFs but not arbitrary custom -`LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and -local build commands. - -### Capsule getters receive the session they are installed on - -`__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and -`__datafusion_physical_extension_codec__` all take the `SessionContext` the object is -being installed on, the same way `__datafusion_table_provider__` does: - -```rust -fn __datafusion_physical_extension_codec__<'py>( - &self, - py: Python<'py>, - session: Bound<'py, PyAny>, -) -> PyResult> { - let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; - let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); - PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") -} -``` - -This exists because the FFI constructors need things an extension library does not -have. `FFI_{Logical,Physical}ExtensionCodec::new` needs a `TaskContextProvider` for the -decode callbacks the codec will receive, and `FFI_QueryPlanner::new` needs both codecs -on top of that. Taking them from the session is what keeps a library from constructing -a `SessionContext` purely to satisfy a parameter — an empty one resolves nothing, and -`FFI_TaskContextProvider` holds it weakly, so a context built inline in the getter is -already dropped by the time the capsule is used. - -A planner uses `FFI_QueryPlanner::new_with_ffi_codecs` with the two codecs it takes off -the session, and never touches a provider directly. That also matches what installation -does anyway: `set_query_planner` builds the planner against the codecs of the session -that will run the query. - -`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 -`ctx.__datafusion_query_planner__(ctx)` are both fine. - -### A codec decodes against the session that is running the query - -Because the provider comes from the host, a decode callback running inside an extension -library resolves names against the session running the query. A function registered with -`ctx.register_udf(...)` is visible to a foreign codec decoding a node that references it -by name, and the handle is live rather than a snapshot, so a registration made after the -codec is installed is visible too. - -This is covered in -`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`, -where the example codecs take a `require_udf_on_decode` name and resolve it out of the -task context they are handed. - -### One session, one `Arc` - -Every codec handed to a foreign object carries an `FFI_TaskContextProvider`, and that -type holds its provider **weakly**. A registered catalog provider upgrades the handle on -every `supports_filters_pushdown` and every `scan`. Those handles are bound to one -particular `Arc` allocation, not to the logical session, so anything that -replaces the allocation orphans all of them and the next query fails with -`TaskContextProvider went out of scope over FFI boundary`. - -So a `PySessionContext` keeps the `Arc` it was created with for its whole -life. Installing a query planner writes the new `SessionState` back through -`state_ref()`, exactly as `add_physical_optimizer_rule` does, rather than deriving a -replacement context. The session id is carried across that rewrite — `SessionStateBuilder` -mints a fresh one otherwise — so `session_id()` and every `TaskContext` the session hands -out keep agreeing. - -Repairing the damage instead of avoiding it does not work in general. A context can -rebuild the codecs it holds in its own fields, but a codec already embedded in a -registered `FFI_CatalogProvider` — and in every `FFI_SchemaProvider` and -`FFI_TableProvider` minted from it — is not reachable from Python at all. Nor can the -codec simply retain the session that built it: a codec handed to a provider is routinely -registered straight back into that same session, which would close the cycle -`SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and leak it. - -`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. - -### 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. - -`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 -visible to every context sharing that session — including ones a `with_*` call returned -earlier. Installing a codec on a session that already has a foreign planner rebuilds -that planner against the new chain for the same reason: there is one planner, and it has -to carry the codecs currently in force. This happens on the shared session, so it takes -effect even if the returned context is discarded — `ctx.with_python_udf_inlining(...)` -whose result is thrown away still leaves the session's planner carrying the codecs of -that discarded handle. A call that changes nothing is exempt: asking for the inlining -setting a context already has returns a handle without touching the session. - -The rule that falls out of this is worth stating on its own, because it is the one thing -that surprises people: - -> The session's query planner carries the codecs of the handle that most recently -> installed one. Every other path — `Expr.to_bytes(ctx)`, `ExecutionPlan.to_bytes(ctx)`, -> registering a provider — uses the codecs of the handle you call it on. - -Those can be different handles, and then one session has two codec chains in effect at -once: - -```python -ctx = ctx.with_logical_extension_codec(codec_a) -ctx.set_query_planner(planner) -ctx.with_logical_extension_codec(codec_b) # discarded - -Expr.to_bytes(expr, ctx) # encodes with [codec_a, default] -- ctx's own field -ctx.sql(...).collect() # plans with [codec_b, codec_a, default] -- the discarded - # handle's chain, installed on the shared session -``` - -Chaining `ctx = ctx.with_...(...)`, as the example below does, keeps the two in step. -`test_the_planner_and_the_handle_can_hold_different_codecs` pins the divergence. - -```python -ctx = SessionContext(config) -ctx = ctx.with_logical_extension_codec(provider_logical_codec) -ctx = ctx.with_physical_extension_codec(provider_physical_codec) -ctx.set_query_planner(planner) -ctx.register_udf(my_udf) -``` - -Order is a readability preference rather than a requirement — installing a codec after a -planner rebuilds the planner against it. - -A session holds exactly one query planner. Calling `set_query_planner` again replaces the -installed planner instead of layering another one. To chain planners, have the new -planner wrap the capsule returned by `SessionContext.__datafusion_query_planner__()`, -captured before the new planner is installed, and delegate to it explicitly. - -### Rebinding a planner's codecs is one level deep - -The rebuild above swaps the codecs on the installed `ForeignQueryPlanner` handle, and -only that handle. A planner that wraps a fallback resolved that fallback when *it* was -installed, and holds the result inside its own library's private data — behind a -`create_physical_plan` function pointer, with no Python-side handle. A codec installed -afterwards therefore reaches the outer planner and not the fallback, which keeps -whichever codecs were in force when it was imported. - -Neither side can repair that: - -- **The host cannot reach it.** `FFI_QueryPlanner::new_with_ffi_codecs` unwraps exactly - one `ForeignQueryPlanner` layer. There is no deeper handle to unwrap — the same - situation as a codec embedded in a registered `FFI_CatalogProvider`. -- **The planner library cannot re-derive it.** `FFI_QueryPlanner` holds its codecs by - value, and `Session` exposes no accessor for the ones the host currently has, so - `create_physical_plan` cannot pick them up from the session it is handed. The rebuild - has to be eager, and an eager rebuild only sees the top layer. - -A fix has to come from upstream, and is tracked in -[apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762). - -The stale codecs stay usable rather than dangling — they hold weak handles to the one -`Arc` that Rule 6 keeps alive — so the effect is a fallback hop -serializing with an older codec, not a failure. It is also invisible to the examples -here, which use one fallback in the same cdylib as its wrapper; `datafusion-ffi` -short-circuits a same-library hop rather than serializing, so no codec runs. A fallback -in a *different* library would serialize, and would do it with the codecs it was -imported with. - -So install the codecs before a layered planner. If a codec has to go in afterwards, -install the outer planner again *on the handle that holds the new codec* — that re-runs -its getter, which re-imports the fallback against that handle's codecs. Re-installing on -the original handle rebinds the session's planner back to the original handle's codecs -instead, which is the trap -`test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` pins. - -## Alternative Approach - -Suppose you needed to expose some other features of DataFusion and you could not wait -for the upstream repository to implement the FFI approach we describe. In this case -you decide to create your dependency on the `datafusion-python` crate instead. - -As we discussed, this is not guaranteed to work across different compiler versions and -optimization levels. If you wish to go down this route, there are two approaches we -have identified you can use. - -1. Re-export all of `datafusion-python` yourself with your extensions built in. -2. Carefully synchronize your software releases with the `datafusion-python` CI build - system so that your libraries use the exact same compiler, features, and - optimization level. - -We currently do not recommend either of these approaches as they are difficult to -maintain over a long period. Additionally, they require a tight version coupling -between libraries. - -## Status of Work - -At the time of this writing, the FFI features are under active development. To see -the latest status, we recommend reviewing the code in the [datafusion-ffi] crate. - -[datafusion-ffi]: https://crates.io/crates/datafusion-ffi -[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example -[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index a989a068d..45787f98f 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -19,11 +19,17 @@ # Contributor Guide -Guides for contributors to the DataFusion in Python project. +Guides for contributors to the DataFusion in Python project — people changing +this repository itself. + +If you are writing a separate library that plugs into datafusion-python, you +want the {ref}`Extension Guide ` instead. If you are using +one, see {ref}`user_guide_extensions`. ```{toctree} :maxdepth: 2 introduction -ffi +pyo3-guidelines +ffi-internals ``` diff --git a/docs/source/contributor-guide/pyo3-guidelines.md b/docs/source/contributor-guide/pyo3-guidelines.md new file mode 100644 index 000000000..4f81e569c --- /dev/null +++ b/docs/source/contributor-guide/pyo3-guidelines.md @@ -0,0 +1,74 @@ + + +(pyo3_guidelines)= + +# PyO3 binding guidelines + +Conventions for the `#[pyclass]` bindings in `crates/`. These are review +policy for this repository, not part of the extension protocol — an extension +library is free to shape its own classes differently. + +(ffi_pyclass_mutability)= + +## Class mutability + +PyO3 bindings should present immutable wrappers whenever a struct stores shared +or interior-mutable state. In practice this means that any `#[pyclass]` +containing an `Arc>` or similar synchronized primitive must opt into +`#[pyclass(frozen)]` unless there is a compelling reason not to. + +The execution context illustrates the preferred pattern. `PySessionContext` in +{file}`src/context.rs` stays frozen even though it shares mutable state +internally via `SessionContext`. This ensures PyO3 tracks borrows correctly +while Python-facing APIs clone the inner `SessionContext` or return new +wrappers instead of mutating the existing instance in place: + +```rust +#[pyclass(from_py_object, frozen, name = "SessionContext", module = "datafusion", subclass)] +#[derive(Clone)] +pub struct PySessionContext { + pub ctx: SessionContext, +} +``` + +Occasionally a type must remain mutable—for example when PyO3 attribute setters +need to update fields directly. In these rare cases add an inline justification +so reviewers and future contributors understand why `frozen` is unsafe to +enable. `DataTypeMap` in {file}`src/common/data_type.rs` includes such a +comment because PyO3 still needs to track field updates: + +```rust +// TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now +#[derive(Debug, Clone)] +#[pyclass(from_py_object, name = "DataTypeMap", module = "datafusion.common", subclass)] +pub struct DataTypeMap { + #[pyo3(get, set)] + pub arrow_type: PyDataType, + #[pyo3(get, set)] + pub python_type: PythonType, + #[pyo3(get, set)] + pub sql_type: SqlType, +} +``` + +When reviewers encounter a mutable `#[pyclass]` without a comment, they should +request an explanation or ask that `frozen` be added. Keeping these wrappers +frozen by default helps avoid subtle bugs stemming from PyO3's interior +mutability tracking. diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md new file mode 100644 index 000000000..791fdfc8e --- /dev/null +++ b/docs/source/extension-guide/bundles.md @@ -0,0 +1,305 @@ + + +(extension_bundles)= + +# Extension bundles + +If your library ships codecs, or a query planner, or both, expose a **bundle** +and let callers install it with +{py:meth}`~datafusion.SessionContext.with_extensions`. This is the recommended +way to package an extension, and the rest of this page explains what the +bundle protocol asks of you and why. + +## Why not have callers install the pieces + +Installing the pieces by hand 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. + +`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_components__(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. The caller then +writes: + +```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())) +``` + +`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 +{py:class}`~datafusion.SessionExtensionComponents`. + +## A bundle you can run + +Before wiring up a cdylib, it is worth seeing the protocol work end to end in +pure Python. This runs against the plain wheel — the codec here re-exports the +host session's own capsule, where a real one would return its library's: + +```python +from datafusion import SessionContext, SessionExtensionComponents + + +class Codec: + """Wraps a capsule so it carries an id. A Rust library ships this shape.""" + + def __init__(self, codec_id, capsule): + self.__datafusion_codec_id__ = codec_id + self._capsule = capsule + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule + + +class Bundle: + def __init__(self, codec_id): + self.codec_id = codec_id + + def __datafusion_session_components__(self, ctx): + # Fresh components on every call, bound to the `ctx` handed in. + # Never cache these, and never retain `ctx`. + return SessionExtensionComponents( + logical_extension_codecs=( + Codec(self.codec_id, ctx.__datafusion_logical_extension_codec__()), + ) + ) + + +ctx = SessionContext().with_extensions(Bundle("tables.v1"), Bundle("engine.v1")) +ctx.logical_extension_codec_ids() +# ['tables.v1', 'engine.v1'] +``` + +Three things this makes observable, each pinned by a test in +`python/tests/test_context.py`: + +- **Ids accumulate in bundle order.** Decoding does not depend on that order; + only encoding does. See {ref}`extension_codec_order`. +- **Two bundles claiming one id are refused**, with a `ValueError` naming the + id — not resolved by position, since a positional id would break stored plans + the first time a bundle reordered what it returns. +- **A hook that raises leaves the receiving session untouched.** Add a bundle + whose hook raises and the source context's + {py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` is still + empty afterwards. + +(extension_bundles_two_phases)= + +## 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_components__` 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 {ref}`extension_codec_order`. +- **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 + {ref}`planner_codec_rebinding`), 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. + +## Returning a planner, or not + +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 +) +``` + +(extension_bundles_order_conflict)= + +## 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_components__(self, ctx): + return self.inner.__datafusion_session_components__(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. + +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. + +The low-level +{py:meth}`~datafusion.SessionContext.with_logical_extension_codec` / +{py:meth}`~datafusion.SessionContext.with_physical_extension_codec` / +{py:meth}`~datafusion.SessionContext.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 +{ref}`one-level rebind ` does not reach inside it. +Reach for it last. + +(extension_bundles_codecs_are_objects)= + +## 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. + +## Failure and rollback + +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 {ref}`extension_sessions`. 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. + +`with_extensions` sidesteps the {ref}`one-level rebind ` +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. + +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/extension-guide/capsule-protocol.md b/docs/source/extension-guide/capsule-protocol.md new file mode 100644 index 000000000..3b8331b4b --- /dev/null +++ b/docs/source/extension-guide/capsule-protocol.md @@ -0,0 +1,238 @@ + + +(extension_capsule_protocol)= + +# The capsule protocol + +Every integration point in this section works the same way. Your library +exposes a dunder method, datafusion-python calls it, and it hands back a +`PyCapsule` wrapping an FFI-safe struct. This page describes that mechanism +once; the pages after it describe what goes inside the capsule for each kind +of component. + +The bulk of the code necessary to perform our FFI operations is in the +upstream [DataFusion](https://datafusion.apache.org/) core repository. You can +review the code and documentation in the [datafusion-ffi] crate. + +## FFI-safe types + +Our FFI implementation is narrowly focused on sharing data and functions with +Rust backed libraries. This allows us to use the +[stabby crate](https://crates.io/crates/stabby), which converts between Rust +native types and FFI-safe alternatives. For example, if you needed to pass a +`Vec` via FFI, you can convert it to a +`stabby::vec::Vec` — the crate's own examples alias +these as `SVec` and `SString`, which is the convention [datafusion-ffi] follows +too. + +For `Option` and `Result`, [datafusion-ffi] defines its own `FFI_Option` and +`FFI_Result` rather than using stabby's. Stabby's versions require +`T: IStable` for niche optimization, and many of the `FFI_*` structs hold +self-referential function pointers that cannot implement it. + +## `FFI_` on the provider, `Foreign` on the receiver + +The [datafusion-ffi] crate has been designed to make it easy to convert from +DataFusion traits into their FFI counterparts. For example, if you have +defined a custom +[TableProvider](https://docs.rs/datafusion/45.0.0/datafusion/catalog/trait.TableProvider.html) +and you want to create a sharable FFI counterpart, you could write: + +```rust +let my_provider = Arc::new(MyTableProvider::default()); +let ffi_provider = FFI_TableProvider::new_with_ffi_codec(my_provider, false, None, codec); +``` + +where `codec` is the host's logical codec, read off the argument your getter +was handed — see {ref}`extension_getter_argument`. + +If you were interfacing with a library that provided the above +`FFI_TableProvider` and you needed a usable `TableProvider` back, you convert +it into an `Arc`: + +```rust +let provider: Arc = (&ffi_provider).into(); +``` + +If you review the code in [datafusion-ffi] you will find that each of the +traits we share across the boundary has two portions, one with an `FFI_` +prefix and one with a `Foreign` prefix. This is used to distinguish which side +of the FFI boundary that struct is designed to be used on. The structures with +the `FFI_` prefix are to be used on the **provider** of the structure. In the +example we're showing, this means the code that has written the underlying +`TableProvider` implementation to access your custom data source. The +structures with the `Foreign` prefix are to be used by the receiver. In this +case, it is the `datafusion-python` library. + +Convert to the trait object rather than naming `ForeignTableProvider` yourself. +The conversion compares the provider's library marker against the receiver's: +when both sides turn out to be the same shared library it hands back the +original `Arc` and skips the boundary entirely, and only otherwise wraps it in +a `ForeignTableProvider`. Which one you get is an implementation detail, and +both implement `TableProvider`. + +## Wrapping it in a capsule + +In order to share these FFI structures, we need to wrap them in some kind of +Python object that can be used to interface from one package to another. As +described in {ref}`extension_why_ffi`, we use `PyCapsule`. We can create a +`PyCapsule` for our provider thusly: + +```rust +PyCapsule::new_with_value(py, ffi_provider, cr"datafusion_table_provider") +``` + +On the receiving side, read the `FFI_TableProvider` back out of the capsule and +convert it, which is what `table_provider_from_pycapsule` in `crates/util` does: + +```rust +validate_pycapsule(capsule, "datafusion_table_provider")?; +let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_table_provider"))? + .cast(); +let ffi_provider = unsafe { data.as_ref() }; +check_ffi_version("table provider", unsafe { (ffi_provider.version)() })?; +let provider: Arc = ffi_provider.into(); +``` + +## The naming rule + +The getter's name and the capsule's name are both fixed by the protocol, and +they follow one rule with no exceptions: + +- The method is `__datafusion___`. +- The capsule it returns is named `datafusion_` — the same string + without the underscores. + +So a table provider is reached by calling `__datafusion_table_provider__` and +must return a capsule named `datafusion_table_provider`. Return a capsule with +the wrong name and the import fails with an error naming both the name found +and the name expected, rather than reading the pointer as the wrong type. + +The full list of hooks is in the +{ref}`hook reference `. `TableProvider` was the first +extension written this way and is the most thoroughly implemented; every hook +added since follows the same pattern. + +## Version checking + +Objects imported through this protocol are checked against the major version of +`datafusion-ffi` that datafusion-python was built with. A component produced by +a library built against a different DataFusion major version raises an +`ImportError` naming the version found and the version expected. + +This is a diagnostic rather than a soundness guarantee — reading the version +out of the struct already assumes the local field layout — but it turns the +common "extension library built against the wrong DataFusion" mistake into a +clear message rather than undefined behaviour on first use. See +{ref}`extension_version_mismatch`. + +Three FFI structs carry no version field and so cannot be checked: +`FFI_TaskContextProvider`, `FFI_TableProviderFactory`, and +`FFI_ExtensionOptions`. + +(extension_getter_argument)= + +## What your getter receives + +Most getters take one positional argument beyond `py`. The +{ref}`hook reference ` says which, and the group that does +looks like this: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +This exists because the FFI constructors need things an extension library does +not have. `FFI_{Logical,Physical}ExtensionCodec::new` needs a +`TaskContextProvider` for the decode callbacks the codec will receive, and +`FFI_QueryPlanner::new` needs both codecs on top of that. Taking them from the +argument is what keeps a library from constructing a `SessionContext` purely to +satisfy a parameter — an empty one resolves nothing, and +`FFI_TaskContextProvider` holds it weakly, so a context built inline in the +getter is already dropped by the time the capsule is used. + +### It is not always a session + +The parameter is conventionally named `session`, and for the codec and planner +hooks it genuinely is one. For the provider and catalog hooks it may instead be +a bare `datafusion_logical_extension_codec` capsule: `SessionContext.register_table` +passes the session, while `Schema.register_table`, +`SessionContext.register_catalog_provider`, `register_catalog_provider_list`, +and `register_table_factory` pass the host's codec directly. + +This is why the helpers accept either. `ffi_logical_codec_from_pycapsule` +calls the codec getter if the object has one and returns the object untouched +if it does not, so the same line works in both cases: + +```rust +let codec = ffi_logical_codec_from_pycapsule(session, None)?; +``` + +The rule to hold onto is therefore about capability rather than type: **the +argument is something you can read the host's logical extension codec off.** +The hooks that need more than a codec — the two extension codecs and the query +planner, which need a task-context provider — are exactly the hooks that are +always handed a real session. + +### Duck-type it + +Do not check the argument's type. Beyond the codec-capsule case above, even +when it *is* a session it is the PyO3 context the binding installs through, not +the `datafusion.context.SessionContext` wrapper. 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 are the exception: `__datafusion_session_components__` and +`__datafusion_session_planner__` are dispatched from Python by +{py:meth}`~datafusion.SessionContext.with_extensions`, so they receive the +wrapper. See {doc}`bundles`. + +`SessionContext` accepts the argument on its own codec and planner getters and +ignores it, so a session satisfies the same protocol an extension library +implements. When you export the current planner in order to wrap it, +`ctx.__datafusion_query_planner__()` and `ctx.__datafusion_query_planner__(ctx)` +are both fine. + +(extension_task_context_provider)= + +## `__datafusion_task_context_provider__` + +This is the one hook implemented by the **host** rather than by your library. +You never define it; you read it off the session you were handed, which is what +`ffi_task_context_provider_from_pycapsule` does. Nothing in datafusion-python +calls it on a foreign object. + +Taking the host's provider means your decode callbacks resolve names against +the session that is actually running the query — see +{ref}`extension_codec_decode_session` — and it removes any need for your +library to construct a `SessionContext` of its own. + +[datafusion-ffi]: https://crates.io/crates/datafusion-ffi diff --git a/docs/source/extension-guide/checklist.md b/docs/source/extension-guide/checklist.md new file mode 100644 index 000000000..6e204b3f5 --- /dev/null +++ b/docs/source/extension-guide/checklist.md @@ -0,0 +1,99 @@ + + +(extension_checklist)= + +# Extension author checklist + +The rules in this section, gathered into one list to run through before you +publish. Each links to the page that explains it. + +## Protocol + +- [ ] **Every getter's capsule name matches its method name.** + `__datafusion___` returns a capsule named `datafusion_`. + → {ref}`extension_capsule_protocol` +- [ ] **Your getter does not inspect its argument.** Pass it to + `ffi_logical_codec_from_pycapsule` and move on. It is not always a + session, and when it is, it is not the Python `SessionContext` wrapper. + → {ref}`extension_getter_argument` +- [ ] **You never construct a `SessionContext` inside your library.** Take + what the FFI constructors need off the argument you were handed. A + context built inline is already dropped by the time the capsule is used. + → {ref}`extension_getter_argument` +- [ ] **You do not depend on the `datafusion-python` crate.** + → {ref}`extension_why_ffi` + +## Codecs + +- [ ] **Your codec claims narrowly.** Downcast to your own types. Claiming a + broad category takes nodes from every library installed after you and + makes your library order-sensitive for everyone downstream. + → {ref}`extension_codec_order` +- [ ] **You declare `__datafusion_codec_id__` if you might rename the class.** + The default id is the exporting class's import path, so a rename stops + older plans decoding. → {ref}`extension_codec_ids` +- [ ] **Name-only decoders check `name` before trusting `buf`.** An empty + payload has no id to route on, so your `try_decode_udf` can be called + with another library's function name and an empty buffer. + → {ref}`extension_codecs` +- [ ] **You round-trip a plan in a test and assert *your* codec did the work.** + Both being installed does not mean your node reached you. + → {ref}`extension_codec_order` + +## Bundles and planners + +- [ ] **You ship a bundle, not loose pieces**, if you have codecs or a planner. + → {ref}`extension_bundles` +- [ ] **Your bundle is configuration-only.** Fresh components on every call, + no cached bound components, no retaining the context passed in, no + registering anything on it — a factory that mutates the context is not + rolled back if a later factory raises. + → {ref}`extension_bundles` +- [ ] **Your codecs are objects exposing the getter, not bare capsules.** + `with_extensions` refuses a capsule, because there would be nothing to + name the codec by. → {ref}`extension_bundles_codecs_are_objects` +- [ ] **Your planner hook wraps `fallback` and delegates to it.** Ignoring it + replaces every layer beneath you, which is legal but not composable. + → {ref}`extension_bundles` +- [ ] **Your planner hook returns `None`, not `fallback`, when it has nothing + to contribute.** Returning `fallback` installs the session's own planner + as a foreign one and adds an FFI hop that was not there. + → {ref}`extension_bundles` +- [ ] **You read the host's codec chains in the planner hook, not the extension + hook.** Phase one runs before anything is installed. + → {ref}`extension_bundles_two_phases` +- [ ] **If you also offer the low-level path**, document that codecs go in + before a layered planner. → {ref}`planner_codec_rebinding` + +## Packaging and documentation + +- [ ] **You state which `datafusion` version your release requires.** A + mismatch raises an `ImportError` on import, which is a good failure — but + only if your users know what to install. → {ref}`extension_version_mismatch` +- [ ] **You tell your users to keep a context alive** for as long as anything + derived from it is in use. This is the rule most likely to arrive as a + bug report against your library. → {ref}`extension_sessions` +- [ ] **Your production codec serializes durable metadata**, not a + process-local token. The examples in this repository use tokens to make + ownership observable; that is a demonstration, not a pattern. + → {ref}`extension_codec_durable_metadata` +- [ ] **You have integration tests across a real FFI boundary.** The two + example crates in this repository are the pattern: build the cdylib, + install the wheel, then exercise it from Python. diff --git a/docs/source/extension-guide/codecs.md b/docs/source/extension-guide/codecs.md new file mode 100644 index 000000000..160439927 --- /dev/null +++ b/docs/source/extension-guide/codecs.md @@ -0,0 +1,183 @@ + + +(extension_codecs)= + +# Extension codecs + +A codec is what lets your objects survive being serialized into a plan and +rebuilt somewhere else — another process, or another program. If your library +contributes table providers, functions, or execution plan nodes and those plans +have to leave the process, you need one. + +Codecs are contributed to a session either through an +{ref}`extension bundle `, which is the recommended route, +or one at a time through +{py:meth}`~datafusion.SessionContext.with_logical_extension_codec` and +{py:meth}`~datafusion.SessionContext.with_physical_extension_codec`. + +## Codecs compose + +Each call to `with_logical_extension_codec` or +`with_physical_extension_codec` **appends** the codec to the session's codec +chain rather than replacing prior codecs. One session can therefore carry +codecs from several independent libraries at once. + +**Nothing is asked of the codec itself.** Implement `LogicalExtensionCodec` or +`PhysicalExtensionCodec` exactly as you would for a session that installs only +yours. When your codec writes bytes into a serialized plan, datafusion-python +records which codec wrote them, and strips that record off again before handing +the bytes back. So your codec receives, byte for byte, the payload it wrote, +and is never offered a payload another codec wrote. + +A codec that also ships to hosts which dispatch differently may still want its +own guard against foreign payloads. Keeping one is fine; it is simply not +needed for the datafusion-python path. + +(extension_codec_durable_metadata)= + +## Encode metadata, not a handle to a live object + +Your payload has to be enough to rebuild the object somewhere your process is +not. Write the metadata a fresh instance can be constructed from — a path, a +connection string, a schema, the options the object was created with. + +The example codecs in this repository do not do this, and it is worth knowing +before copying them. They keep a process-local `HashMap` of live providers and +encode an integer token into it: encoding inserts, decoding removes. That makes +Rust type identity observable across three separately loaded libraries in one +test, which is what the examples exist to show. It also means a decode consumes +its token, so the same bytes cannot be decoded twice, one encoded plan cannot +fan out to several readers, and a plan that never reaches a decoder keeps its +provider alive for the life of the process. A real codec has none of those +properties because it does not park the object anywhere. + +(extension_codec_ids)= + +## Codec ids + +That record is the codec's **id**: a short string stored inside the plan, +naming the codec that wrote each payload. Because plans are decoded in another +process — or another program — the id has to name the same codec there as it +did where the plan was written. + +Ids are assigned for you. A codec's id is normally its exporting class's import +path, such as `my_library.Codec`, which is what you will see in +{py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` and in decode +errors. You choose one yourself in 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 + 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. + + {py:meth}`~datafusion.SessionContext.with_extensions` takes no `codec_id=`, + so it refuses a bare capsule outright and tells you to wrap it. See + {ref}`extension_bundles_codecs_are_objects`. +- **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. + +{py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` and its +physical counterpart list the ids installed on a session, which is also what a +decode failure names. + +Installing one context's codec stack on another session composes the two +sessions rather than copying codecs out of one: the imported codecs resolve +their task context against the original and stop working when it is dropped — +see {ref}`ffi_internals_one_arc`. Pass the context itself rather than the +capsule it exports, so its codecs get an id that other sessions can decode. + +## Functions whose name is the whole encoding + +A codec may own functions that need no payload at all, where the name is the +whole encoding: `try_encode_udf` writes nothing and `try_decode_udf` rebuilds +the function from `name`. That is supported and needs no id, because an `Ok` +with an empty buffer is read as "no opinion" and passes the object to the next +codec. `NameOnlyUdfCodec` in [`datafusion-ffi-example`] is the worked case. +Anything no installed codec claims falls through to +`Default{Logical,Physical}ExtensionCodec`. + +This is the one case where your decoder is consulted about something you may +not own, because an empty payload has no id to route on. `try_decode_udf` and +its aggregate and window siblings can therefore be called with an empty `buf` +and a `name` belonging to another library. Decide from `name` and return an +error if it is not yours; do not assume `buf` is non-empty. + +:::{note} +The framing itself — how an id is stored alongside a payload and routed back, +and the two cases that stay unframed — is internal to datafusion-python and +documented in `crates/core/src/codec.rs` for anyone changing it. +::: + +The current FFI logical codec supports providers and UDFs but not arbitrary +custom `LogicalPlan::Extension` nodes. See both example READMEs for the +supported flow and local build commands. + +(extension_codec_order)= + +## When codec order matters + +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 [`datafusion-ffi-example`] claims this way, and +the query-planner example's test suite 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 + {py:meth}`ExecutionPlan.to_bytes ` / + {py:meth}`~datafusion.ExecutionPlan.from_bytes` in a test and assert your + codec did the work. + +(extension_codec_decode_session)= + +## A codec decodes against the session running the query + +Because the task-context provider comes from the host — see +{ref}`extension_getter_argument` — a decode callback running inside an +extension library resolves names against the session running the query. A +function registered with `ctx.register_udf(...)` is visible to a foreign codec +decoding a node that references it by name, and the handle is live rather than +a snapshot, so a registration made after the codec is installed is visible too. + +This is covered in +`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`, +where the example codecs take a `require_udf_on_decode` name and resolve it out +of the task context they are handed. + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/functions.md b/docs/source/extension-guide/functions.md new file mode 100644 index 000000000..fcfef9f1a --- /dev/null +++ b/docs/source/extension-guide/functions.md @@ -0,0 +1,106 @@ + + +(extension_functions)= + +# Functions and table functions + +Four hooks contribute functions written in Rust. Users can also define +functions in pure Python — see +{doc}`../user-guide/common-operations/udf-and-udfa` — and the two roads meet at +the same registration methods. + +| Hook | Contributes | Wrapped by | Registered with | +| --- | --- | --- | --- | +| `__datafusion_scalar_udf__` | scalar function | {py:func}`datafusion.udf` | {py:meth}`~datafusion.SessionContext.register_udf` | +| `__datafusion_aggregate_udf__` | aggregate function | {py:func}`datafusion.udaf` | {py:meth}`~datafusion.SessionContext.register_udaf` | +| `__datafusion_window_udf__` | window function | {py:func}`datafusion.udwf` | {py:meth}`~datafusion.SessionContext.register_udwf` | +| `__datafusion_table_function__` | function returning a table | {py:func}`datafusion.udtf` | {py:meth}`~datafusion.SessionContext.register_udtf` | + +All four are implemented in [`datafusion-ffi-example`], one per file. + +## The three scalar-shaped hooks + +The scalar, aggregate, and window getters take **no argument** beyond `py` — +they need no codec and no task-context provider, because a function is +identified by name and signature rather than by anything session-scoped: + +```rust +#[pymethods] +impl MyScalarUDF { + fn __datafusion_scalar_udf__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let udf = Arc::new(ScalarUDF::from(self.clone())); + let ffi = FFI_ScalarUDF::from(udf); + + PyCapsule::new_with_value(py, ffi, cr"datafusion_scalar_udf") + } +} +``` + +Aggregate and window follow identically with `FFI_AggregateUDF` / +`FFI_WindowUDF` and the matching capsule names. + +Your users wrap the object once and register the result: + +```python +from datafusion import udf + +ctx.register_udf(udf(my_library.MyScalarUDF())) +``` + +## Table functions + +A table function takes literal `Expr` arguments and returns a table provider, +so it needs the host's logical codec the way a +{ref}`table provider ` does: + +```rust +fn __datafusion_table_function__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let func = self.clone(); + let codec = ffi_logical_codec_from_pycapsule(session, None)?; + let provider = FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, codec); + + PyCapsule::new_with_value(py, provider, cr"datafusion_table_function") +} +``` + +Only literal expressions are supported as arguments. The Python side is +described under +{doc}`Table Functions <../user-guide/common-operations/udf-and-udfa>`. + +## Serializing functions + +A function that appears in a plan leaving the process has to be reconstructible +on the far side. Functions are the one case where a codec often needs **no +payload at all**: the name is the whole encoding, `try_encode_udf` writes +nothing, and `try_decode_udf` rebuilds the function from `name`. See +{ref}`extension_codecs` for how that works and for the one obligation it puts +on your decoder — with an empty payload there is no id to route on, so your +`try_decode_udf` can be called with a `name` belonging to another library. + +`NameOnlyUdfCodec` in [`datafusion-ffi-example`] is the worked case. + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/index.md b/docs/source/extension-guide/index.md new file mode 100644 index 000000000..91d42a3c5 --- /dev/null +++ b/docs/source/extension-guide/index.md @@ -0,0 +1,136 @@ + + +(ffi)= +(extension_guide)= + +# Extension Guide + +This section is for people **writing** a library that plugs into +datafusion-python — a package that contributes table providers, functions, a +catalog, extension codecs, or a query planner, usually written in Rust and +exposed through [PyO3](https://pyo3.rs). + +Two neighbouring audiences are served elsewhere: + +- **Using** an extension library someone else published: + {ref}`user_guide_extensions`. +- **Changing datafusion-python itself**, including the framing that makes this + protocol work: {doc}`../contributor-guide/index`. + +The protocol described here is a public, versioned contract. When a hook's +signature changes, the change is documented in +{doc}`../user-guide/upgrade-guides` with a before-and-after, and objects built +against a different DataFusion major version are rejected with a clear error +rather than used as-is. + +## Start here + +If you have not built an extension before, read {doc}`why-ffi` and +{doc}`capsule-protocol` in order. Together they explain why your library must +not depend on the `datafusion-python` crate, and the one convention every hook +follows. After that the pages are independent — go to the one matching what +you are contributing. + +## The three roles in a query + +A single query can involve three independent native libraries, and much of +this section only makes sense once they are distinct in your head: + +- **datafusion-python** — the host. Owns the session, and decodes whatever + comes back from the other two. +- **A provider library** — owns table providers, catalogs, and functions, plus + the codecs that serialize them. +- **A planner library** — owns a query planner and the configuration it needs. + +The worked examples in this repository use two separate crates, +[`datafusion-ffi-example`] and [`datafusion-ffi-query-planner-example`], so +each role has a distinct shared-library identity. A real library may play more +than one role; keeping them separate in the examples is what makes the +boundaries observable. + +The session owns the codecs used for the exchange and supplies them to the +foreign planner. That is what lets the planner decode provider-owned objects, +and lets datafusion-python decode the physical plan the planner returns. + +## Hook reference + +Every integration point is a dunder method named `__datafusion_*__`. The +convention is uniform enough to be worth stating once: your object exposes the +getter, datafusion-python calls it, and it returns a `PyCapsule` wrapping an +FFI-safe struct. See {doc}`capsule-protocol` for what that means and +{ref}`extension_getter_argument` for the argument every getter in the middle +group receives. + +Each hook's capsule name follows from its own name by +{ref}`the naming rule `, so it is not repeated +here. + +| Hook | Contributes | Argument | Documented on | +| --- | --- | --- | --- | +| `__datafusion_table_provider__` | one table | codec source | {doc}`table-providers` | +| `__datafusion_table_provider_factory__` | tables built by `CREATE EXTERNAL TABLE` | codec source | {doc}`table-providers` | +| `__datafusion_catalog_provider__` | a named set of schemas | codec source | {doc}`table-providers` | +| `__datafusion_catalog_provider_list__` | the whole catalog namespace | codec source | {doc}`table-providers` | +| `__datafusion_schema_provider__` | a named set of tables | codec source | {doc}`table-providers` | +| `__datafusion_table_function__` | a table-valued function | session | {doc}`functions` | +| `__datafusion_scalar_udf__` | a scalar function | none | {doc}`functions` | +| `__datafusion_aggregate_udf__` | an aggregate function | none | {doc}`functions` | +| `__datafusion_window_udf__` | a window function | none | {doc}`functions` | +| `__datafusion_logical_extension_codec__` | a logical codec | session | {doc}`codecs` | +| `__datafusion_physical_extension_codec__` | a physical codec | session | {doc}`codecs` | +| `__datafusion_codec_id__` | the wire id a codec's payloads carry | — | {ref}`extension_codec_ids` | +| `__datafusion_query_planner__` | a query planner | session | {doc}`query-planners` | +| `__datafusion_session_components__` | a bundle's codecs | `ctx` | {doc}`bundles` | +| `__datafusion_session_planner__` | a bundle's planner, wrapping `fallback` | `ctx`, `fallback` | {doc}`bundles` | +| `__datafusion_physical_optimizer_rule__` | a physical optimizer rule | none | {ref}`extension_other_hooks` | +| `__datafusion_extension_options__` | typed entries in the session config | none | {ref}`extension_other_hooks` | +| `__datafusion_task_context_provider__` | the host's task context | none | {ref}`extension_task_context_provider` | + +Three rows are not like the others. `__datafusion_task_context_provider__` is +implemented by the **host**, not by your library — you read it off the session +you are handed. `__datafusion_codec_id__` is a plain string attribute rather +than a method returning a capsule. And the two `session_` hooks are dispatched +from Python and return objects rather than capsules. + +"codec source" in the argument column means the value is something you can +read the host's logical extension codec off, which is not always a session. +{ref}`extension_getter_argument` explains why, and what to do with it. + +`python/tests/test_docstrings.py` compares this table against the hook names +the package actually dispatches, so a hook added or renamed without touching +this page fails the suite. + +```{toctree} +:maxdepth: 2 + +why-ffi +capsule-protocol +table-providers +functions +codecs +bundles +query-planners +other-components +sessions +checklist +``` + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/extension-guide/other-components.md b/docs/source/extension-guide/other-components.md new file mode 100644 index 000000000..919b8718e --- /dev/null +++ b/docs/source/extension-guide/other-components.md @@ -0,0 +1,72 @@ + + +(extension_other_hooks)= + +# Optimizer rules and configuration + +Two hooks contribute things that are neither data nor functions: a rewrite pass +over physical plans, and typed entries in the session config. Both take no +argument and both are implemented in [`datafusion-ffi-example`]. + +## A physical optimizer rule + +**`__datafusion_physical_optimizer_rule__`** contributes a rule that rewrites +physical plans, installed with +{py:meth}`~datafusion.SessionContext.add_physical_optimizer_rule`. Reach for +this rather than a {doc}`query planner ` when you want to +adjust the plan DataFusion produced rather than produce it yourself — it is +much the smaller commitment, and rules accumulate where planners nest. + +```rust +fn __datafusion_physical_optimizer_rule__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let rule: Arc = Arc::new(self.clone()); + let runtime = get_tokio_runtime().handle().clone(); + let ffi = FFI_PhysicalOptimizerRule::new(rule, Some(runtime)); + + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_optimizer_rule") +} +``` + +## Typed configuration + +**`__datafusion_extension_options__`** contributes typed configuration entries +that your components can read back out of the session config, installed with +{py:meth}`SessionConfig.with_extension `. +`FFI_ExtensionOptions` carries no version field, so it is one of the three +components that cannot be version-checked on import. + +```rust +fn __datafusion_extension_options__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let mut config = FFI_ExtensionOptions::default(); + config + .add_config(self) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + + PyCapsule::new_with_value(py, config, cr"datafusion_extension_options") +} +``` + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/query-planners.md b/docs/source/extension-guide/query-planners.md new file mode 100644 index 000000000..e373bb552 --- /dev/null +++ b/docs/source/extension-guide/query-planners.md @@ -0,0 +1,104 @@ + + +(extension_planners)= + +# Query planners + +A query planner turns a logical plan into a physical one. Contributing your own +is how a library changes the way queries execute rather than what data they can +reach — a distributed engine is the motivating case, and an optimizing rewriter +is another. + +The session owns the codecs used for the exchange and supplies them to your +planner. That is what lets the planner decode provider-owned objects, and lets +datafusion-python decode the physical plan the planner returns. Your planner +uses `FFI_QueryPlanner::new_with_ffi_codecs` with the two codecs it takes off +the session, and never touches a task-context provider directly. That also +matches what installation does anyway: +{py:meth}`~datafusion.SessionContext.set_query_planner` builds the planner +against the codecs of the session that will run the query. + +`MyQueryPlanner` in [`datafusion-ffi-query-planner-example`] is the worked +implementation. + +## One planner per session + +A session holds exactly one query planner. Calling `set_query_planner` again +**replaces** the installed planner instead of layering another one. + +To chain planners, have the new planner wrap the capsule returned by +{py:meth}`SessionContext.__datafusion_query_planner__ `, +captured before the new planner is installed, and delegate to it explicitly: + +```python +fallback = ctx.__datafusion_query_planner__() +ctx.set_query_planner(MyPlanner(fallback=fallback)) +``` + +`set_query_planner` returns nothing. 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 visible to every context sharing that session. See +{ref}`extension_sessions`. + +If the planners you are composing all ship as +{ref}`bundles `, prefer `with_extensions` — it does the +nesting for you and cannot capture a partial codec chain. + +(planner_codec_rebinding)= + +## Install codecs before a layered planner + +Installing a codec on a session that already has a foreign planner rebuilds +that planner against the new chain: there is one planner, and it has to carry +the codecs currently in force. But that rebuild swaps the codecs on the +installed `ForeignQueryPlanner` handle, and **only that handle**. + +A planner that wraps a fallback resolved that fallback when *it* was installed, +and holds the result inside its own library's private data — behind a +`create_physical_plan` function pointer, with no Python-side handle. A codec +installed afterwards therefore reaches the outer planner and not the fallback, +which keeps whichever codecs were in force when it was imported. + +The stale codecs stay usable rather than dangling — they hold weak handles to +the one `Arc` the session keeps alive — so the effect is a +fallback hop serializing with an older codec, not a failure. Neither side can +repair it; the reasons are in {ref}`ffi_internals_rebinding`, and a fix has to +come from upstream +([apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762)). + +So, three rules: + +- **Install the codecs before a layered planner.** +- If a codec has to go in afterwards, install the outer planner again *on the + handle that holds the new codec* — that re-runs its getter, which re-imports + the fallback against that handle's codecs. Re-installing on the original + handle rebinds the session's planner back to the original handle's codecs + instead, which is the trap the query-planner example's test suite pins. +- Better, use {ref}`extension_bundles`, where there is no "afterwards" within a + call. + +:::{note} +This is invisible in the examples here, which use one fallback in the same +cdylib as its wrapper; `datafusion-ffi` short-circuits a same-library hop +rather than serializing, so no codec runs. A fallback in a *different* library +would serialize, and would do it with the codecs it was imported with. +::: + +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/extension-guide/sessions.md b/docs/source/extension-guide/sessions.md new file mode 100644 index 000000000..21bfb6e35 --- /dev/null +++ b/docs/source/extension-guide/sessions.md @@ -0,0 +1,117 @@ + + +(extension_sessions)= + +# Sessions, handles, and lifetimes + +One `SessionContext` in Python is a *handle* on a session, not the session +itself. Several handles can share one session, and which handle you call +something on sometimes matters. This page is the set of rules that follow from +that — the ones an extension author trips over. + +## The rule that surprises people + +> The session's query planner carries the codecs of the handle that most +> recently installed one. Every other path — `Expr.to_bytes(ctx)`, +> `ExecutionPlan.to_bytes(ctx)`, registering a provider — uses the codecs of +> the handle you call it on. + +Those can be different handles, and then one session has two codec chains in +effect at once: + +```python +ctx = ctx.with_logical_extension_codec(codec_a) +ctx.set_query_planner(planner) +ctx.with_logical_extension_codec(codec_b) # discarded + +Expr.to_bytes(expr, ctx) # encodes with [codec_a, default] -- ctx's own field +ctx.sql(...).collect() # plans with [codec_b, codec_a, default] -- the discarded + # handle's chain, installed on the shared session +``` + +Chaining `ctx = ctx.with_...(...)` keeps the two in step, which is why every +example in this guide does. The query-planner example's test suite pins the +divergence. + +## What a derived context shares + +{py:meth}`~datafusion.SessionContext.with_logical_extension_codec`, +{py:meth}`~datafusion.SessionContext.with_physical_extension_codec`, +{py:meth}`~datafusion.SessionContext.with_python_udf_inlining`, and +{py:meth}`~datafusion.SessionContext.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 +task-context-provider 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. See {ref}`ffi_internals_one_arc` for why the allocation is kept +rather than replaced. + +`set_query_planner` does not return anything, because the query planner lives +in `SessionState` and is therefore a property of the session rather than of a +handle on it. Installing one is visible to every context sharing that session — +including ones a `with_*` call returned earlier. Installing a codec on a +session that already has a foreign planner rebuilds that planner against the +new chain for the same reason. This happens on the shared session, so it takes +effect even if the returned context is discarded: +`ctx.with_python_udf_inlining(...)` whose result is thrown away still leaves +the session's planner carrying the codecs of that discarded handle. A call that +changes nothing is exempt — asking for the inlining setting a context already +has returns a handle without touching the session. + +Order between installing codecs and installing a planner is a readability +preference rather than a requirement, since installing a codec after a planner +rebuilds the planner against it. The exception is a *layered* planner, where +codecs-first is a requirement: see {ref}`planner_codec_rebinding`. + +## Keep a context alive + +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: + +```text +TaskContextProvider went out of scope over FFI boundary +``` + +Keep a context alive for as long as objects derived from it are in use. This is +the rule most likely to reach your users as a bug report against your library, +so it is worth stating in your own documentation too — the user-facing version +is in {ref}`user_guide_extensions`. + +The same rule applies to a capsule you take off a context inside your own code: +a codec capsule taken from a throwaway `SessionContext()` names a session that +is already gone and fails on first use. + +:::{warning} +{py:meth}`~datafusion.SessionContext.enable_url_table` is an exception to the +one-session-one-allocation rule above: it clones the underlying +`SessionContext`, so the returned context has an allocation of its own and must +not 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 build on the behaviour. +::: diff --git a/docs/source/extension-guide/table-providers.md b/docs/source/extension-guide/table-providers.md new file mode 100644 index 000000000..8f5a0ee08 --- /dev/null +++ b/docs/source/extension-guide/table-providers.md @@ -0,0 +1,143 @@ + + +(extension_providers)= + +# Providers and catalogs + +Five hooks expose data to a session, at four levels of granularity. All five +follow the {ref}`capsule protocol ` and all five +receive a {ref}`codec source ` as their single +argument. Every one of them is implemented in [`datafusion-ffi-example`]. + +| Hook | Exposes | Registered with | +| --- | --- | --- | +| `__datafusion_table_provider__` | one table | {py:meth}`SessionContext.register_table ` | +| `__datafusion_table_provider_factory__` | a factory that builds tables from `CREATE EXTERNAL TABLE` | {py:meth}`SessionContext.register_table_factory ` | +| `__datafusion_schema_provider__` | a named set of tables | {py:meth}`Catalog.register_schema ` | +| `__datafusion_catalog_provider__` | a named set of schemas | {py:meth}`SessionContext.register_catalog_provider ` | +| `__datafusion_catalog_provider_list__` | the whole catalog namespace | {py:meth}`SessionContext.register_catalog_provider_list ` | + +A schema provider is the one that does not register on the session: you reach a +{py:class}`~datafusion.catalog.Catalog` first, with +{py:meth}`SessionContext.catalog `, and +register the schema on that. + +Start with a table provider. Reach for the schema and catalog levels when your +data source has its own namespace that should be browsable rather than +registered table by table, and for the provider list only when your library is +replacing the catalog namespace outright. + +## A table provider + +Implement +[TableProvider](https://datafusion.apache.org/library-user-guide/custom-table-providers.html) +in Rust, then expose it: + +```rust +#[pymethods] +impl MyTableProvider { + fn __datafusion_table_provider__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let provider = Arc::new(self.clone()); + let codec = ffi_logical_codec_from_pycapsule(session, None)?; + let provider = FFI_TableProvider::new_with_ffi_codec(provider, false, None, codec); + + PyCapsule::new_with_value(py, provider, cr"datafusion_table_provider") + } +} +``` + +Your users then register it as they would any other table — see +{ref}`io_custom_table_provider` for the Python side. + +## The catalog family + +The three catalog-level hooks have the same shape as each other. Taking the +schema provider as the representative: + +```rust +#[pymethods] +impl MySchemaProvider { + fn __datafusion_schema_provider__<'py>( + &self, + py: Python<'py>, + codec: Bound<'py, PyAny>, + ) -> PyResult> { + let provider = Arc::clone(&self.inner) as Arc; + + let codec = ffi_logical_codec_from_pycapsule(codec, None)?; + let provider = FFI_SchemaProvider::new_with_ffi_codec(provider, None, codec); + + PyCapsule::new_with_value(py, provider, cr"datafusion_schema_provider") + } +} +``` + +Swap `Schema` for `Catalog` or `CatalogProviderList` and the getter, capsule +name, and FFI type change together, following {ref}`the naming rule +`. `catalog_provider.rs` in +[`datafusion-ffi-example`] implements all three in one file, which is the +easiest way to see the symmetry. + +Note the parameter name. These four hooks are handed the host's logical codec +directly rather than a session, so naming it `codec` is more honest than +`session` — but do not rely on either: pass it to +`ffi_logical_codec_from_pycapsule` and do not inspect it. See +{ref}`extension_getter_argument`. + +## A table provider factory + +A factory backs `CREATE EXTERNAL TABLE`: DataFusion hands it the statement's +options and it produces a provider. The getter takes the codec and wraps an +`Arc`: + +```rust +fn __datafusion_table_provider_factory__<'py>( + &self, + py: Python<'py>, + codec: Bound<'py, PyAny>, +) -> PyResult> { + let codec = ffi_logical_codec_from_pycapsule(codec, None)?; + let factory = Arc::clone(&self.inner) as Arc; + let factory = FFI_TableProviderFactory::new_with_ffi_codec(factory, None, codec); + + PyCapsule::new_with_value(py, factory, cr"datafusion_table_provider_factory") +} +``` + +`FFI_TableProviderFactory` carries no version field, so a factory is one of the +three components that cannot be version-checked on import. Be correspondingly +careful about which DataFusion version you build against. + +## Serializing what you expose + +If plans referencing your tables have to leave the process — a distributed +engine will make them — your library also needs a logical extension codec, so +the provider can be rebuilt on the other side. That is {doc}`codecs`, and +{doc}`bundles` is how you ship the two together. + +The codec you pass to `new_with_ffi_codec` above is the **host's**, used to +serialize the parts of a plan the host owns. It is not a substitute for +contributing your own. + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/why-ffi.md b/docs/source/extension-guide/why-ffi.md new file mode 100644 index 000000000..e1321107a --- /dev/null +++ b/docs/source/extension-guide/why-ffi.md @@ -0,0 +1,148 @@ + + +(extension_why_ffi)= + +# Why FFI + +The DataFusion in Python project is designed to allow users to extend its +functionality in a few core areas. Ideally many users would like to package +their extensions as a Python package and easily integrate that package with +this project. This page describes the problem that makes the obvious approach +fail, and the approach this project uses instead. + +## The primary issue + +Suppose you wish to use DataFusion and you have a custom data source that can +produce tables that can then be queried against, similar to how you can +register a {ref}`CSV ` or {ref}`Parquet ` file. In +DataFusion terminology, you likely want to implement a +{ref}`Custom Table Provider `. In an effort to make +your data source as performant as possible and to utilize the features of +DataFusion, you may decide to write your source in Rust and then expose it +through [PyO3](https://pyo3.rs) as a Python library. + +At first glance, it may appear the best way to do this is to add the +`datafusion-python` crate as a dependency, provide a `PyTable`, and then to +register it with the `SessionContext`. Unfortunately, this will not work. + +When you produce your code as a Python library and it needs to interact with +the DataFusion library, at the lowest level they communicate through an +Application Binary Interface (ABI). The acronym sounds similar to API +(Application Programming Interface), but it is distinctly different. + +The ABI sets the standard for how these libraries can share data and functions +between each other. One of the key differences between Rust and other +programming languages is that Rust does not have a stable ABI. What this means +in practice is that if you compile a Rust library with one version of the +`rustc` compiler and I compile another library to interface with it but I use a +different version of the compiler, there is no guarantee the interface will be +the same. + +In practice, this means that a Python library built with `datafusion-python` as +a Rust dependency will generally **not** be compatible with the DataFusion +Python package, even if they reference the same version of +`datafusion-python`. If you attempt to do this, it may work on your local +computer if you have built both packages with the same optimizations. This can +sometimes lead to a false expectation that the code will work, but it +frequently breaks the moment you try to use your package against the released +packages. + +You can find more information about the Rust ABI in their +[online documentation](https://doc.rust-lang.org/reference/abi.html). + +## The FFI approach + +Rust supports interacting with other programming languages through its Foreign +Function Interface (FFI). The advantage of using the FFI is that it enables you +to write data structures and functions that have a stable ABI. That allows you +to use Rust code with C, Python, and other languages. In fact, the +[PyO3](https://pyo3.rs) library uses the FFI to share data and functions +between Python and Rust. + +The approach we are taking in the DataFusion in Python project is to +incrementally expose more portions of the DataFusion project via FFI +interfaces. This allows users to write Rust code that does **not** require the +`datafusion-python` crate as a dependency, expose their code in Python via +PyO3, and have it interact with the DataFusion Python package. + +Early adopters of this approach include +[delta-rs](https://delta-io.github.io/delta-rs/) who has adapted their Table +Provider for use in `datafusion-python` with only a few lines of code. Also, +the DataFusion Python project uses the existing definitions from +[Apache Arrow CStream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html) +to support importing **and** exporting tables. Any Python package that supports +reading the Arrow C Stream interface can work with DataFusion Python out of the +box! You can read more about working with Arrow sources in the +{ref}`Data Sources ` page. + +To learn more about the Foreign Function Interface in Rust, the +[Rustonomicon](https://doc.rust-lang.org/nomicon/ffi.html) is a good resource. + +## Inspiration from Arrow + +DataFusion is built upon [Apache Arrow](https://arrow.apache.org/). The +canonical Python Arrow implementation, +[pyarrow](https://arrow.apache.org/docs/python/index.html), provides an +excellent way to share Arrow data between Python projects without performing +any copy operations on the data, using a well defined set of interfaces — see +their [stream interface](https://arrow.apache.org/docs/format/CStreamInterface.html). +The [Rust Arrow implementation](https://github.com/apache/arrow-rs) also +supports these `C` style definitions via the Foreign Function Interface. Beyond +transferring data, `pyarrow` goes one step further and makes the interfaces +themselves easy to share in Python, by exposing PyCapsules that contain the +expected functionality. + +Two lessons we leverage from the Arrow project in DataFusion Python are: + +- We reuse the existing Arrow FFI functionality wherever possible. +- We expose PyCapsules that contain an FFI stable struct. + +You can learn more about PyCapsules from the official +[Python online documentation](https://docs.python.org/3/c-api/capsule.html). +PyCapsules have excellent support in PyO3 already; the +[PyO3 online documentation](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule) +is a good source for more details on using PyCapsules in Rust. + +## If FFI does not yet cover what you need + +:::{note} +Suppose you needed to expose some other features of DataFusion and you could +not wait for the upstream repository to implement the FFI approach we describe. +In this case you decide to create your dependency on the `datafusion-python` +crate instead. + +As we discussed, this is not guaranteed to work across different compiler +versions and optimization levels. If you wish to go down this route, there are +two approaches we have identified you can use. + +1. Re-export all of `datafusion-python` yourself with your extensions built in. +2. Carefully synchronize your software releases with the `datafusion-python` CI + build system so that your libraries use the exact same compiler, features, + and optimization level. + +We currently do not recommend either of these approaches as they are difficult +to maintain over a long period. Additionally, they require a tight version +coupling between libraries. + +The better path is to open an issue describing what you need exposed. The FFI +surface in the [datafusion-ffi] crate grows in response to these. +::: + +[datafusion-ffi]: https://crates.io/crates/datafusion-ffi diff --git a/docs/source/index.md b/docs/source/index.md index 13f7c7df6..c702fc040 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -59,11 +59,22 @@ df.show() ``` +## Where to go next + +- **[User Guide](user-guide/index)** — reading data, building queries, tuning + execution, and distributing work. +- **[Extension Guide](extension-guide/index)** — writing a library that plugs + into DataFusion in Python: table providers, functions, extension codecs, and + query planners. +- **[Contributor Guide](contributor-guide/index)** — changing + datafusion-python itself. + ```{toctree} :hidden: true :maxdepth: 1 user-guide/index +extension-guide/index contributor-guide/index API Reference links diff --git a/docs/source/llms.txt b/docs/source/llms.txt index 3ff6b3813..47b303ff6 100644 --- a/docs/source/llms.txt +++ b/docs/source/llms.txt @@ -16,6 +16,14 @@ - [Common operations](https://datafusion.apache.org/python/user-guide/common-operations/index.html): select, filter, join, aggregate, window, expressions, and functions. - [SQL](https://datafusion.apache.org/python/user-guide/sql.html): running SQL against registered tables. - [Configuration](https://datafusion.apache.org/python/user-guide/configuration.html): session and runtime options. +- [Using extension libraries](https://datafusion.apache.org/python/user-guide/extensions.html): installing a compiled extension and registering its tables, functions, or execution backend with `SessionContext.with_extensions`. + +## Extensions and distribution + +- [Using extension libraries](https://datafusion.apache.org/python/user-guide/extensions.html): which extensions register directly and which need `with_extensions`, plus the two failure modes (a collected context, a version mismatch). +- [Distributing work](https://datafusion.apache.org/python/user-guide/distributing-work/index.html): choosing between shipping expressions to a worker pool and installing a distributed query engine. +- [Shipping expressions to workers](https://datafusion.apache.org/python/user-guide/distributing-work/expressions.html): pickling `Expr` to `multiprocessing` or Ray, what travels inline vs by name, `datafusion.ipc` worker/sender contexts, and the security model. +- [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html): the PyCapsule protocol for shipping table providers, catalogs, functions, extension codecs, and query planners from a Rust cdylib. Includes the full `__datafusion_*__` hook reference (18 hooks) and an author checklist. ## DataFrame API reference @@ -23,14 +31,18 @@ - [`datafusion.expr`](https://datafusion.apache.org/python/autoapi/datafusion/expr/index.html): expression tree nodes (`Expr`, `Window`, `WindowFrame`, `GroupingSet`). - [`datafusion.functions`](https://datafusion.apache.org/python/autoapi/datafusion/functions/index.html): 290+ scalar, aggregate, and window functions. - [`datafusion.context.SessionContext`](https://datafusion.apache.org/python/autoapi/datafusion/context/index.html): session entry point, data loading, SQL execution. +- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionComponentsExportable`, `SessionPlannerExportable`, `QueryPlannerExportable` — the extension-bundle protocol. +- [`datafusion.ipc`](https://datafusion.apache.org/python/autoapi/datafusion/ipc/index.html): worker and sender context slots for shipping expressions between processes. ## Examples - [TPC-H queries (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples/tpch): canonical translations of TPC-H Q01–Q22 to idiomatic DataFrame code, each with reference SQL embedded in the module docstring. - [Other examples (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples): UDF/UDAF/UDWF, Substrait, Pandas/Polars interop, S3 reads. +- [FFI extension example (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example): a Rust cdylib exporting table providers, catalogs, functions, and extension codecs. Has build commands in its README. +- [FFI query planner example (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example): a Rust cdylib exporting a query planner and a complete extension bundle. ## Optional -- [Contributor guide](https://datafusion.apache.org/python/contributor-guide/introduction.html): building from source, extending the Python bindings. +- [Contributor guide](https://datafusion.apache.org/python/contributor-guide/index.html): building from source and changing datafusion-python itself. - [Upgrade guides](https://datafusion.apache.org/python/user-guide/upgrade-guides.html): migration notes between releases. - [Upstream Rust `DataFusion`](https://datafusion.apache.org/): the underlying query engine. diff --git a/docs/source/user-guide/common-operations/udf-and-udfa.md b/docs/source/user-guide/common-operations/udf-and-udfa.md index 8d1cc876d..07005fd0f 100644 --- a/docs/source/user-guide/common-operations/udf-and-udfa.md +++ b/docs/source/user-guide/common-operations/udf-and-udfa.md @@ -391,7 +391,8 @@ df.select("a", exp_smooth(col("a")).alias("smooth_a")).show() User Defined Table Functions are slightly different than the other functions described here. These functions take any number of `Expr` arguments, but only literal expressions are supported. Table functions must return a Table -Provider as described in the ref:`_io_custom_table_provider` page. +Provider as described in the {ref}`Custom Table Provider ` +page. Once you have a table function, you can register it with the session context by using {py:func}`datafusion.context.SessionContext.register_udtf`. diff --git a/docs/source/user-guide/data-sources.md b/docs/source/user-guide/data-sources.md index 22e666837..ae56ae5a1 100644 --- a/docs/source/user-guide/data-sources.md +++ b/docs/source/user-guide/data-sources.md @@ -214,7 +214,7 @@ Features that are available in PyIceberg but not yet in Iceberg Rust will not be ## Custom Table Provider You can implement a custom Data Provider in Rust and expose it to DataFusion through the -the interface as describe in the {ref}`Custom Table Provider ` +interface described in the {ref}`Custom Table Provider ` section. This is an advanced topic, but a [user example](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example) is provided in the DataFusion repository. diff --git a/docs/source/user-guide/distributing-work.md b/docs/source/user-guide/distributing-work/expressions.md similarity index 86% rename from docs/source/user-guide/distributing-work.md rename to docs/source/user-guide/distributing-work/expressions.md index 527fe8830..252defc24 100644 --- a/docs/source/user-guide/distributing-work.md +++ b/docs/source/user-guide/distributing-work/expressions.md @@ -17,24 +17,19 @@ under the License. --> -# Distributing work +(distributed_expressions)= + +# Shipping expressions to workers DataFusion supports splitting work across processes by shipping serialized expressions to workers: the driver builds an {py:class}`~datafusion.Expr`, each worker evaluates it against its own slice of data. This pattern suits embarrassingly-parallel -workloads where the driver decides partitioning up front. - -Query-level distribution — where the runtime partitions a single -logical or physical plan across worker nodes — is in progress -upstream via [datafusion-distributed](https://github.com/apache/datafusion-distributed) and [Apache -Ballista](https://github.com/apache/datafusion-ballista). Both -have short sections at the end of this page; integration details -will land as those projects become usable from datafusion-python. - -## Expression-level distribution +workloads where the driver decides partitioning up front. If you +would rather have a library partition your query plan for you, see +{ref}`distributed_query_engines` instead. -DataFusion expressions support distribution directly: pass one to a +Expressions support this directly: pass one to a worker process and Python's standard [pickle](https://docs.python.org/3/library/pickle.html) machinery serializes it transparently — the same machinery @@ -43,7 +38,7 @@ similar libraries already use to ship function arguments. Python UDFs — scalar, aggregate, and window — travel inside the serialized expression; the receiver does not need to pre-register them. -### Basic worker-pool example +## Basic worker-pool example Define a worker function that takes the expression plus a batch and returns the evaluated result: @@ -92,7 +87,7 @@ see [Safe importing of main module](https://docs.python.org/3/library/multiproce in the Python docs. ::: -### What travels with the expression +## What travels with the expression - **Built-in functions** (`abs`, `length`, arithmetic, comparisons, etc.) — fully portable. Worker needs nothing pre-registered. @@ -113,7 +108,7 @@ in the Python docs. (distributed_udf_portability)= -### Portability requirements for inline Python UDFs +## Portability requirements for inline Python UDFs Inline Python UDFs ride on [cloudpickle](https://github.com/cloudpipe/cloudpickle), which imposes two requirements on the worker environment: @@ -133,7 +128,7 @@ requirements on the worker environment: Self-contained UDFs (no imports beyond what the worker already has, e.g. `pyarrow`) avoid this entirely. -### Registering shared UDFs on workers +## Registering shared UDFs on workers When an expression references an FFI capsule UDF (or any UDF the worker must resolve from its registered functions), set up the @@ -164,7 +159,7 @@ fine for expressions that only reference built-ins and Python UDFs, but FFI-capsule-backed registrations must be installed on the global context to resolve. -### Python 3.14 default change +## Python 3.14 default change Python 3.14 changed the Linux default start method for {py:mod}`multiprocessing` from `fork` to `forkserver` (macOS has @@ -174,7 +169,7 @@ workers via copy-on-write; with `forkserver` and `spawn` it is not. The {py:func}`~datafusion.ipc.set_worker_ctx` pattern works on every start method — prefer it over relying on inherited state. -### Practical considerations +## Practical considerations - **Serialized size scales with what travels inline.** A serialized expression of just built-ins is small (tens of bytes). An @@ -189,7 +184,7 @@ every start method — prefer it over relying on inherited state. the captured state is large, mutable, or not portable to the worker's environment. See {ref}`Portability requirements for inline Python UDFs ` for the Python-version and imported-module rules. -### Disabling Python UDF inlining +## Disabling Python UDF inlining For a stricter wire format, call {py:meth}`SessionContext.with_python_udf_inlining(enabled=False) @@ -241,7 +236,7 @@ threat model. (distributed_expr_security)= -### Security +## Security :::{warning} Reconstructing an expression containing a Python UDF executes @@ -255,7 +250,9 @@ functions and pre-registered Rust-side UDFs, and avoid {py:func}`pickle.loads` on externally supplied bytes entirely. ::: -### Reference: session context slots +(session_context_slots)= + +## Reference: session context slots There is only one type — {py:class}`SessionContext`. It can occupy up to four *slots* in a running program: @@ -297,32 +294,10 @@ Sharp edges: - The inlining toggle is per-context state, not a global switch. Two contexts with different toggles can coexist in one process. -## Query-level distribution via datafusion-distributed - -🚧 *Work in progress upstream — not yet usable from datafusion-python.* - -[datafusion-distributed](https://github.com/apache/datafusion-distributed) -splits a single physical plan into stages and runs each stage on a -different worker node. The driver writes a SQL or DataFrame query -once; the runtime handles partitioning, shuffles, and reassembly. - -A datafusion-python integration is in development. This section will -document the integration once it lands. In the meantime, the -expression-level approach above covers most use cases that do not -require automatic plan partitioning. - -## Query-level distribution via Apache Ballista - -🚧 *Work in progress upstream — not yet usable from datafusion-python.* - -[Apache Ballista](https://github.com/apache/datafusion-ballista) -provides distributed query execution on top of DataFusion with a -scheduler / executor model better suited to long-lived cluster -deployments. A datafusion-python integration is on the roadmap; this -section will fill in once the integration is usable. - ## See also +- {ref}`distributed_query_engines` — the other road: let a library + partition the plan for you. - {py:mod}`datafusion.ipc` — worker context API. - `examples/multiprocessing_pickle_expr.py` — runnable `multiprocessing.Pool` example that ships a different parametric diff --git a/docs/source/user-guide/distributing-work/index.md b/docs/source/user-guide/distributing-work/index.md new file mode 100644 index 000000000..3047ea39e --- /dev/null +++ b/docs/source/user-guide/distributing-work/index.md @@ -0,0 +1,81 @@ + + +(distributing_work)= + +# Distributing work + +A single {py:class}`~datafusion.SessionContext` already uses every +core on the machine — DataFusion partitions and parallelizes within +a process without being asked. See +{doc}`../configuration` for tuning that. + +This section is about the step after that: getting work onto more +than one process or more than one machine. There are two roads, and +they suit different problems. + +## Pick your road + +**You decide the partitioning → ship expressions.** + +You already know how the data divides — one file per worker, one +customer per worker, one parameter setting per worker. You build an +{py:class}`~datafusion.Expr` in the driver and hand a copy to each +worker along with its slice. Standard Python `pickle` moves it, so +{py:mod}`multiprocessing`, Ray, and anything else that ships function +arguments works with no extra machinery. + +Best for embarrassingly-parallel work: parameter sweeps, per-file +transforms, scoring batches. Available today. + +→ {ref}`distributed_expressions` + +**A library decides the partitioning → install a query engine.** + +You write one ordinary SQL or DataFrame query against a table too +large for one machine, and an engine library splits the plan into +stages, runs them on its workers, and reassembles the result. You do +not partition anything and your queries do not change. + +Best for large-scale analytical queries — joins and aggregations over +data that does not fit on one node. Being built upstream; not yet +usable from datafusion-python. + +→ {ref}`distributed_query_engines` + +## Choosing between them + +The two are not competing implementations of one feature. The +question is who owns the partitioning decision. + +If you can state the partitioning in one line — "one worker per input +file" — the expression road is simpler, has no cluster to operate, +and works now. If stating it requires knowing how a join will +shuffle, that decision belongs to a query planner, which is what an +engine library provides. + +They also compose. An engine handles the query; expressions handle +whatever you want to fan out around it. + +```{toctree} +:maxdepth: 2 + +expressions +query-engines +``` diff --git a/docs/source/user-guide/distributing-work/query-engines.md b/docs/source/user-guide/distributing-work/query-engines.md new file mode 100644 index 000000000..0aa045c79 --- /dev/null +++ b/docs/source/user-guide/distributing-work/query-engines.md @@ -0,0 +1,100 @@ + + +(distributed_query_engines)= + +# Distributed query engines + +A distributed query engine takes the query you already wrote and runs +it across several machines for you. You do not partition anything by +hand and you do not change your queries — the engine installs itself +on your session, rewrites the plan into stages, and executes those +stages on its own workers. + +This is the counterpart to {ref}`distributed_expressions`, where +*you* decide the partitioning and ship one expression per slice of +data. Choose an engine when you want a single query spread across a +cluster; choose expressions when you already know how to split the +work and only need parallelism. + +## How an engine attaches to your session + +An engine library ships an object you hand to +{py:meth}`~datafusion.SessionContext.with_extensions`: + +```python +from datafusion import SessionContext +import my_engine + +ctx = SessionContext().with_extensions(my_engine.Extension("scheduler:50050")) + +ctx.register_table("events", my_engine.TableProvider("s3://bucket/events")) +ctx.sql("SELECT country, count(*) FROM events GROUP BY country").show() +``` + +The `sql` call is unchanged from a single-process program. What +changed is underneath it: installing the engine gave the session a +query planner of the engine's own, and that planner is what turns +your plan into distributed stages. + +Two consequences worth knowing: + +- **The engine must be installed before you run the query**, not + before you register tables. Registration order does not matter; + `with_extensions` is a session-level setup step. +- **The engine has to be able to reach your tables and functions.** + It ships the plan to its workers, so anything the plan references + has to be reconstructible there. Table providers from the engine's + own library always are. A Python UDF may or may not be — the engine + library documents what it supports, and + {ref}`distributed_udf_portability` describes the constraints that + apply to Python callables crossing a process boundary in general. + +If you install more than one library, pass them in one +`with_extensions` call so they can see each other. See +{ref}`user_guide_extensions` for the details of installing extension +libraries, and {ref}`ffi` if you want to write an engine yourself. + +## Available engines + +Query-level distribution is being built upstream. Neither project +below is usable from datafusion-python yet; both sections will fill +in as the integrations land. + +### datafusion-distributed + +🚧 *Work in progress upstream — not yet usable from datafusion-python.* + +[datafusion-distributed](https://github.com/apache/datafusion-distributed) +splits a single physical plan into stages and runs each stage on a +different worker node. The driver writes a SQL or DataFrame query +once; the runtime handles partitioning, shuffles, and reassembly. + +A datafusion-python integration is in development. In the meantime, +{ref}`distributed_expressions` covers most use cases that do not +require automatic plan partitioning. + +### Apache Ballista + +🚧 *Work in progress upstream — not yet usable from datafusion-python.* + +[Apache Ballista](https://github.com/apache/datafusion-ballista) +provides distributed query execution on top of DataFusion with a +scheduler / executor model better suited to long-lived cluster +deployments. A datafusion-python integration is on the roadmap. diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md new file mode 100644 index 000000000..8c86b4a67 --- /dev/null +++ b/docs/source/user-guide/extensions.md @@ -0,0 +1,142 @@ + + +(user_guide_extensions)= + +# Using extension libraries + +An extension library is a separate package that teaches a +{py:class}`~datafusion.SessionContext` something it does not know on its own — +a new data source, extra functions, or a different way of executing your +queries. You install it with `pip`, hand it to your session, and keep writing +the same SQL and DataFrame code. + +Examples in the wild include [delta-rs](https://delta-io.github.io/delta-rs/), +which exposes Delta Lake tables to DataFusion, and the two worked examples in +this repository under +[`examples/`](https://github.com/apache/datafusion-python/tree/main/examples). + +## Two kinds of extension + +Which one you have determines how much setup you do. + +**Tables and functions register directly.** If the library gives you a table +or a function, register it the same way you would register a CSV file. No +extra setup: + +```python +from datafusion import SessionContext +import my_tables + +ctx = SessionContext() +ctx.register_table("events", my_tables.TableProvider("s3://bucket/events")) +ctx.sql("SELECT count(*) FROM events").show() +``` + +**Libraries that change how queries run need to be installed on the session.** +A distributed engine, or anything that rewrites your query plan, has to be +attached to the session before it can do its work. That is what +{py:meth}`~datafusion.SessionContext.with_extensions` is for. The library +documents an object — often called `Extension` — that you pass to it: + +```python +from datafusion import SessionContext +import my_engine + +ctx = SessionContext().with_extensions(my_engine.Extension("scheduler:50050")) +ctx.register_table("events", my_engine.TableProvider("s3://bucket/events")) +ctx.sql("SELECT count(*) FROM events").show() +``` + +`with_extensions` returns a context; use the returned one. It shares +everything else with the context you called it on, so tables you registered +before the call are still there. + +## Using more than one library + +Pass them all to a single call: + +```python +ctx = SessionContext().with_extensions( + my_tables.Extension(), + my_engine.Extension("scheduler:50050"), +) +``` + +One call rather than several is worth preferring: it lets the libraries see +each other, which they cannot do if you install them one at a time. Order +rarely matters. When a library needs a particular position — usually "list me +last" for something that wraps the others — it says so in its own +documentation. + +## Two things that will bite you + +**Keep your context alive.** A `DataFrame` or a plan does not keep its session +alive on its own. If a context is garbage-collected while something built from +it is still in use, the next query fails with: + +```text +TaskContextProvider went out of scope over FFI boundary +``` + +Almost always this is a helper that built a context locally and returned a +DataFrame: + +```python +# Wrong — ctx is collected when the function returns. +def load(): + ctx = SessionContext().with_extensions(my_engine.Extension()) + return ctx.sql("SELECT * FROM events") + +# Right — hand back the context too, or keep it on an object that lives +# as long as the frames derived from it. +def load(): + ctx = SessionContext().with_extensions(my_engine.Extension()) + return ctx, ctx.sql("SELECT * FROM events") +``` + +**Versions have to match.** An extension library is compiled against one +DataFusion version. A mismatch raises an `ImportError` naming the version it +found and the version expected, at the moment you register or install the +library — not silently at query time. If you see one, upgrade or downgrade the +extension library so its DataFusion version matches this package's. See +{ref}`extension_version_mismatch`. + +## Checking what a session knows about + +{py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` and +{py:meth}`~datafusion.SessionContext.physical_extension_codec_ids` list which +libraries a session has been taught about. Useful when a query fails and you +want to confirm the library actually got installed: + +```python +ctx = SessionContext().with_extensions(my_engine.Extension()) +ctx.logical_extension_codec_ids() +# ['my_engine.LogicalCodec'] +``` + +An empty list means nothing extra is installed. + +## Next steps + +- {ref}`distributed_query_engines` — running your queries across several + machines with an engine library. +- {ref}`distributed_expressions` — the other road to parallelism, where you + decide the partitioning and ship expressions to a worker pool. +- {ref}`ffi` — writing an extension library of your own. diff --git a/docs/source/user-guide/index.md b/docs/source/user-guide/index.md index 0155c4e91..07b912a8b 100644 --- a/docs/source/user-guide/index.md +++ b/docs/source/user-guide/index.md @@ -20,7 +20,11 @@ # User Guide The user guide walks through installing DataFusion in Python, building queries -with the DataFrame API or SQL, reading and writing data, and tuning execution. +with the DataFrame API or SQL, reading and writing data, tuning execution, and +spreading work across processes or machines. + +If you are writing an extension library rather than using one, see the +{ref}`Extension Guide `. ```{toctree} :maxdepth: 2 @@ -32,7 +36,8 @@ dataframe/index common-operations/index io/index configuration -distributing-work +extensions +distributing-work/index sql upgrade-guides ai-coding-assistants diff --git a/docs/source/user-guide/io/table_provider.md b/docs/source/user-guide/io/table_provider.md index 5dc2dc086..1b6987e7f 100644 --- a/docs/source/user-guide/io/table_provider.md +++ b/docs/source/user-guide/io/table_provider.md @@ -28,29 +28,8 @@ you must use DataFusion 43.0.0 or later and expose a [FFI_TableProvider](https:/ via [PyCapsule](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule). A complete example can be found in the [examples folder](https://github.com/apache/datafusion-python/tree/main/examples). - -The method takes the `SessionContext` it is being registered on. Take whatever -the FFI constructor needs from that session — here the logical extension codec — -rather than building one inside your library. See the {ref}`ffi` guide for the -full capsule protocol. - -```rust -#[pymethods] -impl MyTableProvider { - - fn __datafusion_table_provider__<'py>( - &self, - py: Python<'py>, - session: Bound<'py, PyAny>, - ) -> PyResult> { - let provider = Arc::new(self.clone()); - let codec = ffi_logical_codec_from_pycapsule(session, None)?; - let provider = FFI_TableProvider::new_with_ffi_codec(provider, false, None, codec); - - PyCapsule::new_with_value(py, provider, cr"datafusion_table_provider") - } -} -``` +For how to write one — the getter, what it receives, and how to serialize what +it exposes — see {ref}`extension_providers` in the Extension Guide. Once you have this library available, you can construct a {py:class}`~datafusion.Table` in Python and register it with the diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 257749c3a..f98590b0a 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -21,8 +21,9 @@ ## DataFusion 55.0.0 -This release extends the change made in 52.0.0 to the remaining {ref}`ffi` hook -methods. Users who contribute their own `LogicalExtensionCodec` or +This release extends the change made in 52.0.0 to the remaining +{ref}`extension_capsule_protocol` hook methods. Users who contribute their own +`LogicalExtensionCodec` or `PhysicalExtensionCodec` via FFI must update `__datafusion_logical_extension_codec__` and `__datafusion_physical_extension_codec__` to accept an additional @@ -75,13 +76,37 @@ used. it, so existing calls such as `ctx.__datafusion_logical_extension_codec__()` continue to work unchanged. +The provider and catalog getters — +`__datafusion_table_provider_factory__`, `__datafusion_catalog_provider__`, +`__datafusion_catalog_provider_list__`, and `__datafusion_schema_provider__` — +now receive the host's logical extension codec as a bare +`datafusion_logical_extension_codec` capsule rather than a session. +`__datafusion_table_provider__` receives a session from +`SessionContext.register_table` and a codec capsule from +`Schema.register_table`. No code change is needed in either case: pass the +argument to `ffi_logical_codec_from_pycapsule`, which calls the codec getter +when the object has one and returns the object unchanged when it does not. Do +not inspect the argument's type. See {ref}`extension_getter_argument`. + New in this release, `__datafusion_query_planner__` follows the same protocol. It receives the session and takes both extension codecs from it, so a planner library never builds a `TaskContextProvider` at all. Install one with `SessionContext.set_query_planner(planner)`, which mutates the session the same 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. +particular handle on it. See {ref}`extension_planners` 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_components__` 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 {ref}`extension_bundles`. + +(extension_version_mismatch)= ### Mismatched extension libraries now fail loudly diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index 401eb610a..aadea909f 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -49,6 +49,6 @@ ctx = ctx.with_physical_extension_codec(provider_physical_codec) ctx.set_query_planner(planner) ``` -Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. The exception is a planner that wraps a fallback: the rebuild reaches the installed planner only, not the fallback inside it, so codecs-first is a requirement there. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on. +Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. The exception is a planner that wraps a fallback: the rebuild reaches the installed planner only, not the fallback inside it, so codecs-first is a requirement there. See [Rebinding a planner's codecs is one level deep](https://datafusion.apache.org/python/extension-guide/query-planners.html#install-codecs-before-a-layered-planner), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on. -For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. +For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see the [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html). 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..af128886f 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_components__` +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)) @@ -55,6 +78,6 @@ ctx.set_query_planner(MyQueryPlanner()) `MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. -The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). +The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](https://datafusion.apache.org/python/extension-guide/query-planners.html#install-codecs-before-a-layered-planner). -For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. +For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see the [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html). 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..8984b8ac4 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_components__( + 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_components__( + 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_components__( + 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_components__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return self.inner.__datafusion_session_components__(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_components__( + 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_components__( + 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,94 @@ 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_components__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + codecs = self._codecs.__datafusion_session_components__(ctx) + planner = self._planner.__datafusion_session_components__(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 the ``+SKIP`` statements are taken. The docstring also carries a + runnable example above them, which the main suite already executes under + ``--doctest-modules``; running it again here would need its own namespace + and prove nothing. + + 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 = [ + example + for example in doctest.DocTestParser().get_examples( + inspect.getdoc(SessionContext.with_extensions) + ) + if example.options.pop(doctest.SKIP, False) + ] + assert examples, "with_extensions docstring has no skipped examples to check" + + 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..542b36ea3 --- /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_components__` 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_components__<'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/examples/multiprocessing_pickle_expr.py b/examples/multiprocessing_pickle_expr.py index 73a99c2db..db1a4c5fd 100644 --- a/examples/multiprocessing_pickle_expr.py +++ b/examples/multiprocessing_pickle_expr.py @@ -19,7 +19,7 @@ For background — the shipped-expression model, what travels inline vs by name, portability requirements, and the security threat model — -see ``docs/source/user-guide/distributing-work.rst``. +see https://datafusion.apache.org/python/user-guide/distributing-work/expressions.html. Builds a list of parametric expressions in the driver — each closing over a different threshold value — ships one per worker via diff --git a/examples/ray_pickle_expr.py b/examples/ray_pickle_expr.py index 04cea463d..e8be06b0f 100644 --- a/examples/ray_pickle_expr.py +++ b/examples/ray_pickle_expr.py @@ -19,7 +19,7 @@ For background — the shipped-expression model, what travels inline vs by name, portability requirements, and the security threat model — -see ``docs/source/user-guide/distributing-work.rst``. +see https://datafusion.apache.org/python/user-guide/distributing-work/expressions.html. Build an expression in the driver, ship it to a pool of Ray actors, and have each actor evaluate it against its own slice of data. Python UDFs diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9c55f446c..3696d92a8 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, + SessionComponentsExportable, + SessionExtensionComponents, + SessionPlannerExportable, +) from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions from .plan import ExecutionPlan, LogicalPlan, Metric, MetricsSet @@ -127,13 +133,17 @@ "MetricsSet", "ParquetColumnOptions", "ParquetWriterOptions", + "QueryPlannerExportable", "RecordBatch", "RecordBatchStream", "RuntimeEnvBuilder", "SQLOptions", "ScalarUDF", + "SessionComponentsExportable", "SessionConfig", "SessionContext", + "SessionExtensionComponents", + "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..bbf08e84e 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, + SessionComponentsExportable, + SessionExtensionComponents, + SessionPlannerExportable, +) from datafusion.options import ( DEFAULT_MAX_INFER_SCHEMA, CsvReadOptions, @@ -129,7 +135,17 @@ def __arrow_c_array__( # noqa: D105 class TableProviderExportable(Protocol): """Type hint for object that has __datafusion_table_provider__ PyCapsule. - https://datafusion.apache.org/python/user-guide/io/table_provider.html + See :ref:`io_custom_table_provider` for registering one, and + :ref:`extension_providers` for writing one. + + Args: + session: See + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`. + For this getter it is a session when the provider is registered + through :py:meth:`SessionContext.register_table` and the host's + logical codec when it goes through + :py:meth:`datafusion.catalog.Schema.register_table`, so + duck-typing it is not optional. """ def __datafusion_table_provider__(self, session: Any) -> object: ... # noqa: D105 @@ -145,18 +161,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.""" @@ -545,6 +549,27 @@ class SessionContext: """This is the main interface for executing queries and creating DataFrames. See :ref:`user_guide_concepts` in the online documentation for more information. + + **A context is a handle on a session, not the session itself.** The + ``with_*`` methods — :py:meth:`with_logical_extension_codec`, + :py:meth:`with_physical_extension_codec`, + :py:meth:`with_python_udf_inlining`, and :py:meth:`with_extensions` — + return a new context 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 through + either handle is visible to both. + + A few things therefore belong to the session rather than to a handle, and + take effect even if the handle that set them is discarded: the query + planner (see :py:meth:`set_query_planner`), and the rebuild of an installed + foreign planner that follows installing a codec. :ref:`extension_sessions` + in the online documentation works through when that matters. + + **Keep a context alive for as long as anything derived from it is in use.** + A :py:class:`~datafusion.DataFrame`, logical plan, or exported capsule does + not extend the session's lifetime. Once the last context on a session is + collected, any operation that reaches an extension codec fails with + ``TaskContextProvider went out of scope over FFI boundary``. """ def __init__( @@ -1775,48 +1800,202 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non """Install a custom query planner on this session. The planner is imported through its ``__datafusion_query_planner__`` - PyCapsule and installed on this context, in the same way + PyCapsule, in the same way :meth:`~SessionContext.add_physical_optimizer_rule` installs a rule. - The query planner is part of the session state, so it applies to this - context and to every context sharing its session — including ones - already returned by - :meth:`~SessionContext.with_logical_extension_codec` and friends. - A session holds exactly one planner, so calling this again replaces the - previous one rather than layering. To chain planners, have the new + Returns nothing, because the planner lives in the session rather than + in a handle on it — installing one is visible to every context sharing + that session, including ones an earlier ``with_*`` call returned. See + :py:class:`SessionContext`. + + A session holds exactly one planner, so calling this again **replaces** + the previous one rather than layering. To chain planners, have the new planner wrap the capsule from - :meth:`~SessionContext.__datafusion_query_planner__`, captured - *before* the new planner is installed. + :meth:`~SessionContext.__datafusion_query_planner__`, captured *before* + the new planner is installed. - Install any extension codecs before a layered planner. Installing a - codec afterwards rebuilds the installed planner against it, but not the - fallback inside it, which keeps the codecs it was imported with. Note - also that the planner is built against the codecs of the context this - method is called on, so installing the same planner again on a different - handle rebinds the session's planner to *that* handle's codecs. See the - FFI extensions guide for the full multi-library registration recipe. + Install any extension codecs before a layered planner; the rebuild that + follows a later codec install does not reach the fallback inside one. + See :ref:`planner_codec_rebinding`, or prefer + :py:meth:`with_extensions`, which cannot capture a partial chain. 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. + Raises: + ValueError: If the capsule is not named ``datafusion_query_planner``. + Examples: - >>> from my_extension import DistributedQueryPlanner # doctest: +SKIP + A session exports its own planner, which is what you capture to + wrap: + + >>> from datafusion import SessionContext >>> ctx = SessionContext() - >>> ctx.set_query_planner(DistributedQueryPlanner()) # doctest: +SKIP - >>> ctx.sql("SELECT * FROM remote_table").collect() # doctest: +SKIP + >>> fallback = ctx.__datafusion_query_planner__() + >>> type(fallback).__name__ + 'PyCapsule' - Layer a planner on top of the one already installed by capturing - the existing planner first: + Skipped here (needs a built extension library): - >>> fallback = ctx.__datafusion_query_planner__() # doctest: +SKIP + >>> from my_extension import DistributedQueryPlanner # doctest: +SKIP >>> ctx.set_query_planner( ... DistributedQueryPlanner(fallback=fallback) ... ) # doctest: +SKIP + >>> ctx.sql("SELECT * FROM remote_table").collect() # doctest: +SKIP """ self.ctx.set_query_planner(planner) + def with_extensions( + self, *extensions: SessionComponentsExportable | SessionPlannerExportable + ) -> SessionContext: + """Create a new session context with the given extension bundles. + + This is the preferred way to install extension codecs and query + planners, because it removes the ordering question that installing them + by hand creates. + + Each argument is called twice, in two phases: + + 1. ``__datafusion_session_components__(ctx)`` on every extension, then + all the returned codecs are installed at once. + 2. ``__datafusion_session_planner__(ctx, fallback)`` on every + extension, **in argument order**, each handed the planner built so + far and the context carrying every bundle's codecs. The last + extension listed ends up outermost. + + An extension implements either hook or both. Return ``None`` from the + planner hook to contribute no planner; see + :py:class:`~datafusion.extensions.SessionPlannerExportable`. + + Nothing is written to the session until every hook has returned and + every capsule has been validated, so a hook that raises leaves the + session as it was. A hook that *mutates* the context it is handed — + registering a table, say — is not rolled back, which is why bundle + objects must be configuration-only. + + Shares its session with this context — see :py:class:`SessionContext`. + + See :ref:`extension_bundles` in the online documentation for why the + phases are split, how to contribute a bundle's two halves at different + positions, and a worked Rust implementation. + + 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 a hook + returns the wrong type, or if a codec is contributed as a bare + ``PyCapsule`` rather than an object exposing the getter. + ValueError: If two codecs claim the same id, or a getter returns a + capsule of the wrong kind. See + :py:meth:`with_logical_extension_codec` for how ids are + assigned. + + Examples: + The returned handle is a different object sharing one session, and + an empty call is legal: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> derived = ctx.with_extensions() + >>> derived is ctx + False + >>> ctx.from_pydict({"a": [1, 2]}, name="t") # doctest: +ELLIPSIS + DataFrame()... + >>> derived.table_exist("t") + True + >>> derived.logical_extension_codec_ids() + [] + + A runnable multi-bundle example, showing what a bundle returns and + how its codec ids accumulate, is in :ref:`extension_bundles`. + + Real usage. Skipped here (needs a built extension library); run + verbatim by ``test_with_extensions_docstring_example_still_runs``. + + >>> 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, (SessionComponentsExportable, SessionPlannerExportable) + ): + msg = ( + "Extension implements neither " + "__datafusion_session_components__ 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, SessionComponentsExportable): + continue + components = extension.__datafusion_session_components__(self) + if not isinstance(components, SessionExtensionComponents): + msg = ( + "__datafusion_session_components__ 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. @@ -2245,10 +2424,10 @@ def __datafusion_codec_id__(self) -> str: session, so two contexts can be installed on one session and a plan 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. + Contexts derived from the same session 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 — see :py:class:`SessionContext` — so + their payloads would be indistinguishable on decode. Examples: >>> from datafusion import SessionContext @@ -2265,15 +2444,50 @@ def __datafusion_logical_extension_codec__(self, session: Any = None) -> Any: ``session`` is accepted so a context satisfies the same protocol an extension library implements, where the argument is how the library - reaches the session it is being installed on. A context already is one, - so the argument is ignored. + reaches the host's codec. A context already carries one, so the + argument is ignored. See + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable` + for what an extension library does with it. + + Args: + session: Accepted and ignored. + + Returns: + A ``datafusion_logical_extension_codec`` PyCapsule. + + Examples: + >>> from datafusion import SessionContext + >>> type(SessionContext().__datafusion_logical_extension_codec__()).__name__ + 'PyCapsule' """ return self.ctx.__datafusion_logical_extension_codec__(session) def __datafusion_query_planner__(self, session: Any = None) -> Any: """Access the ``FFI_QueryPlanner`` PyCapsule for the current planner. - See :meth:`__datafusion_logical_extension_codec__` for ``session``. + This is how you capture the planner a session already has in order to + wrap it. Capture it *before* installing the new one, since + :py:meth:`set_query_planner` replaces rather than layers. + + Args: + session: Accepted and ignored. See + :meth:`__datafusion_logical_extension_codec__`. + + Returns: + A ``datafusion_query_planner`` PyCapsule wrapping the session's + current planner, exported for a foreign planner to delegate to. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> fallback = ctx.__datafusion_query_planner__() + >>> type(fallback).__name__ + 'PyCapsule' + + Both spellings work, since the argument is ignored: + + >>> type(ctx.__datafusion_query_planner__(ctx)).__name__ + 'PyCapsule' """ return self.ctx.__datafusion_query_planner__(session) @@ -2294,33 +2508,54 @@ def with_logical_extension_codec( affect decoding. A serialized plan records which codec wrote each payload, as a short id - taken from the codec's class. ``codec_id`` overrides that id and is - normally unnecessary. Pass it when installing from a bare ``PyCapsule``, - which has no class to take an id from, or when installing two instances - of one class, which otherwise claim the same id and raise ``ValueError``. + taken from the codec's class. - The returned context shares its session state with the original, so a - later registration on either is visible to both, and an installed query - planner is rebound on the shared session even if the returned context is - discarded. + Shares its session with this context — see :py:class:`SessionContext`. - See :ref:`ffi` in the online documentation for how ids are assigned, - what an extension codec has to implement, and a worked multi-library - registration recipe. + See :ref:`extension_codec_ids` in the online documentation for how ids + are assigned and what an extension codec has to implement, and + :py:meth:`with_extensions` for installing a library's codecs and planner + together. + + Args: + codec: Object implementing ``__datafusion_logical_extension_codec__`` + (see + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`), + or a raw ``datafusion_logical_extension_codec`` PyCapsule. + codec_id: Overrides the id the codec's payloads are tagged with. + Normally unnecessary. Pass it when installing two instances of + one class, which otherwise claim the same id, and when + installing a bare ``PyCapsule``: a capsule has no class to take + an id from, so it is given a random ``anon:`` id that differs on + every install, and plans it encodes can never be decoded + elsewhere. + + Returns: + A new context carrying this codec in addition to any already + installed. + + Raises: + ValueError: If the resolved id is already installed on this session. Examples: + A context exports its own codec, which stands in here for a real + library's: + >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx = ctx.with_logical_extension_codec( - ... my_library.Codec() - ... ) # doctest: +SKIP + >>> host = SessionContext() + >>> capsule = host.__datafusion_logical_extension_codec__() + >>> ctx = SessionContext().with_logical_extension_codec( + ... capsule, codec_id="my_library.Codec" + ... ) + >>> ctx.logical_extension_codec_ids() + ['my_library.Codec'] - Installing from a bare capsule, pinning the id so encoded - plans remain decodable on another session: + Without ``codec_id`` a bare capsule gets an anonymous id, which is + fine only if its plans never leave this session: - >>> ctx = ctx.with_logical_extension_codec( - ... capsule, codec_id="my_library.Codec" - ... ) # doctest: +SKIP + >>> ctx = SessionContext().with_logical_extension_codec(capsule) + >>> ctx.logical_extension_codec_ids()[0].startswith("anon:") + True """ new_internal = self.ctx.with_logical_extension_codec(codec, codec_id) new = SessionContext.__new__(SessionContext) @@ -2338,15 +2573,20 @@ def logical_extension_codec_ids(self) -> list[str]: DataFusion's own default codec is not listed. It handles whatever no installed codec claims, and it carries no identity to list. + Returns: + The installed codec ids, in install order. Empty if none are + installed. + Examples: >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx.logical_extension_codec_ids() + >>> host = SessionContext() + >>> host.logical_extension_codec_ids() [] - >>> ctx = ctx.with_logical_extension_codec( - ... my_library.Codec() - ... ) # doctest: +SKIP - >>> ctx.logical_extension_codec_ids() # doctest: +SKIP + >>> capsule = host.__datafusion_logical_extension_codec__() + >>> ctx = SessionContext().with_logical_extension_codec( + ... capsule, codec_id="my_library.Codec" + ... ) + >>> ctx.logical_extension_codec_ids() ['my_library.Codec'] """ return self.ctx.logical_extension_codec_ids() @@ -2355,6 +2595,18 @@ def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: """Access the PyCapsule FFI_PhysicalExtensionCodec. See :meth:`__datafusion_logical_extension_codec__` for ``session``. + + Args: + session: Accepted and ignored. + + Returns: + A ``datafusion_physical_extension_codec`` PyCapsule. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> type(ctx.__datafusion_physical_extension_codec__()).__name__ + 'PyCapsule' """ return self.ctx.__datafusion_physical_extension_codec__(session) @@ -2363,11 +2615,21 @@ def physical_extension_codec_ids(self) -> list[str]: See :py:meth:`logical_extension_codec_ids`. + Returns: + The installed codec ids, in install order. Empty if none are + installed. + Examples: >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx.physical_extension_codec_ids() + >>> host = SessionContext() + >>> host.physical_extension_codec_ids() [] + >>> capsule = host.__datafusion_physical_extension_codec__() + >>> ctx = SessionContext().with_physical_extension_codec( + ... capsule, codec_id="my_library.PhysicalCodec" + ... ) + >>> ctx.physical_extension_codec_ids() + ['my_library.PhysicalCodec'] """ return self.ctx.physical_extension_codec_ids() @@ -2386,16 +2648,27 @@ def with_physical_extension_codec( :py:meth:`with_logical_extension_codec` does, including when to pass ``codec_id`` and what the returned context shares. See that method. + Args: + codec: As :py:meth:`with_logical_extension_codec`, for the physical + getter. + codec_id: As :py:meth:`with_logical_extension_codec`. + + Returns: + A new context carrying this codec in addition to any already + installed. + + Raises: + ValueError: If the resolved id is already installed on this session. + Examples: >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx = ctx.with_physical_extension_codec( - ... my_library.PhysicalCodec() - ... ) # doctest: +SKIP - - >>> ctx = ctx.with_physical_extension_codec( + >>> host = SessionContext() + >>> capsule = host.__datafusion_physical_extension_codec__() + >>> ctx = SessionContext().with_physical_extension_codec( ... capsule, codec_id="my_library.PhysicalCodec" - ... ) # doctest: +SKIP + ... ) + >>> ctx.physical_extension_codec_ids() + ['my_library.PhysicalCodec'] """ new_internal = self.ctx.with_physical_extension_codec(codec, codec_id) new = SessionContext.__new__(SessionContext) @@ -2405,10 +2678,6 @@ def with_physical_extension_codec( def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext: """Control whether Python UDFs are embedded in serialized expressions. - ``enabled`` is keyword-only and required: callers must pick a - mode explicitly. Fresh sessions inline UDFs (``enabled=True`` - behavior) until this method overrides the toggle. - With ``enabled=True``, serialized expressions carry the Python code for any scalar, aggregate, or window UDFs they reference. The receiver rebuilds the UDFs from those bytes and does not @@ -2437,14 +2706,18 @@ def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext: :func:`pickle.loads` on untrusted bytes remains unsafe regardless of the toggle. - Returns a new :class:`SessionContext` with the toggle applied; - the original context's own codec settings are unchanged. The - returned context shares its session state with the original, so - a later registration on either is visible to both. If a custom - query planner is installed, it is rebuilt against the new codecs - on the shared session, so the original context plans with them - too. This happens on the shared session, so it takes effect even - if the returned context is discarded. + Shares its session with this context — see + :py:class:`SessionContext`. The original context's own codec + settings are unchanged. + + Args: + enabled: Whether to embed Python UDFs in serialized + expressions. Keyword-only and required, so callers must + pick a mode explicitly. Fresh sessions behave as + ``enabled=True`` until this method overrides the toggle. + + Returns: + A new :class:`SessionContext` with the toggle applied. Examples: >>> import pyarrow as pa diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index c198b646c..18ce3554d 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -576,6 +576,13 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr: (sufficient for built-ins and Python UDFs, plus any UDFs registered on the global context). + Note that the global context carries no extension codecs, so + falling back to it cannot decode a payload written by one. A + worker that must decode extension-codec payloads has to install a + context carrying those codecs — see + :func:`datafusion.ipc.set_worker_ctx` and + :meth:`SessionContext.with_extensions`. + .. warning:: Security Decoding may invoke ``cloudpickle.loads`` on bytes embedded in the payload, which executes arbitrary Python code. Treat @@ -627,19 +634,19 @@ def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]: portable across minor versions. See :meth:`to_bytes` for details on what travels by value vs. by reference. - Examples: - >>> import pickle - >>> from datafusion import col, lit - >>> e = col("a") * lit(2) - >>> pickle.loads(pickle.dumps(e)).canonical_name() - 'a * Int64(2)' - The encoding side honors a driver-side sender context installed via :func:`datafusion.ipc.set_sender_ctx` — that is how :meth:`SessionContext.with_python_udf_inlining` propagates through ``pickle.dumps``. The sender context is read by ``__reduce__``, so :func:`copy.copy` and :func:`copy.deepcopy` — which also go through ``__reduce__`` — pick it up too. + + Examples: + >>> import pickle + >>> from datafusion import col, lit + >>> e = col("a") * lit(2) + >>> pickle.loads(pickle.dumps(e)).canonical_name() + 'a * Int64(2)' """ return (Expr._reconstruct, (self.to_bytes(get_sender_ctx()),)) diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py new file mode 100644 index 000000000..77ae92fc2 --- /dev/null +++ b/python/datafusion/extensions.py @@ -0,0 +1,338 @@ +# 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:`SessionComponentsExportable` 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()) + +Codecs and planners install in two phases: every +:py:class:`SessionComponentsExportable` runs first and its codecs are installed, +then every :py:class:`SessionPlannerExportable` runs in argument order. A bundle +implements either hook or both. Bundle order is significant for planners, which +nest, and irrelevant for codecs, which accumulate. + +Of the four names here, only the two bundle hooks are ``@runtime_checkable``, +because :py:meth:`~datafusion.context.SessionContext.with_extensions` +dispatches on them from Python. :py:class:`QueryPlannerExportable` is a type +hint only, matching the other capsule-getter protocols in +:py:mod:`datafusion.user_defined` and :py:mod:`datafusion.catalog`. + +See :ref:`extension_bundles` in the online documentation for why the phases are +split and for a worked implementation. +""" + +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", + "SessionComponentsExportable", + "SessionExtensionComponents", + "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, and duck-type it — see + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable` for + ``session``. + + Unlike the two bundle hooks in this module, this protocol is a type hint + only: it is not ``@runtime_checkable``, so ``isinstance`` against it raises + ``TypeError``. + + Examples: + A :py:class:`~datafusion.context.SessionContext` satisfies this + protocol, which is what lets a foreign planner wrap the one a session + already has: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> type(ctx.__datafusion_query_planner__(ctx)).__name__ + 'PyCapsule' + + The protocol itself is not runtime-checkable: + + >>> from datafusion import QueryPlannerExportable + >>> try: + ... isinstance(ctx, QueryPlannerExportable) + ... except TypeError as e: + ... print("runtime_checkable" in str(e)) + True + """ + + 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:`SessionComponentsExportable.__datafusion_session_components__` + 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`. + + 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 + () + + 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 + + >>> ctx = SessionContext() + >>> components = SessionExtensionComponents( + ... logical_extension_codecs=(NamedCodec( + ... ctx.__datafusion_logical_extension_codec__() + ... ),) + ... ) + >>> components.logical_extension_codecs[0].__datafusion_codec_id__ + 'my_library.v1' + >>> components.physical_extension_codecs + () + + 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(ctx)) + 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. + + Objects exposing ``__datafusion_logical_extension_codec__``, never bare + ``PyCapsule`` objects — a codec's id is read off the object it is handed + over as. Any iterable is accepted and stored as a tuple. See + :ref:`extension_bundles_codecs_are_objects`. + """ + + physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = () + """Physical codecs to add to the session's codec chain, in declaration order. + + As :py:attr:`logical_extension_codecs`, for + ``__datafusion_physical_extension_codec__``. + """ + + def __post_init__(self) -> None: + """Normalize each codec 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 `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 `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 SessionComponentsExportable(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. + + A bundle that also contributes a query planner implements + :py:class:`SessionPlannerExportable` alongside this protocol. + + Args: + ctx: The session the components will run on. Take the task-context + provider off it. Do **not** read its codec chains expecting to find + this call's codecs, including your own: this hook runs before + anything is installed, so ``ctx`` still carries whatever chains the + receiver had. :py:class:`SessionPlannerExportable` is the hook that + sees the completed chains — see :ref:`extension_bundles_two_phases`. + + Returns: + The codecs this bundle contributes. + + Examples: + >>> from datafusion import ( + ... SessionExtensionComponents, + ... SessionComponentsExportable, + ... ) + >>> class MyLibraryExtension: + ... def __datafusion_session_components__(self, ctx): + ... return SessionExtensionComponents() + >>> isinstance(MyLibraryExtension(), SessionComponentsExportable) + True + >>> isinstance(object(), SessionComponentsExportable) + False + """ + + def __datafusion_session_components__( # 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. See + :ref:`extension_bundles_two_phases`. + + **Return ``None`` to contribute no planner**, leaving ``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 that planner as a foreign one and every + later plan crosses an FFI boundary that was not there before. A bundle that + decides at runtime it has 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. + + Returns: + A planner wrapping ``fallback``, or ``None`` to contribute none. + + Examples: + A real library returns its own planner wrapping ``fallback``, e.g. + ``my_library.Planner(fallback=fallback)``. The two degenerate cases are + worth contrasting, because both plan queries successfully and only one + of them is the no-op: + + >>> from datafusion import SessionContext, SessionPlannerExportable + >>> class Contributes: + ... def __datafusion_session_planner__(self, ctx, fallback): + ... return None # the no-op: session keeps its own planner + >>> class Replaces: + ... def __datafusion_session_planner__(self, ctx, fallback): + ... return fallback # installs it as a *foreign* planner + >>> for bundle in (Contributes(), Replaces()): + ... ctx = SessionContext().with_extensions(bundle) + ... ctx.sql("SELECT 1 AS n").collect()[0].column(0).to_pylist() + [1] + [1] + + >>> isinstance(Contributes(), SessionPlannerExportable) + True + >>> isinstance(object(), SessionPlannerExportable) + False + """ + + def __datafusion_session_planner__( # noqa: D105 + self, ctx: SessionContext, fallback: _PyCapsule + ) -> QueryPlannerExportable | _PyCapsule | None: ... diff --git a/python/datafusion/ipc.py b/python/datafusion/ipc.py index 487abd4c3..d9f63e267 100644 --- a/python/datafusion/ipc.py +++ b/python/datafusion/ipc.py @@ -175,11 +175,16 @@ def set_sender_ctx(ctx: SessionContext) -> None: Examples: >>> from datafusion import SessionContext - >>> from datafusion.ipc import set_sender_ctx, get_sender_ctx + >>> from datafusion.ipc import ( + ... clear_sender_ctx, + ... get_sender_ctx, + ... set_sender_ctx, + ... ) >>> driver = SessionContext().with_python_udf_inlining(enabled=False) >>> set_sender_ctx(driver) >>> get_sender_ctx() is driver True + >>> clear_sender_ctx() """ _local.sender_ctx = ctx @@ -224,7 +229,9 @@ def _resolve_ctx( Priority: explicit argument > worker context > global context. Falling back to the global :class:`SessionContext` (instead of a freshly constructed one) preserves any registrations the user has - installed on it. + installed on it. It carries no extension codecs, though, so a worker + that must decode payloads written by one has to install a context + carrying them via :func:`set_worker_ctx`. Examples: >>> from datafusion import SessionContext diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index b2c6eab3e..8d03bae2c 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -82,9 +82,9 @@ def display_indent_schema(self) -> str: def display_graphviz(self) -> str: """Print the graph visualization of the logical plan. - Returns a `format`able structure that produces lines meant for graphical display - using the `DOT` language. This format can be visualized using software from - [`graphviz`](https://graphviz.org/) + Returns a formattable structure that produces lines meant for graphical + display using the ``DOT`` language. This format can be visualized using + software from `graphviz `_. """ return self._raw_plan.display_graphviz() @@ -92,21 +92,33 @@ def display_graphviz(self) -> str: def from_bytes(ctx: SessionContext, data: bytes) -> LogicalPlan: """Create a LogicalPlan from serialized protobuf bytes. - Decoding routes through the session's installed - `LogicalExtensionCodec`. Tables created in memory from record - batches are currently not supported. + Decoding routes through the codecs installed on ``ctx`` with + :py:meth:`~datafusion.SessionContext.with_logical_extension_codec`. + Tables created in memory from record batches are currently not + supported. + + Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and + positional, and there is no fallback to a worker or global context. + + See Also: + :py:meth:`to_bytes`, :py:meth:`ExecutionPlan.from_bytes`, + :py:meth:`datafusion.Expr.from_bytes`. """ return LogicalPlan(df_internal.LogicalPlan.from_bytes(ctx.ctx, data)) def to_bytes(self, ctx: SessionContext | None = None) -> bytes: """Convert a LogicalPlan to serialized protobuf bytes. - When ``ctx`` is supplied, encoding routes through the session's - installed `LogicalExtensionCodec` so user FFI codecs (registered - via :py:meth:`SessionContext.with_logical_extension_codec`) see - the encode path. With ``ctx=None`` a default codec is used. - Tables created in memory from record batches are currently not - supported. + When ``ctx`` is supplied, encoding routes through the codecs + installed on it with + :py:meth:`~datafusion.SessionContext.with_logical_extension_codec`, + so extension codecs see the encode path. With ``ctx=None`` a + default codec is used. Tables created in memory from record + batches are currently not supported. + + See Also: + :py:meth:`from_bytes`, :py:meth:`ExecutionPlan.to_bytes`, + :py:meth:`datafusion.Expr.to_bytes`. """ ctx_arg = ctx.ctx if ctx is not None else None return self._raw_plan.to_bytes(ctx_arg) @@ -145,7 +157,7 @@ def __init__(self, plan: df_internal.ExecutionPlan) -> None: self._raw_plan = plan def children(self) -> list[ExecutionPlan]: - """Get a list of children `ExecutionPlan` that act as inputs to this plan. + """Get a list of children ``ExecutionPlan`` that act as inputs to this plan. The returned list will be empty for leaf nodes such as scans, will contain a single value for unary nodes, or two values for binary nodes (such as joins). @@ -173,18 +185,35 @@ def partition_count(self) -> int: def from_bytes(ctx: SessionContext, data: bytes) -> ExecutionPlan: """Create an ExecutionPlan from serialized protobuf bytes. - Decoding routes through the session's installed - `PhysicalExtensionCodec`. Tables created in memory from record - batches are currently not supported. + Decoding routes through the codecs installed on ``ctx`` with + :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`. + Tables created in memory from record batches are currently not + supported. + + Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and + positional, and there is no fallback to a worker or global context. + + See Also: + :py:meth:`to_bytes`, :py:meth:`LogicalPlan.from_bytes`. """ return ExecutionPlan(df_internal.ExecutionPlan.from_bytes(ctx.ctx, data)) def to_bytes(self, ctx: SessionContext | None = None) -> bytes: """Convert an ExecutionPlan into serialized protobuf bytes. - When ``ctx`` is supplied, encoding routes through the session's - installed `PhysicalExtensionCodec`. Tables created in memory - from record batches are currently not supported. + When ``ctx`` is supplied, encoding routes through the codecs + installed on it with + :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`. + Tables created in memory from record batches are currently not + supported. + + Round-tripping through this method and :py:meth:`from_bytes` is how + an extension library checks that its own codec claimed its nodes, + rather than a codec installed earlier in the chain — see + :ref:`extension_codec_order`. + + See Also: + :py:meth:`from_bytes`, :py:meth:`LogicalPlan.to_bytes`. """ ctx_arg = ctx.ctx if ctx is not None else None return self._raw_plan.to_bytes(ctx_arg) diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index eafcefdaf..f8d273177 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -116,11 +116,6 @@ def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]: class LogicalExtensionCodecExportable(Protocol): """Type hint for objects exposing ``__datafusion_logical_extension_codec__``. - ``session`` is the :py:class:`~datafusion.context.SessionContext` the codec - is being installed on. Take the task context provider from it rather than - building a session of your own, so the decode callbacks resolve names - against the session that runs the query. - Implement the codec itself exactly as you would for a session that installs only yours. A session may hold several codecs, but each payload records the codec that wrote it and is only ever handed back to that codec, so there is @@ -132,6 +127,43 @@ class LogicalExtensionCodecExportable(Protocol): of this protocol and is rarely needed: declare it when renaming your class must not stop older plans from decoding, or when one library installs two instances that own disjoint slices of the wire format. + + Args: + session: A handle on the session this codec is being installed on. Take + the task context provider from it rather than building a session of + your own, so the decode callbacks resolve names against the session + that runs the query. + + **Duck-type it; do not check its type.** This argument is the + reference for every capsule getter in the protocol, and across them + it is not always a session at all: the provider and catalog getters + receive the host's logical codec as a bare ``PyCapsule``. Even when + it *is* a session 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__`` — + everything the protocol asks of it — but + ``isinstance(session, SessionContext)`` is ``False`` even though its + ``repr`` reads ``datafusion.SessionContext``. Pass it to + ``ffi_logical_codec_from_pycapsule``, which handles both. + + The two bundle hooks are the exception: + :py:class:`~datafusion.extensions.SessionComponentsExportable` and + :py:class:`~datafusion.extensions.SessionPlannerExportable` are + dispatched from Python and receive the wrapper. See + :ref:`extension_getter_argument`. + + Returns: + A ``datafusion_logical_extension_codec`` PyCapsule. + + Examples: + A :py:class:`~datafusion.context.SessionContext` satisfies this + protocol, which is how you obtain a real capsule to test against: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> type(ctx.__datafusion_logical_extension_codec__(ctx)).__name__ + 'PyCapsule' """ def __datafusion_logical_extension_codec__( # noqa: D105 diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 3c95835af..1a15e4a54 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_components__(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_components__(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_components__(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_components__(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_components__(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_components__(self, ctx): + return self.inner.__datafusion_session_components__(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_components__(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_components__(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_components__(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_components__(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_docstrings.py b/python/tests/test_docstrings.py new file mode 100644 index 000000000..9ac1f0b76 --- /dev/null +++ b/python/tests/test_docstrings.py @@ -0,0 +1,246 @@ +# 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. + +"""Shape checks on the docstrings that ship in the wheel. + +These are about *form*, not content: a docstring that has grown into an essay, +a documented API with no example, or a pointer at the guide that the rendered +HTML cannot follow. See "Python Function Docstrings" in ``AGENTS.md``. + +Content is checked elsewhere. ``--doctest-modules`` executes the examples, and +Sphinx builds with ``--fail-on-warning`` so a broken ``:ref:`` fails the docs +build. Neither of those notices a docstring that is simply too long, or one +that never had an example to run. +""" + +from __future__ import annotations + +import ast +import inspect +import pathlib +import re + +import datafusion +import pytest +from datafusion import SessionContext, extensions + +PACKAGE_ROOT = pathlib.Path(datafusion.__file__).parent +REPO_ROOT = PACKAGE_ROOT.parents[1] + +# The hook reference in the extension guide is a hand-written table of every +# `__datafusion_*__` name. Nothing about adding a hook forces it to be updated, +# so the table is compared against the names the package actually dispatches. +# +# "Dispatches" means a *site*, not an occurrence: the name spelled where the +# host looks the hook up or defines its own, not everywhere the name is +# written. Scanning raw text would make the table's contents depend on prose — +# a doc-comment contrasting a hook with one that was removed, or naming a +# hypothetical, would have to be either deleted or added to the table, and +# neither is right. Structure answers the question text cannot. +HOOK_REFERENCE = REPO_ROOT / "docs" / "source" / "extension-guide" / "index.md" +HOOK_NAME = re.compile(r"^__datafusion_[a-z_]+__$") +HOOK_TABLE_ROW = re.compile(r"^\| `(__datafusion_[a-z_]+__)`") + +# A Rust dispatch site is a string literal holding nothing but the hook name — +# what `hasattr`, `getattr`, and `call_capsule_getter` are handed — or a `fn` +# of that name, which is a hook the host itself implements. An error message +# that merely embeds the name (`"__datafusion_scalar_udf__ does not exist"`) +# is prose and does not count; every such message sits beside a real lookup. +RUST_HOOK_SITE = re.compile( + r'"(__datafusion_[a-z_]+__)"|\bfn\s+(__datafusion_[a-z_]+__)\b' +) + +# Above this, a docstring has stopped being a contract and become a +# narrative. Move the argument to a guide page under `docs/source/` and leave +# a one-line pointer. Raising this is not the fix. +# +# Deliberately has no waiver list: at the time of writing the longest +# docstring in the package is `udaf` at 92 lines, which is a legitimately +# long UDF-authoring reference with many worked examples. If a genuinely +# necessary docstring ever exceeds this, prefer splitting the API. +MAX_DOCSTRING_LINES = 95 + +# An *inclusion* list, to be grown — not an exclusion list to be shrunk. +# Enforcing "every public callable has an example" package-wide is a separate +# project; several hundred callables do not have one yet. These are the +# extension-protocol surface, where an undocumented method is the difference +# between an extension author succeeding and filing an issue. +DOCSTRINGS_REQUIRING_EXAMPLES: list[tuple[str, object]] = [ + *[ + (f"SessionContext.{name}", getattr(SessionContext, name)) + for name in ( + "with_extensions", + "set_query_planner", + "with_logical_extension_codec", + "with_physical_extension_codec", + "logical_extension_codec_ids", + "physical_extension_codec_ids", + "with_python_udf_inlining", + "__datafusion_codec_id__", + "__datafusion_logical_extension_codec__", + "__datafusion_physical_extension_codec__", + "__datafusion_query_planner__", + ) + ], + *[(name, getattr(extensions, name)) for name in extensions.__all__], +] + + +def _iter_docstrings() -> list[tuple[str, int]]: + """Yield ``(location, line count)`` for every docstring under ``python/``.""" + found = [] + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if not isinstance( + node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef) + ): + continue + doc = ast.get_docstring(node, clean=True) + if not doc: + continue + name = getattr(node, "name", "") + rel = path.relative_to(PACKAGE_ROOT.parent) + line = getattr(node, "lineno", 1) + found.append((f"{rel}:{line} {name}", len(doc.splitlines()))) + return found + + +def test_no_docstring_is_an_essay() -> None: + """No docstring exceeds the length ceiling.""" + too_long = [ + (where, count) + for where, count in _iter_docstrings() + if count > MAX_DOCSTRING_LINES + ] + assert not too_long, ( + "These docstrings exceed " + f"{MAX_DOCSTRING_LINES} lines and have become narrative rather than " + "contract:\n" + + "\n".join(f" {where} — {count} lines" for where, count in too_long) + + "\n\nMove the argument to a guide page under docs/source/ and leave a " + "one-line pointer with a :ref:. See 'One canonical home per claim' in " + "AGENTS.md." + ) + + +@pytest.mark.parametrize( + ("name", "obj"), + DOCSTRINGS_REQUIRING_EXAMPLES, + ids=[name for name, _ in DOCSTRINGS_REQUIRING_EXAMPLES], +) +def test_extension_api_has_a_doctest(name: str, obj: object) -> None: + """Every extension-protocol member carries at least one example.""" + doc = inspect.getdoc(obj) + assert doc, f"{name} has no docstring" + assert ">>>" in doc, ( + f"{name} has a docstring but no example. Everything on the extension " + "protocol needs one, even if the realistic usage has to be marked " + "+SKIP — in which case a runnable example goes above it. See " + "'Examples that need a compiled extension' in AGENTS.md." + ) + + +def _rust_hook_sites() -> set[str]: + """Hook names Rust looks up or defines. See :data:`RUST_HOOK_SITE`.""" + found: set[str] = set() + for path in sorted((REPO_ROOT / "crates").rglob("*.rs")): + for match in RUST_HOOK_SITE.finditer(path.read_text()): + found.add(match.group(1) or match.group(2)) + return found + + +def _python_hook_sites() -> set[str]: + """Hook names the Python wrappers call, declare, or look up by string. + + Read off the syntax tree rather than the text, so a name is counted when + it is an attribute being accessed (``extension.__datafusion_x__(ctx)``), a + method being defined (the protocol stubs), or a string standing alone (a + ``getattr`` argument) — and not when it merely appears inside a docstring. + """ + found: set[str] = set() + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + name = node.name + elif isinstance(node, ast.Attribute): + name = node.attr + elif isinstance(node, ast.Constant) and isinstance(node.value, str): + name = node.value + else: + continue + if HOOK_NAME.match(name): + found.add(name) + return found + + +def test_hook_reference_table_lists_every_hook() -> None: + """The guide's hook table matches the hooks the package dispatches.""" + if not HOOK_REFERENCE.is_file(): + pytest.skip("running against an installed wheel, without docs/ or crates/") + + dispatched = _rust_hook_sites() | _python_hook_sites() + + documented = { + match.group(1) + for line in HOOK_REFERENCE.read_text().splitlines() + if (match := HOOK_TABLE_ROW.match(line)) + } + + problems = [ + f" dispatched but not in the table: {name}" + for name in sorted(dispatched - documented) + ] + [ + f" in the table but dispatched from nowhere: {name}" + for name in sorted(documented - dispatched) + ] + assert not problems, ( + f"{HOOK_REFERENCE.relative_to(REPO_ROOT)} is out of sync with the " + "hooks in crates/ and python/datafusion/:\n" + + "\n".join(problems) + + "\n\nOnly dispatch sites count — a name looked up by string, an " + "attribute accessed, or a method defined. Naming a hook in prose does " + "not put it in this set, and does not belong in the table either." + ) + + +def test_no_dead_pointers_at_the_guide() -> None: + """A docstring naming the guide must link it, not just mention it.""" + role = re.compile(r":(ref|doc|py:\w+):`") + dead = [] + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if not isinstance( + node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef) + ): + continue + doc = ast.get_docstring(node, clean=True) + if not doc or "guide" not in doc.lower(): + continue + if role.search(doc): + continue + rel = path.relative_to(PACKAGE_ROOT.parent) + name = getattr(node, "name", "") + dead.append(f"{rel}:{getattr(node, 'lineno', 1)} {name}") + assert not dead, ( + "These docstrings send the reader to a guide without a resolvable " + "link, which is a dead end in the rendered HTML:\n" + + "\n".join(f" {where}" for where in dead) + + "\n\nUse :ref:`some_label` naming the specific section." + ) diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py index fea4cc91f..0e0a3965f 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", + "SessionComponentsExportable", + "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)