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/28 18:08:42 UTC
(datafusion-python) branch feat/ffi-query-planner-core updated (7652a336 -> d3ce2b20)
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a change to branch feat/ffi-query-planner-core
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
from 7652a336 Update rev for upstream datafusion to pre-release of 55.1.0
new 1153af4b refactor!: install the query planner in place instead of forking
new 06c9b9b8 test: install codecs in the provider-survival test
new d3ce2b20 fix: keep the UDF inlining setting when installing a codec
The 3 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.
Summary of changes:
.ai/skills/ffi-capsule-protocol/SKILL.md | 59 +++--
crates/core/src/codec.rs | 78 ++-----
crates/core/src/context.rs | 250 ++++++++-------------
docs/source/contributor-guide/ffi.md | 98 ++++----
docs/source/user-guide/upgrade-guides.md | 7 +-
examples/datafusion-ffi-example/README.md | 4 +-
.../datafusion-ffi-query-planner-example/README.md | 4 +-
.../tests/_test_three_library_query_planner.py | 204 +++++++++--------
.../datafusion-ffi-query-planner-example/uv.lock | 7 +
python/datafusion/context.py | 70 +++---
python/tests/test_context.py | 75 ++++---
python/tests/test_pickle_expr.py | 39 ++++
python/tests/test_plans.py | 43 ++++
13 files changed, 482 insertions(+), 456 deletions(-)
create mode 100644 examples/datafusion-ffi-query-planner-example/uv.lock
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 03/03: fix: keep the UDF inlining setting when installing a codec
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-query-planner-core
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit d3ce2b20fb37cff5bfa89f09b4cf150de51ee68e
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 28 13:48:06 2026 -0400
fix: keep the UDF inlining setting when installing a codec
`with_logical_extension_codec` and `with_physical_extension_codec` built
the replacement wrapper with `Python{Logical,Physical}Codec::new`, whose
constructor defaults `python_udf_inlining` to true. Installing a codec on
a context that had opted out therefore turned inlining back on without
saying so:
ctx = SessionContext().with_python_udf_inlining(enabled=False)
ctx = ctx.with_logical_extension_codec(codec) # inlining silently back on
That matters beyond a stale flag. Inlining is what embeds a cloudpickled
callable in the wire format, and it is opt-out precisely because that is
not portable across interpreters and not something every deployment wants
to ship. A codec install is not a request to change it.
Carry the receiver's setting across instead. Both new tests fail on the
prior build with `DFPYUDF` reappearing in the blob, and the paired
`..._preserves_inlining_when_enabled` case pins the default-on direction so
the fix cannot degenerate into hard-coding it off.
The physical case is covered in `test_plans.py` rather than alongside the
logical one: `Expr.to_bytes` only routes through the logical codec, so an
assertion there would pass with the physical bug still present. It takes an
`ExecutionPlan` to observe.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---
crates/core/src/context.rs | 15 ++++++++++++--
python/tests/test_pickle_expr.py | 39 ++++++++++++++++++++++++++++++++++++
python/tests/test_plans.py | 43 ++++++++++++++++++++++++++++++++++++++++
3 files changed, 95 insertions(+), 2 deletions(-)
diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs
index da8e08a2..c60d9ab2 100644
--- a/crates/core/src/context.rs
+++ b/crates/core/src/context.rs
@@ -1445,9 +1445,16 @@ impl PySessionContext {
) -> PyDataFusionResult<Self> {
let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
let inner: Arc<dyn LogicalExtensionCodec> = (&inner_ffi).into();
- let logical_codec = Arc::new(PythonLogicalCodec::new(inner));
let this = slf.borrow();
+ // Carry the receiver's inlining setting over. `PythonLogicalCodec::new`
+ // defaults it to on, so building the replacement without this would
+ // silently re-enable inline Python UDF encoding on a context that had
+ // opted out with `with_python_udf_inlining(enabled=False)`.
+ let logical_codec = Arc::new(
+ PythonLogicalCodec::new(inner)
+ .with_python_udf_inlining(this.logical_codec.python_udf_inlining()),
+ );
let derived = Self {
ctx: Arc::clone(&this.ctx),
logical_codec,
@@ -1476,9 +1483,13 @@ impl PySessionContext {
) -> PyDataFusionResult<Self> {
let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?;
let inner: Arc<dyn PhysicalExtensionCodec> = (&inner_ffi).into();
- let physical_codec = Arc::new(PythonPhysicalCodec::new(inner));
let this = slf.borrow();
+ // See `with_logical_extension_codec` for why the flag is carried over.
+ let physical_codec = Arc::new(
+ PythonPhysicalCodec::new(inner)
+ .with_python_udf_inlining(this.physical_codec.python_udf_inlining()),
+ );
let derived = Self {
ctx: Arc::clone(&this.ctx),
logical_codec: Arc::clone(&this.logical_codec),
diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py
index 588caa21..451f5d21 100644
--- a/python/tests/test_pickle_expr.py
+++ b/python/tests/test_pickle_expr.py
@@ -387,6 +387,45 @@ class TestPythonUdfInliningToggle:
decoded = Expr.from_bytes(blob_toggled, ctx=SessionContext())
assert "double" in decoded.canonical_name()
+ def test_installing_a_codec_preserves_strict_mode(self):
+ """Installing an extension codec must not re-enable inlining.
+
+ `with_{logical,physical}_extension_codec` builds a replacement
+ `Python{Logical,Physical}Codec` around the imported one, and the
+ constructor defaults inlining to on. A context that opted out has to
+ keep its setting, or a codec install silently starts shipping
+ cloudpickled callables again.
+
+ The session's own capsule getters are a convenient stand-in for a real
+ extension codec here: they return a genuine FFI codec, so the import
+ path under test is the same one an extension library exercises.
+
+ This covers the logical codec. The physical one is checked in
+ ``test_plans.py``, since it takes an ``ExecutionPlan`` to observe.
+ """
+ strict = SessionContext().with_python_udf_inlining(enabled=False)
+ e = self._build_double_udf()(col("a"))
+ assert b"DFPYUDF" not in e.to_bytes(strict)
+
+ with_logical = strict.with_logical_extension_codec(
+ strict.__datafusion_logical_extension_codec__()
+ )
+ assert b"DFPYUDF" not in e.to_bytes(with_logical)
+
+ def test_installing_a_codec_preserves_inlining_when_enabled(self):
+ """The converse: the default stays on across a codec install.
+
+ Guards against 'fixing' the above by hard-coding inlining off.
+ """
+ ctx = SessionContext()
+ e = self._build_double_udf()(col("a"))
+ assert b"DFPYUDF" in e.to_bytes(ctx)
+
+ installed = ctx.with_logical_extension_codec(
+ ctx.__datafusion_logical_extension_codec__()
+ )
+ assert b"DFPYUDF" in e.to_bytes(installed)
+
def test_strict_roundtrip_via_registry(self):
"""When both sender and receiver disable inlining, the UDF
travels by name only and the receiver resolves it from its
diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py
index 11e709f6..0145d123 100644
--- a/python/tests/test_plans.py
+++ b/python/tests/test_plans.py
@@ -17,6 +17,7 @@
import datetime
+import pyarrow as pa
import pytest
from datafusion import (
ExecutionPlan,
@@ -24,6 +25,8 @@ from datafusion import (
Metric,
MetricsSet,
SessionContext,
+ col,
+ udf,
)
@@ -92,6 +95,46 @@ def test_session_with_logical_extension_codec_roundtrip(ctx, df) -> None:
assert df.collect() == df_round_trip.collect()
+def test_installing_a_physical_codec_preserves_strict_mode() -> None:
+ """Installing a physical extension codec must not re-enable inlining.
+
+ `with_physical_extension_codec` builds a replacement `PythonPhysicalCodec`
+ around the imported one, and the constructor defaults inlining to on. A
+ context that opted out via `with_python_udf_inlining(enabled=False)` has to
+ keep its setting, or installing a codec silently starts embedding
+ cloudpickled callables in serialized execution plans.
+
+ The logical counterpart lives in `test_pickle_expr.py`; this one needs an
+ `ExecutionPlan` because only the physical codec encodes it. `DFPYUDF` is
+ the scalar Python-UDF family prefix, shared by both layers; see
+ `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs.
+ """
+ identity = udf(
+ lambda arr: arr,
+ [pa.string()],
+ pa.string(),
+ volatility="immutable",
+ name="identity_str",
+ )
+
+ def plan_bytes(ctx: SessionContext) -> bytes:
+ df = ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select(
+ identity(col("c1"))
+ )
+ return df.execution_plan().to_bytes(ctx)
+
+ # The inlining default is on, so the strict blob is what has to differ.
+ assert b"DFPYUDF" in plan_bytes(SessionContext())
+
+ strict = SessionContext().with_python_udf_inlining(enabled=False)
+ assert b"DFPYUDF" not in plan_bytes(strict)
+
+ installed = strict.with_physical_extension_codec(
+ strict.__datafusion_physical_extension_codec__()
+ )
+ assert b"DFPYUDF" not in plan_bytes(installed)
+
+
def test_session_codec_capsule_getters(ctx) -> None:
"""SessionContext exposes both logical and physical codec capsules."""
logical = ctx.ctx.__datafusion_logical_extension_codec__()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 01/03: refactor!: install the query planner in place instead of forking
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-query-planner-core
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit 1153af4b81752fad336db69359f9eee60d97e711
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 28 13:42:30 2026 -0400
refactor!: install the query planner in place instead of forking
`with_query_planner` derived a new `SessionContext` to install a planner,
on the grounds that the receiver must not be modified. That mints a fresh
`Arc<SessionContext>` allocation, and every FFI handle in play is bound to
an allocation rather than to the logical session: `FFI_TaskContextProvider`
holds its provider weakly, and a registered catalog provider upgrades that
handle on every `supports_filters_pushdown` and every `scan`.
So the natural `ctx = ctx.with_query_planner(planner)` dropped the session
a foreign catalog had been registered on, and the next query failed with
`TaskContextProvider went out of scope over FFI boundary`. Reproduced with
a `MyCatalogProvider` registered before the install and a `WHERE` clause to
force pushdown during logical optimization.
Rebinding cannot cover this. It reaches the codecs `PySessionContext` holds
in its own fields; a codec embedded in a registered `FFI_CatalogProvider` —
and in every `FFI_SchemaProvider` and `FFI_TableProvider` minted from it —
has no Python-side handle. Nor can a codec retain the session that built
it: codecs are routinely handed to a provider that is registered straight
back into that session, closing `SessionContext -> catalog -> FFI provider
-> FFI codec -> SessionContext`.
Install in place instead, writing `SessionState` back through `state_ref()`
exactly as `add_physical_optimizer_rule` already did. A session keeps one
`Arc<SessionContext>` for life, so no handle is ever orphaned and the bug
cannot occur. This deletes the fork and everything that existed to repair
it: `ancestors`, `rebound_{logical,physical}_codec`, `exported_session` on
both codecs, and the `exported_ffi_*` builders.
The query planner lives in `SessionState`, so it belongs to the session
rather than to a handle on it. `with_query_planner(planner) ->
SessionContext` therefore becomes `set_query_planner(planner) -> None`,
matching `add_physical_optimizer_rule`.
The 55.1.0 pin may no longer be needed — its stated reason in Cargo.toml is
the rebinding this removes — but that is left alone pending a check of the
rest of the PR.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---
.ai/skills/ffi-capsule-protocol/SKILL.md | 59 ++++--
crates/core/src/codec.rs | 78 ++-----
crates/core/src/context.rs | 235 +++++++--------------
docs/source/contributor-guide/ffi.md | 98 ++++-----
docs/source/user-guide/upgrade-guides.md | 7 +-
examples/datafusion-ffi-example/README.md | 4 +-
.../datafusion-ffi-query-planner-example/README.md | 4 +-
.../tests/_test_three_library_query_planner.py | 208 ++++++++++--------
python/datafusion/context.py | 70 +++---
python/tests/test_context.py | 75 ++++---
10 files changed, 387 insertions(+), 451 deletions(-)
diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md
index fc20b2b8..1cbb934d 100644
--- a/.ai/skills/ffi-capsule-protocol/SKILL.md
+++ b/.ai/skills/ffi-capsule-protocol/SKILL.md
@@ -116,22 +116,49 @@ one of them, and the failure is a bare `TypeError` from a `call1`. So:
`python/datafusion/user_defined.py`, where the `Protocol` type hints for
these methods live.
-## Rule 6 — a fork must rebind the codecs it carries
-
-Installing a foreign query planner **forks** the session
-(`PySessionContext::derived_parts`), because installing one writes to
-`SessionState` and the receiver must not be modified. A foreign codec holds an
-`FFI_TaskContextProvider` pointing at the session it was installed on, so the
-fork rebinds each one onto itself via
-`PySessionContext::rebound_{logical,physical}_codec`. Skip that and decode
-callbacks answer from the pre-fork registry.
-
-Rebinding relies on `FFI_{Logical,Physical}ExtensionCodec::new` adopting the
-provider on the already-foreign path, which needs DataFusion 55.1.0 or newer
-(apache/datafusion#24722). It clones the handle before overwriting, so the
-receiver keeps its own binding — assert that, not just that the fork works.
-A codec this library owns round-trips unchanged, so the rebind is safe to apply
-unconditionally.
+## Rule 6 — a session keeps one `Arc<SessionContext>` for life
+
+`FFI_TaskContextProvider` holds its provider **weakly**, and every codec handed
+to a foreign object carries one. A registered catalog provider upgrades that
+handle on every `supports_filters_pushdown` and every `scan`. The handle is
+bound to an `Arc<SessionContext>` *allocation*, so anything that replaces the
+allocation orphans every handle bound to the old one:
+`TaskContextProvider went out of scope over FFI boundary`.
+
+So mutate `SessionState` in place — `*self.ctx.state_ref().write() = ...`, the
+way `add_physical_optimizer_rule` and `set_session_query_planner` both do —
+rather than deriving a replacement `SessionContext`. Carry the session id
+across the rewrite; `SessionStateBuilder::new_from_existing` drops it and
+`build` mints a fresh one, which desyncs `session_id()` from every
+`TaskContext` the session hands out.
+
+Do not try to repair it after the fact:
+
+- **You cannot rebind what you cannot reach.** A codec embedded in a registered
+ `FFI_CatalogProvider`, and in every `FFI_SchemaProvider` and
+ `FFI_TableProvider` minted from it, has no Python-side handle.
+- **A codec must not retain its session.** Codecs are routinely handed to a
+ provider that is registered straight back into the session that built them,
+ closing `SessionContext -> catalog -> FFI provider -> FFI codec ->
+ SessionContext`.
+
+`test_registered_providers_survive_a_planner_install` in
+`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`
+guards this. Its `WHERE` clause is load-bearing: filter pushdown upgrades the
+weak handle during logical optimization, before plan serialization could fail
+first for an unrelated reason.
+
+`SessionContext.enable_url_table` is the one method that mints a second
+allocation for a session. Its result must not outlive the receiver.
+
+## Rule 7 — installing a planner mutates the session, and says so
+
+`set_query_planner` returns `None`, matching `add_physical_optimizer_rule`. The
+query planner lives in `SessionState`, so it belongs to the session and not to
+a handle on it; every context sharing that session plans through it. Do not
+reintroduce a `with_query_planner` that pretends otherwise — the only way to
+give a handle its own planner is a fresh `Arc<SessionContext>`, which is what
+Rule 6 forbids.
## Where the truth is
diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs
index 1c665893..94942a2d 100644
--- a/crates/core/src/codec.rs
+++ b/crates/core/src/codec.rs
@@ -103,7 +103,6 @@ use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
use datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
use datafusion::physical_plan::ExecutionPlan;
-use datafusion::prelude::SessionContext;
use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec};
use datafusion_proto::physical_plan::{
DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension,
@@ -234,31 +233,18 @@ fn strip_wire_header<'a>(
/// Sitting at the top of the session's logical codec stack means
/// every serializer that reads `session.logical_codec()` automatically
/// picks up Python-aware encoding for free.
+///
+/// A codec deliberately does **not** retain the session it was built from.
+/// Codecs are routinely handed to a provider that is then registered back into
+/// that same session, so retaining here would close a cycle:
+/// `SessionContext -> catalog -> FFI provider -> FFI codec -> here`. Keeping
+/// the weak `FFI_TaskContextProvider` valid is instead a matter of never
+/// replacing the session's `Arc<SessionContext>`; see
+/// `PySessionContext::set_session_query_planner`.
+#[derive(Debug)]
pub struct PythonLogicalCodec {
inner: Arc<dyn LogicalExtensionCodec>,
python_udf_inlining: bool,
- /// Keeps the exporting session alive for a codec handed across the FFI
- /// boundary.
- ///
- /// `FFI_TaskContextProvider` stores its provider in a `Weak`, so an exported
- /// codec stops working the moment the object that produced it goes out of
- /// scope. Retaining the session in the inner codec survives both that and
- /// `clone`, which clones the inner codec's `Arc` and so carries this along.
- ///
- /// Set this only on codecs that are leaving this library. A codec installed
- /// *in* a session must not hold one, or the session would own itself:
- /// `SessionContext -> SessionState -> query planner -> FFI codec -> here`.
- exported_session: Option<Arc<SessionContext>>,
-}
-
-impl std::fmt::Debug for PythonLogicalCodec {
- fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- formatter
- .debug_struct("PythonLogicalCodec")
- .field("inner", &self.inner)
- .field("python_udf_inlining", &self.python_udf_inlining)
- .finish_non_exhaustive()
- }
}
impl PythonLogicalCodec {
@@ -266,19 +252,9 @@ impl PythonLogicalCodec {
Self {
inner,
python_udf_inlining: true,
- exported_session: None,
}
}
- /// Retain `ctx` so this codec keeps working after the exporting
- /// `SessionContext` goes out of scope. Only for codecs being exported over
- /// FFI; see the `exported_session` field for why installed codecs must not
- /// use this.
- pub fn with_exported_session(mut self, ctx: Arc<SessionContext>) -> Self {
- self.exported_session = Some(ctx);
- self
- }
-
pub fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> {
&self.inner
}
@@ -475,31 +451,13 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio
/// would round-trip at the logical level but break at the physical
/// level. Both layers reuse the shared payload framing
/// ([`PY_SCALAR_UDF_FAMILY`] et al.) so the wire format is identical.
+///
+/// Like [`PythonLogicalCodec`], this does not retain the session it was built
+/// from; see that type for why.
+#[derive(Debug)]
pub struct PythonPhysicalCodec {
inner: Arc<dyn PhysicalExtensionCodec>,
python_udf_inlining: bool,
- /// Keeps the exporting session alive for a codec handed across the FFI
- /// boundary.
- ///
- /// `FFI_TaskContextProvider` stores its provider in a `Weak`, so an exported
- /// codec stops working the moment the object that produced it goes out of
- /// scope. Retaining the session in the inner codec survives both that and
- /// `clone`, which clones the inner codec's `Arc` and so carries this along.
- ///
- /// Set this only on codecs that are leaving this library. A codec installed
- /// *in* a session must not hold one, or the session would own itself:
- /// `SessionContext -> SessionState -> query planner -> FFI codec -> here`.
- exported_session: Option<Arc<SessionContext>>,
-}
-
-impl std::fmt::Debug for PythonPhysicalCodec {
- fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- formatter
- .debug_struct("PythonPhysicalCodec")
- .field("inner", &self.inner)
- .field("python_udf_inlining", &self.python_udf_inlining)
- .finish_non_exhaustive()
- }
}
impl PythonPhysicalCodec {
@@ -507,19 +465,9 @@ impl PythonPhysicalCodec {
Self {
inner,
python_udf_inlining: true,
- exported_session: None,
}
}
- /// Retain `ctx` so this codec keeps working after the exporting
- /// `SessionContext` goes out of scope. Only for codecs being exported over
- /// FFI; see the `exported_session` field for why installed codecs must not
- /// use this.
- pub fn with_exported_session(mut self, ctx: Arc<SessionContext>) -> Self {
- self.exported_session = Some(ctx);
- self
- }
-
pub fn inner(&self) -> &Arc<dyn PhysicalExtensionCodec> {
&self.inner
}
diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs
index 865ebd58..da8e08a2 100644
--- a/crates/core/src/context.rs
+++ b/crates/core/src/context.rs
@@ -409,6 +409,11 @@ impl PySessionContext {
}
pub fn enable_url_table(&self) -> PyResult<Self> {
+ // Pre-existing caveat, unrelated to query planners: this is the one
+ // method that mints a second `Arc<SessionContext>` for a session. Any
+ // weak `FFI_TaskContextProvider` handed out by the receiver stays bound
+ // to the receiver, so the returned context must not outlive it. See
+ // `set_session_query_planner` for why everything else mutates in place.
Ok(PySessionContext {
ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()),
logical_codec: Arc::clone(&self.logical_codec),
@@ -1207,9 +1212,9 @@ impl PySessionContext {
let mut guard = state_ref.write();
// Rebuilding through the builder mints a fresh session id, but this
// mutates the caller's own session rather than deriving a new one, so
- // the id has to survive. See `derived_parts` for why losing it leaves
- // `session_id()` disagreeing with every `TaskContext` the session
- // hands out.
+ // the id has to survive. See `set_session_query_planner` for why losing
+ // it leaves `session_id()` disagreeing with every `TaskContext` the
+ // session hands out.
let new_state = SessionStateBuilder::new_from_existing(guard.clone())
.with_session_id(guard.session_id().to_string())
.with_physical_optimizer_rule(rule)
@@ -1218,23 +1223,19 @@ impl PySessionContext {
Ok(())
}
- pub fn with_query_planner(
+ /// Install a foreign query planner on this session.
+ ///
+ /// A session holds exactly one planner, so this replaces whichever one is
+ /// installed. It mutates the session rather than deriving a new one, the
+ /// same way `add_physical_optimizer_rule` does; see
+ /// [`Self::set_session_query_planner`] for why that matters to FFI handles.
+ pub fn set_query_planner(
slf: &Bound<'_, Self>,
planner: Bound<'_, PyAny>,
- ) -> PyDataFusionResult<Self> {
+ ) -> PyDataFusionResult<()> {
let planner = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?;
- let this = slf.borrow();
- let (ctx, logical_codec, physical_codec) = this.derived_parts(
- Arc::clone(&this.logical_codec),
- Arc::clone(&this.physical_codec),
- Some(planner),
- );
-
- Ok(Self {
- ctx,
- logical_codec,
- physical_codec,
- })
+ slf.borrow().set_session_query_planner(Some(planner));
+ Ok(())
}
pub fn table_provider(&self, name: &str, py: Python) -> PyResult<PyTable> {
@@ -1413,7 +1414,7 @@ impl PySessionContext {
session: Option<Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyCapsule>> {
let _ = session;
- create_logical_extension_capsule(py, &self.exported_ffi_logical_codec())
+ create_logical_extension_capsule(py, self.ffi_logical_codec().as_ref())
}
/// See [`Self::__datafusion_logical_extension_codec__`] for `session`.
@@ -1432,8 +1433,8 @@ impl PySessionContext {
let planner = Arc::clone(self.ctx.state().query_planner());
let ffi = FFI_QueryPlanner::new_with_ffi_codecs(
planner,
- self.exported_ffi_logical_codec(),
- self.exported_ffi_physical_codec(),
+ Self::ffi_logical_codec_for(&self.ctx, &self.logical_codec),
+ Self::ffi_physical_codec_for(&self.ctx, &self.physical_codec),
);
create_query_planner_capsule(py, &ffi)
}
@@ -1447,13 +1448,15 @@ impl PySessionContext {
let logical_codec = Arc::new(PythonLogicalCodec::new(inner));
let this = slf.borrow();
- let (ctx, logical_codec, physical_codec) =
- this.derived_parts(logical_codec, Arc::clone(&this.physical_codec), None);
- Ok(Self {
- ctx,
+ let derived = Self {
+ ctx: Arc::clone(&this.ctx),
logical_codec,
- physical_codec,
- })
+ physical_codec: Arc::clone(&this.physical_codec),
+ };
+ // The session is shared, and an installed foreign planner holds the
+ // codecs it was built with, so it has to be rebuilt against this one.
+ derived.set_session_query_planner(None);
+ Ok(derived)
}
/// See [`Self::__datafusion_logical_extension_codec__`] for `session`.
@@ -1464,7 +1467,7 @@ impl PySessionContext {
session: Option<Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyCapsule>> {
let _ = session;
- create_physical_extension_capsule(py, &self.exported_ffi_physical_codec())
+ create_physical_extension_capsule(py, self.ffi_physical_codec().as_ref())
}
pub fn with_physical_extension_codec<'py>(
@@ -1476,13 +1479,14 @@ impl PySessionContext {
let physical_codec = Arc::new(PythonPhysicalCodec::new(inner));
let this = slf.borrow();
- let (ctx, logical_codec, physical_codec) =
- this.derived_parts(Arc::clone(&this.logical_codec), physical_codec, None);
- Ok(Self {
- ctx,
- logical_codec,
+ let derived = Self {
+ ctx: Arc::clone(&this.ctx),
+ logical_codec: Arc::clone(&this.logical_codec),
physical_codec,
- })
+ };
+ // See `with_logical_extension_codec`.
+ derived.set_session_query_planner(None);
+ Ok(derived)
}
pub fn with_python_udf_inlining(&self, enabled: bool) -> Self {
@@ -1494,128 +1498,70 @@ impl PySessionContext {
PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner()))
.with_python_udf_inlining(enabled),
);
- let (ctx, logical_codec, physical_codec) =
- self.derived_parts(logical_codec, physical_codec, None);
- Self {
- ctx,
+ let derived = Self {
+ ctx: Arc::clone(&self.ctx),
logical_codec,
physical_codec,
- }
+ };
+ // See `with_logical_extension_codec`.
+ derived.set_session_query_planner(None);
+ derived
}
}
impl PySessionContext {
- /// Return the pieces a derived `PySessionContext` should hold, binding any
- /// foreign query planner and foreign codecs to the session that will run
- /// the query.
+ /// Write the session's query planner, in place.
///
- /// Pass `Some(planner)` to install one, or `None` to rebind whichever
- /// planner the session already holds.
+ /// Pass `Some(planner)` to install one, or `None` to rebuild whichever
+ /// foreign planner is already installed against this context's current
+ /// codecs. A foreign planner carries the FFI codecs it encodes and decodes
+ /// with, so replacing a codec means handing the planner new ones. A planner
+ /// this library owns carries no codecs and needs neither.
///
- /// With no foreign planner in play there is nothing to rebind, so the
- /// existing context and codecs are shared unchanged. A foreign planner does
- /// have to be rebound, and that forks the session, because installing it
- /// writes to `SessionState` and the receiver must not be modified. The
- /// codecs are built from the fork and the fork's state is then overwritten
- /// in place, since rebuilding the context afterwards would leave them
- /// pointing at a session that is no longer used for planning.
- ///
- /// A fork is not a deep copy. `SessionState` keeps its catalog list behind
- /// an `Arc`, so catalogs and tables stay shared with the original session,
- /// while registered functions and the configuration are snapshotted at the
- /// time of the call. The session id is explicitly carried over, so the fork
- /// and its `TaskContext`s all report the id the original session had.
- fn derived_parts(
- &self,
- logical_codec: Arc<PythonLogicalCodec>,
- physical_codec: Arc<PythonPhysicalCodec>,
- planner: Option<FFI_QueryPlanner>,
- ) -> (
- Arc<SessionContext>,
- Arc<PythonLogicalCodec>,
- Arc<PythonPhysicalCodec>,
- ) {
- let state = self.ctx.state();
-
- // When the caller is only replacing codecs, recover the handle behind
- // the installed planner so it can be rebound below. A planner this
- // library owns is not a `ForeignQueryPlanner` and needs no rebinding,
- // because it does not carry codecs of its own.
+ /// This mutates `SessionState` through `state_ref()` rather than deriving a
+ /// new `SessionContext`, and that is the whole point. Every codec handed to
+ /// a foreign object carries an `FFI_TaskContextProvider`, which holds its
+ /// provider *weakly* and upgrades it on each callback — a registered
+ /// catalog provider does so on every `supports_filters_pushdown` and every
+ /// `scan`. Those handles are bound to one `Arc<SessionContext>` allocation.
+ /// Deriving a replacement would orphan every handle bound to the old one
+ /// and leave callbacks failing with `TaskContextProvider went out of scope
+ /// over FFI boundary`, and only the codecs this struct happens to hold
+ /// could be moved across; one already embedded in a registered
+ /// `FFI_CatalogProvider` is unreachable. Keeping the single allocation for
+ /// the life of the session avoids the problem rather than repairing it.
+ fn set_session_query_planner(&self, planner: Option<FFI_QueryPlanner>) {
let planner = planner.or_else(|| {
+ let state = self.ctx.state();
let installed: &dyn std::any::Any = state.query_planner().as_ref();
installed
.downcast_ref::<ForeignQueryPlanner>()
.map(|planner| planner.0.clone())
});
-
let Some(planner) = planner else {
- return (Arc::clone(&self.ctx), logical_codec, physical_codec);
+ return;
};
- let ctx = Arc::new(SessionContext::new_with_state(state));
- let logical_codec = Arc::new(Self::rebound_logical_codec(&ctx, &logical_codec));
- let physical_codec = Arc::new(Self::rebound_physical_codec(&ctx, &physical_codec));
-
let inner: Arc<dyn QueryPlanner + Send + Sync> = (&planner).into();
let planner: Arc<dyn QueryPlanner + Send + Sync> = (&FFI_QueryPlanner::new_with_ffi_codecs(
inner,
- Self::ffi_logical_codec_for(&ctx, &logical_codec),
- Self::ffi_physical_codec_for(&ctx, &physical_codec),
+ Self::ffi_logical_codec_for(&self.ctx, &self.logical_codec),
+ Self::ffi_physical_codec_for(&self.ctx, &self.physical_codec),
))
.into();
+
+ let state_ref = self.ctx.state_ref();
+ let mut guard = state_ref.write();
// `with_session_id` is load-bearing, not redundant.
// `SessionStateBuilder::new_from_existing` drops the id and `build`
// mints a fresh one, while `SessionContext` cached the original in a
- // field of its own at `new_with_state`. Without this the fork would
- // report one id from `session_id()` and another from every
- // `TaskContext` it hands out.
- let state = SessionStateBuilder::new_from_existing(ctx.state())
- .with_session_id(ctx.session_id())
+ // field of its own. Without this the session would report one id from
+ // `session_id()` and another from every `TaskContext` it hands out.
+ // Same reason `add_physical_optimizer_rule` carries it over.
+ *guard = SessionStateBuilder::new_from_existing(guard.clone())
+ .with_session_id(guard.session_id().to_string())
.with_query_planner(planner)
.build();
- *ctx.state_ref().write() = state;
-
- (ctx, logical_codec, physical_codec)
- }
-
- /// Rebind a foreign logical codec to `ctx`.
- ///
- /// A codec imported from another library holds an `FFI_TaskContextProvider`
- /// pointing at whichever session it was installed on. Its decode callbacks
- /// resolve names against that session, so a fork has to move them onto the
- /// fork or they keep answering from the pre-fork registry.
- ///
- /// `FFI_LogicalExtensionCodec::new` adopts the provider supplied here when
- /// the codec is already foreign, returning a *clone* of the handle, so the
- /// context this codec was derived from keeps its own binding. A codec this
- /// library owns has no stored provider, and round-trips back to the same
- /// `Arc` unchanged, so this is safe to apply unconditionally.
- fn rebound_logical_codec(
- ctx: &Arc<SessionContext>,
- codec: &Arc<PythonLogicalCodec>,
- ) -> PythonLogicalCodec {
- let ctx_provider = Arc::clone(ctx) as Arc<dyn TaskContextProvider>;
- let runtime = get_tokio_runtime().handle().clone();
- let rebound =
- FFI_LogicalExtensionCodec::new(Arc::clone(codec.inner()), Some(runtime), &ctx_provider);
- PythonLogicalCodec::new((&rebound).into())
- .with_python_udf_inlining(codec.python_udf_inlining())
- }
-
- /// Physical counterpart of [`Self::rebound_logical_codec`].
- fn rebound_physical_codec(
- ctx: &Arc<SessionContext>,
- codec: &Arc<PythonPhysicalCodec>,
- ) -> PythonPhysicalCodec {
- let ctx_provider = Arc::clone(ctx) as Arc<dyn TaskContextProvider>;
- let runtime = get_tokio_runtime().handle().clone();
- let rebound = FFI_PhysicalExtensionCodec::new(
- Arc::clone(codec.inner()),
- Some(runtime),
- &ctx_provider,
- );
- PythonPhysicalCodec::new((&rebound).into())
- .with_python_udf_inlining(codec.python_udf_inlining())
}
async fn _table(&self, name: &str) -> datafusion::common::Result<DataFrame> {
@@ -1704,33 +1650,12 @@ impl PySessionContext {
FFI_PhysicalExtensionCodec::new(codec, Some(runtime), &ctx_provider)
}
- /// Build an FFI-wrapped logical codec for handing out in a PyCapsule.
- ///
- /// Same as [`Self::ffi_logical_codec`] except the inner codec retains this
- /// session. The FFI task-context handle is weak, so without that a capsule
- /// stops working as soon as the exporting `SessionContext` goes out of
- /// scope, which the natural `ctx = ctx.with_query_planner(planner)` does.
- ///
- /// Only for codecs leaving this library. The plain builder is still correct
- /// for codecs attached to objects that end up back inside this session,
- /// which would otherwise make the session own itself.
- fn exported_ffi_logical_codec(&self) -> FFI_LogicalExtensionCodec {
- let codec = Arc::new(
- PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner()))
- .with_python_udf_inlining(self.logical_codec.python_udf_inlining())
- .with_exported_session(Arc::clone(&self.ctx)),
- );
- Self::ffi_logical_codec_for(&self.ctx, &codec)
- }
-
- /// Physical companion to [`Self::exported_ffi_logical_codec`].
- fn exported_ffi_physical_codec(&self) -> FFI_PhysicalExtensionCodec {
- let codec = Arc::new(
- PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner()))
- .with_python_udf_inlining(self.physical_codec.python_udf_inlining())
- .with_exported_session(Arc::clone(&self.ctx)),
- );
- Self::ffi_physical_codec_for(&self.ctx, &codec)
+ /// Build an FFI-wrapped clone of the session's physical codec on demand.
+ pub(crate) fn ffi_physical_codec(&self) -> Arc<FFI_PhysicalExtensionCodec> {
+ Arc::new(Self::ffi_physical_codec_for(
+ &self.ctx,
+ &self.physical_codec,
+ ))
}
}
diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md
index f107cfbe..78855c8f 100644
--- a/docs/source/contributor-guide/ffi.md
+++ b/docs/source/contributor-guide/ffi.md
@@ -286,7 +286,7 @@ already dropped by the time the capsule is used.
A planner uses `FFI_QueryPlanner::new_with_ffi_codecs` with the two codecs it takes off
the session, and never touches a provider directly. That also matches what installation
-does anyway: `with_query_planner` rebinds a foreign planner to the codecs of the session
+does anyway: `set_query_planner` builds the planner against the codecs of the session
that will run the query.
`SessionContext` accepts the argument on all three getters and ignores it, so a session
@@ -302,67 +302,69 @@ library resolves names against the session running the query. A function registe
by name, and the handle is live rather than a snapshot, so a registration made after the
codec is installed is visible too.
-This survives a fork. Installing a foreign query planner forks the session, and the fork
-rebinds every foreign codec it carries onto the new session, so a function registered
-after the fork is still visible to them. `FFI_LogicalExtensionCodec::new` adopts the
-provider supplied to it when the codec is already foreign, returning a clone of the
-handle, so the context the fork was derived from keeps its own binding and continues to
-resolve against its own session.
-
-That behaviour requires DataFusion 55.1.0 or newer. Before it, those constructors
-silently discarded the provider, a foreign codec could not be rebound, and a fork left
-it resolving against the pre-fork session
-([apache/datafusion#24722](https://github.com/apache/datafusion/issues/24722)).
-
-All three properties are covered in
+This is covered in
`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`,
where the example codecs take a `require_udf_on_decode` name and resolve it out of the
task context they are handed.
-### What a derived context shares
-
-`with_query_planner`, `with_logical_extension_codec`, `with_physical_extension_codec`,
-and `with_python_udf_inlining` all return a new `SessionContext` rather than mutating
-the receiver. How much the two contexts then share depends on whether a foreign query
-planner is involved.
+### One session, one `Arc<SessionContext>`
+
+Every codec handed to a foreign object carries an `FFI_TaskContextProvider`, and that
+type holds its provider **weakly**. A registered catalog provider upgrades the handle on
+every `supports_filters_pushdown` and every `scan`. Those handles are bound to one
+particular `Arc<SessionContext>` allocation, not to the logical session, so anything that
+replaces the allocation orphans all of them and the next query fails with
+`TaskContextProvider went out of scope over FFI boundary`.
+
+So a `PySessionContext` keeps the `Arc<SessionContext>` it was created with for its whole
+life. Installing a query planner writes the new `SessionState` back through
+`state_ref()`, exactly as `add_physical_optimizer_rule` does, rather than deriving a
+replacement context. The session id is carried across that rewrite — `SessionStateBuilder`
+mints a fresh one otherwise — so `session_id()` and every `TaskContext` the session hands
+out keep agreeing.
+
+Repairing the damage instead of avoiding it does not work in general. A context can
+rebuild the codecs it holds in its own fields, but a codec already embedded in a
+registered `FFI_CatalogProvider` — and in every `FFI_SchemaProvider` and
+`FFI_TableProvider` minted from it — is not reachable from Python at all. Nor can the
+codec simply retain the session that built it: a codec handed to a provider is routinely
+registered straight back into that same session, which would close the cycle
+`SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and leak it.
+
+`SessionContext.enable_url_table` is the one exception. It clones the underlying
+`SessionContext`, so the returned context has an allocation of its own and must not
+outlive the receiver.
-Without one, the derived context wraps the *same* underlying session, so a registration
-on either side is visible to both.
-
-`with_query_planner` is different, and so is any codec change made on a session that
-already has a foreign planner installed. A foreign planner holds the FFI codecs it was
-built with, so changing the codecs means rebuilding the planner against the context
-that will actually run the query. That forks the session state, and the two halves of
-the fork behave differently:
-
-- **Shared.** Catalogs, schemas, and tables. `SessionState` holds its catalog list
- behind an `Arc`, so a table registered on either context is visible to both. The
- runtime environment is shared for the same reason.
-- **Copied.** Registered scalar, aggregate, and window functions, table functions, the
- session configuration, and the analyzer and optimizer rule lists. These are
- snapshotted when the derived context is created, so a UDF registered on the original
- context afterwards is not visible to the derived one, and a `SET` applied to one does
- not reach the other.
+### What a derived context shares
-The session id is carried over to the fork, so both contexts report the same id, and so
-does every `TaskContext` either one hands to a foreign codec. A library that identifies a
-session by its id — to correlate host-side and worker-side state, or to assert which
-session a decode callback was bound to — keeps working across a fork.
+`with_logical_extension_codec`, `with_physical_extension_codec`, and
+`with_python_udf_inlining` return a new `SessionContext` wrapping the *same* underlying
+session. Only the Python-side codec settings differ; catalogs, tables, registered
+functions, and configuration are the one shared session, so a registration on either
+side is visible to both.
-Register functions before deriving, or register them directly on the derived context:
+`set_query_planner` does not return anything. The query planner lives in `SessionState`,
+so it is a property of the session rather than of a handle on it, and installing one is
+visible to every context sharing that session — including ones a `with_*` call returned
+earlier. Installing a codec on a session that already has a foreign planner rebuilds
+that planner against the new codec for the same reason: there is one planner, and it has
+to carry the codecs currently in force.
```python
ctx = SessionContext(config)
ctx = ctx.with_logical_extension_codec(provider_logical_codec)
ctx = ctx.with_physical_extension_codec(provider_physical_codec)
-ctx = ctx.with_query_planner(planner)
-ctx.register_udf(my_udf) # registered on the context that will run the query
+ctx.set_query_planner(planner)
+ctx.register_udf(my_udf)
```
-A session holds exactly one query planner. Calling `with_query_planner` again replaces
-the installed planner instead of layering another one. To chain planners, have the new
-planner wrap the capsule returned by `SessionContext.__datafusion_query_planner__()`
-and delegate to it explicitly.
+Order is a readability preference rather than a requirement — installing a codec after a
+planner rebuilds the planner against it.
+
+A session holds exactly one query planner. Calling `set_query_planner` again replaces the
+installed planner instead of layering another one. To chain planners, have the new
+planner wrap the capsule returned by `SessionContext.__datafusion_query_planner__()`,
+captured before the new planner is installed, and delegate to it explicitly.
## Alternative Approach
diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md
index 1ded2192..81e2b128 100644
--- a/docs/source/user-guide/upgrade-guides.md
+++ b/docs/source/user-guide/upgrade-guides.md
@@ -77,8 +77,11 @@ continue to work unchanged.
New in this release, `__datafusion_query_planner__` follows the same protocol.
It receives the session and takes both extension codecs from it, so a planner
-library never builds a `TaskContextProvider` at all. See the {ref}`ffi` guide
-for the full protocol.
+library never builds a `TaskContextProvider` at all. Install one with
+`SessionContext.set_query_planner(planner)`, which mutates the session the same
+way `add_physical_optimizer_rule` does and returns nothing — the query planner
+lives in `SessionState`, so it belongs to the session rather than to a
+particular handle on it. See the {ref}`ffi` guide for the full protocol.
### Changes to the `datafusion-python-util` crate
diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md
index f283d161..c46710da 100644
--- a/examples/datafusion-ffi-example/README.md
+++ b/examples/datafusion-ffi-example/README.md
@@ -40,9 +40,9 @@ This example makes the provider library the sole external codec owner. Register
```python
ctx = ctx.with_logical_extension_codec(provider_logical_codec)
ctx = ctx.with_physical_extension_codec(provider_physical_codec)
-ctx = ctx.with_query_planner(planner)
+ctx.set_query_planner(planner)
```
-Derived contexts also rebind an installed planner when codecs change, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly.
+Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly.
For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md
index af4d05b1..50eeb960 100644
--- a/examples/datafusion-ffi-query-planner-example/README.md
+++ b/examples/datafusion-ffi-query-planner-example/README.md
@@ -50,11 +50,11 @@ ctx = ctx.with_logical_extension_codec(provider_logical_codec)
ctx = ctx.with_physical_extension_codec(provider_physical_codec)
ctx.register_table("numbers", provider)
ctx.register_udf(provider_udf)
-ctx = ctx.with_query_planner(MyQueryPlanner())
+ctx.set_query_planner(MyQueryPlanner())
```
`MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit.
-The provider's codec pair is attached to the planner when the derived context is created and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit.
+The provider's codec pair is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against it, but planner-last order is easier to audit.
For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide.
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 9a5d78c4..51c561d3 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
@@ -23,6 +23,7 @@ import pytest
from datafusion import SessionConfig, SessionContext, udf
from datafusion_ffi_example import (
IsNullUDF,
+ MyCatalogProvider,
MyLogicalExtensionCodec,
MyPhysicalExtensionCodec,
MyPhysicalOptimizerRule,
@@ -70,7 +71,7 @@ def probe_context(
ctx = ctx.with_physical_extension_codec(physical_codec)
ctx.register_table("numbers", MyTableProvider(1, 6, 1))
ctx.register_udf(udf(IsNullUDF()))
- ctx = ctx.with_query_planner(MyQueryPlanner())
+ ctx.set_query_planner(MyQueryPlanner())
return ctx, logical_codec, physical_codec
@@ -121,20 +122,21 @@ def test_codec_sees_a_udf_registered_after_it_was_installed():
ctx.register_table("numbers", MyTableProvider(1, 6, 1))
# Registered after the codec was installed and bound to this session.
ctx.register_udf(udf(IsNullUDF()))
- ctx = ctx.with_query_planner(MyQueryPlanner())
+ ctx.set_query_planner(MyQueryPlanner())
batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert batches[0].column(0).to_pylist() == [0, 1, 2]
assert logical_codec.task_context_udf_resolutions() > 0
-def test_codec_follows_the_session_across_a_planner_fork():
- """Installing a planner forks the session, and the codec moves with it.
+def test_codec_sees_a_udf_registered_after_the_planner():
+ """Installing a planner does not detach the codec from the session.
- The codec is bound to the session before the fork and the function is
- registered on the fork afterwards, so resolving it proves the codec's task
- context provider was rebound to the forked session rather than left
- pointing at the one it was installed on.
+ The codec is bound to the session before the planner is installed and the
+ function is registered afterwards. Installing writes through
+ ``state_ref()`` rather than deriving a new ``SessionContext``, so the
+ codec's task context provider still points at the one live session and
+ sees the later registration.
"""
config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF)
@@ -142,8 +144,8 @@ def test_codec_follows_the_session_across_a_planner_fork():
ctx = ctx.with_logical_extension_codec(logical_codec)
ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
ctx.register_table("numbers", MyTableProvider(1, 6, 1))
- ctx = ctx.with_query_planner(MyQueryPlanner())
- # Registered on the fork, after the codec was installed on its parent.
+ ctx.set_query_planner(MyQueryPlanner())
+ # Registered after the planner, on the same session the codec is bound to.
ctx.register_udf(udf(IsNullUDF()))
batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
@@ -152,32 +154,64 @@ def test_codec_follows_the_session_across_a_planner_fork():
assert logical_codec.task_context_udf_resolutions() > 0
-def test_rebinding_a_fork_leaves_the_receiver_bound_to_its_own_session():
- """Rebinding the fork's codec must not disturb the context it came from.
+def test_registered_providers_survive_a_planner_install():
+ """A planner install must not orphan a previously registered provider.
- ``FFI_LogicalExtensionCodec::new`` clones the handle before adopting the
- new provider, so each context keeps its own binding. Both contexts here
- have a planner, so both exercise the codec; the function is registered only
- on the second, and only the second can resolve it.
+ A foreign catalog provider is handed a codec carrying a *weak*
+ ``FFI_TaskContextProvider`` pointing at the session it was registered on,
+ and upgrades it on every ``supports_filters_pushdown`` and every ``scan``.
+ That codec lives inside the registered ``FFI_CatalogProvider``, so nothing
+ on the Python side can reach it to rebind it. Deriving a replacement
+ ``SessionContext`` here would drop the allocation those handles point at
+ and the next query would fail with ``TaskContextProvider went out of scope
+ over FFI boundary``; installing in place keeps the one allocation alive.
+
+ The ``WHERE`` clause is load-bearing -- it forces filter pushdown, which
+ upgrades the weak handle during logical optimization, before any plan
+ serialization could fail first for an unrelated reason.
"""
- config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
- logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF)
+ config = SessionConfig().with_extension(MyPlannerConfig(max_rows=10))
ctx = SessionContext(config)
- ctx = ctx.with_logical_extension_codec(logical_codec)
- ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
- ctx.register_table("numbers", MyTableProvider(1, 6, 1))
+ ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider())
- first = ctx.with_query_planner(MyQueryPlanner())
- second = first.with_query_planner(MyQueryPlanner())
- second.register_udf(udf(IsNullUDF()))
+ ctx.set_query_planner(MyQueryPlanner())
+ gc.collect()
- batches = second.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
- assert batches[0].column(0).to_pylist() == [0, 1, 2]
+ batches = ctx.sql(
+ "SELECT units FROM ffi_catalog.my_schema.my_table WHERE units > 5"
+ ).collect()
+ assert sorted(v for b in batches for v in b.column(0).to_pylist()) == [
+ 7,
+ 10,
+ 20,
+ 30,
+ ]
- with pytest.raises(
- Exception, match=rf"could not resolve scalar function '{HOST_ONLY_UDF}'"
- ):
- first.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+
+def test_registered_providers_survive_a_codec_install_after_a_planner():
+ """The same, for the other write path.
+
+ Installing a codec on a session that already has a foreign planner rebuilds
+ that planner against the new codec. That rebuild must also happen in place.
+ """
+ config = SessionConfig().with_extension(MyPlannerConfig(max_rows=10))
+ ctx = SessionContext(config)
+ ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider())
+
+ ctx.set_query_planner(MyQueryPlanner())
+ ctx = ctx.with_logical_extension_codec(MyLogicalExtensionCodec())
+ ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
+ gc.collect()
+
+ batches = ctx.sql(
+ "SELECT units FROM ffi_catalog.my_schema.my_table WHERE units > 5"
+ ).collect()
+ assert sorted(v for b in batches for v in b.column(0).to_pylist()) == [
+ 7,
+ 10,
+ 20,
+ 30,
+ ]
def codec_context(max_rows: int = 3):
@@ -185,7 +219,7 @@ def codec_context(max_rows: int = 3):
Unlike :func:`probe_context` the codecs ask for no function, so the only
thing they record is the session id of the task context they are handed.
- No planner yet -- the caller installs one, since that is what forks.
+ No planner yet -- the caller installs one.
"""
config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows))
logical_codec = MyLogicalExtensionCodec()
@@ -197,13 +231,13 @@ def codec_context(max_rows: int = 3):
return ctx, logical_codec, physical_codec
-def test_a_fork_decodes_against_the_session_id_it_reports():
- """A fork and its decode callbacks must agree on the session id.
+def test_installing_a_planner_keeps_the_session_id():
+ """A session and its decode callbacks must agree on the session id.
- ``with_query_planner`` forks the session state by rebuilding it through
- ``SessionStateBuilder``, which mints a fresh id unless handed one, and
+ Installing a planner rebuilds ``SessionState`` through
+ ``SessionStateBuilder``, which mints a fresh id unless handed one, while
``SessionContext`` caches its id in a field of its own. Dropping the id
- there leaves the fork reporting one id from ``session_id()`` and a
+ there leaves the session reporting one id from ``session_id()`` and a
different one from every ``TaskContext`` it gives a foreign codec.
Asserting on ``session_id()`` alone cannot catch that: it reads the cached
@@ -214,55 +248,46 @@ def test_a_fork_decodes_against_the_session_id_it_reports():
ctx, logical_codec, physical_codec = codec_context()
session_id = ctx.session_id()
- fork = ctx.with_query_planner(MyQueryPlanner())
- assert fork.session_id() == session_id
+ ctx.set_query_planner(MyQueryPlanner())
+ assert ctx.session_id() == session_id
- fork.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+ ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert logical_codec.last_task_context_session_id() == session_id
assert physical_codec.last_task_context_session_id() == session_id
def test_adding_a_physical_optimizer_rule_keeps_the_session_id():
- """Mutating a session in place must not move the id it decodes against.
+ """The same guarantee for the other in-place ``SessionState`` rewrite.
- ``add_physical_optimizer_rule`` rebuilds ``SessionState`` and writes it
- back into the caller's own session rather than deriving a new one, so a
- regenerated id would desync a context from itself with no fork to explain
- it.
+ ``add_physical_optimizer_rule`` also rebuilds ``SessionState`` and writes
+ it back, so a regenerated id would desync a context from itself.
"""
ctx, logical_codec, physical_codec = codec_context()
- fork = ctx.with_query_planner(MyQueryPlanner())
- session_id = fork.session_id()
+ ctx.set_query_planner(MyQueryPlanner())
+ session_id = ctx.session_id()
- fork.add_physical_optimizer_rule(MyPhysicalOptimizerRule())
- assert fork.session_id() == session_id
+ ctx.add_physical_optimizer_rule(MyPhysicalOptimizerRule())
+ assert ctx.session_id() == session_id
- fork.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+ ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert logical_codec.last_task_context_session_id() == session_id
assert physical_codec.last_task_context_session_id() == session_id
-def test_rebinding_a_fork_does_not_move_the_parent_session_id():
- """Each context in a fork chain decodes against its own id.
+def test_replacing_a_planner_keeps_the_session_id():
+ """Installing repeatedly must not drift the id.
- Two forks off one parent share the parent's id, so a test that only
- compared against the parent would pass even if rebinding leaked one
- context's provider into the other. Registering a function on just one fork
- is what distinguishes them.
+ Each install rebuilds ``SessionState``, so an id carried over only on the
+ first write would still be lost by the second.
"""
ctx, logical_codec, _physical_codec = codec_context()
session_id = ctx.session_id()
- first = ctx.with_query_planner(MyQueryPlanner())
- second = first.with_query_planner(MyQueryPlanner())
-
- assert first.session_id() == session_id
- assert second.session_id() == session_id
-
- first.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
- assert logical_codec.last_task_context_session_id() == session_id
+ ctx.set_query_planner(MyQueryPlanner())
+ ctx.set_query_planner(MyQueryPlanner())
+ assert ctx.session_id() == session_id
- second.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+ ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert logical_codec.last_task_context_session_id() == session_id
@@ -274,7 +299,7 @@ def test_three_library_query_planner(raw_capsule: bool):
exported_planner = (
planner.__datafusion_query_planner__(ctx) if raw_capsule else planner
)
- ctx = ctx.with_query_planner(exported_planner)
+ ctx.set_query_planner(exported_planner)
batches = ctx.sql(
'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"'
@@ -315,7 +340,7 @@ def test_spawning_plan_across_three_libraries():
ctx.register_table("numbers", MyTableProvider(1, 6, 3))
planner = MyQueryPlanner()
- ctx = ctx.with_query_planner(planner)
+ ctx.set_query_planner(planner)
batches = ctx.sql(
'SELECT "A" % 2 AS parity, count(*) AS n FROM numbers GROUP BY 1 ORDER BY 1'
@@ -343,11 +368,11 @@ def test_planner_layers_on_the_session_planner():
ctx, logical_codec, physical_codec = configured_context(max_rows=3)
fallback = ctx.__datafusion_query_planner__()
planner = MyQueryPlanner(fallback=fallback)
- # Rebinding `ctx` drops the context that produced the capsule. The exported
- # codecs retain it, so the capsule stays usable. Without that the FFI
- # task-context handle is weak and planning fails with "TaskContextProvider
- # went out of scope over FFI boundary".
- ctx = ctx.with_query_planner(planner)
+ # The capsule's FFI codecs hold weak handles to this session. Installing in
+ # place keeps that session alive, so the capsule stays usable; deriving a
+ # replacement here would fail with "TaskContextProvider went out of scope
+ # over FFI boundary".
+ ctx.set_query_planner(planner)
gc.collect()
batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
@@ -370,7 +395,7 @@ def test_a_planner_can_fall_back_to_another_planner_library():
ctx, logical_codec, physical_codec = configured_context(max_rows=3)
inner = MyQueryPlanner()
outer = MyQueryPlanner(fallback=inner)
- ctx = ctx.with_query_planner(outer)
+ ctx.set_query_planner(outer)
batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert batches[0].column(0).to_pylist() == [0, 1, 2]
@@ -387,10 +412,10 @@ def test_a_session_fallback_delegates_to_its_installed_planner():
"""Passing a SessionContext delegates to whatever planner it holds."""
ctx, _logical_codec, _physical_codec = configured_context(max_rows=3)
first = MyQueryPlanner()
- ctx = ctx.with_query_planner(first)
+ ctx.set_query_planner(first)
second = MyQueryPlanner(fallback=ctx)
- ctx = ctx.with_query_planner(second)
+ ctx.set_query_planner(second)
batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert batches[0].column(0).to_pylist() == [0, 1, 2]
@@ -410,7 +435,7 @@ def test_observations_accumulate_across_queries():
"""
ctx, _logical_codec, _physical_codec = configured_context(max_rows=3)
planner = MyQueryPlanner()
- ctx = ctx.with_query_planner(planner)
+ ctx.set_query_planner(planner)
ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert planner.foreign_session_observed()
@@ -433,7 +458,8 @@ def test_second_planner_replaces_the_first():
ctx, _logical_codec, _physical_codec = configured_context(max_rows=2)
first = MyQueryPlanner()
second = MyQueryPlanner()
- ctx = ctx.with_query_planner(first).with_query_planner(second)
+ ctx.set_query_planner(first)
+ ctx.set_query_planner(second)
batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert batches[0].column(0).to_pylist() == [0, 1]
@@ -441,16 +467,23 @@ def test_second_planner_replaces_the_first():
assert first.plan_calls() == 0
-def test_planner_is_not_installed_on_the_original_context():
- """``with_query_planner`` returns a fork; the receiver keeps its planner."""
+def test_the_planner_reaches_every_handle_on_the_session():
+ """The planner is session state, so all handles on that session use it.
+
+ ``set_query_planner`` writes through ``state_ref()``. A context returned by
+ an earlier ``with_*`` call shares that session, so it plans through the new
+ planner too -- there is one session and one planner, not a family of
+ diverging copies.
+ """
ctx, _logical_codec, _physical_codec = configured_context(max_rows=2)
- planner = MyQueryPlanner()
- derived = ctx.with_query_planner(planner)
+ # Shares the session with `ctx`; predates the planner.
+ sibling = ctx.with_python_udf_inlining(enabled=False)
- ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
- assert planner.plan_calls() == 0
+ planner = MyQueryPlanner()
+ ctx.set_query_planner(planner)
- derived.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+ batches = sibling.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+ assert batches[0].column(0).to_pylist() == [0, 1]
assert planner.plan_calls() > 0
@@ -459,7 +492,7 @@ def test_installed_codecs_outlive_python_exporters():
del logical_codec, physical_codec
gc.collect()
- ctx = ctx.with_query_planner(MyQueryPlanner())
+ ctx.set_query_planner(MyQueryPlanner())
batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
assert batches[0].column(0).to_pylist() == [0, 1]
@@ -469,7 +502,8 @@ def test_provider_codecs_can_be_installed_after_planner():
planner = MyQueryPlanner()
logical_codec = MyLogicalExtensionCodec()
physical_codec = MyPhysicalExtensionCodec()
- ctx = SessionContext(config).with_query_planner(planner)
+ ctx = SessionContext(config)
+ ctx.set_query_planner(planner)
ctx = ctx.with_logical_extension_codec(logical_codec)
ctx = ctx.with_physical_extension_codec(physical_codec)
ctx.register_table("numbers", MyTableProvider(1, 4, 1))
@@ -485,7 +519,7 @@ def test_query_planner_requires_provider_codec():
config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2))
ctx = SessionContext(config)
ctx.register_table("numbers", MyTableProvider(1, 3, 1))
- ctx = ctx.with_query_planner(MyQueryPlanner())
+ ctx.set_query_planner(MyQueryPlanner())
with pytest.raises(Exception, match=r"LogicalExtensionCodec|TableProvider"):
ctx.sql('SELECT "A" FROM numbers').collect()
@@ -494,7 +528,7 @@ def test_query_planner_requires_provider_codec():
@pytest.mark.parametrize("max_rows", ["0", "oops"])
def test_query_planner_rejects_invalid_config(max_rows: str):
ctx, _logical_codec, _physical_codec = configured_context(max_rows=2)
- ctx = ctx.with_query_planner(MyQueryPlanner())
+ ctx.set_query_planner(MyQueryPlanner())
with pytest.raises(Exception, match=r"max_rows|Invalid value"):
ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect()
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index 6e0560f1..0f9d5328 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -1771,49 +1771,43 @@ class SessionContext:
"""
self.ctx.add_physical_optimizer_rule(rule)
- def with_query_planner(
- self, planner: QueryPlannerExportable | _PyCapsule
- ) -> SessionContext:
- """Create a new session context with a custom query planner.
+ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> None:
+ """Install a custom query planner on this session.
The planner is imported through its ``__datafusion_query_planner__``
- PyCapsule. The returned context carries over the current session state
- and the logical and physical extension codec settings. Codec changes
- made on a derived context are rebound to the planner before planning.
+ PyCapsule and installed on this context, in the same way
+ :meth:`~SessionContext.add_physical_optimizer_rule` installs a rule.
+ The query planner is part of the session state, so it applies to this
+ context and to every context sharing its session — including ones
+ already returned by
+ :meth:`~SessionContext.with_logical_extension_codec` and friends.
A session holds exactly one planner, so calling this again replaces the
previous one rather than layering. To chain planners, have the new
planner wrap the capsule from
- :meth:`~SessionContext.__datafusion_query_planner__`.
-
- .. note:: Derived contexts share catalogs, not registrations
- The returned context is a fork. Catalogs, schemas, and tables stay
- shared with the original context, but registered functions and
- configuration are copied at the time of the call. A UDF registered
- on the original context afterwards is **not** visible here, while a
- table registered on either context is visible to both. Register
- functions before deriving, or register them on the derived context.
+ :meth:`~SessionContext.__datafusion_query_planner__`, captured
+ *before* the new planner is installed.
Args:
planner: Object exposing ``__datafusion_query_planner__`` (see
:class:`QueryPlannerExportable`) or a raw
``datafusion_query_planner`` PyCapsule.
- Returns:
- A new context that uses the specified query planner.
-
Examples:
>>> from my_extension import DistributedQueryPlanner # doctest: +SKIP
>>> ctx = SessionContext()
- >>> planner = DistributedQueryPlanner() # doctest: +SKIP
- >>> planner_ctx = ctx.with_query_planner(planner) # doctest: +SKIP
- >>> query = planner_ctx.sql("SELECT * FROM remote_table") # doctest: +SKIP
- >>> query.collect() # doctest: +SKIP
+ >>> ctx.set_query_planner(DistributedQueryPlanner()) # doctest: +SKIP
+ >>> ctx.sql("SELECT * FROM remote_table").collect() # doctest: +SKIP
+
+ Layer a planner on top of the one already installed by capturing
+ the existing planner first:
+
+ >>> fallback = ctx.__datafusion_query_planner__() # doctest: +SKIP
+ >>> ctx.set_query_planner(
+ ... DistributedQueryPlanner(fallback=fallback)
+ ... ) # doctest: +SKIP
"""
- new_internal = self.ctx.with_query_planner(planner)
- new = SessionContext.__new__(SessionContext)
- new.ctx = new_internal
- return new
+ self.ctx.set_query_planner(planner)
def table_provider(self, name: str) -> Table:
"""Return the :py:class:`~datafusion.catalog.Table` for the given table name.
@@ -2261,10 +2255,9 @@ class SessionContext:
:py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`).
The returned context shares its session state with the original, so a
- later registration on either is visible to both. The exception is a
- session with a custom query planner installed: that planner has to be
- rebound to the new codec, which forks the state. See
- :meth:`~SessionContext.with_query_planner` for what a fork shares.
+ later registration on either is visible to both. If a custom query
+ planner is installed, it is rebuilt against the new codec on the shared
+ session, so the original context plans with the new codec too.
"""
new_internal = self.ctx.with_logical_extension_codec(codec)
new = SessionContext.__new__(SessionContext)
@@ -2288,10 +2281,9 @@ class SessionContext:
:py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`).
The returned context shares its session state with the original, so a
- later registration on either is visible to both. The exception is a
- session with a custom query planner installed: that planner has to be
- rebound to the new codec, which forks the state. See
- :meth:`~SessionContext.with_query_planner` for what a fork shares.
+ later registration on either is visible to both. If a custom query
+ planner is installed, it is rebuilt against the new codec on the shared
+ session, so the original context plans with the new codec too.
"""
new_internal = self.ctx.with_physical_extension_codec(codec)
new = SessionContext.__new__(SessionContext)
@@ -2336,11 +2328,9 @@ class SessionContext:
Returns a new :class:`SessionContext` with the toggle applied;
the original session is unchanged. The returned context shares
its session state with the original, so a later registration on
- either is visible to both. The exception is a session with a
- custom query planner installed: that planner has to be rebound
- to the new codecs, which forks the state. See
- :meth:`~SessionContext.with_query_planner` for what a fork
- shares.
+ either is visible to both. If a custom query planner is
+ installed, it is rebuilt against the new codecs on the shared
+ session, so the original context plans with them too.
Examples:
>>> import pyarrow as pa
diff --git a/python/tests/test_context.py b/python/tests/test_context.py
index a8d5a116..89759eb4 100644
--- a/python/tests/test_context.py
+++ b/python/tests/test_context.py
@@ -732,9 +732,9 @@ def test_remove_optimizer_rule(ctx):
assert ctx.remove_optimizer_rule("nonexistent_rule") is False
-def test_with_query_planner_rejects_wrong_capsule(ctx):
+def test_set_query_planner_rejects_wrong_capsule(ctx):
with pytest.raises(ValueError, match="datafusion_query_planner"):
- ctx.with_query_planner(ctx.__datafusion_task_context_provider__())
+ ctx.set_query_planner(ctx.__datafusion_task_context_provider__())
def test_pre_55_codec_signature_reports_an_upgrade(ctx):
@@ -788,7 +788,7 @@ def test_non_type_errors_from_a_getter_propagate(ctx):
ctx.with_logical_extension_codec(RaisesValueError())
-def test_with_query_planner_capsule(ctx):
+def test_set_query_planner_capsule(ctx):
capsule = ctx.__datafusion_query_planner__()
get_name = ctypes.pythonapi.PyCapsule_GetName
get_name.argtypes = [ctypes.py_object]
@@ -799,62 +799,69 @@ def test_with_query_planner_capsule(ctx):
"query_planner_test",
[[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]],
)
- planner_context = ctx.with_query_planner(capsule)
- assert planner_context.table_exist("query_planner_test")
- batches = planner_context.sql("SELECT 1 AS value").collect()
+ ctx.set_query_planner(capsule)
+ assert ctx.table_exist("query_planner_test")
+ batches = ctx.sql("SELECT 1 AS value").collect()
assert batches[0].column(0) == pa.array([1])
-def test_derived_context_shares_catalogs(ctx):
- """Catalogs live behind an Arc, so tables cross the fork in both directions."""
- derived = ctx.with_query_planner(ctx.__datafusion_query_planner__())
+def test_installing_a_planner_leaves_the_session_intact(ctx):
+ """The planner is written into the existing session, not a copy of it.
+ Registrations made before the install are still visible afterwards, and
+ ones made after are visible too -- there is a single session throughout,
+ so neither the catalogs nor the function registry are snapshotted.
+ """
ctx.register_record_batches(
- "registered_on_parent",
+ "registered_before",
[[pa.RecordBatch.from_pydict({"value": [1]})]],
)
- derived.register_record_batches(
- "registered_on_derived",
- [[pa.RecordBatch.from_pydict({"value": [2]})]],
- )
-
- assert derived.table_exist("registered_on_parent")
- assert ctx.table_exist("registered_on_derived")
-
-
-def test_derived_context_snapshots_functions(ctx):
- """Function registries are copied at fork time, unlike catalogs.
-
- A UDF registered on the parent before the fork is carried over; one
- registered afterwards is not. Guards the caveat documented on
- ``SessionContext.with_query_planner``.
- """
before = udf(
lambda arr: arr,
[pa.int64()],
pa.int64(),
volatility="immutable",
- name="registered_before_fork",
+ name="registered_before",
)
ctx.register_udf(before)
- derived = ctx.with_query_planner(ctx.__datafusion_query_planner__())
+ ctx.set_query_planner(ctx.__datafusion_query_planner__())
+ ctx.register_record_batches(
+ "registered_after",
+ [[pa.RecordBatch.from_pydict({"value": [2]})]],
+ )
after = udf(
lambda arr: arr,
[pa.int64()],
pa.int64(),
volatility="immutable",
- name="registered_after_fork",
+ name="registered_after",
)
ctx.register_udf(after)
- assert derived.sql("SELECT registered_before_fork(1)").collect()
- with pytest.raises(Exception, match="registered_after_fork"):
- derived.sql("SELECT registered_after_fork(1)").collect()
+ assert ctx.table_exist("registered_before")
+ assert ctx.table_exist("registered_after")
+ assert ctx.sql("SELECT registered_before(1)").collect()
+ assert ctx.sql("SELECT registered_after(1)").collect()
+
+
+def test_contexts_sharing_a_session_share_the_planner(ctx):
+ """A context derived before the install still plans through the planner.
+
+ ``with_python_udf_inlining`` returns a handle on the same session, and the
+ query planner lives in that session's state.
+ """
+ sibling = ctx.with_python_udf_inlining(enabled=False)
+ ctx.register_record_batches(
+ "shared_planner_test",
+ [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]],
+ )
+
+ ctx.set_query_planner(ctx.__datafusion_query_planner__())
- # The parent is unaffected by the fork.
- assert ctx.sql("SELECT registered_after_fork(1)").collect()
+ assert sibling.table_exist("shared_planner_test")
+ assert sibling.session_id() == ctx.session_id()
def test_table_provider(ctx):
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-python) 02/03: test: install codecs in the provider-survival 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-query-planner-core
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit 06c9b9b8c6d4cb4f90646910754220017d2193c0
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Aug 28 13:42:30 2026 -0400
test: install codecs in the provider-survival test
The test registered a foreign catalog provider and no codecs, so it could
never finish: once the planner install stopped orphaning the provider, the
query got past filter pushdown and then failed at plan serialization with
`LogicalExtensionCodec is not provided`. That is the same unrelated failure
`test_query_planner_requires_provider_codec` already covers, and it would
mask a dangling handle rather than expose one.
Install both provider codecs, and fold the codec-install-after-planner case
in as a parameter rather than a near-duplicate test. Both orderings write
`SessionState` — one installs the planner, the other rebuilds it against a
new codec — so both exercise the path that must not replace the session's
`Arc<SessionContext>`.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---
.../tests/_test_three_library_query_planner.py | 56 +++++++++-------------
.../datafusion-ffi-query-planner-example/uv.lock | 7 +++
2 files changed, 30 insertions(+), 33 deletions(-)
diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py
index 51c561d3..f3982091 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
@@ -154,53 +154,43 @@ def test_codec_sees_a_udf_registered_after_the_planner():
assert logical_codec.task_context_udf_resolutions() > 0
-def test_registered_providers_survive_a_planner_install():
- """A planner install must not orphan a previously registered provider.
+@pytest.mark.parametrize("codecs_first", [True, False])
+def test_registered_providers_survive_a_planner_install(codecs_first: bool):
+ """Neither write path may orphan a previously registered provider.
A foreign catalog provider is handed a codec carrying a *weak*
``FFI_TaskContextProvider`` pointing at the session it was registered on,
and upgrades it on every ``supports_filters_pushdown`` and every ``scan``.
That codec lives inside the registered ``FFI_CatalogProvider``, so nothing
on the Python side can reach it to rebind it. Deriving a replacement
- ``SessionContext`` here would drop the allocation those handles point at
- and the next query would fail with ``TaskContextProvider went out of scope
- over FFI boundary``; installing in place keeps the one allocation alive.
+ ``SessionContext`` would drop the allocation those handles point at, and
+ the next query would fail with ``TaskContextProvider went out of scope over
+ FFI boundary``. Installing in place keeps the one allocation alive.
+
+ Both parameters exercise that, because both write ``SessionState``:
+ ``set_query_planner`` installs the planner, and installing a codec on a
+ session that already has one rebuilds that planner against the new codec.
The ``WHERE`` clause is load-bearing -- it forces filter pushdown, which
- upgrades the weak handle during logical optimization, before any plan
- serialization could fail first for an unrelated reason.
+ upgrades the weak handle during logical optimization. It is also why both
+ codecs have to be installed: without them the query fails at plan
+ serialization with ``LogicalExtensionCodec is not provided``, which would
+ mask a dangling handle rather than expose it.
"""
config = SessionConfig().with_extension(MyPlannerConfig(max_rows=10))
ctx = SessionContext(config)
ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider())
- ctx.set_query_planner(MyQueryPlanner())
- gc.collect()
-
- batches = ctx.sql(
- "SELECT units FROM ffi_catalog.my_schema.my_table WHERE units > 5"
- ).collect()
- assert sorted(v for b in batches for v in b.column(0).to_pylist()) == [
- 7,
- 10,
- 20,
- 30,
- ]
-
+ def install_codecs(ctx):
+ ctx = ctx.with_logical_extension_codec(MyLogicalExtensionCodec())
+ return ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
-def test_registered_providers_survive_a_codec_install_after_a_planner():
- """The same, for the other write path.
-
- Installing a codec on a session that already has a foreign planner rebuilds
- that planner against the new codec. That rebuild must also happen in place.
- """
- config = SessionConfig().with_extension(MyPlannerConfig(max_rows=10))
- ctx = SessionContext(config)
- ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider())
-
- ctx.set_query_planner(MyQueryPlanner())
- ctx = ctx.with_logical_extension_codec(MyLogicalExtensionCodec())
- ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec())
+ if codecs_first:
+ ctx = install_codecs(ctx)
+ ctx.set_query_planner(MyQueryPlanner())
+ else:
+ ctx.set_query_planner(MyQueryPlanner())
+ ctx = install_codecs(ctx)
gc.collect()
batches = ctx.sql(
diff --git a/examples/datafusion-ffi-query-planner-example/uv.lock b/examples/datafusion-ffi-query-planner-example/uv.lock
new file mode 100644
index 00000000..0c546a7e
--- /dev/null
+++ b/examples/datafusion-ffi-query-planner-example/uv.lock
@@ -0,0 +1,7 @@
+version = 1
+revision = 3
+requires-python = ">=3.10"
+
+[[package]]
+name = "datafusion-ffi-query-planner-example"
+source = { editable = "." }
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org