Let extension bundles declare physical optimizer rules

[timsaucer]

Which issue does this PR close?

Part 2 of 4 toward #1676.

1. scalar, aggregate, and window functions
2. This PR -> physical optimizer rules
3. table and table functions
4. catalog providers

Rationale for this change

#1738 added the function fields and, with them, the rule that installing a component splits into a fallible resolve step and an infallible commit step. Physical optimizer rules are the other half of the group that rule was written for: like the three function getters, __datafusion_physical_optimizer_rule__ takes no argument, so resolution needs nothing from the session and has no ordering constraint. Landing them here finishes that group before the stack moves on to the components that do need the session.

There is a second reason that is independent of bundles. add_physical_optimizer_rule does a full SessionStateBuilder::new_from_existing(guard.clone()).build() per call. A library contributing three rules clones the entire session state three times, and a failure on the third leaves the first two installed. Batching removes both.

What changes are included in this PR?

SessionExtensionComponents.physical_optimizer_rules. Objects exposing __datafusion_physical_optimizer_rule__, installed in declaration order.

Rules never collide. This is the one component kind with no collision rule at all. Two functions claiming a name is a ValueError because a function registry has no fall-through; rules accumulate, so two libraries each contributing one is the normal case and there is nothing to refuse. extension-guide/other-components.md already said "rules accumulate where planners nest" — this makes the bundle path agree with it.

Two new private Rust primitives, split to keep the commit infallible. _resolve_extension_physical_optimizer_rules imports every capsule and writes nothing; _install_extension_physical_optimizer_rules applies them all in one SessionState rebuild and returns () rather than a Result, because by then there is nothing left that can fail. This is the contract #1738 stated at the commit boundary, and it matters concretely here: a rule that failed to import after the planner was bound would leave the session half-installed with nothing to roll back to, since the returned handle and the receiver are one session.

The imported rules travel from resolve to commit in PyPhysicalOptimizerRules, a pyclass deliberately not added to the module — with_extensions is its only producer and only consumer. Both primitives are added to the private-method allowlist in test_wrapper_coverage.py.

The session id is carried across the rebuild 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 — which is what a codec's decode callbacks resolve against.

PhysicalOptimizerRuleExportable moves to datafusion.extensions. It was the one member of the protocol family still living in datafusion.context, absent from datafusion.__all__, and with no example. It is now beside its siblings, exported from the package root, and documented with a runnable block showing the refusal a non-capsule gets plus a +SKIP block naming the test that runs the real thing. It remains importable from datafusion.context, so this is purely additive.

Real-FFI coverage. MyRuleExtension declares two MyPhysicalOptimizerRule instances, which is what makes accumulation observable — each carries its own call counter and both fire on the next query. The new tests cover that, the session id surviving the rebuild, functions and rules arriving from two different bundles in one call, and a failure after the hooks leaving neither rule installed. That last one asserts on the counters rather than on a registry, since an installed rule is invisible to session_id() and to udf().

Are there any user-facing changes?

One new optional field on SessionExtensionComponents, defaulting to (). PhysicalOptimizerRuleExportable is now importable from datafusion as well as datafusion.context. Nothing is removed or renamed, no hook signatures change, so there is no upgrade-guide entry and no api change label.

Worth naming: add_physical_optimizer_rule still rebuilds once per call. Only the bundle path batches. Changing the public method's behaviour is not needed for this issue and would be a separate judgement about whether anyone depends on the per-call rebuild.

Review notes

The infallible-commit split is the part to push on. If you think importing the capsules could just as well happen inside the commit call, the case against it is that _install_extension_planner runs first: a Result returned after that point is a partially-installed session, and the only reason it would ever be one is convenience. Making the commit method return () rather than PyDataFusionResult<()> is the enforcement — it cannot silently grow a failure path later.

PyPhysicalOptimizerRules is not exported. If you would rather it were a plain opaque capsule than a pyclass, say so; the pyclass was chosen because it carries Vec<Arc<dyn PhysicalOptimizerRule>> without a second unsafe boundary.

🤖 Generated with Claude Code