From 4facb0c4de609c9de74701554d9e2be2716ce278 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:09:16 -0400 Subject: [PATCH 01/12] Add atomic SessionContext.with_extensions API Installing FFI extension codecs and query planners by chaining the existing with_* methods can bind task-context providers to intermediate contexts that are later collected, breaking the weak provider reference over the FFI boundary. with_extensions creates one destination context, passes it to each extension factory so components bind to that exact context, and installs everything in a single state write. Co-Authored-By: Claude Fable 5 --- crates/core/src/context.rs | 75 +++++++++++++++++ python/datafusion/__init__.py | 4 + python/datafusion/context.py | 146 ++++++++++++++++++++++++++++++++++ python/tests/test_context.py | 117 +++++++++++++++++++++++++++ 4 files changed, 342 insertions(+) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 84182ff19..04427f49c 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1608,6 +1608,81 @@ impl PySessionContext { derived.set_session_query_planner(None); derived } + + /// Create the destination context for a `with_extensions` transaction. + /// + /// Private support method for `SessionContext.with_extensions`. The + /// returned context is the single `Arc` that every FFI + /// task-context provider created during the transaction must target; + /// `_install_extensions` later mutates its state in place rather than + /// deriving a new context. + pub fn _derive_for_extensions(&self) -> Self { + Self { + ctx: Arc::new(SessionContext::new_with_state(self.ctx.state())), + logical_codec: Arc::clone(&self.logical_codec), + physical_codec: Arc::clone(&self.physical_codec), + } + } + + /// Commit a `with_extensions` transaction onto this context. + /// + /// Private support method for `SessionContext.with_extensions`; `self` + /// must be a context produced by `_derive_for_extensions`. Codec capsules + /// are imported and validated before any state change, so a failure + /// leaves the context untouched. The final state is written through this + /// context's own `state_ref()`, never a derived context, so FFI + /// task-context 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>, + physical_codecs: Vec>, + 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(), + ) + }; + + for codec in logical_codecs { + let id = resolve_codec_id(&codec, None, &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_codec_id(&codec, None, &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 derived = Self { + ctx: Arc::clone(&slf.borrow().ctx), + logical_codec, + physical_codec, + }; + // Bind the planner only once the codec chains are final, and through + // the derived 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. + derived.set_session_query_planner(planner); + + Ok(derived) + } } impl PySessionContext { diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9c55f446c..86d0054b3 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -80,6 +80,8 @@ RuntimeEnvBuilder, SessionConfig, SessionContext, + SessionExtensionComponents, + SessionExtensionExportable, SQLOptions, ) from .dataframe import ( @@ -134,6 +136,8 @@ "ScalarUDF", "SessionConfig", "SessionContext", + "SessionExtensionComponents", + "SessionExtensionExportable", "Table", "TableFunction", "TableProviderFactory", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 644c7b445..ded8e8f39 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -46,6 +46,7 @@ import uuid import warnings +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol try: @@ -157,6 +158,49 @@ class QueryPlannerExportable(Protocol): 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:`SessionContext.with_extensions`. Every component + must be created against the context passed to that method; components bound + to any other context hold a task-context provider for the wrong session and + cannot be rebound. + + Attributes: + logical_extension_codecs: Logical codecs to add to the session's codec + chain, in declaration order. + physical_extension_codecs: Physical codecs to add to the session's + codec chain, in declaration order. + query_planner: Optional query planner. At most one extension per + :py:meth:`SessionContext.with_extensions` call may supply one. + """ + + logical_extension_codecs: tuple[ + LogicalExtensionCodecExportable | _PyCapsule, ... + ] = () + physical_extension_codecs: tuple[ + PhysicalExtensionCodecExportable | _PyCapsule, ... + ] = () + query_planner: QueryPlannerExportable | _PyCapsule | None = None + + +class SessionExtensionExportable(Protocol): + """Type hint for extension bundles installable via ``with_extensions``. + + Implementations are reusable configuration objects: they must not retain a + :py:class:`SessionContext` and must create fresh components on every call + using the context supplied by :py:meth:`SessionContext.with_extensions`. + They should also avoid mutating global state during binding, since a + failed installation discards the destination context. + """ + + def __datafusion_session_extension__( # noqa: D105 + self, ctx: SessionContext + ) -> SessionExtensionComponents: ... + + class SessionConfig: """Session configuration options.""" @@ -1817,6 +1861,108 @@ 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 + the destination context so it can bind its components to that exact + context, 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 components can end up + bound to an intermediate context that is later garbage collected. + + 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 on the source context is rebound to the + final codec chains. + + If any extension raises or returns invalid components, the source + context's state is left unchanged and the partially built destination + is discarded. Extension factories must treat the context they receive + as configuration-only: catalogs are shared with the source context, so + registering tables or otherwise mutating the context during binding is + not rolled back on failure. + + Args: + extensions: One or more objects implementing + ``__datafusion_session_extension__`` (see + :py:class:`SessionExtensionExportable`). + + 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. Ids + are derived the same way :py:meth:`with_logical_extension_codec` + derives them, so an extension that contributes two instances of + one codec class must declare ``__datafusion_codec_id__`` on at + least one of them. + + Examples: + >>> from my_extension import DistributedEngineExtension # doctest: +SKIP + >>> ctx = SessionContext().with_extensions( + ... DistributedEngineExtension("scheduler:50050") + ... ) # doctest: +SKIP + >>> ctx.sql("SELECT 1").collect() # doctest: +SKIP + """ + if not extensions: + msg = "with_extensions requires at least one extension" + raise ValueError(msg) + for extension in extensions: + if not hasattr(extension, "__datafusion_session_extension__"): + msg = ( + "Extension does not implement __datafusion_session_extension__: " + f"{extension!r}" + ) + raise TypeError(msg) + + # Single destination context. Every component the extensions create + # must bind to this context; _install_extensions later mutates its + # state in place so those bindings stay valid. + destination = SessionContext.__new__(SessionContext) + destination.ctx = self.ctx._derive_for_extensions() + + logical_codecs: list[LogicalExtensionCodecExportable | _PyCapsule] = [] + physical_codecs: list[PhysicalExtensionCodecExportable | _PyCapsule] = [] + planner: QueryPlannerExportable | _PyCapsule | None = None + for extension in extensions: + components = extension.__datafusion_session_extension__(destination) + if not isinstance(components, SessionExtensionComponents): + msg = ( + "__datafusion_session_extension__ must return " + "SessionExtensionComponents, got " + f"{type(components).__name__} from {extension!r}" + ) + raise TypeError(msg) + logical_codecs.extend(components.logical_extension_codecs) + physical_codecs.extend(components.physical_extension_codecs) + 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 = destination.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. diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 3c95835af..f509d8afa 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,121 @@ 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.""" + + 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 destination context'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_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_returned_context(ctx): + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension) + + # The context passed to the factory shares the same underlying session + # as the returned context: registrations made through it are visible. + extension.bound_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]]) From a62e9672d591d58b77a756fd6ec5d4025afe03a6 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:17:45 -0400 Subject: [PATCH 02/12] Add extension-bundle example and with_extensions FFI tests MyPlannerExtension in the query-planner example crate implements the __datafusion_session_extension__ protocol from Rust: it extracts the destination context's task-context provider, binds fresh observing codecs and a planner to it, and returns SessionExtensionComponents. Its codecs record the max_rows config value resolved through the weak provider, letting tests prove the provider targets the returned context rather than the source. Documents with_extensions as the preferred API in the FFI guide. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 1 + docs/source/contributor-guide/ffi.md | 51 +++ .../Cargo.toml | 1 + .../README.md | 17 +- .../_test_three_library_query_planner.py | 168 +++++++++- .../src/extension.rs | 299 ++++++++++++++++++ .../src/lib.rs | 3 + .../src/planner.rs | 26 +- 8 files changed, 552 insertions(+), 14 deletions(-) create mode 100644 examples/datafusion-ffi-query-planner-example/src/extension.rs 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/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index d86858a83..8902947a8 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -343,6 +343,54 @@ 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 two things that are +easy to get wrong: keeping every intermediate context alive, and installing the codecs +before the planner. Every codec and planner capsule carries an +`FFI_TaskContextProvider` holding a *weak* reference to the context it was built +against, so a component bound to a `with_*` result that is then discarded fails at +query time with `TaskContextProvider went out of scope over FFI boundary`. + +`SessionContext.with_extensions` removes both hazards. 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 + # exact context the host will return from with_extensions. + 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 creates one destination context, passes it to every factory, installs all the +codecs, binds the planner against the final codec chains, and returns that context 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. If any factory raises, the source +context is left exactly as it was. + +Bundle objects must be configuration-only: create fresh components on each call, never +cache bound components, and do not retain the context passed in. Catalogs are shared +with the source context, so registrations made during binding are not rolled back on +failure. + +`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 @@ -516,6 +564,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/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..597a25142 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 destination context, binds fresh codec and planner components to that +context's task-context provider, and returns them as `SessionExtensionComponents`. +The host installs everything in one step, so no component can end up bound to an +intermediate context that is later collected. + +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..cb68cfb30 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 @@ -21,7 +21,14 @@ 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 +37,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 +699,159 @@ 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,), + ) + + +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_provider_targets_returned_context(): + """The bundle's task-context provider reads current state from the + returned context, not the source it was derived from.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + planner_ext = MyPlannerExtension() + result = source.with_extensions(ProviderCodecsExtension(), planner_ext) + + # Diverge the two live contexts. Config state is copied at derivation, + # so after these statements source and result disagree. + source.sql("SET ffi_query_planner.max_rows = 5").collect() + result.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 + + # Resolving the provider bound during with_extensions is what a codec's + # decode callback does. Seeing 2 (never 5) proves the provider targets the + # returned context rather than the source. + assert planner_ext.max_rows_through_provider() == 2 + + +def test_with_extensions_survives_dropping_source_and_bundles(): + """Neither the source context nor the bundle objects are needed to keep + the installed components' task-context provider alive.""" + 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_composed_codecs_with_query_planner(): """A second pair of codecs installed on top of the provider codecs composes with them instead of replacing them. 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..3f60cc819 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -0,0 +1,299 @@ +// 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 the destination context 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 destination context supplied by the + // host. Components must not be cached across calls: each installation + // targets a different context. + // + // 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] From e8c0855e571a3278a211468c7bef328cafd99e2f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:19:06 -0400 Subject: [PATCH 03/12] Document and test the context-outlives-DataFrame contract A DataFrame does not keep its SessionContext alive. FFI components hold a weak task-context provider, so operations that reach an FFI codec after the context is collected fail with a clean out-of-scope error rather than crashing. Lock that behavior in with a test and document the ownership contract in the FFI guide and with_extensions docstring. Co-Authored-By: Claude Fable 5 --- docs/source/contributor-guide/ffi.md | 7 +++++++ .../_test_three_library_query_planner.py | 19 +++++++++++++++++++ python/datafusion/context.py | 5 +++++ 3 files changed, 31 insertions(+) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 8902947a8..c3d781126 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -387,6 +387,13 @@ cache bound components, and do not retain the context passed in. Catalogs are sh with the source context, so registrations made during binding are not rolled back on failure. +The returned context is the strong owner of every installed component's task-context +provider, and dependent objects do not extend its lifetime. A `DataFrame`, logical +plan, or capsule can outlive the context, but any operation that reaches an FFI codec +after the context is collected fails with `TaskContextProvider went out of scope over +FFI boundary`. Keep the 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`. 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 cb68cfb30..1fb453124 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 @@ -852,6 +852,25 @@ def __datafusion_session_extension__( assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] +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. diff --git a/python/datafusion/context.py b/python/datafusion/context.py index ded8e8f39..b20e84ae9 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1890,6 +1890,11 @@ def with_extensions( registering tables or otherwise mutating the context during binding is not rolled back on failure. + The returned context is the strong owner of the installed components' + task-context providers. Keep it alive for as long as DataFrames or + plans derived from it are in use; FFI operations after the context is + collected raise an error. + Args: extensions: One or more objects implementing ``__datafusion_session_extension__`` (see From a4435bbedd9b169c8a247874650307fe714d7c96 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:20:13 -0400 Subject: [PATCH 04/12] Skip private internal methods in wrapper coverage test Single-underscore methods on internal pyo3 classes (such as SessionContext._install_extensions) are private support methods for the Python wrappers and do not require a public wrapper. Co-Authored-By: Claude Fable 5 --- python/tests/test_wrapper_coverage.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index cf6719ecf..b1afd6832 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -67,6 +67,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): + # Single-underscore names are private support methods for the + # wrappers (e.g. SessionContext._install_extensions) and are not + # part of the public surface that requires a wrapper. + if internal_attr_name.startswith("_") and not internal_attr_name.startswith( + "__" + ): + continue + wrapped_attr_name = internal_attr_name.removeprefix("Raw") assert wrapped_attr_name in dir(wrapped_obj) From 25fee05bf16451fde5e25441115d4a040bec7b44 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:23:14 -0400 Subject: [PATCH 05/12] Test planner rebinding and codec ids in with_extensions A codec-only bundle installed on a context that already holds an FFI planner must rebind that planner to the new chains, so the planner decodes through the bundle's codecs. Codec ids are derived from the exporting class, so two bundles shipping the same codec class collide and the install is refused. Declaring __datafusion_codec_id__ on the object a bundle hands over resolves it, and both chains then install. Co-Authored-By: Claude Fable 5 Co-Authored-By: Claude Opus 5 (1M context) --- .../_test_three_library_query_planner.py | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) 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 1fb453124..0036cd65c 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 @@ -720,6 +720,43 @@ def __datafusion_session_extension__( ) +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.""" @@ -852,6 +889,66 @@ def __datafusion_session_extension__( 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 From 56ca5014b000f66bee19127ff421cb9ac280b0f8 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 8 Aug 2026 18:40:06 -0400 Subject: [PATCH 06/12] Fix duplicate attribute docs in SessionExtensionComponents The docs build runs Sphinx with --fail-on-warning. SessionExtensionComponents documented its fields in both a napoleon `Attributes:` section and the dataclass class-body annotations, so autoapi emitted each field twice and the build failed with six "duplicate object description" warnings. Move each field's description to a per-field docstring under its annotation so autoapi renders exactly one entry per field. Co-Authored-By: Claude Opus 5 (1M context) --- python/datafusion/context.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index b20e84ae9..8cec49ca4 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -167,23 +167,24 @@ class SessionExtensionComponents: must be created against the context passed to that method; components bound to any other context hold a task-context provider for the wrong session and cannot be rebound. - - Attributes: - logical_extension_codecs: Logical codecs to add to the session's codec - chain, in declaration order. - physical_extension_codecs: Physical codecs to add to the session's - codec chain, in declaration order. - query_planner: Optional query planner. At most one extension per - :py:meth:`SessionContext.with_extensions` call may supply one. """ 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:`SessionContext.with_extensions` call may + supply one. + """ class SessionExtensionExportable(Protocol): From 8c9328adcbebe58ef7ea816ea979d916cf6d6ef3 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sun, 9 Aug 2026 12:26:40 -0400 Subject: [PATCH 07/12] Move session extension types into datafusion.extensions QueryPlannerExportable, SessionExtensionComponents, and SessionExtensionExportable describe how an extension library plugs into a session, not how a SessionContext behaves. Give them their own module so context.py does not keep absorbing the extension surface as it grows. extensions.py imports SessionContext, the codec protocols, and CapsuleType under TYPE_CHECKING only, so context.py can import from it at runtime without a cycle. All three names remain importable from datafusion and datafusion.context; QueryPlannerExportable stays out of the top-level __all__ as before. Co-Authored-By: Claude Opus 5 (1M context) --- python/datafusion/__init__.py | 7 +- python/datafusion/context.py | 62 ++--------------- python/datafusion/extensions.py | 114 ++++++++++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 59 deletions(-) create mode 100644 python/datafusion/extensions.py diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 86d0054b3..4b02a383e 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -80,8 +80,6 @@ RuntimeEnvBuilder, SessionConfig, SessionContext, - SessionExtensionComponents, - SessionExtensionExportable, SQLOptions, ) from .dataframe import ( @@ -94,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 @@ -150,6 +152,7 @@ "common", "configure_formatter", "expr", + "extensions", "functions", "ipc", "lit", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 8cec49ca4..202ff14ba 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -46,7 +46,6 @@ import uuid import warnings -from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol try: @@ -70,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, @@ -146,62 +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 - - -@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:`SessionContext.with_extensions`. Every component - must be created against the context passed to that method; components bound - to any other context hold a task-context provider for the wrong session and - cannot be rebound. - """ - - 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:`SessionContext.with_extensions` call may - supply one. - """ - - -class SessionExtensionExportable(Protocol): - """Type hint for extension bundles installable via ``with_extensions``. - - Implementations are reusable configuration objects: they must not retain a - :py:class:`SessionContext` and must create fresh components on every call - using the context supplied by :py:meth:`SessionContext.with_extensions`. - They should also avoid mutating global state during binding, since a - failed installation discards the destination context. - """ - - def __datafusion_session_extension__( # noqa: D105 - self, ctx: SessionContext - ) -> SessionExtensionComponents: ... - - class SessionConfig: """Session configuration options.""" diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py new file mode 100644 index 000000000..e228a96de --- /dev/null +++ b/python/datafusion/extensions.py @@ -0,0 +1,114 @@ +# 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 destination context so every component binds to +the session that is actually returned. 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 + +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 any other context hold a task-context provider for the + wrong session and cannot be rebound. + """ + + 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. + """ + + +class SessionExtensionExportable(Protocol): + """Type hint for extension bundles installable via ``with_extensions``. + + Implementations are reusable configuration objects: they must not retain a + :py:class:`~datafusion.context.SessionContext` and must create fresh + components on every call using the context supplied by + :py:meth:`~datafusion.context.SessionContext.with_extensions`. They should + also avoid mutating global state during binding, since a failed + installation discards the destination context. + """ + + def __datafusion_session_extension__( # noqa: D105 + self, ctx: SessionContext + ) -> SessionExtensionComponents: ... From ceaf752c02f180537ff4627a9b1a23c9ffb061e3 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 4 Sep 2026 14:05:18 -0400 Subject: [PATCH 08/12] Run the with_extensions docstring example in CI The example was marked `+SKIP` because the main suite has no built FFI extension to import, which is exactly how such an example rots. Parse the statements out of the live docstring in the query-planner example suite, drop the skip, and execute each one against a real extension bundle. Only names are redirected: `my_extension` resolves to a stand-in combining this repository's provider codecs and planner, and `SessionContext` supplies the config that planner reads. A renamed method, a changed signature, or a wrong expected output now fails CI, which already runs this suite. Also drop the `extensions` Args entry's restatement of the type hint and say instead what the hint does not: install order is chain order. Co-Authored-By: Claude Opus 5 (1M context) --- .../_test_three_library_query_planner.py | 83 +++++++++++++++++++ python/datafusion/context.py | 16 +++- 2 files changed, 95 insertions(+), 4 deletions(-) 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 0036cd65c..8dd760bbb 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,7 +17,12 @@ from __future__ import annotations +import doctest import gc +import inspect +import io +import sys +import types import pyarrow as pa import pytest @@ -994,3 +999,81 @@ 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() -> SessionContext: + return SessionContext( + SessionConfig().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/python/datafusion/context.py b/python/datafusion/context.py index 202ff14ba..eaee17f89 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1845,9 +1845,8 @@ def with_extensions( collected raise an error. Args: - extensions: One or more objects implementing - ``__datafusion_session_extension__`` (see - :py:class:`SessionExtensionExportable`). + extensions: Extension bundles to install, in the order their + codecs join the chain. Returns: A new context with all extension components installed. @@ -1864,11 +1863,20 @@ def with_extensions( least one of them. 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 - >>> ctx.sql("SELECT 1").collect() # 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" From bb26fc5aec545935d642ae458f3b2da557cbb87d Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 4 Sep 2026 15:43:15 -0400 Subject: [PATCH 09/12] Share the session in with_extensions instead of forking it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_derive_for_extensions` minted a new `Arc` via `new_with_state(self.ctx.state())`. Every other `with_*` method shares `Arc::clone(&self.ctx)`, and `new_with_state` carries the session id over, so `with_extensions` returned a second live session claiming the same `session_id()` as the source while holding independent `SessionState`. Configuration and the function registry diverged, catalogs stayed shared, and both handles reported the same `__datafusion_codec_id__` — which is `session:` and exists precisely to distinguish codec chains, so installing both on a third session was refused as a duplicate id. The fork also bought nothing. It was introduced to keep components from binding to an intermediate context that could be collected, but there is one `Arc` per session, so no such intermediate exists; deriving one is what creates the hazard. Rule 6 of the ffi-capsule-protocol skill already said to mutate `SessionState` in place rather than derive a replacement. Delete `_derive_for_extensions` and hand the receiver to the extension factories. `_install_extensions` already returned a handle sharing `Arc::clone(&slf.borrow().ctx)`, so removing the fork upstream of it is the whole change. Atomicity is unaffected: both codec chains are built as locals and state is written exactly once, at the end, in `set_session_query_planner`. Replace `test_with_extensions_provider_targets_returned_context`, which is vacuous once the session is shared, with `test_with_extensions_shares_the_session_with_the_source`. It asserts matching session ids and that a `SET` issued through the source after installation is visible to the provider the bundle bound. Reintroducing the fork fails it. Update the prose that described the fork-era design: the `with_extensions` docstring and `SessionExtensionComponents` / `SessionExtensionExportable` in `datafusion.extensions`, the `with_extensions` and "What a derived context shares" sections of the FFI guide, the query planner example's README and `extension.rs` comments, and two test docstrings. Note the shared-session mechanism in Rule 6 of the skill, since `with_extensions` is where it is easiest to get wrong. `enable_url_table` is once again the only method that mints a second `Arc` for a session; its comment, the FFI guide, and the skill now also record that it forks state while keeping the session id, tracked as a bug in #1708. Also add the missing doctest to `SessionExtensionComponents` and a pointer to `with_extensions` from the upgrade guide, which described only the low-level install path. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 13 +++- crates/core/src/context.rs | 56 ++++++-------- docs/source/contributor-guide/ffi.md | 76 +++++++++++-------- docs/source/user-guide/upgrade-guides.md | 7 ++ .../README.md | 8 +- .../_test_three_library_query_planner.py | 40 ++++++---- .../src/extension.rs | 9 ++- python/datafusion/context.py | 65 ++++++++-------- python/datafusion/extensions.py | 40 +++++++--- python/tests/test_context.py | 15 ++-- 10 files changed, 197 insertions(+), 132 deletions(-) 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/crates/core/src/context.rs b/crates/core/src/context.rs index 04427f49c..06d5a20c0 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()) @@ -1609,29 +1614,16 @@ impl PySessionContext { derived } - /// Create the destination context for a `with_extensions` transaction. - /// - /// Private support method for `SessionContext.with_extensions`. The - /// returned context is the single `Arc` that every FFI - /// task-context provider created during the transaction must target; - /// `_install_extensions` later mutates its state in place rather than - /// deriving a new context. - pub fn _derive_for_extensions(&self) -> Self { - Self { - ctx: Arc::new(SessionContext::new_with_state(self.ctx.state())), - logical_codec: Arc::clone(&self.logical_codec), - physical_codec: Arc::clone(&self.physical_codec), - } - } - /// Commit a `with_extensions` transaction onto this context. /// - /// Private support method for `SessionContext.with_extensions`; `self` - /// must be a context produced by `_derive_for_extensions`. Codec capsules - /// are imported and validated before any state change, so a failure - /// leaves the context untouched. The final state is written through this - /// context's own `state_ref()`, never a derived context, so FFI - /// task-context providers bound to it stay valid. + /// 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 any state change, + /// so a failure leaves the session untouched. The final state is written + /// through this context's own `state_ref()`, so those providers stay + /// valid. #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))] pub fn _install_extensions<'py>( slf: &Bound<'py, Self>, @@ -1670,18 +1662,18 @@ impl PySessionContext { .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))) .transpose()?; - let derived = Self { + 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 derived handle so it carries them. Passing `None` still rebuilds + // 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. - derived.set_session_query_planner(planner); + installed.set_session_query_planner(planner); - Ok(derived) + Ok(installed) } } diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index c3d781126..e3754b98c 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -345,21 +345,19 @@ local build commands. ### Extension bundles: `with_extensions` -The chaining above works, but it makes the caller responsible for two things that are -easy to get wrong: keeping every intermediate context alive, and installing the codecs -before the planner. Every codec and planner capsule carries an -`FFI_TaskContextProvider` holding a *weak* reference to the context it was built -against, so a component bound to a `with_*` result that is then discarded fails at -query time with `TaskContextProvider went out of scope over FFI boundary`. +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 both hazards. An extension library exposes a -bundle object implementing `__datafusion_session_extension__`: +`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 - # exact context the host will return from with_extensions. + # session the components will run on. return SessionExtensionComponents( logical_extension_codecs=(self._make_logical_codec(ctx),), physical_extension_codecs=(self._make_physical_codec(ctx),), @@ -367,9 +365,9 @@ class MyEngineExtension: ) ``` -The host creates one destination context, passes it to every factory, installs all the -codecs, binds the planner against the final codec chains, and returns that context in -a single step: +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()) @@ -379,20 +377,25 @@ 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. If any factory raises, the source -context is left exactly as it was. - -Bundle objects must be configuration-only: create fresh components on each call, never -cache bound components, and do not retain the context passed in. Catalogs are shared -with the source context, so registrations made during binding are not rolled back on -failure. - -The returned context is the strong owner of every installed component's task-context -provider, and dependent objects do not extend its lifetime. A `DataFrame`, logical -plan, or capsule can outlive the context, but any operation that reaches an FFI codec -after the context is collected fails with `TaskContextProvider went out of scope over -FFI boundary`. Keep the context alive for as long as objects derived from it are in -use. +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 @@ -474,15 +477,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 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/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 597a25142..7cce6ab60 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -51,10 +51,10 @@ ctx.register_udf(provider_udf) ``` `MyPlannerExtension` implements the `__datafusion_session_extension__` protocol: it -receives the destination context, binds fresh codec and planner components to that -context's task-context provider, and returns them as `SessionExtensionComponents`. -The host installs everything in one step, so no component can end up bound to an -intermediate context that is later collected. +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: 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 8dd760bbb..21432a1ad 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 @@ -788,33 +788,45 @@ def test_with_extensions_three_library_query(): assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 -def test_with_extensions_provider_targets_returned_context(): - """The bundle's task-context provider reads current state from the - returned context, not the source it was derived from.""" +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) - source.register_table("numbers", MyTableProvider(1, 6, 1)) planner_ext = MyPlannerExtension() result = source.with_extensions(ProviderCodecsExtension(), planner_ext) - # Diverge the two live contexts. Config state is copied at derivation, - # so after these statements source and result disagree. - source.sql("SET ffi_query_planner.max_rows = 5").collect() - result.sql("SET ffi_query_planner.max_rows = 2").collect() + 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 - - # Resolving the provider bound during with_extensions is what a codec's - # decode callback does. Seeing 2 (never 5) proves the provider targets the - # returned context rather than the source. 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(): - """Neither the source context nor the bundle objects are needed to keep - the installed components' task-context provider alive.""" + """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() diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs index 3f60cc819..0a7be7453 100644 --- a/examples/datafusion-ffi-query-planner-example/src/extension.rs +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -56,7 +56,7 @@ 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 the destination context alive. +/// here does not keep that session alive. type BoundProvider = Arc>>; fn record_task_ctx(observed: &ObservedMaxRows, ctx: &TaskContext) { @@ -244,9 +244,10 @@ impl MyPlannerExtension { py: Python<'py>, ctx: Bound<'py, PyAny>, ) -> PyResult> { - // Bind every component to the destination context supplied by the - // host. Components must not be cached across calls: each installation - // targets a different context. + // 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 diff --git a/python/datafusion/context.py b/python/datafusion/context.py index eaee17f89..6c5c047c4 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1818,31 +1818,37 @@ def with_extensions( 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 - the destination context so it can bind its components to that exact - context, then all components are installed in one step. This avoids - the pitfalls of chaining :py:meth:`with_logical_extension_codec`, + 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 components can end up - bound to an intermediate context that is later garbage collected. + :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 on the source context is rebound to the - final codec chains. + does, an existing FFI planner is rebound to the final codec chains. - If any extension raises or returns invalid components, the source - context's state is left unchanged and the partially built destination - is discarded. Extension factories must treat the context they receive - as configuration-only: catalogs are shared with the source context, so - registering tables or otherwise mutating the context during binding is - not rolled back on failure. + 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. - The returned context is the strong owner of the installed components' - task-context providers. Keep it alive for as long as DataFrames or - plans derived from it are in use; FFI operations after the context is - collected raise an error. + 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 @@ -1889,17 +1895,16 @@ def with_extensions( ) raise TypeError(msg) - # Single destination context. Every component the extensions create - # must bind to this context; _install_extensions later mutates its - # state in place so those bindings stay valid. - destination = SessionContext.__new__(SessionContext) - destination.ctx = self.ctx._derive_for_extensions() - + # 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. logical_codecs: list[LogicalExtensionCodecExportable | _PyCapsule] = [] physical_codecs: list[PhysicalExtensionCodecExportable | _PyCapsule] = [] planner: QueryPlannerExportable | _PyCapsule | None = None for extension in extensions: - components = extension.__datafusion_session_extension__(destination) + components = extension.__datafusion_session_extension__(self) if not isinstance(components, SessionExtensionComponents): msg = ( "__datafusion_session_extension__ must return " @@ -1920,9 +1925,7 @@ def with_extensions( planner = components.query_planner new = SessionContext.__new__(SessionContext) - new.ctx = destination.ctx._install_extensions( - logical_codecs, physical_codecs, planner - ) + new.ctx = self.ctx._install_extensions(logical_codecs, physical_codecs, planner) return new def table_provider(self, name: str) -> Table: @@ -2354,9 +2357,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 index e228a96de..768c61da0 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -28,9 +28,10 @@ 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 destination context so every component binds to -the session that is actually returned. See the FFI extensions guide in the -contributor documentation for the full rationale. +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 @@ -75,8 +76,26 @@ class SessionExtensionComponents: 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 any other context hold a task-context provider for the - wrong session and cannot be rebound. + components bound to a different session hold a task-context provider for + that other session and cannot be rebound. + + 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[ @@ -101,12 +120,13 @@ class SessionExtensionComponents: class SessionExtensionExportable(Protocol): """Type hint for extension bundles installable via ``with_extensions``. - Implementations are reusable configuration objects: they must not retain a - :py:class:`~datafusion.context.SessionContext` and must create fresh + 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`. They should - also avoid mutating global state during binding, since a failed - installation discards the destination context. + :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. """ def __datafusion_session_extension__( # noqa: D105 diff --git a/python/tests/test_context.py b/python/tests/test_context.py index f509d8afa..86bcb5728 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -901,7 +901,7 @@ def __datafusion_session_extension__(self, ctx): class _PlannerExtension: - """Contributes the destination context's own exported planner.""" + """Contributes the receiving session's own exported planner.""" def __datafusion_session_extension__(self, ctx): return SessionExtensionComponents( @@ -961,13 +961,18 @@ def test_with_extensions_installs_codecs_and_planner(ctx): assert batches[0].column(0) == pa.array([1]) -def test_with_extensions_binds_to_returned_context(ctx): +def test_with_extensions_binds_to_the_receiving_session(ctx): extension = _CodecOnlyExtension() result = ctx.with_extensions(extension) - # The context passed to the factory shares the same underlying session - # as the returned context: registrations made through it are visible. - extension.bound_ctx.register_record_batches( + # 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]})]], ) From 352950803ac267a74ef91d0497e1c1caf1e317da Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 4 Sep 2026 15:59:19 -0400 Subject: [PATCH 10/12] Name a bundle's bare capsules after the bundle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A codec handed to `with_extensions` as a bare `PyCapsule` fell through to `anon:`, an id private to the session that installed it. Plans written through it are undecodable anywhere else, and `with_extensions` accepts no `codec_id=` to override that — so the workaround was to wrap the capsule in an object declaring `__datafusion_codec_id__`, which nothing documented. A distributed engine has to decode its plans in another process, so the shape it would naturally ship — a Rust bundle handing over capsules, as `MyPlannerExtension` does — was the one shape that could not work. The bundle is the stable name that was missing. It is a plain Python object, so its `module.QualName` is library-owned and exactly as stable across processes as an exporting codec class's, which arm 3 of `derive_codec_id` already trusts. The capsule was unnameable only because a capsule carries no type of its own, not because nothing stable was in reach. Resolve a capsule's id through the contributing bundle, using `derive_codec_id` itself so the bundle inherits the same `__datafusion_codec_id__` escape hatch against a class rename. The fallback applies only where randomness would have: an id declared on the handed-over object, or that object's own class, still wins, so an extension can name a codec directly. Two bare capsules of one kind from one bundle collide and are refused. Numbering them by position would be exactly the id `codec.rs` rejects for `anon:` — one another library can mint the same value from — and would break stored plans the first time the bundle reordered what it returns. `resolve_codec_id` gains the bundle argument, `_install_extensions` takes (codec, bundle) pairs, and the collision message now names both routes to a distinct identity; it previously offered only `codec_id=`, which is unreachable from `with_extensions`. Covered in `python/tests/test_context.py`, which reaches every arm without a built extension library: the bundle-derived name, an extension pinning its own id, an id on the handed-over object winning, an exporting object keeping its own, and the two-capsule collision. The cross-FFI case is pinned in the query planner example, where a Rust bundle's capsules must report `datafusion_ffi_query_planner_example.MyPlannerExtension` and no id may be `anon:`. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 65 ++++++++--- docs/source/contributor-guide/ffi.md | 18 ++- .../_test_three_library_query_planner.py | 30 +++++ python/datafusion/context.py | 41 +++++-- python/datafusion/extensions.py | 7 ++ python/tests/test_context.py | 105 ++++++++++++++++++ 6 files changed, 235 insertions(+), 31 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 06d5a20c0..4c0605e74 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1490,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(); @@ -1555,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(); @@ -1627,8 +1627,8 @@ impl PySessionContext { #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))] pub fn _install_extensions<'py>( slf: &Bound<'py, Self>, - logical_codecs: Vec>, - physical_codecs: Vec>, + 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 -- @@ -1642,16 +1642,21 @@ impl PySessionContext { ) }; - for codec in logical_codecs { - let id = resolve_codec_id(&codec, None, &logical_codec.codec_ids())?; + // 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 in physical_codecs { - let id = resolve_codec_id(&codec, None, &physical_codec.codec_ids())?; + 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); @@ -1843,13 +1848,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 @@ -1859,20 +1872,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); } @@ -1882,6 +1903,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 e3754b98c..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. 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 21432a1ad..4ee9b3935 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 @@ -788,6 +788,36 @@ def test_with_extensions_three_library_query(): 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. diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 6c5c047c4..1ea937ac8 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1831,6 +1831,15 @@ def with_extensions( 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 @@ -1862,11 +1871,13 @@ def with_extensions( 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. Ids - are derived the same way :py:meth:`with_logical_extension_codec` - derives them, so an extension that contributes two instances of - one codec class must declare ``__datafusion_codec_id__`` on at - least one of them. + 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 @@ -1900,8 +1911,16 @@ def with_extensions( # bound here holds a task-context provider that the returned handle # keeps alive, and `_install_extensions` writes the final state through # that same session. - logical_codecs: list[LogicalExtensionCodecExportable | _PyCapsule] = [] - physical_codecs: list[PhysicalExtensionCodecExportable | _PyCapsule] = [] + # + # 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) @@ -1912,8 +1931,12 @@ def with_extensions( f"{type(components).__name__} from {extension!r}" ) raise TypeError(msg) - logical_codecs.extend(components.logical_extension_codecs) - physical_codecs.extend(components.physical_extension_codecs) + 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 = ( diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 768c61da0..c12e631ca 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -79,6 +79,13 @@ class SessionExtensionComponents: 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: diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 86bcb5728..ff53c1776 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -946,6 +946,111 @@ def __datafusion_session_extension__(self, ctx): 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", From d8f32cca09747ea4d4d1eac3e6d244d3d0d2f6b6 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 4 Sep 2026 16:03:29 -0400 Subject: [PATCH 11/12] Stop claiming _install_extensions always writes state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The doc comment said "the final state is written through this context's own `state_ref()`", which overstates it. `set_session_query_planner` returns early when there is no planner to bind, and the codec chains live on the returned `PySessionContext` fields rather than in `SessionState` — so a codec-only install onto a session with no FFI planner writes nothing. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 4c0605e74..4fb40ae3d 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1620,9 +1620,15 @@ impl PySessionContext { /// 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 any state change, - /// so a failure leaves the session untouched. The final state is written - /// through this context's own `state_ref()`, so those providers stay + /// 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>( From e06662acb7007c5493dea5f457fdef26108ecd50 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 4 Sep 2026 16:07:49 -0400 Subject: [PATCH 12/12] Tidy the loose ends from review of with_extensions Mark `SessionExtensionExportable` `@runtime_checkable` and have `with_extensions` check it with `isinstance` rather than `hasattr`, so the annotation and the runtime check are the same statement, and callers can ask the question too. Covered by a doctest on the protocol. Replace the leading-underscore skip in `test_wrapper_coverage` with a named allowlist. The pattern also excused `DataFrame._repr_html_`, which a wrapper does have to provide, so a two-method need was weakening coverage for every private name. Removing `_install_extensions` from the allowlist fails the test, so the entry is load-bearing rather than decorative. Say in `_CodecOnlyExtension` that retaining the context is what the protocol tells real extensions not to do, and that it is kept only so a test can assert which context the factory was handed. Let the docstring-example shim in the query planner example accept a config positionally, the way the real constructor does. Editing the docstring to `SessionContext(config)` now fails as a doctest diff rather than as a `TypeError` inside the harness. Co-Authored-By: Claude Opus 5 (1M context) --- .../_test_three_library_query_planner.py | 10 ++++---- python/datafusion/context.py | 2 +- python/datafusion/extensions.py | 21 ++++++++++++++++- python/tests/test_context.py | 8 ++++++- python/tests/test_wrapper_coverage.py | 23 ++++++++++++++----- 5 files changed, 51 insertions(+), 13 deletions(-) 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 4ee9b3935..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 @@ -1099,10 +1099,12 @@ def test_with_extensions_docstring_example_still_runs(): module = types.ModuleType("my_extension") module.DistributedEngineExtension = _DocstringExampleExtension - def make_context() -> SessionContext: - return SessionContext( - SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) - ) + 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, diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 1ea937ac8..ac453bb72 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1899,7 +1899,7 @@ def with_extensions( msg = "with_extensions requires at least one extension" raise ValueError(msg) for extension in extensions: - if not hasattr(extension, "__datafusion_session_extension__"): + if not isinstance(extension, SessionExtensionExportable): msg = ( "Extension does not implement __datafusion_session_extension__: " f"{extension!r}" diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index c12e631ca..8e35c720d 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -37,7 +37,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: from _typeshed import CapsuleType as _PyCapsule @@ -124,9 +124,15 @@ class SessionExtensionComponents: """ +@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 @@ -134,6 +140,19 @@ class SessionExtensionExportable(Protocol): 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 diff --git a/python/tests/test_context.py b/python/tests/test_context.py index ff53c1776..0418acde2 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -882,7 +882,13 @@ def test_contexts_sharing_a_session_share_the_planner(ctx): class _CodecOnlyExtension: - """Contributes decline-all codecs exported from an unrelated session.""" + """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() diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index b1afd6832..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,12 +78,12 @@ def missing_exports(internal_obj, wrapped_obj) -> None: pytest.fail(f"Missing __repr__: {internal_obj.__name__}") for internal_attr_name in dir(internal_obj): - # Single-underscore names are private support methods for the - # wrappers (e.g. SessionContext._install_extensions) and are not - # part of the public surface that requires a wrapper. - if internal_attr_name.startswith("_") and not internal_attr_name.startswith( - "__" - ): + # 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")