From 94cb55fe545b85818911511794c76af4ec2095b4 Mon Sep 17 00:00:00 2001
From: Tim Saucer <timsaucer@gmail.com>
Date: Tue, 22 Sep 2026 16:30:13 -0400
Subject: [PATCH 1/6] chore: harden the extension API before it ships

`SessionContext.with_extensions` and `SessionExtensionComponents` landed
in #1679 and have not shipped in a release yet. The bundle stack
(#1738-#1741) reshapes them substantially, and a release would freeze
three surfaces in their current form.

`PhysicalOptimizerRuleExportable` was defined in `datafusion.context` and
not exported from the package root, so `datafusion.context` would become
its canonical import path. Move it to `datafusion.extensions` beside the
rest of the `*Exportable` family, re-export it from `datafusion.context`
so the old path keeps working, and export it from the package root. The
move brings it under `test_extension_api_has_a_doctest`, which drives off
`extensions.__all__`, so it gains the example it was missing.

`SessionExtensionComponents` was positionally constructible with two
fields. The stack takes it to nine, three of them pair-shaped. Make
construction keyword-only so every later field addition is additive; no
call site in the repository constructed it positionally. This is a new
convention rather than a backport, so it has to be applied forward to the
stack as well.

The ordering that makes `with_extensions` transactional was stated in
three docstrings with no canonical home to point at. Record it under
`ffi_internals_commit_order` in the contributor guide, and label the
existing "Failure and rollback" section `extension_bundles_transaction`,
matching the names the stack links to.

No released behaviour changes, so no `api change` label and no
upgrade-guide entry.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .../source/contributor-guide/ffi-internals.md | 40 ++++++++++++++
 docs/source/extension-guide/bundles.md        |  8 +++
 python/datafusion/__init__.py                 |  2 +
 python/datafusion/context.py                  | 19 +++----
 python/datafusion/extensions.py               | 53 ++++++++++++++++---
 5 files changed, 105 insertions(+), 17 deletions(-)

diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md
index 75c32444c..7fdfd6a6e 100644
--- a/docs/source/contributor-guide/ffi-internals.md
+++ b/docs/source/contributor-guide/ffi-internals.md
@@ -111,6 +111,46 @@ library would serialize, and would do it with the codecs it was imported with.
 The extension-facing consequence — install codecs before a layered planner, and
 prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`.
 
+(ffi_internals_commit_order)=
+
+## Why `with_extensions` commits last
+
+`with_extensions` promises that a bundle which raises leaves the session as it
+was. Keeping that promise is an ordering constraint on the implementation, not
+a property of any one step, because the planner is bound on the shared
+`SessionState` rather than on the returned handle.
+
+A call therefore splits into a part that may fail and a part that may not:
+
+1. **Collect.** Every `__datafusion_session_components__` runs and its codecs
+   are gathered. Nothing is installed yet, so a hook that raises here has
+   touched nothing.
+2. **Chains.** The codecs are assembled into the returned handle. Codec chains
+   live on that handle rather than on the session, so this step writes nothing
+   to the session even though it can fail on a bad capsule or a duplicate id.
+3. **Resolve.** Every `__datafusion_session_planner__` runs, in argument order,
+   against the completed chains, and each supplied planner is exported to a
+   capsule. Everything that can raise has raised by the end of this step.
+4. **Commit.** The accumulated planner is bound, in a single `SessionState`
+   rebuild. The bind is skipped entirely when the call installed nothing, so an
+   empty call does not drag a planner sitting on another handle's codecs onto
+   this one's.
+
+Only step 4 touches the session. This is a rule for the next field added to
+`SessionExtensionComponents`, not only a description of the current code: a new
+kind of component must do its fallible work — importing a capsule, resolving a
+name — in step 3, so that step 4 cannot raise part-way through.
+
+There is nothing to roll back to if it does. The returned handle shares one
+session with the receiver, so the damage is visible from every other handle;
+and undoing a registration is not the same as restoring what it displaced,
+because deregistering a function that shadowed a built-in removes the built-in
+too. The split is cheaper than an undo log that cannot be written correctly.
+
+The extension-facing statement of this is
+{ref}`extension_bundles_transaction`, which says only that declaring a
+component is safe where registering one during the hook is not.
+
 ## Two argument kinds for one convention
 
 `CapsuleGetterArg` in `crates/util/src/lib.rs` distinguishes three cases: no
diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md
index 791fdfc8e..437deb165 100644
--- a/docs/source/extension-guide/bundles.md
+++ b/docs/source/extension-guide/bundles.md
@@ -281,6 +281,8 @@ for direct ones. The wrapper travels with the codec; the bundle does not.
 The query planner is exempt — it carries no wire id, so it may be an object or
 a capsule.
 
+(extension_bundles_transaction)=
+
 ## Failure and rollback
 
 Nothing is written to the session until every factory has returned and every
@@ -290,6 +292,12 @@ table, say — is **not** rolled back, which is why bundle objects must be
 configuration-only: create fresh components on each call, never cache bound
 components, and do not retain the context passed in.
 
+Declaring a component is what buys you that guarantee. Anything you return from
+your hook is validated while a failure still costs nothing, and is written only
+after every bundle in the call has succeeded. Anything you register yourself is
+written immediately, before the other bundles have even run. The ordering that
+makes this hold is recorded at {ref}`ffi_internals_commit_order`.
+
 Like every other derivation, the returned context is a handle on the *same*
 session as the receiver — see {ref}`extension_sessions`. Only the Python-side
 codec chains belong to the returned handle; the planner is installed on the
diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py
index 1b44f8a73..f5a2b3849 100644
--- a/python/datafusion/__init__.py
+++ b/python/datafusion/__init__.py
@@ -93,6 +93,7 @@
 from .dataframe_formatter import configure_formatter
 from .expr import Expr, WindowFrame
 from .extensions import (
+    PhysicalOptimizerRuleExportable,
     QueryPlannerExportable,
     SessionComponentsExportable,
     SessionExtensionComponents,
@@ -139,6 +140,7 @@
     "MetricsSet",
     "ParquetColumnOptions",
     "ParquetWriterOptions",
+    "PhysicalOptimizerRuleExportable",
     "PhysicalPartitioning",
     "QueryPlannerExportable",
     "RecordBatch",
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index 63cdfd487..1dedfa6dd 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -70,6 +70,7 @@
 from datafusion.dataframe import DataFrame
 from datafusion.expr import sort_list_to_raw_sort_list
 from datafusion.extensions import (
+    PhysicalOptimizerRuleExportable,
     QueryPlannerExportable,
     SessionComponentsExportable,
     SessionExtensionComponents,
@@ -151,16 +152,6 @@ class TableProviderExportable(Protocol):
     def __datafusion_table_provider__(self, session: Any) -> object: ...  # noqa: D105
 
 
-class PhysicalOptimizerRuleExportable(Protocol):
-    """Type hint for object that has __datafusion_physical_optimizer_rule__ PyCapsule.
-
-    The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``,
-    typically produced by a separate compiled extension.
-    """
-
-    def __datafusion_physical_optimizer_rule__(self) -> object: ...  # noqa: D105
-
-
 class SessionConfig:
     """Session configuration options."""
 
@@ -1918,7 +1909,8 @@ def with_extensions(
         every capsule has been validated, so a hook that raises leaves the
         session as it was. A hook that *mutates* the context it is handed —
         registering a table, say — is not rolled back, which is why bundle
-        objects must be configuration-only.
+        objects must be configuration-only. See
+        :ref:`extension_bundles_transaction`.
 
         Shares its session with this context — see :py:class:`SessionContext`.
 
@@ -2030,6 +2022,11 @@ def with_extensions(
                 continue
             planner = new.ctx._export_query_planner(supplied)
 
+        # The commit step. Everything above is allowed to raise; this is not.
+        # See docs/source/contributor-guide/ffi-internals.md, "Why
+        # `with_extensions` commits last", for what a new component kind has to
+        # do to keep that true.
+        #
         # Rebinding the session's planner is a side effect on state shared with
         # every other handle, so do not pay it for a call that installs nothing
         # -- the same guard `with_python_udf_inlining` carries. With no codec
diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py
index 77ae92fc2..196e25ad0 100644
--- a/python/datafusion/extensions.py
+++ b/python/datafusion/extensions.py
@@ -32,11 +32,12 @@
 implements either hook or both. Bundle order is significant for planners, which
 nest, and irrelevant for codecs, which accumulate.
 
-Of the four names here, only the two bundle hooks are ``@runtime_checkable``,
+Of the five names here, only the two bundle hooks are ``@runtime_checkable``,
 because :py:meth:`~datafusion.context.SessionContext.with_extensions`
-dispatches on them from Python. :py:class:`QueryPlannerExportable` is a type
-hint only, matching the other capsule-getter protocols in
-:py:mod:`datafusion.user_defined` and :py:mod:`datafusion.catalog`.
+dispatches on them from Python. :py:class:`QueryPlannerExportable` and
+:py:class:`PhysicalOptimizerRuleExportable` are type hints only, matching the
+other capsule-getter protocols in :py:mod:`datafusion.user_defined` and
+:py:mod:`datafusion.catalog`.
 
 See :ref:`extension_bundles` in the online documentation for why the phases are
 split and for a worked implementation.
@@ -57,6 +58,7 @@
     )
 
 __all__ = [
+    "PhysicalOptimizerRuleExportable",
     "QueryPlannerExportable",
     "SessionComponentsExportable",
     "SessionExtensionComponents",
@@ -64,6 +66,41 @@
 ]
 
 
+class PhysicalOptimizerRuleExportable(Protocol):
+    """Type hint for object that has __datafusion_physical_optimizer_rule__ PyCapsule.
+
+    The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``,
+    typically produced by a separate compiled extension. It takes **no
+    argument**: a rule needs neither a codec nor a task-context provider, so
+    there is nothing session-scoped to hand it.
+
+    Rules accumulate rather than replace. Install one with
+    :py:meth:`~datafusion.context.SessionContext.add_physical_optimizer_rule`
+    — see :ref:`extension_other_hooks`.
+
+    Examples:
+        The getter is the whole protocol, and a capsule is what it must return
+        — anything else is refused where it is installed rather than at plan
+        time:
+
+        >>> from datafusion import SessionContext
+        >>> ctx = SessionContext()
+        >>> ctx.add_physical_optimizer_rule(object())
+        Traceback (most recent call last):
+            ...
+        RuntimeError: "Invalid datafusion_physical_optimizer_rule...
+
+        Real usage. Skipped here (needs a built extension library); run for
+        real by ``test_ffi_physical_optimizer_rule_runs_during_planning`` in
+        ``datafusion-ffi-example``.
+
+        >>> from datafusion_ffi_example import MyPhysicalOptimizerRule  # doctest: +SKIP
+        >>> ctx.add_physical_optimizer_rule(MyPhysicalOptimizerRule())  # doctest: +SKIP
+    """
+
+    def __datafusion_physical_optimizer_rule__(self) -> object: ...  # noqa: D105
+
+
 class QueryPlannerExportable(Protocol):
     """Type hint for object that has a __datafusion_query_planner__ PyCapsule.
 
@@ -110,7 +147,7 @@ def _not_a_codec_iterable(field: str, value: object) -> str:
     )
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, kw_only=True)
 class SessionExtensionComponents:
     """Components an extension contributes to a session context.
 
@@ -121,6 +158,9 @@ class SessionExtensionComponents:
     components bound to a different session hold a task-context provider for
     that other session and cannot be rebound.
 
+    Construction is keyword-only, so later releases can add component kinds
+    without changing what an existing call means.
+
     Query planners are not listed here. They install in a second phase so each
     can wrap the one before it — see :py:class:`SessionPlannerExportable`.
 
@@ -233,7 +273,8 @@ class SessionComponentsExportable(Protocol):
     retain that context or cache the components they bound to it, since the
     next call may install onto a different session. They should also avoid
     mutating the context they are handed — a registration made during binding
-    is not rolled back if a later extension fails.
+    is not rolled back if a later extension fails. See
+    :ref:`extension_bundles_transaction`.
 
     A bundle that also contributes a query planner implements
     :py:class:`SessionPlannerExportable` alongside this protocol.

From 593c29d5b551f6476dbe36acf10721d534d17059 Mon Sep 17 00:00:00 2001
From: Tim Saucer <timsaucer@gmail.com>
Date: Wed, 23 Sep 2026 07:53:04 -0400
Subject: [PATCH 2/6] refactor: narrow the extension protocol surface

The capsule-getter protocols are annotations, never arguments: a bundle
author constructs a SessionExtensionComponents but only ever names
SessionComponentsExportable in a type hint. Exporting the hints from the
package root made this one family the exception among sixteen such
protocols, every other one of which is reached through its defining
module.

Drop PhysicalOptimizerRuleExportable, QueryPlannerExportable,
SessionComponentsExportable, and SessionPlannerExportable from the root
(__all__ 58 -> 54), keeping SessionExtensionComponents, which is the one
name a bundle constructs. The three bundle protocols are new in 55.0.0,
so no import path is lost.

PhysicalOptimizerRuleExportable shipped in 54.0.0 from datafusion.context,
so its move is a break: context.py now imports it under TYPE_CHECKING only,
and the upgrade guide records the new path. Nothing else changes for a rule
author -- the protocol is structural and not runtime-checkable, and
add_physical_optimizer_rule is untouched.

Also removes three now-dead autoapi skip entries, repoints three doctests
that imported from the root, and fixes the add_physical_optimizer_rule
cross-reference, which stopped resolving once the class left context.py.

test_extension_protocols_are_exported_together asserted the premise this
reverses, so it goes. The five doctests in extensions.py already prove the
classes exist there, and SessionExtensionComponents' own docstring pins
the remaining root export.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 docs/source/conf.py                      |  3 ---
 docs/source/user-guide/upgrade-guides.md | 28 ++++++++++++++++++++++++
 python/datafusion/__init__.py            | 12 +---------
 python/datafusion/context.py             | 10 +++++++--
 python/datafusion/extensions.py          | 11 +++++-----
 python/tests/test_imports.py             | 18 ---------------
 6 files changed, 42 insertions(+), 40 deletions(-)

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 14cdc58d8..efb547652 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -122,10 +122,7 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool:  # noqa
         # Re-exports
         ("class", "datafusion.DataFrame"),
         ("class", "datafusion.SessionContext"),
-        ("class", "datafusion.QueryPlannerExportable"),
         ("class", "datafusion.SessionExtensionComponents"),
-        ("class", "datafusion.SessionComponentsExportable"),
-        ("class", "datafusion.SessionPlannerExportable"),
         ("module", "datafusion.common"),
         # Duplicate modules (skip module-level docs to avoid duplication)
         ("module", "datafusion.col"),
diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md
index d60653a21..0b8e94bb5 100644
--- a/docs/source/user-guide/upgrade-guides.md
+++ b/docs/source/user-guide/upgrade-guides.md
@@ -159,6 +159,34 @@ installed produces the same bytes as before, as do functions encoded by name.
 Regenerate any plan you serialized with an earlier release and stored for later
 use, if it was produced by a session with an extension codec installed.
 
+### Capsule-getter protocols moved to `datafusion.extensions`
+
+`PhysicalOptimizerRuleExportable` now lives in `datafusion.extensions`, next to
+the other protocols an extension library implements against. It was previously
+importable from `datafusion.context`, and that path is gone.
+
+```python
+from datafusion.context import PhysicalOptimizerRuleExportable  # before
+from datafusion.extensions import PhysicalOptimizerRuleExportable  # after
+```
+
+This affects type annotations only. The protocol is structural and not
+`@runtime_checkable`, so nothing imports it to call `isinstance`, and
+`SessionContext.add_physical_optimizer_rule` is unchanged — a rule object that
+worked before still works, whether or not its library names the protocol
+anywhere.
+
+The bundle protocols added in this release —  `QueryPlannerExportable`,
+`SessionComponentsExportable`, and `SessionPlannerExportable` — are reached the
+same way, through `datafusion.extensions` rather than the package root. They are
+new in 55.0.0, so no earlier import path existed. `SessionExtensionComponents`
+stays at the root, because a bundle constructs one rather than merely naming it:
+
+```python
+from datafusion import SessionExtensionComponents
+from datafusion.extensions import SessionComponentsExportable
+```
+
 ### `SessionContext.execute` renamed its second parameter
 
 The parameter is a single partition index, not a count, and is now named
diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py
index f5a2b3849..76f647796 100644
--- a/python/datafusion/__init__.py
+++ b/python/datafusion/__init__.py
@@ -92,13 +92,7 @@
 )
 from .dataframe_formatter import configure_formatter
 from .expr import Expr, WindowFrame
-from .extensions import (
-    PhysicalOptimizerRuleExportable,
-    QueryPlannerExportable,
-    SessionComponentsExportable,
-    SessionExtensionComponents,
-    SessionPlannerExportable,
-)
+from .extensions import SessionExtensionComponents
 from .io import read_avro, read_csv, read_json, read_parquet
 from .options import CsvReadOptions
 from .plan import (
@@ -140,19 +134,15 @@
     "MetricsSet",
     "ParquetColumnOptions",
     "ParquetWriterOptions",
-    "PhysicalOptimizerRuleExportable",
     "PhysicalPartitioning",
-    "QueryPlannerExportable",
     "RecordBatch",
     "RecordBatchStream",
     "RuntimeEnvBuilder",
     "SQLOptions",
     "ScalarUDF",
-    "SessionComponentsExportable",
     "SessionConfig",
     "SessionContext",
     "SessionExtensionComponents",
-    "SessionPlannerExportable",
     "Table",
     "TableFunction",
     "TableProviderFactory",
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index 1dedfa6dd..10a8f928a 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -70,7 +70,6 @@
 from datafusion.dataframe import DataFrame
 from datafusion.expr import sort_list_to_raw_sort_list
 from datafusion.extensions import (
-    PhysicalOptimizerRuleExportable,
     QueryPlannerExportable,
     SessionComponentsExportable,
     SessionExtensionComponents,
@@ -100,6 +99,12 @@
     from datafusion.catalog import CatalogProvider, Table
     from datafusion.common import DFSchema
     from datafusion.expr import Expr, SortKey
+
+    # Type-only on purpose. `datafusion.extensions` is the one home for the
+    # capsule-getter protocols; importing this at runtime would restore
+    # `datafusion.context.PhysicalOptimizerRuleExportable`, the 54.0.0 path
+    # that 55.0.0 drops.
+    from datafusion.extensions import PhysicalOptimizerRuleExportable
     from datafusion.plan import ExecutionPlan, LogicalPlan
     from datafusion.user_defined import (
         AggregateUDF,
@@ -1821,7 +1826,8 @@ def add_physical_optimizer_rule(
 
         Args:
             rule: Object exposing ``__datafusion_physical_optimizer_rule__``,
-                a :class:`PhysicalOptimizerRuleExportable`.
+                a
+                :py:class:`~datafusion.extensions.PhysicalOptimizerRuleExportable`.
 
         Examples:
             >>> from datafusion import SessionContext
diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py
index 196e25ad0..18fe62099 100644
--- a/python/datafusion/extensions.py
+++ b/python/datafusion/extensions.py
@@ -127,7 +127,7 @@ class QueryPlannerExportable(Protocol):
 
         The protocol itself is not runtime-checkable:
 
-        >>> from datafusion import QueryPlannerExportable
+        >>> from datafusion.extensions import QueryPlannerExportable
         >>> try:
         ...     isinstance(ctx, QueryPlannerExportable)
         ... except TypeError as e:
@@ -291,10 +291,8 @@ class SessionComponentsExportable(Protocol):
         The codecs this bundle contributes.
 
     Examples:
-        >>> from datafusion import (
-        ...     SessionExtensionComponents,
-        ...     SessionComponentsExportable,
-        ... )
+        >>> from datafusion import SessionExtensionComponents
+        >>> from datafusion.extensions import SessionComponentsExportable
         >>> class MyLibraryExtension:
         ...     def __datafusion_session_components__(self, ctx):
         ...         return SessionExtensionComponents()
@@ -355,7 +353,8 @@ class SessionPlannerExportable(Protocol):
         worth contrasting, because both plan queries successfully and only one
         of them is the no-op:
 
-        >>> from datafusion import SessionContext, SessionPlannerExportable
+        >>> from datafusion import SessionContext
+        >>> from datafusion.extensions import SessionPlannerExportable
         >>> class Contributes:
         ...     def __datafusion_session_planner__(self, ctx, fallback):
         ...         return None  # the no-op: session keeps its own planner
diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py
index 0e0a3965f..fea4cc91f 100644
--- a/python/tests/test_imports.py
+++ b/python/tests/test_imports.py
@@ -94,24 +94,6 @@ def test_datafusion_python_version():
     assert datafusion.__version__ is not None
 
 
-def test_extension_protocols_are_exported_together():
-    """The extension protocol family is reachable from the package root.
-
-    ``QueryPlannerExportable`` types the planner a
-    ``__datafusion_session_planner__`` hook returns, so a bundle author needs
-    it exactly as much as the other three; leaving it in the submodule made
-    one member of one family import differently from the rest.
-    """
-    for name in [
-        "QueryPlannerExportable",
-        "SessionExtensionComponents",
-        "SessionComponentsExportable",
-        "SessionPlannerExportable",
-    ]:
-        assert name in datafusion.__all__, f"{name} missing from datafusion.__all__"
-        assert getattr(datafusion, name) is getattr(datafusion.extensions, name)
-
-
 def test_class_module_is_datafusion():
     # context
     for klass in [

From 4ff228064382d9f7c1032f5adbd238df84b3d91f Mon Sep 17 00:00:00 2001
From: Tim Saucer <timsaucer@gmail.com>
Date: Wed, 23 Sep 2026 08:13:11 -0400
Subject: [PATCH 3/6] docs: state the commit-order invariant correctly

The ordering rule said step 4 "cannot raise part-way through", and the
comment in `with_extensions` said "everything above is allowed to raise;
this is not". Neither holds: `_install_extension_planner` runs
`ffi_query_planner_from_pycapsule` before it calls
`set_session_query_planner`, so the commit step has fallible work of its
own.

The guarantee survives, because that import happens before the write. But
the passage is written as a rule for whoever adds the next component kind,
and as phrased it asks them to preserve a property the code does not have.
Restate it as what actually holds: every fallible operation, including the
ones inside the commit, completes before the first write.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 .../source/contributor-guide/ffi-internals.md | 33 +++++++++++--------
 python/datafusion/context.py                  | 11 ++++---
 2 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md
index 7fdfd6a6e..be63cda69 100644
--- a/docs/source/contributor-guide/ffi-internals.md
+++ b/docs/source/contributor-guide/ffi-internals.md
@@ -120,7 +120,7 @@ was. Keeping that promise is an ordering constraint on the implementation, not
 a property of any one step, because the planner is bound on the shared
 `SessionState` rather than on the returned handle.
 
-A call therefore splits into a part that may fail and a part that may not:
+A call therefore splits into four steps, of which only the last writes:
 
 1. **Collect.** Every `__datafusion_session_components__` runs and its codecs
    are gathered. Nothing is installed yet, so a hook that raises here has
@@ -130,18 +130,25 @@ A call therefore splits into a part that may fail and a part that may not:
    to the session even though it can fail on a bad capsule or a duplicate id.
 3. **Resolve.** Every `__datafusion_session_planner__` runs, in argument order,
    against the completed chains, and each supplied planner is exported to a
-   capsule. Everything that can raise has raised by the end of this step.
-4. **Commit.** The accumulated planner is bound, in a single `SessionState`
-   rebuild. The bind is skipped entirely when the call installed nothing, so an
-   empty call does not drag a planner sitting on another handle's codecs onto
-   this one's.
-
-Only step 4 touches the session. This is a rule for the next field added to
-`SessionExtensionComponents`, not only a description of the current code: a new
-kind of component must do its fallible work — importing a capsule, resolving a
-name — in step 3, so that step 4 cannot raise part-way through.
-
-There is nothing to roll back to if it does. The returned handle shares one
+   capsule. Every hook that can raise has run by the end of this step.
+4. **Commit.** The accumulated planner is re-imported from its capsule and
+   bound, in a single `SessionState` rebuild. The bind is skipped entirely when
+   the call installed nothing, so an empty call does not drag a planner sitting
+   on another handle's codecs onto this one's.
+
+Only step 4 touches the session, and it is not itself infallible:
+`_install_extension_planner` runs `ffi_query_planner_from_pycapsule` before it
+calls `set_session_query_planner`, which cannot fail. The property that keeps
+the promise is therefore about order, not about any step being incapable of
+raising — every fallible operation, including the ones inside the commit,
+completes before the first write.
+
+That is the rule for the next field added to `SessionExtensionComponents`, not
+only a description of the current code: a new kind of component must do its
+fallible work — importing a capsule, resolving a name — before anything is
+written, so no failure can leave the session half-updated.
+
+There would be nothing to roll back to if one did. The returned handle shares one
 session with the receiver, so the damage is visible from every other handle;
 and undoing a registration is not the same as restoring what it displaced,
 because deregistering a function that shadowed a built-in removes the built-in
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index 10a8f928a..aa002d6c3 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -2028,10 +2028,13 @@ def with_extensions(
                 continue
             planner = new.ctx._export_query_planner(supplied)
 
-        # The commit step. Everything above is allowed to raise; this is not.
-        # See docs/source/contributor-guide/ffi-internals.md, "Why
-        # `with_extensions` commits last", for what a new component kind has to
-        # do to keep that true.
+        # The commit step, and the only one that writes to the session. It can
+        # still raise -- `_install_extension_planner` re-imports the capsule
+        # before binding it -- so what keeps the promise is the order, not any
+        # step being incapable of failing: every fallible operation finishes
+        # before the first write. See docs/source/contributor-guide/
+        # ffi-internals.md, "Why `with_extensions` commits last", for what a new
+        # component kind has to do to keep that true.
         #
         # Rebinding the session's planner is a side effect on state shared with
         # every other handle, so do not pay it for a call that installs nothing

From 5266516fc73b6231e626af49e1792163b4996e34 Mon Sep 17 00:00:00 2001
From: Tim Saucer <timsaucer@gmail.com>
Date: Wed, 23 Sep 2026 08:23:47 -0400
Subject: [PATCH 4/6] test: verify the physical optimizer rule docstring
 example

The `+SKIP` block in `PhysicalOptimizerRuleExportable` named
`test_ffi_physical_optimizer_rule_runs_during_planning` as the test that
runs it for real, but that test never reads the docstring. It is a
separately written test that happens to call the same two APIs, so it
catches a renamed method or module only by coincidence, and cannot see an
edit to the docstring at all.

Add the mirror the convention actually asks for, in the shape of
`test_with_extensions_docstring_example_still_runs`: parse the live
docstring, keep only the skipped statements, drop the skip, and run them.
Only `ctx` is supplied, because the skipped statements go on using the
context the runnable block above them opened.

Verified by mutation. Renaming the imported class in the docstring fails
with `NameError: name 'MyPhysicalOptimizerRule' is not defined`, and
deleting the block fails the `assert examples` guard rather than passing
vacuously.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 .../tests/_test_physical_optimizer_rule.py    | 44 +++++++++++++++++++
 python/datafusion/extensions.py               |  5 ++-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py b/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py
index 0c877d78e..9271c290a 100644
--- a/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py
+++ b/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py
@@ -17,8 +17,13 @@
 
 from __future__ import annotations
 
+import doctest
+import inspect
+import io
+
 import pyarrow as pa
 from datafusion import SessionContext
+from datafusion.extensions import PhysicalOptimizerRuleExportable
 from datafusion_ffi_example import MyPhysicalOptimizerRule
 
 
@@ -43,3 +48,42 @@ def test_ffi_physical_optimizer_rule_runs_during_planning():
         f"before={before} after={after}"
     )
     assert result[0].column(0).to_pylist() == [1, 2, 3]
+
+
+def test_physical_optimizer_rule_docstring_example_still_runs():
+    """Run the ``PhysicalOptimizerRuleExportable`` docstring example verbatim.
+
+    The example is marked ``+SKIP`` because the main suite has no built FFI
+    extension to import, which is exactly how such an example rots. Here the
+    statements are parsed out of the live docstring, the skip is dropped, and
+    each one is executed.
+
+    Only the ``+SKIP`` statements are taken. The docstring also carries a
+    runnable example above them, which the main suite already executes under
+    ``--doctest-modules``; running it again here would prove nothing.
+
+    Only ``ctx`` is supplied, because the skipped statements go on using the
+    context the runnable block opened. Everything else resolves for real: a
+    renamed class, a changed signature, or a wrong expected output in the
+    docstring fails here.
+    """
+    examples = [
+        example
+        for example in doctest.DocTestParser().get_examples(
+            inspect.getdoc(PhysicalOptimizerRuleExportable)
+        )
+        if example.options.pop(doctest.SKIP, False)
+    ]
+    assert examples, "docstring has no skipped examples to check"
+
+    test = doctest.DocTest(
+        examples,
+        {"ctx": SessionContext()},
+        "PhysicalOptimizerRuleExportable",
+        None,
+        None,
+        None,
+    )
+    output = io.StringIO()
+    results = doctest.DocTestRunner().run(test, out=output.write)
+    assert results.failed == 0, output.getvalue()
diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py
index 18fe62099..edb9a513e 100644
--- a/python/datafusion/extensions.py
+++ b/python/datafusion/extensions.py
@@ -90,8 +90,9 @@ class PhysicalOptimizerRuleExportable(Protocol):
             ...
         RuntimeError: "Invalid datafusion_physical_optimizer_rule...
 
-        Real usage. Skipped here (needs a built extension library); run for
-        real by ``test_ffi_physical_optimizer_rule_runs_during_planning`` in
+        Real usage. Skipped here (needs a built extension library); parsed out
+        of this docstring and run for real by
+        ``test_physical_optimizer_rule_docstring_example_still_runs`` in
         ``datafusion-ffi-example``.
 
         >>> from datafusion_ffi_example import MyPhysicalOptimizerRule  # doctest: +SKIP

From e1cdb64ede2a9aeff042a76089b9ca1e80f71345 Mon Sep 17 00:00:00 2001
From: Tim Saucer <timsaucer@gmail.com>
Date: Thu, 24 Sep 2026 15:22:00 -0400
Subject: [PATCH 5/6] refactor: keep bundle protocols out of datafusion.context

datafusion.context imported QueryPlannerExportable,
SessionComponentsExportable and SessionPlannerExportable at runtime,
so they were reachable as datafusion.context.* despite
datafusion.extensions being their one home. Move them under
TYPE_CHECKING and route the runtime isinstance checks through a
private _extensions module alias. The protocols are new in 55.0.0,
so no released import path is dropped.

Add a test pinning that all four capsule-getter protocols live only
in datafusion.extensions, and list PhysicalOptimizerRuleExportable
in llms.txt.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/source/llms.txt         |  2 +-
 python/datafusion/context.py | 33 +++++++++++++++++++--------------
 python/tests/test_imports.py | 18 ++++++++++++++++++
 3 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/docs/source/llms.txt b/docs/source/llms.txt
index 47b303ff6..597ce059d 100644
--- a/docs/source/llms.txt
+++ b/docs/source/llms.txt
@@ -31,7 +31,7 @@
 - [`datafusion.expr`](https://datafusion.apache.org/python/autoapi/datafusion/expr/index.html): expression tree nodes (`Expr`, `Window`, `WindowFrame`, `GroupingSet`).
 - [`datafusion.functions`](https://datafusion.apache.org/python/autoapi/datafusion/functions/index.html): 290+ scalar, aggregate, and window functions.
 - [`datafusion.context.SessionContext`](https://datafusion.apache.org/python/autoapi/datafusion/context/index.html): session entry point, data loading, SQL execution.
-- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionComponentsExportable`, `SessionPlannerExportable`, `QueryPlannerExportable` — the extension-bundle protocol.
+- [`datafusion.extensions`](https://datafusion.apache.org/python/autoapi/datafusion/extensions/index.html): `SessionExtensionComponents`, `SessionComponentsExportable`, `SessionPlannerExportable`, `QueryPlannerExportable`, `PhysicalOptimizerRuleExportable` — the extension-bundle and capsule-getter protocols.
 - [`datafusion.ipc`](https://datafusion.apache.org/python/autoapi/datafusion/ipc/index.html): worker and sender context slots for shipping expressions between processes.
 
 ## Examples
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index aa002d6c3..4043e357b 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -58,6 +58,7 @@
 
 import pyarrow as pa
 
+from datafusion import extensions as _extensions
 from datafusion.catalog import (
     Catalog,
     CatalogList,
@@ -69,12 +70,6 @@
 )
 from datafusion.dataframe import DataFrame
 from datafusion.expr import sort_list_to_raw_sort_list
-from datafusion.extensions import (
-    QueryPlannerExportable,
-    SessionComponentsExportable,
-    SessionExtensionComponents,
-    SessionPlannerExportable,
-)
 from datafusion.options import (
     DEFAULT_MAX_INFER_SCHEMA,
     CsvReadOptions,
@@ -101,10 +96,16 @@
     from datafusion.expr import Expr, SortKey
 
     # Type-only on purpose. `datafusion.extensions` is the one home for the
-    # capsule-getter protocols; importing this at runtime would restore
-    # `datafusion.context.PhysicalOptimizerRuleExportable`, the 54.0.0 path
-    # that 55.0.0 drops.
-    from datafusion.extensions import PhysicalOptimizerRuleExportable
+    # capsule-getter protocols; importing these at runtime would make them
+    # reachable as `datafusion.context.*`, and for
+    # `PhysicalOptimizerRuleExportable` would restore the 54.0.0 path that
+    # 55.0.0 drops. Runtime checks go through the private `_extensions` alias.
+    from datafusion.extensions import (
+        PhysicalOptimizerRuleExportable,
+        QueryPlannerExportable,
+        SessionComponentsExportable,
+        SessionPlannerExportable,
+    )
     from datafusion.plan import ExecutionPlan, LogicalPlan
     from datafusion.user_defined import (
         AggregateUDF,
@@ -1975,7 +1976,11 @@ def with_extensions(
         """
         for extension in extensions:
             if not isinstance(
-                extension, (SessionComponentsExportable, SessionPlannerExportable)
+                extension,
+                (
+                    _extensions.SessionComponentsExportable,
+                    _extensions.SessionPlannerExportable,
+                ),
             ):
                 msg = (
                     "Extension implements neither "
@@ -1991,10 +1996,10 @@ def with_extensions(
         logical_codecs: list[LogicalExtensionCodecExportable] = []
         physical_codecs: list[PhysicalExtensionCodecExportable] = []
         for extension in extensions:
-            if not isinstance(extension, SessionComponentsExportable):
+            if not isinstance(extension, _extensions.SessionComponentsExportable):
                 continue
             components = extension.__datafusion_session_components__(self)
-            if not isinstance(components, SessionExtensionComponents):
+            if not isinstance(components, _extensions.SessionExtensionComponents):
                 msg = (
                     "__datafusion_session_components__ must return "
                     "SessionExtensionComponents, got "
@@ -2016,7 +2021,7 @@ def with_extensions(
         # rather than wrapping the session's default in an FFI hop.
         planner: _PyCapsule | None = None
         for extension in extensions:
-            if not isinstance(extension, SessionPlannerExportable):
+            if not isinstance(extension, _extensions.SessionPlannerExportable):
                 continue
             fallback = (
                 planner
diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py
index fea4cc91f..c1c0b8243 100644
--- a/python/tests/test_imports.py
+++ b/python/tests/test_imports.py
@@ -192,6 +192,24 @@ def test_import_from_functions_submodule():
         from datafusion.functions import foobar  # noqa: F401
 
 
+@pytest.mark.parametrize(
+    "name",
+    [
+        "PhysicalOptimizerRuleExportable",
+        "QueryPlannerExportable",
+        "SessionComponentsExportable",
+        "SessionPlannerExportable",
+    ],
+)
+def test_extension_protocols_live_only_in_extensions(name):
+    import datafusion.context
+    import datafusion.extensions
+
+    assert getattr(datafusion.extensions, name).__module__ == "datafusion.extensions"
+    assert not hasattr(datafusion.context, name)
+    assert not hasattr(datafusion, name)
+
+
 def test_classes_are_inheritable():
     class MyExecContext(SessionContext):
         pass

From cc3ff0c33250fade3419ec02d419c6eb7a3c66d8 Mon Sep 17 00:00:00 2001
From: Tim Saucer <timsaucer@gmail.com>
Date: Thu, 24 Sep 2026 15:31:59 -0400
Subject: [PATCH 6/6] test: pin keyword-only construction; scope the rollback
 promise

Add a test that SessionExtensionComponents rejects positional
arguments, so dropping kw_only=True fails the suite. Every existing
caller passes keywords and would stay green without it.

Drop the absence assertions from the extension protocol import test.
Re-exporting a name is additive and breaks no caller, so asserting a
name is missing only adds friction for a later deliberate export.
Keep the positive check that each protocol imports from
datafusion.extensions.

The contributor guide promised that a raising bundle leaves the
session as it was without the carve-out for writes a hook makes to
the context it is handed. Add it with a ref to the bundles guide,
which is where the exception is explained.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/source/contributor-guide/ffi-internals.md |  3 ++-
 python/tests/test_context.py                   | 15 +++++++++++++++
 python/tests/test_imports.py                   |  5 +----
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md
index be63cda69..bbf6b65e7 100644
--- a/docs/source/contributor-guide/ffi-internals.md
+++ b/docs/source/contributor-guide/ffi-internals.md
@@ -116,7 +116,8 @@ prefer `with_extensions` — is documented at {ref}`planner_codec_rebinding`.
 ## Why `with_extensions` commits last
 
 `with_extensions` promises that a bundle which raises leaves the session as it
-was. Keeping that promise is an ordering constraint on the implementation, not
+was, apart from anything a hook writes to the context it is handed
+({ref}`extension_bundles_transaction`). Keeping that promise is an ordering constraint on the implementation, not
 a property of any one step, because the planner is bound on the shared
 `SessionState` rather than on the returned handle.
 
diff --git a/python/tests/test_context.py b/python/tests/test_context.py
index 9fbc4744f..274f2a9b1 100644
--- a/python/tests/test_context.py
+++ b/python/tests/test_context.py
@@ -1086,6 +1086,21 @@ def test_session_extension_components_rejects_a_single_codec(field):
         SessionExtensionComponents(**{field: codec})
 
 
+def test_session_extension_components_is_keyword_only():
+    codecs = (
+        _NamedCodec(
+            SessionContext().__datafusion_logical_extension_codec__(),
+            "my_library.logical",
+        ),
+    )
+
+    with pytest.raises(TypeError, match=r"positional argument"):
+        SessionExtensionComponents(codecs)
+
+    components = SessionExtensionComponents(logical_extension_codecs=codecs)
+    assert components.logical_extension_codecs == codecs
+
+
 def test_session_extension_components_rejects_a_string():
     """A str is iterable, so it needs refusing on its own.
 
diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py
index c1c0b8243..e1ef0d38c 100644
--- a/python/tests/test_imports.py
+++ b/python/tests/test_imports.py
@@ -201,13 +201,10 @@ def test_import_from_functions_submodule():
         "SessionPlannerExportable",
     ],
 )
-def test_extension_protocols_live_only_in_extensions(name):
-    import datafusion.context
+def test_extension_protocols_import_from_extensions(name):
     import datafusion.extensions
 
     assert getattr(datafusion.extensions, name).__module__ == "datafusion.extensions"
-    assert not hasattr(datafusion.context, name)
-    assert not hasattr(datafusion, name)
 
 
 def test_classes_are_inheritable():