refactor: fold the planner hooks and the commit into one Rust call

The tail of `with_extensions` was three private pymethods and two
Python loops: nest the planner hooks, re-export each return as a
capsule, bind the planner, then walk a table of bound `register_*`
methods. Replace all of it with one `_commit_extensions` primitive
that runs the hooks and commits everything the bundles declared, so
the ordering contract lives in a single function next to the reasoning
it answers to, and the boundary drops `_export_query_planner` and
`_install_extension_planner`.

The Python side keeps everything that reads better in Python: protocol
dispatch, collision messages naming argument positions, and the
resolve step. Each function kind is now its own `_commit_extensions`
parameter, so `_FunctionKind` loses its `register` member -- a kind
that resolves but never commits cannot be written, because the call's
arity refuses it. The hook loop dispatches on the presence of
`__datafusion_session_planner__`, which is the same question the
runtime-checkable protocol asked.

Behaviour is pinned unchanged: no test assertion moves beyond the
private-method allowlist and the dropped `register` line in the
meta-test.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: timsaucer

crates/core/src/context.rs

@@ -1663,10 +1663,9 @@ impl PySessionContext {
     /// **Writes nothing.** The codec chains belong to the returned handle
     /// rather than to `SessionState`, so this phase is transactional for free:
     /// a codec that fails to import, or that collides with an installed id,
-    /// leaves the caller's context exactly as it was. Binding the planner is
-    /// the only step that touches the session, and it is deferred to
-    /// [`Self::_install_extension_planner`] so the planner hooks can run
-    /// against the final chains.
+    /// leaves the caller's context exactly as it was. Everything that touches
+    /// the session is deferred to [`Self::_commit_extensions`] so the planner
+    /// hooks can run against the final chains.
     ///
     /// Codecs must arrive as objects exposing the capsule getter, never as
     /// bare capsules — see [`resolve_bundle_codec_id`].
@@ -1717,49 +1716,87 @@ impl PySessionContext {
         })
     }
 
-    /// Re-export a planner a `__datafusion_session_planner__` hook returned as
-    /// a capsule, so the next hook in the chain receives one either way.
+    /// Run the planner hooks and commit a `with_extensions` call.
     ///
-    /// A hook may hand back an object exposing `__datafusion_query_planner__`
-    /// or a raw capsule; the next hook wraps whatever it is given and should
-    /// not have to branch on which. Importing here also surfaces a malformed
-    /// planner at the hook that produced it rather than at the final install.
-    /// Writes nothing.
-    pub fn _export_query_planner<'py>(
-        slf: &Bound<'py, Self>,
-        planner: Bound<'py, PyAny>,
-    ) -> PyDataFusionResult<Bound<'py, PyCapsule>> {
-        let ffi = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?;
-        Ok(create_query_planner_capsule(slf.py(), &ffi)?)
-    }
-
-    /// Commit the query planner for a `with_extensions` call.
+    /// The second phase, run on the handle carrying the completed chains —
+    /// `session` is that same handle as the Python-level wrapper, which is
+    /// what each `__datafusion_session_planner__` hook receives. The hooks
+    /// run first, **in argument order**, each handed the planner built so
+    /// far as a capsule; a hook may hand back an object exposing
+    /// `__datafusion_query_planner__` or a raw capsule, and each return is
+    /// imported here so a malformed planner surfaces at the hook that
+    /// produced it rather than at the install. Returning `None` contributes
+    /// no planner. All of that writes nothing, so a hook that raises leaves
+    /// the session exactly as it was.
     ///
-    /// The second phase, run once every codec is installed and every planner
-    /// hook has returned, so the planner is bound against the final chains.
-    /// This is the one call in `with_extensions` that writes to the session,
-    /// and it goes through this context's own `state_ref()`, so providers
-    /// bound to it stay valid.
+    /// Everything after the hooks is the commit, and none of it can fail: a
+    /// registration whose commit can fail belongs in the resolve step, split
+    /// into an import that returns a resolved object and an insert that
+    /// cannot raise. There is one session here, shared with the receiver, so
+    /// a failure part-way through would have nothing to roll back to. The
+    /// reasoning is in docs/source/contributor-guide/ffi-internals.md, under
+    /// "Why `with_extensions` commits last".
     ///
-    /// `None` means no bundle supplied a planner. That still rebuilds
-    /// whichever planner the session already holds against the new chains,
-    /// exactly as `with_logical_extension_codec` does, and writes nothing at
-    /// all if the session has no FFI planner to rebuild.
+    /// The planner is bound through this context's own `state_ref()`, so
+    /// providers bound to it stay valid. With no planner supplied the bind
+    /// still rebuilds whichever planner the session already holds against
+    /// the new chains, exactly as `with_logical_extension_codec` does —
+    /// unless `rebind_planner` is also false, meaning the call installed no
+    /// codec either. Then the bind is skipped entirely, the same way
+    /// [`Self::with_python_udf_inlining`] returns early for a no-op toggle:
+    /// there is nothing to rebind against, and the rebuild would drag a
+    /// planner sitting on another handle's codecs onto this one's.
     ///
-    /// The caller skips this step entirely when the call installed no codec
-    /// and no planner, the same way [`Self::with_python_udf_inlining`] returns
-    /// early for a no-op toggle: there is nothing to rebind against, and the
-    /// rebuild would drag a planner sitting on another handle's codecs onto
-    /// this one's.
-    #[pyo3(signature = (planner=None))]
-    pub fn _install_extension_planner<'py>(
+    /// The functions are registered *after* the planner hooks have run, so a
+    /// hook never sees this call's functions in the registry — the
+    /// registrations have no fall-through, and a name is free to shadow one
+    /// the session already had.
+    pub fn _commit_extensions<'py>(
         slf: &Bound<'py, Self>,
-        planner: Option<Bound<'py, PyAny>>,
+        extensions: Vec<Bound<'py, PyAny>>,
+        session: Bound<'py, PyAny>,
+        rebind_planner: bool,
+        udfs: Vec<PyScalarUDF>,
+        udafs: Vec<PyAggregateUDF>,
+        udwfs: Vec<PyWindowUDF>,
     ) -> PyDataFusionResult<()> {
-        let planner = planner
-            .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any())))
-            .transpose()?;
-        slf.borrow().set_session_query_planner(planner);
+        let py = slf.py();
+        // Nest the planners, outermost last. `planner` stays `None` when no
+        // bundle supplies one, which leaves an already-installed planner in
+        // place rather than wrapping the session's default in an FFI hop.
+        let mut planner: Option<FFI_QueryPlanner> = None;
+        for extension in &extensions {
+            if !extension.hasattr("__datafusion_session_planner__")? {
+                continue;
+            }
+            let fallback = match &planner {
+                Some(ffi) => create_query_planner_capsule(py, ffi)?,
+                None => slf.borrow().__datafusion_query_planner__(py, None)?,
+            };
+            let supplied =
+                extension.call_method1("__datafusion_session_planner__", (&session, fallback))?;
+            if supplied.is_none() {
+                continue;
+            }
+            planner = Some(ffi_query_planner_from_pycapsule(
+                &supplied,
+                Some(slf.as_any()),
+            )?);
+        }
+
+        if planner.is_some() || rebind_planner {
+            slf.borrow().set_session_query_planner(planner);
+        }
+        let this = slf.borrow();
+        for udf in udfs {
+            this.ctx.register_udf(udf.function);
+        }
+        for udaf in udafs {
+            this.ctx.register_udaf(udaf.function);
+        }
+        for udwf in udwfs {
+            this.ctx.register_udwf(udwf.function);
+        }
         Ok(())
     }
 }

python/datafusion/context.py

@@ -162,13 +162,15 @@ def __datafusion_physical_optimizer_rule__(self) -> object: ...  # noqa: D105
 
 
 class _FunctionKind(NamedTuple):
-    """How one kind of declared function is resolved and registered.
+    """How one kind of declared function is resolved.
 
     One row per function field on
     :py:class:`~datafusion.extensions.SessionExtensionComponents`, so adding a
     kind is adding a row rather than editing three places. The dataclass
     metadata says which fields are collections to normalize; this table says
-    which of them are functions and what to do with one.
+    which of them are functions and how to wrap one. Committing is not here:
+    each kind is its own ``_commit_extensions`` parameter, so a kind that
+    resolves but never commits cannot be written.
     ``test_every_component_field_has_an_installer`` pins the two together.
 
     The members naming a ``datafusion.user_defined`` object hold its name
@@ -192,9 +194,6 @@ class _FunctionKind(NamedTuple):
     label: str
     """What to call this sort of function in an error message."""
 
-    register: str
-    """:py:class:`SessionContext` method that commits one to the session."""
-
 
 _FUNCTION_KINDS = (
     _FunctionKind(
@@ -203,23 +202,20 @@ class _FunctionKind(NamedTuple):
         getter="__datafusion_scalar_udf__",
         factory="udf",
         label="scalar function",
-        register="register_udf",
     ),
     _FunctionKind(
         field="udafs",
         wrapper="AggregateUDF",
         getter="__datafusion_aggregate_udf__",
         factory="udaf",
         label="aggregate function",
-        register="register_udaf",
     ),
     _FunctionKind(
         field="udwfs",
         wrapper="WindowUDF",
         getter="__datafusion_window_udf__",
         factory="udwf",
         label="window function",
-        register="register_udwf",
     ),
 )
 """Every kind of function a bundle can declare, in registration order."""
@@ -2197,63 +2193,38 @@ def with_extensions(
         # Resolve every declared function to the wrapper that registers it, and
         # settle name collisions, while a failure still costs nothing. None of
         # these getters take an argument, so unlike a provider they do not care
-        # which handle they are resolved against. The bound `register_*` method
-        # is looked up here too, leaving the commit below nothing but calls.
+        # which handle they are resolved against.
         from datafusion import user_defined as _user_defined  # noqa: PLC0415
 
-        resolved: list[tuple[Any, list[Any]]] = [
-            (
-                getattr(new, kind.register),
-                _resolve_declared_functions(
-                    declared[kind.field],
-                    getattr(_user_defined, kind.wrapper),
-                    kind.getter,
-                    getattr(_user_defined, kind.factory),
-                    kind.label,
-                ),
+        resolved: dict[str, list[Any]] = {
+            kind.field: _resolve_declared_functions(
+                declared[kind.field],
+                getattr(_user_defined, kind.wrapper),
+                kind.getter,
+                getattr(_user_defined, kind.factory),
+                kind.label,
             )
             for kind in _FUNCTION_KINDS
-        ]
-
-        # Phase two: nest the planners, outermost last. Each hook runs against
-        # `new`, which carries the final chains, so a planner captured here
-        # never sees a partial codec set. `planner` stays None when no bundle
-        # supplies one, which leaves an already-installed planner in place
-        # rather than wrapping the session's default in an FFI hop.
-        planner: _PyCapsule | None = None
-        for extension in extensions:
-            if not isinstance(extension, SessionPlannerExportable):
-                continue
-            fallback = (
-                planner
-                if planner is not None
-                else new.ctx.__datafusion_query_planner__()
-            )
-            supplied = extension.__datafusion_session_planner__(new, fallback)
-            if supplied is None:
-                continue
-            planner = new.ctx._export_query_planner(supplied)
-
-        # 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
-        # installed the chains the planner would be rebuilt against are the ones
-        # it already holds, so the rebuild is unobservable except in the one case
-        # where it does harm: a planner sitting on some *other* handle's codecs
-        # gets dragged onto this handle's, silently undoing that install.
-
-        # Commit. Everything below this line must be infallible. A registration
-        # whose commit can fail belongs above, split into an import step that
-        # returns a resolved object and an insert step that cannot raise --
-        # there is one session here, shared with the receiver, so a failure
-        # part-way through has nothing to roll back to. The reasoning is in
-        # docs/source/contributor-guide/ffi-internals.md, under "Why
-        # `with_extensions` commits last".
-        if planner is not None or logical_codecs or physical_codecs:
-            new.ctx._install_extension_planner(planner)
-        for register, functions in resolved:
-            for function in functions:
-                register(function)
+        }
+
+        # Phase two: run the planner hooks and commit, in one call. Each hook
+        # runs against `new`, which carries the final chains, so a planner
+        # captured there never sees a partial codec set. The hook loop, the
+        # ordering of the commit, and the guard that skips the planner rebind
+        # for a call that installs nothing all live on the Rust side -- see
+        # `_commit_extensions` and docs/source/contributor-guide/
+        # ffi-internals.md, under "Why `with_extensions` commits last". A new
+        # component field must be resolved above and given its own
+        # `_commit_extensions` parameter; the call's arity is what keeps a
+        # declared component from being quietly dropped.
+        new.ctx._commit_extensions(
+            list(extensions),
+            new,
+            bool(logical_codecs or physical_codecs),
+            [function._udf for function in resolved["udfs"]],
+            [function._udaf for function in resolved["udafs"]],
+            [function._udwf for function in resolved["udwfs"]],
+        )
         return new
 
     def table_provider(self, name: str) -> Table:

python/tests/test_context.py

@@ -1714,7 +1714,6 @@ def test_every_function_kind_names_something_real():
     for kind in _FUNCTION_KINDS:
         assert isinstance(getattr(user_defined, kind.wrapper), type)
         assert callable(getattr(user_defined, kind.factory))
-        assert callable(getattr(SessionContext, kind.register))
         assert kind.field in {spec.name for spec in fields(SessionExtensionComponents)}
 
 

python/tests/test_wrapper_coverage.py

@@ -33,12 +33,11 @@
 # gap in coverage.
 PRIVATE_SUPPORT_METHODS = frozenset(
     {
-        # The three steps of SessionContext.with_extensions: install the
-        # codecs, re-export each planner hook's return value as a capsule,
-        # commit the planner.
+        # The two steps of SessionContext.with_extensions: install the
+        # codecs, then run the planner hooks and commit everything the
+        # bundles declared.
         "_install_extension_codecs",
-        "_export_query_planner",
-        "_install_extension_planner",
+        "_commit_extensions",
     }
 )