diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 294ebfb3a..3e39f3d36 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -154,8 +154,19 @@ guards this. Its `WHERE` clause is load-bearing: filter pushdown upgrades the weak handle during logical optimization, before plan serialization could fail first for an unrelated reason. +`SessionContext.with_extensions` is where this rule is easiest to get wrong, +because "bind the components to the context you are about to return" reads like +an instruction to derive one first. It is not: the factories are handed the +receiver, and the returned handle shares its allocation. There is nothing to +keep alive separately and nothing to garbage-collect out from under a provider. + `SessionContext.enable_url_table` is the one method that mints a second -allocation for a session. Its result must not outlive the receiver. +allocation for a session. Its result must not outlive the receiver, and it also +forks the session's `SessionState` while keeping its id, so two handles report +one `session_id()` with divergent configuration. That is a bug rather than a +design — tracked in +[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708) +— so do not cite it as precedent for deriving a replacement context. ## Rule 7 — installing a planner mutates the session, and says so diff --git a/Cargo.lock b/Cargo.lock index c7632732a..6a7f68438 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1209,6 +1209,7 @@ dependencies = [ "datafusion-catalog", "datafusion-common", "datafusion-ffi", + "datafusion-proto", "datafusion-python-util", "datafusion-session", "pyo3", diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 84182ff19..4fb40ae3d 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -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`, `_install_extensions` — 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()) @@ -1485,7 +1490,7 @@ impl PySessionContext { ) -> PyDataFusionResult { let id = { let this = slf.borrow(); - resolve_codec_id(&codec, codec_id, &this.logical_codec.codec_ids())? + resolve_codec_id(&codec, codec_id, None, &this.logical_codec.codec_ids())? }; let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); @@ -1550,7 +1555,7 @@ impl PySessionContext { ) -> PyDataFusionResult { let id = { let this = slf.borrow(); - resolve_codec_id(&codec, codec_id, &this.physical_codec.codec_ids())? + resolve_codec_id(&codec, codec_id, None, &this.physical_codec.codec_ids())? }; let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); @@ -1608,6 +1613,79 @@ impl PySessionContext { derived.set_session_query_planner(None); derived } + + /// Commit a `with_extensions` transaction onto this context. + /// + /// Private support method for `SessionContext.with_extensions`. `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. Codec capsules are imported and validated before anything is + /// committed, so a failure leaves the session untouched. + /// + /// The codec chains belong to the returned handle rather than to + /// `SessionState`, so a codec-only install onto a session with no FFI + /// planner writes no state at all. The session is written only when there + /// is a planner to bind — one a bundle supplied, or one already installed + /// that has to be rebuilt against the new chains — and that write goes + /// through this context's own `state_ref()`, so providers bound to it stay + /// valid. + #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))] + pub fn _install_extensions<'py>( + slf: &Bound<'py, Self>, + logical_codecs: Vec<(Bound<'py, PyAny>, Bound<'py, PyAny>)>, + physical_codecs: Vec<(Bound<'py, PyAny>, Bound<'py, PyAny>)>, + planner: Option>, + ) -> 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(), + ) + }; + + // Each codec arrives paired with the bundle that contributed it. A + // bundle is a plain object, so its identity is as stable across + // processes as an exporting codec class's, which is what lets a bare + // capsule from a bundle be named instead of randomized. See + // `resolve_codec_id`. + for (codec, bundle) in logical_codecs { + let id = resolve_codec_id(&codec, None, Some(&bundle), &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, bundle) in physical_codecs { + let id = resolve_codec_id(&codec, None, Some(&bundle), &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); + + let planner = planner + .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))) + .transpose()?; + + let installed = Self { + ctx: Arc::clone(&slf.borrow().ctx), + logical_codec, + physical_codec, + }; + // Bind the planner only once the codec chains are final, and through + // the new handle so it carries them. Passing `None` still rebuilds + // whichever planner the session already holds against the new chains, + // exactly as `with_logical_extension_codec` does. + installed.set_session_query_planner(planner); + + Ok(installed) + } } impl PySessionContext { @@ -1776,13 +1854,21 @@ impl PySessionContext { /// 3. The exporting object's `module.QualName`, which is the library's own /// import path and therefore already stable across processes. This is the /// common case and asks nothing of existing extension libraries. -/// 4. For a bare `PyCapsule` there is nothing stable to read — every capsule -/// reports the same type — so mint a fresh random id. Payloads tagged this -/// way decode correctly within the session lineage that installed the -/// codec, because the chain is cloned along with the id, and fail with a -/// pointed error everywhere else. Randomness is the point: an id drawn from -/// a namespace another session can mint the same value from — a counter, a -/// chain position — would let an unrelated codec answer for these bytes. +/// 4. For a bare `PyCapsule` contributed through `with_extensions`, the +/// identity of the bundle that contributed it, resolved by arms 2 and 3 +/// above. A bundle is a plain object, so its import path is library-owned +/// and exactly as stable as an exporting codec class's — the capsule was +/// only unnameable because a capsule carries no type of its own, not +/// because nothing stable was in reach. +/// 5. For a bare `PyCapsule` with no bundle behind it there is nothing stable +/// to read — every capsule reports the same type — so mint a fresh random +/// id. Payloads tagged this way decode correctly within the session lineage +/// that installed the codec, because the chain is cloned along with the id, +/// and fail with a pointed error everywhere else. Randomness is the point: +/// an id drawn from a namespace another session can mint the same value +/// from — a counter, a chain position — would let an unrelated codec answer +/// for these bytes. That is also why arm 4 does not disambiguate two +/// capsules from one bundle by position; it lets them collide instead. /// /// 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 @@ -1792,20 +1878,28 @@ impl PySessionContext { fn resolve_codec_id( codec: &Bound<'_, PyAny>, explicit: Option, + bundle: Option<&Bound<'_, PyAny>>, existing: &[&str], ) -> PyResult { - let id = derive_codec_id(codec, explicit)?; + let id = derive_codec_id(codec, explicit, bundle)?; if existing.contains(&id.as_str()) { return Err(PyValueError::new_err(format!( "An extension codec with id '{id}' is already installed on this session. Two \ codecs cannot share an id, because a payload names its codec by id when it is \ - decoded. Pass `codec_id=` to give this one a distinct identity." + 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) } -fn derive_codec_id(codec: &Bound<'_, PyAny>, explicit: Option) -> PyResult { +fn derive_codec_id( + codec: &Bound<'_, PyAny>, + explicit: Option, + bundle: Option<&Bound<'_, PyAny>>, +) -> PyResult { if let Some(id) = explicit { return Ok(id); } @@ -1815,6 +1909,14 @@ fn derive_codec_id(codec: &Bound<'_, PyAny>, explicit: Option) -> PyResu return declared.extract::(); } if codec.is_instance_of::() { + // Name the capsule after whoever handed it over, if anyone did. The + // bundle goes through the same resolution, so a bundle that declares + // `__datafusion_codec_id__` pins an id that survives renaming its + // class, exactly as an exporting codec can. A bundle is never itself a + // capsule, so this cannot recurse into the random arm below. + if let Some(bundle) = bundle { + return derive_codec_id(bundle, None, None); + } return Ok(format!( "{ANONYMOUS_CODEC_ID_PREFIX}{}", Uuid::new_v4() diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index d86858a83..b9f2f541c 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -277,10 +277,20 @@ 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 bare `PyCapsule`.** A capsule has no class to take a name from. Installed + through `with_extensions`, it is named after the extension that contributed it — + an extension is a plain object, so its import path is library-owned and just as + stable across processes as a codec class's. Installed directly through + `with_logical_extension_codec` or `with_physical_extension_codec` there is nothing + to fall back on, 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. + + One extension contributing two bare capsules of the same kind is refused, because + both resolve to that one extension's id. 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 name one of them by + wrapping it in an object declaring `__datafusion_codec_id__`. - **A class you intend to rename.** The id follows the class name, so renaming stops older plans from decoding. Declare `__datafusion_codec_id__` on the exporting object to pin an id that survives the rename. @@ -343,6 +353,64 @@ The current FFI logical codec supports providers and UDFs but not arbitrary cust `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and local build commands. +### Extension bundles: `with_extensions` + +The chaining above works, but it makes the caller responsible for ordering: the codecs +have to be installed before the planner, because a planner is built against whatever +codec chains exist when it is installed, and a codec added afterwards rebinds it. Get +that wrong and the planner encodes through a chain that is missing a library. + +`SessionContext.with_extensions` removes the ordering question. An extension library +exposes a bundle object implementing `__datafusion_session_extension__`: + +```python +class MyEngineExtension: + def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: + # Create fresh components bound to `ctx` on every call. `ctx` is the + # session the components will run on. + return SessionExtensionComponents( + logical_extension_codecs=(self._make_logical_codec(ctx),), + physical_extension_codecs=(self._make_physical_codec(ctx),), + query_planner=self._make_planner(ctx), + ) +``` + +The host passes the context to every factory, installs all the codecs, binds the +planner against the final codec chains, and returns a handle on that session in a +single step: + +```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())) +``` + +Extensions are processed left to right and their codecs are appended to the chain in +that order. As above, order affects only encoding — decoding routes by id. At most one +extension per call may supply a query planner. + +Nothing is written to the session until every factory has returned and every capsule +has been validated, so a factory that raises leaves the session exactly as it was. A +factory that mutates the context it is handed — registering a table, say — is not +rolled back, which is why bundle objects must be configuration-only: create fresh +components on each call, never cache bound components, and do not retain the context +passed in. + +Like every other derivation, the returned context is a handle on the *same* session as +the receiver — see [What a derived context shares](#what-a-derived-context-shares). +Only the Python-side codec chains belong to the returned handle; the planner is +installed on the shared session and takes effect even if that handle is discarded. + +The session owns every installed component's task-context provider, and dependent +objects do not extend its lifetime. A `DataFrame`, logical plan, or capsule can outlive +every context on the session, but any operation that reaches an FFI codec after the +last one is collected fails with `TaskContextProvider went out of scope over FFI +boundary`. Keep a context alive for as long as objects derived from it are in use. + +`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust +implementation of the protocol, including taking the task-context provider off the +supplied context and constructing a Python `SessionExtensionComponents`. + ### Capsule getters receive the session they are installed on `__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and @@ -419,15 +487,24 @@ registered straight back into that same session, which would close the cycle `SessionContext.enable_url_table` is the one exception. It clones the underlying `SessionContext`, so the returned context has an allocation of its own and must not -outlive the receiver. +outlive the receiver. It also forks the session's state while keeping its id, so two +handles report one `session_id()` with divergent configuration. That is a bug rather +than a design, tracked in +[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708); +do not copy the pattern. ### What a derived context shares -`with_logical_extension_codec`, `with_physical_extension_codec`, and -`with_python_udf_inlining` return a new `SessionContext` wrapping the *same* underlying -session. Only the Python-side codec settings differ; catalogs, tables, registered -functions, and configuration are the one shared session, so a registration on either -side is visible to both. +`with_logical_extension_codec`, `with_physical_extension_codec`, +`with_python_udf_inlining`, and `with_extensions` return a new `SessionContext` wrapping +the *same* underlying session. Only the Python-side codec settings differ; catalogs, +tables, registered functions, and configuration are the one shared session, so a +registration on either side is visible to both. + +There is one `Arc` per session, which is what makes the weak +`FFI_TaskContextProvider` scheme work: a component bound through any handle stays valid +while *any* handle on that session is alive, so there is no way to bind a component to +an intermediate handle and have it dangle when that handle is dropped. `set_query_planner` does not return anything. The query planner lives in `SessionState`, so it is a property of the session rather than of a handle on it, and installing one is @@ -516,6 +593,9 @@ the original handle rebinds the session's planner back to the original handle's instead, which is the trap `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` pins. +`with_extensions` sidesteps the ordering question entirely: it installs every codec +before it binds the planner, so there is no "afterwards" for a bundle's own planner. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 257749c3a..4506778e7 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -83,6 +83,13 @@ way `add_physical_optimizer_rule` does and returns nothing — the query planner lives in `SessionState`, so it belongs to the session rather than to a particular handle on it. See the {ref}`ffi` guide for the full protocol. +If a library ships codecs *and* a planner, prefer +`SessionContext.with_extensions(bundle)` over installing each piece by hand. It +installs every codec before it binds the planner, so the planner cannot end up +carrying a chain that a later `with_logical_extension_codec` call has grown. +The library exposes a bundle object implementing +`__datafusion_session_extension__`; see the {ref}`ffi` guide. + ### Mismatched extension libraries now fail loudly Objects imported through the capsule protocol are checked against the major diff --git a/examples/datafusion-ffi-query-planner-example/Cargo.toml b/examples/datafusion-ffi-query-planner-example/Cargo.toml index 4d02c69f1..263f034b8 100644 --- a/examples/datafusion-ffi-query-planner-example/Cargo.toml +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -31,6 +31,7 @@ datafusion = { workspace = true } datafusion-catalog = { workspace = true, default-features = false } datafusion-common = { workspace = true, default-features = false } datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } datafusion-session = { workspace = true } async-trait = { workspace = true } datafusion-python-util.workspace = true diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 66bc45196..7cce6ab60 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -41,7 +41,22 @@ 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 the `__datafusion_session_extension__` protocol: it +receives the session it is being installed on, binds fresh codec and planner +components to that session's task-context provider, and returns them as +`SessionExtensionComponents`. The host installs every codec before it binds the +planner, so the planner cannot be left carrying a chain that has since grown. + +The integration tests also cover the low-level chaining setup: ```python config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index c6ef2072a..ed9e6184a 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,23 @@ 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_ffi_example import ( IsNullUDF, MyCatalogProvider, @@ -30,7 +42,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 +704,317 @@ def test_query_planner_rejects_invalid_config(max_rows: str): ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() +class ProviderCodecsExtension: + """Bundles the provider library's codecs for ``with_extensions``. + + These codecs keep their own private task-context provider, so they only + need to be created once; the bundle can hand out the same exporters on + every call. + """ + + def __init__(self) -> None: + self.logical_codec = MyLogicalExtensionCodec() + self.physical_codec = MyPhysicalExtensionCodec() + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(self.logical_codec,), + physical_extension_codecs=(self.physical_codec,), + ) + + +class _NamedCodec: + """Forwards a codec's capsule getters under a declared id. + + ``with_extensions`` takes no ``codec_id=``, so an extension that ships a + codec class another extension also ships declares + ``__datafusion_codec_id__`` on the object it hands over. Both getters are + forwarded because one wrapper stands in for whichever kind it wraps. + """ + + def __init__(self, codec: object, codec_id: str) -> None: + self._codec = codec + self.__datafusion_codec_id__ = codec_id + + def __datafusion_logical_extension_codec__(self, session: object = None) -> object: + return self._codec.__datafusion_logical_extension_codec__(session) + + def __datafusion_physical_extension_codec__(self, session: object = None) -> object: + return self._codec.__datafusion_physical_extension_codec__(session) + + +class IdentifiedProviderCodecsExtension(ProviderCodecsExtension): + """``ProviderCodecsExtension`` whose codecs carry ids of their own.""" + + def __init__(self, prefix: str) -> None: + super().__init__() + self.logical = _NamedCodec(self.logical_codec, f"{prefix}.logical") + self.physical = _NamedCodec(self.physical_codec, f"{prefix}.physical") + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(self.logical,), + physical_extension_codecs=(self.physical,), + ) + + +def test_with_extensions_three_library_query(): + """One with_extensions call installs provider codecs and a planner bundle, + and a real non-empty plan flows across the three libraries.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + provider_ext = ProviderCodecsExtension() + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(provider_ext, planner_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert batches[0].column(1).to_pylist() == [False, False, False] + assert planner_ext.plan_calls() >= 1 + assert planner_ext.last_max_rows() == 3 + assert planner_ext.foreign_session_observed() + assert planner_ext.foreign_provider_observed() + assert planner_ext.foreign_plan_observed() + assert provider_ext.logical_codec.table_provider_encode_calls() > 0 + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_encode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +def test_with_extensions_names_a_rust_bundles_capsules_after_the_bundle(): + """A Rust bundle hands its codecs over as bare capsules, and they are + named after the bundle's own import path. + + 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 gets when installed directly — would make the plan undecodable + there. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + + bundle_id = "datafusion_ffi_query_planner_example.MyPlannerExtension" + assert bundle_id in ctx.logical_extension_codec_ids() + assert bundle_id in ctx.physical_extension_codec_ids() + + # The provider bundle hands over objects, so those keep their own class + # names rather than picking up 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_shares_the_session_with_the_source(): + """``with_extensions`` returns a handle on the source's session, and the + bundle's task-context provider resolves against that one session. + + There is one ``Arc`` per session, so a component bound + during installation cannot be left pointing at a handle that is dropped + later. A `SET` issued through the *source* after installation is therefore + visible to the provider the bundle bound, which is what a codec's decode + callback resolves through. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + source = SessionContext(config) + planner_ext = MyPlannerExtension() + result = source.with_extensions(ProviderCodecsExtension(), planner_ext) + + assert result.session_id() == source.session_id() + + # Registrations and config changes go through the source handle only. + source.register_table("numbers", MyTableProvider(1, 6, 1)) + source.sql("SET ffi_query_planner.max_rows = 2").collect() + + batches = result.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + assert planner_ext.max_rows_through_provider() == 2 + + # Symmetrically, the codec chains installed on the shared session are in + # force for the source handle too. + batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_with_extensions_survives_dropping_source_and_bundles(): + """The returned handle alone keeps the installed components alive. + + The context ``with_extensions`` was called on is a temporary here, and the + bundle objects are dropped with it. Both share their allocation with the + returned handle, so the components' task-context provider stays valid. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + gc.collect() + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_with_extensions_sees_state_changes_after_install(): + """Tables, UDFs, and config changes made after installation are visible + to the planner and to provider callbacks.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=4)) + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(ProviderCodecsExtension(), planner_ext) + + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + ctx.sql("SET ffi_query_planner.max_rows = 2").collect() + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + assert planner_ext.max_rows_through_provider() == 2 + + +def test_with_extensions_bundle_is_reusable(): + """Installing the same bundle into two contexts binds fresh components to + each destination.""" + planner_ext = MyPlannerExtension() + + config_a = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx_a = SessionContext(config_a).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_a.register_table("numbers", MyTableProvider(1, 6, 1)) + + config_b = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + ctx_b = SessionContext(config_b).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_b.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx_a.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + + batches = ctx_b.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert planner_ext.last_max_rows() == 3 + + +def test_with_extensions_failure_leaves_source_usable(): + """A failing factory after a successful one leaves the source context + fully functional.""" + + class BoomExtension: + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + msg = "boom" + raise RuntimeError(msg) + + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + + with pytest.raises(RuntimeError, match="boom"): + source.with_extensions(MyPlannerExtension(), BoomExtension()) + + # No planner was installed, so the default planner runs unrestricted. + batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] + + +def test_with_extensions_rebinds_existing_planner(): + """Codec-only bundles installed on a context that already has an FFI + planner rebind that planner to the new codec chains.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + planner = MyQueryPlanner() + ctx = SessionContext(config) + ctx.set_query_planner(planner) + provider_ext = ProviderCodecsExtension() + ctx = ctx.with_extensions(provider_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.last_max_rows() == 2 + # The planner only sees these codecs if it was rebound to the chains + # built during with_extensions. + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +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 +1041,83 @@ def test_composed_codecs_with_query_planner(): assert logical_codec.table_provider_encode_calls() > 0 assert logical_codec.table_provider_decode_calls() > 0 assert physical_codec.execution_plan_decode_calls() > 0 + + +class _DocstringExampleExtension: + """Stand-in for the ``my_extension`` bundle named in the docstring. + + The docstring shows a single engine bundle taking a scheduler address, + which is what a real distributed engine ships: one object contributing a + planner *and* the codecs that carry its plans. Here that is assembled from + this repository's two example libraries. The address is accepted and + ignored; everything else the example touches is the real API. + """ + + def __init__(self, endpoint: str) -> None: + self.endpoint = endpoint + self._codecs = ProviderCodecsExtension() + self._planner = MyPlannerExtension() + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + codecs = self._codecs.__datafusion_session_extension__(ctx) + planner = self._planner.__datafusion_session_extension__(ctx) + return SessionExtensionComponents( + logical_extension_codecs=( + *codecs.logical_extension_codecs, + *planner.logical_extension_codecs, + ), + physical_extension_codecs=( + *codecs.physical_extension_codecs, + *planner.physical_extension_codecs, + ), + query_planner=planner.query_planner, + ) + + +def test_with_extensions_docstring_example_still_runs(): + """Run the ``with_extensions`` docstring example verbatim. + + The example is marked ``+SKIP`` because the main suite has no built FFI + extension to import, which is exactly how such an example rots. Here the + statements are parsed out of the live docstring, the skip is dropped, and + each one is executed and its output compared. + + Only names are redirected: ``my_extension`` resolves to the bundle above, + and ``SessionContext`` supplies the config this library's planner reads. + A renamed method, a changed signature, or a wrong expected output in the + docstring fails here. + """ + examples = doctest.DocTestParser().get_examples( + inspect.getdoc(SessionContext.with_extensions) + ) + assert examples, "with_extensions docstring has no examples to check" + for example in examples: + example.options.pop(doctest.SKIP, None) + + module = types.ModuleType("my_extension") + module.DistributedEngineExtension = _DocstringExampleExtension + + def make_context(config: SessionConfig | None = None) -> SessionContext: + # Accept a config so the example is free to pass one. Supplying it + # positionally the way the real constructor does keeps a docstring + # edit failing as a doctest diff rather than as a TypeError in here. + config = SessionConfig() if config is None else config + return SessionContext(config.with_extension(MyPlannerConfig(max_rows=3))) + + test = doctest.DocTest( + examples, + {"SessionContext": make_context}, + "SessionContext.with_extensions", + None, + None, + None, + ) + output = io.StringIO() + sys.modules["my_extension"] = module + try: + results = doctest.DocTestRunner().run(test, out=output.write) + finally: + del sys.modules["my_extension"] + assert results.failed == 0, output.getvalue() diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs new file mode 100644 index 000000000..0a7be7453 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -0,0 +1,300 @@ +// 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::Ordering; +use std::sync::{Arc, Mutex}; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, TableReference}; +use datafusion::datasource::TableProvider; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::{Extension, LogicalPlan}; +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_task_context_provider_from_pycapsule, get_tokio_runtime, +}; +use datafusion_session::QueryPlanner; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +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. +/// +/// Only populated when a codec in this bundle is actually consulted. The host +/// dispatches a framed payload straight to the codec whose id it names, so a +/// decline-all codec like the ones here is normally never asked to decode. The +/// binding itself is proved by [`MyPlannerExtension::max_rows_through_provider`], +/// which reads the provider directly rather than waiting for a callback. +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); + } +} + +/// Records the task context resolved by the FFI wrapper, then declines by +/// delegating to the default codec so the host's codec chain falls through to +/// the codec that owns the payload. +struct ObservingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec, + observed: ObservedMaxRows, +} + +impl fmt::Debug for ObservingLogicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ObservingLogicalExtensionCodec") + .finish_non_exhaustive() + } +} + +impl LogicalExtensionCodec for ObservingLogicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[LogicalPlan], + ctx: &TaskContext, + ) -> Result { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: SchemaRef, + ctx: &TaskContext, + ) -> Result> { + record_task_ctx(&self.observed, ctx); + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode_table_provider(table_ref, node, buf) + } +} + +/// Physical companion to [`ObservingLogicalExtensionCodec`]. +struct ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec, + observed: ObservedMaxRows, +} + +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> { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + self.inner.try_encode(node, buf, proto_converter) + } +} + +/// Extension bundle for `SessionContext.with_extensions`. +/// +/// Mirrors how a distributed engine such as Ballista packages its session +/// extensions: the object itself is reusable configuration, and every +/// `__datafusion_session_extension__` call creates fresh codec and planner +/// components bound to the task-context provider of the context it receives. +#[pyclass( + from_py_object, + name = "MyPlannerExtension", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Default, Clone)] +pub(crate) struct MyPlannerExtension { + observations: Arc, + observed_max_rows: ObservedMaxRows, + 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. + /// + /// Usually empty: the host routes a framed payload to the codec named in + /// it, so codecs that own nothing are not consulted. + fn decode_max_rows_seen(&self) -> Vec { + self.observed_max_rows + .lock() + .map(|observed| observed.clone()) + .unwrap_or_default() + } + + /// `ffi_query_planner.max_rows` read through the task-context provider + /// this bundle was last bound to. + /// + /// Resolving the provider is what a codec's decode callback does, so this + /// answers which session those callbacks would resolve against -- the + /// context `with_extensions` returned, not the one it was called on. + /// Returns ``None`` if the bundle was never installed, or if the context it + /// was bound to has been dropped: the provider holds it weakly. + fn max_rows_through_provider(&self) -> Option { + let provider = self.bound_provider.lock().ok()?.clone()?; + let task_ctx = Arc::::try_from(&provider).ok()?; + planner_config_from_options(task_ctx.session_config().options()) + .ok() + .map(|config| config.max_rows) + } + + fn __datafusion_session_extension__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + // Bind every component to the context supplied by the host, which is + // the session the components will run on. Components must not be + // cached across calls: each installation may target a different + // session. + // + // The task-context provider comes off that context rather than from a + // `SessionContext` built here, so the codecs' decode callbacks resolve + // names against the session that will actually run the query. + let provider = ffi_task_context_provider_from_pycapsule(&ctx)?; + if let Ok(mut bound) = self.bound_provider.lock() { + *bound = Some(provider.clone()); + } + let runtime = get_tokio_runtime().handle().clone(); + + let logical: Arc = Arc::new(ObservingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_logical = + FFI_LogicalExtensionCodec::new(logical, Some(runtime.clone()), provider.clone()); + let logical_capsule = create_logical_extension_capsule(py, &ffi_logical)?; + + let physical: Arc = + Arc::new(ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_physical = + FFI_PhysicalExtensionCodec::new(physical, Some(runtime), provider.clone()); + let physical_capsule = create_physical_extension_capsule(py, &ffi_physical)?; + + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + fallback: None, + }); + // The planner takes the host's codecs, not ones built here. Installing + // the codecs above rebuilds the planner against them anyway, 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); + let planner_capsule = create_query_planner_capsule(py, &ffi_planner)?; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("logical_extension_codecs", (logical_capsule,))?; + kwargs.set_item("physical_extension_codecs", (physical_capsule,))?; + kwargs.set_item("query_planner", planner_capsule)?; + components.call((), Some(&kwargs)) + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs index c505c1ce7..70d4a42c5 100644 --- a/examples/datafusion-ffi-query-planner-example/src/lib.rs +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -18,15 +18,18 @@ use pyo3::prelude::*; use crate::config::MyPlannerConfig; +use crate::extension::MyPlannerExtension; use crate::planner::MyQueryPlanner; mod config; +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::()?; 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..733536d21 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -52,14 +52,14 @@ 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, +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. - used_fallback: AtomicBool, + pub(crate) used_fallback: AtomicBool, } impl fmt::Debug for PlannerObservations { @@ -104,8 +104,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 +147,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 +160,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] diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9c55f446c..4b02a383e 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -92,6 +92,10 @@ ) from .dataframe_formatter import configure_formatter from .expr import Expr, WindowFrame +from .extensions import ( + SessionExtensionComponents, + SessionExtensionExportable, +) from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions from .plan import ExecutionPlan, LogicalPlan, Metric, MetricsSet @@ -134,6 +138,8 @@ "ScalarUDF", "SessionConfig", "SessionContext", + "SessionExtensionComponents", + "SessionExtensionExportable", "Table", "TableFunction", "TableProviderFactory", @@ -146,6 +152,7 @@ "common", "configure_formatter", "expr", + "extensions", "functions", "ipc", "lit", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 644c7b445..ac453bb72 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -69,6 +69,11 @@ ) from datafusion.dataframe import DataFrame from datafusion.expr import sort_list_to_raw_sort_list +from datafusion.extensions import ( + QueryPlannerExportable, + SessionExtensionComponents, + SessionExtensionExportable, +) from datafusion.options import ( DEFAULT_MAX_INFER_SCHEMA, CsvReadOptions, @@ -145,18 +150,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.""" @@ -1817,6 +1810,147 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non """ self.ctx.set_query_planner(planner) + def with_extensions( + self, *extensions: SessionExtensionExportable + ) -> SessionContext: + """Create a new session context with the given extension bundles. + + This is the preferred way to install FFI extensions that need a + task-context provider (extension codecs and query planners). Each + extension's ``__datafusion_session_extension__`` method is called with + this context so it can bind its components to the session they will + run on, then all components are installed in one step. This avoids the + pitfalls of chaining :py:meth:`with_logical_extension_codec`, + :py:meth:`with_physical_extension_codec`, and + :py:meth:`set_query_planner` by hand, where the codecs a planner was + built against can end up stale. + + Codecs compose with the existing chain and with each other: extensions + are processed left to right and their codecs are appended to the chain + in that order. Decoding routes by codec id, so the order matters only + for encoding. At most one extension may supply a query planner. If none + does, an existing FFI planner is rebound to the final codec chains. + + Each codec is named after its exporting class, as + :py:meth:`with_logical_extension_codec` describes. A codec handed over + as a bare ``PyCapsule`` has no class to take a name from, so it is + named after the extension that contributed it — the extension's import + path is library-owned and stable across processes, so plans it writes + stay decodable elsewhere. Declare ``__datafusion_codec_id__`` on the + extension to pin that name against a later class rename, or on the + object handed over to name a codec directly. + + Like the individual ``with_*`` methods, the returned context shares its + session with this one: catalogs, tables, registered functions, and + configuration are the one session, so a registration on either side is + visible to both, and the planner is installed on that shared session + even if the returned context is discarded. Only the Python-side codec + chains are specific to the returned handle. + + No state is written until every extension has run and every capsule has + been validated, so an extension that raises or returns invalid + components leaves the session as it was. The exception is an extension + that mutates the context it is handed — registering a table, say — + which is not rolled back. Extension factories should treat that context + as configuration-only. + + The session owns the installed components' task-context providers, and + dependent objects do not extend its lifetime. Keep a context on the + session alive for as long as DataFrames or plans derived from it are in + use; FFI operations after the last one is collected raise an error. + + Args: + extensions: Extension bundles to install, in the order their + codecs join the chain. + + Returns: + A new context with all extension components installed. + + Raises: + TypeError: If an argument does not implement the protocol or + returns something other than a + :py:class:`SessionExtensionComponents`. + ValueError: If no extensions are given, more than one extension + supplies a query planner, or two codecs claim the same id. An + extension that contributes two instances of one codec class, + or two bare capsules of the same kind, must declare + ``__datafusion_codec_id__`` on at least one of them; the + collision is refused rather than resolved by position, because + a positional id would break stored plans the first time the + extension reordered what it returns. + + Examples: + The example is skipped here because it needs a built FFI + extension library, which this package does not ship. It is run + verbatim against a real one by + ``test_with_extensions_docstring_example_still_runs`` in + ``examples/datafusion-ffi-query-planner-example``, so it cannot + drift from the API. + + >>> from my_extension import DistributedEngineExtension # doctest: +SKIP + >>> ctx = SessionContext().with_extensions( + ... DistributedEngineExtension("scheduler:50050") + ... ) # doctest: +SKIP + >>> batches = ctx.sql("SELECT 1 AS n").collect() # doctest: +SKIP + >>> batches[0].column(0).to_pylist() # doctest: +SKIP + [1] + """ + if not extensions: + msg = "with_extensions requires at least one extension" + raise ValueError(msg) + for extension in extensions: + if not isinstance(extension, SessionExtensionExportable): + msg = ( + "Extension does not implement __datafusion_session_extension__: " + f"{extension!r}" + ) + raise TypeError(msg) + + # Bind every component 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, and `_install_extensions` writes the final state through + # that same session. + # + # Each codec is paired with the extension that contributed it. A codec + # handed over as a bare capsule has no class to take an id from, so it + # is named after that extension rather than randomized. + logical_codecs: list[ + tuple[LogicalExtensionCodecExportable | _PyCapsule, object] + ] = [] + physical_codecs: list[ + tuple[PhysicalExtensionCodecExportable | _PyCapsule, object] + ] = [] + planner: QueryPlannerExportable | _PyCapsule | None = None + for extension in extensions: + components = extension.__datafusion_session_extension__(self) + if not isinstance(components, SessionExtensionComponents): + msg = ( + "__datafusion_session_extension__ must return " + "SessionExtensionComponents, got " + f"{type(components).__name__} from {extension!r}" + ) + raise TypeError(msg) + logical_codecs.extend( + (codec, extension) for codec in components.logical_extension_codecs + ) + physical_codecs.extend( + (codec, extension) for codec in components.physical_extension_codecs + ) + if components.query_planner is not None: + if planner is not None: + msg = ( + "Multiple extensions supplied a query planner; a " + "session context has exactly one. Layer planners " + "explicitly instead." + ) + raise ValueError(msg) + planner = components.query_planner + + new = SessionContext.__new__(SessionContext) + new.ctx = self.ctx._install_extensions(logical_codecs, physical_codecs, planner) + return new + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. @@ -2246,9 +2380,11 @@ def __datafusion_codec_id__(self) -> str: written through one will not be decoded by the other. Contexts derived from the same session — including the ones returned by - :py:meth:`with_logical_extension_codec` and - :py:meth:`with_python_udf_inlining` — report the same id, so only one of - them can be installed on a given session. + :py:meth:`with_logical_extension_codec`, + :py:meth:`with_python_udf_inlining`, and :py:meth:`with_extensions` — + report the same id, so only one of them can be installed on a given + session. That is the intended answer: they are one session, so their + payloads would be indistinguishable on decode. Examples: >>> from datafusion import SessionContext diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py new file mode 100644 index 000000000..8e35c720d --- /dev/null +++ b/python/datafusion/extensions.py @@ -0,0 +1,160 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Protocols and value types for installing extensions on a session context. + +An *extension* is a reusable configuration object — typically shipped by a +separate compiled library — that contributes components to a +:py:class:`~datafusion.context.SessionContext`. It implements +:py:class:`SessionExtensionExportable` by returning a +:py:class:`SessionExtensionComponents` describing what it contributes, and is +installed with :py:meth:`~datafusion.context.SessionContext.with_extensions`:: + + ctx = SessionContext().with_extensions(MyLibraryExtension()) + +Installing through ``with_extensions`` rather than by chaining the individual +``with_*`` methods matters for components that hold a task-context provider: +the extension is handed the session its components will run on, and every +codec is installed before the query planner is bound against them, so no +planner is left carrying a codec chain that has since grown. See the FFI +extensions guide in the contributor documentation for the full rationale. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from _typeshed import CapsuleType as _PyCapsule + + from datafusion.context import SessionContext + from datafusion.user_defined import ( + LogicalExtensionCodecExportable, + PhysicalExtensionCodecExportable, + ) + +__all__ = [ + "QueryPlannerExportable", + "SessionExtensionComponents", + "SessionExtensionExportable", +] + + +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:`~datafusion.context.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 + + +@dataclass(frozen=True) +class SessionExtensionComponents: + """Components an extension contributes to a session context. + + Returned by :py:meth:`SessionExtensionExportable.__datafusion_session_extension__` + and consumed by + :py:meth:`~datafusion.context.SessionContext.with_extensions`. Every + component must be created against the context passed to that method; + components bound to a different session hold a task-context provider for + that other session and cannot be rebound. + + Codecs may be handed over either as objects exposing the capsule getters or + as bare ``PyCapsule`` objects. A bare capsule carries no class to take a + codec id from, so it is named after the extension that contributed it. An + extension contributing two bare capsules of the same kind therefore has to + name at least one of them itself, by wrapping it in an object declaring + ``__datafusion_codec_id__``. + + Examples: + A bundle that contributes nothing is valid, and is what the defaults + describe: + + >>> from datafusion import SessionExtensionComponents + >>> components = SessionExtensionComponents() + >>> components.logical_extension_codecs + () + >>> components.query_planner is None + True + + A bundle that contributes one kind of component names it, leaving + the rest empty: + + >>> components = SessionExtensionComponents( + ... query_planner=my_library.make_planner(ctx) + ... ) # doctest: +SKIP + """ + + logical_extension_codecs: tuple[ + LogicalExtensionCodecExportable | _PyCapsule, ... + ] = () + """Logical codecs to add to the session's codec chain, in declaration order.""" + + physical_extension_codecs: tuple[ + PhysicalExtensionCodecExportable | _PyCapsule, ... + ] = () + """Physical codecs to add to the session's codec chain, in declaration order.""" + + query_planner: QueryPlannerExportable | _PyCapsule | None = None + """Optional query planner. + + At most one extension per + :py:meth:`~datafusion.context.SessionContext.with_extensions` call may + supply one. + """ + + +@runtime_checkable +class SessionExtensionExportable(Protocol): + """Type hint for extension bundles installable via ``with_extensions``. + + Runtime-checkable, so ``isinstance`` answers whether an object implements + the protocol. Only the presence of the method is checked, which is the same + question :py:meth:`~datafusion.context.SessionContext.with_extensions` asks + before calling it. + + Implementations are reusable configuration objects: they must create fresh + components on every call using the context supplied by + :py:meth:`~datafusion.context.SessionContext.with_extensions`, and must not + retain that context or cache the components they bound to it, since the + next call may install onto a different session. They should also avoid + mutating the context they are handed — a registration made during binding + is not rolled back if a later extension fails. + + Examples: + >>> from datafusion import ( + ... SessionExtensionComponents, + ... SessionExtensionExportable, + ... ) + >>> class MyLibraryExtension: + ... def __datafusion_session_extension__(self, ctx): + ... return SessionExtensionComponents() + >>> isinstance(MyLibraryExtension(), SessionExtensionExportable) + True + >>> isinstance(object(), SessionExtensionExportable) + False + """ + + def __datafusion_session_extension__( # noqa: D105 + self, ctx: SessionContext + ) -> SessionExtensionComponents: ... diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 3c95835af..0418acde2 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,237 @@ def test_contexts_sharing_a_session_share_the_planner(ctx): assert sibling.session_id() == ctx.session_id() +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): + self.exporter = SessionContext() + self.bound_ctx = None + + def __datafusion_session_extension__(self, ctx): + self.bound_ctx = ctx + return SessionExtensionComponents( + logical_extension_codecs=( + self.exporter.__datafusion_logical_extension_codec__(), + ), + physical_extension_codecs=( + self.exporter.__datafusion_physical_extension_codec__(), + ), + ) + + +class _PlannerExtension: + """Contributes the receiving session's own exported planner.""" + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + query_planner=ctx.__datafusion_query_planner__() + ) + + +def test_with_extensions_requires_an_extension(ctx): + with pytest.raises(ValueError, match="at least one extension"): + ctx.with_extensions() + + +def test_with_extensions_rejects_non_extension(ctx): + with pytest.raises(TypeError, match="__datafusion_session_extension__"): + ctx.with_extensions(object()) + + +def test_with_extensions_rejects_bad_components(ctx): + class BadExtension: + def __datafusion_session_extension__(self, ctx): + return 42 + + with pytest.raises(TypeError, match="SessionExtensionComponents"): + ctx.with_extensions(BadExtension()) + + +def test_with_extensions_rejects_multiple_planners(ctx): + with pytest.raises(ValueError, match="query planner"): + ctx.with_extensions(_PlannerExtension(), _PlannerExtension()) + + +def test_with_extensions_rejects_bad_codec_capsule(ctx): + class BadCodecExtension: + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=(ctx.__datafusion_task_context_provider__(),), + ) + + with pytest.raises( + ValueError, match="Expected name 'datafusion_logical_extension_codec'" + ): + ctx.with_extensions(BadCodecExtension()) + + +def test_with_extensions_names_bare_capsules_after_the_extension(ctx): + """A capsule has no class to take a codec id from, so it is named after + the extension that contributed it. + + The extension's import path is library-owned and stable across processes, + so plans written through the codec stay decodable on another session — a + session-private random id would not be. + """ + result = ctx.with_extensions(_CodecOnlyExtension()) + + expected = f"{_CodecOnlyExtension.__module__}._CodecOnlyExtension" + assert result.logical_extension_codec_ids() == [expected] + assert result.physical_extension_codec_ids() == [expected] + + +def test_with_extensions_extension_can_pin_its_codec_id(ctx): + """``__datafusion_codec_id__`` on the extension survives a class rename.""" + + class PinnedExtension(_CodecOnlyExtension): + __datafusion_codec_id__ = "my_library.v1" + + result = ctx.with_extensions(PinnedExtension()) + assert result.logical_extension_codec_ids() == ["my_library.v1"] + + +def test_with_extensions_codec_id_on_the_codec_beats_the_extension(ctx): + """Naming the handed-over object wins over the extension's name. + + This is how an extension contributing more than one bare capsule of a kind + tells them apart. + """ + + class NamedCapsule: + 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 + + class TwoNamedCodecs: + def __init__(self): + self.exporter = SessionContext() + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=( + NamedCapsule( + self.exporter.__datafusion_logical_extension_codec__(), + "my_library.first", + ), + NamedCapsule( + 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_bare_capsules_from_one_extension(ctx): + """Both capsules resolve to the one extension's 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 TwoCapsuleExtension: + def __init__(self): + self.exporter = SessionContext() + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=( + self.exporter.__datafusion_logical_extension_codec__(), + self.exporter.__datafusion_logical_extension_codec__(), + ), + ) + + with pytest.raises(ValueError, match="__datafusion_codec_id__"): + ctx.with_extensions(TwoCapsuleExtension()) + + +def test_with_extensions_leaves_an_exporting_object_its_own_id(ctx): + """A codec handed over as an object keeps its own identity. + + The extension's name is a fallback for capsules only; it never overrides + an id the codec itself carries. + """ + exporter = SessionContext() + + class ObjectCodecExtension: + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents(logical_extension_codecs=(exporter,)) + + result = ctx.with_extensions(ObjectCodecExtension()) + assert result.logical_extension_codec_ids() == [exporter.__datafusion_codec_id__] + + +def test_with_extensions_installs_codecs_and_planner(ctx): + ctx.register_record_batches( + "extensions_test", + [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]], + ) + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension, _PlannerExtension()) + + assert result.table_exist("extensions_test") + # In-memory tables need a real extension codec to round-trip through the + # FFI planner, so query plans that don't serialize a table provider. + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_binds_to_the_receiving_session(ctx): + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension) + + # Factories are handed the receiver itself, so a component bound during + # installation targets the session the returned handle also wraps. There + # is no intermediate context that could be collected out from under it. + assert extension.bound_ctx is ctx + assert result.session_id() == ctx.session_id() + + # One session: a registration through either handle is visible to both. + ctx.register_record_batches( + "bound_test", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + assert result.table_exist("bound_test") + + +def test_with_extensions_survives_source_collection(): + extension = _CodecOnlyExtension() + result = SessionContext().with_extensions(extension, _PlannerExtension()) + gc.collect() + + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_failure_leaves_source_usable(ctx): + class BoomExtension: + def __datafusion_session_extension__(self, ctx): + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions(_CodecOnlyExtension(), BoomExtension()) + + batches = ctx.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index cf6719ecf..7e8bd3f2d 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -28,6 +28,17 @@ 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( + { + # Support method for SessionContext.with_extensions. + "_install_extensions", + } +) + + 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 +78,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)