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>

Author: timsaucer

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

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

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",

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

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,13 +58,49 @@
     )
 
 __all__ = [
+    "PhysicalOptimizerRuleExportable",
     "QueryPlannerExportable",
     "SessionComponentsExportable",
     "SessionExtensionComponents",
     "SessionPlannerExportable",
 ]
 
 
+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.