-
Notifications
You must be signed in to change notification settings - Fork 166
Add FFI query planner support #1677
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
Open
Changes from all commits
Commits
Show all changes
33 commits
Select commit
Hold shift + click to select a range
9ba147c
Add FFI query planner support
timsaucer 47aea03
Add three-library FFI planner example
timsaucer 05ab5da
Update FFI query planner integration
timsaucer 27147f3
add rat
timsaucer 7982ee5
fix: install FFI test wheels from nested artifact paths
timsaucer c744579
Merge branch 'main' into feat/ffi-query-planner-core
timsaucer 81f68ed
refactor: address review of FFI query planner support
timsaucer 6ef7cf9
refactor: drop the runtime adapter and fix exported capsule lifetimes
timsaucer c6ba98f
add override for datafusion version to pre-release testing of upstrea…
timsaucer 752243b
remove unintentionally committed files
timsaucer e9dc22c
Empty commit to trigger CI
timsaucer 2169e39
fix: keep the example planner's exported task context alive
timsaucer 5bffd5e
test: cover which session a foreign codec decodes against
timsaucer b0800fe
feat: pass the session to the capsule getters that need it
timsaucer 68d4e6d
docs: record the FFI capsule protocol as a convention
timsaucer ffb3521
fix: put skill frontmatter before the license header
timsaucer 1e1d1aa
Update temporary DF version with corrections in FFI
timsaucer 9696956
Merge branch 'main' into feat/ffi-query-planner-core
timsaucer 282e3a5
feat: rebind foreign codecs when a planner install forks the session
timsaucer b7fbd34
fix: restore base64 0.23.1 in Cargo.lock
timsaucer 7ada7c6
fix: preserve session id across a planner fork
timsaucer 283745a
fix: preserve session id in add_physical_optimizer_rule
timsaucer 12b35fb
test: guard the session id a codec decodes against
timsaucer 9fa9a83
fix: stop reporting any getter TypeError as an outdated library
timsaucer 658e4ea
fix!: remove physical_codec_from_pycapsule
timsaucer a75c808
fix: check the FFI major version on every importer that can
timsaucer 8d138f9
docs: say where to relax the FFI version check
timsaucer 3eb8630
fix: route every capsule getter through call_capsule_getter
timsaucer 5d68f7b
docs: explain the dropped & in the codec migration snippet
timsaucer 5448e8d
docs: spell out the token registry lifecycle in the example codecs
timsaucer ca3264e
fix: resolve a planner fallback when it is installed, not constructed
timsaucer a2a2bb2
fix: accumulate planner observations instead of overwriting them
timsaucer 7652a33
Update rev for upstream datafusion to pre-release of 55.1.0
timsaucer File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,146 @@ | ||
| --- | ||
| name: ffi-capsule-protocol | ||
| description: "TRIGGER — read before adding, changing, or reviewing any __datafusion_*__ capsule getter, any FFI_* export that asks for a TaskContextProvider or an extension codec, or any code that calls FFI_QueryPlanner::new / FFI_TableProvider::new / FFI_{Logical,Physical}ExtensionCodec::new. These methods are one protocol with a settled convention. Do not design it fresh; do not construct a SessionContext inside an extension library." | ||
| argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", or omit to review the whole family)" | ||
| --- | ||
|
|
||
| <!--- | ||
| 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. | ||
| --> | ||
|
|
||
| # FFI Capsule Protocol | ||
|
|
||
| `datafusion-python` shares Rust objects with extension libraries through | ||
| PyCapsules. Every hook is a dunder method named `__datafusion_<thing>__` that | ||
| returns a capsule wrapping an FFI-safe struct. They are **one protocol**, not a | ||
| collection of unrelated methods, and they have a settled convention that has | ||
| already been migrated once (see `docs/source/user-guide/upgrade-guides.md`, | ||
| DataFusion 52.0.0 and 55.0.0). | ||
|
|
||
| ## Rule 1 — enumerate the family before you change a member | ||
|
|
||
| Do this first, every time. It takes one command and it is the whole point of | ||
| this skill: | ||
|
|
||
| ```bash | ||
| grep -rn "__datafusion_[a-z_]*__" --include="*.rs" crates/ examples/*/src/ | ||
| ``` | ||
|
|
||
| Compare the signature you are about to write against what the others already | ||
| do. If yours is shaped differently, that is a finding about your design, not | ||
| about theirs. | ||
|
|
||
| ## Rule 2 — a getter takes the session it is being installed on | ||
|
|
||
| ```rust | ||
| fn __datafusion_physical_extension_codec__<'py>( | ||
| &self, | ||
| py: Python<'py>, | ||
| session: Bound<'py, PyAny>, | ||
| ) -> PyResult<Bound<'py, PyCapsule>> { ... } | ||
| ``` | ||
|
|
||
| The host calls the getter and passes itself. That argument is how an extension | ||
| library reaches things only the session has. | ||
|
|
||
| `SessionContext` implements the same getters and ignores the argument, so a | ||
| session satisfies the protocol too — `ctx.__datafusion_query_planner__()` and | ||
| `ctx.__datafusion_query_planner__(ctx)` are both valid. | ||
|
|
||
| ## Rule 3 — never construct a `SessionContext` in an extension library | ||
|
|
||
| The FFI constructors ask for things a library does not have: | ||
|
|
||
| | Constructor | Wants | Take it from | | ||
| |---|---|---| | ||
| | `FFI_{Logical,Physical}ExtensionCodec::new` | `TaskContextProvider` | `ffi_task_context_provider_from_pycapsule(&session)` | | ||
| | `FFI_TableProvider::new_with_ffi_codec` | logical codec | `ffi_logical_codec_from_pycapsule(session, None)` | | ||
| | `FFI_QueryPlanner::new_with_ffi_codecs` | both codecs | `ffi_{logical,physical}_codec_from_pycapsule(session, None)` | | ||
|
|
||
| `Arc::new(SessionContext::new())` is the wrong answer to all three, for two | ||
| independent reasons: | ||
|
|
||
| 1. **It is the wrong registry.** Decode callbacks resolve names against | ||
| whatever provider the codec carries. An empty session resolves nothing, so a | ||
| function the host registered with `register_udf` is invisible to a node that | ||
| references it by name. | ||
| 2. **It dangles.** `FFI_TaskContextProvider` downgrades its provider to a | ||
| `Weak`. A context built inline in the getter is dropped before the capsule | ||
| is ever used, and every callback then fails with `TaskContextProvider went | ||
| out of scope over FFI boundary`. | ||
|
|
||
| Prefer the `*_with_ffi_codec(s)` constructors when they exist. They take | ||
| prebuilt codecs that already carry the host's provider, so there is no provider | ||
| parameter to get wrong. | ||
|
|
||
| ## Rule 4 — the helpers live in `crates/util/src/lib.rs` | ||
|
|
||
| `ffi_logical_codec_from_pycapsule`, `ffi_physical_codec_from_pycapsule`, | ||
| `ffi_query_planner_from_pycapsule`, `ffi_task_context_provider_from_pycapsule`, | ||
| `table_provider_from_pycapsule`. Each takes the object and, where relevant, an | ||
| `Option<&Bound<PyAny>>` session: | ||
|
|
||
| - `Some(session)` — importing a *foreign* object; the getter needs the session. | ||
| - `None` — the object already *is* a session and is being asked for what it | ||
| holds. | ||
|
|
||
| Adding a getter means adding a helper here, not hand-rolling capsule | ||
| extraction at the call site. | ||
|
|
||
| ## Rule 5 — changing a getter's signature is a breaking change | ||
|
|
||
| Extension libraries implement these methods. A signature change breaks every | ||
| one of them, and the failure is a bare `TypeError` from a `call1`. So: | ||
|
|
||
| - Add a section to `docs/source/user-guide/upgrade-guides.md` with before/after | ||
| Rust, matching the 52.0.0 and 55.0.0 entries. | ||
| - Add the `api change` label to the PR. | ||
| - Map the `TypeError` to a diagnosable message. `call_capsule_getter` in | ||
| `crates/util/src/lib.rs` already does this; reuse it. | ||
| - Update `python/datafusion/context.py` and | ||
| `python/datafusion/user_defined.py`, where the `Protocol` type hints for | ||
| these methods live. | ||
|
|
||
| ## Rule 6 — a fork must rebind the codecs it carries | ||
|
|
||
| Installing a foreign query planner **forks** the session | ||
| (`PySessionContext::derived_parts`), because installing one writes to | ||
| `SessionState` and the receiver must not be modified. A foreign codec holds an | ||
| `FFI_TaskContextProvider` pointing at the session it was installed on, so the | ||
| fork rebinds each one onto itself via | ||
| `PySessionContext::rebound_{logical,physical}_codec`. Skip that and decode | ||
| callbacks answer from the pre-fork registry. | ||
|
|
||
| Rebinding relies on `FFI_{Logical,Physical}ExtensionCodec::new` adopting the | ||
| provider on the already-foreign path, which needs DataFusion 55.1.0 or newer | ||
| (apache/datafusion#24722). It clones the handle before overwriting, so the | ||
| receiver keeps its own binding — assert that, not just that the fork works. | ||
| A codec this library owns round-trips unchanged, so the rebind is safe to apply | ||
| unconditionally. | ||
|
|
||
| ## Where the truth is | ||
|
|
||
| - `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. | ||
| - `docs/source/user-guide/upgrade-guides.md` — every past migration. | ||
| - `examples/datafusion-ffi-example/src/` — provider, catalog, function, codec | ||
| getters, all in current form. | ||
| - `examples/datafusion-ffi-query-planner-example/src/planner.rs` — planner | ||
| getter. | ||
| - `examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` | ||
| — `require_udf_on_decode` proves which session a decode callback resolves | ||
| against. Extend these when touching the protocol. |
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
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
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
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
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In order to prove that the 3 library approach works where we have different codecs and different execution plans provided, we are adding a second test library. This way we can make sure there is no accidental ability to reach into a foreign code block.