feat: import FFI physical optimizer rules; drop Python logical rules

Replace the Python-defined OptimizerRule/AnalyzerRule approach with FFI-imported physical optimizer rules.

The Python logical-rule approach could observe plans but not transform them: there are no Python constructors for LogicalPlan node variants, so a rule could only return None or the input plan unchanged. The audience for custom rules also overlaps strongly with people who can write Rust.

DataFusion exposes no FFI bridge for the logical OptimizerRule/AnalyzerRule traits, but it does export FFI_PhysicalOptimizerRule for the physical PhysicalOptimizerRule trait. This commit imports those instead.

Changes:

* Remove crates/core/src/optimizer_rules.rs, python/datafusion/optimizer.py, python/tests/test_optimizer.py, and the SessionContext.add_optimizer_rule / add_analyzer_rule methods. remove_optimizer_rule is unchanged (pre-existing).
* New crates/core/src/physical_optimizer.rs reads a __datafusion_physical_optimizer_rule__ capsule and converts it via Arc<dyn PhysicalOptimizerRule>::from(&FFI_PhysicalOptimizerRule).
* SessionContext gains a physical_optimizer_rules constructor argument. Upstream offers no API to add physical rules to a live context, so they are appended to the builder at construction time only.
* The datafusion-ffi-example crate gains MyPhysicalOptimizerRule, a counter-backed rule used by _test_physical_optimizer_rule.py to prove the rule fires over FFI during physical planning.

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

Author: timsaucer

crates/core/src/context.rs

@@ -375,11 +375,12 @@ pub struct PySessionContext {
 
 #[pymethods]
 impl PySessionContext {
-    #[pyo3(signature = (config=None, runtime=None))]
+    #[pyo3(signature = (config=None, runtime=None, physical_optimizer_rules=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
@@ -392,11 +393,20 @@ impl PySessionContext {
             RuntimeEnvBuilder::default()
         };
         let runtime = Arc::new(runtime_env_builder.build()?);
-        let session_state = SessionStateBuilder::new()
+        let mut state_builder = SessionStateBuilder::new()
             .with_config(config)
             .with_runtime_env(runtime)
-            .with_default_features()
-            .build();
+            .with_default_features();
+        // DataFusion exposes no FFI bridge for the logical optimizer or
+        // analyzer, so only physical optimizer rules can be supplied from
+        // another library. They are appended after the default rules at
+        // construction time; there is no upstream API to add them to a live
+        // `SessionContext`.
+        for rule in physical_optimizer_rules.unwrap_or_default() {
+            let rule = crate::physical_optimizer::physical_optimizer_rule_from_pyobject(&rule)?;
+            state_builder = state_builder.with_physical_optimizer_rule(rule);
+        }
+        let session_state = state_builder.build();
         let ctx = Arc::new(SessionContext::new_with_state(session_state));
         Ok(PySessionContext {
             ctx,
@@ -1145,18 +1155,6 @@ impl PySessionContext {
         self.ctx.remove_optimizer_rule(name)
     }
 
-    pub fn add_optimizer_rule(&self, rule: Bound<'_, PyAny>) -> PyResult<()> {
-        let adapter = crate::optimizer_rules::build_optimizer_rule(rule)?;
-        self.ctx.add_optimizer_rule(adapter);
-        Ok(())
-    }
-
-    pub fn add_analyzer_rule(&self, rule: Bound<'_, PyAny>) -> PyResult<()> {
-        let adapter = crate::optimizer_rules::build_analyzer_rule(rule)?;
-        self.ctx.add_analyzer_rule(adapter);
-        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

crates/core/src/lib.rs

@@ -45,8 +45,8 @@ pub mod expr;
 #[allow(clippy::borrow_deref_ref)]
 mod functions;
 pub mod metrics;
-pub mod optimizer_rules;
 mod options;
+pub mod physical_optimizer;
 pub mod physical_plan;
 mod pyarrow_filter_expression;
 pub mod pyarrow_util;

crates/core/src/optimizer_rules.rs

@@ -0,0 +1,70 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Imports physical optimizer rules supplied by another library over FFI.
+//!
+//! DataFusion has no FFI bridge for the logical [`OptimizerRule`] /
+//! [`AnalyzerRule`] traits, but it does export
+//! [`FFI_PhysicalOptimizerRule`] for the physical
+//! [`PhysicalOptimizerRule`] trait. A producer crate (typically a separate
+//! compiled extension) exposes an object with a
+//! ``__datafusion_physical_optimizer_rule__`` method returning a
+//! :class:`PyCapsule` that wraps an [`FFI_PhysicalOptimizerRule`]. This
+//! module reads that capsule and converts it into an
+//! ``Arc<dyn PhysicalOptimizerRule>`` so it can be registered with a
+//! [`SessionContext`](datafusion::prelude::SessionContext) at construction
+//! time.
+//!
+//! [`OptimizerRule`]: datafusion::optimizer::optimizer::OptimizerRule
+//! [`AnalyzerRule`]: datafusion::optimizer::analyzer::AnalyzerRule
+
+use std::ptr::NonNull;
+use std::sync::Arc;
+
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion_ffi::physical_optimizer::FFI_PhysicalOptimizerRule;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+use crate::errors::{PyDataFusionError, PyDataFusionResult, to_datafusion_err};
+
+/// Convert a Python object exposing ``__datafusion_physical_optimizer_rule__``
+/// into an ``Arc<dyn PhysicalOptimizerRule>`` by reading its FFI capsule.
+pub(crate) fn physical_optimizer_rule_from_pyobject(
+    obj: &Bound<'_, PyAny>,
+) -> PyDataFusionResult<Arc<dyn PhysicalOptimizerRule + Send + Sync>> {
+    if !obj.hasattr("__datafusion_physical_optimizer_rule__")? {
+        return Err(PyDataFusionError::Common(
+            "Expected physical optimizer rule object to define \
+             __datafusion_physical_optimizer_rule__()"
+                .to_string(),
+        ));
+    }
+
+    let capsule = obj
+        .getattr("__datafusion_physical_optimizer_rule__")?
+        .call0()?;
+    let capsule = capsule.cast::<PyCapsule>().map_err(to_datafusion_err)?;
+    let data: NonNull<FFI_PhysicalOptimizerRule> = capsule
+        .pointer_checked(Some(c"datafusion_physical_optimizer_rule"))?
+        .cast();
+    let ffi_rule = unsafe { data.as_ref() };
+
+    Ok(Arc::<dyn PhysicalOptimizerRule + Send + Sync>::from(
+        ffi_rule,
+    ))
+}

crates/core/src/physical_optimizer.rs

@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import pyarrow as pa
+from datafusion import SessionContext
+from datafusion_ffi_example import MyPhysicalOptimizerRule
+
+
+def _setup_session_with_rule() -> tuple[SessionContext, MyPhysicalOptimizerRule]:
+    rule = MyPhysicalOptimizerRule()
+    ctx = SessionContext(physical_optimizer_rules=[rule])
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3])],
+        names=["a"],
+    )
+    ctx.register_record_batches("t", [[batch]])
+    return ctx, rule
+
+
+def test_ffi_physical_optimizer_rule_runs_during_planning():
+    """A rule supplied via physical_optimizer_rules is invoked while the
+    physical plan is built, and the query still returns correct results."""
+    ctx, rule = _setup_session_with_rule()
+
+    before = rule.optimize_calls()
+    result = ctx.sql("SELECT a FROM t").collect()
+    after = rule.optimize_calls()
+
+    assert after > before, (
+        f"Expected user FFI physical optimizer rule to fire, "
+        f"before={before} after={after}"
+    )
+    assert result[0].column(0).to_pylist() == [1, 2, 3]
+
+
+def test_ffi_physical_optimizer_rule_export():
+    """The rule object exposes the FFI capsule entry point."""
+    rule = MyPhysicalOptimizerRule()
+    capsule = rule.__datafusion_physical_optimizer_rule__()
+    assert capsule is not None

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

@@ -22,6 +22,7 @@ use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogP
 use crate::config::MyConfig;
 use crate::logical_extension_codec::MyLogicalExtensionCodec;
 use crate::physical_extension_codec::MyPhysicalExtensionCodec;
+use crate::physical_optimizer::MyPhysicalOptimizerRule;
 use crate::scalar_udf::IsNullUDF;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
@@ -33,6 +34,7 @@ pub(crate) mod catalog_provider;
 pub(crate) mod config;
 pub(crate) mod logical_extension_codec;
 pub(crate) mod physical_extension_codec;
+pub(crate) mod physical_optimizer;
 pub(crate) mod scalar_udf;
 pub(crate) mod table_function;
 pub(crate) mod table_provider;
@@ -55,5 +57,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<MyConfig>()?;
     m.add_class::<MyLogicalExtensionCodec>()?;
     m.add_class::<MyPhysicalExtensionCodec>()?;
+    m.add_class::<MyPhysicalOptimizerRule>()?;
     Ok(())
 }