refactor: register physical optimizer rules via live add method

Drop the `physical_optimizer_rules` constructor argument on
`SessionContext` and replace it with `add_physical_optimizer_rule`,
matching the existing `register_*` shape on the same class. The new
method rebuilds the session state via `SessionStateBuilder::new_from_existing`
so previously registered tables, UDFs, and catalogs are preserved.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: timsaucer

crates/core/src/context.rs

@@ -376,12 +376,11 @@ pub struct PySessionContext {
 
 #[pymethods]
 impl PySessionContext {
-    #[pyo3(signature = (config=None, runtime=None, physical_optimizer_rules=None))]
+    #[pyo3(signature = (config=None, runtime=None))]
     #[new]
     pub fn new(
         config: Option<PySessionConfig>,
         runtime: Option<PyRuntimeEnvBuilder>,
-        physical_optimizer_rules: Option<Vec<Bound<'_, PyAny>>>,
     ) -> PyDataFusionResult<Self> {
         let config = if let Some(c) = config {
             c.config
@@ -394,15 +393,11 @@ impl PySessionContext {
             RuntimeEnvBuilder::default()
         };
         let runtime = Arc::new(runtime_env_builder.build()?);
-        let mut state_builder = SessionStateBuilder::new()
+        let session_state = SessionStateBuilder::new()
             .with_config(config)
             .with_runtime_env(runtime)
-            .with_default_features();
-        for rule in physical_optimizer_rules.unwrap_or_default() {
-            let rule = physical_optimizer_rule_from_pycapsule(&rule)?;
-            state_builder = state_builder.with_physical_optimizer_rule(rule);
-        }
-        let session_state = state_builder.build();
+            .with_default_features()
+            .build();
         let ctx = Arc::new(SessionContext::new_with_state(session_state));
         Ok(PySessionContext {
             ctx,
@@ -1151,6 +1146,17 @@ impl PySessionContext {
         self.ctx.remove_optimizer_rule(name)
     }
 
+    pub fn add_physical_optimizer_rule(&self, rule: Bound<'_, PyAny>) -> PyDataFusionResult<()> {
+        let rule = physical_optimizer_rule_from_pycapsule(&rule)?;
+        let state_ref = self.ctx.state_ref();
+        let mut guard = state_ref.write();
+        let new_state = SessionStateBuilder::new_from_existing(guard.clone())
+            .with_physical_optimizer_rule(rule)
+            .build();
+        *guard = new_state;
+        Ok(())
+    }
+
     pub fn table_provider(&self, name: &str, py: Python) -> PyResult<PyTable> {
         let provider = wait_for_future(py, self.ctx.table_provider(name))
             // Outer error: runtime/async failure

examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py

@@ -23,10 +23,11 @@
 
 
 def test_ffi_physical_optimizer_rule_runs_during_planning():
-    """A rule supplied via physical_optimizer_rules is invoked while the
+    """A rule added via add_physical_optimizer_rule is invoked while the
     physical plan is built, and the query still returns correct results."""
     rule = MyPhysicalOptimizerRule()
-    ctx = SessionContext(physical_optimizer_rules=[rule])
+    ctx = SessionContext()
+    ctx.add_physical_optimizer_rule(rule)
     batch = pa.RecordBatch.from_arrays(
         [pa.array([1, 2, 3])],
         names=["a"],

python/datafusion/context.py

@@ -534,7 +534,6 @@ def __init__(
         self,
         config: SessionConfig | None = None,
         runtime: RuntimeEnvBuilder | None = None,
-        physical_optimizer_rules: list[PhysicalOptimizerRuleExportable] | None = None,
     ) -> None:
         """Main interface for executing queries with DataFusion.
 
@@ -545,11 +544,6 @@ def __init__(
         Args:
             config: Session configuration options.
             runtime: Runtime configuration options.
-            physical_optimizer_rules: User-defined physical optimizer rules to
-                append to the default set, each a
-                :class:`PhysicalOptimizerRuleExportable`. There is no upstream
-                API to add physical rules to a live context, so these can only
-                be supplied at construction time.
 
         Example usage:
 
@@ -560,21 +554,11 @@ def __init__(
 
             ctx = SessionContext()
             df = ctx.read_csv("data.csv")
-
-        To register a physical optimizer rule supplied by a compiled
-        extension, pass it via ``physical_optimizer_rules``::
-
-            from datafusion import SessionContext
-            from my_extension import MyPhysicalOptimizerRule
-
-            ctx = SessionContext(
-                physical_optimizer_rules=[MyPhysicalOptimizerRule()]
-            )
         """
         config = config.config_internal if config is not None else None
         runtime = runtime.config_internal if runtime is not None else None
 
-        self.ctx = SessionContextInternal(config, runtime, physical_optimizer_rules)
+        self.ctx = SessionContextInternal(config, runtime)
 
     def __repr__(self) -> str:
         """Print a string representation of the Session Context."""
@@ -1404,6 +1388,30 @@ def remove_optimizer_rule(self, name: str) -> bool:
         """
         return self.ctx.remove_optimizer_rule(name)
 
+    def add_physical_optimizer_rule(
+        self, rule: PhysicalOptimizerRuleExportable
+    ) -> None:
+        """Append a user-defined physical optimizer rule to the session.
+
+        The rule is imported via its ``__datafusion_physical_optimizer_rule__``
+        PyCapsule, typically produced by a separate compiled extension. The
+        underlying :class:`SessionState` is rebuilt from its current state
+        with the new rule appended, so previously registered tables, UDFs,
+        and catalogs are preserved.
+
+        Args:
+            rule: Object exposing ``__datafusion_physical_optimizer_rule__``,
+                a :class:`PhysicalOptimizerRuleExportable`.
+
+        Examples:
+            >>> from datafusion import SessionContext
+            >>> ctx = SessionContext()
+            >>> from my_extension import MyPhysicalOptimizerRule  # doctest: +SKIP
+            >>> rule = MyPhysicalOptimizerRule()  # doctest: +SKIP
+            >>> ctx.add_physical_optimizer_rule(rule)  # doctest: +SKIP
+        """
+        self.ctx.add_physical_optimizer_rule(rule)
+
     def table_provider(self, name: str) -> Table:
         """Return the :py:class:`~datafusion.catalog.Table` for the given table name.