Let extension bundles declare physical optimizer rules

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/contributor-guide/ffi-internals.md

@@ -128,9 +128,10 @@ A call therefore splits into a part that may fail and a part that may not:
    live on that handle rather than on the session, so this step writes nothing
    even though it can fail on a bad capsule or a duplicate id.
 3. **Resolve.** Every declared function is wrapped and every name is checked,
-   and every `__datafusion_session_planner__` runs against the completed
-   chains.
-4. **Commit.** The planner is bound and the functions are registered.
+   every declared physical optimizer rule has its capsule imported, and every
+   `__datafusion_session_planner__` runs against the completed chains.
+4. **Commit.** The planner is bound, the functions are registered, and the
+   optimizer rules are installed in a single `SessionState` rebuild.
 
 Only step 4 touches the session, and every step that can fail happens before
 it. This is a rule for the next field added to

docs/source/extension-guide/bundles.md

@@ -324,6 +324,10 @@ Two cases this does *not* catch:
   DataFusion function, and replacing one by name is a supported thing to do —
   `enable_spark_functions` works that way.
 
+Physical optimizer rules are exempt from all of this: 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`.
+
 Your caller cannot rename your function, so stay out of the way: prefix the
 names with something tied to your library.
 

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,93 @@ 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_declared_rules_run_in_declaration_order():
+    """The order a bundle lists its rules in is the order they install in.
+
+    Rules rewrite the plan one after another, so the order is part of what a
+    bundle declares. The counters cannot show it — each rule has its own — so
+    the two here append to a log they share.
+    """
+    extension = MyRuleExtension()
+    ctx = SessionContext().with_extensions(extension)
+
+    _query(ctx)
+
+    assert extension.run_order() == [0, 1]
+
+
+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

@@ -15,10 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::{Arc, Mutex};
+
 use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods};
 use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods};
 
 use crate::aggregate_udf::MySumUDF;
+use crate::physical_optimizer::MyPhysicalOptimizerRule;
 use crate::scalar_udf::IsNullUDF;
 use crate::window_udf::MyRankUDF;
 
@@ -62,3 +65,75 @@ 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, and
+/// both append to one run log, which is how a test sees the order they
+/// installed in.
+#[pyclass(
+    from_py_object,
+    name = "MyRuleExtension",
+    module = "datafusion_ffi_example",
+    subclass
+)]
+#[derive(Debug, Clone)]
+pub(crate) struct MyRuleExtension {
+    first: MyPhysicalOptimizerRule,
+    second: MyPhysicalOptimizerRule,
+    run_log: Arc<Mutex<Vec<usize>>>,
+}
+
+#[pymethods]
+impl MyRuleExtension {
+    #[new]
+    fn new() -> Self {
+        let run_log = Arc::new(Mutex::new(Vec::new()));
+        Self {
+            first: MyPhysicalOptimizerRule::with_run_log(0, Arc::clone(&run_log)),
+            second: MyPhysicalOptimizerRule::with_run_log(1, Arc::clone(&run_log)),
+            run_log,
+        }
+    }
+
+    /// 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()
+    }
+
+    /// The labels of the two declared rules, in the order they ran: `0` is the
+    /// first declared and `1` the second.
+    fn run_order(&self) -> Vec<usize> {
+        self.run_log.lock().expect("run log poisoned").clone()
+    }
+
+    /// `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(())
 }