refactor: commit the optimizer rules through _commit_extensions

The rules keep their separate resolve step -- importing the capsules
must stay fallible-and-early -- but the install half was a second
private pymethod whose only caller was the line after
`_commit_extensions`. Make it a parameter instead: the commit stays a
single call, the rules go on in the same one-rebuild block, and the
boundary does not grow an installer per component kind.

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

Author: timsaucer

crates/core/src/context.rs

@@ -1752,6 +1752,7 @@ impl PySessionContext {
     /// 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.
+    #[allow(clippy::too_many_arguments)]
     pub fn _commit_extensions<'py>(
         slf: &Bound<'py, Self>,
         extensions: Vec<Bound<'py, PyAny>>,
@@ -1760,6 +1761,7 @@ impl PySessionContext {
         udfs: Vec<PyScalarUDF>,
         udafs: Vec<PyAggregateUDF>,
         udwfs: Vec<PyWindowUDF>,
+        rules: PyRef<'_, PyPhysicalOptimizerRules>,
     ) -> PyDataFusionResult<()> {
         let py = slf.py();
         // Nest the planners, outermost last. `planner` stays `None` when no
@@ -1798,17 +1800,38 @@ impl PySessionContext {
         for udwf in udwfs {
             this.ctx.register_udwf(udwf.function);
         }
+        // Rules accumulate rather than replace, so unlike a planner there is
+        // no composition order to get right and no collision to refuse. All
+        // of them go on in **one** `SessionState` rebuild.
+        // [`Self::add_physical_optimizer_rule`] rebuilds per call, which for
+        // a bundle contributing several would clone the whole state that many
+        // times. Nothing here can fail: the capsules were imported by
+        // [`Self::_resolve_extension_physical_optimizer_rules`].
+        if !rules.rules.is_empty() {
+            let state_ref = this.ctx.state_ref();
+            let mut guard = state_ref.write();
+            // The session id has to be carried over for the same reason
+            // `add_physical_optimizer_rule` carries it: the builder mints a
+            // fresh one, and losing it leaves `session_id()` disagreeing with
+            // every `TaskContext` the session has already handed out.
+            let mut builder = SessionStateBuilder::new_from_existing(guard.clone())
+                .with_session_id(guard.session_id().to_string());
+            for rule in rules.rules.iter().cloned() {
+                builder = builder.with_physical_optimizer_rule(rule);
+            }
+            *guard = builder.build();
+        }
         Ok(())
     }
 
     /// Import the physical optimizer rules a `with_extensions` call declared.
     ///
     /// The fallible half of installing them, run while the call can still fail
     /// harmlessly. Every capsule is imported here so that
-    /// [`Self::_install_extension_physical_optimizer_rules`] has nothing left
-    /// that can raise — a rule that failed to import after the planner was
-    /// bound would leave the session half-installed, and there is no derived
-    /// context to roll back to.
+    /// [`Self::_commit_extensions`] has nothing left that can raise — a rule
+    /// that failed to import after the planner was bound would leave the
+    /// session half-installed, and there is no derived context to roll back
+    /// to.
     ///
     /// **Writes nothing.**
     pub fn _resolve_extension_physical_optimizer_rules(
@@ -1821,38 +1844,6 @@ impl PySessionContext {
             .collect::<Result<Vec<_>, _>>()?;
         Ok(PyPhysicalOptimizerRules { rules })
     }
-
-    /// Commit the physical optimizer rules for a `with_extensions` call.
-    ///
-    /// Rules accumulate rather than replace, so unlike a planner there is no
-    /// composition order to get right and no collision to refuse.
-    ///
-    /// All of them go on in **one** `SessionState` rebuild.
-    /// [`Self::add_physical_optimizer_rule`] rebuilds per call, which for a
-    /// bundle contributing several would clone the whole state that many times
-    /// and, worse, leave the earlier rules installed if a later one failed.
-    /// Nothing here can fail: the capsules were imported by
-    /// [`Self::_resolve_extension_physical_optimizer_rules`].
-    pub fn _install_extension_physical_optimizer_rules(
-        &self,
-        resolved: PyRef<'_, PyPhysicalOptimizerRules>,
-    ) {
-        if resolved.rules.is_empty() {
-            return;
-        }
-        let state_ref = self.ctx.state_ref();
-        let mut guard = state_ref.write();
-        // The session id has to be carried over for the same reason
-        // `add_physical_optimizer_rule` carries it: the builder mints a fresh
-        // one, and losing it leaves `session_id()` disagreeing with every
-        // `TaskContext` the session has already handed out.
-        let mut builder = SessionStateBuilder::new_from_existing(guard.clone())
-            .with_session_id(guard.session_id().to_string());
-        for rule in resolved.rules.iter().cloned() {
-            builder = builder.with_physical_optimizer_rule(rule);
-        }
-        *guard = builder.build();
-    }
 }
 
 /// Physical optimizer rules imported for a `with_extensions` call.

python/datafusion/context.py

@@ -2267,8 +2267,8 @@ def with_extensions(
             [function._udf for function in resolved["udfs"]],
             [function._udaf for function in resolved["udafs"]],
             [function._udwf for function in resolved["udwfs"]],
+            resolved_rules,
         )
-        new.ctx._install_extension_physical_optimizer_rules(resolved_rules)
         return new
 
     def table_provider(self, name: str) -> Table:

python/tests/test_wrapper_coverage.py

@@ -38,10 +38,11 @@
         # bundles declared.
         "_install_extension_codecs",
         "_commit_extensions",
-        # Physical optimizer rules, split so the capsule import happens while
-        # the call can still fail without leaving the session half-installed.
+        # Physical optimizer rules import their capsules in a separate step,
+        # so the call can still fail without leaving the session
+        # half-installed; the commit itself is a `_commit_extensions`
+        # parameter.
         "_resolve_extension_physical_optimizer_rules",
-        "_install_extension_physical_optimizer_rules",
     }
 )