chore: harden the extension API before it ships (#1759)

* 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>

* 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: 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>

* 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>

* 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>

* 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>

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: timsaucer

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

docs/source/contributor-guide/ffi-internals.md

@@ -111,6 +111,54 @@ 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, 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.
+
+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
+   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. 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
+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

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

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

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()

python/datafusion/__init__.py

@@ -92,12 +92,7 @@
 )
 from .dataframe_formatter import configure_formatter
 from .expr import Expr, WindowFrame
-from .extensions import (
-    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,17 +135,14 @@
     "ParquetColumnOptions",
     "ParquetWriterOptions",
     "PhysicalPartitioning",
-    "QueryPlannerExportable",
     "RecordBatch",
     "RecordBatchStream",
     "RuntimeEnvBuilder",
     "SQLOptions",
     "ScalarUDF",
-    "SessionComponentsExportable",
     "SessionConfig",
     "SessionContext",
     "SessionExtensionComponents",
-    "SessionPlannerExportable",
     "Table",
     "TableFunction",
     "TableProviderFactory",

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,
@@ -99,6 +94,18 @@
     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 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,
@@ -151,16 +158,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."""
 
@@ -1830,7 +1827,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
@@ -1918,7 +1916,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`.
 
@@ -1977,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 "
@@ -1993,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 "
@@ -2018,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
@@ -2030,6 +2033,14 @@ def with_extensions(
                 continue
             planner = new.ctx._export_query_planner(supplied)
 
+        # 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
         # -- the same guard `with_python_udf_inlining` carries. With no codec