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>
diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 7775a4a..554fc65 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs
@@ -1752,6 +1752,7 @@ /// 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 @@ 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,6 +1800,27 @@ 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(()) } @@ -1805,10 +1828,10 @@ /// /// 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 @@ .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.
diff --git a/python/datafusion/context.py b/python/datafusion/context.py index b6f2d69..3f2c0f8 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py
@@ -2267,8 +2267,8 @@ [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:
diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index 4276ffe..50a53a7 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/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", } )