You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ti...@apache.org on 2026/08/07 19:21:13 UTC
(datafusion-python) branch feat/ffi-with-extensions created (now f1526210)
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a change to branch feat/ffi-with-extensions
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
at f1526210 Test planner rebinding and codec precedence in with_extensions
This branch includes the following new commits:
new 1c99c1ee Add atomic SessionContext.with_extensions API
new b644f974 Add extension-bundle example and with_extensions FFI tests
new 53453ea5 Document and test the context-outlives-DataFrame contract
new 665f8fe6 Skip private internal methods in wrapper coverage test
new f1526210 Test planner rebinding and codec precedence in with_extensions
The 5 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 02/05: Add extension-bundle example and with_extensions FFI tests
Posted by ti...@apache.org.
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 b644f97428ab57bf52a157103395f22bc99d4f88
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 7 14:17:45 2026 -0400
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 <no...@anthropic.com>
---
docs/source/contributor-guide/ffi.md | 62 ++++-
.../datafusion-ffi-query-planner-example/README.md | 17 +-
.../tests/_test_three_library_query_planner.py | 170 ++++++++++++-
.../src/extension.rs | 263 +++++++++++++++++++++
.../src/lib.rs | 3 +
.../src/planner.rs | 22 +-
6 files changed, 523 insertions(+), 14 deletions(-)
diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md
index 3b5b8b91..17c58826 100644
--- a/docs/source/contributor-guide/ffi.md
+++ b/docs/source/contributor-guide/ffi.md
@@ -283,8 +283,60 @@ 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.
+
+`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 +359,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 7e1ec899..55e89869 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 0991c5c9..5846661f 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 @@ from __future__ import annotations
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,163 @@ 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_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 00000000..714d3b32
--- /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<Mutex<Vec<usize>>>;
+
+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<Extension> {
+ record_task_ctx(&self.observed, ctx);
+ self.inner.try_decode(buf, inputs, ctx)
+ }
+
+ fn try_encode(&self, node: &Extension, buf: &mut Vec<u8>) -> Result<()> {
+ self.inner.try_encode(node, buf)
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ buf: &[u8],
+ table_ref: &TableReference,
+ schema: SchemaRef,
+ ctx: &TaskContext,
+ ) -> Result<Arc<dyn TableProvider>> {
+ 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<dyn TableProvider>,
+ buf: &mut Vec<u8>,
+ ) -> 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<dyn ExecutionPlan>],
+ ctx: &TaskContext,
+ proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ record_task_ctx(&self.observed, ctx);
+ self.inner.try_decode(buf, inputs, ctx, proto_converter)
+ }
+
+ fn try_encode(
+ &self,
+ node: Arc<dyn ExecutionPlan>,
+ buf: &mut Vec<u8>,
+ proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<()> {
+ self.inner.try_encode(node, buf, proto_converter)
+ }
+}
+
+fn task_ctx_provider_from_session(ctx: &Bound<'_, PyAny>) -> PyResult<FFI_TaskContextProvider> {
+ let capsule = ctx.call_method0("__datafusion_task_context_provider__")?;
+ let capsule = capsule.cast::<PyCapsule>()?;
+ let provider: NonNull<FFI_TaskContextProvider> = 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<PlannerObservations>,
+ 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<usize> {
+ 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<Bound<'py, PyAny>> {
+ // 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<dyn LogicalExtensionCodec> = 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<dyn PhysicalExtensionCodec + Send> =
+ 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<dyn QueryPlanner + Send + Sync> = 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 7635c299..ff8ce784 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::<PlannerConfig>()?;
+ m.add_class::<MyPlannerExtension>()?;
m.add_class::<MyQueryPlanner>()?;
Ok(())
}
diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs
index cb767ffa..3477f8c0 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<dyn ExecutionPlan>) -> bool {
}
fn planner_config(session: &dyn Session) -> datafusion::common::Result<PlannerConfig> {
- 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<PlannerConfig> {
// 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<PlannerCo
}
#[derive(Debug)]
-struct DistributedQueryPlanner {
- observations: Arc<PlannerObservations>,
+pub(crate) struct DistributedQueryPlanner {
+ pub(crate) observations: Arc<PlannerObservations>,
}
#[async_trait]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 05/05: Test planner rebinding and codec precedence in with_extensions
Posted by ti...@apache.org.
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 f1526210338d75ef105246e5c948e0e4a4a0f5ce
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 7 14:23:14 2026 -0400
Test planner rebinding and codec precedence in with_extensions
Co-Authored-By: Claude Fable 5 <no...@anthropic.com>
---
.../tests/_test_three_library_query_planner.py | 39 ++++++++++++++++++++++
1 file changed, 39 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 2fde3e2d..3dd73730 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
@@ -286,6 +286,45 @@ def test_with_extensions_failure_leaves_source_usable():
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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 04/05: Skip private internal methods in wrapper coverage test
Posted by ti...@apache.org.
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 665f8fe6fc4f0cbf23e731b01ac23d785aa694cc
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 7 14:20:13 2026 -0400
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 <no...@anthropic.com>
---
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 cf6719ec..b1afd683 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)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 03/05: Document and test the context-outlives-DataFrame contract
Posted by ti...@apache.org.
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 53453ea5bb7662f334bf8c2a098f3bafa6d383ca
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 7 14:19:06 2026 -0400
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 <no...@anthropic.com>
---
docs/source/contributor-guide/ffi.md | 7 +++++++
.../python/tests/_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 17c58826..a51e979e 100644
--- a/docs/source/contributor-guide/ffi.md
+++ b/docs/source/contributor-guide/ffi.md
@@ -327,6 +327,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 this protocol, including extracting the task-context provider from
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 5846661f..2fde3e2d 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
@@ -286,6 +286,25 @@ def test_with_extensions_failure_leaves_source_usable():
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(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/python/datafusion/context.py b/python/datafusion/context.py
index ce6bbd53..19b3191c 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -1880,6 +1880,11 @@ class SessionContext:
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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 01/05: Add atomic SessionContext.with_extensions API
Posted by ti...@apache.org.
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 <ti...@gmail.com>
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 <no...@anthropic.com>
---
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: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org