chore: harden the extension API before it ships

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