fix: name the bundle that declared something that is not a rule

A rule the importer refuses used to raise from Rust with nothing but the
capsule name. A caller who passed four bundles could not tell which one was at
fault, and the resolve step is the last place that is known.
`_resolve_declared_rules` looks for the getter in Python first, mirroring what
`_resolve_declared_functions` already does for a declared function:

    TypeError A declared optimizer rule must expose
    __datafusion_physical_optimizer_rule__, got <object ...> from
    <RuleExtension ...>

Scoped to match the sibling rather than to go past it. A getter that is present
but returns a non-capsule still falls through to the importer's `RuntimeError`,
exactly as a declared function does, and `with_extensions` now documents that
case instead of listing only the two errors it raises itself.

`MyRuleExtension`'s two rules append to a run log they share, so the order they
installed in is observable. The counters cannot show it: each rule has its own,
so they say how often a rule ran but not when.

`ffi-internals.md` describes the four-step commit order as a rule for the next
field added to `SessionExtensionComponents`. `physical_optimizer_rules` is that
field, so steps three and four name it rather than leaving the enumeration
stale on the commit that invoked it.

Why rules never collide is now argued once, in the extension guide. The
protocol docstring and the field docstring state it and link there, and the
docstring naming the test that runs its skipped example names the whole test.

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

Author: timsaucer

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

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

@@ -133,6 +133,21 @@ def test_declared_rules_all_fire():
     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.
 

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

@@ -15,6 +15,8 @@
 // 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};
 
@@ -68,24 +70,32 @@ impl MyFunctionExtension {
 ///
 /// 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.
+/// 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, Default)]
+#[derive(Debug, Clone)]
 pub(crate) struct MyRuleExtension {
     first: MyPhysicalOptimizerRule,
     second: MyPhysicalOptimizerRule,
+    run_log: Arc<Mutex<Vec<usize>>>,
 }
 
 #[pymethods]
 impl MyRuleExtension {
     #[new]
     fn new() -> Self {
-        Self::default()
+        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.
@@ -98,6 +108,12 @@ impl MyRuleExtension {
         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>(

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

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
 
 use datafusion::common::Result;
 use datafusion::common::config::ConfigOptions;
@@ -31,9 +31,15 @@ use pyo3::types::PyCapsule;
 /// shared counter each time it runs. Tests use the counter to prove that a
 /// session built with this rule actually routed physical planning through a
 /// user-supplied [`PhysicalOptimizerRule`] over FFI.
+///
+/// A rule declared alongside siblings also appends its label to a log they
+/// all share, which is what lets a test see the order they ran in. Counters
+/// alone cannot: each rule has its own, so they say how often but not when.
 #[derive(Debug)]
 struct CountingPhysicalOptimizerRule {
     optimize_calls: Arc<AtomicUsize>,
+    label: usize,
+    run_log: Option<Arc<Mutex<Vec<usize>>>>,
 }
 
 impl PhysicalOptimizerRule for CountingPhysicalOptimizerRule {
@@ -43,6 +49,9 @@ impl PhysicalOptimizerRule for CountingPhysicalOptimizerRule {
         _config: &ConfigOptions,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         self.optimize_calls.fetch_add(1, Ordering::SeqCst);
+        if let Some(run_log) = &self.run_log {
+            run_log.lock().expect("run log poisoned").push(self.label);
+        }
         Ok(plan)
     }
 
@@ -67,6 +76,21 @@ impl PhysicalOptimizerRule for CountingPhysicalOptimizerRule {
 #[derive(Debug, Default, Clone)]
 pub(crate) struct MyPhysicalOptimizerRule {
     optimize_calls: Arc<AtomicUsize>,
+    label: usize,
+    run_log: Option<Arc<Mutex<Vec<usize>>>>,
+}
+
+impl MyPhysicalOptimizerRule {
+    /// A rule that records where it ran relative to the siblings sharing
+    /// `run_log`. `label` is what it appends. Not exposed to Python: only a
+    /// bundle declaring several rules at once has siblings to order against.
+    pub(crate) fn with_run_log(label: usize, run_log: Arc<Mutex<Vec<usize>>>) -> Self {
+        Self {
+            optimize_calls: Arc::new(AtomicUsize::new(0)),
+            label,
+            run_log: Some(run_log),
+        }
+    }
 }
 
 #[pymethods]
@@ -87,6 +111,8 @@ impl MyPhysicalOptimizerRule {
         let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
             Arc::new(CountingPhysicalOptimizerRule {
                 optimize_calls: Arc::clone(&self.optimize_calls),
+                label: self.label,
+                run_log: self.run_log.clone(),
             });
 
         let runtime = get_tokio_runtime().handle().clone();

python/datafusion/context.py

@@ -356,6 +356,40 @@ def _resolve_declared_functions(
     return resolved
 
 
+def _resolve_declared_rules(
+    declared: list[tuple[int, object, Any]], resolve: Any
+) -> Any:
+    """Import the capsule of every physical optimizer rule an extension declared.
+
+    There is no name to check — rules accumulate — so unlike
+    :py:func:`_resolve_declared_functions` this only refuses a declaration that
+    cannot be a rule. The getter is looked for here rather than left to the
+    importer so that the error names the bundle that declared it; a caller who
+    passed four bundles cannot otherwise tell which one is at fault.
+
+    Args:
+        declared: ``(position, extension, rule)`` triples in declaration order.
+        resolve: The primitive that imports a list of rules at once, returning
+            an opaque object for the commit step.
+
+    Returns:
+        The imported rules, opaque, in declaration order.
+
+    Raises:
+        TypeError: If a declaration does not expose
+            ``__datafusion_physical_optimizer_rule__``.
+    """
+    for _, extension, rule in declared:
+        if not hasattr(rule, "__datafusion_physical_optimizer_rule__"):
+            msg = (
+                "A declared optimizer rule must expose "
+                f"__datafusion_physical_optimizer_rule__, got {rule!r} "
+                f"from {extension!r}"
+            )
+            raise TypeError(msg)
+    return resolve([rule for _, _, rule in declared])
+
+
 class SessionConfig:
     """Session configuration options."""
 
@@ -2137,13 +2171,17 @@ def with_extensions(
             TypeError: If an argument implements neither hook, if a hook
                 returns the wrong type, if a codec is contributed as a bare
                 ``PyCapsule`` rather than an object exposing the getter, or if
-                a declared function is neither a wrapper nor exposes its
-                capsule getter.
+                a declared function or optimizer rule does not expose its
+                capsule getter and is not already a wrapper.
             ValueError: If two codecs claim the same id, if two extensions
                 declare a function of one kind under the same name, or if a
                 getter returns a capsule of the wrong kind. See
                 :py:meth:`with_logical_extension_codec` for how ids are
                 assigned.
+            RuntimeError: If a getter is present but returns something that is
+                not a ``PyCapsule`` at all. The message comes from the importer
+                and does not name the bundle, because by then the declaration
+                has already been accepted as the right shape.
 
         Examples:
             The returned handle is a different object sharing one session, and
@@ -2210,8 +2248,9 @@ def with_extensions(
         ]
         # 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"]]
+        resolved_rules = _resolve_declared_rules(
+            declared["physical_optimizer_rules"],
+            new.ctx._resolve_extension_physical_optimizer_rules,
         )
 
         # Phase two: nest the planners, outermost last. Each hook runs against

python/datafusion/extensions.py

@@ -79,11 +79,11 @@ class PhysicalOptimizerRuleExportable(Protocol):
     argument**: a rule needs neither a codec nor a task-context provider, so
     there is nothing session-scoped to hand it.
 
-    Rules accumulate rather than replace, so several libraries may each
-    contribute one and none of them has to know about the others. Install one
-    with :py:meth:`~datafusion.context.SessionContext.add_physical_optimizer_rule`,
+    Rules accumulate rather than replace. Install one with
+    :py:meth:`~datafusion.context.SessionContext.add_physical_optimizer_rule`,
     or declare it on a bundle as
-    :py:attr:`SessionExtensionComponents.physical_optimizer_rules`.
+    :py:attr:`SessionExtensionComponents.physical_optimizer_rules` — see
+    :ref:`extension_other_hooks`.
 
     Examples:
         The getter is the whole protocol, and a capsule is what it must return
@@ -98,7 +98,7 @@ class PhysicalOptimizerRuleExportable(Protocol):
         RuntimeError: "Invalid datafusion_physical_optimizer_rule...
 
         Real usage. Skipped here (needs a built extension library); run for
-        real by ``test_ffi_physical_optimizer_rule`` in
+        real by ``test_ffi_physical_optimizer_rule_runs_during_planning`` in
         ``datafusion-ffi-example``.
 
         >>> from datafusion_ffi_example import MyPhysicalOptimizerRule  # doctest: +SKIP
@@ -299,9 +299,8 @@ class SessionExtensionComponents:
     """Physical optimizer rules to install on the session.
 
     Objects exposing ``__datafusion_physical_optimizer_rule__``. Unlike
-    functions these never collide: rules accumulate, so two extensions may each
-    contribute one without either having to know about the other. All the rules
-    in one call install together, in declaration order.
+    functions these never collide — they accumulate. All the rules in one call
+    install together, in declaration order. See :ref:`extension_other_hooks`.
     """
 
     def __post_init__(self) -> None:

python/tests/test_context.py

@@ -1678,20 +1678,20 @@ def test_session_extension_components_rejects_a_single_optimizer_rule():
         SessionExtensionComponents(physical_optimizer_rules=object())
 
 
-def test_with_extensions_rejects_a_rule_that_is_not_a_capsule(ctx):
-    """A rule that will not import is refused, and nothing is installed.
+def test_with_extensions_rejects_a_rule_that_is_not_a_rule(ctx):
+    """A declaration that is not a rule at all names the bundle that made it.
 
-    Importing the capsules is the only part of installing a rule that can
-    fail, so it happens during resolution. A failure here has to leave the
-    session alone even though the extension ahead of it declared a function
-    that was perfectly good.
+    Which of several bundles is at fault is the whole content of the message,
+    and the only place it is still known is here. A failure also has to leave
+    the session alone even though the extension ahead of it declared a
+    function that was perfectly good.
     """
 
     class RuleExtension:
         def __datafusion_session_components__(self, ctx):
             return SessionExtensionComponents(physical_optimizer_rules=(object(),))
 
-    with pytest.raises(RuntimeError, match="datafusion_physical_optimizer_rule"):
+    with pytest.raises(TypeError, match=r"got .* from .*RuleExtension"):
         ctx.with_extensions(
             _FunctionExtension(udfs=(_doubler(),)),
             RuleExtension(),
@@ -1701,13 +1701,36 @@ def __datafusion_session_components__(self, ctx):
         ctx.udf("double")
 
 
+def test_with_extensions_rejects_a_rule_whose_getter_returns_a_non_capsule(ctx):
+    """A rule shaped right but returning junk is refused by the importer.
+
+    The bundle is past the point where it can be named — it declared the right
+    shape — so this is the one rule failure that surfaces as a ``RuntimeError``
+    from the import rather than a ``TypeError`` from the resolve.
+    """
+
+    class NotACapsule:
+        def __datafusion_physical_optimizer_rule__(self):
+            return object()
+
+    class RuleExtension:
+        def __datafusion_session_components__(self, ctx):
+            return SessionExtensionComponents(physical_optimizer_rules=(NotACapsule(),))
+
+    with pytest.raises(RuntimeError, match="datafusion_physical_optimizer_rule"):
+        ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),)), RuleExtension())
+
+    with pytest.raises(KeyError):
+        ctx.udf("double")
+
+
 def test_with_extensions_declaring_no_rules_leaves_the_session_id(ctx):
-    """Installing rules rebuilds ``SessionState``; declaring none must not.
+    """Installing rules rebuilds ``SessionState``; the id has to survive it.
 
-    The rebuild mints a fresh session id unless it is carried over, and a
-    changed id would break every ``TaskContext`` the session has handed out.
-    Asserted for the empty case too, because that is the one where the rebuild
-    would be pure cost.
+    The rebuild mints a fresh id unless it is carried over, and a changed id
+    would break every ``TaskContext`` the session has handed out. Asserted for
+    a call declaring no rules as well, so the guarantee does not depend on
+    whether the rebuild was skipped.
     """
     before = ctx.session_id()
     result = ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),)))