This is an automated email from the ASF dual-hosted git repository. timsaucer pushed a commit to branch feat/ffi-with-extensions in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit 1c99c1ee01dda117c5b7b97506fa3a16ed7bab82 Author: Tim Saucer <[email protected]> AuthorDate: Fri Aug 7 14:09:16 2026 -0400 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 <[email protected]> --- crates/core/src/context.rs | 84 +++++++++++++++++++++++++ python/datafusion/__init__.py | 4 ++ python/datafusion/context.py | 141 ++++++++++++++++++++++++++++++++++++++++++ python/tests/test_context.py | 115 ++++++++++++++++++++++++++++++++++ 4 files changed, 344 insertions(+) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index ab64f001..cdbc929b 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1534,6 +1534,90 @@ impl PySessionContext { physical_codec, } } + + /// Create the destination context for a `with_extensions` transaction. + /// + /// Private support method for `SessionContext.with_extensions`. The + /// returned context is the single `Arc<SessionContext>` 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>( + &self, + logical_codecs: Vec<Bound<'py, PyAny>>, + physical_codecs: Vec<Bound<'py, PyAny>>, + planner: Option<Bound<'py, PyAny>>, + ) -> PyDataFusionResult<Self> { + let mut logical_codec = self.logical_codec.as_ref().clone(); + for codec in logical_codecs { + let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?; + let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into(); + logical_codec = logical_codec.with_additional_codec(inner); + } + let logical_codec = Arc::new(logical_codec); + + let mut physical_codec = self.physical_codec.as_ref().clone(); + for codec in physical_codecs { + let inner = physical_codec_from_pycapsule(&codec)?; + physical_codec = physical_codec.with_additional_codec(inner); + } + let physical_codec = Arc::new(physical_codec); + + // Bind the planner only after the codec chains are final. Both FFI + // codec wrappers target this exact context so their weak task-context + // providers stay valid for as long as the returned context lives. + let ffi_logical = Self::ffi_logical_codec_for(&self.ctx, &logical_codec); + let ffi_physical = Self::ffi_physical_codec_for(&self.ctx, &physical_codec); + let query_planner: Option<Arc<dyn QueryPlanner + Send + Sync>> = match planner { + Some(planner) => { + let planner = ffi_query_planner_from_pycapsule(&planner)?; + let planner: Arc<dyn QueryPlanner + Send + Sync> = (&planner).into(); + let planner = + FFI_QueryPlanner::new_with_ffi_codecs(planner, ffi_logical, ffi_physical); + Some(Arc::new(RuntimeAwareQueryPlanner { planner })) + } + None => { + let state = self.ctx.state(); + let planner_any: &dyn std::any::Any = state.query_planner().as_ref(); + planner_any + .downcast_ref::<RuntimeAwareQueryPlanner>() + .map(|p| { + Arc::new(p.with_ffi_codecs(ffi_logical, ffi_physical)) + as Arc<dyn QueryPlanner + Send + Sync> + }) + } + }; + + if let Some(query_planner) = query_planner { + let state = SessionStateBuilder::new_from_existing(self.ctx.state()) + .with_query_planner(query_planner) + .build(); + *self.ctx.state_ref().write() = state; + } + + Ok(Self { + ctx: Arc::clone(&self.ctx), + logical_codec, + physical_codec, + }) + } } impl PySessionContext { diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9c55f446..86d0054b 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -80,6 +80,8 @@ from .context import ( RuntimeEnvBuilder, SessionConfig, SessionContext, + SessionExtensionComponents, + SessionExtensionExportable, SQLOptions, ) from .dataframe import ( @@ -134,6 +136,8 @@ __all__ = [ "ScalarUDF", "SessionConfig", "SessionContext", + "SessionExtensionComponents", + "SessionExtensionExportable", "Table", "TableFunction", "TableProviderFactory", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index b4214fdd..ce6bbd53 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -46,6 +46,7 @@ from __future__ import annotations import uuid import warnings +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol try: @@ -155,6 +156,49 @@ class QueryPlannerExportable(Protocol): def __datafusion_query_planner__(self) -> 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.""" @@ -1808,6 +1852,103 @@ class SessionContext: new.ctx = new_internal return new + 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:`with_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 prepend to the codec chain, so codecs + from later extensions are consulted first. 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 or more than one extension + supplies a query planner. + + 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 6e6eaadb..20723485 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 @@ from datafusion import ( RuntimeEnvBuilder, SessionConfig, SessionContext, + SessionExtensionComponents, SQLOptions, Table, column, @@ -754,6 +756,119 @@ def test_with_query_planner_capsule(ctx): assert batches[0].column(0) == pa.array([1]) +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="incorrect name"): + 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]]) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
