chore: harden the extension API before it ships - #1759
Merged
timsaucer merged 6 commits intoSep 24, 2026
Merged
Conversation
`SessionContext.with_extensions` and `SessionExtensionComponents` landed in apache#1679 and have not shipped in a release yet. The bundle stack (apache#1738-apache#1741) reshapes them substantially, and a release would freeze three surfaces in their current form. `PhysicalOptimizerRuleExportable` was defined in `datafusion.context` and not exported from the package root, so `datafusion.context` would become its canonical import path. Move it to `datafusion.extensions` beside the rest of the `*Exportable` family, re-export it from `datafusion.context` so the old path keeps working, and export it from the package root. The move brings it under `test_extension_api_has_a_doctest`, which drives off `extensions.__all__`, so it gains the example it was missing. `SessionExtensionComponents` was positionally constructible with two fields. The stack takes it to nine, three of them pair-shaped. Make construction keyword-only so every later field addition is additive; no call site in the repository constructed it positionally. This is a new convention rather than a backport, so it has to be applied forward to the stack as well. The ordering that makes `with_extensions` transactional was stated in three docstrings with no canonical home to point at. Record it under `ffi_internals_commit_order` in the contributor guide, and label the existing "Failure and rollback" section `extension_bundles_transaction`, matching the names the stack links to. No released behaviour changes, so no `api change` label and no upgrade-guide entry. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The capsule-getter protocols are annotations, never arguments: a bundle author constructs a SessionExtensionComponents but only ever names SessionComponentsExportable in a type hint. Exporting the hints from the package root made this one family the exception among sixteen such protocols, every other one of which is reached through its defining module. Drop PhysicalOptimizerRuleExportable, QueryPlannerExportable, SessionComponentsExportable, and SessionPlannerExportable from the root (__all__ 58 -> 54), keeping SessionExtensionComponents, which is the one name a bundle constructs. The three bundle protocols are new in 55.0.0, so no import path is lost. PhysicalOptimizerRuleExportable shipped in 54.0.0 from datafusion.context, so its move is a break: context.py now imports it under TYPE_CHECKING only, and the upgrade guide records the new path. Nothing else changes for a rule author -- the protocol is structural and not runtime-checkable, and add_physical_optimizer_rule is untouched. Also removes three now-dead autoapi skip entries, repoints three doctests that imported from the root, and fixes the add_physical_optimizer_rule cross-reference, which stopped resolving once the class left context.py. test_extension_protocols_are_exported_together asserted the premise this reverses, so it goes. The five doctests in extensions.py already prove the classes exist there, and SessionExtensionComponents' own docstring pins the remaining root export. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The ordering rule said step 4 "cannot raise part-way through", and the comment in `with_extensions` said "everything above is allowed to raise; this is not". Neither holds: `_install_extension_planner` runs `ffi_query_planner_from_pycapsule` before it calls `set_session_query_planner`, so the commit step has fallible work of its own. The guarantee survives, because that import happens before the write. But the passage is written as a rule for whoever adds the next component kind, and as phrased it asks them to preserve a property the code does not have. Restate it as what actually holds: every fallible operation, including the ones inside the commit, completes before the first write. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The `+SKIP` block in `PhysicalOptimizerRuleExportable` named `test_ffi_physical_optimizer_rule_runs_during_planning` as the test that runs it for real, but that test never reads the docstring. It is a separately written test that happens to call the same two APIs, so it catches a renamed method or module only by coincidence, and cannot see an edit to the docstring at all. Add the mirror the convention actually asks for, in the shape of `test_with_extensions_docstring_example_still_runs`: parse the live docstring, keep only the skipped statements, drop the skip, and run them. Only `ctx` is supplied, because the skipped statements go on using the context the runnable block above them opened. Verified by mutation. Renaming the imported class in the docstring fails with `NameError: name 'MyPhysicalOptimizerRule' is not defined`, and deleting the block fails the `assert examples` guard rather than passing vacuously. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Public API boundaries need regression coverage, and the transaction guarantee and agent-facing documentation remain inconsistent.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 2
Open (4)
What changed in this PR
Hardens the extension API before the 55.0.0 release.
Changes:
- Consolidates extension protocols under
datafusion.extensions. - Makes
SessionExtensionComponentskeyword-only and narrows root exports. - Documents transaction ordering and adds a mirrored FFI docstring test.
| File | Description |
|---|---|
python/tests/test_imports.py |
Removes the former root-export assertion. |
python/datafusion/extensions.py |
Adds the optimizer protocol and keyword-only components. |
python/datafusion/context.py |
Updates annotations and commit-order guidance. |
python/datafusion/__init__.py |
Narrows package-root exports. |
examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py |
Mirrors the skipped docstring example. |
docs/source/user-guide/upgrade-guides.md |
Documents the protocol import migration. |
docs/source/extension-guide/bundles.md |
Clarifies failure and rollback behavior. |
docs/source/contributor-guide/ffi-internals.md |
Documents commit ordering. |
docs/source/conf.py |
Removes obsolete AutoAPI re-export exclusions. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
datafusion.context imported QueryPlannerExportable, SessionComponentsExportable and SessionPlannerExportable at runtime, so they were reachable as datafusion.context.* despite datafusion.extensions being their one home. Move them under TYPE_CHECKING and route the runtime isinstance checks through a private _extensions module alias. The protocols are new in 55.0.0, so no released import path is dropped. Add a test pinning that all four capsule-getter protocols live only in datafusion.extensions, and list PhysicalOptimizerRuleExportable in llms.txt. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add a test that SessionExtensionComponents rejects positional arguments, so dropping kw_only=True fails the suite. Every existing caller passes keywords and would stay green without it. Drop the absence assertions from the extension protocol import test. Re-exporting a name is additive and breaks no caller, so asserting a name is missing only adds friction for a later deliberate export. Keep the positive check that each protocol imports from datafusion.extensions. The contributor guide promised that a raising bundle leaves the session as it was without the carve-out for writes a hook makes to the context it is handed. Add it with a ref to the bundles guide, which is where the exception is explained. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Related to #1676
Rationale for this change
We need to release
datafusion-python 55.0.0but the follow on work linked in #1676 will not be reviewed and ready to merge for a while. This PR brings in a smaller set of changes from those in order to have fewer breaking changes in the following PR merges.The extension surface is new in this release, so now is the only cheap moment to settle its shape. Anything that ships in 55.0.0 — a name at the package root, a constructor that accepts positional arguments — becomes a compatibility obligation that costs a deprecation cycle to undo later. Narrowing it now costs nothing.
What changes are included in this PR?
with_extensionscommits last and an extension-guide section on what the failure guarantee does and does not cover.SessionExtensionComponentskeyword-only, so a later release can add component kinds without changing what an existing call means.PhysicalOptimizerRuleExportableout ofdatafusion.contextand intodatafusion.extensions, alongside the other capsule-getter protocols.PhysicalOptimizerRuleExportable,QueryPlannerExportable,SessionComponentsExportable, andSessionPlannerExportableare no longer exported fromdatafusionitself (__all__goes from 58 names to 54). They are annotations, never arguments, and every other capsule-getter protocol in the package is already reached through its defining module.SessionExtensionComponentsstays at the root, because a bundle constructs one rather than merely naming it._install_extension_plannerrunsffi_query_planner_from_pycapsulebefore it callsset_session_query_planner. The guarantee survives, because that import happens before the write, but the passage is written as a rule for whoever adds the next component kind and as phrased asked them to preserve a property the code does not have. It now states what actually holds: every fallible operation, including the ones inside the commit, completes before the first write.+SKIPconvention asks for. ThePhysicalOptimizerRuleExportableexample named a test that never read the docstring, so it caught a renamed method only by coincidence and could not see a docstring edit at all.test_physical_optimizer_rule_docstring_example_still_runsparses the live docstring, drops the skip, and runs it.Are there any user-facing changes?
Yes, one breaking change, which is why this carries the
api changelabel.PhysicalOptimizerRuleExportableshipped in 54.0.0 asdatafusion.context.PhysicalOptimizerRuleExportableand that import path is gone.docs/source/user-guide/upgrade-guides.mdgains a section showing the before and after. The effect is limited to type annotations: the protocol is structural and not@runtime_checkable, so nothing imports it to callisinstance, andSessionContext.add_physical_optimizer_ruleis unchanged — a rule object that worked before still works whether or not its library names the protocol anywhere.Everything else here is new in 55.0.0 and so breaks no released path. The four names dropped from the package root are all still importable from
datafusion.extensions, and three of them have never been importable from anywhere else.