|
| 1 | +--- |
| 2 | +name: ffi-capsule-protocol |
| 3 | +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." |
| 4 | +argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", or omit to review the whole family)" |
| 5 | +--- |
| 6 | + |
| 7 | +<!--- |
| 8 | + Licensed to the Apache Software Foundation (ASF) under one |
| 9 | + or more contributor license agreements. See the NOTICE file |
| 10 | + distributed with this work for additional information |
| 11 | + regarding copyright ownership. The ASF licenses this file |
| 12 | + to you under the Apache License, Version 2.0 (the |
| 13 | + "License"); you may not use this file except in compliance |
| 14 | + with the License. You may obtain a copy of the License at |
| 15 | +
|
| 16 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 17 | +
|
| 18 | + Unless required by applicable law or agreed to in writing, |
| 19 | + software distributed under the License is distributed on an |
| 20 | + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 21 | + KIND, either express or implied. See the License for the |
| 22 | + specific language governing permissions and limitations |
| 23 | + under the License. |
| 24 | +--> |
| 25 | + |
| 26 | +# FFI Capsule Protocol |
| 27 | + |
| 28 | +`datafusion-python` shares Rust objects with extension libraries through |
| 29 | +PyCapsules. Every hook is a dunder method named `__datafusion_<thing>__` that |
| 30 | +returns a capsule wrapping an FFI-safe struct. They are **one protocol**, not a |
| 31 | +collection of unrelated methods, and they have a settled convention that has |
| 32 | +already been migrated once (see `docs/source/user-guide/upgrade-guides.md`, |
| 33 | +DataFusion 52.0.0 and 55.0.0). |
| 34 | + |
| 35 | +## Rule 1 — enumerate the family before you change a member |
| 36 | + |
| 37 | +Do this first, every time. It takes one command and it is the whole point of |
| 38 | +this skill: |
| 39 | + |
| 40 | +```bash |
| 41 | +grep -rn "__datafusion_[a-z_]*__" --include="*.rs" crates/ examples/*/src/ |
| 42 | +``` |
| 43 | + |
| 44 | +Compare the signature you are about to write against what the others already |
| 45 | +do. If yours is shaped differently, that is a finding about your design, not |
| 46 | +about theirs. |
| 47 | + |
| 48 | +## Rule 2 — a getter takes the session it is being installed on |
| 49 | + |
| 50 | +```rust |
| 51 | +fn __datafusion_physical_extension_codec__<'py>( |
| 52 | + &self, |
| 53 | + py: Python<'py>, |
| 54 | + session: Bound<'py, PyAny>, |
| 55 | +) -> PyResult<Bound<'py, PyCapsule>> { ... } |
| 56 | +``` |
| 57 | + |
| 58 | +The host calls the getter and passes itself. That argument is how an extension |
| 59 | +library reaches things only the session has. |
| 60 | + |
| 61 | +`SessionContext` implements the same getters and ignores the argument, so a |
| 62 | +session satisfies the protocol too — `ctx.__datafusion_query_planner__()` and |
| 63 | +`ctx.__datafusion_query_planner__(ctx)` are both valid. |
| 64 | + |
| 65 | +## Rule 3 — never construct a `SessionContext` in an extension library |
| 66 | + |
| 67 | +The FFI constructors ask for things a library does not have: |
| 68 | + |
| 69 | +| Constructor | Wants | Take it from | |
| 70 | +|---|---|---| |
| 71 | +| `FFI_{Logical,Physical}ExtensionCodec::new` | `TaskContextProvider` | `ffi_task_context_provider_from_pycapsule(&session)` | |
| 72 | +| `FFI_TableProvider::new_with_ffi_codec` | logical codec | `ffi_logical_codec_from_pycapsule(session, None)` | |
| 73 | +| `FFI_QueryPlanner::new_with_ffi_codecs` | both codecs | `ffi_{logical,physical}_codec_from_pycapsule(session, None)` | |
| 74 | + |
| 75 | +`Arc::new(SessionContext::new())` is the wrong answer to all three, for two |
| 76 | +independent reasons: |
| 77 | + |
| 78 | +1. **It is the wrong registry.** Decode callbacks resolve names against |
| 79 | + whatever provider the codec carries. An empty session resolves nothing, so a |
| 80 | + function the host registered with `register_udf` is invisible to a node that |
| 81 | + references it by name. |
| 82 | +2. **It dangles.** `FFI_TaskContextProvider` downgrades its provider to a |
| 83 | + `Weak`. A context built inline in the getter is dropped before the capsule |
| 84 | + is ever used, and every callback then fails with `TaskContextProvider went |
| 85 | + out of scope over FFI boundary`. |
| 86 | + |
| 87 | +Prefer the `*_with_ffi_codec(s)` constructors when they exist. They take |
| 88 | +prebuilt codecs that already carry the host's provider, so there is no provider |
| 89 | +parameter to get wrong. |
| 90 | + |
| 91 | +## Rule 4 — the helpers live in `crates/util/src/lib.rs` |
| 92 | + |
| 93 | +`ffi_logical_codec_from_pycapsule`, `ffi_physical_codec_from_pycapsule`, |
| 94 | +`ffi_query_planner_from_pycapsule`, `ffi_task_context_provider_from_pycapsule`, |
| 95 | +`table_provider_from_pycapsule`. Each takes the object and, where relevant, an |
| 96 | +`Option<&Bound<PyAny>>` session: |
| 97 | + |
| 98 | +- `Some(session)` — importing a *foreign* object; the getter needs the session. |
| 99 | +- `None` — the object already *is* a session and is being asked for what it |
| 100 | + holds. |
| 101 | + |
| 102 | +Adding a getter means adding a helper here, not hand-rolling capsule |
| 103 | +extraction at the call site. |
| 104 | + |
| 105 | +## Rule 5 — changing a getter's signature is a breaking change |
| 106 | + |
| 107 | +Extension libraries implement these methods. A signature change breaks every |
| 108 | +one of them, and the failure is a bare `TypeError` from a `call1`. So: |
| 109 | + |
| 110 | +- Add a section to `docs/source/user-guide/upgrade-guides.md` with before/after |
| 111 | + Rust, matching the 52.0.0 and 55.0.0 entries. |
| 112 | +- Add the `api change` label to the PR. |
| 113 | +- Map the `TypeError` to a diagnosable message. `call_capsule_getter` in |
| 114 | + `crates/util/src/lib.rs` already does this; reuse it. |
| 115 | +- Update `python/datafusion/context.py` and |
| 116 | + `python/datafusion/user_defined.py`, where the `Protocol` type hints for |
| 117 | + these methods live. |
| 118 | + |
| 119 | +Changing what a codec puts *on the wire* is equally breaking, and easier to |
| 120 | +miss because no signature moves and nothing fails to compile. Serialized plans |
| 121 | +outlive the process that wrote them, so the same checklist applies: upgrade |
| 122 | +guide, `api change` label, and a statement of exactly which sessions produce |
| 123 | +different bytes. |
| 124 | + |
| 125 | +## Rule 6 — a session keeps one `Arc<SessionContext>` for life |
| 126 | + |
| 127 | +`FFI_TaskContextProvider` holds its provider **weakly**, and every codec handed |
| 128 | +to a foreign object carries one. A registered catalog provider upgrades that |
| 129 | +handle on every `supports_filters_pushdown` and every `scan`. The handle is |
| 130 | +bound to an `Arc<SessionContext>` *allocation*, so anything that replaces the |
| 131 | +allocation orphans every handle bound to the old one: |
| 132 | +`TaskContextProvider went out of scope over FFI boundary`. |
| 133 | + |
| 134 | +So mutate `SessionState` in place — `*self.ctx.state_ref().write() = ...`, the |
| 135 | +way `add_physical_optimizer_rule` and `set_session_query_planner` both do — |
| 136 | +rather than deriving a replacement `SessionContext`. Carry the session id |
| 137 | +across the rewrite; `SessionStateBuilder::new_from_existing` drops it and |
| 138 | +`build` mints a fresh one, which desyncs `session_id()` from every |
| 139 | +`TaskContext` the session hands out. |
| 140 | + |
| 141 | +Do not try to repair it after the fact: |
| 142 | + |
| 143 | +- **You cannot rebind what you cannot reach.** A codec embedded in a registered |
| 144 | + `FFI_CatalogProvider`, and in every `FFI_SchemaProvider` and |
| 145 | + `FFI_TableProvider` minted from it, has no Python-side handle. |
| 146 | +- **A codec must not retain its session.** Codecs are routinely handed to a |
| 147 | + provider that is registered straight back into the session that built them, |
| 148 | + closing `SessionContext -> catalog -> FFI provider -> FFI codec -> |
| 149 | + SessionContext`. |
| 150 | + |
| 151 | +`test_registered_providers_survive_a_planner_install` in |
| 152 | +`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` |
| 153 | +guards this. Its `WHERE` clause is load-bearing: filter pushdown upgrades the |
| 154 | +weak handle during logical optimization, before plan serialization could fail |
| 155 | +first for an unrelated reason. |
| 156 | + |
| 157 | +`SessionContext.enable_url_table` is the one method that mints a second |
| 158 | +allocation for a session. Its result must not outlive the receiver. |
| 159 | + |
| 160 | +## Rule 7 — installing a planner mutates the session, and says so |
| 161 | + |
| 162 | +`set_query_planner` returns `None`, matching `add_physical_optimizer_rule`. The |
| 163 | +query planner lives in `SessionState`, so it belongs to the session and not to |
| 164 | +a handle on it; every context sharing that session plans through it. Do not |
| 165 | +reintroduce a `with_query_planner` that pretends otherwise — the only way to |
| 166 | +give a handle its own planner is a fresh `Arc<SessionContext>`, which is what |
| 167 | +Rule 6 forbids. |
| 168 | + |
| 169 | +Installing a codec rebuilds the installed planner against it, and that rebuild |
| 170 | +reaches exactly one layer. `FFI_QueryPlanner::new_with_ffi_codecs` unwraps one |
| 171 | +`ForeignQueryPlanner`; a fallback that planner resolved at install time sits in |
| 172 | +its library's private data with no handle on this side, and cannot re-derive |
| 173 | +codecs itself because `FFI_QueryPlanner` holds them by value and `Session` |
| 174 | +exposes no accessor for the host's current ones. So do not promise that install |
| 175 | +order is free — for a layered planner it is not. The examples cannot show this: |
| 176 | +their fallback lives in the same cdylib as its wrapper, and `datafusion-ffi` |
| 177 | +short-circuits a same-library hop rather than serializing. A fix has to come |
| 178 | +from upstream; tracked in |
| 179 | +[apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762). |
| 180 | + |
| 181 | +The session's planner also tracks whichever handle wrote it last, so |
| 182 | +re-installing a planner on the original handle rebinds the session back to that |
| 183 | +handle's codecs. `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` |
| 184 | +pins that; changing it should be deliberate. |
| 185 | + |
| 186 | +## Where the truth is |
| 187 | + |
| 188 | +- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. |
| 189 | +- `docs/source/user-guide/upgrade-guides.md` — every past migration. |
| 190 | +- `crates/core/src/codec.rs` — the codec chain: the envelope, identity dispatch, |
| 191 | + and the two unframed cases from Rule 8. |
| 192 | +- `examples/datafusion-ffi-example/src/` — provider, catalog, function, codec |
| 193 | + getters, all in current form. `name_only_codec.rs` is the codec that encodes |
| 194 | + nothing. |
| 195 | +- `examples/datafusion-ffi-query-planner-example/src/planner.rs` — planner |
| 196 | + getter. |
| 197 | +- `examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` |
| 198 | + — `require_udf_on_decode` proves which session a decode callback resolves |
| 199 | + against. Extend these when touching the protocol. |
0 commit comments