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>

Author: timsaucer

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

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

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