diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index ab64f0010..cdbc929b2 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` 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>, + physical_codecs: Vec>, + planner: Option>, + ) -> PyDataFusionResult { + 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 = (&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> = match planner { + Some(planner) => { + let planner = ffi_query_planner_from_pycapsule(&planner)?; + let planner: Arc = (&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::() + .map(|p| { + Arc::new(p.with_ffi_codecs(ffi_logical, ffi_physical)) + as Arc + }) + } + }; + + 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/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 3b5b8b91b..a51e979e6 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -283,8 +283,67 @@ planner are rebound automatically, but a planner wrapped inside another planner fallback is opaque and keeps the codecs it was exported with. **Install all extension codecs before exporting or chaining planners.** -Putting it together for a session using two extension libraries that each provide -tables, functions, and a query planner: +### Extension bundles: `with_extensions` + +Codec and planner capsules carry an `FFI_TaskContextProvider` holding a *weak* +reference to the `SessionContext` they were created against. A capsule does not keep +that context alive, and a component bound to one context cannot be rebound to another. +Chaining the low-level `with_*` methods by hand therefore risks binding components to +an intermediate context that is later garbage collected, which fails at query time +with `TaskContextProvider went out of scope over FFI boundary` — or worse, silently +reads stale session state. + +`SessionContext.with_extensions` avoids this by construction. An extension library +exposes a bundle object implementing the `__datafusion_session_extension__` protocol: + +```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 +codecs, binds the planner against the final codec chains, and returns the 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 prepend to the codec chain, so codecs from +later extensions are consulted first. At most one extension may supply a query +planner. If any factory fails, the source context's state is unchanged. + +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. + +`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust +implementation of this protocol, including extracting the task-context provider from +the supplied context and constructing a Python `SessionExtensionComponents`. + +### Advanced: chaining the low-level methods + +The `with_logical_extension_codec`, `with_physical_extension_codec`, and +`with_query_planner` methods remain available for advanced use. Putting them together +for a session using two extension libraries that each provide tables, functions, and +a query planner: ```python ctx = SessionContext(config) @@ -307,6 +366,12 @@ ctx.register_table("t", lib_a.TableProvider()) ctx.register_udf(udf(lib_b.SomeUDF())) ``` +When chaining by hand, keep the final context assigned to `ctx` as the single owner: +components created against earlier intermediate contexts (for example a codec +constructed with a context that is later discarded) hold weak references that break +once that intermediate context is collected. Prefer `with_extensions` whenever the +extension library provides a bundle. + ## 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/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 7e1ec899a..55e898693 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(PlannerConfig(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(PlannerConfig(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 0991c5c9e..3dd737301 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 @@ -20,14 +20,23 @@ import gc import pytest -from datafusion import SessionConfig, SessionContext, udf +from datafusion import ( + SessionConfig, + SessionContext, + SessionExtensionComponents, + udf, +) from datafusion_ffi_example import ( IsNullUDF, MyLogicalExtensionCodec, MyPhysicalExtensionCodec, MyTableProvider, ) -from datafusion_ffi_query_planner_example import MyQueryPlanner, PlannerConfig +from datafusion_ffi_query_planner_example import ( + MyPlannerExtension, + MyQueryPlanner, + PlannerConfig, +) def configured_context(max_rows: int): @@ -120,6 +129,221 @@ 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(PlannerConfig(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(PlannerConfig(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 + + # Codec decode calls resolve the task context through the weak provider + # bound during with_extensions. Seeing 2 (never 5) proves the provider + # targets the returned context rather than the source. + seen = planner_ext.decode_max_rows_seen() + assert seen, "expected the bundle codecs to observe at least one decode" + assert set(seen) == {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(PlannerConfig(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(PlannerConfig(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 + seen = planner_ext.decode_max_rows_seen() + assert seen + assert set(seen) == {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(PlannerConfig(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(PlannerConfig(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(PlannerConfig(max_rows=2)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + + with pytest.raises(RuntimeError, match="boom"): + source.with_extensions(MyPlannerExtension(), BoomExtension()) + + # No planner was installed, so the default planner runs unrestricted. + batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] + + +def test_with_extensions_rebinds_existing_planner(): + """Codec-only bundles installed on a context that already has an FFI + planner rebind that planner to the new codec chains.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + planner = MyQueryPlanner() + ctx = SessionContext(config).with_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_codec_precedence(): + """Extensions are processed left to right and prepend to the codec + chain, so codecs from later extensions are consulted first. Both + bundles' codecs can handle the payload (they share the provider + library's token registry); only the one consulted first is used.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + ext_a = ProviderCodecsExtension() + ext_b = ProviderCodecsExtension() + 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] + + assert ext_b.logical_codec.table_provider_encode_calls() > 0 + assert ext_b.logical_codec.table_provider_decode_calls() > 0 + assert ext_a.logical_codec.table_provider_encode_calls() == 0 + assert ext_a.logical_codec.table_provider_decode_calls() == 0 + + +def test_dataframe_outliving_context_fails_cleanly(): + """A DataFrame does not keep its SessionContext alive. FFI components + resolve the task context through a weak reference, so using the + DataFrame after dropping the context raises a clean error instead of + crashing. This locks in the documented ownership contract: the context + must outlive DataFrames that depend on FFI codecs.""" + config = SessionConfig().with_extension(PlannerConfig(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. The extra codecs 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..714d3b324 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -0,0 +1,263 @@ +// 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::ptr::NonNull; +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::get_tokio_runtime; +use datafusion_session::QueryPlanner; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, 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. Decoding reads the provider's current +/// session state, so these prove which context the provider targets. +type ObservedMaxRows = 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) + } +} + +fn task_ctx_provider_from_session(ctx: &Bound<'_, PyAny>) -> PyResult { + let capsule = ctx.call_method0("__datafusion_task_context_provider__")?; + let capsule = capsule.cast::()?; + let provider: NonNull = capsule + .pointer_checked(Some(c"datafusion_task_context_provider"))? + .cast(); + // The FFI provider holds a weak reference; cloning it does not keep the + // session context alive. The host's returned context is the strong owner. + Ok(unsafe { provider.as_ref() }.clone()) +} + +/// 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(Debug, Default, Clone)] +pub(crate) struct MyPlannerExtension { + observations: Arc, + observed_max_rows: ObservedMaxRows, +} + +#[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. + fn decode_max_rows_seen(&self) -> Vec { + self.observed_max_rows + .lock() + .map(|observed| observed.clone()) + .unwrap_or_default() + } + + 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. + let provider = task_ctx_provider_from_session(&ctx)?; + 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 = + PyCapsule::new_with_value(py, ffi_logical, cr"datafusion_logical_extension_codec")?; + + let physical: Arc = + Arc::new(ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_physical = + FFI_PhysicalExtensionCodec::new(physical, Some(runtime.clone()), provider.clone()); + let physical_capsule = + PyCapsule::new_with_value(py, ffi_physical, cr"datafusion_physical_extension_codec")?; + + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + }); + let ffi_planner = FFI_QueryPlanner::new( + planner, + Some(runtime), + provider, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ); + let planner_capsule = + PyCapsule::new_with_value(py, ffi_planner, cr"datafusion_query_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 7635c2992..ff8ce7848 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::PlannerConfig; +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 cb767ffa5..3477f8c0f 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -41,12 +41,12 @@ use pyo3::types::PyCapsule; use crate::config::PlannerConfig; #[derive(Debug, 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, } fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool { @@ -70,8 +70,12 @@ fn physical_plan_has_foreign_plan(plan: &Arc) -> bool { } 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 { // Read the flattened entry first. Some DataFusion revisions add an extra // `datafusion_ffi` namespace while reconstructing a ForeignSession. Parsing // it directly also ensures malformed values are reported instead of being @@ -107,8 +111,8 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result, +pub(crate) struct DistributedQueryPlanner { + pub(crate) observations: Arc, } #[async_trait] 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 b4214fdd5..19b3191c0 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: @@ -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,108 @@ def with_query_planner( 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. + + 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 + :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 6e6eaadbe..207234851 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, @@ -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]]) 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)