timsaucer opened a new pull request, #1739: URL: https://github.com/apache/datafusion-python/pull/1739
# Which issue does this PR close? Part 2 of 4 toward #1676. Stacked on #1738 — review that one first; this PR's diff against it is the second commit only. # 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](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
