From 4facb0c4de609c9de74701554d9e2be2716ce278 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:09:16 -0400 Subject: [PATCH 01/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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") From 3924704701a93fdfe0a6c2daad35b50dec828fe3 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 4 Sep 2026 16:37:35 -0400 Subject: [PATCH 13/33] Require with_extensions codecs to be objects, not capsules Reverses "Name a bundle's bare capsules after the bundle". Deriving a capsule's id from the bundle that contributed it reads the identity off the wrong object: the bundle is whatever the caller passed to with_extensions, so an application that packages several libraries as one bundle of its own stamps its identity onto the inner libraries' codecs. Their pinned __datafusion_codec_id__ is discarded and there is nothing the inner library can do about it, since its object never reaches _install_extensions. Nothing fails at install time; the mismatch surfaces as an undecodable plan in the process that reads it, naming an id nobody wrote in source. So with_extensions now refuses a bare capsule and names the getter to implement. An id read off the handed-over object is composition-stable by construction, which the new tests pin at both layers. This also decouples a codec's wire identity from the bundle's Python class name, which is what __datafusion_codec_id__ exists for, and closes the case where a bundle built by a factory function contributed a wire id containing "". The low-level methods keep accepting capsules: they take codec_id=, so the random anon: arm still has an escape hatch. Query planners are unaffected, carrying no wire id. MyPlannerExtension gains BundledLogicalCodec and BundledPhysicalCodec, small pyclasses holding the bound FFI codec and declaring pinned ids, as the reference shape for a library whose plans leave the process. Also, unrelated to the above but adjacent in the docs: with_extensions never said that a bundle-supplied planner replaces an installed one rather than layering, and the SessionExtensionComponents example that showed a codec was fully skipped with undefined names. Both fixed. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 13 ++ crates/core/src/context.rs | 112 +++++++----- docs/source/contributor-guide/ffi.md | 63 +++++-- .../README.md | 6 + .../_test_three_library_query_planner.py | 75 ++++++-- .../src/extension.rs | 95 +++++++++- .../src/lib.rs | 4 +- python/datafusion/context.py | 58 +++---- python/datafusion/extensions.py | 45 +++-- python/tests/test_context.py | 164 +++++++++++++----- 10 files changed, 468 insertions(+), 167 deletions(-) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 3e39f3d36..826763ce1 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -62,6 +62,19 @@ library reaches things only the session has. session satisfies the protocol too — `ctx.__datafusion_query_planner__()` and `ctx.__datafusion_query_planner__(ctx)` are both valid. +A *codec* must always be handed over as an object implementing its getter, never +as the bare capsule the getter returns; `with_extensions` refuses a capsule. +A codec's wire id — the string a payload names on decode, which has to mean the +same thing in whichever process decodes — is read off the object it arrives as, +and a capsule has no type to read one from. Deriving the id from the bundle that +contributed the capsule is not the fix: the bundle is whatever object the caller +passed, so an application packaging your library inside a bundle of its own +would re-tag your payloads and they would stop decoding where they are read. If +the object's class name is not the identity you want on the wire, declare +`__datafusion_codec_id__` on it. `BundledLogicalCodec` in +`examples/datafusion-ffi-query-planner-example/src/extension.rs` is the shape. +This applies only to codecs — a query planner carries no wire id. + ## Rule 3 — never construct a `SessionContext` in an extension library The FFI constructors ask for things a library does not have: diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 4fb40ae3d..8ade72397 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -66,7 +66,7 @@ use datafusion_python_util::{ }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; -use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; +use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; use url::Url; @@ -1490,7 +1490,7 @@ impl PySessionContext { ) -> PyDataFusionResult { let id = { let this = slf.borrow(); - resolve_codec_id(&codec, codec_id, None, &this.logical_codec.codec_ids())? + resolve_codec_id(&codec, codec_id, &this.logical_codec.codec_ids())? }; let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); @@ -1555,7 +1555,7 @@ impl PySessionContext { ) -> PyDataFusionResult { let id = { let this = slf.borrow(); - resolve_codec_id(&codec, codec_id, None, &this.physical_codec.codec_ids())? + resolve_codec_id(&codec, codec_id, &this.physical_codec.codec_ids())? }; let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); @@ -1630,11 +1630,15 @@ impl PySessionContext { /// 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. + /// + /// Codecs must arrive as objects exposing the capsule getter, never as + /// bare capsules — see [`resolve_bundle_codec_id`]. The planner has no + /// wire id, so it may still be a capsule. #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))] pub fn _install_extensions<'py>( slf: &Bound<'py, Self>, - logical_codecs: Vec<(Bound<'py, PyAny>, Bound<'py, PyAny>)>, - physical_codecs: Vec<(Bound<'py, PyAny>, Bound<'py, PyAny>)>, + logical_codecs: Vec>, + physical_codecs: Vec>, planner: Option>, ) -> PyDataFusionResult { // Chains are built as local values, so a codec that fails to import -- @@ -1648,21 +1652,24 @@ impl PySessionContext { ) }; - // 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())?; + for codec in logical_codecs { + let id = resolve_bundle_codec_id( + &codec, + "__datafusion_logical_extension_codec__", + &logical_codec.codec_ids(), + )?; let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); logical_codec = logical_codec.with_additional_codec(id, inner); } let logical_codec = Arc::new(logical_codec); - for (codec, bundle) in physical_codecs { - let id = resolve_codec_id(&codec, None, Some(&bundle), &physical_codec.codec_ids())?; + for codec in physical_codecs { + let id = resolve_bundle_codec_id( + &codec, + "__datafusion_physical_extension_codec__", + &physical_codec.codec_ids(), + )?; let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); physical_codec = physical_codec.with_additional_codec(id, inner); @@ -1854,21 +1861,17 @@ 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` 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. +/// 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. +/// Reachable only from `with_logical_extension_codec` and +/// `with_physical_extension_codec`, where `codec_id=` is the way out; +/// `with_extensions` refuses bare capsules outright rather than naming them +/// after something that is not the codec. See [`resolve_bundle_codec_id`]. /// /// An id already in use is rejected rather than shadowed. Two codecs sharing an /// id are indistinguishable on decode, and the API cannot tell whether two @@ -1878,10 +1881,9 @@ 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, bundle)?; + let id = derive_codec_id(codec, explicit)?; if existing.contains(&id.as_str()) { return Err(PyValueError::new_err(format!( "An extension codec with id '{id}' is already installed on this session. Two \ @@ -1895,11 +1897,45 @@ fn resolve_codec_id( Ok(id) } -fn derive_codec_id( +/// Resolve the wire id for a codec contributed through `with_extensions`, +/// requiring an object that can name itself. +/// +/// `with_extensions` takes no `codec_id=`, so the only naming channels are the +/// ones [`derive_codec_id`] reads off the handed-over object: a declared +/// `__datafusion_codec_id__`, or its class's `module.QualName`. A bare capsule +/// has neither. Naming it after the bundle that contributed it looks like an +/// answer and is not one: the bundle is whatever object the caller passed to +/// `with_extensions`, so a bundle that wraps another library's bundle — the +/// natural way for an application to package several libraries as one — would +/// stamp its own identity onto the inner library's codecs and silently change +/// the wire format. The inner library cannot defend against that no matter what +/// it declares, and the mismatch does not surface until a plan fails to decode +/// in another process. +/// +/// So the capsule is refused here, where the author can fix it by wrapping it +/// in an object. Wrapping also decouples the codec's wire identity from the +/// bundle's Python class name, which is the whole point of +/// `__datafusion_codec_id__`. +fn resolve_bundle_codec_id( codec: &Bound<'_, PyAny>, - explicit: Option, - bundle: Option<&Bound<'_, PyAny>>, + getter: &str, + existing: &[&str], ) -> PyResult { + if codec.is_instance_of::() { + return Err(PyTypeError::new_err(format!( + "A codec contributed through `with_extensions` must be an object exposing \ + `{getter}`, not a bare PyCapsule. A capsule carries no type of its own, so \ + there is nothing to name the codec by, and a payload names its codec by id \ + when it is decoded — an id that has to mean the same thing in whichever \ + process decodes. Wrap the capsule in an object that exposes `{getter}` and, \ + if the class name is not the identity you want on the wire, declares \ + `__datafusion_codec_id__`." + ))); + } + resolve_codec_id(codec, None, existing) +} + +fn derive_codec_id(codec: &Bound<'_, PyAny>, explicit: Option) -> PyResult { if let Some(id) = explicit { return Ok(id); } @@ -1909,14 +1945,6 @@ fn derive_codec_id( 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 b9f2f541c..3ed8ccfd4 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -277,20 +277,15 @@ three cases: - **Two instances of one class.** Both get the same id, so the second install raises `ValueError`. Pass `codec_id=` to tell them apart. -- **A bare `PyCapsule`.** A capsule has no class to take a name from. 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 bare `PyCapsule`.** A capsule has no class to take a name from, so installing + one through `with_logical_extension_codec` or `with_physical_extension_codec` gives + it an id private to the session that installed it; plans it encodes fail with a + clear error on any other session rather than being decoded by the wrong codec. Pass + `codec_id=` if those plans have to cross sessions. + + `with_extensions` takes no `codec_id=`, so it refuses a bare capsule outright and + tells you to wrap it. See + [Extension bundles: `with_extensions`](#extension-bundles-with_extensions). - **A class you intend to rename.** The id follows the class name, so renaming stops older plans from decoding. Declare `__datafusion_codec_id__` on the exporting object to pin an id that survives the rename. @@ -387,7 +382,42 @@ 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. +extension per call may supply a query planner. Supplying one replaces whatever planner +the session already has; to layer instead, capture the existing planner from +`__datafusion_query_planner__` first and have yours fall back to it. + +#### Codecs are objects, not capsules + +`with_extensions` requires each codec to be an object exposing the capsule getter, and +refuses a bare `PyCapsule`. A codec's id is read off the object it is handed over as, +and a capsule has no type to read one from; since this method takes no `codec_id=`, +there would be nothing left to name it by. A library holding a raw capsule — which is +what a Rust implementation has — wraps it: + +```python +class MyLogicalCodec: + # Optional. Without it the id is this class's import path, which is already + # stable; declare it if you may rename the class and need old plans to decode. + __datafusion_codec_id__ = "my_library.logical.v1" + + def __init__(self, capsule): + self._capsule = capsule + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule +``` + +Wrapping is not just bookkeeping. It ties the id to the codec rather than to the bundle +that contributed it, and that difference is load-bearing: an application commonly +presents several libraries as one bundle of its own, and the id has to survive that. +Were the id taken from the contributing bundle, wrapping `my_engine.Extension` inside +`my_app.Extension` would silently re-tag the engine's payloads, and a scheduler that +installs the engine's codec by its documented id would fail to decode plans from +composed clients while succeeding for direct ones. The wrapper travels with the codec; +the bundle does not. + +The query planner is exempt — it carries no wire id, so it may be an object or a +capsule. Nothing is written to the session until every factory has returned and every capsule has been validated, so a factory that raises leaves the session exactly as it was. A @@ -409,7 +439,8 @@ boundary`. Keep a context alive for as long as objects derived from it are in us `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`. +supplied context, wrapping its codecs in `BundledLogicalCodec` / `BundledPhysicalCodec` +so they carry declared ids, and constructing a Python `SessionExtensionComponents`. ### Capsule getters receive the session they are installed on diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 7cce6ab60..efcb3376f 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -56,6 +56,12 @@ 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. +Its codecs are handed over as `BundledLogicalCodec` and `BundledPhysicalCodec` rather +than as bare capsules. `with_extensions` requires an object, because a codec's wire id +is read off the object it arrives as and a capsule has no type to read one from. Each +wrapper declares `__datafusion_codec_id__`, so the id belongs to this library and does +not change when the bundle is nested inside an application's own bundle. + The integration tests also cover the low-level chaining setup: ```python 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 ed9e6184a..d272bd998 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 @@ -704,6 +704,14 @@ def test_query_planner_rejects_invalid_config(max_rows: str): ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() +# Ids `MyPlannerExtension`'s codec wrappers declare, mirroring +# `LOGICAL_CODEC_ID` / `PHYSICAL_CODEC_ID` in the crate's `extension.rs`. A +# scheduler decoding this library's plans installs codecs under these names, so +# they are part of its wire format rather than an implementation detail. +LOGICAL_CODEC_ID = "datafusion_ffi_query_planner_example.logical.v1" +PHYSICAL_CODEC_ID = "datafusion_ffi_query_planner_example.physical.v1" + + class ProviderCodecsExtension: """Bundles the provider library's codecs for ``with_extensions``. @@ -788,27 +796,25 @@ 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. +def test_with_extensions_names_a_rust_bundles_codecs_by_their_declared_id(): + """A Rust bundle wraps each codec in an object that declares its own id. This is the identity that has to survive leaving the process: a plan a distributed engine writes here is decoded by its scheduler, which installs a codec under the same id. A session-private random id — what a bare - capsule gets when installed directly — would make the plan undecodable - there. + capsule would get if `with_extensions` accepted one — would make the plan + undecodable there, which is why bare capsules are refused. """ config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) ctx = SessionContext(config).with_extensions( ProviderCodecsExtension(), MyPlannerExtension() ) - 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() + assert LOGICAL_CODEC_ID in ctx.logical_extension_codec_ids() + assert PHYSICAL_CODEC_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. + # The provider bundle's codecs declare no id, so they fall back to their + # own class names — never to the bundle's. assert ( "datafusion_ffi_example.MyLogicalExtensionCodec" in ctx.logical_extension_codec_ids() @@ -818,6 +824,55 @@ def test_with_extensions_names_a_rust_bundles_capsules_after_the_bundle(): ) +def test_with_extensions_rejects_a_rust_bundles_bare_capsule(): + """A bundle handing over a raw capsule is refused, with the fix named. + + This is the shape a Rust library reaches for first — `MyPlannerExtension` + wraps its capsules precisely to avoid it. + """ + + class BareCapsuleExtension: + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=( + ctx.__datafusion_logical_extension_codec__(), + ), + ) + + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + with pytest.raises(TypeError, match="must be an object exposing"): + SessionContext(config).with_extensions(BareCapsuleExtension()) + + +def test_with_extensions_codec_ids_survive_bundle_composition(): + """Nesting a bundle inside another does not re-tag its codecs. + + An application that presents several libraries as one bundle is the + natural shape, and it must not change what the inner libraries write on + the wire — a scheduler installing `MyPlannerExtension`'s codec by id has + no idea which application wrapper the client used. Reading the id off the + handed-over object rather than off the contributing bundle is what makes + that hold. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + direct = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + composed = SessionContext(config).with_extensions( + _DocstringExampleExtension("scheduler:50050") + ) + + assert composed.logical_extension_codec_ids() == ( + direct.logical_extension_codec_ids() + ) + assert composed.physical_extension_codec_ids() == ( + direct.physical_extension_codec_ids() + ) + assert LOGICAL_CODEC_ID in composed.logical_extension_codec_ids() + + def test_with_extensions_shares_the_session_with_the_source(): """``with_extensions`` returns a handle on the source's session, and the bundle's task-context provider resolves against that one session. diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs index 0a7be7453..8307b5761 100644 --- a/examples/datafusion-ffi-query-planner-example/src/extension.rs +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -40,7 +40,7 @@ use datafusion_python_util::{ }; use datafusion_session::QueryPlanner; use pyo3::prelude::*; -use pyo3::types::PyDict; +use pyo3::types::{PyCapsule, PyDict}; use crate::planner::{DistributedQueryPlanner, PlannerObservations, planner_config_from_options}; @@ -156,6 +156,84 @@ impl PhysicalExtensionCodec for ObservingPhysicalExtensionCodec { } } +/// Wire id this library's logical codec claims, pinned so that renaming the +/// Rust or Python types does not invalidate plans already encoded. +const LOGICAL_CODEC_ID: &str = "datafusion_ffi_query_planner_example.logical.v1"; + +/// Physical companion to [`LOGICAL_CODEC_ID`]. +const PHYSICAL_CODEC_ID: &str = "datafusion_ffi_query_planner_example.physical.v1"; + +/// Carries this bundle's logical codec as an object rather than a bare capsule. +/// +/// `with_extensions` requires an object: a codec's wire id is read off the +/// thing it is handed over as, and a capsule has no type to read one from. +/// Wrapping is also what keeps the id *this library's*. An id derived from the +/// contributing bundle would follow whichever object the caller passed to +/// `with_extensions`, so an application that packages this library inside a +/// bundle of its own would silently re-tag these payloads and they would stop +/// decoding in the process that reads them. The wrapper travels with the codec; +/// the bundle does not. +/// +/// Declaring `__datafusion_codec_id__` is optional — the class's +/// `module.QualName` would serve — but a library whose plans leave the process +/// should pin the id rather than let a refactor move it. +#[pyclass( + name = "BundledLogicalCodec", + module = "datafusion_ffi_query_planner_example" +)] +pub(crate) struct BundledLogicalCodec { + codec: FFI_LogicalExtensionCodec, +} + +#[pymethods] +impl BundledLogicalCodec { + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + LOGICAL_CODEC_ID + } + + /// `session` is unused: the codec was bound to its task-context provider + /// when the bundle was installed, which is the whole reason the bundle + /// receives the context. + #[pyo3(signature = (session=None))] + fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + create_logical_extension_capsule(py, &self.codec) + } +} + +/// Physical companion to [`BundledLogicalCodec`]. +#[pyclass( + name = "BundledPhysicalCodec", + module = "datafusion_ffi_query_planner_example" +)] +pub(crate) struct BundledPhysicalCodec { + codec: FFI_PhysicalExtensionCodec, +} + +#[pymethods] +impl BundledPhysicalCodec { + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + PHYSICAL_CODEC_ID + } + + /// See [`BundledLogicalCodec::__datafusion_logical_extension_codec__`]. + #[pyo3(signature = (session=None))] + fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + create_physical_extension_capsule(py, &self.codec) + } +} + /// Extension bundle for `SessionContext.with_extensions`. /// /// Mirrors how a distributed engine such as Ballista packages its session @@ -264,7 +342,9 @@ impl MyPlannerExtension { }); let ffi_logical = FFI_LogicalExtensionCodec::new(logical, Some(runtime.clone()), provider.clone()); - let logical_capsule = create_logical_extension_capsule(py, &ffi_logical)?; + // Handed over as an object, not a capsule, so the codec carries an id + // of its own. See `BundledLogicalCodec`. + let logical_codec = Py::new(py, BundledLogicalCodec { codec: ffi_logical })?; let physical: Arc = Arc::new(ObservingPhysicalExtensionCodec { @@ -273,7 +353,12 @@ impl MyPlannerExtension { }); let ffi_physical = FFI_PhysicalExtensionCodec::new(physical, Some(runtime), provider.clone()); - let physical_capsule = create_physical_extension_capsule(py, &ffi_physical)?; + let physical_codec = Py::new( + py, + BundledPhysicalCodec { + codec: ffi_physical, + }, + )?; let planner: Arc = Arc::new(DistributedQueryPlanner { observations: Arc::clone(&self.observations), @@ -292,8 +377,8 @@ impl MyPlannerExtension { .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("logical_extension_codecs", (logical_codec,))?; + kwargs.set_item("physical_extension_codecs", (physical_codec,))?; 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 70d4a42c5..50cc4831e 100644 --- a/examples/datafusion-ffi-query-planner-example/src/lib.rs +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -18,7 +18,7 @@ use pyo3::prelude::*; use crate::config::MyPlannerConfig; -use crate::extension::MyPlannerExtension; +use crate::extension::{BundledLogicalCodec, BundledPhysicalCodec, MyPlannerExtension}; use crate::planner::MyQueryPlanner; mod config; @@ -28,6 +28,8 @@ mod planner; #[pymodule] fn datafusion_ffi_query_planner_example(m: &Bound<'_, PyModule>) -> PyResult<()> { pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/datafusion/context.py b/python/datafusion/context.py index ac453bb72..5d6037cef 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1831,14 +1831,18 @@ 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. + Codecs must be handed over as objects exposing the capsule getter, not + as bare ``PyCapsule`` objects, and are named after their exporting + class as :py:meth:`with_logical_extension_codec` describes. Declare + ``__datafusion_codec_id__`` on the object to pin an id that survives a + later class rename. A capsule carries no type of its own, so there + would be nothing to name the codec by, and this method takes no + ``codec_id=``; wrap it in an object instead. That also keeps a codec's + wire identity independent of the extension that ships it, so an + extension composed inside another one still writes the same ids. + + The planner is exempt — it carries no wire id, so it may be an object + or a capsule. Like the individual ``with_*`` methods, the returned context shares its session with this one: catalogs, tables, registered functions, and @@ -1867,17 +1871,17 @@ def with_extensions( 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`. + TypeError: If an argument does not implement the protocol, returns + something other than a + :py:class:`SessionExtensionComponents`, or contributes a codec + as a bare ``PyCapsule``. ValueError: If no extensions are given, more than one extension supplies a query planner, or two codecs claim the same id. An - extension that contributes two instances of one codec class, - or two bare capsules of the same kind, must declare - ``__datafusion_codec_id__`` on at least one of them; the - collision is refused rather than resolved by position, because - a positional id would break stored plans the first time the - extension reordered what it returns. + extension that contributes two instances of one codec class + must declare ``__datafusion_codec_id__`` on at least one of + them; the collision is refused rather than resolved by + position, because a positional id would break stored plans the + first time the extension reordered what it returns. Examples: The example is skipped here because it needs a built FFI @@ -1911,16 +1915,8 @@ 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. - # - # 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] - ] = [] + logical_codecs: list[LogicalExtensionCodecExportable] = [] + physical_codecs: list[PhysicalExtensionCodecExportable] = [] planner: QueryPlannerExportable | _PyCapsule | None = None for extension in extensions: components = extension.__datafusion_session_extension__(self) @@ -1931,12 +1927,8 @@ def with_extensions( f"{type(components).__name__} from {extension!r}" ) raise TypeError(msg) - logical_codecs.extend( - (codec, extension) for codec in components.logical_extension_codecs - ) - physical_codecs.extend( - (codec, extension) for codec in components.physical_extension_codecs - ) + 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 = ( diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 8e35c720d..6c95b9003 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -79,12 +79,12 @@ 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__``. + Codecs must be objects exposing the capsule getters, never bare + ``PyCapsule`` objects: a codec's id is read off the object it is handed + over as, and a capsule has no type to read. A library holding a raw capsule + wraps it in an object, which is also what gives the codec an identity of + its own — one that does not change when the codec is contributed through a + different extension. Examples: A bundle that contributes nothing is valid, and is what the defaults @@ -97,22 +97,33 @@ class SessionExtensionComponents: >>> components.query_planner is None True - A bundle that contributes one kind of component names it, leaving - the rest empty: - + A bundle that contributes one kind of component names it, leaving the + rest empty. Here the codec is a capsule wrapped in an object that + declares the id its payloads will carry: + + >>> from datafusion import SessionContext + >>> class NamedCodec: + ... __datafusion_codec_id__ = "my_library.v1" + ... + ... def __init__(self, capsule): + ... self._capsule = capsule + ... + ... def __datafusion_logical_extension_codec__(self, session=None): + ... return self._capsule + >>> capsule = SessionContext().__datafusion_logical_extension_codec__() >>> components = SessionExtensionComponents( - ... query_planner=my_library.make_planner(ctx) - ... ) # doctest: +SKIP + ... logical_extension_codecs=(NamedCodec(capsule),) + ... ) + >>> components.logical_extension_codecs[0].__datafusion_codec_id__ + 'my_library.v1' + >>> components.physical_extension_codecs + () """ - logical_extension_codecs: tuple[ - LogicalExtensionCodecExportable | _PyCapsule, ... - ] = () + logical_extension_codecs: tuple[LogicalExtensionCodecExportable, ...] = () """Logical codecs to add to the session's codec chain, in declaration order.""" - physical_extension_codecs: tuple[ - PhysicalExtensionCodecExportable | _PyCapsule, ... - ] = () + physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = () """Physical codecs to add to the session's codec chain, in declaration order.""" query_planner: QueryPlannerExportable | _PyCapsule | None = None diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 0418acde2..d41173102 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -881,6 +881,25 @@ def test_contexts_sharing_a_session_share_the_planner(ctx): assert sibling.session_id() == ctx.session_id() +class _NamedCodec: + """Wraps a codec capsule in an object that can name itself. + + ``with_extensions`` requires objects rather than bare capsules, because a + codec's wire id is read off the object it is handed over as. This is the + shape a library holding a raw capsule hands over. + """ + + def __init__(self, capsule, codec_id): + self._capsule = capsule + self.__datafusion_codec_id__ = codec_id + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule + + def __datafusion_physical_extension_codec__(self, session=None): + return self._capsule + + class _CodecOnlyExtension: """Contributes decline-all codecs exported from an unrelated session. @@ -890,18 +909,25 @@ class _CodecOnlyExtension: context the factory was handed. """ - def __init__(self): + def __init__(self, prefix="my_library"): self.exporter = SessionContext() + self.prefix = prefix self.bound_ctx = None def __datafusion_session_extension__(self, ctx): self.bound_ctx = ctx return SessionExtensionComponents( logical_extension_codecs=( - self.exporter.__datafusion_logical_extension_codec__(), + _NamedCodec( + self.exporter.__datafusion_logical_extension_codec__(), + f"{self.prefix}.logical", + ), ), physical_extension_codecs=( - self.exporter.__datafusion_physical_extension_codec__(), + _NamedCodec( + self.exporter.__datafusion_physical_extension_codec__(), + f"{self.prefix}.physical", + ), ), ) @@ -940,10 +966,15 @@ def test_with_extensions_rejects_multiple_planners(ctx): def test_with_extensions_rejects_bad_codec_capsule(ctx): + """A correctly shaped object still has to return the right capsule.""" + class BadCodecExtension: def __datafusion_session_extension__(self, ctx): + wrong_capsule = ctx.__datafusion_task_context_provider__() return SessionExtensionComponents( - logical_extension_codecs=(ctx.__datafusion_task_context_provider__(),), + logical_extension_codecs=( + _NamedCodec(wrong_capsule, "my_library.logical"), + ), ) with pytest.raises( @@ -952,45 +983,88 @@ 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. +def test_with_extensions_rejects_a_bare_capsule_codec(ctx): + """A codec must be an object that can name itself, not a bare capsule. - 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. + An id is read off the object a codec is handed over as, and a capsule has + no type to read one from. ``with_extensions`` takes no ``codec_id=``, so + the capsule is refused here rather than given an id derived from something + that is not the codec. """ - result = ctx.with_extensions(_CodecOnlyExtension()) - expected = f"{_CodecOnlyExtension.__module__}._CodecOnlyExtension" - assert result.logical_extension_codec_ids() == [expected] - assert result.physical_extension_codec_ids() == [expected] + class BareCapsuleExtension: + def __init__(self): + self.exporter = SessionContext() + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=( + self.exporter.__datafusion_logical_extension_codec__(), + ), + ) -def test_with_extensions_extension_can_pin_its_codec_id(ctx): - """``__datafusion_codec_id__`` on the extension survives a class rename.""" + with pytest.raises( + TypeError, + match="must be an object exposing `__datafusion_logical_extension_codec__`", + ): + ctx.with_extensions(BareCapsuleExtension()) + + +def test_with_extensions_rejects_a_bare_physical_capsule_codec(ctx): + """The physical getter is named in its own diagnostic.""" + + class BareCapsuleExtension: + def __init__(self): + self.exporter = SessionContext() - class PinnedExtension(_CodecOnlyExtension): - __datafusion_codec_id__ = "my_library.v1" + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + physical_extension_codecs=( + self.exporter.__datafusion_physical_extension_codec__(), + ), + ) - result = ctx.with_extensions(PinnedExtension()) - assert result.logical_extension_codec_ids() == ["my_library.v1"] + with pytest.raises( + TypeError, + match="must be an object exposing `__datafusion_physical_extension_codec__`", + ): + ctx.with_extensions(BareCapsuleExtension()) -def test_with_extensions_codec_id_on_the_codec_beats_the_extension(ctx): - """Naming the handed-over object wins over the extension's name. +def test_with_extensions_codec_ids_survive_composition(ctx): + """A codec keeps its id when its extension is nested inside another one. - This is how an extension contributing more than one bare capsule of a kind - tells them apart. + Wire ids have to mean the same thing in whichever process decodes, so + packaging one extension inside another — the natural way for an + application to present several libraries as one — must not re-tag the + inner library's payloads. Reading the id off the handed-over object rather + than off the contributing extension is what guarantees that. """ - class NamedCapsule: - def __init__(self, capsule, codec_id): - self._capsule = capsule - self.__datafusion_codec_id__ = codec_id + class ComposedExtension: + """Presents another extension's components as its own.""" - def __datafusion_logical_extension_codec__(self, session=None): - return self._capsule + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_extension__(self, ctx): + return self.inner.__datafusion_session_extension__(ctx) + + direct = ctx.with_extensions(_CodecOnlyExtension()) + wrapped = SessionContext().with_extensions(ComposedExtension(_CodecOnlyExtension())) + + assert direct.logical_extension_codec_ids() == ["my_library.logical"] + assert wrapped.logical_extension_codec_ids() == ["my_library.logical"] + assert direct.physical_extension_codec_ids() == ["my_library.physical"] + assert wrapped.physical_extension_codec_ids() == ["my_library.physical"] + + +def test_with_extensions_uses_ids_declared_on_the_codec(ctx): + """``__datafusion_codec_id__`` on the handed-over object names the codec. + + This is how an extension contributing more than one codec of a kind tells + them apart. + """ class TwoNamedCodecs: def __init__(self): @@ -999,11 +1073,11 @@ def __init__(self): def __datafusion_session_extension__(self, ctx): return SessionExtensionComponents( logical_extension_codecs=( - NamedCapsule( + _NamedCodec( self.exporter.__datafusion_logical_extension_codec__(), "my_library.first", ), - NamedCapsule( + _NamedCodec( self.exporter.__datafusion_logical_extension_codec__(), "my_library.second", ), @@ -1017,36 +1091,40 @@ def __datafusion_session_extension__(self, ctx): ] -def test_with_extensions_rejects_two_bare_capsules_from_one_extension(ctx): - """Both capsules resolve to the one extension's id, so they collide. +def test_with_extensions_rejects_two_codecs_of_one_class(ctx): + """Two wrappers of one class claim one class-derived id, so they collide. Numbering them by position would be an id another library can mint the same value from, and would break stored plans the first time the extension reordered what it returns, so the ambiguity is refused instead. """ - class TwoCapsuleExtension: + class UnnamedCodec: + def __init__(self, capsule): + self._capsule = capsule + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule + + class TwoUnnamedCodecs: def __init__(self): self.exporter = SessionContext() def __datafusion_session_extension__(self, ctx): + capsule = self.exporter.__datafusion_logical_extension_codec__ return SessionExtensionComponents( logical_extension_codecs=( - self.exporter.__datafusion_logical_extension_codec__(), - self.exporter.__datafusion_logical_extension_codec__(), + UnnamedCodec(capsule()), + UnnamedCodec(capsule()), ), ) with pytest.raises(ValueError, match="__datafusion_codec_id__"): - ctx.with_extensions(TwoCapsuleExtension()) + ctx.with_extensions(TwoUnnamedCodecs()) 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. - """ + """A codec handed over as an object keeps the identity it declares.""" exporter = SessionContext() class ObjectCodecExtension: From 31a468793424a9252febf63fa0cdf29d6f23c65a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 5 Sep 2026 10:46:15 -0400 Subject: [PATCH 14/33] Install extension codecs and planners in two phases MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A session chains many codecs and dispatches between them by id, so codecs accumulate and their order does not affect decoding. A session holds exactly one query planner, so planners cannot accumulate — they compose by nesting, each wrapping the one before it. Collecting both from a single hook forced with_extensions to refuse more than one planner per call, because every factory ran before anything was installed and so no bundle could see another bundle's planner to wrap it. Two libraries that each ship a planner could not be installed together at all, and splitting them across two calls silently discarded the first. Codecs now come from __datafusion_session_extension__ and planners from a new __datafusion_session_planner__(ctx, fallback), which runs once per bundle in argument order after every codec is installed. Each receives the planner built so far; wrapping it nests this bundle outside the previous one, so the last bundle listed ends up outermost. A bundle implements either hook or both, which also lets a library that ships only an optimizing planner stop returning empty components. SessionExtensionComponents loses its query_planner field. Running the planner hooks after every codec is installed is what makes a nested planner safe. The rebuild that follows a later codec install reaches only the outermost layer, so a fallback captured against a partial chain would stay stale; there is now no "afterwards" within a call. Atomicity is unchanged. _install_extension_codecs writes nothing — the chains belong to the returned handle — so phase one is transactional for free, and the nest is built in memory with _install_extension_planner performing the single session write after the last hook returns. A hook that raises in either phase leaves the caller's context as it was. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 8 ++ crates/core/src/context.rs | 86 +++++++---- docs/source/contributor-guide/ffi.md | 67 +++++++-- .../README.md | 12 +- .../_test_three_library_query_planner.py | 70 ++++++++- .../src/extension.rs | 47 ++++-- python/datafusion/__init__.py | 2 + python/datafusion/context.py | 135 +++++++++++------- python/datafusion/extensions.py | 89 ++++++++++-- python/tests/test_context.py | 81 +++++++++-- python/tests/test_wrapper_coverage.py | 8 +- 11 files changed, 467 insertions(+), 138 deletions(-) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 826763ce1..3dfaf353f 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -62,6 +62,14 @@ library reaches things only the session has. session satisfies the protocol too — `ctx.__datafusion_query_planner__()` and `ctx.__datafusion_query_planner__(ctx)` are both valid. +`__datafusion_session_planner__(ctx, fallback)` is the exception to the shape +above: it takes a second argument, the planner assembled so far. A session has +one planner slot, so planners compose by nesting rather than by chaining, and +the host hands each bundle the previous layer instead of letting it capture one. +Wrap `fallback` and delegate to it; returning a planner that ignores it discards +every layer beneath, including one the session already had. It runs after every +bundle's codecs are installed, so `ctx` carries the final chains. + A *codec* must always be handed over as an object implementing its getter, never as the bare capsule the getter returns; `with_extensions` refuses a capsule. A codec's wire id — the string a payload names on decode, which has to mean the diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 8ade72397..4224d1574 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1614,32 +1614,29 @@ impl PySessionContext { derived } - /// Commit a `with_extensions` transaction onto this context. + /// Build the codec chains for a `with_extensions` call. /// - /// 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 + /// Private support method for `SessionContext.with_extensions`, and the + /// first of the two phases that method runs. `self` is the context the + /// extensions bound their components against, and is also the + /// `Arc` every FFI task-context provider they created /// targets, so the returned handle shares it rather than deriving a new - /// one. Codec capsules are imported and validated before anything is - /// committed, so a failure leaves the session untouched. + /// one. /// - /// 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. + /// **Writes nothing.** The codec chains belong to the returned handle + /// rather than to `SessionState`, so this phase is transactional for free: + /// a codec that fails to import, or that collides with an installed id, + /// leaves the caller's context exactly as it was. Binding the planner is + /// the only step that touches the session, and it is deferred to + /// [`Self::_install_extension_planner`] so the planner hooks can run + /// against the final chains. /// /// Codecs must arrive as objects exposing the capsule getter, never as - /// bare capsules — see [`resolve_bundle_codec_id`]. The planner has no - /// wire id, so it may still be a capsule. - #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))] - pub fn _install_extensions<'py>( + /// bare capsules — see [`resolve_bundle_codec_id`]. + pub fn _install_extension_codecs<'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 @@ -1676,22 +1673,51 @@ impl PySessionContext { } let physical_codec = Arc::new(physical_codec); - let planner = planner - .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))) - .transpose()?; - - let installed = Self { + Ok(Self { ctx: Arc::clone(&slf.borrow().ctx), logical_codec, physical_codec, - }; - // Bind the planner only once the codec chains are final, and through - // the new handle so it carries them. Passing `None` still rebuilds - // whichever planner the session already holds against the new chains, - // exactly as `with_logical_extension_codec` does. - installed.set_session_query_planner(planner); + }) + } + + /// Re-export a planner a `__datafusion_session_planner__` hook returned as + /// a capsule, so the next hook in the chain receives one either way. + /// + /// A hook may hand back an object exposing `__datafusion_query_planner__` + /// or a raw capsule; the next hook wraps whatever it is given and should + /// not have to branch on which. Importing here also surfaces a malformed + /// planner at the hook that produced it rather than at the final install. + /// Writes nothing. + pub fn _rebind_query_planner<'py>( + slf: &Bound<'py, Self>, + planner: Bound<'py, PyAny>, + ) -> PyDataFusionResult> { + let ffi = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?; + Ok(create_query_planner_capsule(slf.py(), &ffi)?) + } - Ok(installed) + /// Commit the query planner for a `with_extensions` call. + /// + /// The second phase, run once every codec is installed and every planner + /// hook has returned, so the planner is bound against the final chains. + /// This is the one call in `with_extensions` that writes to the session, + /// and it goes through this context's own `state_ref()`, so providers + /// bound to it stay valid. + /// + /// `None` means no bundle supplied a planner. That still rebuilds + /// whichever planner the session already holds against the new chains, + /// exactly as `with_logical_extension_codec` does, and writes nothing at + /// all if the session has no FFI planner to rebuild. + #[pyo3(signature = (planner=None))] + pub fn _install_extension_planner<'py>( + slf: &Bound<'py, Self>, + planner: Option>, + ) -> PyDataFusionResult<()> { + let planner = planner + .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))) + .transpose()?; + slf.borrow().set_session_query_planner(planner); + Ok(()) } } diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 3ed8ccfd4..c38f705cf 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -356,23 +356,26 @@ codec chains exist when it is installed, and a codec added afterwards rebinds it that wrong and the planner encodes through a chain that is missing a library. `SessionContext.with_extensions` removes the ordering question. An extension library -exposes a bundle object implementing `__datafusion_session_extension__`: +exposes a bundle object implementing one or both of two hooks: ```python class MyEngineExtension: def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: - # Create fresh components bound to `ctx` on every call. `ctx` is the - # session the components will run on. + # Phase one. Create fresh components bound to `ctx` on every call. return SessionExtensionComponents( logical_extension_codecs=(self._make_logical_codec(ctx),), physical_extension_codecs=(self._make_physical_codec(ctx),), - query_planner=self._make_planner(ctx), ) + + def __datafusion_session_planner__(self, ctx: SessionContext, fallback): + # Phase two. `ctx` now carries every bundle's codecs, and `fallback` is + # the planner built so far. Wrapping it is what makes this library + # compose with the other planners in the call. + return self._make_planner(ctx, fallback=fallback) ``` -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: +Implement whichever apply: a codec-only library defines the first, a library that ships +only an optimizing planner defines the second. ```python ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) @@ -380,11 +383,44 @@ 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. Supplying one replaces whatever planner -the session already has; to layer instead, capture the existing planner from -`__datafusion_query_planner__` first and have yours fall back to it. +#### Two phases, because codecs and planners compose differently + +A session chains **many** codecs and dispatches between them by id. Codecs therefore +just accumulate: order affects encoding only, and decoding always routes to the codec +that wrote the payload. A session holds exactly **one** query planner, so planners +cannot accumulate — they compose by *nesting*, each wrapping the one before it and +delegating to it for work it does not handle. + +So `with_extensions` runs every `__datafusion_session_extension__` and installs all the +codecs, and only then runs each `__datafusion_session_planner__`, in argument order, +handing each the planner built so far. Two consequences worth holding onto: + +- **Bundle order is irrelevant for codecs and significant for planners.** The last + extension listed ends up outermost and is consulted first. +- **A planner is always built against the complete codec set**, including codecs from + bundles listed after it. This is what the low-level chaining cannot give you, and it + matters most for a nested planner: the rebuild that follows a later codec install + reaches only the outermost layer (see [Rebinding a planner's codecs is one level + deep](#rebinding-a-planners-codecs-is-one-level-deep)), so a fallback captured before + the codecs were complete would stay stale forever. + +An extension that ignores `fallback` and returns an unrelated planner replaces every +layer beneath it, including any planner the session already had. That is legal — a +library that must be the only planner does it deliberately — but it is not composable, +and nothing detects it. Returning `None` contributes no planner and leaves `fallback` +in place. + +Three libraries that each ship a planner therefore install like this, with the +outermost last: + +```python +ctx = SessionContext(config).with_extensions( + tables.Extension(), # codecs only + functions.Extension(), # codecs only + optimizer.Extension(), # planner, wrapping the session default + distributed.Extension(), # planner, wrapping the optimizer +) +``` #### Codecs are objects, not capsules @@ -624,8 +660,11 @@ 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. +`with_extensions` sidesteps this entirely, and for nested planners too: every codec from +every bundle is installed before the first planner hook runs, so no layer — outer or +fallback — is ever captured against a partial chain. There is no "afterwards" within a +call. Prefer it over hand-layering whenever the planners you are composing all ship as +bundles. ## Alternative Approach diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index efcb3376f..4b9140b53 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -50,11 +50,13 @@ ctx.register_table("numbers", provider) ctx.register_udf(provider_udf) ``` -`MyPlannerExtension` implements the `__datafusion_session_extension__` protocol: it -receives the session it is being installed on, binds fresh codec and planner -components to that session's task-context provider, and returns them as -`SessionExtensionComponents`. The host installs every codec before it binds the -planner, so the planner cannot be left carrying a chain that has since grown. +`MyPlannerExtension` implements both extension hooks. `__datafusion_session_extension__` +receives the session it is being installed on, binds fresh codecs to that session's +task-context provider, and returns them as `SessionExtensionComponents`. +`__datafusion_session_planner__` then runs in the host's second phase, after every +bundle's codecs are installed, and builds a planner that delegates to the `fallback` it +is handed — so several libraries that each ship a planner nest instead of displacing one +another, and no planner is left carrying a chain that has since grown. Its codecs are handed over as `BundledLogicalCodec` and `BundledPhysicalCodec` rather than as bare capsules. `with_extensions` requires an object, because a codec's wire id 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 d272bd998..684c96404 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 @@ -846,6 +846,70 @@ def __datafusion_session_extension__( SessionContext(config).with_extensions(BareCapsuleExtension()) +class PlannerOnlyExtension: + """A library that ships a planner and no codecs. + + Implements only the planner hook — there is nothing to contribute in phase + one, and the protocol should not make it say so. + """ + + def __init__(self) -> None: + self.planner = None + + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + self.planner = MyQueryPlanner(fallback=fallback) + return self.planner + + +def test_with_extensions_nests_planners_in_argument_order(): + """Two planner-shipping libraries compose instead of displacing each other. + + This is the four-library case: A and C contribute codecs, B an optimizing + planner, D a distributed one that should sit outside B. Both planners run + for one query, which is only possible if D delegates to B rather than + replacing it — a session holds exactly one planner, so the nesting is the + only way both are reachable. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + codecs = ProviderCodecsExtension() + inner = PlannerOnlyExtension() + outer = PlannerOnlyExtension() + ctx = SessionContext(config).with_extensions(codecs, inner, outer) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + + assert inner.planner.plan_calls() >= 1 + assert outer.planner.plan_calls() >= 1 + # The last extension listed is outermost, so it is the one that had to + # delegate. The inner planner is the fallback, and reaches the session's + # original planner through its own. + assert outer.planner.used_fallback() + + +def test_with_extensions_planner_sees_every_bundles_codecs(): + """Phase two runs after phase one, for every bundle. + + A planner contributed by an early argument is still built against codecs a + later argument installed — the ordering trap that chaining the low-level + methods by hand leaves to the caller. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + planner_first = PlannerOnlyExtension() + codecs_last = ProviderCodecsExtension() + ctx = SessionContext(config).with_extensions(planner_first, codecs_last) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + # The provider codecs were installed after the planner was listed, and the + # query still round-trips its table provider through them. + assert codecs_last.logical_codec.table_provider_decode_calls() > 0 + + def test_with_extensions_codec_ids_survive_bundle_composition(): """Nesting a bundle inside another does not re-tag its codecs. @@ -1127,9 +1191,13 @@ def __datafusion_session_extension__( *codecs.physical_extension_codecs, *planner.physical_extension_codecs, ), - query_planner=planner.query_planner, ) + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + return self._planner.__datafusion_session_planner__(ctx, fallback) + def test_with_extensions_docstring_example_still_runs(): """Run the ``with_extensions`` docstring example verbatim. diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs index 8307b5761..f491e3114 100644 --- a/examples/datafusion-ffi-query-planner-example/src/extension.rs +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -36,7 +36,8 @@ use datafusion_proto::physical_plan::{ 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, + ffi_physical_codec_from_pycapsule, ffi_query_planner_from_pycapsule, + ffi_task_context_provider_from_pycapsule, get_tokio_runtime, }; use datafusion_session::QueryPlanner; use pyo3::prelude::*; @@ -360,26 +361,42 @@ impl MyPlannerExtension { }, )?; - 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_codec,))?; kwargs.set_item("physical_extension_codecs", (physical_codec,))?; - kwargs.set_item("query_planner", planner_capsule)?; components.call((), Some(&kwargs)) } + + /// Contribute this library's planner, nesting it on whatever came before. + /// + /// Runs in the host's second phase, after every bundle's codecs are + /// installed, so `ctx` carries the final chains and the planner this + /// builds is not left encoding through a partial set. `fallback` is the + /// planner assembled so far — the session's existing one for the first + /// bundle, the previous bundle's for the rest — and delegating to it is + /// what makes several planner-shipping libraries composable. Returning a + /// planner that ignored it would discard every layer beneath. + fn __datafusion_session_planner__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + fallback: Bound<'py, PyAny>, + ) -> PyResult> { + let fallback = ffi_query_planner_from_pycapsule(&fallback, Some(&ctx))?; + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + fallback: Some((&fallback).into()), + }); + // The planner takes the host's codecs, not ones built here. By now + // those are the final chains, and this library has no business minting + // a provider of its own. + let host_logical = ffi_logical_codec_from_pycapsule(ctx.clone(), None)?; + let host_physical = ffi_physical_codec_from_pycapsule(ctx, None)?; + let ffi_planner = + FFI_QueryPlanner::new_with_ffi_codecs(planner, host_logical, host_physical); + create_query_planner_capsule(py, &ffi_planner) + } } diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 4b02a383e..9484f2f50 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -95,6 +95,7 @@ from .extensions import ( SessionExtensionComponents, SessionExtensionExportable, + SessionPlannerExportable, ) from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions @@ -140,6 +141,7 @@ "SessionContext", "SessionExtensionComponents", "SessionExtensionExportable", + "SessionPlannerExportable", "Table", "TableFunction", "TableProviderFactory", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 5d6037cef..f7838030c 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -73,6 +73,7 @@ QueryPlannerExportable, SessionExtensionComponents, SessionExtensionExportable, + SessionPlannerExportable, ) from datafusion.options import ( DEFAULT_MAX_INFER_SCHEMA, @@ -1816,20 +1817,36 @@ def with_extensions( """Create a new session context with the given extension bundles. This is the preferred way to install FFI extensions that need a - task-context provider (extension codecs and query planners). Each - extension's ``__datafusion_session_extension__`` method is called with - this context so it can bind its components to the session they will - run on, then all components are installed in one step. This avoids the - pitfalls of chaining :py:meth:`with_logical_extension_codec`, + task-context provider (extension codecs and query planners). It avoids + the pitfalls of chaining :py:meth:`with_logical_extension_codec`, :py:meth:`with_physical_extension_codec`, and :py:meth:`set_query_planner` by hand, where the codecs a planner was built against can end up stale. - Codecs compose with the existing chain and with each other: extensions - are processed left to right and their codecs are appended to the chain - in that order. Decoding routes by codec id, so the order matters only - for encoding. At most one extension may supply a query planner. If none - does, an existing FFI planner is rebound to the final codec chains. + Installation runs in two phases, because codecs and planners compose + differently: + + 1. Every extension's ``__datafusion_session_extension__`` is called + with this context and its codecs are collected, then all of them are + installed at once. A session chains many codecs and dispatches + between them by id, so they merely accumulate; order affects + encoding only. + 2. Every extension's ``__datafusion_session_planner__`` is then called, + **in argument order**, each receiving the planner built so far. A + session holds exactly one planner, so planners compose by *nesting*: + each wraps the previous one and delegates to it. The last extension + listed ends up outermost and is consulted first. + + An extension implements either hook or both. Phase two runs after every + codec is installed and receives a context carrying the final chains, so + a nested planner is never left encoding through a chain a later + extension has grown. + + If no extension supplies a planner, an existing FFI planner is rebound + to the final codec chains and the session's planner is otherwise left + alone. An extension that ignores the ``fallback`` it is handed replaces + the planners before it instead of nesting on them, including any the + session already had. Codecs must be handed over as objects exposing the capsule getter, not as bare ``PyCapsule`` objects, and are named after their exporting @@ -1841,8 +1858,8 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare wire identity independent of the extension that ships it, so an extension composed inside another one still writes the same ids. - The planner is exempt — it carries no wire id, so it may be an object - or a capsule. + Planners are exempt — a planner carries no wire id, so a hook may + return an object or a capsule. Like the individual ``with_*`` methods, the returned context shares its session with this one: catalogs, tables, registered functions, and @@ -1852,11 +1869,13 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare chains are specific to the returned handle. No state is written until every extension has run and every capsule has - been validated, so an extension that raises or returns invalid - components leaves the session as it was. The exception is an extension - that mutates the context it is handed — registering a table, say — - which is not rolled back. Extension factories should treat that context - as configuration-only. + been validated, so an extension that raises or returns something + invalid — in either phase — leaves the session as it was. Codec chains + belong to the returned handle, and the single session write happens + after the last planner hook returns. The exception is an extension that + mutates the context it is handed — registering a table, say — which is + not rolled back. Extension factories should treat that context as + configuration-only. The session owns the installed components' task-context providers, and dependent objects do not extend its lifetime. Keep a context on the @@ -1864,24 +1883,25 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare use; FFI operations after the last one is collected raise an error. Args: - extensions: Extension bundles to install, in the order their - codecs join the chain. + extensions: Extension bundles to install. Order is irrelevant for + codecs and significant for planners, which nest in this order + with the last one outermost. Returns: A new context with all extension components installed. Raises: - TypeError: If an argument does not implement the protocol, returns - something other than a - :py:class:`SessionExtensionComponents`, or contributes a codec - as a bare ``PyCapsule``. - ValueError: If no extensions are given, more than one extension - supplies a query planner, or two codecs claim the same id. An - extension that contributes two instances of one codec class - 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. + TypeError: If an argument implements neither hook, if + ``__datafusion_session_extension__`` returns something other + than a :py:class:`SessionExtensionComponents`, or if an + extension contributes a codec as a bare ``PyCapsule``. + ValueError: If no extensions are given, or two codecs claim the + same id. An extension that contributes two instances of one + codec class must declare ``__datafusion_codec_id__`` on at + least one of them; the collision is refused rather than + resolved by position, because a positional id would break + stored plans the first time the extension reordered what it + returns. Examples: The example is skipped here because it needs a built FFI @@ -1903,22 +1923,25 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare msg = "with_extensions requires at least one extension" raise ValueError(msg) for extension in extensions: - if not isinstance(extension, SessionExtensionExportable): + if not isinstance( + extension, (SessionExtensionExportable, SessionPlannerExportable) + ): msg = ( - "Extension does not implement __datafusion_session_extension__: " - f"{extension!r}" + "Extension implements neither " + "__datafusion_session_extension__ nor " + f"__datafusion_session_planner__: {extension!r}" ) raise TypeError(msg) - # Bind every component against this context, not a context derived from - # it. There is one `Arc` per session, so a component - # bound here holds a task-context provider that the returned handle - # keeps alive, and `_install_extensions` writes the final state through - # that same session. + # Phase one: collect every bundle's codecs. Components are bound + # against this context, not a context derived from it. There is one + # `Arc` per session, so a component bound here holds a + # task-context provider that the returned handle keeps alive. logical_codecs: list[LogicalExtensionCodecExportable] = [] physical_codecs: list[PhysicalExtensionCodecExportable] = [] - planner: QueryPlannerExportable | _PyCapsule | None = None for extension in extensions: + if not isinstance(extension, SessionExtensionExportable): + continue components = extension.__datafusion_session_extension__(self) if not isinstance(components, SessionExtensionComponents): msg = ( @@ -1929,18 +1952,32 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare 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 + # Writes nothing: the chains belong to the new handle, so a failure + # above or below leaves this context as it was. new = SessionContext.__new__(SessionContext) - new.ctx = self.ctx._install_extensions(logical_codecs, physical_codecs, planner) + new.ctx = self.ctx._install_extension_codecs(logical_codecs, physical_codecs) + + # Phase two: nest the planners, outermost last. Each hook runs against + # `new`, which carries the final chains, so a planner captured here + # never sees a partial codec set. `planner` stays None when no bundle + # supplies one, which leaves an already-installed planner in place + # rather than wrapping the session's default in an FFI hop. + planner: _PyCapsule | None = None + for extension in extensions: + if not isinstance(extension, SessionPlannerExportable): + continue + fallback = ( + planner + if planner is not None + else new.ctx.__datafusion_query_planner__() + ) + supplied = extension.__datafusion_session_planner__(new, fallback) + if supplied is None: + continue + planner = new.ctx._rebind_query_planner(supplied) + + new.ctx._install_extension_planner(planner) return new def table_provider(self, name: str) -> Table: diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 6c95b9003..e9f2f7156 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -29,9 +29,21 @@ Installing through ``with_extensions`` rather than by chaining the individual ``with_*`` methods matters for components that hold a task-context provider: the extension is handed the session its components will run on, and every -codec is installed before the query planner is bound against them, so no +codec is installed before any query planner is bound against them, so no planner is left carrying a codec chain that has since grown. See the FFI extensions guide in the contributor documentation for the full rationale. + +Codecs and planners install in two phases, because they compose differently. A +session's codec chain holds many codecs and dispatches between them by id, so +codecs merely accumulate and their order does not affect decoding. A session +holds exactly *one* query planner, so planners compose by nesting: each wraps +the one before it. Phase one collects every bundle's codecs through +:py:class:`SessionExtensionExportable` and installs them; phase two runs +:py:class:`SessionPlannerExportable` once per bundle, in argument order, +handing each the planner built so far. + +That split is what lets several libraries that each ship a planner coexist. It +also means bundle order is significant for planners and irrelevant for codecs. """ from __future__ import annotations @@ -52,6 +64,7 @@ "QueryPlannerExportable", "SessionExtensionComponents", "SessionExtensionExportable", + "SessionPlannerExportable", ] @@ -79,6 +92,9 @@ class SessionExtensionComponents: components bound to a different session hold a task-context provider for that other session and cannot be rebound. + Query planners are not listed here. They install in a second phase so each + can wrap the one before it — see :py:class:`SessionPlannerExportable`. + Codecs must be objects exposing the capsule getters, never bare ``PyCapsule`` objects: a codec's id is read off the object it is handed over as, and a capsule has no type to read. A library holding a raw capsule @@ -87,15 +103,15 @@ class SessionExtensionComponents: different extension. Examples: - A bundle that contributes nothing is valid, and is what the defaults - describe: + A bundle that contributes no codecs is valid — a planner-only library + returns this, or omits the hook entirely: >>> from datafusion import SessionExtensionComponents >>> components = SessionExtensionComponents() >>> components.logical_extension_codecs () - >>> components.query_planner is None - True + >>> components.physical_extension_codecs + () A bundle that contributes one kind of component names it, leaving the rest empty. Here the codec is a capsule wrapped in an object that @@ -126,14 +142,6 @@ class SessionExtensionComponents: physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = () """Physical codecs to add to the session's codec chain, in declaration order.""" - query_planner: QueryPlannerExportable | _PyCapsule | None = None - """Optional query planner. - - At most one extension per - :py:meth:`~datafusion.context.SessionContext.with_extensions` call may - supply one. - """ - @runtime_checkable class SessionExtensionExportable(Protocol): @@ -152,6 +160,11 @@ class SessionExtensionExportable(Protocol): mutating the context they are handed — a registration made during binding is not rolled back if a later extension fails. + A bundle that also contributes a query planner implements + :py:class:`SessionPlannerExportable` alongside this protocol. Planners are + installed in a second phase, so they are not part of the components + returned here. + Examples: >>> from datafusion import ( ... SessionExtensionComponents, @@ -169,3 +182,53 @@ class SessionExtensionExportable(Protocol): def __datafusion_session_extension__( # noqa: D105 self, ctx: SessionContext ) -> SessionExtensionComponents: ... + + +@runtime_checkable +class SessionPlannerExportable(Protocol): + """Type hint for extension bundles that contribute a query planner. + + A session holds exactly one query planner, so planners compose by nesting + rather than by chaining: each wraps the one before it and delegates to it + for the work it does not handle. + :py:meth:`~datafusion.context.SessionContext.with_extensions` runs this + hook once per bundle that implements it, **in argument order**, handing + each the planner built so far. Returning a planner that wraps ``fallback`` + puts this bundle *outside* the previous one, so the last bundle listed ends + up outermost and is consulted first. + + The hook runs after every codec from every bundle is installed, and ``ctx`` + is the context carrying those final chains. That ordering is the point: a + planner captured here sees the complete codec set, so a nested planner is + not left encoding through a chain that a later bundle has grown. + + Return ``None`` to contribute no planner and leave ``fallback`` in place. + Ignoring ``fallback`` and returning a planner that does not delegate to it + is legal and means "replace" — but it discards every planner listed before + this one, including any the session already had. + + Args: + ctx: The session the planner will run on, carrying the final codec + chains. + fallback: The planner built so far, as a ``PyCapsule``. For the first + bundle this is the session's existing planner, which is the + DataFusion default unless one was installed earlier. + + Examples: + >>> from datafusion import SessionPlannerExportable + >>> class MyEngineExtension: + ... def __datafusion_session_planner__(self, ctx, fallback): + ... # A real library returns its own planner wrapping + ... # `fallback`, e.g. ``my_library.Planner(fallback=fallback)``. + ... # Handing it straight back is the degenerate wrap: valid, + ... # and contributes nothing. + ... return fallback + >>> isinstance(MyEngineExtension(), SessionPlannerExportable) + True + >>> isinstance(object(), SessionPlannerExportable) + False + """ + + def __datafusion_session_planner__( # noqa: D105 + self, ctx: SessionContext, fallback: _PyCapsule + ) -> QueryPlannerExportable | _PyCapsule | None: ... diff --git a/python/tests/test_context.py b/python/tests/test_context.py index d41173102..56b90b307 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -933,12 +933,21 @@ def __datafusion_session_extension__(self, ctx): class _PlannerExtension: - """Contributes the receiving session's own exported planner.""" + """Contributes a planner, recording the fallback it was handed. - def __datafusion_session_extension__(self, ctx): - return SessionExtensionComponents( - query_planner=ctx.__datafusion_query_planner__() - ) + Passing ``fallback`` straight back through is the degenerate wrap: the + resulting session plans exactly as it did before, which is what lets a + pure-Python test assert the threading without a real layering planner. + """ + + def __init__(self): + self.fallbacks = [] + self.planner_ctx = None + + def __datafusion_session_planner__(self, ctx, fallback): + self.planner_ctx = ctx + self.fallbacks.append(fallback) + return fallback def test_with_extensions_requires_an_extension(ctx): @@ -947,7 +956,7 @@ def test_with_extensions_requires_an_extension(ctx): def test_with_extensions_rejects_non_extension(ctx): - with pytest.raises(TypeError, match="__datafusion_session_extension__"): + with pytest.raises(TypeError, match="__datafusion_session_planner__"): ctx.with_extensions(object()) @@ -960,9 +969,63 @@ def __datafusion_session_extension__(self, ctx): 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_accepts_a_planner_only_extension(ctx): + """An extension may implement the planner hook alone. + + A library that ships an optimizing planner and no codecs — nothing to + contribute in phase one — should not have to return empty components. + """ + extension = _PlannerExtension() + result = ctx.with_extensions(extension) + + assert len(extension.fallbacks) == 1 + assert result.session_id() == ctx.session_id() + + +def test_with_extensions_threads_the_planner_through_in_order(ctx): + """Each planner hook receives what the previous one returned. + + Planners nest rather than chain, so the host hands each bundle the planner + built so far. Argument order is nesting order, last one outermost. + """ + first, second = _PlannerExtension(), _PlannerExtension() + ctx.with_extensions(first, second) + + assert len(first.fallbacks) == 1 + assert len(second.fallbacks) == 1 + # `first` returned its fallback unchanged, and the host normalizes each + # hook's return value before passing it on, so `second` sees a capsule + # standing for the same planner rather than the session's original. + assert second.fallbacks[0] is not None + + +def test_with_extensions_planner_hook_sees_the_new_handle(ctx): + """Phase two runs against the handle carrying the final codec chains. + + A planner captured against the pre-install handle would encode through a + chain missing every codec this call installed. + """ + codecs = _CodecOnlyExtension() + planner = _PlannerExtension() + result = ctx.with_extensions(codecs, planner) + + assert planner.planner_ctx.logical_extension_codec_ids() == ["my_library.logical"] + assert result.logical_extension_codec_ids() == ["my_library.logical"] + + +def test_with_extensions_skips_a_planner_hook_returning_none(ctx): + """Returning ``None`` contributes no planner and keeps the fallback.""" + + class NoPlanner: + def __datafusion_session_planner__(self, ctx, fallback): + return None + + downstream = _PlannerExtension() + ctx.with_extensions(NoPlanner(), downstream) + + # The skipped hook did not become `downstream`'s fallback; it got the + # session's own planner instead. + assert len(downstream.fallbacks) == 1 def test_with_extensions_rejects_bad_codec_capsule(ctx): diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index 7e8bd3f2d..dc6dfea11 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -33,8 +33,12 @@ # gap in coverage. PRIVATE_SUPPORT_METHODS = frozenset( { - # Support method for SessionContext.with_extensions. - "_install_extensions", + # The three phases of SessionContext.with_extensions: install the + # codecs, normalize each planner hook's return value, commit the + # planner. + "_install_extension_codecs", + "_rebind_query_planner", + "_install_extension_planner", } ) From 74ee4668c0c5cf44a8b261a90ff50a7c0adad0d2 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 5 Sep 2026 11:03:18 -0400 Subject: [PATCH 15/33] Give the planner example a node only its own codec can carry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bundle shipped a planner and codecs, but the halves never met: the planner emitted a stock GlobalLimitExec and the codecs delegated everything to the default codec. So the example asserted by structure that a planner and its codecs belong together without demonstrating why, and no test would have caught a bundle whose planner emits a node its own codec cannot encode. DistributedQueryPlanner now wraps its result in a DistributedExec, a type private to this library, and ObservingPhysicalExtensionCodec claims it by downcast and rebuilds it from its inputs. Nothing else in the session knows the type, which is the reason the two ship as one bundle. The observing codecs stop being dead weight in the process — they were previously never consulted, and decode_max_rows_seen had no caller. Also documents what codec order does and does not control, which building this surfaced. Decoding routes by id and is never order-dependent. Encoding stops at the first codec that claims the node, so a codec claiming a broad category — MyPhysicalExtensionCodec claims any ForeignExecutionPlan — takes nodes from any library installed after it. The query still succeeds; only the library that wrote the bytes changes, which breaks a plan that has to decode elsewhere. That gives a bundle two reasons to want different positions for its two halves. The guide now says to contribute each half at its own position with a small adapter rather than reordering, since the hooks are independent, and treats the low-level sequence as the last resort it is: it works, but it hands back responsibility for codec-before-planner ordering and leaves a hand-layered fallback holding the codecs it captured. No attempt is made to express every permutation from one call. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/contributor-guide/ffi.md | 82 ++++++++++- .../_test_three_library_query_planner.py | 128 ++++++++++++++++++ .../src/distributed_exec.rs | 108 +++++++++++++++ .../src/extension.rs | 86 ++++++++++-- .../src/lib.rs | 1 + .../src/planner.rs | 14 +- python/datafusion/context.py | 11 ++ 7 files changed, 413 insertions(+), 17 deletions(-) create mode 100644 examples/datafusion-ffi-query-planner-example/src/distributed_exec.rs diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index c38f705cf..54ef53b72 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -395,8 +395,10 @@ So `with_extensions` runs every `__datafusion_session_extension__` and installs codecs, and only then runs each `__datafusion_session_planner__`, in argument order, handing each the planner built so far. Two consequences worth holding onto: -- **Bundle order is irrelevant for codecs and significant for planners.** The last - extension listed ends up outermost and is consulted first. +- **Bundle order matters differently for each.** For planners it sets the nesting: the + last extension listed ends up outermost and is consulted first. For codecs it never + affects decoding, and affects encoding only when two codecs would claim the same node + — see [When codec order does matter](#when-codec-order-does-matter). - **A planner is always built against the complete codec set**, including codecs from bundles listed after it. This is what the low-level chaining cannot give you, and it matters most for a nested planner: the rebuild that follows a later codec install @@ -422,6 +424,82 @@ ctx = SessionContext(config).with_extensions( ) ``` +#### When codec order does matter + +Decoding is never order-dependent: a payload names its codec by id and the chain +dispatches straight to it. Encoding walks the chain in install order and stops at the +first codec that claims the node. Most of the time that is invisible, because libraries +claim disjoint things — one owns its table providers, another its UDFs, a third its own +execution plan nodes. + +It stops being invisible when a codec claims *broadly*. A node that came from another +library arrives as an opaque `ForeignExecutionPlan`, and a codec that claims any of +those will take nodes it does not own from any library installed after it. The query +still succeeds. What changes is which library wrote the bytes — so a plan that has to +decode in another process now needs whichever library happened to win, not the one whose +node it is. `MyPhysicalExtensionCodec` in the provider example claims this way, and +`test_a_greedy_codec_installed_first_claims_another_librarys_node` pins the consequence. + +Two rules of thumb: + +- **Writing a codec, claim narrowly.** Downcast to your own types. Claiming a broad + category makes your library order-sensitive for everyone downstream of it. +- **Shipping plans out of the process, verify.** Do not assume your node reached your + codec just because both are installed. Round-trip a plan through + `ExecutionPlan.to_bytes` / `from_bytes` in a test and assert your codec did the work. + +#### When the two orders conflict + +Because codec position and planner position both come from one argument list, a library +can in principle need to be early for one and late for the other: its codec must precede +a broad claimer, while its planner must nest outside that library's planner. + +Do not try to satisfy both by reordering — contribute each half at its own position. The +two hooks are independent, so a three-line adapter each is enough: + +```python +class CodecsOf: + """Contribute only the codec half of a bundle, at this position.""" + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_extension__(self, ctx): + return self.inner.__datafusion_session_extension__(ctx) + + +class PlannerOf: + """Contribute only the planner half of a bundle, at this position.""" + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_planner__(self, ctx, fallback): + return self.inner.__datafusion_session_planner__(ctx, fallback) + + +ctx = SessionContext(config).with_extensions( + CodecsOf(engine), CodecsOf(tables), # engine's codec first + PlannerOf(tables), PlannerOf(engine), # engine's planner outermost +) +``` + +This keeps everything `with_extensions` guarantees: one transaction, codecs complete +before any planner is built, codec ids untouched — an id is read off the codec object, +not off the extension that contributed it, so splitting a bundle cannot re-tag its +payloads. A library that expects to be composed this way should expose the halves itself +rather than make callers write the adapters. + +Falling back to the low-level `with_logical_extension_codec` / +`with_physical_extension_codec` / `set_query_planner` sequence also works, and it is the +right answer when the pieces do not come as bundles at all. But it is a real downgrade, +not just a more verbose spelling: you take back responsibility for installing every +codec before every planner, and a planner you layer by hand keeps the codecs it captured +— the [one-level rebind](#rebinding-a-planners-codecs-is-one-level-deep) does not reach +inside it. Reach for it last. + +There is no attempt here to make every permutation expressible from one call. Two +positions per bundle covers the cases that arise; anything stranger is a sign the +libraries disagree about what they own, which is better fixed there. + #### Codecs are objects, not capsules `with_extensions` requires each codec to be an object exposing the capsule getter, and 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 684c96404..891c91296 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 @@ -34,6 +34,7 @@ col, udf, ) +from datafusion.plan import ExecutionPlan from datafusion_ffi_example import ( IsNullUDF, MyCatalogProvider, @@ -846,6 +847,133 @@ def __datafusion_session_extension__( SessionContext(config).with_extensions(BareCapsuleExtension()) +def test_bundle_codec_carries_its_own_planners_node(): + """The two halves of a bundle meet: its codec serializes its planner's node. + + ``DistributedQueryPlanner`` emits a ``DistributedExec``, a type private to + this library. No other codec in the session knows it, so the planner is + only useful alongside the codec that carries it — which is why the two ship + as one bundle, and why ``with_extensions`` installs every codec before it + binds any planner. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(bundle, ProviderCodecsExtension()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert bundle.distributed_exec_encode_calls() > 0 + assert bundle.distributed_exec_decode_calls() > 0 + # Reaching the session config from inside those decode callbacks is what + # shows the provider bound at installation resolves against this session. + assert bundle.decode_max_rows_seen() == [2] * len(bundle.decode_max_rows_seen()) + assert bundle.decode_max_rows_seen() + + +def test_bundle_planners_node_survives_a_plan_round_trip(): + """A plan carrying the node serializes and comes back intact. + + This is the path a distributed engine takes to ship a plan to a remote + executor, and the reason its node's id has to mean the same thing there. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(bundle, ProviderCodecsExtension()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + plan = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').execution_plan() + assert "DistributedExec" in plan.display() + + before = bundle.distributed_exec_encode_calls() + restored = ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx)) + + assert bundle.distributed_exec_encode_calls() > before + assert "DistributedExec" in restored.display() + + +def test_a_greedy_codec_installed_first_claims_another_librarys_node(): + """Encode order decides *which* library serializes a node. + + Decoding routes by id, so codec order never affects it. Encoding walks the + chain in install order and stops at the first codec that claims the node — + and a codec may claim broadly. ``MyPhysicalExtensionCodec`` claims any + ``ForeignExecutionPlan``, which is what a node from another library looks + like once it crosses the boundary, so installing it ahead of this bundle + takes the bundle's own node away from it. + + Nothing detects this. A library whose plans must decode elsewhere should + not assume its node reached its own codec just because both are installed. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + # Provider codecs first, so their broad claim wins. + ctx = SessionContext(config).with_extensions(ProviderCodecsExtension(), bundle) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + # The query still runs — the node was carried, just not by its own library. + assert bundle.distributed_exec_encode_calls() == 0 + + +class CodecsOf: + """Contribute only the codec half of a bundle, at this position.""" + + def __init__(self, inner: object) -> None: + self.inner = inner + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return self.inner.__datafusion_session_extension__(ctx) + + +class PlannerOf: + """Contribute only the planner half of a bundle, at this position.""" + + def __init__(self, inner: object) -> None: + self.inner = inner + + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + return self.inner.__datafusion_session_planner__(ctx, fallback) + + +def test_splitting_a_bundle_resolves_conflicting_orders(): + """A bundle can take one position for its codec and another for its planner. + + Codec position and planner position both come from one argument list, so a + library can need to be early for one and late for the other: here the + bundle's codec must precede the provider's broad claim, while its planner + must stay outermost. Splitting the halves satisfies both without giving up + what ``with_extensions`` guarantees. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + bundle = MyPlannerExtension() + provider = ProviderCodecsExtension() + ctx = SessionContext(config).with_extensions( + CodecsOf(bundle), + CodecsOf(provider), + PlannerOf(bundle), + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + # Its codec ran ahead of the provider's broad claim, so the bundle kept its + # own node... + assert bundle.distributed_exec_encode_calls() > 0 + # ...and splitting did not re-tag anything: an id is read off the codec + # object, never off the extension that contributed it. + assert PHYSICAL_CODEC_ID in ctx.physical_extension_codec_ids() + assert ( + "datafusion_ffi_example.MyPhysicalExtensionCodec" + in ctx.physical_extension_codec_ids() + ) + + class PlannerOnlyExtension: """A library that ships a planner and no codecs. diff --git a/examples/datafusion-ffi-query-planner-example/src/distributed_exec.rs b/examples/datafusion-ffi-query-planner-example/src/distributed_exec.rs new file mode 100644 index 000000000..a0a347246 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/distributed_exec.rs @@ -0,0 +1,108 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A custom execution plan node owned by this library. +//! +//! This is the half of an extension bundle that makes the other half +//! necessary. A query planner that only rearranges stock DataFusion nodes needs +//! no codec of its own; one that emits a node *it* defines does, because +//! nothing else in the process knows how to serialize it. Shipping the planner +//! and the codec that carries its nodes as one bundle is the normal case, and +//! it is why `with_extensions` installs every codec before it binds any +//! planner. +//! +//! The node itself is deliberately trivial — it passes its child's stream +//! through untouched. A real distributed engine would ship the child plan to a +//! remote executor here; what matters for the example is that the node exists, +//! that this library's planner produces it, and that only this library's codec +//! can encode and decode it. + +use std::fmt; +use std::sync::Arc; + +use datafusion::common::Result; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; + +/// Marks a subtree this library claims for remote execution. +#[derive(Debug)] +pub(crate) struct DistributedExec { + input: Arc, + properties: Arc, +} + +impl DistributedExec { + pub(crate) fn new(input: Arc) -> Self { + // The node is pass-through, so it inherits its child's properties + // rather than describing anything of its own. + let properties = Arc::clone(input.properties()); + Self { input, properties } + } +} + +impl DisplayAs for DistributedExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "DistributedExec") + } +} + +impl ExecutionPlan for DistributedExec { + fn name(&self) -> &str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Owns no physical expressions of its own; the child holds them all. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return datafusion::common::internal_err!( + "DistributedExec expects exactly one child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new(children.swap_remove(0)))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs index f491e3114..4e917b236 100644 --- a/examples/datafusion-ffi-query-planner-example/src/extension.rs +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -16,11 +16,11 @@ // under the License. use std::fmt; -use std::sync::atomic::Ordering; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::{Result, TableReference}; +use datafusion::common::{Result, TableReference, internal_err}; use datafusion::datasource::TableProvider; use datafusion::execution::TaskContext; use datafusion::logical_expr::{Extension, LogicalPlan}; @@ -43,16 +43,16 @@ use datafusion_session::QueryPlanner; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict}; +use crate::distributed_exec::DistributedExec; use crate::planner::{DistributedQueryPlanner, PlannerObservations, planner_config_from_options}; /// Values of `ffi_query_planner.max_rows` observed through the task-context /// provider bound at installation time. /// -/// 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. +/// Recorded on every decode call the chain makes to this bundle's codecs, +/// including ones they decline. Reaching the session config from inside a +/// decode callback is what proves the provider bound at installation resolves +/// against the session running the query. type ObservedMaxRows = Arc>>; /// The task-context provider handed to this bundle's components, if it has been @@ -121,12 +121,39 @@ impl LogicalExtensionCodec for ObservingLogicalExtensionCodec { } } -/// Physical companion to [`ObservingLogicalExtensionCodec`]. +/// Carries this library's own [`DistributedExec`] nodes. +/// +/// This is the codec half of the bundle, and the reason the bundle ships both. +/// `DistributedQueryPlanner` emits a `DistributedExec`; no other codec in the +/// session knows the type, so without this one the plans that planner produces +/// cannot be serialized at all. Anything else is declined by delegating to the +/// default codec, so the host's chain falls through to whichever library owns +/// the node. +/// +/// The payload is a marker rather than a serialized node. `DistributedExec` is +/// pass-through and its child arrives already decoded in `inputs`, so there is +/// nothing else to write down; a node with state of its own would encode that +/// state here. struct ObservingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec, observed: ObservedMaxRows, + claims: Arc, } +/// How often this bundle's codec claimed one of its own nodes. +/// +/// Distinct from [`ObservedMaxRows`], which counts every call the chain made, +/// including ones this codec declined. +#[derive(Default, Debug)] +pub(crate) struct DistributedExecClaims { + encoded: AtomicUsize, + decoded: AtomicUsize, +} + +/// Payload written for a [`DistributedExec`]. See +/// [`ObservingPhysicalExtensionCodec`]. +const DISTRIBUTED_EXEC_MARKER: &[u8] = b"datafusion_ffi_query_planner_example:DistributedExec"; + impl fmt::Debug for ObservingPhysicalExtensionCodec { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter @@ -143,7 +170,21 @@ impl PhysicalExtensionCodec for ObservingPhysicalExtensionCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { + // Reading the config through `ctx` is what proves the task-context + // provider bound at installation resolves against the session running + // the query. It happens on the real decode path now, not a synthetic + // one. record_task_ctx(&self.observed, ctx); + if buf == DISTRIBUTED_EXEC_MARKER { + let [input] = inputs else { + return internal_err!( + "DistributedExec expects exactly one input, got {}", + inputs.len() + ); + }; + self.claims.decoded.fetch_add(1, Ordering::SeqCst); + return Ok(Arc::new(DistributedExec::new(Arc::clone(input)))); + } self.inner.try_decode(buf, inputs, ctx, proto_converter) } @@ -153,6 +194,11 @@ impl PhysicalExtensionCodec for ObservingPhysicalExtensionCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { + if node.is::() { + self.claims.encoded.fetch_add(1, Ordering::SeqCst); + buf.extend_from_slice(DISTRIBUTED_EXEC_MARKER); + return Ok(()); + } self.inner.try_encode(node, buf, proto_converter) } } @@ -251,6 +297,7 @@ impl BundledPhysicalCodec { pub(crate) struct MyPlannerExtension { observations: Arc, observed_max_rows: ObservedMaxRows, + claims: Arc, bound_provider: BoundProvider, } @@ -293,8 +340,8 @@ impl MyPlannerExtension { /// `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. + /// One entry per decode call the chain routed to this bundle, so a query + /// whose plan carries a `DistributedExec` leaves several. fn decode_max_rows_seen(&self) -> Vec { self.observed_max_rows .lock() @@ -302,6 +349,24 @@ impl MyPlannerExtension { .unwrap_or_default() } + /// How often this bundle's physical codec encoded one of the + /// `DistributedExec` nodes its own planner produced. + /// + /// Non-zero only when the plan was actually serialized — running a query + /// does not do that, because an FFI planner hands its result back as an + /// opaque plan handle. A distributed engine shipping the plan to a remote + /// executor does, which is the case the pairing exists for; in this + /// repository `ExecutionPlan.to_bytes` stands in for it. + fn distributed_exec_encode_calls(&self) -> usize { + self.claims.encoded.load(Ordering::SeqCst) + } + + /// Companion to [`Self::distributed_exec_encode_calls`], counting the + /// nodes rebuilt on the way back in. + fn distributed_exec_decode_calls(&self) -> usize { + self.claims.decoded.load(Ordering::SeqCst) + } + /// `ffi_query_planner.max_rows` read through the task-context provider /// this bundle was last bound to. /// @@ -351,6 +416,7 @@ impl MyPlannerExtension { Arc::new(ObservingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec {}, observed: Arc::clone(&self.observed_max_rows), + claims: Arc::clone(&self.claims), }); let ffi_physical = FFI_PhysicalExtensionCodec::new(physical, Some(runtime), provider.clone()); diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs index 50cc4831e..30a2d5e4f 100644 --- a/examples/datafusion-ffi-query-planner-example/src/lib.rs +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -22,6 +22,7 @@ use crate::extension::{BundledLogicalCodec, BundledPhysicalCodec, MyPlannerExten use crate::planner::MyQueryPlanner; mod config; +mod distributed_exec; mod extension; mod planner; diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 733536d21..d3587790b 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -40,6 +40,7 @@ use pyo3::prelude::*; use pyo3::types::PyCapsule; use crate::config::MyPlannerConfig; +use crate::distributed_exec::DistributedExec; /// What the planner saw, accumulated across every call rather than reset each /// time. @@ -205,11 +206,14 @@ impl QueryPlanner for DistributedQueryPlanner { .foreign_plan .fetch_or(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst); - Ok(Arc::new(GlobalLimitExec::new( - plan, - 0, - Some(config.max_rows), - ))) + // Wrap the result in a node this library owns. Nothing else in the + // process can serialize a `DistributedExec`, so a session that installs + // this planner without the matching physical codec cannot round-trip + // the plans it produces -- which is the reason the two ship as one + // bundle. See `ObservingPhysicalExtensionCodec`. + Ok(Arc::new(DistributedExec::new(Arc::new( + GlobalLimitExec::new(plan, 0, Some(config.max_rows)), + )))) } } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index f7838030c..48d3e4849 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1848,6 +1848,17 @@ def with_extensions( the planners before it instead of nesting on them, including any the session already had. + Codec order never affects decoding, which routes by codec id. It + affects encoding only when two codecs would claim the same node: the + chain stops at the first that does, so a codec claiming a broad + category can take nodes belonging to a library installed after it. The + query still succeeds, but the plan is written by the wrong library and + may not decode elsewhere. If an extension needs to be early for its + codec and late for its planner, contribute each half at its own + position rather than reordering — the two hooks are independent, so a + small adapter implementing one of them and delegating is enough. The + FFI extensions guide shows the pattern. + Codecs must be handed over as objects exposing the capsule getter, not as bare ``PyCapsule`` objects, and are named after their exporting class as :py:meth:`with_logical_extension_codec` describes. Declare From 54c8863d5891ecadf2c2a6454a9f302c649f86ff Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 14:48:49 -0400 Subject: [PATCH 16/33] Drop the planner example's logical codec observer The physical observer earns its place now that it claims the bundle's own DistributedExec, but the logical one never did and cannot: it declines all four methods to the default codec, and this library defines no logical extension node for it to claim. The FFI logical codec does not carry arbitrary LogicalPlan::Extension nodes anyway, so there is no logical analogue to give it. Measuring a query confirms it: every record_task_ctx firing comes from the physical decode path and none from the logical one. BundledLogicalCodec now wraps DefaultLogicalExtensionCodec directly, which keeps what the logical half actually demonstrated -- a bundle contributing both codec kinds under ids it declares -- and drops 54 lines of trait impl and hand-written Debug that fed an accessor nothing could observe. Also narrows PlannerObservations::used_fallback back to private. It is read only through MyQueryPlanner::used_fallback in the same module; its neighbours need pub(crate) because extension.rs reads them, and it does not. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/extension.rs | 75 +++---------------- .../src/planner.rs | 6 +- 2 files changed, 14 insertions(+), 67 deletions(-) diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs index 4e917b236..c9a42dd6e 100644 --- a/examples/datafusion-ffi-query-planner-example/src/extension.rs +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -19,11 +19,8 @@ use std::fmt; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; -use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::{Result, TableReference, internal_err}; -use datafusion::datasource::TableProvider; +use datafusion::common::{Result, internal_err}; 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; @@ -49,10 +46,10 @@ use crate::planner::{DistributedQueryPlanner, PlannerObservations, planner_confi /// Values of `ffi_query_planner.max_rows` observed through the task-context /// provider bound at installation time. /// -/// Recorded on every decode call the chain makes to this bundle's codecs, -/// including ones they decline. Reaching the session config from inside a -/// decode callback is what proves the provider bound at installation resolves -/// against the session running the query. +/// Recorded on every decode call the chain routes to this bundle's physical +/// codec, including ones it declines. Reaching the session config from inside +/// a decode callback is what proves the provider bound at installation +/// resolves against the session running the query. type ObservedMaxRows = Arc>>; /// The task-context provider handed to this bundle's components, if it has been @@ -68,59 +65,6 @@ fn record_task_ctx(observed: &ObservedMaxRows, ctx: &TaskContext) { } } -/// 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) - } -} - /// Carries this library's own [`DistributedExec`] nodes. /// /// This is the codec half of the bundle, and the reason the bundle ships both. @@ -402,10 +346,11 @@ impl MyPlannerExtension { } let runtime = get_tokio_runtime().handle().clone(); - let logical: Arc = Arc::new(ObservingLogicalExtensionCodec { - inner: DefaultLogicalExtensionCodec {}, - observed: Arc::clone(&self.observed_max_rows), - }); + // Plain default: this library defines no logical extension node, so + // there is nothing for a logical codec of its own to claim. It is still + // contributed so the bundle carries both codec kinds under ids it + // declares -- see `BundledLogicalCodec`. + let logical: Arc = Arc::new(DefaultLogicalExtensionCodec {}); let ffi_logical = FFI_LogicalExtensionCodec::new(logical, Some(runtime.clone()), provider.clone()); // Handed over as an object, not a capsule, so the codec carries an id diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index d3587790b..744b67458 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -59,8 +59,10 @@ pub(crate) struct PlannerObservations { 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. - pub(crate) used_fallback: AtomicBool, + /// Only ever set to `true`, so it is already cumulative. Read only through + /// `MyQueryPlanner::used_fallback` in this module, so unlike its + /// neighbours it needs no wider visibility. + used_fallback: AtomicBool, } impl fmt::Debug for PlannerObservations { From 22d7c8b00d82259f281212180b64f6db3f3d0f1d Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 14:53:05 -0400 Subject: [PATCH 17/33] Export QueryPlannerExportable and drop the empty-extensions guard Two loose ends from review. QueryPlannerExportable was the only member of the extension protocol family left in the submodule while SessionExtensionComponents, SessionExtensionExportable and SessionPlannerExportable were exported from the package root. It types the planner a __datafusion_session_planner__ hook returns, so a bundle author needs it just as much, and one family member importing differently from the rest is a papercut with no upside. Also fixes two doc references that stopped resolving when these classes moved out of context.py: a bare :class:`QueryPlannerExportable` and a bare :py:class:`SessionExtensionComponents`, both now spelled with their module the way the neighbouring datafusion.user_defined references are. with_extensions() with no arguments raised instead of installing nothing. That put it out of family: every sibling varargs method -- DataFrame.select, filter, sort, drop, window -- accepts zero arguments and returns a no-op result, and the two existing "at least one" guards in the codebase both cover cases with no meaningful identity element, which this is not. Installing no extensions has an obvious answer, and a caller assembling the list from a plugin registry should not have to special-case it being empty. Phase two still runs, so the empty case rebinds an existing FFI planner to unchanged chains; a test pins that the planner survives it. Co-Authored-By: Claude Opus 5 (1M context) --- python/datafusion/__init__.py | 2 ++ python/datafusion/context.py | 27 ++++++++++++------------- python/tests/test_context.py | 38 ++++++++++++++++++++++++++++++++--- python/tests/test_imports.py | 18 +++++++++++++++++ 4 files changed, 68 insertions(+), 17 deletions(-) diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9484f2f50..fd6f27b82 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -93,6 +93,7 @@ from .dataframe_formatter import configure_formatter from .expr import Expr, WindowFrame from .extensions import ( + QueryPlannerExportable, SessionExtensionComponents, SessionExtensionExportable, SessionPlannerExportable, @@ -132,6 +133,7 @@ "MetricsSet", "ParquetColumnOptions", "ParquetWriterOptions", + "QueryPlannerExportable", "RecordBatch", "RecordBatchStream", "RuntimeEnvBuilder", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 48d3e4849..a4386567c 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1792,7 +1792,7 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non Args: planner: Object exposing ``__datafusion_query_planner__`` (see - :class:`QueryPlannerExportable`) or a raw + :py:class:`~datafusion.extensions.QueryPlannerExportable`) or a raw ``datafusion_query_planner`` PyCapsule. Examples: @@ -1896,7 +1896,9 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare Args: extensions: Extension bundles to install. Order is irrelevant for codecs and significant for planners, which nest in this order - with the last one outermost. + with the last one outermost. Passing none installs nothing and + returns a handle on this session, so a caller assembling the + list at runtime need not special-case it being empty. Returns: A new context with all extension components installed. @@ -1904,15 +1906,15 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare Raises: TypeError: If an argument implements neither hook, if ``__datafusion_session_extension__`` returns something other - than a :py:class:`SessionExtensionComponents`, or if an - extension contributes a codec as a bare ``PyCapsule``. - ValueError: If no extensions are given, or two codecs claim the - same id. An extension that contributes two instances of one - codec class must declare ``__datafusion_codec_id__`` on at - least one of them; the collision is refused rather than - resolved by position, because a positional id would break - stored plans the first time the extension reordered what it - returns. + than a + :py:class:`~datafusion.extensions.SessionExtensionComponents`, + or if an extension contributes a codec as a bare ``PyCapsule``. + ValueError: If two codecs claim the same id. An extension that + contributes two instances of one codec class must declare + ``__datafusion_codec_id__`` on at least one of them; the + collision is refused rather than resolved by position, because + a positional id would break stored plans the first time the + extension reordered what it returns. Examples: The example is skipped here because it needs a built FFI @@ -1930,9 +1932,6 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare >>> batches[0].column(0).to_pylist() # doctest: +SKIP [1] """ - if not extensions: - msg = "with_extensions requires at least one extension" - raise ValueError(msg) for extension in extensions: if not isinstance( extension, (SessionExtensionExportable, SessionPlannerExportable) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 56b90b307..f836f64e0 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -950,9 +950,41 @@ def __datafusion_session_planner__(self, ctx, fallback): return fallback -def test_with_extensions_requires_an_extension(ctx): - with pytest.raises(ValueError, match="at least one extension"): - ctx.with_extensions() +def test_with_extensions_accepts_no_extensions(ctx): + """No extensions installs nothing and returns a handle on this session. + + A caller assembling the list at runtime — from a plugin registry, say — + should not have to special-case it being empty, and every sibling varargs + method on ``DataFrame`` accepts zero arguments the same way. + """ + ctx.register_record_batches( + "empty_extensions_test", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + result = ctx.with_extensions() + + assert result.session_id() == ctx.session_id() + assert result.table_exist("empty_extensions_test") + assert result.logical_extension_codec_ids() == ctx.logical_extension_codec_ids() + + +def test_with_extensions_no_extensions_keeps_an_installed_planner(ctx): + """The empty case must not disturb a planner the session already has. + + Phase two still runs with nothing to install, which rebinds an existing + FFI planner to the codec chains — unchanged here, so the planner has to + come through intact. + """ + extension = _CodecOnlyExtension() + installed = ctx.with_extensions(extension, _PlannerExtension()) + + result = installed.with_extensions() + + assert result.logical_extension_codec_ids() == ( + installed.logical_extension_codec_ids() + ) + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) def test_with_extensions_rejects_non_extension(ctx): diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py index fea4cc91f..9764e2973 100644 --- a/python/tests/test_imports.py +++ b/python/tests/test_imports.py @@ -94,6 +94,24 @@ def test_datafusion_python_version(): assert datafusion.__version__ is not None +def test_extension_protocols_are_exported_together(): + """The extension protocol family is reachable from the package root. + + ``QueryPlannerExportable`` types the planner a + ``__datafusion_session_planner__`` hook returns, so a bundle author needs + it exactly as much as the other three; leaving it in the submodule made + one member of one family import differently from the rest. + """ + for name in [ + "QueryPlannerExportable", + "SessionExtensionComponents", + "SessionExtensionExportable", + "SessionPlannerExportable", + ]: + assert name in datafusion.__all__, f"{name} missing from datafusion.__all__" + assert getattr(datafusion, name) is getattr(datafusion.extensions, name) + + def test_class_module_is_datafusion(): # context for klass in [ From 3e7f13e873c32ec961618550be1a167f01b46ba0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 14:57:47 -0400 Subject: [PATCH 18/33] Describe both extension hooks in the upgrade guide The planner-install section still said a bundle exposes __datafusion_session_extension__, which stopped being the whole protocol when planners moved to a hook of their own. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/user-guide/upgrade-guides.md | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 4506778e7..4bd42c192 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -85,10 +85,13 @@ 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 +installs every codec before it binds any planner, so a planner cannot end up carrying a chain that a later `with_logical_extension_codec` call has grown. The library exposes a bundle object implementing -`__datafusion_session_extension__`; see the {ref}`ffi` guide. +`__datafusion_session_extension__` for its codecs and +`__datafusion_session_planner__` for its planner — the latter is handed the +planner installed so far, so several libraries that each ship one nest instead +of displacing each other. See the {ref}`ffi` guide. ### Mismatched extension libraries now fail loudly From 734673ab6281d7513774bfa01f8442c7fda537ce Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 15:19:41 -0400 Subject: [PATCH 19/33] Generate heading anchors for h4 so the FFI guide's links resolve `myst_heading_anchors` was 3, but the extension-bundles section added in this branch cross-references its own `####` subsections. Sphinx warns `'myst' cross-reference target not found: 'when-codec-order-does-matter'` and renders that link as plain text; the docs build does not pass `-W`, so it went unnoticed. Bumping to 4 rather than promoting the heading keeps the four subsections nested under `### Extension bundles: with_extensions`, where they belong. Only one other `####` heading exists under `docs/source/`. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/conf.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 22bace809..8b3dbd0dd 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -199,8 +199,10 @@ def setup(sphinx) -> None: "**": ["sidebar-globaltoc.html"], } -# tell myst_parser to auto-generate anchor links for headers h1, h2, h3 -myst_heading_anchors = 3 +# tell myst_parser to auto-generate anchor links for headers h1 through h4. +# h4 is included because the FFI guide cross-references its own `####` +# subsections; without an anchor those links render but resolve nowhere. +myst_heading_anchors = 4 # MyST extensions: # - tasklist: GitHub-style `- [x]` checkboxes From f57feec079ef5a9a2995e2173bb5933711f14df1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 15:22:38 -0400 Subject: [PATCH 20/33] Skip the planner commit when with_extensions installs nothing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `with_extensions` ended every call with `_install_extension_planner`, which rebuilds `SessionState` to rebind an existing FFI planner to this handle's codec chains. When the call installed no codec there is nothing to rebind against, so the rebuild is at best churn — and at worst it drags a planner that is sitting on another handle's codecs onto this one's, silently undoing that install. `with_python_udf_inlining` already guards its no-op toggle for exactly this reason; `with_extensions` now guards the same way. `test_with_extensions_installing_nothing_leaves_the_planner_alone` covers both shapes of "installed nothing": no arguments at all, and a bundle whose hooks answer empty. Both fail without the guard, with the planner left on an empty logical chain and the query erroring out. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 6 +++ .../_test_three_library_query_planner.py | 51 +++++++++++++++++++ python/datafusion/context.py | 19 +++++-- 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 4224d1574..8b11709e2 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1708,6 +1708,12 @@ impl PySessionContext { /// whichever planner the session already holds against the new chains, /// exactly as `with_logical_extension_codec` does, and writes nothing at /// all if the session has no FFI planner to rebuild. + /// + /// The caller skips this step entirely when the call installed no codec + /// and no planner, the same way [`Self::with_python_udf_inlining`] returns + /// early for a no-op toggle: there is nothing to rebind against, and the + /// rebuild would drag a planner sitting on another handle's codecs onto + /// this one's. #[pyo3(signature = (planner=None))] pub fn _install_extension_planner<'py>( slf: &Bound<'py, Self>, 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 891c91296..80ba7c724 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -1203,6 +1203,57 @@ def test_with_extensions_rebinds_existing_planner(): assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 +class NoOpExtension: + """A bundle that turns out to contribute nothing. + + A plugin that finds no work to do -- an engine pointed at no scheduler, a + codec pack for a feature the session did not enable -- still gets listed, + and both its hooks answer empty rather than the caller having to filter it + out. + """ + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents() + + def __datafusion_session_planner__( + self, ctx: SessionContext, fallback: object + ) -> object: + return None + + +@pytest.mark.parametrize("bundles", [(), (NoOpExtension(),)], ids=["empty", "no_op"]) +def test_with_extensions_installing_nothing_leaves_the_planner_alone(bundles): + """A call that installs nothing must not rebind the session's planner. + + The sibling of ``test_an_unchanged_inlining_setting_leaves_the_planner_alone``, + and the same hazard: committing the planner rebuilds ``SessionState`` to + rebind an existing FFI planner to *this handle's* chains. When no codec was + installed there is nothing to rebind against, so the rebuild buys nothing + and can only do harm. + + Observable only once the planner holds some *other* handle's codec, which + is what the discarded ``with_logical_extension_codec`` below arranges. + Without the guard the no-op call drags the planner back onto ``ctx``'s + codecs -- and ``ctx`` has no logical codec, so the planner is left with an + empty chain and the query fails outright instead of quietly using the wrong + codec. + """ + ctx, _physical_codec = physical_only_context() + ctx.set_query_planner(MyQueryPlanner()) + + planner_codec = MyLogicalExtensionCodec() + ctx.with_logical_extension_codec(planner_codec) # discarded; planner keeps it + gc.collect() + + ctx.with_extensions(*bundles) + gc.collect() + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner_codec.table_provider_encode_calls() > 0 + + def test_with_extensions_rejects_two_bundles_of_the_same_codec_class(): """Two bundles contributing the same codec class collide on id. diff --git a/python/datafusion/context.py b/python/datafusion/context.py index a4386567c..8d26ba181 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1842,10 +1842,11 @@ def with_extensions( a nested planner is never left encoding through a chain a later extension has grown. - If no extension supplies a planner, an existing FFI planner is rebound - to the final codec chains and the session's planner is otherwise left - alone. An extension that ignores the ``fallback`` it is handed replaces - the planners before it instead of nesting on them, including any the + If no extension supplies a planner but codecs were installed, an + existing FFI planner is rebound to the final chains; if the call + installed nothing at all, the session's planner is not touched. An + extension that ignores the ``fallback`` it is handed replaces the + planners before it instead of nesting on them, including any the session already had. Codec order never affects decoding, which routes by codec id. It @@ -1987,7 +1988,15 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare continue planner = new.ctx._rebind_query_planner(supplied) - new.ctx._install_extension_planner(planner) + # Rebinding the session's planner is a side effect on state shared with + # every other handle, so do not pay it for a call that installs nothing + # -- the same guard `with_python_udf_inlining` carries. With no codec + # installed the chains the planner would be rebuilt against are the ones + # it already holds, so the rebuild is unobservable except in the one case + # where it does harm: a planner sitting on some *other* handle's codecs + # gets dragged onto this handle's, silently undoing that install. + if planner is not None or logical_codecs or physical_codecs: + new.ctx._install_extension_planner(planner) return new def table_provider(self, name: str) -> Table: From 7b6193e2cadc1057639c4fc98b180771518cb29e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 15:30:31 -0400 Subject: [PATCH 21/33] Reject a lone codec where SessionExtensionComponents wants an iterable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `logical_extension_codecs=codec` instead of `(codec,)` is the easy mistake to make, and it surfaced as `'MyCodec' object is not iterable` raised by an `extend` call inside `with_extensions` — naming neither the field nor the hook that built the value. `__post_init__` now checks it, so the error lands in the extension library's own frame and says which field is wrong and how to spell one codec. It also normalizes each field to a tuple. The declared type is a tuple and the class is frozen, so a list left in place would be a mutable member of an immutable value, and a generator would be exhausted by the first read. A str is refused rather than normalized: it is iterable, so it would otherwise become a tuple of characters and fail much later as that many bogus codecs. Co-Authored-By: Claude Opus 5 (1M context) --- python/datafusion/extensions.py | 54 +++++++++++++++++++++++++++++++++ python/tests/test_context.py | 31 +++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index e9f2f7156..c2faf9e92 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -81,6 +81,15 @@ class QueryPlannerExportable(Protocol): def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 +def _not_a_codec_iterable(field: str, value: object) -> str: + """Message for a codec field that cannot be read as a collection.""" + return ( + f"{field} must be an iterable of codec objects, not a single " + f"{type(value).__name__}. A lone codec is written as a one-element " + f"tuple — {field}=(codec,) — and the trailing comma is what makes it one." + ) + + @dataclass(frozen=True) class SessionExtensionComponents: """Components an extension contributes to a session context. @@ -134,6 +143,23 @@ class SessionExtensionComponents: 'my_library.v1' >>> components.physical_extension_codecs () + + Any iterable is accepted and stored as a tuple, so a bundle that builds + its codecs with a list comprehension does not have to convert: + + >>> components = SessionExtensionComponents( + ... logical_extension_codecs=[NamedCodec(capsule)] + ... ) + >>> type(components.logical_extension_codecs).__name__ + 'tuple' + + A single codec is not an iterable of codecs, and forgetting the + trailing comma is the easy way to write one by accident: + + >>> SessionExtensionComponents(logical_extension_codecs=NamedCodec(capsule)) + Traceback (most recent call last): + ... + TypeError: logical_extension_codecs must be an iterable of codec objects... """ logical_extension_codecs: tuple[LogicalExtensionCodecExportable, ...] = () @@ -142,6 +168,34 @@ class SessionExtensionComponents: physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = () """Physical codecs to add to the session's codec chain, in declaration order.""" + def __post_init__(self) -> None: + """Normalize each field to a tuple, rejecting what cannot become one. + + A bundle that writes ``logical_extension_codecs=codec`` instead of + ``(codec,)`` is contributing one codec, not an iterable of them. + Without this, the mistake surfaces inside + :py:meth:`~datafusion.context.SessionContext.with_extensions` as + ``'MyCodec' object is not iterable``, which names neither the field + nor the hook that built it. Checking here puts the error in the + extension library's own frame. + + Normalizing is worth doing on its own: the declared type is a tuple + and the class is frozen, so a list left in place would be a mutable + member of an immutable value, and a generator would be exhausted by + the first read. + """ + for name in ("logical_extension_codecs", "physical_extension_codecs"): + value = getattr(self, name) + # A str is iterable, so it would otherwise normalize into a tuple + # of characters and fail much later as that many bogus codecs. + if isinstance(value, (str, bytes)): + raise TypeError(_not_a_codec_iterable(name, value)) + try: + codecs = tuple(value) + except TypeError: + raise TypeError(_not_a_codec_iterable(name, value)) from None + object.__setattr__(self, name, codecs) + @runtime_checkable class SessionExtensionExportable(Protocol): diff --git a/python/tests/test_context.py b/python/tests/test_context.py index f836f64e0..830ab3c97 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1001,6 +1001,37 @@ def __datafusion_session_extension__(self, ctx): ctx.with_extensions(BadExtension()) +@pytest.mark.parametrize( + "field", ["logical_extension_codecs", "physical_extension_codecs"] +) +def test_session_extension_components_rejects_a_single_codec(field): + """A lone codec is not an iterable of codecs. + + Dropping the trailing comma is the easy way to write one by accident. The + check lives on the value type so the error lands in the extension + library's own frame, naming the field it got wrong, rather than surfacing + later inside ``with_extensions`` as ``'_NamedCodec' object is not + iterable``. + """ + codec = _NamedCodec( + SessionContext().__datafusion_logical_extension_codec__(), + "my_library.logical", + ) + + with pytest.raises(TypeError, match=r"must be an iterable of codec objects"): + SessionExtensionComponents(**{field: codec}) + + +def test_session_extension_components_rejects_a_string(): + """A str is iterable, so it needs refusing on its own. + + Left alone it would normalize into a tuple of characters and fail much + later as that many bogus codecs. + """ + with pytest.raises(TypeError, match=r"not a single str"): + SessionExtensionComponents(logical_extension_codecs="my_library.logical") + + def test_with_extensions_accepts_a_planner_only_extension(ctx): """An extension may implement the planner hook alone. From 0868f391a903f375842b30f783bd564165b4535e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 15:50:05 -0400 Subject: [PATCH 22/33] Widen the with_extensions annotation to both hooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `with_extensions` accepts a bundle implementing either hook — the runtime check tests against both protocols, `test_with_extensions_accepts_a_planner_only_extension` pins it, and the FFI guide's `PlannerOf` adapter recommends contributing only the planner half. The annotation named `SessionExtensionExportable` alone, so a type checker rejected the very shape the guide tells authors to write. Two doc comments also went stale. `test_with_extensions_no_extensions_keeps_an_installed_planner` still described phase two running and rebinding an existing planner, which the no-op guard now skips outright, and `__datafusion_codec_id__` listed a `_install_extensions` method that never existed under that name. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 10 +++++----- python/datafusion/context.py | 2 +- python/tests/test_context.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 8b11709e2..4300c1f57 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1435,11 +1435,11 @@ impl PySessionContext { /// decode. See [`SESSION_CODEC_ID_PREFIX`]. /// /// Handles derived from one session — `with_python_udf_inlining`, - /// `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 + /// `with_logical_extension_codec`, [`Self::_install_extension_codecs`] — + /// report the same id even though their codec chains differ, so installing + /// two of them on one target is refused. That is the intended answer: they + /// share a `state_ref`, so their payloads would resolve against the same + /// session and are indistinguishable on decode. Every derivation shares the /// session for exactly this reason; `enable_url_table` is the one that does /// not, and it is tracked as a bug. #[getter] diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 8d26ba181..7da8ec3bb 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1812,7 +1812,7 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non self.ctx.set_query_planner(planner) def with_extensions( - self, *extensions: SessionExtensionExportable + self, *extensions: SessionExtensionExportable | SessionPlannerExportable ) -> SessionContext: """Create a new session context with the given extension bundles. diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 830ab3c97..c67c52870 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -971,9 +971,9 @@ def test_with_extensions_accepts_no_extensions(ctx): def test_with_extensions_no_extensions_keeps_an_installed_planner(ctx): """The empty case must not disturb a planner the session already has. - Phase two still runs with nothing to install, which rebinds an existing - FFI planner to the codec chains — unchanged here, so the planner has to - come through intact. + A call that installs no codec and no planner skips the planner commit + entirely, so the installed planner keeps the chains it was bound to and the + session still plans through it. """ extension = _CodecOnlyExtension() installed = ctx.with_extensions(extension, _PlannerExtension()) From f7c448e7d498f19e876f577bbd101b51ef28814f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 15:52:18 -0400 Subject: [PATCH 23/33] Correct three claims in the extension bundle docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `SessionExtensionComponents` doctest took its codec capsule off a `SessionContext()` that was dropped on the same line. An FFI codec holds its task-context provider weakly, so that capsule names a session that is already gone — the doctest only reads an id back so it passes, but it is the exact shape the FFI guide warns against. It now keeps the context in a name. `SessionPlannerExportable` called returning `fallback` a wrap that "contributes nothing". It is not: the capsule the first bundle receives wraps the session's planner for export, so handing it back installs it as a foreign planner and every later plan crosses an FFI boundary that was not there before. `None` is the no-op. Corrected in the protocol docstring, the FFI guide's canonical section, and the `_PlannerExtension` test helper that repeated the claim. `__post_init__` walked a written-out list of field names. It now walks `dataclasses.fields`, filtered on the `_codecs` suffix so a codec field added later is normalized without anyone remembering to name it, and a future field that is not a codec collection is left alone. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/contributor-guide/ffi.md | 5 ++++ python/datafusion/extensions.py | 34 ++++++++++++++++++++++++---- python/tests/test_context.py | 8 ++++--- 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 54ef53b72..c663cf823 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -412,6 +412,11 @@ library that must be the only planner does it deliberately — but it is not com and nothing detects it. Returning `None` contributes no planner and leaves `fallback` in place. +`None` is the no-op, not `fallback`. The capsule handed to the first bundle wraps the +session's planner for export, so returning it unchanged installs that planner as a +foreign one and every plan built afterwards crosses an FFI boundary it did not before. +A bundle that decides at runtime it has nothing to contribute returns `None`. + Three libraries that each ship a planner therefore install like this, with the outermost last: diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index c2faf9e92..01012d128 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -48,7 +48,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: @@ -135,7 +135,15 @@ class SessionExtensionComponents: ... ... def __datafusion_logical_extension_codec__(self, session=None): ... return self._capsule - >>> capsule = SessionContext().__datafusion_logical_extension_codec__() + + The context stays in scope for as long as the codec does. An + ``FFI_LogicalExtensionCodec`` holds its task-context provider *weakly*, + so a capsule taken off a throwaway ``SessionContext()`` names a session + that is already gone and fails on first use with ``TaskContextProvider + went out of scope over FFI boundary``: + + >>> ctx = SessionContext() + >>> capsule = ctx.__datafusion_logical_extension_codec__() >>> components = SessionExtensionComponents( ... logical_extension_codecs=(NamedCodec(capsule),) ... ) @@ -183,8 +191,17 @@ def __post_init__(self) -> None: and the class is frozen, so a list left in place would be a mutable member of an immutable value, and a generator would be exhausted by the first read. + + Driven off :py:func:`dataclasses.fields` rather than a written-out + list, so a codec field added later is normalized without anyone + remembering to name it here. The ``_codecs`` suffix is what marks a + field as one of them, leaving room for a future field that is not a + codec collection and must not be turned into a tuple. """ - for name in ("logical_extension_codecs", "physical_extension_codecs"): + for field in fields(self): + name = field.name + if not name.endswith("_codecs"): + continue value = getattr(self, name) # A str is iterable, so it would otherwise normalize into a tuple # of characters and fail much later as that many bogus codecs. @@ -257,6 +274,12 @@ class SessionPlannerExportable(Protocol): not left encoding through a chain that a later bundle has grown. Return ``None`` to contribute no planner and leave ``fallback`` in place. + That is the no-op, and it is not the same as returning ``fallback``: the + capsule the first bundle receives wraps the session's planner for export, so + handing it back installs it as a foreign planner and every later plan crosses + an FFI boundary that was not there before. A bundle with nothing to + contribute returns ``None``. + Ignoring ``fallback`` and returning a planner that does not delegate to it is legal and means "replace" — but it discards every planner listed before this one, including any the session already had. @@ -274,8 +297,9 @@ class SessionPlannerExportable(Protocol): ... def __datafusion_session_planner__(self, ctx, fallback): ... # A real library returns its own planner wrapping ... # `fallback`, e.g. ``my_library.Planner(fallback=fallback)``. - ... # Handing it straight back is the degenerate wrap: valid, - ... # and contributes nothing. + ... # Handing it straight back is the degenerate wrap: legal, + ... # but it still installs `fallback` as a foreign planner. + ... # Return None instead to contribute nothing. ... return fallback >>> isinstance(MyEngineExtension(), SessionPlannerExportable) True diff --git a/python/tests/test_context.py b/python/tests/test_context.py index c67c52870..df751d054 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -935,9 +935,11 @@ def __datafusion_session_extension__(self, ctx): class _PlannerExtension: """Contributes a planner, recording the fallback it was handed. - Passing ``fallback`` straight back through is the degenerate wrap: the - resulting session plans exactly as it did before, which is what lets a - pure-Python test assert the threading without a real layering planner. + Passing ``fallback`` straight back through is the degenerate wrap: it plans + the same queries to the same plans, which is what lets a pure-Python test + assert the threading without a real layering planner. It is not a no-op — + the capsule gets installed, so the session ends up planning through a + foreign planner — but nothing here depends on that either way. """ def __init__(self): From 0712eff15469f15656684ca9982f4fe4be1a50a2 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 15:53:43 -0400 Subject: [PATCH 24/33] Rename _rebind_query_planner to _export_query_planner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The method rebinds nothing. It imports whatever a `__datafusion_session_planner__` hook returned — an object exposing the getter or a raw capsule — and hands back a capsule, so the next hook receives one either way; its own doc comment already said "re-export". Meanwhile "rebind" means something specific and different in this file: rebuilding an installed planner against a handle's codec chains, which is what `set_session_query_planner` does. Freeing the word keeps the two apart. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 2 +- python/datafusion/context.py | 2 +- python/tests/test_wrapper_coverage.py | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 4300c1f57..c711a62dc 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1688,7 +1688,7 @@ impl PySessionContext { /// not have to branch on which. Importing here also surfaces a malformed /// planner at the hook that produced it rather than at the final install. /// Writes nothing. - pub fn _rebind_query_planner<'py>( + pub fn _export_query_planner<'py>( slf: &Bound<'py, Self>, planner: Bound<'py, PyAny>, ) -> PyDataFusionResult> { diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 7da8ec3bb..58aa8938e 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1986,7 +1986,7 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare supplied = extension.__datafusion_session_planner__(new, fallback) if supplied is None: continue - planner = new.ctx._rebind_query_planner(supplied) + planner = new.ctx._export_query_planner(supplied) # Rebinding the session's planner is a side effect on state shared with # every other handle, so do not pay it for a call that installs nothing diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index dc6dfea11..6927632b9 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -33,11 +33,11 @@ # gap in coverage. PRIVATE_SUPPORT_METHODS = frozenset( { - # The three phases of SessionContext.with_extensions: install the - # codecs, normalize each planner hook's return value, commit the - # planner. + # The three steps of SessionContext.with_extensions: install the + # codecs, re-export each planner hook's return value as a capsule, + # commit the planner. "_install_extension_codecs", - "_rebind_query_planner", + "_export_query_planner", "_install_extension_planner", } ) From d13787b4edd25d30b0ba7f2f4b82ac509cc642f2 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 15:54:38 -0400 Subject: [PATCH 25/33] Say which codec chains each extension hook's context carries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `__datafusion_session_planner__` was documented as receiving a context that carries the final codec chains, and `MyPlannerExtension` relies on exactly that when it takes the host's codecs off `ctx` instead of minting its own. The other side was never stated: `__datafusion_session_extension__` runs before anything is installed, so its `ctx` is the same session with the chains the receiver already had — missing this call's codecs, including the bundle's own. Both hooks hand back a valid task-context provider, which is what components actually need, so the difference only bites an author who reads codec chains off the context. Recorded on `SessionExtensionExportable`, in the FFI guide's two-phase section, and in Rule 2 of the capsule-protocol skill. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 6 ++++++ docs/source/contributor-guide/ffi.md | 9 +++++++++ python/datafusion/extensions.py | 9 +++++++++ 3 files changed, 24 insertions(+) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 3dfaf353f..1efc93ec5 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -70,6 +70,12 @@ Wrap `fallback` and delegate to it; returning a planner that ignores it discards every layer beneath, including one the session already had. It runs after every bundle's codecs are installed, so `ctx` carries the final chains. +That is also the only hook where it does. `__datafusion_session_extension__` +runs before anything is installed, so its `ctx` still carries the chains the +receiver had — the same session, and the same task-context provider, but not +this call's codecs, not even your own. Read the host's codec chains in the +planner hook, never in the extension hook. + A *codec* must always be handed over as an object implementing its getter, never as the bare capsule the getter returns; `with_extensions` refuses a capsule. A codec's wire id — the string a payload names on decode, which has to mean the diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index c663cf823..cc1c5618d 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -406,6 +406,15 @@ handing each the planner built so far. Two consequences worth holding onto: deep](#rebinding-a-planners-codecs-is-one-level-deep)), so a fallback captured before the codecs were complete would stay stale forever. +The two hooks therefore see the same session through different chains. Both receive a +handle on the one session, so the task-context provider taken off either is the same and +stays valid — but the `ctx` in phase one still carries the chains the receiver had, since +nothing is installed yet, while the `ctx` in phase two carries every bundle's codecs. A +bundle that reads the host's codec chains — `MyPlannerExtension` does, to give its +planner the host's codecs rather than minting its own — must do that in the planner hook. +Reading them in phase one gets the chains from before the call, missing even the bundle's +own codecs. + An extension that ignores `fallback` and returns an unrelated planner replaces every layer beneath it, including any planner the session already had. That is legal — a library that must be the only planner does it deliberately — but it is not composable, diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 01012d128..7823aa427 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -231,6 +231,15 @@ class SessionExtensionExportable(Protocol): mutating the context they are handed — a registration made during binding is not rolled back if a later extension fails. + ``ctx`` is the right session but not yet the final codec chains: this hook + runs before anything is installed, so ``ctx`` still carries whatever chains + the receiver had. Take the task-context provider off it — that is bound to + the session and is what the components need — but do not read its codec + chains expecting to find this call's codecs, including your own. + :py:class:`SessionPlannerExportable` is the hook that sees the completed + chains, which is why a planner that wraps the host's codecs builds them + there rather than here. + A bundle that also contributes a query planner implements :py:class:`SessionPlannerExportable` alongside this protocol. Planners are installed in a second phase, so they are not part of the components From 07792ad275d28002d0e2a1c5cab7ca37842d8968 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 16:30:42 -0400 Subject: [PATCH 26/33] Make the planner-hook tests assert what their names claim `test_with_extensions_threads_the_planner_through_in_order` never checked an order: it asserted each hook recorded one fallback and that the second was not None, both of which hold for a host that ran the hooks backwards. `test_with_extensions_skips_a_planner_hook_returning_none` asserted only that downstream ran, while its comment claimed the skipped hook had not become downstream's fallback. `_PlannerExtension` now takes an optional shared list the hooks append themselves to, so order is observable. The threading test asserts that list, plus that the second hook's fallback is not the object the first was handed -- the host re-exports every return value before passing it on. A capsule is opaque from Python, so that cannot separate a re-export of the first planner from a fresh read of the session's; the comment says so and points at the FFI suite's `test_with_extensions_nests_planners_in_argument_order`, which pins the nesting by asserting the outer planner delegated. The skip test records the skipped hook's fallback too, and asserts both hooks ran, that downstream was handed a different capsule, and that the resulting context still queries. Each new assertion was mutation-tested: reversing the planner loop, dropping the `if supplied is None: continue`, and replacing the re-export with a straight pass-through each fail exactly one of these tests. Co-Authored-By: Claude Opus 5 (1M context) --- python/tests/test_context.py | 60 ++++++++++++++++++++++++++++-------- 1 file changed, 47 insertions(+), 13 deletions(-) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index df751d054..c1aa19a4e 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -942,11 +942,15 @@ class _PlannerExtension: foreign planner — but nothing here depends on that either way. """ - def __init__(self): + def __init__(self, calls=None): self.fallbacks = [] self.planner_ctx = None + # Shared list the hooks append themselves to, so a test can assert the + # order they ran in rather than only that each ran. + self.calls = [] if calls is None else calls def __datafusion_session_planner__(self, ctx, fallback): + self.calls.append(self) self.planner_ctx = ctx self.fallbacks.append(fallback) return fallback @@ -1053,15 +1057,25 @@ def test_with_extensions_threads_the_planner_through_in_order(ctx): Planners nest rather than chain, so the host hands each bundle the planner built so far. Argument order is nesting order, last one outermost. """ - first, second = _PlannerExtension(), _PlannerExtension() + calls = [] + first, second = _PlannerExtension(calls), _PlannerExtension(calls) ctx.with_extensions(first, second) - assert len(first.fallbacks) == 1 - assert len(second.fallbacks) == 1 - # `first` returned its fallback unchanged, and the host normalizes each - # hook's return value before passing it on, so `second` sees a capsule - # standing for the same planner rather than the session's original. - assert second.fallbacks[0] is not None + # Argument order, once each. Nothing else pins the order: both hooks + # return capsules, and a host that ran them backwards would still leave + # each with one fallback recorded. + assert calls == [first, second] + + # `first` returned its fallback unchanged, but the host re-exports every + # hook's return value before handing it on, so `second` receives a capsule + # of its own rather than the object `first` was handed. + assert second.fallbacks[0] is not first.fallbacks[0] + + # That is as far as pure Python reaches: a capsule is opaque, so this + # cannot tell a re-export of `first`'s planner from a fresh read of the + # session's. `test_with_extensions_nests_planners_in_argument_order` in + # examples/datafusion-ffi-query-planner-example is what pins the nesting, + # by asserting the outer planner delegated to the inner one. def test_with_extensions_planner_hook_sees_the_new_handle(ctx): @@ -1079,18 +1093,38 @@ def test_with_extensions_planner_hook_sees_the_new_handle(ctx): def test_with_extensions_skips_a_planner_hook_returning_none(ctx): - """Returning ``None`` contributes no planner and keeps the fallback.""" + """Returning ``None`` contributes no planner and keeps the fallback. + + The skip is what lets the call succeed at all: a host that treated the + ``None`` as a contribution would hand it to the export step and fail with + ``'None' is not an instance of 'PyCapsule'`` before ``downstream`` ran. + """ class NoPlanner: + def __init__(self): + self.fallbacks = [] + def __datafusion_session_planner__(self, ctx, fallback): - return None + self.fallbacks.append(fallback) + # Spelled out rather than left to fall off the end: `None` is the + # protocol's "contribute no planner", which is what this test is + # about, and an implicit one would read as an oversight. + return None # noqa: RET501, PLR1711 + skipped = NoPlanner() downstream = _PlannerExtension() - ctx.with_extensions(NoPlanner(), downstream) + result = ctx.with_extensions(skipped, downstream) - # The skipped hook did not become `downstream`'s fallback; it got the - # session's own planner instead. + # Both hooks ran, and `downstream` was handed a planner rather than the + # `None` in front of it. It is a fresh read of the session's planner, not + # the object `skipped` was given, so a host that fell back by reusing the + # previous hook's *input* is ruled out too. + assert len(skipped.fallbacks) == 1 assert len(downstream.fallbacks) == 1 + assert downstream.fallbacks[0] is not skipped.fallbacks[0] + + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) def test_with_extensions_rejects_bad_codec_capsule(ctx): From 4296c5ee515e128db38b97b761508ce5f5774959 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 16:30:52 -0400 Subject: [PATCH 27/33] Correct three more claims in the extension bundle docs `QueryPlannerExportable` said `session` is the `datafusion.context.SessionContext` the planner is being installed on. It is not. The capsule getters are called from Rust and receive the PyO3 context, so `isinstance(session, SessionContext)` is False -- while its repr reads `datafusion.SessionContext`, because the pyclass declares `module = "datafusion"`. It carries every capsule getter and `__datafusion_codec_id__`, which is all the protocol needs, so the fix is to say duck-type it rather than to change what is passed. The two bundle hooks are the exception and do receive the wrapper, since `with_extensions` dispatches them from Python; the ffi.md section on capsule getters now draws the same distinction. The `with_extensions` `Raises:` section listed ValueError for colliding codec ids only. A getter returning a capsule of the wrong kind also raises it -- `Expected name 'datafusion_query_planner' in PyCapsule, instead got 'datafusion_logical_extension_codec'` -- which `test_with_extensions_rejects_bad_codec_capsule` already pins. The `datafusion.extensions` module docstring said phase two runs the planner hook "once per bundle". Once per bundle that implements it; a bundle implements either hook or both. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/contributor-guide/ffi.md | 8 ++++++++ python/datafusion/context.py | 6 +++++- python/datafusion/extensions.py | 25 ++++++++++++++++++------- 3 files changed, 31 insertions(+), 8 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index cc1c5618d..244505c2c 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -602,6 +602,14 @@ the session, and never touches a provider directly. That also matches what insta does anyway: `set_query_planner` builds the planner against the codecs of the session that will run the query. +Duck-type `session`, and do not check its type. It is the PyO3 context the binding +installs through, not the `datafusion.context.SessionContext` wrapper, so it carries +every capsule getter and `__datafusion_codec_id__` — everything the protocol asks of it +— but `isinstance(session, SessionContext)` is `False` in Python even though its `repr` +reads `datafusion.SessionContext`. The two bundle hooks +`__datafusion_session_extension__` and `__datafusion_session_planner__` are the +exception: `with_extensions` dispatches them from Python and hands them the wrapper. + `SessionContext` accepts the argument on all three getters and ignores it, so a session satisfies the same protocol an extension library implements. When you export the current planner to wrap it, `ctx.__datafusion_query_planner__()` and diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 58aa8938e..69339a33f 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1915,7 +1915,11 @@ class as :py:meth:`with_logical_extension_codec` describes. 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. + extension reordered what it returns. Also if a capsule getter + returns a capsule of the wrong kind — a physical codec handed + over under ``__datafusion_logical_extension_codec__``, say — + which is reported against the name the getter should have + produced. Examples: The example is skipped here because it needs a built FFI diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 7823aa427..631fe239b 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -37,10 +37,11 @@ session's codec chain holds many codecs and dispatches between them by id, so codecs merely accumulate and their order does not affect decoding. A session holds exactly *one* query planner, so planners compose by nesting: each wraps -the one before it. Phase one collects every bundle's codecs through +the one before it. Phase one collects the codecs of every bundle implementing :py:class:`SessionExtensionExportable` and installs them; phase two runs -:py:class:`SessionPlannerExportable` once per bundle, in argument order, -handing each the planner built so far. +:py:class:`SessionPlannerExportable` once for each bundle that implements it, +in argument order, handing each the planner built so far. A bundle implements +either hook or both, and one it does not implement is simply not called. That split is what lets several libraries that each ship a planner coexist. It also means bundle order is significant for planners and irrelevant for codecs. @@ -72,10 +73,20 @@ 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. + produced by a separate compiled extension. ``session`` is a handle on the + session the planner is being installed on; take the extension codecs from + it rather than building your own. + + Duck-type that handle rather than checking its type. It is the PyO3 + context from ``datafusion._internal``, not the + :py:class:`~datafusion.context.SessionContext` wrapper, so it exposes every + capsule getter and ``__datafusion_codec_id__`` — which is all the protocol + asks of it — but ``isinstance(session, SessionContext)`` is ``False`` even + though its ``repr`` reads ``datafusion.SessionContext``. The same is true + of the codec getters in :py:mod:`datafusion.user_defined`. The two bundle + hooks are the exception: :py:class:`SessionExtensionExportable` and + :py:class:`SessionPlannerExportable` are dispatched from Python and receive + the wrapper. """ def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 From 781e980acbb479e24f9c392ddaaa7460ee122404 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 08:20:51 -0400 Subject: [PATCH 28/33] docs: add a user-facing extensions page and split distributing-work MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The user guide had no page for someone who installs an extension library and wants to run queries with it. `with_extensions` was described only in a 123-line docstring and in the contributor-guide FFI page, which opens by explaining that Rust has no stable ABI — the wrong altitude for a reader who just wants their queries to run somewhere else. Adds `user-guide/extensions.md`: what an extension library is, which kinds register directly versus needing `with_extensions`, the two failure modes that actually bite (a collected context, a version mismatch), and how to check what a session was taught. It names no capsule, no ABI, and no codec id; the only occurrences of `FFI` and `TaskContextProvider` are inside the error string a reader would be searching for. Splits `distributing-work.md` into a directory. The page was entirely about pickling expressions to worker pools and treated query-level distribution as two stubs at the end, so nothing on the site connected `with_extensions` to distribution at all — which is the road a data scientist is actually looking for. The index now asks who owns the partitioning decision and routes accordingly; `query-engines.md` carries the missing bridge and absorbs the two upstream work-in-progress sections. Wires `sphinx-reredirects`, which has been a declared dependency since #1578 without ever being enabled, so the old `distributing-work.html` URL keeps working. Also fixes three pointers that went stale in the MyST migration and still named `.rst` files, a malformed `ref:` role in udf-and-udfa.md that rendered as literal text, and a doubled "the" in data-sources.md. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/codec.rs | 2 +- docs/source/conf.py | 12 ++ .../common-operations/udf-and-udfa.md | 3 +- docs/source/user-guide/data-sources.md | 2 +- .../expressions.md} | 65 +++----- .../user-guide/distributing-work/index.md | 81 ++++++++++ .../distributing-work/query-engines.md | 100 ++++++++++++ docs/source/user-guide/extensions.md | 142 ++++++++++++++++++ docs/source/user-guide/index.md | 9 +- docs/source/user-guide/upgrade-guides.md | 2 + examples/multiprocessing_pickle_expr.py | 2 +- examples/ray_pickle_expr.py | 2 +- 12 files changed, 370 insertions(+), 52 deletions(-) rename docs/source/user-guide/{distributing-work.md => distributing-work/expressions.md} (86%) create mode 100644 docs/source/user-guide/distributing-work/index.md create mode 100644 docs/source/user-guide/distributing-work/query-engines.md create mode 100644 docs/source/user-guide/extensions.md diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 8f43bdb5e..cca45147d 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -660,7 +660,7 @@ impl PythonLogicalCodec { /// `cloudpickle.loads` on the inline `DFPY*` payload. It does /// **not** make `pickle.loads(untrusted_bytes)` safe; treat every /// `pickle.loads` on untrusted input as unsafe regardless of this - /// setting. See `docs/source/user-guide/io/distributing_work.rst` + /// setting. See `docs/source/user-guide/distributing-work/expressions.md` /// (Security section) for the full threat model, and Python's /// [pickle module security warning][1] for why `pickle.loads` is /// unsafe in general. diff --git a/docs/source/conf.py b/docs/source/conf.py index 8b3dbd0dd..7f9c4b89a 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -54,8 +54,20 @@ # raises an extension conflict. "myst_nb", "autoapi.extension", + # Emits a meta-refresh stub at each old docname listed in `redirects` + # below, so inbound links to pages that have moved keep working. + "sphinx_reredirects", ] +# Old page URLs that have moved. The site is single-version (each release +# overwrites asf-site wholesale), so these exist purely for inbound external +# links — issue comments, release blog posts, and README links that pinned a +# heading anchor. Keys are docnames without a suffix; the source file must be +# gone, or Sphinx builds the real page and the stub is never written. +redirects = { + "user-guide/distributing-work": "distributing-work/index.html", +} + # NOTE: .rst stays alongside .md because sphinx-autoapi generates RST # under autoapi/ and Sphinx needs the suffix to parse it. The human- # authored docs are all MyST .md now. ".md" is routed through myst-nb so diff --git a/docs/source/user-guide/common-operations/udf-and-udfa.md b/docs/source/user-guide/common-operations/udf-and-udfa.md index 8d1cc876d..07005fd0f 100644 --- a/docs/source/user-guide/common-operations/udf-and-udfa.md +++ b/docs/source/user-guide/common-operations/udf-and-udfa.md @@ -391,7 +391,8 @@ df.select("a", exp_smooth(col("a")).alias("smooth_a")).show() User Defined Table Functions are slightly different than the other functions described here. These functions take any number of `Expr` arguments, but only literal expressions are supported. Table functions must return a Table -Provider as described in the ref:`_io_custom_table_provider` page. +Provider as described in the {ref}`Custom Table Provider ` +page. Once you have a table function, you can register it with the session context by using {py:func}`datafusion.context.SessionContext.register_udtf`. diff --git a/docs/source/user-guide/data-sources.md b/docs/source/user-guide/data-sources.md index 22e666837..ae56ae5a1 100644 --- a/docs/source/user-guide/data-sources.md +++ b/docs/source/user-guide/data-sources.md @@ -214,7 +214,7 @@ Features that are available in PyIceberg but not yet in Iceberg Rust will not be ## Custom Table Provider You can implement a custom Data Provider in Rust and expose it to DataFusion through the -the interface as describe in the {ref}`Custom Table Provider ` +interface described in the {ref}`Custom Table Provider ` section. This is an advanced topic, but a [user example](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example) is provided in the DataFusion repository. diff --git a/docs/source/user-guide/distributing-work.md b/docs/source/user-guide/distributing-work/expressions.md similarity index 86% rename from docs/source/user-guide/distributing-work.md rename to docs/source/user-guide/distributing-work/expressions.md index 527fe8830..252defc24 100644 --- a/docs/source/user-guide/distributing-work.md +++ b/docs/source/user-guide/distributing-work/expressions.md @@ -17,24 +17,19 @@ under the License. --> -# Distributing work +(distributed_expressions)= + +# Shipping expressions to workers DataFusion supports splitting work across processes by shipping serialized expressions to workers: the driver builds an {py:class}`~datafusion.Expr`, each worker evaluates it against its own slice of data. This pattern suits embarrassingly-parallel -workloads where the driver decides partitioning up front. - -Query-level distribution — where the runtime partitions a single -logical or physical plan across worker nodes — is in progress -upstream via [datafusion-distributed](https://github.com/apache/datafusion-distributed) and [Apache -Ballista](https://github.com/apache/datafusion-ballista). Both -have short sections at the end of this page; integration details -will land as those projects become usable from datafusion-python. - -## Expression-level distribution +workloads where the driver decides partitioning up front. If you +would rather have a library partition your query plan for you, see +{ref}`distributed_query_engines` instead. -DataFusion expressions support distribution directly: pass one to a +Expressions support this directly: pass one to a worker process and Python's standard [pickle](https://docs.python.org/3/library/pickle.html) machinery serializes it transparently — the same machinery @@ -43,7 +38,7 @@ similar libraries already use to ship function arguments. Python UDFs — scalar, aggregate, and window — travel inside the serialized expression; the receiver does not need to pre-register them. -### Basic worker-pool example +## Basic worker-pool example Define a worker function that takes the expression plus a batch and returns the evaluated result: @@ -92,7 +87,7 @@ see [Safe importing of main module](https://docs.python.org/3/library/multiproce in the Python docs. ::: -### What travels with the expression +## What travels with the expression - **Built-in functions** (`abs`, `length`, arithmetic, comparisons, etc.) — fully portable. Worker needs nothing pre-registered. @@ -113,7 +108,7 @@ in the Python docs. (distributed_udf_portability)= -### Portability requirements for inline Python UDFs +## Portability requirements for inline Python UDFs Inline Python UDFs ride on [cloudpickle](https://github.com/cloudpipe/cloudpickle), which imposes two requirements on the worker environment: @@ -133,7 +128,7 @@ requirements on the worker environment: Self-contained UDFs (no imports beyond what the worker already has, e.g. `pyarrow`) avoid this entirely. -### Registering shared UDFs on workers +## Registering shared UDFs on workers When an expression references an FFI capsule UDF (or any UDF the worker must resolve from its registered functions), set up the @@ -164,7 +159,7 @@ fine for expressions that only reference built-ins and Python UDFs, but FFI-capsule-backed registrations must be installed on the global context to resolve. -### Python 3.14 default change +## Python 3.14 default change Python 3.14 changed the Linux default start method for {py:mod}`multiprocessing` from `fork` to `forkserver` (macOS has @@ -174,7 +169,7 @@ workers via copy-on-write; with `forkserver` and `spawn` it is not. The {py:func}`~datafusion.ipc.set_worker_ctx` pattern works on every start method — prefer it over relying on inherited state. -### Practical considerations +## Practical considerations - **Serialized size scales with what travels inline.** A serialized expression of just built-ins is small (tens of bytes). An @@ -189,7 +184,7 @@ every start method — prefer it over relying on inherited state. the captured state is large, mutable, or not portable to the worker's environment. See {ref}`Portability requirements for inline Python UDFs ` for the Python-version and imported-module rules. -### Disabling Python UDF inlining +## Disabling Python UDF inlining For a stricter wire format, call {py:meth}`SessionContext.with_python_udf_inlining(enabled=False) @@ -241,7 +236,7 @@ threat model. (distributed_expr_security)= -### Security +## Security :::{warning} Reconstructing an expression containing a Python UDF executes @@ -255,7 +250,9 @@ functions and pre-registered Rust-side UDFs, and avoid {py:func}`pickle.loads` on externally supplied bytes entirely. ::: -### Reference: session context slots +(session_context_slots)= + +## Reference: session context slots There is only one type — {py:class}`SessionContext`. It can occupy up to four *slots* in a running program: @@ -297,32 +294,10 @@ Sharp edges: - The inlining toggle is per-context state, not a global switch. Two contexts with different toggles can coexist in one process. -## Query-level distribution via datafusion-distributed - -🚧 *Work in progress upstream — not yet usable from datafusion-python.* - -[datafusion-distributed](https://github.com/apache/datafusion-distributed) -splits a single physical plan into stages and runs each stage on a -different worker node. The driver writes a SQL or DataFrame query -once; the runtime handles partitioning, shuffles, and reassembly. - -A datafusion-python integration is in development. This section will -document the integration once it lands. In the meantime, the -expression-level approach above covers most use cases that do not -require automatic plan partitioning. - -## Query-level distribution via Apache Ballista - -🚧 *Work in progress upstream — not yet usable from datafusion-python.* - -[Apache Ballista](https://github.com/apache/datafusion-ballista) -provides distributed query execution on top of DataFusion with a -scheduler / executor model better suited to long-lived cluster -deployments. A datafusion-python integration is on the roadmap; this -section will fill in once the integration is usable. - ## See also +- {ref}`distributed_query_engines` — the other road: let a library + partition the plan for you. - {py:mod}`datafusion.ipc` — worker context API. - `examples/multiprocessing_pickle_expr.py` — runnable `multiprocessing.Pool` example that ships a different parametric diff --git a/docs/source/user-guide/distributing-work/index.md b/docs/source/user-guide/distributing-work/index.md new file mode 100644 index 000000000..3047ea39e --- /dev/null +++ b/docs/source/user-guide/distributing-work/index.md @@ -0,0 +1,81 @@ + + +(distributing_work)= + +# Distributing work + +A single {py:class}`~datafusion.SessionContext` already uses every +core on the machine — DataFusion partitions and parallelizes within +a process without being asked. See +{doc}`../configuration` for tuning that. + +This section is about the step after that: getting work onto more +than one process or more than one machine. There are two roads, and +they suit different problems. + +## Pick your road + +**You decide the partitioning → ship expressions.** + +You already know how the data divides — one file per worker, one +customer per worker, one parameter setting per worker. You build an +{py:class}`~datafusion.Expr` in the driver and hand a copy to each +worker along with its slice. Standard Python `pickle` moves it, so +{py:mod}`multiprocessing`, Ray, and anything else that ships function +arguments works with no extra machinery. + +Best for embarrassingly-parallel work: parameter sweeps, per-file +transforms, scoring batches. Available today. + +→ {ref}`distributed_expressions` + +**A library decides the partitioning → install a query engine.** + +You write one ordinary SQL or DataFrame query against a table too +large for one machine, and an engine library splits the plan into +stages, runs them on its workers, and reassembles the result. You do +not partition anything and your queries do not change. + +Best for large-scale analytical queries — joins and aggregations over +data that does not fit on one node. Being built upstream; not yet +usable from datafusion-python. + +→ {ref}`distributed_query_engines` + +## Choosing between them + +The two are not competing implementations of one feature. The +question is who owns the partitioning decision. + +If you can state the partitioning in one line — "one worker per input +file" — the expression road is simpler, has no cluster to operate, +and works now. If stating it requires knowing how a join will +shuffle, that decision belongs to a query planner, which is what an +engine library provides. + +They also compose. An engine handles the query; expressions handle +whatever you want to fan out around it. + +```{toctree} +:maxdepth: 2 + +expressions +query-engines +``` diff --git a/docs/source/user-guide/distributing-work/query-engines.md b/docs/source/user-guide/distributing-work/query-engines.md new file mode 100644 index 000000000..0aa045c79 --- /dev/null +++ b/docs/source/user-guide/distributing-work/query-engines.md @@ -0,0 +1,100 @@ + + +(distributed_query_engines)= + +# Distributed query engines + +A distributed query engine takes the query you already wrote and runs +it across several machines for you. You do not partition anything by +hand and you do not change your queries — the engine installs itself +on your session, rewrites the plan into stages, and executes those +stages on its own workers. + +This is the counterpart to {ref}`distributed_expressions`, where +*you* decide the partitioning and ship one expression per slice of +data. Choose an engine when you want a single query spread across a +cluster; choose expressions when you already know how to split the +work and only need parallelism. + +## How an engine attaches to your session + +An engine library ships an object you hand to +{py:meth}`~datafusion.SessionContext.with_extensions`: + +```python +from datafusion import SessionContext +import my_engine + +ctx = SessionContext().with_extensions(my_engine.Extension("scheduler:50050")) + +ctx.register_table("events", my_engine.TableProvider("s3://bucket/events")) +ctx.sql("SELECT country, count(*) FROM events GROUP BY country").show() +``` + +The `sql` call is unchanged from a single-process program. What +changed is underneath it: installing the engine gave the session a +query planner of the engine's own, and that planner is what turns +your plan into distributed stages. + +Two consequences worth knowing: + +- **The engine must be installed before you run the query**, not + before you register tables. Registration order does not matter; + `with_extensions` is a session-level setup step. +- **The engine has to be able to reach your tables and functions.** + It ships the plan to its workers, so anything the plan references + has to be reconstructible there. Table providers from the engine's + own library always are. A Python UDF may or may not be — the engine + library documents what it supports, and + {ref}`distributed_udf_portability` describes the constraints that + apply to Python callables crossing a process boundary in general. + +If you install more than one library, pass them in one +`with_extensions` call so they can see each other. See +{ref}`user_guide_extensions` for the details of installing extension +libraries, and {ref}`ffi` if you want to write an engine yourself. + +## Available engines + +Query-level distribution is being built upstream. Neither project +below is usable from datafusion-python yet; both sections will fill +in as the integrations land. + +### datafusion-distributed + +🚧 *Work in progress upstream — not yet usable from datafusion-python.* + +[datafusion-distributed](https://github.com/apache/datafusion-distributed) +splits a single physical plan into stages and runs each stage on a +different worker node. The driver writes a SQL or DataFrame query +once; the runtime handles partitioning, shuffles, and reassembly. + +A datafusion-python integration is in development. In the meantime, +{ref}`distributed_expressions` covers most use cases that do not +require automatic plan partitioning. + +### Apache Ballista + +🚧 *Work in progress upstream — not yet usable from datafusion-python.* + +[Apache Ballista](https://github.com/apache/datafusion-ballista) +provides distributed query execution on top of DataFusion with a +scheduler / executor model better suited to long-lived cluster +deployments. A datafusion-python integration is on the roadmap. diff --git a/docs/source/user-guide/extensions.md b/docs/source/user-guide/extensions.md new file mode 100644 index 000000000..8c86b4a67 --- /dev/null +++ b/docs/source/user-guide/extensions.md @@ -0,0 +1,142 @@ + + +(user_guide_extensions)= + +# Using extension libraries + +An extension library is a separate package that teaches a +{py:class}`~datafusion.SessionContext` something it does not know on its own — +a new data source, extra functions, or a different way of executing your +queries. You install it with `pip`, hand it to your session, and keep writing +the same SQL and DataFrame code. + +Examples in the wild include [delta-rs](https://delta-io.github.io/delta-rs/), +which exposes Delta Lake tables to DataFusion, and the two worked examples in +this repository under +[`examples/`](https://github.com/apache/datafusion-python/tree/main/examples). + +## Two kinds of extension + +Which one you have determines how much setup you do. + +**Tables and functions register directly.** If the library gives you a table +or a function, register it the same way you would register a CSV file. No +extra setup: + +```python +from datafusion import SessionContext +import my_tables + +ctx = SessionContext() +ctx.register_table("events", my_tables.TableProvider("s3://bucket/events")) +ctx.sql("SELECT count(*) FROM events").show() +``` + +**Libraries that change how queries run need to be installed on the session.** +A distributed engine, or anything that rewrites your query plan, has to be +attached to the session before it can do its work. That is what +{py:meth}`~datafusion.SessionContext.with_extensions` is for. The library +documents an object — often called `Extension` — that you pass to it: + +```python +from datafusion import SessionContext +import my_engine + +ctx = SessionContext().with_extensions(my_engine.Extension("scheduler:50050")) +ctx.register_table("events", my_engine.TableProvider("s3://bucket/events")) +ctx.sql("SELECT count(*) FROM events").show() +``` + +`with_extensions` returns a context; use the returned one. It shares +everything else with the context you called it on, so tables you registered +before the call are still there. + +## Using more than one library + +Pass them all to a single call: + +```python +ctx = SessionContext().with_extensions( + my_tables.Extension(), + my_engine.Extension("scheduler:50050"), +) +``` + +One call rather than several is worth preferring: it lets the libraries see +each other, which they cannot do if you install them one at a time. Order +rarely matters. When a library needs a particular position — usually "list me +last" for something that wraps the others — it says so in its own +documentation. + +## Two things that will bite you + +**Keep your context alive.** A `DataFrame` or a plan does not keep its session +alive on its own. If a context is garbage-collected while something built from +it is still in use, the next query fails with: + +```text +TaskContextProvider went out of scope over FFI boundary +``` + +Almost always this is a helper that built a context locally and returned a +DataFrame: + +```python +# Wrong — ctx is collected when the function returns. +def load(): + ctx = SessionContext().with_extensions(my_engine.Extension()) + return ctx.sql("SELECT * FROM events") + +# Right — hand back the context too, or keep it on an object that lives +# as long as the frames derived from it. +def load(): + ctx = SessionContext().with_extensions(my_engine.Extension()) + return ctx, ctx.sql("SELECT * FROM events") +``` + +**Versions have to match.** An extension library is compiled against one +DataFusion version. A mismatch raises an `ImportError` naming the version it +found and the version expected, at the moment you register or install the +library — not silently at query time. If you see one, upgrade or downgrade the +extension library so its DataFusion version matches this package's. See +{ref}`extension_version_mismatch`. + +## Checking what a session knows about + +{py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` and +{py:meth}`~datafusion.SessionContext.physical_extension_codec_ids` list which +libraries a session has been taught about. Useful when a query fails and you +want to confirm the library actually got installed: + +```python +ctx = SessionContext().with_extensions(my_engine.Extension()) +ctx.logical_extension_codec_ids() +# ['my_engine.LogicalCodec'] +``` + +An empty list means nothing extra is installed. + +## Next steps + +- {ref}`distributed_query_engines` — running your queries across several + machines with an engine library. +- {ref}`distributed_expressions` — the other road to parallelism, where you + decide the partitioning and ship expressions to a worker pool. +- {ref}`ffi` — writing an extension library of your own. diff --git a/docs/source/user-guide/index.md b/docs/source/user-guide/index.md index 0155c4e91..07b912a8b 100644 --- a/docs/source/user-guide/index.md +++ b/docs/source/user-guide/index.md @@ -20,7 +20,11 @@ # User Guide The user guide walks through installing DataFusion in Python, building queries -with the DataFrame API or SQL, reading and writing data, and tuning execution. +with the DataFrame API or SQL, reading and writing data, tuning execution, and +spreading work across processes or machines. + +If you are writing an extension library rather than using one, see the +{ref}`Extension Guide `. ```{toctree} :maxdepth: 2 @@ -32,7 +36,8 @@ dataframe/index common-operations/index io/index configuration -distributing-work +extensions +distributing-work/index sql upgrade-guides ai-coding-assistants diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 4bd42c192..b5cf8d16c 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -93,6 +93,8 @@ The library exposes a bundle object implementing planner installed so far, so several libraries that each ship one nest instead of displacing each other. See the {ref}`ffi` guide. +(extension_version_mismatch)= + ### Mismatched extension libraries now fail loudly Objects imported through the capsule protocol are checked against the major diff --git a/examples/multiprocessing_pickle_expr.py b/examples/multiprocessing_pickle_expr.py index 73a99c2db..db1a4c5fd 100644 --- a/examples/multiprocessing_pickle_expr.py +++ b/examples/multiprocessing_pickle_expr.py @@ -19,7 +19,7 @@ For background — the shipped-expression model, what travels inline vs by name, portability requirements, and the security threat model — -see ``docs/source/user-guide/distributing-work.rst``. +see https://datafusion.apache.org/python/user-guide/distributing-work/expressions.html. Builds a list of parametric expressions in the driver — each closing over a different threshold value — ships one per worker via diff --git a/examples/ray_pickle_expr.py b/examples/ray_pickle_expr.py index 04cea463d..e8be06b0f 100644 --- a/examples/ray_pickle_expr.py +++ b/examples/ray_pickle_expr.py @@ -19,7 +19,7 @@ For background — the shipped-expression model, what travels inline vs by name, portability requirements, and the security threat model — -see ``docs/source/user-guide/distributing-work.rst``. +see https://datafusion.apache.org/python/user-guide/distributing-work/expressions.html. Build an expression in the driver, ship it to a pool of Ray actors, and have each actor evaluate it against its own slice of data. Python UDFs From 7c86fb063cc7791961d2f50c929fa27bb2f6780f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 08:39:54 -0400 Subject: [PATCH 29/33] docs: split the FFI guide into one section per audience MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `contributor-guide/ffi.md` had grown to 795 lines serving three different readers at once, filed under a section whose index says it is for people contributing to this repository. An engineer at delta-rs or a distributed-engine vendor is neither a contributor nor an end user; they consume a published, versioned protocol, and the only description of it lived behind a heading telling them the page was not for them. Adds a third top-level section, `extension-guide/`, so the sidebar reads User Guide / Extension Guide / Contributor Guide — one per audience. It carries the `(ffi)=` label, so every existing reference keeps resolving, including the `:ref:`ffi`` inside `context.py`'s docstring that ships in the wheel. The maintainer-facing rationale moves to `contributor-guide/ffi-internals.md`: the weak-`Arc` scheme, why repairing an orphaned provider cannot work, why planner codec rebinding is one level deep, and the two upstream issues. An extension vendor should not be reading "that is a bug rather than a design, do not copy the pattern" halfway down their integration guide. The PyO3 `frozen` policy moves to `contributor-guide/pyo3-guidelines.md`. It is project review policy with nothing to do with FFI, and extracting it repairs a prose bug: it had been spliced into the middle of "Implementation Details", so the sentence "If you were interfacing with a library that provided the above `FFI_TableProvider`" resumed 46 lines after the snippet it referred to. Those two halves are rejoined on `capsule-protocol.md`. Fills the coverage gap the split exposed. The codebase exports 18 capsule getters and the old page documented 7; the remaining 11 appeared on no page that even listed them. The section index now carries a table of all 18, and `table-providers.md` and `functions.md` document the catalog family, table functions, physical optimizer rules, and extension options for the first time. Corrects the argument rule while moving it. The old section asserted that capsule getters "receive the SessionContext they are being installed on", which is true for the codec, planner, and table-function hooks but not for the catalog family: `CapsuleGetterArg::LogicalCodec` passes the host's logical codec as a bare capsule, and `__datafusion_table_provider__` gets a session from `SessionContext.register_table` but a codec capsule from `Schema.register_table`. Nothing breaks, because every implementation passes the argument to `ffi_logical_codec_from_pycapsule`, which handles both — so the rule is now stated by capability rather than by type, and the upgrade guide says which hooks changed. Also converts the guide's five in-page heading links to labelled `{ref}` targets, since heading anchors rot silently on rewording, and redirects the old `contributor-guide/ffi.html` URL. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/conf.py | 1 + .../source/contributor-guide/ffi-internals.md | 137 +++ docs/source/contributor-guide/ffi.md | 795 ------------------ docs/source/contributor-guide/index.md | 10 +- .../contributor-guide/pyo3-guidelines.md | 74 ++ docs/source/extension-guide/bundles.md | 252 ++++++ .../extension-guide/capsule-protocol.md | 224 +++++ docs/source/extension-guide/checklist.md | 99 +++ docs/source/extension-guide/codecs.md | 166 ++++ docs/source/extension-guide/functions.md | 153 ++++ docs/source/extension-guide/index.md | 131 +++ docs/source/extension-guide/query-planners.md | 106 +++ docs/source/extension-guide/sessions.md | 118 +++ .../source/extension-guide/table-providers.md | 138 +++ docs/source/extension-guide/why-ffi.md | 148 ++++ docs/source/index.md | 11 + docs/source/user-guide/io/table_provider.md | 25 +- docs/source/user-guide/upgrade-guides.md | 21 +- 18 files changed, 1785 insertions(+), 824 deletions(-) create mode 100644 docs/source/contributor-guide/ffi-internals.md delete mode 100644 docs/source/contributor-guide/ffi.md create mode 100644 docs/source/contributor-guide/pyo3-guidelines.md create mode 100644 docs/source/extension-guide/bundles.md create mode 100644 docs/source/extension-guide/capsule-protocol.md create mode 100644 docs/source/extension-guide/checklist.md create mode 100644 docs/source/extension-guide/codecs.md create mode 100644 docs/source/extension-guide/functions.md create mode 100644 docs/source/extension-guide/index.md create mode 100644 docs/source/extension-guide/query-planners.md create mode 100644 docs/source/extension-guide/sessions.md create mode 100644 docs/source/extension-guide/table-providers.md create mode 100644 docs/source/extension-guide/why-ffi.md diff --git a/docs/source/conf.py b/docs/source/conf.py index 7f9c4b89a..0369d7fe7 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -66,6 +66,7 @@ # gone, or Sphinx builds the real page and the stub is never written. redirects = { "user-guide/distributing-work": "distributing-work/index.html", + "contributor-guide/ffi": "../extension-guide/index.html", } # NOTE: .rst stays alongside .md because sphinx-autoapi generates RST diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md new file mode 100644 index 000000000..75c32444c --- /dev/null +++ b/docs/source/contributor-guide/ffi-internals.md @@ -0,0 +1,137 @@ + + +(ffi_internals)= + +# FFI framing internals + +Read this before changing how datafusion-python frames FFI components. You do +**not** need it to write an extension library — that is the +{ref}`Extension Guide `. + +The invariants here are the reasons the extension-facing rules are shaped the +way they are. Each one has an "obvious" simplification that does not work, and +the point of this page is to record why, so the next person does not spend a +week rediscovering it. + +For the wire format itself — how a codec id is stored alongside a payload, +routed back on decode, and which two cases stay unframed — see the module +documentation in `crates/core/src/codec.rs`. It is the authority; this page +does not restate it. + +(ffi_internals_one_arc)= + +## One session, one `Arc` + +Every codec handed to a foreign object carries an `FFI_TaskContextProvider`, +and that type holds its provider **weakly**. A registered catalog provider +upgrades the handle on every `supports_filters_pushdown` and every `scan`. +Those handles are bound to one particular `Arc` allocation, not +to the logical session, so anything that replaces the allocation orphans all of +them and the next query fails with `TaskContextProvider went out of scope over +FFI boundary`. + +So a `PySessionContext` keeps the `Arc` it was created with for +its whole life. Installing a query planner writes the new `SessionState` back +through `state_ref()`, exactly as `add_physical_optimizer_rule` does, rather +than deriving a replacement context. The session id is carried across that +rewrite — `SessionStateBuilder` mints a fresh one otherwise — so `session_id()` +and every `TaskContext` the session hands out keep agreeing. + +Repairing the damage instead of avoiding it does not work in general. A context +can rebuild the codecs it holds in its own fields, but a codec already embedded +in a registered `FFI_CatalogProvider` — and in every `FFI_SchemaProvider` and +`FFI_TableProvider` minted from it — is not reachable from Python at all. Nor +can the codec simply retain the session that built it: a codec handed to a +provider is routinely registered straight back into that same session, which +would close the cycle +`SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and +leak it. + +`SessionContext.enable_url_table` is the one exception. It clones the +underlying `SessionContext`, so the returned context has an allocation of its +own and must not outlive the receiver. It also forks the session's state while +keeping its id, so two handles report one `session_id()` with divergent +configuration. That is a bug rather than a design, tracked in +[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708); +do not copy the pattern. + +(ffi_internals_rebinding)= + +## Why planner codec rebinding is one level deep + +Installing a codec on a session that already has a foreign planner rebuilds +that planner against the new chain. The rebuild swaps the codecs on the +installed `ForeignQueryPlanner` handle, and only that handle. A planner that +wraps a fallback resolved that fallback when *it* was installed, and holds the +result inside its own library's private data — behind a `create_physical_plan` +function pointer, with no Python-side handle. A codec installed afterwards +therefore reaches the outer planner and not the fallback, which keeps whichever +codecs were in force when it was imported. + +Neither side can repair that: + +- **The host cannot reach it.** `FFI_QueryPlanner::new_with_ffi_codecs` unwraps + exactly one `ForeignQueryPlanner` layer. There is no deeper handle to unwrap + — the same situation as a codec embedded in a registered + `FFI_CatalogProvider`. +- **The planner library cannot re-derive it.** `FFI_QueryPlanner` holds its + codecs by value, and `Session` exposes no accessor for the ones the host + currently has, so `create_physical_plan` cannot pick them up from the session + it is handed. The rebuild has to be eager, and an eager rebuild only sees the + top layer. + +A fix has to come from upstream, and is tracked in +[apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762). + +The stale codecs stay usable rather than dangling — they hold weak handles to +the one `Arc` that the previous section keeps alive — so the +effect is a fallback hop serializing with an older codec, not a failure. It is +also invisible to the examples in this repository, which use one fallback in +the same cdylib as its wrapper; `datafusion-ffi` short-circuits a same-library +hop rather than serializing, so no codec runs. A fallback in a *different* +library would serialize, and would do it with the codecs it was imported with. + +The extension-facing consequence — install codecs before a layered planner, and +prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`. + +## Two argument kinds for one convention + +`CapsuleGetterArg` in `crates/util/src/lib.rs` distinguishes three cases: no +argument, the session, and the host's logical extension codec as a bare +capsule. Provider and catalog registration passes the codec; the extension +codecs, the query planner, and table functions get the session. + +The distinction exists so that a `TypeError` from a getter that refused its +argument can be rewritten into an `ImportError` naming what the getter *should* +accept. Getting that message wrong sends an extension author to fix the wrong +signature, which is why every capsule getter routes through +`call_capsule_getter` rather than calling `getattr` directly. Three importers +previously each had their own copy of that logic and each missed later +corrections to it. + +When adding a hook, decide which arm it needs by what its FFI constructor +requires: a task-context provider means it must have the session, since you +cannot get a provider off a capsule. A codec alone means either works, and +passing the codec keeps the host from handing out session handles it does not +need to. + +The extension-facing statement of this is {ref}`extension_getter_argument`, +which deliberately describes the argument by capability — "something you can +read the host's logical codec off" — rather than by type. diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md deleted file mode 100644 index 244505c2c..000000000 --- a/docs/source/contributor-guide/ffi.md +++ /dev/null @@ -1,795 +0,0 @@ - - -(ffi)= - -# Python Extensions - -The DataFusion in Python project is designed to allow users to extend its functionality in a few core -areas. Ideally many users would like to package their extensions as a Python package and easily -integrate that package with this project. This page serves to describe some of the challenges we face -when doing these integrations and the approach our project uses. - -## The Primary Issue - -Suppose you wish to use DataFusion and you have a custom data source that can produce tables that -can then be queried against, similar to how you can register a {ref}`CSV ` or -{ref}`Parquet ` file. In DataFusion terminology, you likely want to implement a -{ref}`Custom Table Provider `. In an effort to make your data source -as performant as possible and to utilize the features of DataFusion, you may decide to write -your source in Rust and then expose it through [PyO3](https://pyo3.rs) as a Python library. - -At first glance, it may appear the best way to do this is to add the `datafusion-python` -crate as a dependency, provide a `PyTable`, and then to register it with the -`SessionContext`. Unfortunately, this will not work. - -When you produce your code as a Python library and it needs to interact with the DataFusion -library, at the lowest level they communicate through an Application Binary Interface (ABI). -The acronym sounds similar to API (Application Programming Interface), but it is distinctly -different. - -The ABI sets the standard for how these libraries can share data and functions between each -other. One of the key differences between Rust and other programming languages is that Rust -does not have a stable ABI. What this means in practice is that if you compile a Rust library -with one version of the `rustc` compiler and I compile another library to interface with it -but I use a different version of the compiler, there is no guarantee the interface will be -the same. - -In practice, this means that a Python library built with `datafusion-python` as a Rust -dependency will generally **not** be compatible with the DataFusion Python package, even -if they reference the same version of `datafusion-python`. If you attempt to do this, it may -work on your local computer if you have built both packages with the same optimizations. -This can sometimes lead to a false expectation that the code will work, but it frequently -breaks the moment you try to use your package against the released packages. - -You can find more information about the Rust ABI in their -[online documentation](https://doc.rust-lang.org/reference/abi.html). - -## The FFI Approach - -Rust supports interacting with other programming languages through it's Foreign Function -Interface (FFI). The advantage of using the FFI is that it enables you to write data structures -and functions that have a stable ABI. The allows you to use Rust code with C, Python, and -other languages. In fact, the [PyO3](https://pyo3.rs) library uses the FFI to share data -and functions between Python and Rust. - -The approach we are taking in the DataFusion in Python project is to incrementally expose -more portions of the DataFusion project via FFI interfaces. This allows users to write Rust -code that does **not** require the `datafusion-python` crate as a dependency, expose their -code in Python via PyO3, and have it interact with the DataFusion Python package. - -Early adopters of this approach include [delta-rs](https://delta-io.github.io/delta-rs/) -who has adapted their Table Provider for use in `` `datafusion-python` `` with only a few lines -of code. Also, the DataFusion Python project uses the existing definitions from -[Apache Arrow CStream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html) -to support importing **and** exporting tables. Any Python package that supports reading -the Arrow C Stream interface can work with DataFusion Python out of the box! You can read -more about working with Arrow sources in the {ref}`Data Sources ` -page. - -To learn more about the Foreign Function Interface in Rust, the -[Rustonomicon](https://doc.rust-lang.org/nomicon/ffi.html) is a good resource. - -## Inspiration from Arrow - -DataFusion is built upon [Apache Arrow](https://arrow.apache.org/). The canonical Python -Arrow implementation, [pyarrow](https://arrow.apache.org/docs/python/index.html) provides -an excellent way to share Arrow data between Python projects without performing any copy -operations on the data. They do this by using a well defined set of interfaces. You can -find the details about their stream interface -[here](https://arrow.apache.org/docs/format/CStreamInterface.html). The -[Rust Arrow Implementation](https://github.com/apache/arrow-rs) also supports these -`C` style definitions via the Foreign Function Interface. - -In addition to using these interfaces to transfer Arrow data between libraries, `pyarrow` -goes one step further to make sharing the interfaces easier in Python. They do this -by exposing PyCapsules that contain the expected functionality. - -You can learn more about PyCapsules from the official -[Python online documentation](https://docs.python.org/3/c-api/capsule.html). PyCapsules -have excellent support in PyO3 already. The -[PyO3 online documentation](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule) is a good source -for more details on using PyCapsules in Rust. - -Two lessons we leverage from the Arrow project in DataFusion Python are: - -- We reuse the existing Arrow FFI functionality wherever possible. -- We expose PyCapsules that contain a FFI stable struct. - -## Implementation Details - -The bulk of the code necessary to perform our FFI operations is in the upstream -[DataFusion](https://datafusion.apache.org/) core repository. You can review the code and -documentation in the [datafusion-ffi] crate. - -Our FFI implementation is narrowly focused at sharing data and functions with Rust backed -libraries. This allows us to use the [abi_stable crate](https://crates.io/crates/abi_stable). -This is an excellent crate that allows for easy conversion between Rust native types -and FFI-safe alternatives. For example, if you needed to pass a `Vec` via FFI, -you can simply convert it to a `RVec` in an intuitive manner. It also supports -features like `RResult` and `ROption` that do not have an obvious translation to a -C equivalent. - -The [datafusion-ffi] crate has been designed to make it easy to convert from DataFusion -traits into their FFI counterparts. For example, if you have defined a custom -[TableProvider](https://docs.rs/datafusion/45.0.0/datafusion/catalog/trait.TableProvider.html) -and you want to create a sharable FFI counterpart, you could write: - -```rust -let my_provider = MyTableProvider::default(); -let ffi_provider = FFI_TableProvider::new(Arc::new(my_provider), false, None); -``` - -(ffi_pyclass_mutability)= - -## PyO3 class mutability guidelines - -PyO3 bindings should present immutable wrappers whenever a struct stores shared or -interior-mutable state. In practice this means that any `#[pyclass]` containing an -`Arc>` or similar synchronized primitive must opt into `#[pyclass(frozen)]` -unless there is a compelling reason not to. - -The execution context illustrates the preferred pattern. `PySessionContext` in -{file}`src/context.rs` stays frozen even though it shares mutable state internally via -`SessionContext`. This ensures PyO3 tracks borrows correctly while Python-facing APIs -clone the inner `SessionContext` or return new wrappers instead of mutating the -existing instance in place: - -```rust -#[pyclass(from_py_object, frozen, name = "SessionContext", module = "datafusion", subclass)] -#[derive(Clone)] -pub struct PySessionContext { - pub ctx: SessionContext, -} -``` - -Occasionally a type must remain mutable—for example when PyO3 attribute setters need to -update fields directly. In these rare cases add an inline justification so reviewers and -future contributors understand why `frozen` is unsafe to enable. `DataTypeMap` in -{file}`src/common/data_type.rs` includes such a comment because PyO3 still needs to track -field updates: - -```rust -// TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now -#[derive(Debug, Clone)] -#[pyclass(from_py_object, name = "DataTypeMap", module = "datafusion.common", subclass)] -pub struct DataTypeMap { - #[pyo3(get, set)] - pub arrow_type: PyDataType, - #[pyo3(get, set)] - pub python_type: PythonType, - #[pyo3(get, set)] - pub sql_type: SqlType, -} -``` - -When reviewers encounter a mutable `#[pyclass]` without a comment, they should request -an explanation or ask that `frozen` be added. Keeping these wrappers frozen by default -helps avoid subtle bugs stemming from PyO3's interior mutability tracking. - -If you were interfacing with a library that provided the above `FFI_TableProvider` and -you needed to turn it back into an `TableProvider`, you can turn it into a -`ForeignTableProvider` with implements the `TableProvider` trait. - -```rust -let foreign_provider: ForeignTableProvider = ffi_provider.into(); -``` - -If you review the code in [datafusion-ffi] you will find that each of the traits we share -across the boundary has two portions, one with a `FFI_` prefix and one with a `Foreign` -prefix. This is used to distinguish which side of the FFI boundary that struct is -designed to be used on. The structures with the `FFI_` prefix are to be used on the -**provider** of the structure. In the example we're showing, this means the code that has -written the underlying `TableProvider` implementation to access your custom data source. -The structures with the `Foreign` prefix are to be used by the receiver. In this case, -it is the `datafusion-python` library. - -In order to share these FFI structures, we need to wrap them in some kind of Python object -that can be used to interface from one package to another. As described in the above -section on our inspiration from Arrow, we use `PyCapsule`. We can create a `PyCapsule` -for our provider thusly: - -```rust -let name = CString::new("datafusion_table_provider")?; -let my_capsule = PyCapsule::new_bound(py, provider, Some(name))?; -``` - -On the receiving side, turn this pycapsule object into the `FFI_TableProvider`, which -can then be turned into a `ForeignTableProvider` the associated code is: - -```rust -let capsule = capsule.cast::()?; -let data: NonNull = capsule - .pointer_checked(Some(name))? - .cast(); -let codec = unsafe { data.as_ref() }; -``` - -By convention the `datafusion-python` library expects a Python object that has a -`TableProvider` PyCapsule to have this capsule accessible by calling a function named -`__datafusion_table_provider__`. You can see a complete working example of how to -share a `TableProvider` from one python library to DataFusion Python in the -[repository examples folder](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example). - -This section has been written using `TableProvider` as an example. It is the first -extension that has been written using this approach and the most thoroughly implemented. -As we continue to expose more of the DataFusion features, we intend to follow this same -design pattern. - -## Query Planners Across Multiple Libraries - -A query can involve three independent native libraries: `datafusion-python`, a library -that owns table providers or functions, and a library that owns the query planner. The -examples use two separate extension crates so each role has a distinct shared-library -identity: - -- [`datafusion-ffi-example`] owns providers, functions, and their codecs. -- [`datafusion-ffi-query-planner-example`] owns the planner and its configuration. - -The `SessionContext` owns the codecs used for the exchange and supplies them to the -foreign planner. This lets the planner decode provider-owned objects and lets -`datafusion-python` decode the physical plan returned by the planner. The examples use -process-local tokens to demonstrate ownership; production codecs should serialize -durable metadata instead. - -### Composable codecs - -Extension codecs compose. Each call to `with_logical_extension_codec` or -`with_physical_extension_codec` appends the codec to the session's codec chain -rather than replacing prior codecs. - -**Nothing is asked of the codec itself.** Implement `LogicalExtensionCodec` or -`PhysicalExtensionCodec` exactly as you would for a session that installs only -yours. When your codec writes bytes into a serialized plan, datafusion-python -records which codec wrote them, and strips that record off again before handing the -bytes back. So your codec receives, byte for byte, the payload it wrote, and is -never offered a payload another codec wrote. - -A codec that also ships to hosts which dispatch differently may still want its own -guard against foreign payloads. Keeping one is fine; it is simply not needed for the -datafusion-python path. - -That record is the codec's **id**: a short string stored inside the plan, naming the -codec that wrote each payload. Because plans are decoded in another process — or -another program — the id has to name the same codec there as it did where the plan -was written. - -Ids are assigned for you. A codec's id is normally its exporting class's import -path, such as `my_library.Codec`, which is what you will see in -`logical_extension_codec_ids()` and in decode errors. You choose one yourself in -three cases: - -- **Two instances of one class.** Both get the same id, so the second install - raises `ValueError`. Pass `codec_id=` to tell them apart. -- **A bare `PyCapsule`.** A capsule has no class to take a name from, so installing - one through `with_logical_extension_codec` or `with_physical_extension_codec` gives - it an id private to the session that installed it; plans it encodes fail with a - clear error on any other session rather than being decoded by the wrong codec. Pass - `codec_id=` if those plans have to cross sessions. - - `with_extensions` takes no `codec_id=`, so it refuses a bare capsule outright and - tells you to wrap it. See - [Extension bundles: `with_extensions`](#extension-bundles-with_extensions). -- **A class you intend to rename.** The id follows the class name, so renaming stops - older plans from decoding. Declare `__datafusion_codec_id__` on the exporting - object to pin an id that survives the rename. - -`SessionContext.logical_extension_codec_ids()` and its physical counterpart list the -ids installed on a session, which is also what a decode failure names. - -Installing one context's codec stack on another session composes the two sessions -rather than copying codecs out of one: the imported codecs resolve their task context -against the original and stop working when it is dropped — see -[One session, one `Arc`](#one-session-one-arcsessioncontext). Pass -the context itself rather than the capsule it exports, so its codecs get an id that -other sessions can decode. - -Because decoding keys off the id rather than install position, registration order -between independent libraries does not affect decoding at all. It is visible only -on encoding, where codecs are consulted in install order and the first to claim an -object wins — so installing a library can claim objects nothing else claimed, but -never takes over an object an earlier codec was already encoding. Two libraries -that each own tables, functions, and a planner register like this: - -```python -ctx = SessionContext(config) - -# Codecs from both libraries. Order between libraries does not matter. -ctx = ctx.with_logical_extension_codec(lib_a.codec()) -ctx = ctx.with_logical_extension_codec(lib_b.codec()) -ctx = ctx.with_physical_extension_codec(lib_a.physical_codec()) -ctx = ctx.with_physical_extension_codec(lib_b.physical_codec()) - -# A session holds one planner, so layering is explicit delegation. Install the -# codecs first: the fallback captured here keeps the codecs it was exported -# with. See "Rebinding a planner's codecs is one level deep" below. -ctx.set_query_planner(lib_a.Planner()) -ctx.set_query_planner(lib_b.Planner(fallback=ctx.__datafusion_query_planner__())) - -# Tables and functions — any time before the first query. -ctx.register_table("t", lib_a.TableProvider()) -ctx.register_udf(udf(lib_b.SomeUDF())) -``` - -A codec may own functions that need no payload at all, where the name is the whole -encoding: `try_encode_udf` writes nothing and `try_decode_udf` rebuilds the function -from `name`. That is supported and needs no id, because an `Ok` with an empty -buffer is read as "no opinion" and passes the object to the next codec. -`NameOnlyUdfCodec` in the FFI example is the worked case. Anything no installed -codec claims falls through to `Default{Logical,Physical}ExtensionCodec`. - -This is the one case where your decoder is consulted about something you may not -own, because an empty payload has no id to route on. `try_decode_udf` and -its aggregate and window siblings can therefore be called with an empty `buf` and a -`name` belonging to another library. Decide from `name` and return an error if it is -not yours; do not assume `buf` is non-empty. - -The framing itself — how an id is stored alongside a payload and routed back, and the two cases -that stay unframed — is internal to datafusion-python and documented in -`crates/core/src/codec.rs` for anyone changing it. - -The current FFI logical codec supports providers and UDFs but not arbitrary custom -`LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and -local build commands. - -### Extension bundles: `with_extensions` - -The chaining above works, but it makes the caller responsible for ordering: the codecs -have to be installed before the planner, because a planner is built against whatever -codec chains exist when it is installed, and a codec added afterwards rebinds it. Get -that wrong and the planner encodes through a chain that is missing a library. - -`SessionContext.with_extensions` removes the ordering question. An extension library -exposes a bundle object implementing one or both of two hooks: - -```python -class MyEngineExtension: - def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: - # Phase one. Create fresh components bound to `ctx` on every call. - return SessionExtensionComponents( - logical_extension_codecs=(self._make_logical_codec(ctx),), - physical_extension_codecs=(self._make_physical_codec(ctx),), - ) - - def __datafusion_session_planner__(self, ctx: SessionContext, fallback): - # Phase two. `ctx` now carries every bundle's codecs, and `fallback` is - # the planner built so far. Wrapping it is what makes this library - # compose with the other planners in the call. - return self._make_planner(ctx, fallback=fallback) -``` - -Implement whichever apply: a codec-only library defines the first, a library that ships -only an optimizing planner defines the second. - -```python -ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) -ctx.register_table("t", lib_a.TableProvider()) -ctx.register_udf(udf(lib_b.SomeUDF())) -``` - -#### Two phases, because codecs and planners compose differently - -A session chains **many** codecs and dispatches between them by id. Codecs therefore -just accumulate: order affects encoding only, and decoding always routes to the codec -that wrote the payload. A session holds exactly **one** query planner, so planners -cannot accumulate — they compose by *nesting*, each wrapping the one before it and -delegating to it for work it does not handle. - -So `with_extensions` runs every `__datafusion_session_extension__` and installs all the -codecs, and only then runs each `__datafusion_session_planner__`, in argument order, -handing each the planner built so far. Two consequences worth holding onto: - -- **Bundle order matters differently for each.** For planners it sets the nesting: the - last extension listed ends up outermost and is consulted first. For codecs it never - affects decoding, and affects encoding only when two codecs would claim the same node - — see [When codec order does matter](#when-codec-order-does-matter). -- **A planner is always built against the complete codec set**, including codecs from - bundles listed after it. This is what the low-level chaining cannot give you, and it - matters most for a nested planner: the rebuild that follows a later codec install - reaches only the outermost layer (see [Rebinding a planner's codecs is one level - deep](#rebinding-a-planners-codecs-is-one-level-deep)), so a fallback captured before - the codecs were complete would stay stale forever. - -The two hooks therefore see the same session through different chains. Both receive a -handle on the one session, so the task-context provider taken off either is the same and -stays valid — but the `ctx` in phase one still carries the chains the receiver had, since -nothing is installed yet, while the `ctx` in phase two carries every bundle's codecs. A -bundle that reads the host's codec chains — `MyPlannerExtension` does, to give its -planner the host's codecs rather than minting its own — must do that in the planner hook. -Reading them in phase one gets the chains from before the call, missing even the bundle's -own codecs. - -An extension that ignores `fallback` and returns an unrelated planner replaces every -layer beneath it, including any planner the session already had. That is legal — a -library that must be the only planner does it deliberately — but it is not composable, -and nothing detects it. Returning `None` contributes no planner and leaves `fallback` -in place. - -`None` is the no-op, not `fallback`. The capsule handed to the first bundle wraps the -session's planner for export, so returning it unchanged installs that planner as a -foreign one and every plan built afterwards crosses an FFI boundary it did not before. -A bundle that decides at runtime it has nothing to contribute returns `None`. - -Three libraries that each ship a planner therefore install like this, with the -outermost last: - -```python -ctx = SessionContext(config).with_extensions( - tables.Extension(), # codecs only - functions.Extension(), # codecs only - optimizer.Extension(), # planner, wrapping the session default - distributed.Extension(), # planner, wrapping the optimizer -) -``` - -#### When codec order does matter - -Decoding is never order-dependent: a payload names its codec by id and the chain -dispatches straight to it. Encoding walks the chain in install order and stops at the -first codec that claims the node. Most of the time that is invisible, because libraries -claim disjoint things — one owns its table providers, another its UDFs, a third its own -execution plan nodes. - -It stops being invisible when a codec claims *broadly*. A node that came from another -library arrives as an opaque `ForeignExecutionPlan`, and a codec that claims any of -those will take nodes it does not own from any library installed after it. The query -still succeeds. What changes is which library wrote the bytes — so a plan that has to -decode in another process now needs whichever library happened to win, not the one whose -node it is. `MyPhysicalExtensionCodec` in the provider example claims this way, and -`test_a_greedy_codec_installed_first_claims_another_librarys_node` pins the consequence. - -Two rules of thumb: - -- **Writing a codec, claim narrowly.** Downcast to your own types. Claiming a broad - category makes your library order-sensitive for everyone downstream of it. -- **Shipping plans out of the process, verify.** Do not assume your node reached your - codec just because both are installed. Round-trip a plan through - `ExecutionPlan.to_bytes` / `from_bytes` in a test and assert your codec did the work. - -#### When the two orders conflict - -Because codec position and planner position both come from one argument list, a library -can in principle need to be early for one and late for the other: its codec must precede -a broad claimer, while its planner must nest outside that library's planner. - -Do not try to satisfy both by reordering — contribute each half at its own position. The -two hooks are independent, so a three-line adapter each is enough: - -```python -class CodecsOf: - """Contribute only the codec half of a bundle, at this position.""" - def __init__(self, inner): - self.inner = inner - - def __datafusion_session_extension__(self, ctx): - return self.inner.__datafusion_session_extension__(ctx) - - -class PlannerOf: - """Contribute only the planner half of a bundle, at this position.""" - def __init__(self, inner): - self.inner = inner - - def __datafusion_session_planner__(self, ctx, fallback): - return self.inner.__datafusion_session_planner__(ctx, fallback) - - -ctx = SessionContext(config).with_extensions( - CodecsOf(engine), CodecsOf(tables), # engine's codec first - PlannerOf(tables), PlannerOf(engine), # engine's planner outermost -) -``` - -This keeps everything `with_extensions` guarantees: one transaction, codecs complete -before any planner is built, codec ids untouched — an id is read off the codec object, -not off the extension that contributed it, so splitting a bundle cannot re-tag its -payloads. A library that expects to be composed this way should expose the halves itself -rather than make callers write the adapters. - -Falling back to the low-level `with_logical_extension_codec` / -`with_physical_extension_codec` / `set_query_planner` sequence also works, and it is the -right answer when the pieces do not come as bundles at all. But it is a real downgrade, -not just a more verbose spelling: you take back responsibility for installing every -codec before every planner, and a planner you layer by hand keeps the codecs it captured -— the [one-level rebind](#rebinding-a-planners-codecs-is-one-level-deep) does not reach -inside it. Reach for it last. - -There is no attempt here to make every permutation expressible from one call. Two -positions per bundle covers the cases that arise; anything stranger is a sign the -libraries disagree about what they own, which is better fixed there. - -#### Codecs are objects, not capsules - -`with_extensions` requires each codec to be an object exposing the capsule getter, and -refuses a bare `PyCapsule`. A codec's id is read off the object it is handed over as, -and a capsule has no type to read one from; since this method takes no `codec_id=`, -there would be nothing left to name it by. A library holding a raw capsule — which is -what a Rust implementation has — wraps it: - -```python -class MyLogicalCodec: - # Optional. Without it the id is this class's import path, which is already - # stable; declare it if you may rename the class and need old plans to decode. - __datafusion_codec_id__ = "my_library.logical.v1" - - def __init__(self, capsule): - self._capsule = capsule - - def __datafusion_logical_extension_codec__(self, session=None): - return self._capsule -``` - -Wrapping is not just bookkeeping. It ties the id to the codec rather than to the bundle -that contributed it, and that difference is load-bearing: an application commonly -presents several libraries as one bundle of its own, and the id has to survive that. -Were the id taken from the contributing bundle, wrapping `my_engine.Extension` inside -`my_app.Extension` would silently re-tag the engine's payloads, and a scheduler that -installs the engine's codec by its documented id would fail to decode plans from -composed clients while succeeding for direct ones. The wrapper travels with the codec; -the bundle does not. - -The query planner is exempt — it carries no wire id, so it may be an object or a -capsule. - -Nothing is written to the session until every factory has returned and every capsule -has been validated, so a factory that raises leaves the session exactly as it was. A -factory that mutates the context it is handed — registering a table, say — is not -rolled back, which is why bundle objects must be configuration-only: create fresh -components on each call, never cache bound components, and do not retain the context -passed in. - -Like every other derivation, the returned context is a handle on the *same* session as -the receiver — see [What a derived context shares](#what-a-derived-context-shares). -Only the Python-side codec chains belong to the returned handle; the planner is -installed on the shared session and takes effect even if that handle is discarded. - -The session owns every installed component's task-context provider, and dependent -objects do not extend its lifetime. A `DataFrame`, logical plan, or capsule can outlive -every context on the session, but any operation that reaches an FFI codec after the -last one is collected fails with `TaskContextProvider went out of scope over FFI -boundary`. Keep a context alive for as long as objects derived from it are in use. - -`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust -implementation of the protocol, including taking the task-context provider off the -supplied context, wrapping its codecs in `BundledLogicalCodec` / `BundledPhysicalCodec` -so they carry declared ids, and constructing a Python `SessionExtensionComponents`. - -### Capsule getters receive the session they are installed on - -`__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and -`__datafusion_physical_extension_codec__` all take the `SessionContext` the object is -being installed on, the same way `__datafusion_table_provider__` does: - -```rust -fn __datafusion_physical_extension_codec__<'py>( - &self, - py: Python<'py>, - session: Bound<'py, PyAny>, -) -> PyResult> { - let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; - let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); - PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") -} -``` - -This exists because the FFI constructors need things an extension library does not -have. `FFI_{Logical,Physical}ExtensionCodec::new` needs a `TaskContextProvider` for the -decode callbacks the codec will receive, and `FFI_QueryPlanner::new` needs both codecs -on top of that. Taking them from the session is what keeps a library from constructing -a `SessionContext` purely to satisfy a parameter — an empty one resolves nothing, and -`FFI_TaskContextProvider` holds it weakly, so a context built inline in the getter is -already dropped by the time the capsule is used. - -A planner uses `FFI_QueryPlanner::new_with_ffi_codecs` with the two codecs it takes off -the session, and never touches a provider directly. That also matches what installation -does anyway: `set_query_planner` builds the planner against the codecs of the session -that will run the query. - -Duck-type `session`, and do not check its type. It is the PyO3 context the binding -installs through, not the `datafusion.context.SessionContext` wrapper, so it carries -every capsule getter and `__datafusion_codec_id__` — everything the protocol asks of it -— but `isinstance(session, SessionContext)` is `False` in Python even though its `repr` -reads `datafusion.SessionContext`. The two bundle hooks -`__datafusion_session_extension__` and `__datafusion_session_planner__` are the -exception: `with_extensions` dispatches them from Python and hands them the wrapper. - -`SessionContext` accepts the argument on all three getters and ignores it, so a session -satisfies the same protocol an extension library implements. When you export the current -planner to wrap it, `ctx.__datafusion_query_planner__()` and -`ctx.__datafusion_query_planner__(ctx)` are both fine. - -### A codec decodes against the session that is running the query - -Because the provider comes from the host, a decode callback running inside an extension -library resolves names against the session running the query. A function registered with -`ctx.register_udf(...)` is visible to a foreign codec decoding a node that references it -by name, and the handle is live rather than a snapshot, so a registration made after the -codec is installed is visible too. - -This is covered in -`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`, -where the example codecs take a `require_udf_on_decode` name and resolve it out of the -task context they are handed. - -### One session, one `Arc` - -Every codec handed to a foreign object carries an `FFI_TaskContextProvider`, and that -type holds its provider **weakly**. A registered catalog provider upgrades the handle on -every `supports_filters_pushdown` and every `scan`. Those handles are bound to one -particular `Arc` allocation, not to the logical session, so anything that -replaces the allocation orphans all of them and the next query fails with -`TaskContextProvider went out of scope over FFI boundary`. - -So a `PySessionContext` keeps the `Arc` it was created with for its whole -life. Installing a query planner writes the new `SessionState` back through -`state_ref()`, exactly as `add_physical_optimizer_rule` does, rather than deriving a -replacement context. The session id is carried across that rewrite — `SessionStateBuilder` -mints a fresh one otherwise — so `session_id()` and every `TaskContext` the session hands -out keep agreeing. - -Repairing the damage instead of avoiding it does not work in general. A context can -rebuild the codecs it holds in its own fields, but a codec already embedded in a -registered `FFI_CatalogProvider` — and in every `FFI_SchemaProvider` and -`FFI_TableProvider` minted from it — is not reachable from Python at all. Nor can the -codec simply retain the session that built it: a codec handed to a provider is routinely -registered straight back into that same session, which would close the cycle -`SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and leak it. - -`SessionContext.enable_url_table` is the one exception. It clones the underlying -`SessionContext`, so the returned context has an allocation of its own and must not -outlive the receiver. 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`, -`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 -visible to every context sharing that session — including ones a `with_*` call returned -earlier. Installing a codec on a session that already has a foreign planner rebuilds -that planner against the new chain for the same reason: there is one planner, and it has -to carry the codecs currently in force. This happens on the shared session, so it takes -effect even if the returned context is discarded — `ctx.with_python_udf_inlining(...)` -whose result is thrown away still leaves the session's planner carrying the codecs of -that discarded handle. A call that changes nothing is exempt: asking for the inlining -setting a context already has returns a handle without touching the session. - -The rule that falls out of this is worth stating on its own, because it is the one thing -that surprises people: - -> The session's query planner carries the codecs of the handle that most recently -> installed one. Every other path — `Expr.to_bytes(ctx)`, `ExecutionPlan.to_bytes(ctx)`, -> registering a provider — uses the codecs of the handle you call it on. - -Those can be different handles, and then one session has two codec chains in effect at -once: - -```python -ctx = ctx.with_logical_extension_codec(codec_a) -ctx.set_query_planner(planner) -ctx.with_logical_extension_codec(codec_b) # discarded - -Expr.to_bytes(expr, ctx) # encodes with [codec_a, default] -- ctx's own field -ctx.sql(...).collect() # plans with [codec_b, codec_a, default] -- the discarded - # handle's chain, installed on the shared session -``` - -Chaining `ctx = ctx.with_...(...)`, as the example below does, keeps the two in step. -`test_the_planner_and_the_handle_can_hold_different_codecs` pins the divergence. - -```python -ctx = SessionContext(config) -ctx = ctx.with_logical_extension_codec(provider_logical_codec) -ctx = ctx.with_physical_extension_codec(provider_physical_codec) -ctx.set_query_planner(planner) -ctx.register_udf(my_udf) -``` - -Order is a readability preference rather than a requirement — installing a codec after a -planner rebuilds the planner against it. - -A session holds exactly one query planner. Calling `set_query_planner` again replaces the -installed planner instead of layering another one. To chain planners, have the new -planner wrap the capsule returned by `SessionContext.__datafusion_query_planner__()`, -captured before the new planner is installed, and delegate to it explicitly. - -### Rebinding a planner's codecs is one level deep - -The rebuild above swaps the codecs on the installed `ForeignQueryPlanner` handle, and -only that handle. A planner that wraps a fallback resolved that fallback when *it* was -installed, and holds the result inside its own library's private data — behind a -`create_physical_plan` function pointer, with no Python-side handle. A codec installed -afterwards therefore reaches the outer planner and not the fallback, which keeps -whichever codecs were in force when it was imported. - -Neither side can repair that: - -- **The host cannot reach it.** `FFI_QueryPlanner::new_with_ffi_codecs` unwraps exactly - one `ForeignQueryPlanner` layer. There is no deeper handle to unwrap — the same - situation as a codec embedded in a registered `FFI_CatalogProvider`. -- **The planner library cannot re-derive it.** `FFI_QueryPlanner` holds its codecs by - value, and `Session` exposes no accessor for the ones the host currently has, so - `create_physical_plan` cannot pick them up from the session it is handed. The rebuild - has to be eager, and an eager rebuild only sees the top layer. - -A fix has to come from upstream, and is tracked in -[apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762). - -The stale codecs stay usable rather than dangling — they hold weak handles to the one -`Arc` that Rule 6 keeps alive — so the effect is a fallback hop -serializing with an older codec, not a failure. It is also invisible to the examples -here, which use one fallback in the same cdylib as its wrapper; `datafusion-ffi` -short-circuits a same-library hop rather than serializing, so no codec runs. A fallback -in a *different* library would serialize, and would do it with the codecs it was -imported with. - -So install the codecs before a layered planner. If a codec has to go in afterwards, -install the outer planner again *on the handle that holds the new codec* — that re-runs -its getter, which re-imports the fallback against that handle's codecs. Re-installing on -the original handle rebinds the session's planner back to the original handle's codecs -instead, which is the trap -`test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` pins. - -`with_extensions` sidesteps this entirely, and for nested planners too: every codec from -every bundle is installed before the first planner hook runs, so no layer — outer or -fallback — is ever captured against a partial chain. There is no "afterwards" within a -call. Prefer it over hand-layering whenever the planners you are composing all ship as -bundles. - -## Alternative Approach - -Suppose you needed to expose some other features of DataFusion and you could not wait -for the upstream repository to implement the FFI approach we describe. In this case -you decide to create your dependency on the `datafusion-python` crate instead. - -As we discussed, this is not guaranteed to work across different compiler versions and -optimization levels. If you wish to go down this route, there are two approaches we -have identified you can use. - -1. Re-export all of `datafusion-python` yourself with your extensions built in. -2. Carefully synchronize your software releases with the `datafusion-python` CI build - system so that your libraries use the exact same compiler, features, and - optimization level. - -We currently do not recommend either of these approaches as they are difficult to -maintain over a long period. Additionally, they require a tight version coupling -between libraries. - -## Status of Work - -At the time of this writing, the FFI features are under active development. To see -the latest status, we recommend reviewing the code in the [datafusion-ffi] crate. - -[datafusion-ffi]: https://crates.io/crates/datafusion-ffi -[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example -[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index a989a068d..45787f98f 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -19,11 +19,17 @@ # Contributor Guide -Guides for contributors to the DataFusion in Python project. +Guides for contributors to the DataFusion in Python project — people changing +this repository itself. + +If you are writing a separate library that plugs into datafusion-python, you +want the {ref}`Extension Guide ` instead. If you are using +one, see {ref}`user_guide_extensions`. ```{toctree} :maxdepth: 2 introduction -ffi +pyo3-guidelines +ffi-internals ``` diff --git a/docs/source/contributor-guide/pyo3-guidelines.md b/docs/source/contributor-guide/pyo3-guidelines.md new file mode 100644 index 000000000..4f81e569c --- /dev/null +++ b/docs/source/contributor-guide/pyo3-guidelines.md @@ -0,0 +1,74 @@ + + +(pyo3_guidelines)= + +# PyO3 binding guidelines + +Conventions for the `#[pyclass]` bindings in `crates/`. These are review +policy for this repository, not part of the extension protocol — an extension +library is free to shape its own classes differently. + +(ffi_pyclass_mutability)= + +## Class mutability + +PyO3 bindings should present immutable wrappers whenever a struct stores shared +or interior-mutable state. In practice this means that any `#[pyclass]` +containing an `Arc>` or similar synchronized primitive must opt into +`#[pyclass(frozen)]` unless there is a compelling reason not to. + +The execution context illustrates the preferred pattern. `PySessionContext` in +{file}`src/context.rs` stays frozen even though it shares mutable state +internally via `SessionContext`. This ensures PyO3 tracks borrows correctly +while Python-facing APIs clone the inner `SessionContext` or return new +wrappers instead of mutating the existing instance in place: + +```rust +#[pyclass(from_py_object, frozen, name = "SessionContext", module = "datafusion", subclass)] +#[derive(Clone)] +pub struct PySessionContext { + pub ctx: SessionContext, +} +``` + +Occasionally a type must remain mutable—for example when PyO3 attribute setters +need to update fields directly. In these rare cases add an inline justification +so reviewers and future contributors understand why `frozen` is unsafe to +enable. `DataTypeMap` in {file}`src/common/data_type.rs` includes such a +comment because PyO3 still needs to track field updates: + +```rust +// TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now +#[derive(Debug, Clone)] +#[pyclass(from_py_object, name = "DataTypeMap", module = "datafusion.common", subclass)] +pub struct DataTypeMap { + #[pyo3(get, set)] + pub arrow_type: PyDataType, + #[pyo3(get, set)] + pub python_type: PythonType, + #[pyo3(get, set)] + pub sql_type: SqlType, +} +``` + +When reviewers encounter a mutable `#[pyclass]` without a comment, they should +request an explanation or ask that `frozen` be added. Keeping these wrappers +frozen by default helps avoid subtle bugs stemming from PyO3's interior +mutability tracking. diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md new file mode 100644 index 000000000..2a0314196 --- /dev/null +++ b/docs/source/extension-guide/bundles.md @@ -0,0 +1,252 @@ + + +(extension_bundles)= + +# Extension bundles + +If your library ships codecs, or a query planner, or both, expose a **bundle** +and let callers install it with +{py:meth}`~datafusion.SessionContext.with_extensions`. This is the recommended +way to package an extension, and the rest of this page explains what the +bundle protocol asks of you and why. + +## Why not have callers install the pieces + +Installing the pieces by hand works, but it makes the caller responsible for +ordering: the codecs have to be installed before the planner, because a planner +is built against whatever codec chains exist when it is installed, and a codec +added afterwards rebinds it. Get that wrong and the planner encodes through a +chain that is missing a library. + +`with_extensions` removes the ordering question. An extension library exposes a +bundle object implementing one or both of two hooks: + +```python +class MyEngineExtension: + def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: + # Phase one. Create fresh components bound to `ctx` on every call. + return SessionExtensionComponents( + logical_extension_codecs=(self._make_logical_codec(ctx),), + physical_extension_codecs=(self._make_physical_codec(ctx),), + ) + + def __datafusion_session_planner__(self, ctx: SessionContext, fallback): + # Phase two. `ctx` now carries every bundle's codecs, and `fallback` is + # the planner built so far. Wrapping it is what makes this library + # compose with the other planners in the call. + return self._make_planner(ctx, fallback=fallback) +``` + +Implement whichever apply: a codec-only library defines the first, a library +that ships only an optimizing planner defines the second. The caller then +writes: + +```python +ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) +ctx.register_table("t", lib_a.TableProvider()) +ctx.register_udf(udf(lib_b.SomeUDF())) +``` + +`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete +Rust implementation of the protocol, including taking the task-context provider +off the supplied context, wrapping its codecs in `BundledLogicalCodec` / +`BundledPhysicalCodec` so they carry declared ids, and constructing a Python +{py:class}`~datafusion.SessionExtensionComponents`. + +(extension_bundles_two_phases)= + +## Two phases, because codecs and planners compose differently + +A session chains **many** codecs and dispatches between them by id. Codecs +therefore just accumulate: order affects encoding only, and decoding always +routes to the codec that wrote the payload. A session holds exactly **one** +query planner, so planners cannot accumulate — they compose by *nesting*, each +wrapping the one before it and delegating to it for work it does not handle. + +So `with_extensions` runs every `__datafusion_session_extension__` and installs +all the codecs, and only then runs each `__datafusion_session_planner__`, in +argument order, handing each the planner built so far. Two consequences worth +holding onto: + +- **Bundle order matters differently for each.** For planners it sets the + nesting: the last extension listed ends up outermost and is consulted first. + For codecs it never affects decoding, and affects encoding only when two + codecs would claim the same node — see {ref}`extension_codec_order`. +- **A planner is always built against the complete codec set**, including + codecs from bundles listed after it. This is what the low-level chaining + cannot give you, and it matters most for a nested planner: the rebuild that + follows a later codec install reaches only the outermost layer (see + {ref}`planner_codec_rebinding`), so a fallback captured before the codecs + were complete would stay stale forever. + +The two hooks therefore see the same session through different chains. Both +receive a handle on the one session, so the task-context provider taken off +either is the same and stays valid — but the `ctx` in phase one still carries +the chains the receiver had, since nothing is installed yet, while the `ctx` in +phase two carries every bundle's codecs. A bundle that reads the host's codec +chains — `MyPlannerExtension` does, to give its planner the host's codecs +rather than minting its own — must do that in the planner hook. Reading them in +phase one gets the chains from before the call, missing even the bundle's own +codecs. + +## Returning a planner, or not + +An extension that ignores `fallback` and returns an unrelated planner replaces +every layer beneath it, including any planner the session already had. That is +legal — a library that must be the only planner does it deliberately — but it +is not composable, and nothing detects it. Returning `None` contributes no +planner and leaves `fallback` in place. + +`None` is the no-op, not `fallback`. The capsule handed to the first bundle +wraps the session's planner for export, so returning it unchanged installs that +planner as a foreign one and every plan built afterwards crosses an FFI +boundary it did not before. A bundle that decides at runtime it has nothing to +contribute returns `None`. + +Three libraries that each ship a planner therefore install like this, with the +outermost last: + +```python +ctx = SessionContext(config).with_extensions( + tables.Extension(), # codecs only + functions.Extension(), # codecs only + optimizer.Extension(), # planner, wrapping the session default + distributed.Extension(), # planner, wrapping the optimizer +) +``` + +(extension_bundles_order_conflict)= + +## When the two orders conflict + +Because codec position and planner position both come from one argument list, a +library can in principle need to be early for one and late for the other: its +codec must precede a broad claimer, while its planner must nest outside that +library's planner. + +Do not try to satisfy both by reordering — contribute each half at its own +position. The two hooks are independent, so a three-line adapter each is +enough: + +```python +class CodecsOf: + """Contribute only the codec half of a bundle, at this position.""" + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_extension__(self, ctx): + return self.inner.__datafusion_session_extension__(ctx) + + +class PlannerOf: + """Contribute only the planner half of a bundle, at this position.""" + def __init__(self, inner): + self.inner = inner + + def __datafusion_session_planner__(self, ctx, fallback): + return self.inner.__datafusion_session_planner__(ctx, fallback) + + +ctx = SessionContext(config).with_extensions( + CodecsOf(engine), CodecsOf(tables), # engine's codec first + PlannerOf(tables), PlannerOf(engine), # engine's planner outermost +) +``` + +This keeps everything `with_extensions` guarantees: one transaction, codecs +complete before any planner is built, codec ids untouched — an id is read off +the codec object, not off the extension that contributed it, so splitting a +bundle cannot re-tag its payloads. A library that expects to be composed this +way should expose the halves itself rather than make callers write the +adapters. + +There is no attempt here to make every permutation expressible from one call. +Two positions per bundle covers the cases that arise; anything stranger is a +sign the libraries disagree about what they own, which is better fixed there. + +The low-level +{py:meth}`~datafusion.SessionContext.with_logical_extension_codec` / +{py:meth}`~datafusion.SessionContext.with_physical_extension_codec` / +{py:meth}`~datafusion.SessionContext.set_query_planner` sequence also works, +and it is the right answer when the pieces do not come as bundles at all. But +it is a real downgrade, not just a more verbose spelling: you take back +responsibility for installing every codec before every planner, and a planner +you layer by hand keeps the codecs it captured — the +{ref}`one-level rebind ` does not reach inside it. +Reach for it last. + +(extension_bundles_codecs_are_objects)= + +## Codecs are objects, not capsules + +`with_extensions` requires each codec to be an object exposing the capsule +getter, and refuses a bare `PyCapsule`. A codec's id is read off the object it +is handed over as, and a capsule has no type to read one from; since this +method takes no `codec_id=`, there would be nothing left to name it by. A +library holding a raw capsule — which is what a Rust implementation has — +wraps it: + +```python +class MyLogicalCodec: + # Optional. Without it the id is this class's import path, which is already + # stable; declare it if you may rename the class and need old plans to decode. + __datafusion_codec_id__ = "my_library.logical.v1" + + def __init__(self, capsule): + self._capsule = capsule + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule +``` + +Wrapping is not just bookkeeping. It ties the id to the codec rather than to +the bundle that contributed it, and that difference is load-bearing: an +application commonly presents several libraries as one bundle of its own, and +the id has to survive that. Were the id taken from the contributing bundle, +wrapping `my_engine.Extension` inside `my_app.Extension` would silently re-tag +the engine's payloads, and a scheduler that installs the engine's codec by its +documented id would fail to decode plans from composed clients while succeeding +for direct ones. The wrapper travels with the codec; the bundle does not. + +The query planner is exempt — it carries no wire id, so it may be an object or +a capsule. + +## Failure and rollback + +Nothing is written to the session until every factory has returned and every +capsule has been validated, so a factory that raises leaves the session exactly +as it was. A factory that mutates the context it is handed — registering a +table, say — is **not** rolled back, which is why bundle objects must be +configuration-only: create fresh components on each call, never cache bound +components, and do not retain the context passed in. + +Like every other derivation, the returned context is a handle on the *same* +session as the receiver — see {ref}`extension_sessions`. Only the Python-side +codec chains belong to the returned handle; the planner is installed on the +shared session and takes effect even if that handle is discarded. + +`with_extensions` sidesteps the {ref}`one-level rebind ` +entirely, and for nested planners too: every codec from every bundle is +installed before the first planner hook runs, so no layer — outer or fallback — +is ever captured against a partial chain. There is no "afterwards" within a +call. Prefer it over hand-layering whenever the planners you are composing all +ship as bundles. + +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/extension-guide/capsule-protocol.md b/docs/source/extension-guide/capsule-protocol.md new file mode 100644 index 000000000..ada5a6f57 --- /dev/null +++ b/docs/source/extension-guide/capsule-protocol.md @@ -0,0 +1,224 @@ + + +(extension_capsule_protocol)= + +# The capsule protocol + +Every integration point in this section works the same way. Your library +exposes a dunder method, datafusion-python calls it, and it hands back a +`PyCapsule` wrapping an FFI-safe struct. This page describes that mechanism +once; the pages after it describe what goes inside the capsule for each kind +of component. + +The bulk of the code necessary to perform our FFI operations is in the +upstream [DataFusion](https://datafusion.apache.org/) core repository. You can +review the code and documentation in the [datafusion-ffi] crate. + +## FFI-safe types + +Our FFI implementation is narrowly focused on sharing data and functions with +Rust backed libraries. This allows us to use the +[abi_stable crate](https://crates.io/crates/abi_stable). This is an excellent +crate that allows for easy conversion between Rust native types and FFI-safe +alternatives. For example, if you needed to pass a `Vec` via FFI, you +can simply convert it to an `RVec` in an intuitive manner. It also +supports features like `RResult` and `ROption` that do not have an obvious +translation to a C equivalent. + +## `FFI_` on the provider, `Foreign` on the receiver + +The [datafusion-ffi] crate has been designed to make it easy to convert from +DataFusion traits into their FFI counterparts. For example, if you have +defined a custom +[TableProvider](https://docs.rs/datafusion/45.0.0/datafusion/catalog/trait.TableProvider.html) +and you want to create a sharable FFI counterpart, you could write: + +```rust +let my_provider = MyTableProvider::default(); +let ffi_provider = FFI_TableProvider::new(Arc::new(my_provider), false, None); +``` + +If you were interfacing with a library that provided the above +`FFI_TableProvider` and you needed to turn it back into a `TableProvider`, you +can turn it into a `ForeignTableProvider`, which implements the `TableProvider` +trait: + +```rust +let foreign_provider: ForeignTableProvider = ffi_provider.into(); +``` + +If you review the code in [datafusion-ffi] you will find that each of the +traits we share across the boundary has two portions, one with an `FFI_` +prefix and one with a `Foreign` prefix. This is used to distinguish which side +of the FFI boundary that struct is designed to be used on. The structures with +the `FFI_` prefix are to be used on the **provider** of the structure. In the +example we're showing, this means the code that has written the underlying +`TableProvider` implementation to access your custom data source. The +structures with the `Foreign` prefix are to be used by the receiver. In this +case, it is the `datafusion-python` library. + +## Wrapping it in a capsule + +In order to share these FFI structures, we need to wrap them in some kind of +Python object that can be used to interface from one package to another. As +described in {ref}`extension_why_ffi`, we use `PyCapsule`. We can create a +`PyCapsule` for our provider thusly: + +```rust +let name = CString::new("datafusion_table_provider")?; +let my_capsule = PyCapsule::new_bound(py, provider, Some(name))?; +``` + +On the receiving side, turn this pycapsule object into the +`FFI_TableProvider`, which can then be turned into a `ForeignTableProvider`; +the associated code is: + +```rust +let capsule = capsule.cast::()?; +let data: NonNull = capsule + .pointer_checked(Some(name))? + .cast(); +let codec = unsafe { data.as_ref() }; +``` + +## The naming rule + +The getter's name and the capsule's name are both fixed by the protocol, and +they follow one rule with no exceptions: + +- The method is `__datafusion___`. +- The capsule it returns is named `datafusion_` — the same string + without the underscores. + +So a table provider is reached by calling `__datafusion_table_provider__` and +must return a capsule named `datafusion_table_provider`. Return a capsule with +the wrong name and the import fails with an error naming both the name found +and the name expected, rather than reading the pointer as the wrong type. + +The full list of hooks and their capsule names is in the +{ref}`hook reference `. `TableProvider` was the first +extension written this way and is the most thoroughly implemented; every hook +added since follows the same pattern. + +## Version checking + +Objects imported through this protocol are checked against the major version of +`datafusion-ffi` that datafusion-python was built with. A component produced by +a library built against a different DataFusion major version raises an +`ImportError` naming the version found and the version expected. + +This is a diagnostic rather than a soundness guarantee — reading the version +out of the struct already assumes the local field layout — but it turns the +common "extension library built against the wrong DataFusion" mistake into a +clear message rather than undefined behaviour on first use. See +{ref}`extension_version_mismatch`. + +Three FFI structs carry no version field and so cannot be checked: +`FFI_TaskContextProvider`, `FFI_TableProviderFactory`, and +`FFI_ExtensionOptions`. + +(extension_getter_argument)= + +## What your getter receives + +Most getters take one positional argument beyond `py`. The +{ref}`hook reference ` says which, and the group that does +looks like this: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +This exists because the FFI constructors need things an extension library does +not have. `FFI_{Logical,Physical}ExtensionCodec::new` needs a +`TaskContextProvider` for the decode callbacks the codec will receive, and +`FFI_QueryPlanner::new` needs both codecs on top of that. Taking them from the +argument is what keeps a library from constructing a `SessionContext` purely to +satisfy a parameter — an empty one resolves nothing, and +`FFI_TaskContextProvider` holds it weakly, so a context built inline in the +getter is already dropped by the time the capsule is used. + +### It is not always a session + +The parameter is conventionally named `session`, and for the codec and planner +hooks it genuinely is one. For the provider and catalog hooks it may instead be +a bare `datafusion_logical_extension_codec` capsule: `SessionContext.register_table` +passes the session, while `Schema.register_table`, +`SessionContext.register_catalog_provider`, `register_catalog_provider_list`, +and `register_table_factory` pass the host's codec directly. + +This is why the helpers accept either. `ffi_logical_codec_from_pycapsule` +calls the codec getter if the object has one and returns the object untouched +if it does not, so the same line works in both cases: + +```rust +let codec = ffi_logical_codec_from_pycapsule(session, None)?; +``` + +The rule to hold onto is therefore about capability rather than type: **the +argument is something you can read the host's logical extension codec off.** +The hooks that need more than a codec — the two extension codecs and the query +planner, which need a task-context provider — are exactly the hooks that are +always handed a real session. + +### Duck-type it + +Do not check the argument's type. Beyond the codec-capsule case above, even +when it *is* a session it is the PyO3 context the binding installs through, not +the `datafusion.context.SessionContext` wrapper. It carries every capsule +getter and `__datafusion_codec_id__` — everything the protocol asks of it — but +`isinstance(session, SessionContext)` is `False` in Python even though its +`repr` reads `datafusion.SessionContext`. + +The two bundle hooks are the exception: `__datafusion_session_extension__` and +`__datafusion_session_planner__` are dispatched from Python by +{py:meth}`~datafusion.SessionContext.with_extensions`, so they receive the +wrapper. See {doc}`bundles`. + +`SessionContext` accepts the argument on its own codec and planner getters and +ignores it, so a session satisfies the same protocol an extension library +implements. When you export the current planner in order to wrap it, +`ctx.__datafusion_query_planner__()` and `ctx.__datafusion_query_planner__(ctx)` +are both fine. + +(extension_task_context_provider)= + +## `__datafusion_task_context_provider__` + +This is the one hook implemented by the **host** rather than by your library. +You never define it; you read it off the session you were handed, which is what +`ffi_task_context_provider_from_pycapsule` does. Nothing in datafusion-python +calls it on a foreign object. + +Taking the host's provider means your decode callbacks resolve names against +the session that is actually running the query — see +{ref}`extension_codec_decode_session` — and it removes any need for your +library to construct a `SessionContext` of its own. + +[datafusion-ffi]: https://crates.io/crates/datafusion-ffi diff --git a/docs/source/extension-guide/checklist.md b/docs/source/extension-guide/checklist.md new file mode 100644 index 000000000..944fa6abf --- /dev/null +++ b/docs/source/extension-guide/checklist.md @@ -0,0 +1,99 @@ + + +(extension_checklist)= + +# Extension author checklist + +The rules in this section, gathered into one list to run through before you +publish. Each links to the page that explains it. + +## Protocol + +- [ ] **Every getter's capsule name matches its method name.** + `__datafusion___` returns a capsule named `datafusion_`. + → {ref}`extension_capsule_protocol` +- [ ] **Your getter does not inspect its argument.** Pass it to + `ffi_logical_codec_from_pycapsule` and move on. It is not always a + session, and when it is, it is not the Python `SessionContext` wrapper. + → {ref}`extension_getter_argument` +- [ ] **You never construct a `SessionContext` inside your library.** Take + what the FFI constructors need off the argument you were handed. A + context built inline is already dropped by the time the capsule is used. + → {ref}`extension_getter_argument` +- [ ] **You do not depend on the `datafusion-python` crate.** + → {ref}`extension_why_ffi` + +## Codecs + +- [ ] **Your codec claims narrowly.** Downcast to your own types. Claiming a + broad category takes nodes from every library installed after you and + makes your library order-sensitive for everyone downstream. + → {ref}`extension_codec_order` +- [ ] **You declare `__datafusion_codec_id__` if you might rename the class.** + The default id is the exporting class's import path, so a rename stops + older plans decoding. → {ref}`extension_codec_ids` +- [ ] **Name-only decoders check `name` before trusting `buf`.** An empty + payload has no id to route on, so your `try_decode_udf` can be called + with another library's function name and an empty buffer. + → {ref}`extension_codecs` +- [ ] **You round-trip a plan in a test and assert *your* codec did the work.** + Both being installed does not mean your node reached you. + → {ref}`extension_codec_order` + +## Bundles and planners + +- [ ] **You ship a bundle, not loose pieces**, if you have codecs or a planner. + → {ref}`extension_bundles` +- [ ] **Your bundle is configuration-only.** Fresh components on every call, + no cached bound components, no retaining the context passed in, no + registering anything on it — a factory that mutates the context is not + rolled back if a later factory raises. + → {ref}`extension_bundles` +- [ ] **Your codecs are objects exposing the getter, not bare capsules.** + `with_extensions` refuses a capsule, because there would be nothing to + name the codec by. → {ref}`extension_bundles_codecs_are_objects` +- [ ] **Your planner hook wraps `fallback` and delegates to it.** Ignoring it + replaces every layer beneath you, which is legal but not composable. + → {ref}`extension_bundles` +- [ ] **Your planner hook returns `None`, not `fallback`, when it has nothing + to contribute.** Returning `fallback` installs the session's own planner + as a foreign one and adds an FFI hop that was not there. + → {ref}`extension_bundles` +- [ ] **You read the host's codec chains in the planner hook, not the extension + hook.** Phase one runs before anything is installed. + → {ref}`extension_bundles_two_phases` +- [ ] **If you also offer the low-level path**, document that codecs go in + before a layered planner. → {ref}`planner_codec_rebinding` + +## Packaging and documentation + +- [ ] **You state which `datafusion` version your release requires.** A + mismatch raises an `ImportError` on import, which is a good failure — but + only if your users know what to install. → {ref}`extension_version_mismatch` +- [ ] **You tell your users to keep a context alive** for as long as anything + derived from it is in use. This is the rule most likely to arrive as a + bug report against your library. → {ref}`extension_sessions` +- [ ] **Your production codec serializes durable metadata**, not a + process-local token. The examples in this repository use tokens to make + ownership observable; that is a demonstration, not a pattern. + → {ref}`extension_guide` +- [ ] **You have integration tests across a real FFI boundary.** The two + example crates in this repository are the pattern: build the cdylib, + install the wheel, then exercise it from Python. diff --git a/docs/source/extension-guide/codecs.md b/docs/source/extension-guide/codecs.md new file mode 100644 index 000000000..c7f051324 --- /dev/null +++ b/docs/source/extension-guide/codecs.md @@ -0,0 +1,166 @@ + + +(extension_codecs)= + +# Extension codecs + +A codec is what lets your objects survive being serialized into a plan and +rebuilt somewhere else — another process, or another program. If your library +contributes table providers, functions, or execution plan nodes and those plans +have to leave the process, you need one. + +Codecs are contributed to a session either through an +{ref}`extension bundle `, which is the recommended route, +or one at a time through +{py:meth}`~datafusion.SessionContext.with_logical_extension_codec` and +{py:meth}`~datafusion.SessionContext.with_physical_extension_codec`. + +## Codecs compose + +Each call to `with_logical_extension_codec` or +`with_physical_extension_codec` **appends** the codec to the session's codec +chain rather than replacing prior codecs. One session can therefore carry +codecs from several independent libraries at once. + +**Nothing is asked of the codec itself.** Implement `LogicalExtensionCodec` or +`PhysicalExtensionCodec` exactly as you would for a session that installs only +yours. When your codec writes bytes into a serialized plan, datafusion-python +records which codec wrote them, and strips that record off again before handing +the bytes back. So your codec receives, byte for byte, the payload it wrote, +and is never offered a payload another codec wrote. + +A codec that also ships to hosts which dispatch differently may still want its +own guard against foreign payloads. Keeping one is fine; it is simply not +needed for the datafusion-python path. + +(extension_codec_ids)= + +## Codec ids + +That record is the codec's **id**: a short string stored inside the plan, +naming the codec that wrote each payload. Because plans are decoded in another +process — or another program — the id has to name the same codec there as it +did where the plan was written. + +Ids are assigned for you. A codec's id is normally its exporting class's import +path, such as `my_library.Codec`, which is what you will see in +{py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` and in decode +errors. You choose one yourself in three cases: + +- **Two instances of one class.** Both get the same id, so the second install + raises `ValueError`. Pass `codec_id=` to tell them apart. +- **A bare `PyCapsule`.** A capsule has no class to take a name from, so + installing one through `with_logical_extension_codec` or + `with_physical_extension_codec` gives it an id private to the session that + installed it; plans it encodes fail with a clear error on any other session + rather than being decoded by the wrong codec. Pass `codec_id=` if those plans + have to cross sessions. + + {py:meth}`~datafusion.SessionContext.with_extensions` takes no `codec_id=`, + so it refuses a bare capsule outright and tells you to wrap it. See + {ref}`extension_bundles_codecs_are_objects`. +- **A class you intend to rename.** The id follows the class name, so renaming + stops older plans from decoding. Declare `__datafusion_codec_id__` on the + exporting object to pin an id that survives the rename. + +{py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` and its +physical counterpart list the ids installed on a session, which is also what a +decode failure names. + +Installing one context's codec stack on another session composes the two +sessions rather than copying codecs out of one: the imported codecs resolve +their task context against the original and stop working when it is dropped — +see {ref}`ffi_internals_one_arc`. Pass the context itself rather than the +capsule it exports, so its codecs get an id that other sessions can decode. + +## Functions whose name is the whole encoding + +A codec may own functions that need no payload at all, where the name is the +whole encoding: `try_encode_udf` writes nothing and `try_decode_udf` rebuilds +the function from `name`. That is supported and needs no id, because an `Ok` +with an empty buffer is read as "no opinion" and passes the object to the next +codec. `NameOnlyUdfCodec` in [`datafusion-ffi-example`] is the worked case. +Anything no installed codec claims falls through to +`Default{Logical,Physical}ExtensionCodec`. + +This is the one case where your decoder is consulted about something you may +not own, because an empty payload has no id to route on. `try_decode_udf` and +its aggregate and window siblings can therefore be called with an empty `buf` +and a `name` belonging to another library. Decide from `name` and return an +error if it is not yours; do not assume `buf` is non-empty. + +:::{note} +The framing itself — how an id is stored alongside a payload and routed back, +and the two cases that stay unframed — is internal to datafusion-python and +documented in `crates/core/src/codec.rs` for anyone changing it. +::: + +The current FFI logical codec supports providers and UDFs but not arbitrary +custom `LogicalPlan::Extension` nodes. See both example READMEs for the +supported flow and local build commands. + +(extension_codec_order)= + +## When codec order matters + +Decoding is never order-dependent: a payload names its codec by id and the +chain dispatches straight to it. Encoding walks the chain in install order and +stops at the first codec that claims the node. Most of the time that is +invisible, because libraries claim disjoint things — one owns its table +providers, another its UDFs, a third its own execution plan nodes. + +It stops being invisible when a codec claims *broadly*. A node that came from +another library arrives as an opaque `ForeignExecutionPlan`, and a codec that +claims any of those will take nodes it does not own from any library installed +after it. The query still succeeds. What changes is which library wrote the +bytes — so a plan that has to decode in another process now needs whichever +library happened to win, not the one whose node it is. +`MyPhysicalExtensionCodec` in [`datafusion-ffi-example`] claims this way, and +`test_a_greedy_codec_installed_first_claims_another_librarys_node` pins the +consequence. + +Two rules of thumb: + +- **Writing a codec, claim narrowly.** Downcast to your own types. Claiming a + broad category makes your library order-sensitive for everyone downstream of + it. +- **Shipping plans out of the process, verify.** Do not assume your node + reached your codec just because both are installed. Round-trip a plan through + {py:meth}`ExecutionPlan.to_bytes ` / + {py:meth}`~datafusion.ExecutionPlan.from_bytes` in a test and assert your + codec did the work. + +(extension_codec_decode_session)= + +## A codec decodes against the session running the query + +Because the task-context provider comes from the host — see +{ref}`extension_getter_argument` — a decode callback running inside an +extension library resolves names against the session running the query. A +function registered with `ctx.register_udf(...)` is visible to a foreign codec +decoding a node that references it by name, and the handle is live rather than +a snapshot, so a registration made after the codec is installed is visible too. + +This is covered in +`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`, +where the example codecs take a `require_udf_on_decode` name and resolve it out +of the task context they are handed. + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/functions.md b/docs/source/extension-guide/functions.md new file mode 100644 index 000000000..92a30b7d9 --- /dev/null +++ b/docs/source/extension-guide/functions.md @@ -0,0 +1,153 @@ + + +(extension_functions)= + +# Functions and table functions + +Four hooks contribute functions written in Rust. Users can also define +functions in pure Python — see +{doc}`../user-guide/common-operations/udf-and-udfa` — and the two roads meet at +the same registration methods. + +| Hook | Contributes | Wrapped by | Registered with | +| --- | --- | --- | --- | +| `__datafusion_scalar_udf__` | scalar function | {py:func}`datafusion.udf` | {py:meth}`~datafusion.SessionContext.register_udf` | +| `__datafusion_aggregate_udf__` | aggregate function | {py:func}`datafusion.udaf` | {py:meth}`~datafusion.SessionContext.register_udaf` | +| `__datafusion_window_udf__` | window function | {py:func}`datafusion.udwf` | {py:meth}`~datafusion.SessionContext.register_udwf` | +| `__datafusion_table_function__` | function returning a table | {py:func}`datafusion.udtf` | {py:meth}`~datafusion.SessionContext.register_udtf` | + +All four are implemented in [`datafusion-ffi-example`], one per file. + +## The three scalar-shaped hooks + +The scalar, aggregate, and window getters take **no argument** beyond `py` — +they need no codec and no task-context provider, because a function is +identified by name and signature rather than by anything session-scoped: + +```rust +#[pymethods] +impl MyScalarUDF { + fn __datafusion_scalar_udf__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let udf = Arc::new(ScalarUDF::from(self.clone())); + let ffi = FFI_ScalarUDF::from(udf); + + PyCapsule::new_with_value(py, ffi, cr"datafusion_scalar_udf") + } +} +``` + +Aggregate and window follow identically with `FFI_AggregateUDF` / +`FFI_WindowUDF` and the matching capsule names. + +Your users wrap the object once and register the result: + +```python +from datafusion import udf + +ctx.register_udf(udf(my_library.MyScalarUDF())) +``` + +## Table functions + +A table function takes literal `Expr` arguments and returns a table provider, +so it needs the host's logical codec the way a +{ref}`table provider ` does: + +```rust +fn __datafusion_table_function__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let func = self.clone(); + let codec = ffi_logical_codec_from_pycapsule(session, None)?; + let provider = FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, codec); + + PyCapsule::new_with_value(py, provider, cr"datafusion_table_function") +} +``` + +Only literal expressions are supported as arguments. The Python side is +described under +{doc}`Table Functions <../user-guide/common-operations/udf-and-udfa>`. + +## Serializing functions + +A function that appears in a plan leaving the process has to be reconstructible +on the far side. Functions are the one case where a codec often needs **no +payload at all**: the name is the whole encoding, `try_encode_udf` writes +nothing, and `try_decode_udf` rebuilds the function from `name`. See +{ref}`extension_codecs` for how that works and for the one obligation it puts +on your decoder — with an empty payload there is no id to route on, so your +`try_decode_udf` can be called with a `name` belonging to another library. + +`NameOnlyUdfCodec` in [`datafusion-ffi-example`] is the worked case. + +(extension_other_hooks)= + +## Other session components + +Two further hooks contribute things that are neither data nor functions. Both +take no argument and both are implemented in [`datafusion-ffi-example`]. + +**`__datafusion_physical_optimizer_rule__`** contributes a rule that rewrites +physical plans, installed with +{py:meth}`~datafusion.SessionContext.add_physical_optimizer_rule`. Reach for +this rather than a {doc}`query planner ` when you want to +adjust the plan DataFusion produced rather than produce it yourself — it is +much the smaller commitment, and rules accumulate where planners nest. + +```rust +fn __datafusion_physical_optimizer_rule__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let rule: Arc = Arc::new(self.clone()); + let runtime = get_tokio_runtime().handle().clone(); + let ffi = FFI_PhysicalOptimizerRule::new(rule, Some(runtime)); + + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_optimizer_rule") +} +``` + +**`__datafusion_extension_options__`** contributes typed configuration entries +that your components can read back out of the session config, installed with +{py:meth}`SessionConfig.with_extension `. +`FFI_ExtensionOptions` carries no version field, so it is one of the three +components that cannot be version-checked on import. + +```rust +fn __datafusion_extension_options__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let mut config = FFI_ExtensionOptions::default(); + config + .add_config(self) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + + PyCapsule::new_with_value(py, config, cr"datafusion_extension_options") +} +``` + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/index.md b/docs/source/extension-guide/index.md new file mode 100644 index 000000000..62e8a0fc0 --- /dev/null +++ b/docs/source/extension-guide/index.md @@ -0,0 +1,131 @@ + + +(ffi)= +(extension_guide)= + +# Extension Guide + +This section is for people **writing** a library that plugs into +datafusion-python — a package that contributes table providers, functions, a +catalog, extension codecs, or a query planner, usually written in Rust and +exposed through [PyO3](https://pyo3.rs). + +Two neighbouring audiences are served elsewhere: + +- **Using** an extension library someone else published: + {ref}`user_guide_extensions`. +- **Changing datafusion-python itself**, including the framing that makes this + protocol work: {doc}`../contributor-guide/index`. + +The protocol described here is a public, versioned contract. When a hook's +signature changes, the change is documented in +{doc}`../user-guide/upgrade-guides` with a before-and-after, and objects built +against a different DataFusion major version are rejected with a clear error +rather than used as-is. + +## Start here + +If you have not built an extension before, read {doc}`why-ffi` and +{doc}`capsule-protocol` in order. Together they explain why your library must +not depend on the `datafusion-python` crate, and the one convention every hook +follows. After that the pages are independent — go to the one matching what +you are contributing. + +## The three roles in a query + +A single query can involve three independent native libraries, and much of +this section only makes sense once they are distinct in your head: + +- **datafusion-python** — the host. Owns the session, and decodes whatever + comes back from the other two. +- **A provider library** — owns table providers, catalogs, and functions, plus + the codecs that serialize them. +- **A planner library** — owns a query planner and the configuration it needs. + +The worked examples in this repository use two separate crates, +[`datafusion-ffi-example`] and [`datafusion-ffi-query-planner-example`], so +each role has a distinct shared-library identity. A real library may play more +than one role; keeping them separate in the examples is what makes the +boundaries observable. + +The session owns the codecs used for the exchange and supplies them to the +foreign planner. That is what lets the planner decode provider-owned objects, +and lets datafusion-python decode the physical plan the planner returns. + +:::{note} +The example codecs use process-local tokens to demonstrate ownership. +A production codec should serialize durable metadata instead. +::: + +## Hook reference + +Every integration point is a dunder method named `__datafusion_*__`. The +convention is uniform enough to be worth stating once: your object exposes the +getter, datafusion-python calls it, and it returns a `PyCapsule` wrapping an +FFI-safe struct. See {doc}`capsule-protocol` for what that means and +{ref}`extension_getter_argument` for the argument every getter in the middle +group receives. + +| Hook | Capsule name | Argument | Documented on | +| --- | --- | --- | --- | +| `__datafusion_table_provider__` | `datafusion_table_provider` | codec source | {doc}`table-providers` | +| `__datafusion_table_provider_factory__` | `datafusion_table_provider_factory` | codec source | {doc}`table-providers` | +| `__datafusion_catalog_provider__` | `datafusion_catalog_provider` | codec source | {doc}`table-providers` | +| `__datafusion_catalog_provider_list__` | `datafusion_catalog_provider_list` | codec source | {doc}`table-providers` | +| `__datafusion_schema_provider__` | `datafusion_schema_provider` | codec source | {doc}`table-providers` | +| `__datafusion_table_function__` | `datafusion_table_function` | session | {doc}`functions` | +| `__datafusion_scalar_udf__` | `datafusion_scalar_udf` | none | {doc}`functions` | +| `__datafusion_aggregate_udf__` | `datafusion_aggregate_udf` | none | {doc}`functions` | +| `__datafusion_window_udf__` | `datafusion_window_udf` | none | {doc}`functions` | +| `__datafusion_logical_extension_codec__` | `datafusion_logical_extension_codec` | session | {doc}`codecs` | +| `__datafusion_physical_extension_codec__` | `datafusion_physical_extension_codec` | session | {doc}`codecs` | +| `__datafusion_codec_id__` | *not a capsule — a string attribute* | — | {doc}`codecs` | +| `__datafusion_query_planner__` | `datafusion_query_planner` | session | {doc}`query-planners` | +| `__datafusion_session_extension__` | *not a capsule — returns components* | `ctx` | {doc}`bundles` | +| `__datafusion_session_planner__` | `datafusion_query_planner`, or `None` | `ctx`, `fallback` | {doc}`bundles` | +| `__datafusion_physical_optimizer_rule__` | `datafusion_physical_optimizer_rule` | none | {ref}`extension_other_hooks` | +| `__datafusion_extension_options__` | `datafusion_extension_options` | none | {ref}`extension_other_hooks` | +| `__datafusion_task_context_provider__` | `datafusion_task_context_provider` | none | {ref}`extension_task_context_provider` | + +Two rows are not like the others. `__datafusion_task_context_provider__` is +implemented by the **host**, not by your library — you read it off the session +you are handed. And `__datafusion_codec_id__` is a plain string attribute +rather than a method returning a capsule. + +"codec source" in the argument column means the value is something you can +read the host's logical extension codec off, which is not always a session. +{ref}`extension_getter_argument` explains why, and what to do with it. + +```{toctree} +:maxdepth: 2 + +why-ffi +capsule-protocol +table-providers +functions +codecs +bundles +query-planners +sessions +checklist +``` + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/extension-guide/query-planners.md b/docs/source/extension-guide/query-planners.md new file mode 100644 index 000000000..a97214bda --- /dev/null +++ b/docs/source/extension-guide/query-planners.md @@ -0,0 +1,106 @@ + + +(extension_planners)= + +# Query planners + +A query planner turns a logical plan into a physical one. Contributing your own +is how a library changes the way queries execute rather than what data they can +reach — a distributed engine is the motivating case, and an optimizing rewriter +is another. + +The session owns the codecs used for the exchange and supplies them to your +planner. That is what lets the planner decode provider-owned objects, and lets +datafusion-python decode the physical plan the planner returns. Your planner +uses `FFI_QueryPlanner::new_with_ffi_codecs` with the two codecs it takes off +the session, and never touches a task-context provider directly. That also +matches what installation does anyway: +{py:meth}`~datafusion.SessionContext.set_query_planner` builds the planner +against the codecs of the session that will run the query. + +`MyQueryPlanner` in [`datafusion-ffi-query-planner-example`] is the worked +implementation. + +## One planner per session + +A session holds exactly one query planner. Calling `set_query_planner` again +**replaces** the installed planner instead of layering another one. + +To chain planners, have the new planner wrap the capsule returned by +{py:meth}`SessionContext.__datafusion_query_planner__ `, +captured before the new planner is installed, and delegate to it explicitly: + +```python +fallback = ctx.__datafusion_query_planner__() +ctx.set_query_planner(MyPlanner(fallback=fallback)) +``` + +`set_query_planner` returns nothing. The query planner lives in `SessionState`, +so it is a property of the session rather than of a handle on it, and +installing one is visible to every context sharing that session. See +{ref}`extension_sessions`. + +If the planners you are composing all ship as +{ref}`bundles `, prefer `with_extensions` — it does the +nesting for you and cannot capture a partial codec chain. + +(planner_codec_rebinding)= + +## Install codecs before a layered planner + +Installing a codec on a session that already has a foreign planner rebuilds +that planner against the new chain: there is one planner, and it has to carry +the codecs currently in force. But that rebuild swaps the codecs on the +installed `ForeignQueryPlanner` handle, and **only that handle**. + +A planner that wraps a fallback resolved that fallback when *it* was installed, +and holds the result inside its own library's private data — behind a +`create_physical_plan` function pointer, with no Python-side handle. A codec +installed afterwards therefore reaches the outer planner and not the fallback, +which keeps whichever codecs were in force when it was imported. + +The stale codecs stay usable rather than dangling — they hold weak handles to +the one `Arc` the session keeps alive — so the effect is a +fallback hop serializing with an older codec, not a failure. Neither side can +repair it; the reasons are in {ref}`ffi_internals_rebinding`, and a fix has to +come from upstream +([apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762)). + +So, three rules: + +- **Install the codecs before a layered planner.** +- If a codec has to go in afterwards, install the outer planner again *on the + handle that holds the new codec* — that re-runs its getter, which re-imports + the fallback against that handle's codecs. Re-installing on the original + handle rebinds the session's planner back to the original handle's codecs + instead, which is the trap + `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` + pins. +- Better, use {ref}`extension_bundles`, where there is no "afterwards" within a + call. + +:::{note} +This is invisible in the examples here, which use one fallback in the same +cdylib as its wrapper; `datafusion-ffi` short-circuits a same-library hop +rather than serializing, so no codec runs. A fallback in a *different* library +would serialize, and would do it with the codecs it was imported with. +::: + +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/extension-guide/sessions.md b/docs/source/extension-guide/sessions.md new file mode 100644 index 000000000..c4f0bdb01 --- /dev/null +++ b/docs/source/extension-guide/sessions.md @@ -0,0 +1,118 @@ + + +(extension_sessions)= + +# Sessions, handles, and lifetimes + +One `SessionContext` in Python is a *handle* on a session, not the session +itself. Several handles can share one session, and which handle you call +something on sometimes matters. This page is the set of rules that follow from +that — the ones an extension author trips over. + +## The rule that surprises people + +> The session's query planner carries the codecs of the handle that most +> recently installed one. Every other path — `Expr.to_bytes(ctx)`, +> `ExecutionPlan.to_bytes(ctx)`, registering a provider — uses the codecs of +> the handle you call it on. + +Those can be different handles, and then one session has two codec chains in +effect at once: + +```python +ctx = ctx.with_logical_extension_codec(codec_a) +ctx.set_query_planner(planner) +ctx.with_logical_extension_codec(codec_b) # discarded + +Expr.to_bytes(expr, ctx) # encodes with [codec_a, default] -- ctx's own field +ctx.sql(...).collect() # plans with [codec_b, codec_a, default] -- the discarded + # handle's chain, installed on the shared session +``` + +Chaining `ctx = ctx.with_...(...)` keeps the two in step, which is why every +example in this guide does. +`test_the_planner_and_the_handle_can_hold_different_codecs` pins the +divergence. + +## What a derived context shares + +{py:meth}`~datafusion.SessionContext.with_logical_extension_codec`, +{py:meth}`~datafusion.SessionContext.with_physical_extension_codec`, +{py:meth}`~datafusion.SessionContext.with_python_udf_inlining`, and +{py:meth}`~datafusion.SessionContext.with_extensions` return a new +`SessionContext` wrapping the *same* underlying session. Only the Python-side +codec settings differ; catalogs, tables, registered functions, and +configuration are the one shared session, so a registration on either side is +visible to both. + +There is one `Arc` per session, which is what makes the weak +task-context-provider scheme work: a component bound through any handle stays +valid while *any* handle on that session is alive, so there is no way to bind a +component to an intermediate handle and have it dangle when that handle is +dropped. See {ref}`ffi_internals_one_arc` for why the allocation is kept +rather than replaced. + +`set_query_planner` does not return anything, because the query planner lives +in `SessionState` and is therefore a property of the session rather than of a +handle on it. Installing one is visible to every context sharing that session — +including ones a `with_*` call returned earlier. Installing a codec on a +session that already has a foreign planner rebuilds that planner against the +new chain for the same reason. This happens on the shared session, so it takes +effect even if the returned context is discarded: +`ctx.with_python_udf_inlining(...)` whose result is thrown away still leaves +the session's planner carrying the codecs of that discarded handle. A call that +changes nothing is exempt — asking for the inlining setting a context already +has returns a handle without touching the session. + +Order between installing codecs and installing a planner is a readability +preference rather than a requirement, since installing a codec after a planner +rebuilds the planner against it. The exception is a *layered* planner, where +codecs-first is a requirement: see {ref}`planner_codec_rebinding`. + +## Keep a context alive + +The session owns every installed component's task-context provider, and +dependent objects do not extend its lifetime. A `DataFrame`, logical plan, or +capsule can outlive every context on the session, but any operation that +reaches an FFI codec after the last one is collected fails with: + +```text +TaskContextProvider went out of scope over FFI boundary +``` + +Keep a context alive for as long as objects derived from it are in use. This is +the rule most likely to reach your users as a bug report against your library, +so it is worth stating in your own documentation too — the user-facing version +is in {ref}`user_guide_extensions`. + +The same rule applies to a capsule you take off a context inside your own code: +a codec capsule taken from a throwaway `SessionContext()` names a session that +is already gone and fails on first use. + +:::{warning} +{py:meth}`~datafusion.SessionContext.enable_url_table` is an exception to the +one-session-one-allocation rule above: it clones the underlying +`SessionContext`, so the returned context has an allocation of its own and must +not outlive the receiver. It also forks the session's state while keeping its +id, so two handles report one `session_id()` with divergent configuration. That +is a bug rather than a design, tracked in +[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708); +do not build on the behaviour. +::: diff --git a/docs/source/extension-guide/table-providers.md b/docs/source/extension-guide/table-providers.md new file mode 100644 index 000000000..b33f52ec3 --- /dev/null +++ b/docs/source/extension-guide/table-providers.md @@ -0,0 +1,138 @@ + + +(extension_providers)= + +# Providers and catalogs + +Five hooks expose data to a session, at four levels of granularity. All five +follow the {ref}`capsule protocol ` and all five +receive a {ref}`codec source ` as their single +argument. Every one of them is implemented in [`datafusion-ffi-example`]. + +| Hook | Exposes | Registered with | +| --- | --- | --- | +| `__datafusion_table_provider__` | one table | {py:meth}`~datafusion.SessionContext.register_table` | +| `__datafusion_table_provider_factory__` | a factory that builds tables from `CREATE EXTERNAL TABLE` | {py:meth}`~datafusion.SessionContext.register_table_factory` | +| `__datafusion_schema_provider__` | a named set of tables | {py:meth}`datafusion.catalog.Catalog.register_schema` | +| `__datafusion_catalog_provider__` | a named set of schemas | {py:meth}`~datafusion.SessionContext.register_catalog_provider` | +| `__datafusion_catalog_provider_list__` | the whole catalog namespace | {py:meth}`~datafusion.SessionContext.register_catalog_provider_list` | + +Start with a table provider. Reach for the schema and catalog levels when your +data source has its own namespace that should be browsable rather than +registered table by table, and for the provider list only when your library is +replacing the catalog namespace outright. + +## A table provider + +Implement +[TableProvider](https://datafusion.apache.org/library-user-guide/custom-table-providers.html) +in Rust, then expose it: + +```rust +#[pymethods] +impl MyTableProvider { + fn __datafusion_table_provider__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let provider = Arc::new(self.clone()); + let codec = ffi_logical_codec_from_pycapsule(session, None)?; + let provider = FFI_TableProvider::new_with_ffi_codec(provider, false, None, codec); + + PyCapsule::new_with_value(py, provider, cr"datafusion_table_provider") + } +} +``` + +Your users then register it as they would any other table — see +{ref}`io_custom_table_provider` for the Python side. + +## The catalog family + +The three catalog-level hooks have the same shape as each other. Taking the +schema provider as the representative: + +```rust +#[pymethods] +impl MySchemaProvider { + fn __datafusion_schema_provider__<'py>( + &self, + py: Python<'py>, + codec: Bound<'py, PyAny>, + ) -> PyResult> { + let provider = Arc::clone(&self.inner) as Arc; + + let codec = ffi_logical_codec_from_pycapsule(codec, None)?; + let provider = FFI_SchemaProvider::new_with_ffi_codec(provider, None, codec); + + PyCapsule::new_with_value(py, provider, cr"datafusion_schema_provider") + } +} +``` + +Swap `Schema` for `Catalog` or `CatalogProviderList` and the getter, capsule +name, and FFI type change together, following {ref}`the naming rule +`. `catalog_provider.rs` in +[`datafusion-ffi-example`] implements all three in one file, which is the +easiest way to see the symmetry. + +Note the parameter name. These four hooks are handed the host's logical codec +directly rather than a session, so naming it `codec` is more honest than +`session` — but do not rely on either: pass it to +`ffi_logical_codec_from_pycapsule` and do not inspect it. See +{ref}`extension_getter_argument`. + +## A table provider factory + +A factory backs `CREATE EXTERNAL TABLE`: DataFusion hands it the statement's +options and it produces a provider. The getter takes the codec and wraps an +`Arc`: + +```rust +fn __datafusion_table_provider_factory__<'py>( + &self, + py: Python<'py>, + codec: Bound<'py, PyAny>, +) -> PyResult> { + let codec = ffi_logical_codec_from_pycapsule(codec, None)?; + let factory = Arc::clone(&self.inner) as Arc; + let factory = FFI_TableProviderFactory::new_with_ffi_codec(factory, None, codec); + + PyCapsule::new_with_value(py, factory, cr"datafusion_table_provider_factory") +} +``` + +`FFI_TableProviderFactory` carries no version field, so a factory is one of the +three components that cannot be version-checked on import. Be correspondingly +careful about which DataFusion version you build against. + +## Serializing what you expose + +If plans referencing your tables have to leave the process — a distributed +engine will make them — your library also needs a logical extension codec, so +the provider can be rebuilt on the other side. That is {doc}`codecs`, and +{doc}`bundles` is how you ship the two together. + +The codec you pass to `new_with_ffi_codec` above is the **host's**, used to +serialize the parts of a plan the host owns. It is not a substitute for +contributing your own. + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/why-ffi.md b/docs/source/extension-guide/why-ffi.md new file mode 100644 index 000000000..e1321107a --- /dev/null +++ b/docs/source/extension-guide/why-ffi.md @@ -0,0 +1,148 @@ + + +(extension_why_ffi)= + +# Why FFI + +The DataFusion in Python project is designed to allow users to extend its +functionality in a few core areas. Ideally many users would like to package +their extensions as a Python package and easily integrate that package with +this project. This page describes the problem that makes the obvious approach +fail, and the approach this project uses instead. + +## The primary issue + +Suppose you wish to use DataFusion and you have a custom data source that can +produce tables that can then be queried against, similar to how you can +register a {ref}`CSV ` or {ref}`Parquet ` file. In +DataFusion terminology, you likely want to implement a +{ref}`Custom Table Provider `. In an effort to make +your data source as performant as possible and to utilize the features of +DataFusion, you may decide to write your source in Rust and then expose it +through [PyO3](https://pyo3.rs) as a Python library. + +At first glance, it may appear the best way to do this is to add the +`datafusion-python` crate as a dependency, provide a `PyTable`, and then to +register it with the `SessionContext`. Unfortunately, this will not work. + +When you produce your code as a Python library and it needs to interact with +the DataFusion library, at the lowest level they communicate through an +Application Binary Interface (ABI). The acronym sounds similar to API +(Application Programming Interface), but it is distinctly different. + +The ABI sets the standard for how these libraries can share data and functions +between each other. One of the key differences between Rust and other +programming languages is that Rust does not have a stable ABI. What this means +in practice is that if you compile a Rust library with one version of the +`rustc` compiler and I compile another library to interface with it but I use a +different version of the compiler, there is no guarantee the interface will be +the same. + +In practice, this means that a Python library built with `datafusion-python` as +a Rust dependency will generally **not** be compatible with the DataFusion +Python package, even if they reference the same version of +`datafusion-python`. If you attempt to do this, it may work on your local +computer if you have built both packages with the same optimizations. This can +sometimes lead to a false expectation that the code will work, but it +frequently breaks the moment you try to use your package against the released +packages. + +You can find more information about the Rust ABI in their +[online documentation](https://doc.rust-lang.org/reference/abi.html). + +## The FFI approach + +Rust supports interacting with other programming languages through its Foreign +Function Interface (FFI). The advantage of using the FFI is that it enables you +to write data structures and functions that have a stable ABI. That allows you +to use Rust code with C, Python, and other languages. In fact, the +[PyO3](https://pyo3.rs) library uses the FFI to share data and functions +between Python and Rust. + +The approach we are taking in the DataFusion in Python project is to +incrementally expose more portions of the DataFusion project via FFI +interfaces. This allows users to write Rust code that does **not** require the +`datafusion-python` crate as a dependency, expose their code in Python via +PyO3, and have it interact with the DataFusion Python package. + +Early adopters of this approach include +[delta-rs](https://delta-io.github.io/delta-rs/) who has adapted their Table +Provider for use in `datafusion-python` with only a few lines of code. Also, +the DataFusion Python project uses the existing definitions from +[Apache Arrow CStream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html) +to support importing **and** exporting tables. Any Python package that supports +reading the Arrow C Stream interface can work with DataFusion Python out of the +box! You can read more about working with Arrow sources in the +{ref}`Data Sources ` page. + +To learn more about the Foreign Function Interface in Rust, the +[Rustonomicon](https://doc.rust-lang.org/nomicon/ffi.html) is a good resource. + +## Inspiration from Arrow + +DataFusion is built upon [Apache Arrow](https://arrow.apache.org/). The +canonical Python Arrow implementation, +[pyarrow](https://arrow.apache.org/docs/python/index.html), provides an +excellent way to share Arrow data between Python projects without performing +any copy operations on the data, using a well defined set of interfaces — see +their [stream interface](https://arrow.apache.org/docs/format/CStreamInterface.html). +The [Rust Arrow implementation](https://github.com/apache/arrow-rs) also +supports these `C` style definitions via the Foreign Function Interface. Beyond +transferring data, `pyarrow` goes one step further and makes the interfaces +themselves easy to share in Python, by exposing PyCapsules that contain the +expected functionality. + +Two lessons we leverage from the Arrow project in DataFusion Python are: + +- We reuse the existing Arrow FFI functionality wherever possible. +- We expose PyCapsules that contain an FFI stable struct. + +You can learn more about PyCapsules from the official +[Python online documentation](https://docs.python.org/3/c-api/capsule.html). +PyCapsules have excellent support in PyO3 already; the +[PyO3 online documentation](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule) +is a good source for more details on using PyCapsules in Rust. + +## If FFI does not yet cover what you need + +:::{note} +Suppose you needed to expose some other features of DataFusion and you could +not wait for the upstream repository to implement the FFI approach we describe. +In this case you decide to create your dependency on the `datafusion-python` +crate instead. + +As we discussed, this is not guaranteed to work across different compiler +versions and optimization levels. If you wish to go down this route, there are +two approaches we have identified you can use. + +1. Re-export all of `datafusion-python` yourself with your extensions built in. +2. Carefully synchronize your software releases with the `datafusion-python` CI + build system so that your libraries use the exact same compiler, features, + and optimization level. + +We currently do not recommend either of these approaches as they are difficult +to maintain over a long period. Additionally, they require a tight version +coupling between libraries. + +The better path is to open an issue describing what you need exposed. The FFI +surface in the [datafusion-ffi] crate grows in response to these. +::: + +[datafusion-ffi]: https://crates.io/crates/datafusion-ffi diff --git a/docs/source/index.md b/docs/source/index.md index 13f7c7df6..c702fc040 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -59,11 +59,22 @@ df.show() ``` +## Where to go next + +- **[User Guide](user-guide/index)** — reading data, building queries, tuning + execution, and distributing work. +- **[Extension Guide](extension-guide/index)** — writing a library that plugs + into DataFusion in Python: table providers, functions, extension codecs, and + query planners. +- **[Contributor Guide](contributor-guide/index)** — changing + datafusion-python itself. + ```{toctree} :hidden: true :maxdepth: 1 user-guide/index +extension-guide/index contributor-guide/index API Reference links diff --git a/docs/source/user-guide/io/table_provider.md b/docs/source/user-guide/io/table_provider.md index 5dc2dc086..1b6987e7f 100644 --- a/docs/source/user-guide/io/table_provider.md +++ b/docs/source/user-guide/io/table_provider.md @@ -28,29 +28,8 @@ you must use DataFusion 43.0.0 or later and expose a [FFI_TableProvider](https:/ via [PyCapsule](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule). A complete example can be found in the [examples folder](https://github.com/apache/datafusion-python/tree/main/examples). - -The method takes the `SessionContext` it is being registered on. Take whatever -the FFI constructor needs from that session — here the logical extension codec — -rather than building one inside your library. See the {ref}`ffi` guide for the -full capsule protocol. - -```rust -#[pymethods] -impl MyTableProvider { - - fn __datafusion_table_provider__<'py>( - &self, - py: Python<'py>, - session: Bound<'py, PyAny>, - ) -> PyResult> { - let provider = Arc::new(self.clone()); - let codec = ffi_logical_codec_from_pycapsule(session, None)?; - let provider = FFI_TableProvider::new_with_ffi_codec(provider, false, None, codec); - - PyCapsule::new_with_value(py, provider, cr"datafusion_table_provider") - } -} -``` +For how to write one — the getter, what it receives, and how to serialize what +it exposes — see {ref}`extension_providers` in the Extension Guide. Once you have this library available, you can construct a {py:class}`~datafusion.Table` in Python and register it with the diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index b5cf8d16c..7250ac079 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -21,8 +21,9 @@ ## DataFusion 55.0.0 -This release extends the change made in 52.0.0 to the remaining {ref}`ffi` hook -methods. Users who contribute their own `LogicalExtensionCodec` or +This release extends the change made in 52.0.0 to the remaining +{ref}`extension_capsule_protocol` hook methods. Users who contribute their own +`LogicalExtensionCodec` or `PhysicalExtensionCodec` via FFI must update `__datafusion_logical_extension_codec__` and `__datafusion_physical_extension_codec__` to accept an additional @@ -75,13 +76,25 @@ used. it, so existing calls such as `ctx.__datafusion_logical_extension_codec__()` continue to work unchanged. +The provider and catalog getters — +`__datafusion_table_provider_factory__`, `__datafusion_catalog_provider__`, +`__datafusion_catalog_provider_list__`, and `__datafusion_schema_provider__` — +now receive the host's logical extension codec as a bare +`datafusion_logical_extension_codec` capsule rather than a session. +`__datafusion_table_provider__` receives a session from +`SessionContext.register_table` and a codec capsule from +`Schema.register_table`. No code change is needed in either case: pass the +argument to `ffi_logical_codec_from_pycapsule`, which calls the codec getter +when the object has one and returns the object unchanged when it does not. Do +not inspect the argument's type. See {ref}`extension_getter_argument`. + New in this release, `__datafusion_query_planner__` follows the same protocol. It receives the session and takes both extension codecs from it, so a planner library never builds a `TaskContextProvider` at all. Install one with `SessionContext.set_query_planner(planner)`, which mutates the session the same way `add_physical_optimizer_rule` does and returns nothing — the query planner lives in `SessionState`, so it belongs to the session rather than to a -particular handle on it. See the {ref}`ffi` guide for the full protocol. +particular handle on it. See {ref}`extension_planners` for the full protocol. If a library ships codecs *and* a planner, prefer `SessionContext.with_extensions(bundle)` over installing each piece by hand. It @@ -91,7 +104,7 @@ The library exposes a bundle object implementing `__datafusion_session_extension__` for its codecs and `__datafusion_session_planner__` for its planner — the latter is handed the planner installed so far, so several libraries that each ship one nest instead -of displacing each other. See the {ref}`ffi` guide. +of displacing each other. See {ref}`extension_bundles`. (extension_version_mismatch)= From 872afb613175f0b8a02705bb28665ffb80a55fb0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 08:51:58 -0400 Subject: [PATCH 30/33] docs: give each docstring claim one canonical home MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The extension docstrings and the FFI guide had grown the same six claims 3-5 times each, in wording that had already started to drift. `with_extensions` was 123 lines — the longest docstring in the package by 30 — and most of it argued a design rather than stating a contract. Routes each recurring claim to one home and leaves a one-line pointer elsewhere. Session sharing and context lifetime move onto the `SessionContext` class docstring, since they are properties of the type and that is exactly why five methods had each reworded them. The `None`-vs-`fallback` contract moves onto `SessionPlannerExportable`, since it is a return-value contract of one method. The two-phase rationale, the objects-not-capsules argument, and the codec-order argument stay in the guide, which owns the "why". The duck-type-the-session rule moves from `QueryPlannerExportable` to `LogicalExtensionCodecExportable`, which is already the designated `session` reference for that family and where the codec protocols were pointing for it anyway. It also gains the fact the old text was missing: across the protocol this argument is not always a session, so duck-typing it is not a style preference. Not every docstring shrank. `__datafusion_query_planner__` was 4 lines with no example while `set_query_planner` told callers to capture the fallback through it, so it grew to 25. Several others grew by gaining the `Args`/`Returns`/ `Raises` sections they were missing, and by trading `+SKIP` examples for runnable ones — a `SessionContext` satisfies the capsule-getter protocols, so its own exported capsule stands in for a library's without a build step. The total across these sixteen docstrings is roughly flat, at 577 lines before and 629 after; what changed is that the rationale left and the contract arrived. The 21 lines of `dataclasses.fields` rationale on `SessionExtensionComponents.__post_init__` become comments in the method body. That is the fix rather than touching `autoapi_options`: `__post_init__` is a *special* member, so dropping `private-members` would not hide it, and dropping `special-members` would delete the entire `__datafusion_*__` reference surface. `conf.py` now records why that setting is deliberately left alone, and de-duplicates the four re-exported `extensions` classes the way it already did for `DataFrame` and `SessionContext`. Adds `python/tests/test_docstrings.py`, which is what would have caught the 123-liner: a 95-line ceiling with no waiver list, a doctest-presence check over the extension protocol, and a check that a docstring naming the guide actually links it. The last two each found a real defect on first run. Also fixes what this branch made newly load-bearing in the serialization surface: `plan.py`'s single-backtick RST rendered `LogicalExtensionCodec` and friends as italics rather than links into the new API, and the `global_ctx()` fallback in `Expr.from_bytes` now silently yields a context with no extension codecs, which before this branch lost only registrations. Co-Authored-By: Claude Opus 5 (1M context) --- AGENTS.md | 41 ++ docs/source/conf.py | 9 + docs/source/extension-guide/bundles.md | 53 +++ .../_test_three_library_query_planner.py | 19 +- python/datafusion/context.py | 439 +++++++++++------- python/datafusion/expr.py | 21 +- python/datafusion/extensions.py | 224 ++++----- python/datafusion/ipc.py | 11 +- python/datafusion/plan.py | 67 ++- python/datafusion/user_defined.py | 42 +- python/tests/test_docstrings.py | 158 +++++++ 11 files changed, 759 insertions(+), 325 deletions(-) create mode 100644 python/tests/test_docstrings.py diff --git a/AGENTS.md b/AGENTS.md index 659094ec0..f61b05636 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -133,6 +133,47 @@ Every Python function must include a docstring with usage examples. `array_sort`) only need a one-line description and a `See Also` reference to the primary function. They do not need their own examples. +### One canonical home per claim + +A claim lives where the reader already is when they need it — **exactly once**. + +- A property of one callable's arguments, return value, or errors → that + callable's docstring. +- A property of a *type* that several callables share → the class docstring. + (Session sharing and context lifetime live on `SessionContext`, not on each + of the five `with_*` methods.) +- **Why** the API is shaped this way, a multi-library recipe, Rust-side code, a + trade-off, or a limitation with an upstream issue → the narrative guide under + `docs/source/`. + +Everywhere else: one sentence plus one Sphinx role. A docstring may *state* a +claim the guide also makes; it must not *argue* it — no "because", no "the +reason is", no counter-argument. + +**The test:** if a paragraph would survive being moved into the guide unchanged, +move it. `python/tests/test_docstrings.py` enforces a length ceiling, which is +the symptom this rule treats. + +When you point at the guide, use a real `:ref:` to a specific label. Prose +saying "see the extensions guide" with no role is a dead end in the rendered +HTML. + +### Examples that need a compiled extension + +Some APIs cannot be demonstrated without a built FFI extension library, which +this package does not ship. The convention is: + +1. A **runnable** example block first, using only the wheel. A `SessionContext` + satisfies the capsule-getter protocols, so its own exported capsule stands in + for a real library's in a doctest. +2. Then a `# doctest: +SKIP` block showing real usage, at most a few lines, with + one line of prose naming the test that runs it for real. +3. Every `+SKIP` block needs that mirror. See + `test_with_extensions_docstring_example_still_runs` in + `examples/datafusion-ffi-query-planner-example`, which parses the live + docstring, drops the skip, and executes it — so a renamed method or a wrong + expected output fails there. + ## Aggregate and Window Function Documentation When adding or updating an aggregate or window function, ensure the corresponding diff --git a/docs/source/conf.py b/docs/source/conf.py index 0369d7fe7..2463c580b 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -103,6 +103,11 @@ # This pattern also affects html_static_path and html_extra_path. exclude_patterns = [] +# `autoapi_options` is deliberately left at the default. In particular +# `special-members` is load-bearing: it is what publishes the +# `__datafusion_*__` capsule getters, which are the public protocol an +# extension library implements against. Dropping it would delete that +# reference surface along with `DataFrame.__init__` and the Arrow dunders. autoapi_dirs = ["../../python"] autoapi_ignore = ["*tests*"] autoapi_member_order = "groupwise" @@ -116,6 +121,10 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool: # noqa # Re-exports ("class", "datafusion.DataFrame"), ("class", "datafusion.SessionContext"), + ("class", "datafusion.QueryPlannerExportable"), + ("class", "datafusion.SessionExtensionComponents"), + ("class", "datafusion.SessionExtensionExportable"), + ("class", "datafusion.SessionPlannerExportable"), ("module", "datafusion.common"), # Duplicate modules (skip module-level docs to avoid duplication) ("module", "datafusion.col"), diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index 2a0314196..fffbedb59 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -70,6 +70,59 @@ off the supplied context, wrapping its codecs in `BundledLogicalCodec` / `BundledPhysicalCodec` so they carry declared ids, and constructing a Python {py:class}`~datafusion.SessionExtensionComponents`. +## A bundle you can run + +Before wiring up a cdylib, it is worth seeing the protocol work end to end in +pure Python. This runs against the plain wheel — the codec here re-exports the +host session's own capsule, where a real one would return its library's: + +```python +from datafusion import SessionContext, SessionExtensionComponents + + +class Codec: + """Wraps a capsule so it carries an id. A Rust library ships this shape.""" + + def __init__(self, codec_id, capsule): + self.__datafusion_codec_id__ = codec_id + self._capsule = capsule + + def __datafusion_logical_extension_codec__(self, session=None): + return self._capsule + + +class Bundle: + def __init__(self, codec_id): + self.codec_id = codec_id + + def __datafusion_session_extension__(self, ctx): + # Fresh components on every call, bound to the `ctx` handed in. + # Never cache these, and never retain `ctx`. + return SessionExtensionComponents( + logical_extension_codecs=( + Codec(self.codec_id, ctx.__datafusion_logical_extension_codec__()), + ) + ) + + +ctx = SessionContext().with_extensions(Bundle("tables.v1"), Bundle("engine.v1")) +ctx.logical_extension_codec_ids() +# ['tables.v1', 'engine.v1'] +``` + +Three things this makes observable, each pinned by a test in +`python/tests/test_context.py`: + +- **Ids accumulate in bundle order.** Decoding does not depend on that order; + only encoding does. See {ref}`extension_codec_order`. +- **Two bundles claiming one id are refused**, with a `ValueError` naming the + id — not resolved by position, since a positional id would break stored plans + the first time a bundle reordered what it returns. +- **A hook that raises leaves the receiving session untouched.** Add a bundle + whose hook raises and the source context's + {py:meth}`~datafusion.SessionContext.logical_extension_codec_ids` is still + empty afterwards. + (extension_bundles_two_phases)= ## Two phases, because codecs and planners compose differently 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 80ba7c724..cb85e6bf2 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 @@ -1386,17 +1386,24 @@ def test_with_extensions_docstring_example_still_runs(): statements are parsed out of the live docstring, the skip is dropped, and each one is executed and its output compared. + Only the ``+SKIP`` statements are taken. The docstring also carries a + runnable example above them, which the main suite already executes under + ``--doctest-modules``; running it again here would need its own namespace + and prove nothing. + Only names are redirected: ``my_extension`` resolves to the bundle above, and ``SessionContext`` supplies the config this library's planner reads. A renamed method, a changed signature, or a wrong expected output in the docstring fails here. """ - examples = 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) + examples = [ + example + for example in doctest.DocTestParser().get_examples( + inspect.getdoc(SessionContext.with_extensions) + ) + if example.options.pop(doctest.SKIP, False) + ] + assert examples, "with_extensions docstring has no skipped examples to check" module = types.ModuleType("my_extension") module.DistributedEngineExtension = _DocstringExampleExtension diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 69339a33f..96b9f11a8 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -135,7 +135,17 @@ def __arrow_c_array__( # noqa: D105 class TableProviderExportable(Protocol): """Type hint for object that has __datafusion_table_provider__ PyCapsule. - https://datafusion.apache.org/python/user-guide/io/table_provider.html + See :ref:`io_custom_table_provider` for registering one, and + :ref:`extension_providers` for writing one. + + Args: + session: See + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`. + For this getter it is a session when the provider is registered + through :py:meth:`SessionContext.register_table` and the host's + logical codec when it goes through + :py:meth:`datafusion.catalog.Schema.register_table`, so + duck-typing it is not optional. """ def __datafusion_table_provider__(self, session: Any) -> object: ... # noqa: D105 @@ -539,6 +549,27 @@ class SessionContext: """This is the main interface for executing queries and creating DataFrames. See :ref:`user_guide_concepts` in the online documentation for more information. + + **A context is a handle on a session, not the session itself.** The + ``with_*`` methods — :py:meth:`with_logical_extension_codec`, + :py:meth:`with_physical_extension_codec`, + :py:meth:`with_python_udf_inlining`, and :py:meth:`with_extensions` — + return a new context wrapping the *same* underlying session. Only the + Python-side codec settings differ; catalogs, tables, registered functions, + and configuration are the one shared session, so a registration through + either handle is visible to both. + + A few things therefore belong to the session rather than to a handle, and + take effect even if the handle that set them is discarded: the query + planner (see :py:meth:`set_query_planner`), and the rebuild of an installed + foreign planner that follows installing a codec. :ref:`extension_sessions` + in the online documentation works through when that matters. + + **Keep a context alive for as long as anything derived from it is in use.** + A :py:class:`~datafusion.DataFrame`, logical plan, or exported capsule does + not extend the session's lifetime. Once the last context on a session is + collected, any operation that reaches an extension codec fails with + ``TaskContextProvider went out of scope over FFI boundary``. """ def __init__( @@ -1769,45 +1800,50 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non """Install a custom query planner on this session. The planner is imported through its ``__datafusion_query_planner__`` - PyCapsule and installed on this context, in the same way + PyCapsule, in the same way :meth:`~SessionContext.add_physical_optimizer_rule` installs a rule. - The query planner is part of the session state, so it applies to this - context and to every context sharing its session — including ones - already returned by - :meth:`~SessionContext.with_logical_extension_codec` and friends. - A session holds exactly one planner, so calling this again replaces the - previous one rather than layering. To chain planners, have the new + Returns nothing, because the planner lives in the session rather than + in a handle on it — installing one is visible to every context sharing + that session, including ones an earlier ``with_*`` call returned. See + :py:class:`SessionContext`. + + A session holds exactly one planner, so calling this again **replaces** + the previous one rather than layering. To chain planners, have the new planner wrap the capsule from - :meth:`~SessionContext.__datafusion_query_planner__`, captured - *before* the new planner is installed. + :meth:`~SessionContext.__datafusion_query_planner__`, captured *before* + the new planner is installed. - Install any extension codecs before a layered planner. Installing a - codec afterwards rebuilds the installed planner against it, but not the - fallback inside it, which keeps the codecs it was imported with. Note - also that the planner is built against the codecs of the context this - method is called on, so installing the same planner again on a different - handle rebinds the session's planner to *that* handle's codecs. See the - FFI extensions guide for the full multi-library registration recipe. + Install any extension codecs before a layered planner; the rebuild that + follows a later codec install does not reach the fallback inside one. + See :ref:`planner_codec_rebinding`, or prefer + :py:meth:`with_extensions`, which cannot capture a partial chain. Args: planner: Object exposing ``__datafusion_query_planner__`` (see :py:class:`~datafusion.extensions.QueryPlannerExportable`) or a raw ``datafusion_query_planner`` PyCapsule. + Raises: + ValueError: If the capsule is not named ``datafusion_query_planner``. + Examples: - >>> from my_extension import DistributedQueryPlanner # doctest: +SKIP + A session exports its own planner, which is what you capture to + wrap: + + >>> from datafusion import SessionContext >>> ctx = SessionContext() - >>> ctx.set_query_planner(DistributedQueryPlanner()) # doctest: +SKIP - >>> ctx.sql("SELECT * FROM remote_table").collect() # doctest: +SKIP + >>> fallback = ctx.__datafusion_query_planner__() + >>> type(fallback).__name__ + 'PyCapsule' - Layer a planner on top of the one already installed by capturing - the existing planner first: + Skipped here (needs a built extension library): - >>> fallback = ctx.__datafusion_query_planner__() # doctest: +SKIP + >>> from my_extension import DistributedQueryPlanner # doctest: +SKIP >>> ctx.set_query_planner( ... DistributedQueryPlanner(fallback=fallback) ... ) # doctest: +SKIP + >>> ctx.sql("SELECT * FROM remote_table").collect() # doctest: +SKIP """ self.ctx.set_query_planner(planner) @@ -1816,83 +1852,34 @@ def with_extensions( ) -> SessionContext: """Create a new session context with the given extension bundles. - This is the preferred way to install FFI extensions that need a - task-context provider (extension codecs and query planners). It avoids - the pitfalls of chaining :py:meth:`with_logical_extension_codec`, - :py:meth:`with_physical_extension_codec`, and - :py:meth:`set_query_planner` by hand, where the codecs a planner was - built against can end up stale. - - Installation runs in two phases, because codecs and planners compose - differently: - - 1. Every extension's ``__datafusion_session_extension__`` is called - with this context and its codecs are collected, then all of them are - installed at once. A session chains many codecs and dispatches - between them by id, so they merely accumulate; order affects - encoding only. - 2. Every extension's ``__datafusion_session_planner__`` is then called, - **in argument order**, each receiving the planner built so far. A - session holds exactly one planner, so planners compose by *nesting*: - each wraps the previous one and delegates to it. The last extension - listed ends up outermost and is consulted first. - - An extension implements either hook or both. Phase two runs after every - codec is installed and receives a context carrying the final chains, so - a nested planner is never left encoding through a chain a later - extension has grown. - - If no extension supplies a planner but codecs were installed, an - existing FFI planner is rebound to the final chains; if the call - installed nothing at all, the session's planner is not touched. An - extension that ignores the ``fallback`` it is handed replaces the - planners before it instead of nesting on them, including any the - session already had. - - Codec order never affects decoding, which routes by codec id. It - affects encoding only when two codecs would claim the same node: the - chain stops at the first that does, so a codec claiming a broad - category can take nodes belonging to a library installed after it. The - query still succeeds, but the plan is written by the wrong library and - may not decode elsewhere. If an extension needs to be early for its - codec and late for its planner, contribute each half at its own - position rather than reordering — the two hooks are independent, so a - small adapter implementing one of them and delegating is enough. The - FFI extensions guide shows the pattern. - - Codecs must be handed over as objects exposing the capsule getter, not - as bare ``PyCapsule`` objects, and are named after their exporting - class as :py:meth:`with_logical_extension_codec` describes. Declare - ``__datafusion_codec_id__`` on the object to pin an id that survives a - later class rename. A capsule carries no type of its own, so there - would be nothing to name the codec by, and this method takes no - ``codec_id=``; wrap it in an object instead. That also keeps a codec's - wire identity independent of the extension that ships it, so an - extension composed inside another one still writes the same ids. - - Planners are exempt — a planner carries no wire id, so a hook may - return an object or a capsule. - - Like the individual ``with_*`` methods, the returned context shares its - session with this one: catalogs, tables, registered functions, and - configuration are the one session, so a registration on either side is - visible to both, and the planner is installed on that shared session - even if the returned context is discarded. Only the Python-side codec - chains are specific to the returned handle. - - No state is written until every extension has run and every capsule has - been validated, so an extension that raises or returns something - invalid — in either phase — leaves the session as it was. Codec chains - belong to the returned handle, and the single session write happens - after the last planner hook returns. The exception is an extension that - mutates the context it is handed — registering a table, say — which is - not rolled back. Extension factories should treat that context as - configuration-only. - - The session owns the installed components' task-context providers, and - dependent objects do not extend its lifetime. Keep a context on the - session alive for as long as DataFrames or plans derived from it are in - use; FFI operations after the last one is collected raise an error. + This is the preferred way to install extension codecs and query + planners, because it removes the ordering question that installing them + by hand creates. + + Each argument is called twice, in two phases: + + 1. ``__datafusion_session_extension__(ctx)`` on every extension, then + all the returned codecs are installed at once. + 2. ``__datafusion_session_planner__(ctx, fallback)`` on every + extension, **in argument order**, each handed the planner built so + far and the context carrying every bundle's codecs. The last + extension listed ends up outermost. + + An extension implements either hook or both. Return ``None`` from the + planner hook to contribute no planner; see + :py:class:`~datafusion.extensions.SessionPlannerExportable`. + + Nothing is written to the session until every hook has returned and + every capsule has been validated, so a hook that raises leaves the + session as it was. A hook that *mutates* the context it is handed — + registering a table, say — is not rolled back, which is why bundle + objects must be configuration-only. + + Shares its session with this context — see :py:class:`SessionContext`. + + See :ref:`extension_bundles` in the online documentation for why the + phases are split, how to contribute a bundle's two halves at different + positions, and a worked Rust implementation. Args: extensions: Extension bundles to install. Order is irrelevant for @@ -1905,29 +1892,35 @@ class as :py:meth:`with_logical_extension_codec` describes. Declare A new context with all extension components installed. Raises: - TypeError: If an argument implements neither hook, if - ``__datafusion_session_extension__`` returns something other - than a - :py:class:`~datafusion.extensions.SessionExtensionComponents`, - or if an extension contributes a codec as a bare ``PyCapsule``. - ValueError: If two codecs claim the same id. An extension that - contributes two instances of one codec class must declare - ``__datafusion_codec_id__`` on at least one of them; the - collision is refused rather than resolved by position, because - a positional id would break stored plans the first time the - extension reordered what it returns. Also if a capsule getter - returns a capsule of the wrong kind — a physical codec handed - over under ``__datafusion_logical_extension_codec__``, say — - which is reported against the name the getter should have - produced. + TypeError: If an argument implements neither hook, if a hook + returns the wrong type, or if a codec is contributed as a bare + ``PyCapsule`` rather than an object exposing the getter. + ValueError: If two codecs claim the same id, or a getter returns a + capsule of the wrong kind. See + :py:meth:`with_logical_extension_codec` for how ids are + assigned. Examples: - The 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. + The returned handle is a different object sharing one session, and + an empty call is legal: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> derived = ctx.with_extensions() + >>> derived is ctx + False + >>> ctx.from_pydict({"a": [1, 2]}, name="t") # doctest: +ELLIPSIS + DataFrame()... + >>> derived.table_exist("t") + True + >>> derived.logical_extension_codec_ids() + [] + + A runnable multi-bundle example, showing what a bundle returns and + how its codec ids accumulate, is in :ref:`extension_bundles`. + + Real usage. Skipped here (needs a built extension library); run + verbatim by ``test_with_extensions_docstring_example_still_runs``. >>> from my_extension import DistributedEngineExtension # doctest: +SKIP >>> ctx = SessionContext().with_extensions( @@ -2431,12 +2424,10 @@ def __datafusion_codec_id__(self) -> str: session, so two contexts can be installed on one session and a plan written through one will not be decoded by the other. - Contexts derived from the same session — including the ones returned by - :py:meth:`with_logical_extension_codec`, - :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. + Contexts derived from the same session report the same id, so only one + of them can be installed on a given session. That is the intended + answer: they are one session — see :py:class:`SessionContext` — so + their payloads would be indistinguishable on decode. Examples: >>> from datafusion import SessionContext @@ -2453,15 +2444,50 @@ def __datafusion_logical_extension_codec__(self, session: Any = None) -> Any: ``session`` is accepted so a context satisfies the same protocol an extension library implements, where the argument is how the library - reaches the session it is being installed on. A context already is one, - so the argument is ignored. + reaches the host's codec. A context already carries one, so the + argument is ignored. See + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable` + for what an extension library does with it. + + Args: + session: Accepted and ignored. + + Returns: + A ``datafusion_logical_extension_codec`` PyCapsule. + + Examples: + >>> from datafusion import SessionContext + >>> type(SessionContext().__datafusion_logical_extension_codec__()).__name__ + 'PyCapsule' """ return self.ctx.__datafusion_logical_extension_codec__(session) def __datafusion_query_planner__(self, session: Any = None) -> Any: """Access the ``FFI_QueryPlanner`` PyCapsule for the current planner. - See :meth:`__datafusion_logical_extension_codec__` for ``session``. + This is how you capture the planner a session already has in order to + wrap it. Capture it *before* installing the new one, since + :py:meth:`set_query_planner` replaces rather than layers. + + Args: + session: Accepted and ignored. See + :meth:`__datafusion_logical_extension_codec__`. + + Returns: + A ``datafusion_query_planner`` PyCapsule wrapping the session's + current planner, exported for a foreign planner to delegate to. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> fallback = ctx.__datafusion_query_planner__() + >>> type(fallback).__name__ + 'PyCapsule' + + Both spellings work, since the argument is ignored: + + >>> type(ctx.__datafusion_query_planner__(ctx)).__name__ + 'PyCapsule' """ return self.ctx.__datafusion_query_planner__(session) @@ -2482,33 +2508,54 @@ def with_logical_extension_codec( affect decoding. A serialized plan records which codec wrote each payload, as a short id - taken from the codec's class. ``codec_id`` overrides that id and is - normally unnecessary. Pass it when installing from a bare ``PyCapsule``, - which has no class to take an id from, or when installing two instances - of one class, which otherwise claim the same id and raise ``ValueError``. + taken from the codec's class. + + Shares its session with this context — see :py:class:`SessionContext`. - The returned context shares its session state with the original, so a - later registration on either is visible to both, and an installed query - planner is rebound on the shared session even if the returned context is - discarded. + See :ref:`extension_codec_ids` in the online documentation for how ids + are assigned and what an extension codec has to implement, and + :py:meth:`with_extensions` for installing a library's codecs and planner + together. - See :ref:`ffi` in the online documentation for how ids are assigned, - what an extension codec has to implement, and a worked multi-library - registration recipe. + Args: + codec: Object implementing ``__datafusion_logical_extension_codec__`` + (see + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`), + or a raw ``datafusion_logical_extension_codec`` PyCapsule. + codec_id: Overrides the id the codec's payloads are tagged with. + Normally unnecessary. Pass it when installing two instances of + one class, which otherwise claim the same id, and when + installing a bare ``PyCapsule``: a capsule has no class to take + an id from, so it is given a random ``anon:`` id that differs on + every install, and plans it encodes can never be decoded + elsewhere. + + Returns: + A new context carrying this codec in addition to any already + installed. + + Raises: + ValueError: If the resolved id is already installed on this session. Examples: + A context exports its own codec, which stands in here for a real + library's: + >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx = ctx.with_logical_extension_codec( - ... my_library.Codec() - ... ) # doctest: +SKIP + >>> host = SessionContext() + >>> capsule = host.__datafusion_logical_extension_codec__() + >>> ctx = SessionContext().with_logical_extension_codec( + ... capsule, codec_id="my_library.Codec" + ... ) + >>> ctx.logical_extension_codec_ids() + ['my_library.Codec'] - Installing from a bare capsule, pinning the id so encoded - plans remain decodable on another session: + Without ``codec_id`` a bare capsule gets an anonymous id, which is + fine only if its plans never leave this session: - >>> ctx = ctx.with_logical_extension_codec( - ... capsule, codec_id="my_library.Codec" - ... ) # doctest: +SKIP + >>> ctx = SessionContext().with_logical_extension_codec(capsule) + >>> ctx.logical_extension_codec_ids()[0].startswith("anon:") + True """ new_internal = self.ctx.with_logical_extension_codec(codec, codec_id) new = SessionContext.__new__(SessionContext) @@ -2526,15 +2573,20 @@ def logical_extension_codec_ids(self) -> list[str]: DataFusion's own default codec is not listed. It handles whatever no installed codec claims, and it carries no identity to list. + Returns: + The installed codec ids, in install order. Empty if none are + installed. + Examples: >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx.logical_extension_codec_ids() + >>> host = SessionContext() + >>> host.logical_extension_codec_ids() [] - >>> ctx = ctx.with_logical_extension_codec( - ... my_library.Codec() - ... ) # doctest: +SKIP - >>> ctx.logical_extension_codec_ids() # doctest: +SKIP + >>> capsule = host.__datafusion_logical_extension_codec__() + >>> ctx = SessionContext().with_logical_extension_codec( + ... capsule, codec_id="my_library.Codec" + ... ) + >>> ctx.logical_extension_codec_ids() ['my_library.Codec'] """ return self.ctx.logical_extension_codec_ids() @@ -2543,6 +2595,18 @@ def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: """Access the PyCapsule FFI_PhysicalExtensionCodec. See :meth:`__datafusion_logical_extension_codec__` for ``session``. + + Args: + session: Accepted and ignored. + + Returns: + A ``datafusion_physical_extension_codec`` PyCapsule. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> type(ctx.__datafusion_physical_extension_codec__()).__name__ + 'PyCapsule' """ return self.ctx.__datafusion_physical_extension_codec__(session) @@ -2551,11 +2615,21 @@ def physical_extension_codec_ids(self) -> list[str]: See :py:meth:`logical_extension_codec_ids`. + Returns: + The installed codec ids, in install order. Empty if none are + installed. + Examples: >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx.physical_extension_codec_ids() + >>> host = SessionContext() + >>> host.physical_extension_codec_ids() [] + >>> capsule = host.__datafusion_physical_extension_codec__() + >>> ctx = SessionContext().with_physical_extension_codec( + ... capsule, codec_id="my_library.PhysicalCodec" + ... ) + >>> ctx.physical_extension_codec_ids() + ['my_library.PhysicalCodec'] """ return self.ctx.physical_extension_codec_ids() @@ -2574,16 +2648,27 @@ def with_physical_extension_codec( :py:meth:`with_logical_extension_codec` does, including when to pass ``codec_id`` and what the returned context shares. See that method. + Args: + codec: As :py:meth:`with_logical_extension_codec`, for the physical + getter. + codec_id: As :py:meth:`with_logical_extension_codec`. + + Returns: + A new context carrying this codec in addition to any already + installed. + + Raises: + ValueError: If the resolved id is already installed on this session. + Examples: >>> from datafusion import SessionContext - >>> ctx = SessionContext() - >>> ctx = ctx.with_physical_extension_codec( - ... my_library.PhysicalCodec() - ... ) # doctest: +SKIP - - >>> ctx = ctx.with_physical_extension_codec( + >>> host = SessionContext() + >>> capsule = host.__datafusion_physical_extension_codec__() + >>> ctx = SessionContext().with_physical_extension_codec( ... capsule, codec_id="my_library.PhysicalCodec" - ... ) # doctest: +SKIP + ... ) + >>> ctx.physical_extension_codec_ids() + ['my_library.PhysicalCodec'] """ new_internal = self.ctx.with_physical_extension_codec(codec, codec_id) new = SessionContext.__new__(SessionContext) @@ -2593,10 +2678,6 @@ def with_physical_extension_codec( def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext: """Control whether Python UDFs are embedded in serialized expressions. - ``enabled`` is keyword-only and required: callers must pick a - mode explicitly. Fresh sessions inline UDFs (``enabled=True`` - behavior) until this method overrides the toggle. - With ``enabled=True``, serialized expressions carry the Python code for any scalar, aggregate, or window UDFs they reference. The receiver rebuilds the UDFs from those bytes and does not @@ -2625,14 +2706,18 @@ def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext: :func:`pickle.loads` on untrusted bytes remains unsafe regardless of the toggle. - Returns a new :class:`SessionContext` with the toggle applied; - the original context's own codec settings are unchanged. The - returned context shares its session state with the original, so - a later registration on either is visible to both. If a custom - query planner is installed, it is rebuilt against the new codecs - on the shared session, so the original context plans with them - too. This happens on the shared session, so it takes effect even - if the returned context is discarded. + Shares its session with this context — see + :py:class:`SessionContext`. The original context's own codec + settings are unchanged. + + Args: + enabled: Whether to embed Python UDFs in serialized + expressions. Keyword-only and required, so callers must + pick a mode explicitly. Fresh sessions behave as + ``enabled=True`` until this method overrides the toggle. + + Returns: + A new :class:`SessionContext` with the toggle applied. Examples: >>> import pyarrow as pa diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index c198b646c..18ce3554d 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -576,6 +576,13 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr: (sufficient for built-ins and Python UDFs, plus any UDFs registered on the global context). + Note that the global context carries no extension codecs, so + falling back to it cannot decode a payload written by one. A + worker that must decode extension-codec payloads has to install a + context carrying those codecs — see + :func:`datafusion.ipc.set_worker_ctx` and + :meth:`SessionContext.with_extensions`. + .. warning:: Security Decoding may invoke ``cloudpickle.loads`` on bytes embedded in the payload, which executes arbitrary Python code. Treat @@ -627,19 +634,19 @@ def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]: portable across minor versions. See :meth:`to_bytes` for details on what travels by value vs. by reference. - Examples: - >>> import pickle - >>> from datafusion import col, lit - >>> e = col("a") * lit(2) - >>> pickle.loads(pickle.dumps(e)).canonical_name() - 'a * Int64(2)' - The encoding side honors a driver-side sender context installed via :func:`datafusion.ipc.set_sender_ctx` — that is how :meth:`SessionContext.with_python_udf_inlining` propagates through ``pickle.dumps``. The sender context is read by ``__reduce__``, so :func:`copy.copy` and :func:`copy.deepcopy` — which also go through ``__reduce__`` — pick it up too. + + Examples: + >>> import pickle + >>> from datafusion import col, lit + >>> e = col("a") * lit(2) + >>> pickle.loads(pickle.dumps(e)).canonical_name() + 'a * Int64(2)' """ return (Expr._reconstruct, (self.to_bytes(get_sender_ctx()),)) diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 631fe239b..0cc41859e 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -26,25 +26,20 @@ ctx = SessionContext().with_extensions(MyLibraryExtension()) -Installing through ``with_extensions`` rather than by chaining the individual -``with_*`` methods matters for components that hold a task-context provider: -the extension is handed the session its components will run on, and every -codec is installed before any query planner is bound against them, so no -planner is left carrying a codec chain that has since grown. See the FFI -extensions guide in the contributor documentation for the full rationale. - -Codecs and planners install in two phases, because they compose differently. A -session's codec chain holds many codecs and dispatches between them by id, so -codecs merely accumulate and their order does not affect decoding. A session -holds exactly *one* query planner, so planners compose by nesting: each wraps -the one before it. Phase one collects the codecs of every bundle implementing -:py:class:`SessionExtensionExportable` and installs them; phase two runs -:py:class:`SessionPlannerExportable` once for each bundle that implements it, -in argument order, handing each the planner built so far. A bundle implements -either hook or both, and one it does not implement is simply not called. - -That split is what lets several libraries that each ship a planner coexist. It -also means bundle order is significant for planners and irrelevant for codecs. +Codecs and planners install in two phases: every +:py:class:`SessionExtensionExportable` runs first and its codecs are installed, +then every :py:class:`SessionPlannerExportable` runs in argument order. A bundle +implements either hook or both. Bundle order is significant for planners, which +nest, and irrelevant for codecs, which accumulate. + +Of the four names here, only the two bundle hooks are ``@runtime_checkable``, +because :py:meth:`~datafusion.context.SessionContext.with_extensions` +dispatches on them from Python. :py:class:`QueryPlannerExportable` is a type +hint only, matching the other capsule-getter protocols in +:py:mod:`datafusion.user_defined` and :py:mod:`datafusion.catalog`. + +See :ref:`extension_bundles` in the online documentation for why the phases are +split and for a worked implementation. """ from __future__ import annotations @@ -75,18 +70,32 @@ class QueryPlannerExportable(Protocol): The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically produced by a separate compiled extension. ``session`` is a handle on the session the planner is being installed on; take the extension codecs from - it rather than building your own. - - Duck-type that handle rather than checking its type. It is the PyO3 - context from ``datafusion._internal``, not the - :py:class:`~datafusion.context.SessionContext` wrapper, so it exposes every - capsule getter and ``__datafusion_codec_id__`` — which is all the protocol - asks of it — but ``isinstance(session, SessionContext)`` is ``False`` even - though its ``repr`` reads ``datafusion.SessionContext``. The same is true - of the codec getters in :py:mod:`datafusion.user_defined`. The two bundle - hooks are the exception: :py:class:`SessionExtensionExportable` and - :py:class:`SessionPlannerExportable` are dispatched from Python and receive - the wrapper. + it rather than building your own, and duck-type it — see + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable` for + ``session``. + + Unlike the two bundle hooks in this module, this protocol is a type hint + only: it is not ``@runtime_checkable``, so ``isinstance`` against it raises + ``TypeError``. + + Examples: + A :py:class:`~datafusion.context.SessionContext` satisfies this + protocol, which is what lets a foreign planner wrap the one a session + already has: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> type(ctx.__datafusion_query_planner__(ctx)).__name__ + 'PyCapsule' + + The protocol itself is not runtime-checkable: + + >>> from datafusion import QueryPlannerExportable + >>> try: + ... isinstance(ctx, QueryPlannerExportable) + ... except TypeError as e: + ... print("runtime_checkable" in str(e)) + True """ def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 @@ -115,13 +124,6 @@ class SessionExtensionComponents: Query planners are not listed here. They install in a second phase so each can wrap the one before it — see :py:class:`SessionPlannerExportable`. - Codecs must be objects exposing the capsule getters, never bare - ``PyCapsule`` objects: a codec's id is read off the object it is handed - over as, and a capsule has no type to read. A library holding a raw capsule - wraps it in an object, which is also what gives the codec an identity of - its own — one that does not change when the codec is contributed through a - different extension. - Examples: A bundle that contributes no codecs is valid — a planner-only library returns this, or omits the hook entirely: @@ -130,8 +132,6 @@ class SessionExtensionComponents: >>> components = SessionExtensionComponents() >>> components.logical_extension_codecs () - >>> components.physical_extension_codecs - () A bundle that contributes one kind of component names it, leaving the rest empty. Here the codec is a capsule wrapped in an object that @@ -147,68 +147,61 @@ class SessionExtensionComponents: ... def __datafusion_logical_extension_codec__(self, session=None): ... return self._capsule - The context stays in scope for as long as the codec does. An - ``FFI_LogicalExtensionCodec`` holds its task-context provider *weakly*, - so a capsule taken off a throwaway ``SessionContext()`` names a session - that is already gone and fails on first use with ``TaskContextProvider - went out of scope over FFI boundary``: - >>> ctx = SessionContext() - >>> capsule = ctx.__datafusion_logical_extension_codec__() >>> components = SessionExtensionComponents( - ... logical_extension_codecs=(NamedCodec(capsule),) + ... logical_extension_codecs=(NamedCodec( + ... ctx.__datafusion_logical_extension_codec__() + ... ),) ... ) >>> components.logical_extension_codecs[0].__datafusion_codec_id__ 'my_library.v1' >>> components.physical_extension_codecs () - Any iterable is accepted and stored as a tuple, so a bundle that builds - its codecs with a list comprehension does not have to convert: - - >>> components = SessionExtensionComponents( - ... logical_extension_codecs=[NamedCodec(capsule)] - ... ) - >>> type(components.logical_extension_codecs).__name__ - 'tuple' - A single codec is not an iterable of codecs, and forgetting the trailing comma is the easy way to write one by accident: - >>> SessionExtensionComponents(logical_extension_codecs=NamedCodec(capsule)) + >>> SessionExtensionComponents(logical_extension_codecs=NamedCodec(ctx)) Traceback (most recent call last): ... TypeError: logical_extension_codecs must be an iterable of codec objects... """ logical_extension_codecs: tuple[LogicalExtensionCodecExportable, ...] = () - """Logical codecs to add to the session's codec chain, in declaration order.""" + """Logical codecs to add to the session's codec chain, in declaration order. + + Objects exposing ``__datafusion_logical_extension_codec__``, never bare + ``PyCapsule`` objects — a codec's id is read off the object it is handed + over as. Any iterable is accepted and stored as a tuple. See + :ref:`extension_bundles_codecs_are_objects`. + """ physical_extension_codecs: tuple[PhysicalExtensionCodecExportable, ...] = () - """Physical codecs to add to the session's codec chain, in declaration order.""" + """Physical codecs to add to the session's codec chain, in declaration order. + + As :py:attr:`logical_extension_codecs`, for + ``__datafusion_physical_extension_codec__``. + """ def __post_init__(self) -> None: - """Normalize each field to a tuple, rejecting what cannot become one. - - A bundle that writes ``logical_extension_codecs=codec`` instead of - ``(codec,)`` is contributing one codec, not an iterable of them. - Without this, the mistake surfaces inside - :py:meth:`~datafusion.context.SessionContext.with_extensions` as - ``'MyCodec' object is not iterable``, which names neither the field - nor the hook that built it. Checking here puts the error in the - extension library's own frame. - - Normalizing is worth doing on its own: the declared type is a tuple - and the class is frozen, so a list left in place would be a mutable - member of an immutable value, and a generator would be exhausted by - the first read. - - Driven off :py:func:`dataclasses.fields` rather than a written-out - list, so a codec field added later is normalized without anyone - remembering to name it here. The ``_codecs`` suffix is what marks a - field as one of them, leaving room for a future field that is not a - codec collection and must not be turned into a tuple. - """ + """Normalize each codec field to a tuple, rejecting what cannot become one.""" + # A bundle that writes `logical_extension_codecs=codec` instead of + # `(codec,)` is contributing one codec, not an iterable of them. + # Without this, the mistake surfaces inside `with_extensions` as + # `'MyCodec' object is not iterable`, which names neither the field nor + # the hook that built it. Checking here puts the error in the extension + # library's own frame. + # + # Normalizing is worth doing on its own: the declared type is a tuple + # and the class is frozen, so a list left in place would be a mutable + # member of an immutable value, and a generator would be exhausted by + # the first read. + # + # Driven off `dataclasses.fields` rather than a written-out list, so a + # codec field added later is normalized without anyone remembering to + # name it here. The `_codecs` suffix is what marks a field as one of + # them, leaving room for a future field that is not a codec collection + # and must not be turned into a tuple. for field in fields(self): name = field.name if not name.endswith("_codecs"): @@ -242,19 +235,19 @@ class SessionExtensionExportable(Protocol): mutating the context they are handed — a registration made during binding is not rolled back if a later extension fails. - ``ctx`` is the right session but not yet the final codec chains: this hook - runs before anything is installed, so ``ctx`` still carries whatever chains - the receiver had. Take the task-context provider off it — that is bound to - the session and is what the components need — but do not read its codec - chains expecting to find this call's codecs, including your own. - :py:class:`SessionPlannerExportable` is the hook that sees the completed - chains, which is why a planner that wraps the host's codecs builds them - there rather than here. - A bundle that also contributes a query planner implements - :py:class:`SessionPlannerExportable` alongside this protocol. Planners are - installed in a second phase, so they are not part of the components - returned here. + :py:class:`SessionPlannerExportable` alongside this protocol. + + Args: + ctx: The session the components will run on. Take the task-context + provider off it. Do **not** read its codec chains expecting to find + this call's codecs, including your own: this hook runs before + anything is installed, so ``ctx`` still carries whatever chains the + receiver had. :py:class:`SessionPlannerExportable` is the hook that + sees the completed chains — see :ref:`extension_bundles_two_phases`. + + Returns: + The codecs this bundle contributes. Examples: >>> from datafusion import ( @@ -291,14 +284,15 @@ class SessionPlannerExportable(Protocol): The hook runs after every codec from every bundle is installed, and ``ctx`` is the context carrying those final chains. That ordering is the point: a planner captured here sees the complete codec set, so a nested planner is - not left encoding through a chain that a later bundle has grown. + not left encoding through a chain that a later bundle has grown. See + :ref:`extension_bundles_two_phases`. - Return ``None`` to contribute no planner and leave ``fallback`` in place. - That is the no-op, and it is not the same as returning ``fallback``: the - capsule the first bundle receives wraps the session's planner for export, so - handing it back installs it as a foreign planner and every later plan crosses - an FFI boundary that was not there before. A bundle with nothing to - contribute returns ``None``. + **Return ``None`` to contribute no planner**, leaving ``fallback`` in + place. That is the no-op, and it is not the same as returning ``fallback``: + the capsule the first bundle receives wraps the session's planner for + export, so handing it back installs that planner as a foreign one and every + later plan crosses an FFI boundary that was not there before. A bundle that + decides at runtime it has nothing to contribute returns ``None``. Ignoring ``fallback`` and returning a planner that does not delegate to it is legal and means "replace" — but it discards every planner listed before @@ -311,17 +305,29 @@ class SessionPlannerExportable(Protocol): bundle this is the session's existing planner, which is the DataFusion default unless one was installed earlier. + Returns: + A planner wrapping ``fallback``, or ``None`` to contribute none. + Examples: - >>> from datafusion import SessionPlannerExportable - >>> class MyEngineExtension: + A real library returns its own planner wrapping ``fallback``, e.g. + ``my_library.Planner(fallback=fallback)``. The two degenerate cases are + worth contrasting, because both plan queries successfully and only one + of them is the no-op: + + >>> from datafusion import SessionContext, SessionPlannerExportable + >>> class Contributes: + ... def __datafusion_session_planner__(self, ctx, fallback): + ... return None # the no-op: session keeps its own planner + >>> class Replaces: ... def __datafusion_session_planner__(self, ctx, fallback): - ... # A real library returns its own planner wrapping - ... # `fallback`, e.g. ``my_library.Planner(fallback=fallback)``. - ... # Handing it straight back is the degenerate wrap: legal, - ... # but it still installs `fallback` as a foreign planner. - ... # Return None instead to contribute nothing. - ... return fallback - >>> isinstance(MyEngineExtension(), SessionPlannerExportable) + ... return fallback # installs it as a *foreign* planner + >>> for bundle in (Contributes(), Replaces()): + ... ctx = SessionContext().with_extensions(bundle) + ... ctx.sql("SELECT 1 AS n").collect()[0].column(0).to_pylist() + [1] + [1] + + >>> isinstance(Contributes(), SessionPlannerExportable) True >>> isinstance(object(), SessionPlannerExportable) False diff --git a/python/datafusion/ipc.py b/python/datafusion/ipc.py index 487abd4c3..d9f63e267 100644 --- a/python/datafusion/ipc.py +++ b/python/datafusion/ipc.py @@ -175,11 +175,16 @@ def set_sender_ctx(ctx: SessionContext) -> None: Examples: >>> from datafusion import SessionContext - >>> from datafusion.ipc import set_sender_ctx, get_sender_ctx + >>> from datafusion.ipc import ( + ... clear_sender_ctx, + ... get_sender_ctx, + ... set_sender_ctx, + ... ) >>> driver = SessionContext().with_python_udf_inlining(enabled=False) >>> set_sender_ctx(driver) >>> get_sender_ctx() is driver True + >>> clear_sender_ctx() """ _local.sender_ctx = ctx @@ -224,7 +229,9 @@ def _resolve_ctx( Priority: explicit argument > worker context > global context. Falling back to the global :class:`SessionContext` (instead of a freshly constructed one) preserves any registrations the user has - installed on it. + installed on it. It carries no extension codecs, though, so a worker + that must decode payloads written by one has to install a context + carrying them via :func:`set_worker_ctx`. Examples: >>> from datafusion import SessionContext diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index b2c6eab3e..8d03bae2c 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -82,9 +82,9 @@ def display_indent_schema(self) -> str: def display_graphviz(self) -> str: """Print the graph visualization of the logical plan. - Returns a `format`able structure that produces lines meant for graphical display - using the `DOT` language. This format can be visualized using software from - [`graphviz`](https://graphviz.org/) + Returns a formattable structure that produces lines meant for graphical + display using the ``DOT`` language. This format can be visualized using + software from `graphviz `_. """ return self._raw_plan.display_graphviz() @@ -92,21 +92,33 @@ def display_graphviz(self) -> str: def from_bytes(ctx: SessionContext, data: bytes) -> LogicalPlan: """Create a LogicalPlan from serialized protobuf bytes. - Decoding routes through the session's installed - `LogicalExtensionCodec`. Tables created in memory from record - batches are currently not supported. + Decoding routes through the codecs installed on ``ctx`` with + :py:meth:`~datafusion.SessionContext.with_logical_extension_codec`. + Tables created in memory from record batches are currently not + supported. + + Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and + positional, and there is no fallback to a worker or global context. + + See Also: + :py:meth:`to_bytes`, :py:meth:`ExecutionPlan.from_bytes`, + :py:meth:`datafusion.Expr.from_bytes`. """ return LogicalPlan(df_internal.LogicalPlan.from_bytes(ctx.ctx, data)) def to_bytes(self, ctx: SessionContext | None = None) -> bytes: """Convert a LogicalPlan to serialized protobuf bytes. - When ``ctx`` is supplied, encoding routes through the session's - installed `LogicalExtensionCodec` so user FFI codecs (registered - via :py:meth:`SessionContext.with_logical_extension_codec`) see - the encode path. With ``ctx=None`` a default codec is used. - Tables created in memory from record batches are currently not - supported. + When ``ctx`` is supplied, encoding routes through the codecs + installed on it with + :py:meth:`~datafusion.SessionContext.with_logical_extension_codec`, + so extension codecs see the encode path. With ``ctx=None`` a + default codec is used. Tables created in memory from record + batches are currently not supported. + + See Also: + :py:meth:`from_bytes`, :py:meth:`ExecutionPlan.to_bytes`, + :py:meth:`datafusion.Expr.to_bytes`. """ ctx_arg = ctx.ctx if ctx is not None else None return self._raw_plan.to_bytes(ctx_arg) @@ -145,7 +157,7 @@ def __init__(self, plan: df_internal.ExecutionPlan) -> None: self._raw_plan = plan def children(self) -> list[ExecutionPlan]: - """Get a list of children `ExecutionPlan` that act as inputs to this plan. + """Get a list of children ``ExecutionPlan`` that act as inputs to this plan. The returned list will be empty for leaf nodes such as scans, will contain a single value for unary nodes, or two values for binary nodes (such as joins). @@ -173,18 +185,35 @@ def partition_count(self) -> int: def from_bytes(ctx: SessionContext, data: bytes) -> ExecutionPlan: """Create an ExecutionPlan from serialized protobuf bytes. - Decoding routes through the session's installed - `PhysicalExtensionCodec`. Tables created in memory from record - batches are currently not supported. + Decoding routes through the codecs installed on ``ctx`` with + :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`. + Tables created in memory from record batches are currently not + supported. + + Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and + positional, and there is no fallback to a worker or global context. + + See Also: + :py:meth:`to_bytes`, :py:meth:`LogicalPlan.from_bytes`. """ return ExecutionPlan(df_internal.ExecutionPlan.from_bytes(ctx.ctx, data)) def to_bytes(self, ctx: SessionContext | None = None) -> bytes: """Convert an ExecutionPlan into serialized protobuf bytes. - When ``ctx`` is supplied, encoding routes through the session's - installed `PhysicalExtensionCodec`. Tables created in memory - from record batches are currently not supported. + When ``ctx`` is supplied, encoding routes through the codecs + installed on it with + :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`. + Tables created in memory from record batches are currently not + supported. + + Round-tripping through this method and :py:meth:`from_bytes` is how + an extension library checks that its own codec claimed its nodes, + rather than a codec installed earlier in the chain — see + :ref:`extension_codec_order`. + + See Also: + :py:meth:`from_bytes`, :py:meth:`LogicalPlan.to_bytes`. """ ctx_arg = ctx.ctx if ctx is not None else None return self._raw_plan.to_bytes(ctx_arg) diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index eafcefdaf..7e2317ff6 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -116,11 +116,6 @@ def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]: class LogicalExtensionCodecExportable(Protocol): """Type hint for objects exposing ``__datafusion_logical_extension_codec__``. - ``session`` is the :py:class:`~datafusion.context.SessionContext` the codec - is being installed on. Take the task context provider from it rather than - building a session of your own, so the decode callbacks resolve names - against the session that runs the query. - Implement the codec itself exactly as you would for a session that installs only yours. A session may hold several codecs, but each payload records the codec that wrote it and is only ever handed back to that codec, so there is @@ -132,6 +127,43 @@ class LogicalExtensionCodecExportable(Protocol): of this protocol and is rarely needed: declare it when renaming your class must not stop older plans from decoding, or when one library installs two instances that own disjoint slices of the wire format. + + Args: + session: A handle on the session this codec is being installed on. Take + the task context provider from it rather than building a session of + your own, so the decode callbacks resolve names against the session + that runs the query. + + **Duck-type it; do not check its type.** This argument is the + reference for every capsule getter in the protocol, and across them + it is not always a session at all: the provider and catalog getters + receive the host's logical codec as a bare ``PyCapsule``. Even when + it *is* a session it is the PyO3 context from + ``datafusion._internal``, not the + :py:class:`~datafusion.context.SessionContext` wrapper, so it + exposes every capsule getter and ``__datafusion_codec_id__`` — + everything the protocol asks of it — but + ``isinstance(session, SessionContext)`` is ``False`` even though its + ``repr`` reads ``datafusion.SessionContext``. Pass it to + ``ffi_logical_codec_from_pycapsule``, which handles both. + + The two bundle hooks are the exception: + :py:class:`~datafusion.extensions.SessionExtensionExportable` and + :py:class:`~datafusion.extensions.SessionPlannerExportable` are + dispatched from Python and receive the wrapper. See + :ref:`extension_getter_argument`. + + Returns: + A ``datafusion_logical_extension_codec`` PyCapsule. + + Examples: + A :py:class:`~datafusion.context.SessionContext` satisfies this + protocol, which is how you obtain a real capsule to test against: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> type(ctx.__datafusion_logical_extension_codec__(ctx)).__name__ + 'PyCapsule' """ def __datafusion_logical_extension_codec__( # noqa: D105 diff --git a/python/tests/test_docstrings.py b/python/tests/test_docstrings.py new file mode 100644 index 000000000..484698592 --- /dev/null +++ b/python/tests/test_docstrings.py @@ -0,0 +1,158 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Shape checks on the docstrings that ship in the wheel. + +These are about *form*, not content: a docstring that has grown into an essay, +a documented API with no example, or a pointer at the guide that the rendered +HTML cannot follow. See "Python Function Docstrings" in ``AGENTS.md``. + +Content is checked elsewhere. ``--doctest-modules`` executes the examples, and +Sphinx builds with ``--fail-on-warning`` so a broken ``:ref:`` fails the docs +build. Neither of those notices a docstring that is simply too long, or one +that never had an example to run. +""" + +from __future__ import annotations + +import ast +import inspect +import pathlib +import re + +import datafusion +import pytest +from datafusion import SessionContext, extensions + +PACKAGE_ROOT = pathlib.Path(datafusion.__file__).parent + +# Above this, a docstring has stopped being a contract and become a +# narrative. Move the argument to a guide page under `docs/source/` and leave +# a one-line pointer. Raising this is not the fix. +# +# Deliberately has no waiver list: at the time of writing the longest +# docstring in the package is `udaf` at 92 lines, which is a legitimately +# long UDF-authoring reference with many worked examples. If a genuinely +# necessary docstring ever exceeds this, prefer splitting the API. +MAX_DOCSTRING_LINES = 95 + +# An *inclusion* list, to be grown — not an exclusion list to be shrunk. +# Enforcing "every public callable has an example" package-wide is a separate +# project; several hundred callables do not have one yet. These are the +# extension-protocol surface, where an undocumented method is the difference +# between an extension author succeeding and filing an issue. +DOCSTRINGS_REQUIRING_EXAMPLES: list[tuple[str, object]] = [ + *[ + (f"SessionContext.{name}", getattr(SessionContext, name)) + for name in ( + "with_extensions", + "set_query_planner", + "with_logical_extension_codec", + "with_physical_extension_codec", + "logical_extension_codec_ids", + "physical_extension_codec_ids", + "with_python_udf_inlining", + "__datafusion_codec_id__", + "__datafusion_logical_extension_codec__", + "__datafusion_physical_extension_codec__", + "__datafusion_query_planner__", + ) + ], + *[(name, getattr(extensions, name)) for name in extensions.__all__], +] + + +def _iter_docstrings() -> list[tuple[str, int]]: + """Yield ``(location, line count)`` for every docstring under ``python/``.""" + found = [] + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if not isinstance( + node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef) + ): + continue + doc = ast.get_docstring(node, clean=True) + if not doc: + continue + name = getattr(node, "name", "") + rel = path.relative_to(PACKAGE_ROOT.parent) + line = getattr(node, "lineno", 1) + found.append((f"{rel}:{line} {name}", len(doc.splitlines()))) + return found + + +def test_no_docstring_is_an_essay() -> None: + """No docstring exceeds the length ceiling.""" + too_long = [ + (where, count) + for where, count in _iter_docstrings() + if count > MAX_DOCSTRING_LINES + ] + assert not too_long, ( + "These docstrings exceed " + f"{MAX_DOCSTRING_LINES} lines and have become narrative rather than " + "contract:\n" + + "\n".join(f" {where} — {count} lines" for where, count in too_long) + + "\n\nMove the argument to a guide page under docs/source/ and leave a " + "one-line pointer with a :ref:. See 'One canonical home per claim' in " + "AGENTS.md." + ) + + +@pytest.mark.parametrize( + ("name", "obj"), + DOCSTRINGS_REQUIRING_EXAMPLES, + ids=[name for name, _ in DOCSTRINGS_REQUIRING_EXAMPLES], +) +def test_extension_api_has_a_doctest(name: str, obj: object) -> None: + """Every extension-protocol member carries at least one example.""" + doc = inspect.getdoc(obj) + assert doc, f"{name} has no docstring" + assert ">>>" in doc, ( + f"{name} has a docstring but no example. Everything on the extension " + "protocol needs one, even if the realistic usage has to be marked " + "+SKIP — in which case a runnable example goes above it. See " + "'Examples that need a compiled extension' in AGENTS.md." + ) + + +def test_no_dead_pointers_at_the_guide() -> None: + """A docstring naming the guide must link it, not just mention it.""" + role = re.compile(r":(ref|doc|py:\w+):`") + dead = [] + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if not isinstance( + node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef) + ): + continue + doc = ast.get_docstring(node, clean=True) + if not doc or "guide" not in doc.lower(): + continue + if role.search(doc): + continue + rel = path.relative_to(PACKAGE_ROOT.parent) + name = getattr(node, "name", "") + dead.append(f"{rel}:{getattr(node, 'lineno', 1)} {name}") + assert not dead, ( + "These docstrings send the reader to a guide without a resolvable " + "link, which is a dead end in the rendered HTML:\n" + + "\n".join(f" {where}" for where in dead) + + "\n\nUse :ref:`some_label` naming the specific section." + ) From 61ca5ff6d6c4088be692251e4e303828efd5184b Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 08:54:12 -0400 Subject: [PATCH 31/33] docs: point the non-Sphinx consumers at the new pages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four links in the two example READMEs pointed into `docs/source/.../ffi.md` with heading anchors. Sphinx does not link-check those, so deleting the page would have left them silently 404ing. They now use published URLs, which also fixes a second problem: a relative path into `docs/source` only renders on github.com and is broken for anyone reading the README from crates.io, an sdist, or a vendored copy. `grep -rn "docs/source" examples/` is now empty. `.ai/skills/ffi-capsule-protocol/SKILL.md` named the deleted page as "where the truth is", and AGENTS.md sends agents to that skill before they touch a capsule getter. It now points at the specific extension-guide pages and separately at `ffi-internals.md`. `llms.txt` filed the whole subject under "Optional" as "extending the Python bindings" — the wrong shelf for the headline feature of a major release, since that section means "skippable on a tight context budget". It gains an Extensions and distribution section, `datafusion.extensions` and `datafusion.ipc` in the API list, both FFI example crates, and the corrected `distributing-work` URL. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 7 ++++++- docs/source/llms.txt | 14 +++++++++++++- examples/datafusion-ffi-example/README.md | 4 ++-- .../datafusion-ffi-query-planner-example/README.md | 4 ++-- 4 files changed, 23 insertions(+), 6 deletions(-) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 1efc93ec5..ca3a09d35 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -223,7 +223,12 @@ pins that; changing it should be deliberate. ## Where the truth is -- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. +- `docs/source/extension-guide/` — the protocol, for the library author. + `capsule-protocol.md` has the hook convention and what the getter argument + actually is; `codecs.md`, `bundles.md`, and `query-planners.md` have the + per-component rules; `index.md` lists all 18 hooks. +- `docs/source/contributor-guide/ffi-internals.md` — why the framing is shaped + this way, including the weak-`Arc` scheme and the one-level rebind. - `docs/source/user-guide/upgrade-guides.md` — every past migration. - `crates/core/src/codec.rs` — the codec chain: the envelope, identity dispatch, and the two unframed cases from Rule 8. diff --git a/docs/source/llms.txt b/docs/source/llms.txt index 3ff6b3813..24d03ce57 100644 --- a/docs/source/llms.txt +++ b/docs/source/llms.txt @@ -16,6 +16,14 @@ - [Common operations](https://datafusion.apache.org/python/user-guide/common-operations/index.html): select, filter, join, aggregate, window, expressions, and functions. - [SQL](https://datafusion.apache.org/python/user-guide/sql.html): running SQL against registered tables. - [Configuration](https://datafusion.apache.org/python/user-guide/configuration.html): session and runtime options. +- [Using extension libraries](https://datafusion.apache.org/python/user-guide/extensions.html): installing a compiled extension and registering its tables, functions, or execution backend with `SessionContext.with_extensions`. + +## Extensions and distribution + +- [Using extension libraries](https://datafusion.apache.org/python/user-guide/extensions.html): which extensions register directly and which need `with_extensions`, plus the two failure modes (a collected context, a version mismatch). +- [Distributing work](https://datafusion.apache.org/python/user-guide/distributing-work/index.html): choosing between shipping expressions to a worker pool and installing a distributed query engine. +- [Shipping expressions to workers](https://datafusion.apache.org/python/user-guide/distributing-work/expressions.html): pickling `Expr` to `multiprocessing` or Ray, what travels inline vs by name, `datafusion.ipc` worker/sender contexts, and the security model. +- [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html): the PyCapsule protocol for shipping table providers, catalogs, functions, extension codecs, and query planners from a Rust cdylib. Includes the full `__datafusion_*__` hook reference (18 hooks) and an author checklist. ## DataFrame API reference @@ -23,14 +31,18 @@ - [`datafusion.expr`](https://datafusion.apache.org/python/autoapi/datafusion/expr/index.html): expression tree nodes (`Expr`, `Window`, `WindowFrame`, `GroupingSet`). - [`datafusion.functions`](https://datafusion.apache.org/python/autoapi/datafusion/functions/index.html): 290+ scalar, aggregate, and window functions. - [`datafusion.context.SessionContext`](https://datafusion.apache.org/python/autoapi/datafusion/context/index.html): session entry point, data loading, SQL execution. +- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionExtensionExportable`, `SessionPlannerExportable`, `QueryPlannerExportable` — the extension-bundle protocol. +- [`datafusion.ipc`](https://datafusion.apache.org/python/autoapi/datafusion/ipc/index.html): worker and sender context slots for shipping expressions between processes. ## Examples - [TPC-H queries (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples/tpch): canonical translations of TPC-H Q01–Q22 to idiomatic DataFrame code, each with reference SQL embedded in the module docstring. - [Other examples (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples): UDF/UDAF/UDWF, Substrait, Pandas/Polars interop, S3 reads. +- [FFI extension example (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example): a Rust cdylib exporting table providers, catalogs, functions, and extension codecs. Has build commands in its README. +- [FFI query planner example (GitHub)](https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example): a Rust cdylib exporting a query planner and a complete extension bundle. ## Optional -- [Contributor guide](https://datafusion.apache.org/python/contributor-guide/introduction.html): building from source, extending the Python bindings. +- [Contributor guide](https://datafusion.apache.org/python/contributor-guide/index.html): building from source and changing datafusion-python itself. - [Upgrade guides](https://datafusion.apache.org/python/user-guide/upgrade-guides.html): migration notes between releases. - [Upstream Rust `DataFusion`](https://datafusion.apache.org/): the underlying query engine. diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index 401eb610a..aadea909f 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -49,6 +49,6 @@ ctx = ctx.with_physical_extension_codec(provider_physical_codec) ctx.set_query_planner(planner) ``` -Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. The exception is a planner that wraps a fallback: the rebuild reaches the installed planner only, not the fallback inside it, so codecs-first is a requirement there. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on. +Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. The exception is a planner that wraps a fallback: the rebuild reaches the installed planner only, not the fallback inside it, so codecs-first is a requirement there. See [Rebinding a planner's codecs is one level deep](https://datafusion.apache.org/python/extension-guide/query-planners.html#install-codecs-before-a-layered-planner), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on. -For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. +For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see the [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html). diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 4b9140b53..f21aee630 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -78,6 +78,6 @@ ctx.set_query_planner(MyQueryPlanner()) `MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. -The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). +The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](https://datafusion.apache.org/python/extension-guide/query-planners.html#install-codecs-before-a-layered-planner). -For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. +For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see the [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html). From 5a1bfeba1ac9fd27ff3c849c25ec5592b86ebdbe Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 10:31:39 -0400 Subject: [PATCH 32/33] docs: fix extension-guide review findings, rename phase-one bundle hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six findings from a read-through of the new extension guide, plus the API rename one of them turned into. The "process local tokens" note sat in the guide index with nothing around it to explain what a token was or why the reader should care. It moves to a new `extension_codec_durable_metadata` section in `codecs.md`, where the reader is already thinking about what goes in a payload: what to encode, then what the examples do instead, and the three consequences that follow from parking live objects in a process-local map — no double decode, no fan-out, and a leak for any plan that never reaches a decoder. The checklist item now points there instead of at the guide index. The hook reference loses its `Capsule name` column, which restated `datafusion_` for every row when the naming rule already derives it, and gains a `Contributes` column instead. The argument column stays: those four values are protocol, not a signature, and only 4 of the 18 hooks have a Python definition to link at all — the rest are host-side imports, so links into Rust source would rot faster than the table. Staleness is handled by `test_hook_reference_table_lists_every_hook` instead, which greps `crates/` and `python/datafusion/` for `__datafusion_*__` and diffs the set against the table rows. Verified it fails when a row is dropped. `capsule-protocol.md` described `abi_stable`, which datafusion-ffi no longer uses. It now describes `stabby` and the part that is not stabby: `FFI_Option` and `FFI_Result` are datafusion-ffi's own, because stabby's require `T: IStable` and the `FFI_*` structs hold self-referential function pointers. The conversion example converts to `Arc` rather than naming `ForeignTableProvider`, since the `From` impl compares library markers and returns the original `Arc` when both sides are the same library. Three other snippets on that page had gone stale with it: `FFI_TableProvider::new` with three arguments, `PyCapsule::new_bound`, and a receiving snippet whose variable was named `codec`. In `table-providers.md` all five `Registered with` cells now render `Receiver.method`, so the schema row reads `Catalog.register_schema` rather than a bare dotted path, with one sentence on reaching a `Catalog` first. `Other session components` was in `functions.md`, where an optimizer rule and a config struct are neither functions nor tables; it becomes its own page, `other-components.md`, carrying the `extension_other_hooks` label so the index rows still resolve. Three guide pages named individual tests, which invites exactly the divergence the reference is supposed to prevent. They name the suite now. Finally the phase-one bundle hook. `__datafusion_session_extension__` reused the name of the whole thing it is a hook on — `with_extensions` takes extensions and `SessionExtensionExportable` is the bundle protocol — while its sibling `__datafusion_session_planner__` is named for its content. It is now `__datafusion_session_components__`, matching both its sibling and the `SessionExtensionComponents` it returns, which leaves room for the UDF and provider fields that will join the codec fields later. The protocol class follows it to `SessionComponentsExportable`, since that file's convention is one class per hook name. Neither name has shipped, so the upgrade guide needs no before-and-after; it introduces both hooks as new in this release and names the new one. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 2 +- docs/source/conf.py | 2 +- docs/source/extension-guide/bundles.md | 10 +-- .../extension-guide/capsule-protocol.md | 58 +++++++++------ docs/source/extension-guide/checklist.md | 2 +- docs/source/extension-guide/codecs.md | 21 +++++- docs/source/extension-guide/functions.md | 47 ------------ docs/source/extension-guide/index.md | 61 ++++++++-------- .../extension-guide/other-components.md | 72 +++++++++++++++++++ docs/source/extension-guide/query-planners.md | 4 +- docs/source/extension-guide/sessions.md | 3 +- .../source/extension-guide/table-providers.md | 15 ++-- docs/source/llms.txt | 2 +- docs/source/user-guide/upgrade-guides.md | 2 +- .../README.md | 2 +- .../_test_three_library_query_planner.py | 20 +++--- .../src/extension.rs | 4 +- python/datafusion/__init__.py | 4 +- python/datafusion/context.py | 16 ++--- python/datafusion/extensions.py | 20 +++--- python/datafusion/user_defined.py | 2 +- python/tests/test_context.py | 22 +++--- python/tests/test_docstrings.py | 40 +++++++++++ python/tests/test_imports.py | 2 +- 24 files changed, 268 insertions(+), 165 deletions(-) create mode 100644 docs/source/extension-guide/other-components.md diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index ca3a09d35..37be0bc9d 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -70,7 +70,7 @@ Wrap `fallback` and delegate to it; returning a planner that ignores it discards every layer beneath, including one the session already had. It runs after every bundle's codecs are installed, so `ctx` carries the final chains. -That is also the only hook where it does. `__datafusion_session_extension__` +That is also the only hook where it does. `__datafusion_session_components__` runs before anything is installed, so its `ctx` still carries the chains the receiver had — the same session, and the same task-context provider, but not this call's codecs, not even your own. Read the host's codec chains in the diff --git a/docs/source/conf.py b/docs/source/conf.py index 2463c580b..c9f845f12 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -123,7 +123,7 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool: # noqa ("class", "datafusion.SessionContext"), ("class", "datafusion.QueryPlannerExportable"), ("class", "datafusion.SessionExtensionComponents"), - ("class", "datafusion.SessionExtensionExportable"), + ("class", "datafusion.SessionComponentsExportable"), ("class", "datafusion.SessionPlannerExportable"), ("module", "datafusion.common"), # Duplicate modules (skip module-level docs to avoid duplication) diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index fffbedb59..791fdfc8e 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -40,7 +40,7 @@ bundle object implementing one or both of two hooks: ```python class MyEngineExtension: - def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: + def __datafusion_session_components__(self, ctx: SessionContext) -> SessionExtensionComponents: # Phase one. Create fresh components bound to `ctx` on every call. return SessionExtensionComponents( logical_extension_codecs=(self._make_logical_codec(ctx),), @@ -95,7 +95,7 @@ class Bundle: def __init__(self, codec_id): self.codec_id = codec_id - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): # Fresh components on every call, bound to the `ctx` handed in. # Never cache these, and never retain `ctx`. return SessionExtensionComponents( @@ -133,7 +133,7 @@ routes to the codec that wrote the payload. A session holds exactly **one** query planner, so planners cannot accumulate — they compose by *nesting*, each wrapping the one before it and delegating to it for work it does not handle. -So `with_extensions` runs every `__datafusion_session_extension__` and installs +So `with_extensions` runs every `__datafusion_session_components__` and installs all the codecs, and only then runs each `__datafusion_session_planner__`, in argument order, handing each the planner built so far. Two consequences worth holding onto: @@ -204,8 +204,8 @@ class CodecsOf: def __init__(self, inner): self.inner = inner - def __datafusion_session_extension__(self, ctx): - return self.inner.__datafusion_session_extension__(ctx) + def __datafusion_session_components__(self, ctx): + return self.inner.__datafusion_session_components__(ctx) class PlannerOf: diff --git a/docs/source/extension-guide/capsule-protocol.md b/docs/source/extension-guide/capsule-protocol.md index ada5a6f57..3b8331b4b 100644 --- a/docs/source/extension-guide/capsule-protocol.md +++ b/docs/source/extension-guide/capsule-protocol.md @@ -35,12 +35,17 @@ review the code and documentation in the [datafusion-ffi] crate. Our FFI implementation is narrowly focused on sharing data and functions with Rust backed libraries. This allows us to use the -[abi_stable crate](https://crates.io/crates/abi_stable). This is an excellent -crate that allows for easy conversion between Rust native types and FFI-safe -alternatives. For example, if you needed to pass a `Vec` via FFI, you -can simply convert it to an `RVec` in an intuitive manner. It also -supports features like `RResult` and `ROption` that do not have an obvious -translation to a C equivalent. +[stabby crate](https://crates.io/crates/stabby), which converts between Rust +native types and FFI-safe alternatives. For example, if you needed to pass a +`Vec` via FFI, you can convert it to a +`stabby::vec::Vec` — the crate's own examples alias +these as `SVec` and `SString`, which is the convention [datafusion-ffi] follows +too. + +For `Option` and `Result`, [datafusion-ffi] defines its own `FFI_Option` and +`FFI_Result` rather than using stabby's. Stabby's versions require +`T: IStable` for niche optimization, and many of the `FFI_*` structs hold +self-referential function pointers that cannot implement it. ## `FFI_` on the provider, `Foreign` on the receiver @@ -51,17 +56,19 @@ defined a custom and you want to create a sharable FFI counterpart, you could write: ```rust -let my_provider = MyTableProvider::default(); -let ffi_provider = FFI_TableProvider::new(Arc::new(my_provider), false, None); +let my_provider = Arc::new(MyTableProvider::default()); +let ffi_provider = FFI_TableProvider::new_with_ffi_codec(my_provider, false, None, codec); ``` +where `codec` is the host's logical codec, read off the argument your getter +was handed — see {ref}`extension_getter_argument`. + If you were interfacing with a library that provided the above -`FFI_TableProvider` and you needed to turn it back into a `TableProvider`, you -can turn it into a `ForeignTableProvider`, which implements the `TableProvider` -trait: +`FFI_TableProvider` and you needed a usable `TableProvider` back, you convert +it into an `Arc`: ```rust -let foreign_provider: ForeignTableProvider = ffi_provider.into(); +let provider: Arc = (&ffi_provider).into(); ``` If you review the code in [datafusion-ffi] you will find that each of the @@ -74,6 +81,13 @@ example we're showing, this means the code that has written the underlying structures with the `Foreign` prefix are to be used by the receiver. In this case, it is the `datafusion-python` library. +Convert to the trait object rather than naming `ForeignTableProvider` yourself. +The conversion compares the provider's library marker against the receiver's: +when both sides turn out to be the same shared library it hands back the +original `Arc` and skips the boundary entirely, and only otherwise wraps it in +a `ForeignTableProvider`. Which one you get is an implementation detail, and +both implement `TableProvider`. + ## Wrapping it in a capsule In order to share these FFI structures, we need to wrap them in some kind of @@ -82,20 +96,20 @@ described in {ref}`extension_why_ffi`, we use `PyCapsule`. We can create a `PyCapsule` for our provider thusly: ```rust -let name = CString::new("datafusion_table_provider")?; -let my_capsule = PyCapsule::new_bound(py, provider, Some(name))?; +PyCapsule::new_with_value(py, ffi_provider, cr"datafusion_table_provider") ``` -On the receiving side, turn this pycapsule object into the -`FFI_TableProvider`, which can then be turned into a `ForeignTableProvider`; -the associated code is: +On the receiving side, read the `FFI_TableProvider` back out of the capsule and +convert it, which is what `table_provider_from_pycapsule` in `crates/util` does: ```rust -let capsule = capsule.cast::()?; +validate_pycapsule(capsule, "datafusion_table_provider")?; let data: NonNull = capsule - .pointer_checked(Some(name))? + .pointer_checked(Some(c"datafusion_table_provider"))? .cast(); -let codec = unsafe { data.as_ref() }; +let ffi_provider = unsafe { data.as_ref() }; +check_ffi_version("table provider", unsafe { (ffi_provider.version)() })?; +let provider: Arc = ffi_provider.into(); ``` ## The naming rule @@ -112,7 +126,7 @@ must return a capsule named `datafusion_table_provider`. Return a capsule with the wrong name and the import fails with an error naming both the name found and the name expected, rather than reading the pointer as the wrong type. -The full list of hooks and their capsule names is in the +The full list of hooks is in the {ref}`hook reference `. `TableProvider` was the first extension written this way and is the most thoroughly implemented; every hook added since follows the same pattern. @@ -196,7 +210,7 @@ getter and `__datafusion_codec_id__` — everything the protocol asks of it — `isinstance(session, SessionContext)` is `False` in Python even though its `repr` reads `datafusion.SessionContext`. -The two bundle hooks are the exception: `__datafusion_session_extension__` and +The two bundle hooks are the exception: `__datafusion_session_components__` and `__datafusion_session_planner__` are dispatched from Python by {py:meth}`~datafusion.SessionContext.with_extensions`, so they receive the wrapper. See {doc}`bundles`. diff --git a/docs/source/extension-guide/checklist.md b/docs/source/extension-guide/checklist.md index 944fa6abf..6e204b3f5 100644 --- a/docs/source/extension-guide/checklist.md +++ b/docs/source/extension-guide/checklist.md @@ -93,7 +93,7 @@ publish. Each links to the page that explains it. - [ ] **Your production codec serializes durable metadata**, not a process-local token. The examples in this repository use tokens to make ownership observable; that is a demonstration, not a pattern. - → {ref}`extension_guide` + → {ref}`extension_codec_durable_metadata` - [ ] **You have integration tests across a real FFI boundary.** The two example crates in this repository are the pattern: build the cdylib, install the wheel, then exercise it from Python. diff --git a/docs/source/extension-guide/codecs.md b/docs/source/extension-guide/codecs.md index c7f051324..160439927 100644 --- a/docs/source/extension-guide/codecs.md +++ b/docs/source/extension-guide/codecs.md @@ -50,6 +50,24 @@ A codec that also ships to hosts which dispatch differently may still want its own guard against foreign payloads. Keeping one is fine; it is simply not needed for the datafusion-python path. +(extension_codec_durable_metadata)= + +## Encode metadata, not a handle to a live object + +Your payload has to be enough to rebuild the object somewhere your process is +not. Write the metadata a fresh instance can be constructed from — a path, a +connection string, a schema, the options the object was created with. + +The example codecs in this repository do not do this, and it is worth knowing +before copying them. They keep a process-local `HashMap` of live providers and +encode an integer token into it: encoding inserts, decoding removes. That makes +Rust type identity observable across three separately loaded libraries in one +test, which is what the examples exist to show. It also means a decode consumes +its token, so the same bytes cannot be decoded twice, one encoded plan cannot +fan out to several readers, and a plan that never reaches a decoder keeps its +provider alive for the life of the process. A real codec has none of those +properties because it does not park the object anywhere. + (extension_codec_ids)= ## Codec ids @@ -133,8 +151,7 @@ after it. The query still succeeds. What changes is which library wrote the bytes — so a plan that has to decode in another process now needs whichever library happened to win, not the one whose node it is. `MyPhysicalExtensionCodec` in [`datafusion-ffi-example`] claims this way, and -`test_a_greedy_codec_installed_first_claims_another_librarys_node` pins the -consequence. +the query-planner example's test suite pins the consequence. Two rules of thumb: diff --git a/docs/source/extension-guide/functions.md b/docs/source/extension-guide/functions.md index 92a30b7d9..fcfef9f1a 100644 --- a/docs/source/extension-guide/functions.md +++ b/docs/source/extension-guide/functions.md @@ -103,51 +103,4 @@ on your decoder — with an empty payload there is no id to route on, so your `NameOnlyUdfCodec` in [`datafusion-ffi-example`] is the worked case. -(extension_other_hooks)= - -## Other session components - -Two further hooks contribute things that are neither data nor functions. Both -take no argument and both are implemented in [`datafusion-ffi-example`]. - -**`__datafusion_physical_optimizer_rule__`** contributes a rule that rewrites -physical plans, installed with -{py:meth}`~datafusion.SessionContext.add_physical_optimizer_rule`. Reach for -this rather than a {doc}`query planner ` when you want to -adjust the plan DataFusion produced rather than produce it yourself — it is -much the smaller commitment, and rules accumulate where planners nest. - -```rust -fn __datafusion_physical_optimizer_rule__<'py>( - &self, - py: Python<'py>, -) -> PyResult> { - let rule: Arc = Arc::new(self.clone()); - let runtime = get_tokio_runtime().handle().clone(); - let ffi = FFI_PhysicalOptimizerRule::new(rule, Some(runtime)); - - PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_optimizer_rule") -} -``` - -**`__datafusion_extension_options__`** contributes typed configuration entries -that your components can read back out of the session config, installed with -{py:meth}`SessionConfig.with_extension `. -`FFI_ExtensionOptions` carries no version field, so it is one of the three -components that cannot be version-checked on import. - -```rust -fn __datafusion_extension_options__<'py>( - &self, - py: Python<'py>, -) -> PyResult> { - let mut config = FFI_ExtensionOptions::default(); - config - .add_config(self) - .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; - - PyCapsule::new_with_value(py, config, cr"datafusion_extension_options") -} -``` - [`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/index.md b/docs/source/extension-guide/index.md index 62e8a0fc0..91d42a3c5 100644 --- a/docs/source/extension-guide/index.md +++ b/docs/source/extension-guide/index.md @@ -69,11 +69,6 @@ The session owns the codecs used for the exchange and supplies them to the foreign planner. That is what lets the planner decode provider-owned objects, and lets datafusion-python decode the physical plan the planner returns. -:::{note} -The example codecs use process-local tokens to demonstrate ownership. -A production codec should serialize durable metadata instead. -::: - ## Hook reference Every integration point is a dunder method named `__datafusion_*__`. The @@ -83,36 +78,45 @@ FFI-safe struct. See {doc}`capsule-protocol` for what that means and {ref}`extension_getter_argument` for the argument every getter in the middle group receives. -| Hook | Capsule name | Argument | Documented on | +Each hook's capsule name follows from its own name by +{ref}`the naming rule `, so it is not repeated +here. + +| Hook | Contributes | Argument | Documented on | | --- | --- | --- | --- | -| `__datafusion_table_provider__` | `datafusion_table_provider` | codec source | {doc}`table-providers` | -| `__datafusion_table_provider_factory__` | `datafusion_table_provider_factory` | codec source | {doc}`table-providers` | -| `__datafusion_catalog_provider__` | `datafusion_catalog_provider` | codec source | {doc}`table-providers` | -| `__datafusion_catalog_provider_list__` | `datafusion_catalog_provider_list` | codec source | {doc}`table-providers` | -| `__datafusion_schema_provider__` | `datafusion_schema_provider` | codec source | {doc}`table-providers` | -| `__datafusion_table_function__` | `datafusion_table_function` | session | {doc}`functions` | -| `__datafusion_scalar_udf__` | `datafusion_scalar_udf` | none | {doc}`functions` | -| `__datafusion_aggregate_udf__` | `datafusion_aggregate_udf` | none | {doc}`functions` | -| `__datafusion_window_udf__` | `datafusion_window_udf` | none | {doc}`functions` | -| `__datafusion_logical_extension_codec__` | `datafusion_logical_extension_codec` | session | {doc}`codecs` | -| `__datafusion_physical_extension_codec__` | `datafusion_physical_extension_codec` | session | {doc}`codecs` | -| `__datafusion_codec_id__` | *not a capsule — a string attribute* | — | {doc}`codecs` | -| `__datafusion_query_planner__` | `datafusion_query_planner` | session | {doc}`query-planners` | -| `__datafusion_session_extension__` | *not a capsule — returns components* | `ctx` | {doc}`bundles` | -| `__datafusion_session_planner__` | `datafusion_query_planner`, or `None` | `ctx`, `fallback` | {doc}`bundles` | -| `__datafusion_physical_optimizer_rule__` | `datafusion_physical_optimizer_rule` | none | {ref}`extension_other_hooks` | -| `__datafusion_extension_options__` | `datafusion_extension_options` | none | {ref}`extension_other_hooks` | -| `__datafusion_task_context_provider__` | `datafusion_task_context_provider` | none | {ref}`extension_task_context_provider` | - -Two rows are not like the others. `__datafusion_task_context_provider__` is +| `__datafusion_table_provider__` | one table | codec source | {doc}`table-providers` | +| `__datafusion_table_provider_factory__` | tables built by `CREATE EXTERNAL TABLE` | codec source | {doc}`table-providers` | +| `__datafusion_catalog_provider__` | a named set of schemas | codec source | {doc}`table-providers` | +| `__datafusion_catalog_provider_list__` | the whole catalog namespace | codec source | {doc}`table-providers` | +| `__datafusion_schema_provider__` | a named set of tables | codec source | {doc}`table-providers` | +| `__datafusion_table_function__` | a table-valued function | session | {doc}`functions` | +| `__datafusion_scalar_udf__` | a scalar function | none | {doc}`functions` | +| `__datafusion_aggregate_udf__` | an aggregate function | none | {doc}`functions` | +| `__datafusion_window_udf__` | a window function | none | {doc}`functions` | +| `__datafusion_logical_extension_codec__` | a logical codec | session | {doc}`codecs` | +| `__datafusion_physical_extension_codec__` | a physical codec | session | {doc}`codecs` | +| `__datafusion_codec_id__` | the wire id a codec's payloads carry | — | {ref}`extension_codec_ids` | +| `__datafusion_query_planner__` | a query planner | session | {doc}`query-planners` | +| `__datafusion_session_components__` | a bundle's codecs | `ctx` | {doc}`bundles` | +| `__datafusion_session_planner__` | a bundle's planner, wrapping `fallback` | `ctx`, `fallback` | {doc}`bundles` | +| `__datafusion_physical_optimizer_rule__` | a physical optimizer rule | none | {ref}`extension_other_hooks` | +| `__datafusion_extension_options__` | typed entries in the session config | none | {ref}`extension_other_hooks` | +| `__datafusion_task_context_provider__` | the host's task context | none | {ref}`extension_task_context_provider` | + +Three rows are not like the others. `__datafusion_task_context_provider__` is implemented by the **host**, not by your library — you read it off the session -you are handed. And `__datafusion_codec_id__` is a plain string attribute -rather than a method returning a capsule. +you are handed. `__datafusion_codec_id__` is a plain string attribute rather +than a method returning a capsule. And the two `session_` hooks are dispatched +from Python and return objects rather than capsules. "codec source" in the argument column means the value is something you can read the host's logical extension codec off, which is not always a session. {ref}`extension_getter_argument` explains why, and what to do with it. +`python/tests/test_docstrings.py` compares this table against the hook names +the package actually dispatches, so a hook added or renamed without touching +this page fails the suite. + ```{toctree} :maxdepth: 2 @@ -123,6 +127,7 @@ functions codecs bundles query-planners +other-components sessions checklist ``` diff --git a/docs/source/extension-guide/other-components.md b/docs/source/extension-guide/other-components.md new file mode 100644 index 000000000..919b8718e --- /dev/null +++ b/docs/source/extension-guide/other-components.md @@ -0,0 +1,72 @@ + + +(extension_other_hooks)= + +# Optimizer rules and configuration + +Two hooks contribute things that are neither data nor functions: a rewrite pass +over physical plans, and typed entries in the session config. Both take no +argument and both are implemented in [`datafusion-ffi-example`]. + +## A physical optimizer rule + +**`__datafusion_physical_optimizer_rule__`** contributes a rule that rewrites +physical plans, installed with +{py:meth}`~datafusion.SessionContext.add_physical_optimizer_rule`. Reach for +this rather than a {doc}`query planner ` when you want to +adjust the plan DataFusion produced rather than produce it yourself — it is +much the smaller commitment, and rules accumulate where planners nest. + +```rust +fn __datafusion_physical_optimizer_rule__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let rule: Arc = Arc::new(self.clone()); + let runtime = get_tokio_runtime().handle().clone(); + let ffi = FFI_PhysicalOptimizerRule::new(rule, Some(runtime)); + + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_optimizer_rule") +} +``` + +## Typed configuration + +**`__datafusion_extension_options__`** contributes typed configuration entries +that your components can read back out of the session config, installed with +{py:meth}`SessionConfig.with_extension `. +`FFI_ExtensionOptions` carries no version field, so it is one of the three +components that cannot be version-checked on import. + +```rust +fn __datafusion_extension_options__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let mut config = FFI_ExtensionOptions::default(); + config + .add_config(self) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + + PyCapsule::new_with_value(py, config, cr"datafusion_extension_options") +} +``` + +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example diff --git a/docs/source/extension-guide/query-planners.md b/docs/source/extension-guide/query-planners.md index a97214bda..e373bb552 100644 --- a/docs/source/extension-guide/query-planners.md +++ b/docs/source/extension-guide/query-planners.md @@ -90,9 +90,7 @@ So, three rules: handle that holds the new codec* — that re-runs its getter, which re-imports the fallback against that handle's codecs. Re-installing on the original handle rebinds the session's planner back to the original handle's codecs - instead, which is the trap - `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` - pins. + instead, which is the trap the query-planner example's test suite pins. - Better, use {ref}`extension_bundles`, where there is no "afterwards" within a call. diff --git a/docs/source/extension-guide/sessions.md b/docs/source/extension-guide/sessions.md index c4f0bdb01..21bfb6e35 100644 --- a/docs/source/extension-guide/sessions.md +++ b/docs/source/extension-guide/sessions.md @@ -47,8 +47,7 @@ ctx.sql(...).collect() # plans with [codec_b, codec_a, default] -- the disca ``` Chaining `ctx = ctx.with_...(...)` keeps the two in step, which is why every -example in this guide does. -`test_the_planner_and_the_handle_can_hold_different_codecs` pins the +example in this guide does. The query-planner example's test suite pins the divergence. ## What a derived context shares diff --git a/docs/source/extension-guide/table-providers.md b/docs/source/extension-guide/table-providers.md index b33f52ec3..8f5a0ee08 100644 --- a/docs/source/extension-guide/table-providers.md +++ b/docs/source/extension-guide/table-providers.md @@ -28,11 +28,16 @@ argument. Every one of them is implemented in [`datafusion-ffi-example`]. | Hook | Exposes | Registered with | | --- | --- | --- | -| `__datafusion_table_provider__` | one table | {py:meth}`~datafusion.SessionContext.register_table` | -| `__datafusion_table_provider_factory__` | a factory that builds tables from `CREATE EXTERNAL TABLE` | {py:meth}`~datafusion.SessionContext.register_table_factory` | -| `__datafusion_schema_provider__` | a named set of tables | {py:meth}`datafusion.catalog.Catalog.register_schema` | -| `__datafusion_catalog_provider__` | a named set of schemas | {py:meth}`~datafusion.SessionContext.register_catalog_provider` | -| `__datafusion_catalog_provider_list__` | the whole catalog namespace | {py:meth}`~datafusion.SessionContext.register_catalog_provider_list` | +| `__datafusion_table_provider__` | one table | {py:meth}`SessionContext.register_table ` | +| `__datafusion_table_provider_factory__` | a factory that builds tables from `CREATE EXTERNAL TABLE` | {py:meth}`SessionContext.register_table_factory ` | +| `__datafusion_schema_provider__` | a named set of tables | {py:meth}`Catalog.register_schema ` | +| `__datafusion_catalog_provider__` | a named set of schemas | {py:meth}`SessionContext.register_catalog_provider ` | +| `__datafusion_catalog_provider_list__` | the whole catalog namespace | {py:meth}`SessionContext.register_catalog_provider_list ` | + +A schema provider is the one that does not register on the session: you reach a +{py:class}`~datafusion.catalog.Catalog` first, with +{py:meth}`SessionContext.catalog `, and +register the schema on that. Start with a table provider. Reach for the schema and catalog levels when your data source has its own namespace that should be browsable rather than diff --git a/docs/source/llms.txt b/docs/source/llms.txt index 24d03ce57..47b303ff6 100644 --- a/docs/source/llms.txt +++ b/docs/source/llms.txt @@ -31,7 +31,7 @@ - [`datafusion.expr`](https://datafusion.apache.org/python/autoapi/datafusion/expr/index.html): expression tree nodes (`Expr`, `Window`, `WindowFrame`, `GroupingSet`). - [`datafusion.functions`](https://datafusion.apache.org/python/autoapi/datafusion/functions/index.html): 290+ scalar, aggregate, and window functions. - [`datafusion.context.SessionContext`](https://datafusion.apache.org/python/autoapi/datafusion/context/index.html): session entry point, data loading, SQL execution. -- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionExtensionExportable`, `SessionPlannerExportable`, `QueryPlannerExportable` — the extension-bundle protocol. +- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionComponentsExportable`, `SessionPlannerExportable`, `QueryPlannerExportable` — the extension-bundle protocol. - [`datafusion.ipc`](https://datafusion.apache.org/python/autoapi/datafusion/ipc/index.html): worker and sender context slots for shipping expressions between processes. ## Examples diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 7250ac079..f98590b0a 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -101,7 +101,7 @@ If a library ships codecs *and* a planner, prefer installs every codec before it binds any planner, so a planner cannot end up carrying a chain that a later `with_logical_extension_codec` call has grown. The library exposes a bundle object implementing -`__datafusion_session_extension__` for its codecs and +`__datafusion_session_components__` for its codecs and `__datafusion_session_planner__` for its planner — the latter is handed the planner installed so far, so several libraries that each ship one nest instead of displacing each other. See {ref}`extension_bundles`. diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index f21aee630..af128886f 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -50,7 +50,7 @@ ctx.register_table("numbers", provider) ctx.register_udf(provider_udf) ``` -`MyPlannerExtension` implements both extension hooks. `__datafusion_session_extension__` +`MyPlannerExtension` implements both extension hooks. `__datafusion_session_components__` receives the session it is being installed on, binds fresh codecs to that session's task-context provider, and returns them as `SessionExtensionComponents`. `__datafusion_session_planner__` then runs in the host's second phase, after every 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 cb85e6bf2..8984b8ac4 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -725,7 +725,7 @@ def __init__(self) -> None: self.logical_codec = MyLogicalExtensionCodec() self.physical_codec = MyPhysicalExtensionCodec() - def __datafusion_session_extension__( + def __datafusion_session_components__( self, ctx: SessionContext ) -> SessionExtensionComponents: return SessionExtensionComponents( @@ -762,7 +762,7 @@ def __init__(self, prefix: str) -> None: self.logical = _NamedCodec(self.logical_codec, f"{prefix}.logical") self.physical = _NamedCodec(self.physical_codec, f"{prefix}.physical") - def __datafusion_session_extension__( + def __datafusion_session_components__( self, ctx: SessionContext ) -> SessionExtensionComponents: return SessionExtensionComponents( @@ -833,7 +833,7 @@ def test_with_extensions_rejects_a_rust_bundles_bare_capsule(): """ class BareCapsuleExtension: - def __datafusion_session_extension__( + def __datafusion_session_components__( self, ctx: SessionContext ) -> SessionExtensionComponents: return SessionExtensionComponents( @@ -923,10 +923,10 @@ class CodecsOf: def __init__(self, inner: object) -> None: self.inner = inner - def __datafusion_session_extension__( + def __datafusion_session_components__( self, ctx: SessionContext ) -> SessionExtensionComponents: - return self.inner.__datafusion_session_extension__(ctx) + return self.inner.__datafusion_session_components__(ctx) class PlannerOf: @@ -1165,7 +1165,7 @@ def test_with_extensions_failure_leaves_source_usable(): fully functional.""" class BoomExtension: - def __datafusion_session_extension__( + def __datafusion_session_components__( self, ctx: SessionContext ) -> SessionExtensionComponents: msg = "boom" @@ -1212,7 +1212,7 @@ class NoOpExtension: out. """ - def __datafusion_session_extension__( + def __datafusion_session_components__( self, ctx: SessionContext ) -> SessionExtensionComponents: return SessionExtensionComponents() @@ -1356,11 +1356,11 @@ def __init__(self, endpoint: str) -> None: self._codecs = ProviderCodecsExtension() self._planner = MyPlannerExtension() - def __datafusion_session_extension__( + def __datafusion_session_components__( self, ctx: SessionContext ) -> SessionExtensionComponents: - codecs = self._codecs.__datafusion_session_extension__(ctx) - planner = self._planner.__datafusion_session_extension__(ctx) + codecs = self._codecs.__datafusion_session_components__(ctx) + planner = self._planner.__datafusion_session_components__(ctx) return SessionExtensionComponents( logical_extension_codecs=( *codecs.logical_extension_codecs, diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs index c9a42dd6e..542b36ea3 100644 --- a/examples/datafusion-ffi-query-planner-example/src/extension.rs +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -229,7 +229,7 @@ impl BundledPhysicalCodec { /// /// 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 +/// `__datafusion_session_components__` call creates fresh codec and planner /// components bound to the task-context provider of the context it receives. #[pyclass( from_py_object, @@ -327,7 +327,7 @@ impl MyPlannerExtension { .map(|config| config.max_rows) } - fn __datafusion_session_extension__<'py>( + fn __datafusion_session_components__<'py>( &self, py: Python<'py>, ctx: Bound<'py, PyAny>, diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index fd6f27b82..3696d92a8 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -94,8 +94,8 @@ from .expr import Expr, WindowFrame from .extensions import ( QueryPlannerExportable, + SessionComponentsExportable, SessionExtensionComponents, - SessionExtensionExportable, SessionPlannerExportable, ) from .io import read_avro, read_csv, read_json, read_parquet @@ -139,10 +139,10 @@ "RuntimeEnvBuilder", "SQLOptions", "ScalarUDF", + "SessionComponentsExportable", "SessionConfig", "SessionContext", "SessionExtensionComponents", - "SessionExtensionExportable", "SessionPlannerExportable", "Table", "TableFunction", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 96b9f11a8..bbf08e84e 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -71,8 +71,8 @@ from datafusion.expr import sort_list_to_raw_sort_list from datafusion.extensions import ( QueryPlannerExportable, + SessionComponentsExportable, SessionExtensionComponents, - SessionExtensionExportable, SessionPlannerExportable, ) from datafusion.options import ( @@ -1848,7 +1848,7 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non self.ctx.set_query_planner(planner) def with_extensions( - self, *extensions: SessionExtensionExportable | SessionPlannerExportable + self, *extensions: SessionComponentsExportable | SessionPlannerExportable ) -> SessionContext: """Create a new session context with the given extension bundles. @@ -1858,7 +1858,7 @@ def with_extensions( Each argument is called twice, in two phases: - 1. ``__datafusion_session_extension__(ctx)`` on every extension, then + 1. ``__datafusion_session_components__(ctx)`` on every extension, then all the returned codecs are installed at once. 2. ``__datafusion_session_planner__(ctx, fallback)`` on every extension, **in argument order**, each handed the planner built so @@ -1932,11 +1932,11 @@ def with_extensions( """ for extension in extensions: if not isinstance( - extension, (SessionExtensionExportable, SessionPlannerExportable) + extension, (SessionComponentsExportable, SessionPlannerExportable) ): msg = ( "Extension implements neither " - "__datafusion_session_extension__ nor " + "__datafusion_session_components__ nor " f"__datafusion_session_planner__: {extension!r}" ) raise TypeError(msg) @@ -1948,12 +1948,12 @@ def with_extensions( logical_codecs: list[LogicalExtensionCodecExportable] = [] physical_codecs: list[PhysicalExtensionCodecExportable] = [] for extension in extensions: - if not isinstance(extension, SessionExtensionExportable): + if not isinstance(extension, SessionComponentsExportable): continue - components = extension.__datafusion_session_extension__(self) + components = extension.__datafusion_session_components__(self) if not isinstance(components, SessionExtensionComponents): msg = ( - "__datafusion_session_extension__ must return " + "__datafusion_session_components__ must return " "SessionExtensionComponents, got " f"{type(components).__name__} from {extension!r}" ) diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 0cc41859e..77ae92fc2 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -20,14 +20,14 @@ 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:`SessionComponentsExportable` by returning a :py:class:`SessionExtensionComponents` describing what it contributes, and is installed with :py:meth:`~datafusion.context.SessionContext.with_extensions`:: ctx = SessionContext().with_extensions(MyLibraryExtension()) Codecs and planners install in two phases: every -:py:class:`SessionExtensionExportable` runs first and its codecs are installed, +:py:class:`SessionComponentsExportable` runs first and its codecs are installed, then every :py:class:`SessionPlannerExportable` runs in argument order. A bundle implements either hook or both. Bundle order is significant for planners, which nest, and irrelevant for codecs, which accumulate. @@ -58,8 +58,8 @@ __all__ = [ "QueryPlannerExportable", + "SessionComponentsExportable", "SessionExtensionComponents", - "SessionExtensionExportable", "SessionPlannerExportable", ] @@ -114,7 +114,7 @@ def _not_a_codec_iterable(field: str, value: object) -> str: class SessionExtensionComponents: """Components an extension contributes to a session context. - Returned by :py:meth:`SessionExtensionExportable.__datafusion_session_extension__` + Returned by :py:meth:`SessionComponentsExportable.__datafusion_session_components__` and consumed by :py:meth:`~datafusion.context.SessionContext.with_extensions`. Every component must be created against the context passed to that method; @@ -219,7 +219,7 @@ def __post_init__(self) -> None: @runtime_checkable -class SessionExtensionExportable(Protocol): +class SessionComponentsExportable(Protocol): """Type hint for extension bundles installable via ``with_extensions``. Runtime-checkable, so ``isinstance`` answers whether an object implements @@ -252,18 +252,18 @@ class SessionExtensionExportable(Protocol): Examples: >>> from datafusion import ( ... SessionExtensionComponents, - ... SessionExtensionExportable, + ... SessionComponentsExportable, ... ) >>> class MyLibraryExtension: - ... def __datafusion_session_extension__(self, ctx): + ... def __datafusion_session_components__(self, ctx): ... return SessionExtensionComponents() - >>> isinstance(MyLibraryExtension(), SessionExtensionExportable) + >>> isinstance(MyLibraryExtension(), SessionComponentsExportable) True - >>> isinstance(object(), SessionExtensionExportable) + >>> isinstance(object(), SessionComponentsExportable) False """ - def __datafusion_session_extension__( # noqa: D105 + def __datafusion_session_components__( # noqa: D105 self, ctx: SessionContext ) -> SessionExtensionComponents: ... diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index 7e2317ff6..f8d273177 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -148,7 +148,7 @@ class LogicalExtensionCodecExportable(Protocol): ``ffi_logical_codec_from_pycapsule``, which handles both. The two bundle hooks are the exception: - :py:class:`~datafusion.extensions.SessionExtensionExportable` and + :py:class:`~datafusion.extensions.SessionComponentsExportable` and :py:class:`~datafusion.extensions.SessionPlannerExportable` are dispatched from Python and receive the wrapper. See :ref:`extension_getter_argument`. diff --git a/python/tests/test_context.py b/python/tests/test_context.py index c1aa19a4e..1a15e4a54 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -914,7 +914,7 @@ def __init__(self, prefix="my_library"): self.prefix = prefix self.bound_ctx = None - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): self.bound_ctx = ctx return SessionExtensionComponents( logical_extension_codecs=( @@ -1000,7 +1000,7 @@ def test_with_extensions_rejects_non_extension(ctx): def test_with_extensions_rejects_bad_components(ctx): class BadExtension: - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): return 42 with pytest.raises(TypeError, match="SessionExtensionComponents"): @@ -1131,7 +1131,7 @@ def test_with_extensions_rejects_bad_codec_capsule(ctx): """A correctly shaped object still has to return the right capsule.""" class BadCodecExtension: - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): wrong_capsule = ctx.__datafusion_task_context_provider__() return SessionExtensionComponents( logical_extension_codecs=( @@ -1158,7 +1158,7 @@ class BareCapsuleExtension: def __init__(self): self.exporter = SessionContext() - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): return SessionExtensionComponents( logical_extension_codecs=( self.exporter.__datafusion_logical_extension_codec__(), @@ -1179,7 +1179,7 @@ class BareCapsuleExtension: def __init__(self): self.exporter = SessionContext() - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): return SessionExtensionComponents( physical_extension_codecs=( self.exporter.__datafusion_physical_extension_codec__(), @@ -1209,8 +1209,8 @@ class ComposedExtension: def __init__(self, inner): self.inner = inner - def __datafusion_session_extension__(self, ctx): - return self.inner.__datafusion_session_extension__(ctx) + def __datafusion_session_components__(self, ctx): + return self.inner.__datafusion_session_components__(ctx) direct = ctx.with_extensions(_CodecOnlyExtension()) wrapped = SessionContext().with_extensions(ComposedExtension(_CodecOnlyExtension())) @@ -1232,7 +1232,7 @@ class TwoNamedCodecs: def __init__(self): self.exporter = SessionContext() - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): return SessionExtensionComponents( logical_extension_codecs=( _NamedCodec( @@ -1272,7 +1272,7 @@ class TwoUnnamedCodecs: def __init__(self): self.exporter = SessionContext() - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): capsule = self.exporter.__datafusion_logical_extension_codec__ return SessionExtensionComponents( logical_extension_codecs=( @@ -1290,7 +1290,7 @@ def test_with_extensions_leaves_an_exporting_object_its_own_id(ctx): exporter = SessionContext() class ObjectCodecExtension: - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): return SessionExtensionComponents(logical_extension_codecs=(exporter,)) result = ctx.with_extensions(ObjectCodecExtension()) @@ -1341,7 +1341,7 @@ def test_with_extensions_survives_source_collection(): def test_with_extensions_failure_leaves_source_usable(ctx): class BoomExtension: - def __datafusion_session_extension__(self, ctx): + def __datafusion_session_components__(self, ctx): msg = "boom" raise RuntimeError(msg) diff --git a/python/tests/test_docstrings.py b/python/tests/test_docstrings.py index 484698592..2fbb4d699 100644 --- a/python/tests/test_docstrings.py +++ b/python/tests/test_docstrings.py @@ -39,6 +39,14 @@ from datafusion import SessionContext, extensions PACKAGE_ROOT = pathlib.Path(datafusion.__file__).parent +REPO_ROOT = PACKAGE_ROOT.parents[1] + +# The hook reference in the extension guide is a hand-written table of every +# `__datafusion_*__` name. Nothing about adding a hook forces it to be updated, +# so the table is compared against the names the package actually dispatches. +HOOK_REFERENCE = REPO_ROOT / "docs" / "source" / "extension-guide" / "index.md" +HOOK_NAME = re.compile(r"__datafusion_[a-z_]+__") +HOOK_TABLE_ROW = re.compile(r"^\| `(__datafusion_[a-z_]+__)`") # Above this, a docstring has stopped being a contract and become a # narrative. Move the argument to a guide page under `docs/source/` and leave @@ -131,6 +139,38 @@ def test_extension_api_has_a_doctest(name: str, obj: object) -> None: ) +def test_hook_reference_table_lists_every_hook() -> None: + """The guide's hook table matches the hooks the package dispatches.""" + if not HOOK_REFERENCE.is_file(): + pytest.skip("running against an installed wheel, without docs/ or crates/") + + dispatched: set[str] = set() + for directory, suffix in ( + (REPO_ROOT / "crates", "*.rs"), + (PACKAGE_ROOT, "*.py"), + ): + for path in sorted(directory.rglob(suffix)): + dispatched.update(HOOK_NAME.findall(path.read_text())) + + documented = { + match.group(1) + for line in HOOK_REFERENCE.read_text().splitlines() + if (match := HOOK_TABLE_ROW.match(line)) + } + + problems = [ + f" dispatched but not in the table: {name}" + for name in sorted(dispatched - documented) + ] + [ + f" in the table but nowhere in the source: {name}" + for name in sorted(documented - dispatched) + ] + assert not problems, ( + f"{HOOK_REFERENCE.relative_to(REPO_ROOT)} is out of sync with the " + "hooks in crates/ and python/datafusion/:\n" + "\n".join(problems) + ) + + def test_no_dead_pointers_at_the_guide() -> None: """A docstring naming the guide must link it, not just mention it.""" role = re.compile(r":(ref|doc|py:\w+):`") diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py index 9764e2973..0e0a3965f 100644 --- a/python/tests/test_imports.py +++ b/python/tests/test_imports.py @@ -105,7 +105,7 @@ def test_extension_protocols_are_exported_together(): for name in [ "QueryPlannerExportable", "SessionExtensionComponents", - "SessionExtensionExportable", + "SessionComponentsExportable", "SessionPlannerExportable", ]: assert name in datafusion.__all__, f"{name} missing from datafusion.__all__" From 9372e1bf3bae0892628e3daa18d6e261bae60467 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 11:03:31 -0400 Subject: [PATCH 33/33] test: match the hook table against dispatch sites, not raw text `test_hook_reference_table_lists_every_hook` scanned every byte of `crates/**.rs` and `python/datafusion/**.py` for `__datafusion_*__`, so comments and docstrings counted as evidence a hook exists. A doc-comment contrasting a hook with one that was removed, or naming a hypothetical, would have had to be deleted or added to the guide's table, and neither is right. Count sites instead. On the Rust side a site is a string literal holding nothing but the hook name -- what `hasattr`, `getattr`, and `call_capsule_getter` are handed -- or a `fn` of that name, which is a hook the host implements itself. Error strings that merely embed a name no longer count; each already sits beside a real lookup. On the Python side, read the syntax tree rather than the text: a method being defined, an attribute being accessed, or a string standing alone. A docstring is one string node holding the whole docstring, so prose drops out without a rule of its own. The dispatched set is unchanged at 18, still matching the table exactly. Verified in both directions: a comment naming a removed hook now passes, while adding a real lookup for an undocumented name still fails. Co-Authored-By: Claude Opus 5 (1M context) --- python/tests/test_docstrings.py | 68 ++++++++++++++++++++++++++++----- 1 file changed, 58 insertions(+), 10 deletions(-) diff --git a/python/tests/test_docstrings.py b/python/tests/test_docstrings.py index 2fbb4d699..9ac1f0b76 100644 --- a/python/tests/test_docstrings.py +++ b/python/tests/test_docstrings.py @@ -44,10 +44,26 @@ # The hook reference in the extension guide is a hand-written table of every # `__datafusion_*__` name. Nothing about adding a hook forces it to be updated, # so the table is compared against the names the package actually dispatches. +# +# "Dispatches" means a *site*, not an occurrence: the name spelled where the +# host looks the hook up or defines its own, not everywhere the name is +# written. Scanning raw text would make the table's contents depend on prose — +# a doc-comment contrasting a hook with one that was removed, or naming a +# hypothetical, would have to be either deleted or added to the table, and +# neither is right. Structure answers the question text cannot. HOOK_REFERENCE = REPO_ROOT / "docs" / "source" / "extension-guide" / "index.md" -HOOK_NAME = re.compile(r"__datafusion_[a-z_]+__") +HOOK_NAME = re.compile(r"^__datafusion_[a-z_]+__$") HOOK_TABLE_ROW = re.compile(r"^\| `(__datafusion_[a-z_]+__)`") +# A Rust dispatch site is a string literal holding nothing but the hook name — +# what `hasattr`, `getattr`, and `call_capsule_getter` are handed — or a `fn` +# of that name, which is a hook the host itself implements. An error message +# that merely embeds the name (`"__datafusion_scalar_udf__ does not exist"`) +# is prose and does not count; every such message sits beside a real lookup. +RUST_HOOK_SITE = re.compile( + r'"(__datafusion_[a-z_]+__)"|\bfn\s+(__datafusion_[a-z_]+__)\b' +) + # Above this, a docstring has stopped being a contract and become a # narrative. Move the argument to a guide page under `docs/source/` and leave # a one-line pointer. Raising this is not the fix. @@ -139,18 +155,46 @@ def test_extension_api_has_a_doctest(name: str, obj: object) -> None: ) +def _rust_hook_sites() -> set[str]: + """Hook names Rust looks up or defines. See :data:`RUST_HOOK_SITE`.""" + found: set[str] = set() + for path in sorted((REPO_ROOT / "crates").rglob("*.rs")): + for match in RUST_HOOK_SITE.finditer(path.read_text()): + found.add(match.group(1) or match.group(2)) + return found + + +def _python_hook_sites() -> set[str]: + """Hook names the Python wrappers call, declare, or look up by string. + + Read off the syntax tree rather than the text, so a name is counted when + it is an attribute being accessed (``extension.__datafusion_x__(ctx)``), a + method being defined (the protocol stubs), or a string standing alone (a + ``getattr`` argument) — and not when it merely appears inside a docstring. + """ + found: set[str] = set() + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + name = node.name + elif isinstance(node, ast.Attribute): + name = node.attr + elif isinstance(node, ast.Constant) and isinstance(node.value, str): + name = node.value + else: + continue + if HOOK_NAME.match(name): + found.add(name) + return found + + def test_hook_reference_table_lists_every_hook() -> None: """The guide's hook table matches the hooks the package dispatches.""" if not HOOK_REFERENCE.is_file(): pytest.skip("running against an installed wheel, without docs/ or crates/") - dispatched: set[str] = set() - for directory, suffix in ( - (REPO_ROOT / "crates", "*.rs"), - (PACKAGE_ROOT, "*.py"), - ): - for path in sorted(directory.rglob(suffix)): - dispatched.update(HOOK_NAME.findall(path.read_text())) + dispatched = _rust_hook_sites() | _python_hook_sites() documented = { match.group(1) @@ -162,12 +206,16 @@ def test_hook_reference_table_lists_every_hook() -> None: f" dispatched but not in the table: {name}" for name in sorted(dispatched - documented) ] + [ - f" in the table but nowhere in the source: {name}" + f" in the table but dispatched from nowhere: {name}" for name in sorted(documented - dispatched) ] assert not problems, ( f"{HOOK_REFERENCE.relative_to(REPO_ROOT)} is out of sync with the " - "hooks in crates/ and python/datafusion/:\n" + "\n".join(problems) + "hooks in crates/ and python/datafusion/:\n" + + "\n".join(problems) + + "\n\nOnly dispatch sites count — a name looked up by string, an " + "attribute accessed, or a method defined. Naming a hook in prose does " + "not put it in this set, and does not belong in the table either." )