feat: let extension bundles declare physical optimizer rules

`SessionExtensionComponents.physical_optimizer_rules` completes the group of
components whose capsule getter takes no argument, so a library shipping a rule
alongside anything else no longer asks the caller for a separate
`add_physical_optimizer_rule` call.

Rules are the one kind with no collision rule: they accumulate rather than
replace, so two bundles contributing one each is the normal case and there is
nothing to refuse.

Installing them splits across the two new private primitives
`_resolve_extension_physical_optimizer_rules` and
`_install_extension_physical_optimizer_rules`, keeping the commit step
infallible: the capsules are imported during resolution, so a rule that fails
to import cannot leave the session with a planner already bound. All the rules
in one call go on in a single `SessionState` rebuild.
`add_physical_optimizer_rule` rebuilds per call, which for a bundle with
several would clone the whole state that many times and leave the earlier ones
installed if a later one failed. The session id is carried across the rebuild
for the same reason that method carries it.

`PhysicalOptimizerRuleExportable` moves from `datafusion.context` to
`datafusion.extensions`, alongside the rest of the protocol family, and is now
exported from the package root. It stays importable from `datafusion.context`.

`MyRuleExtension` in `datafusion-ffi-example` declares two rules, which is what
makes accumulation observable — each carries its own counter and both fire.

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

Author: timsaucer

crates/core/src/context.rs

@@ -44,6 +44,7 @@ use datafusion::execution::options::{ArrowReadOptions, ReadOptions};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::session_state::SessionStateBuilder;
 use datafusion::execution::{FunctionRegistry, TaskContextProvider};
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
 use datafusion::physical_plan::ExecutionPlanProperties;
 use datafusion::prelude::{
     AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions,
@@ -1762,6 +1763,70 @@ impl PySessionContext {
         slf.borrow().set_session_query_planner(planner);
         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.
+    ///
+    /// **Writes nothing.**
+    pub fn _resolve_extension_physical_optimizer_rules(
+        &self,
+        rules: Vec<Bound<'_, PyAny>>,
+    ) -> PyDataFusionResult<PyPhysicalOptimizerRules> {
+        let rules = rules
+            .iter()
+            .map(physical_optimizer_rule_from_pycapsule)
+            .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.
+///
+/// Opaque to Python, and deliberately not added to the module: it exists only
+/// to carry imported rules from the resolve step to the commit step, so the
+/// import can fail before anything is written. `with_extensions` is its only
+/// producer and its only consumer.
+#[pyclass(name = "PhysicalOptimizerRules", module = "datafusion._internal")]
+pub struct PyPhysicalOptimizerRules {
+    rules: Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>>,
 }
 
 impl PySessionContext {

docs/source/extension-guide/bundles.md

@@ -336,6 +336,10 @@ Shadowing a name the session *already* has is allowed and is not a collision.
 The registry holds every DataFusion built-in, and overriding built-ins by name
 is a supported thing to do — `enable_spark_functions` is built on it.
 
+Physical optimizer rules are exempt: they accumulate rather than replace, so
+two bundles contributing one each is the normal case and there is nothing to
+refuse. See {doc}`other-components`.
+
 Like every other derivation, the returned context is a handle on the *same*
 session as the receiver — see {ref}`extension_sessions`. Only the Python-side
 codec chains belong to the returned handle; the planner is installed on the

docs/source/extension-guide/other-components.md

@@ -47,6 +47,22 @@ fn __datafusion_physical_optimizer_rule__<'py>(
 }
 ```
 
+If your library ships a rule alongside anything else, declare it on your bundle
+as `physical_optimizer_rules` rather than asking the caller for a separate
+`add_physical_optimizer_rule` call:
+
+```python
+return SessionExtensionComponents(physical_optimizer_rules=(MyRule(),))
+```
+
+Rules are the one kind of component with **no collision rule at all**: they
+accumulate, so two libraries may each contribute one and neither has to know
+about the other. Every rule in a call installs in a single `SessionState`
+rebuild, where `add_physical_optimizer_rule` rebuilds once per call — which for
+a bundle contributing several would clone the whole state that many times, and
+would leave the earlier ones installed if a later one failed. See
+{ref}`extension_bundles_transaction`.
+
 ## Typed configuration
 
 **`__datafusion_extension_options__`** contributes typed configuration entries

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

@@ -22,7 +22,7 @@
 import pyarrow as pa
 import pytest
 from datafusion import SessionContext, SessionExtensionComponents
-from datafusion_ffi_example import MyFunctionExtension
+from datafusion_ffi_example import MyFunctionExtension, MyRuleExtension
 
 
 def _session():
@@ -112,6 +112,78 @@ def __datafusion_session_planner__(self, ctx, fallback) -> None:
         ctx.udf("my_custom_is_null")
 
 
+def _query(ctx):
+    batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], names=["a"])
+    ctx.register_record_batches("t", [[batch]])
+    return ctx.sql("SELECT a FROM t").collect()
+
+
+def test_declared_rules_all_fire():
+    """Rules accumulate, so both of a bundle's two rules run.
+
+    Nothing about installing the second displaces the first, which is what
+    makes rules different from a planner and why there is no collision to
+    refuse.
+    """
+    extension = MyRuleExtension()
+    ctx = SessionContext().with_extensions(extension)
+
+    assert _query(ctx)[0].column(0).to_pylist() == [1, 2, 3]
+    assert extension.first_calls() > 0
+    assert extension.second_calls() > 0
+
+
+def test_rules_install_without_changing_the_session_id():
+    """Installing rules rebuilds ``SessionState``; the id has to survive it.
+
+    A fresh id would leave ``session_id()`` disagreeing with every
+    ``TaskContext`` the session already handed out, which is exactly what a
+    codec's decode callbacks resolve against.
+    """
+    ctx = SessionContext()
+    before = ctx.session_id()
+    result = ctx.with_extensions(MyRuleExtension())
+
+    assert result.session_id() == before
+    assert ctx.session_id() == before
+
+
+def test_rules_and_functions_install_together():
+    """Two bundles, one contributing functions and one rules, in one call."""
+    rules = MyRuleExtension()
+    ctx = SessionContext().with_extensions(MyFunctionExtension(), rules)
+    batch = pa.RecordBatch.from_arrays([pa.array([1, 2, None])], names=["a"])
+    ctx.register_record_batches("t", [[batch]])
+
+    result = ctx.sql("SELECT my_custom_is_null(a) FROM t").collect()
+
+    assert result[0].column(0).to_pylist() == [False, False, True]
+    assert rules.first_calls() > 0
+
+
+def test_a_failure_leaves_no_rule_installed():
+    """The transaction covers rules, which write through a state rebuild.
+
+    A rule reaching the session before the failing hook would be invisible to
+    ``session_id()`` and to the function registry, so this asserts on the
+    counter instead: an installed rule fires on the next query.
+    """
+    rules = MyRuleExtension()
+    ctx = SessionContext()
+
+    class BoomPlanner:
+        def __datafusion_session_planner__(self, ctx, fallback) -> None:
+            msg = "boom"
+            raise RuntimeError(msg)
+
+    with pytest.raises(RuntimeError, match="boom"):
+        ctx.with_extensions(rules, BoomPlanner())
+
+    _query(ctx)
+    assert rules.first_calls() == 0
+    assert rules.second_calls() == 0
+
+
 def test_the_hook_returns_the_components_type():
     """The bundle builds a real dataclass, not a duck-typed stand-in.
 

examples/datafusion-ffi-example/src/extension.rs

@@ -19,6 +19,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyDict;
 
 use crate::aggregate_udf::MySumUDF;
+use crate::physical_optimizer::MyPhysicalOptimizerRule;
 use crate::scalar_udf::IsNullUDF;
 use crate::window_udf::MyRankUDF;
 
@@ -62,3 +63,61 @@ impl MyFunctionExtension {
         components.call((), Some(&kwargs))
     }
 }
+
+/// A bundle contributing two physical optimizer rules.
+///
+/// Two, because that is what makes accumulation observable: rules never
+/// collide the way function names do, so both of these install and both fire.
+/// Each carries its own counter, which is how a test tells them apart.
+#[pyclass(
+    from_py_object,
+    name = "MyRuleExtension",
+    module = "datafusion_ffi_example",
+    subclass
+)]
+#[derive(Debug, Clone, Default)]
+pub(crate) struct MyRuleExtension {
+    first: MyPhysicalOptimizerRule,
+    second: MyPhysicalOptimizerRule,
+}
+
+#[pymethods]
+impl MyRuleExtension {
+    #[new]
+    fn new() -> Self {
+        Self::default()
+    }
+
+    /// How many times the first declared rule has run.
+    fn first_calls(&self) -> usize {
+        self.first.optimize_calls()
+    }
+
+    /// How many times the second declared rule has run.
+    fn second_calls(&self) -> usize {
+        self.second.optimize_calls()
+    }
+
+    /// `ctx` is unused: a rule getter takes no argument, so there is nothing
+    /// session-scoped to bind.
+    fn __datafusion_session_components__<'py>(
+        &self,
+        py: Python<'py>,
+        ctx: Bound<'py, PyAny>,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let _ = ctx;
+
+        let components = py
+            .import("datafusion")?
+            .getattr("SessionExtensionComponents")?;
+        let kwargs = PyDict::new(py);
+        kwargs.set_item(
+            "physical_optimizer_rules",
+            (
+                Py::new(py, self.first.clone())?,
+                Py::new(py, self.second.clone())?,
+            ),
+        )?;
+        components.call((), Some(&kwargs))
+    }
+}

examples/datafusion-ffi-example/src/lib.rs

@@ -20,7 +20,7 @@ use pyo3::prelude::*;
 use crate::aggregate_udf::MySumUDF;
 use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList};
 use crate::config::MyConfig;
-use crate::extension::MyFunctionExtension;
+use crate::extension::{MyFunctionExtension, MyRuleExtension};
 use crate::logical_extension_codec::MyLogicalExtensionCodec;
 use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec};
 use crate::physical_extension_codec::MyPhysicalExtensionCodec;
@@ -66,5 +66,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyPhysicalExtensionCodec>()?;
     m.add_class::<MyPhysicalOptimizerRule>()?;
     m.add_class::<MyFunctionExtension>()?;
+    m.add_class::<MyRuleExtension>()?;
     Ok(())
 }

examples/datafusion-ffi-example/src/physical_optimizer.rs

@@ -72,11 +72,11 @@ pub(crate) struct MyPhysicalOptimizerRule {
 #[pymethods]
 impl MyPhysicalOptimizerRule {
     #[new]
-    fn new() -> Self {
+    pub(crate) fn new() -> Self {
         Self::default()
     }
 
-    fn optimize_calls(&self) -> usize {
+    pub(crate) fn optimize_calls(&self) -> usize {
         self.optimize_calls.load(Ordering::SeqCst)
     }
 

python/datafusion/__init__.py

@@ -93,6 +93,7 @@
 from .dataframe_formatter import configure_formatter
 from .expr import Expr, WindowFrame
 from .extensions import (
+    PhysicalOptimizerRuleExportable,
     QueryPlannerExportable,
     SessionComponentsExportable,
     SessionExtensionComponents,
@@ -139,6 +140,7 @@
     "MetricsSet",
     "ParquetColumnOptions",
     "ParquetWriterOptions",
+    "PhysicalOptimizerRuleExportable",
     "PhysicalPartitioning",
     "QueryPlannerExportable",
     "RecordBatch",

python/datafusion/context.py

@@ -70,6 +70,7 @@
 from datafusion.dataframe import DataFrame
 from datafusion.expr import sort_list_to_raw_sort_list
 from datafusion.extensions import (
+    PhysicalOptimizerRuleExportable,
     QueryPlannerExportable,
     SessionComponentsExportable,
     SessionExtensionComponents,
@@ -151,16 +152,6 @@ class TableProviderExportable(Protocol):
     def __datafusion_table_provider__(self, session: Any) -> object: ...  # noqa: D105
 
 
-class PhysicalOptimizerRuleExportable(Protocol):
-    """Type hint for object that has __datafusion_physical_optimizer_rule__ PyCapsule.
-
-    The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``,
-    typically produced by a separate compiled extension.
-    """
-
-    def __datafusion_physical_optimizer_rule__(self) -> object: ...  # noqa: D105
-
-
 def _collect_contributions(
     extensions: tuple[object, ...],
     ctx: SessionContext,
@@ -204,6 +195,7 @@ def _collect_contributions(
         "udfs": [],
         "udafs": [],
         "udwfs": [],
+        "physical_optimizer_rules": [],
     }
     for extension in extensions:
         if not isinstance(extension, SessionComponentsExportable):
@@ -2032,10 +2024,10 @@ def with_extensions(
 
         Nothing is written to the session until every hook has returned and
         every component has been validated, so a hook that raises leaves the
-        session as it was. Declared functions register after the planner is
-        bound, and are visible on every handle sharing this session. A hook
-        that *mutates* the context it is handed — registering a table, say — is
-        not rolled back, which is why bundle objects must be
+        session as it was. Declared functions and optimizer rules install after
+        the planner is bound, and are visible on every handle sharing this
+        session. A hook that *mutates* the context it is handed — registering a
+        table, say — is not rolled back, which is why bundle objects must be
         configuration-only.
 
         Shares its session with this context — see :py:class:`SessionContext`.
@@ -2147,6 +2139,11 @@ def with_extensions(
             _udwf,
             "window function",
         )
+        # Rules accumulate, so there is no name to check and nothing to refuse
+        # -- only the capsules to import while failing is still free.
+        resolved_rules = new.ctx._resolve_extension_physical_optimizer_rules(
+            [rule for _, rule in declared["physical_optimizer_rules"]]
+        )
 
         # Phase two: nest the planners, outermost last. Each hook runs against
         # `new`, which carries the final chains, so a planner captured here
@@ -2187,6 +2184,7 @@ def with_extensions(
             new.register_udaf(function)
         for function in resolved_udwfs:
             new.register_udwf(function)
+        new.ctx._install_extension_physical_optimizer_rules(resolved_rules)
         return new
 
     def table_provider(self, name: str) -> Table: