diff --git a/.ai/skills/audit-skill-md/SKILL.md b/.ai/skills/audit-skill-md/SKILL.md index ba5255a59..5f3177b3d 100644 --- a/.ai/skills/audit-skill-md/SKILL.md +++ b/.ai/skills/audit-skill-md/SKILL.md @@ -1,3 +1,9 @@ +--- +name: audit-skill-md +description: Audit the user-facing skill at skills/datafusion_python/SKILL.md against the current public Python API. Find new APIs that should be documented, stale mentions of removed/renamed APIs, examples that drifted from current idiomatic style, and places that need a "requires datafusion-python NN or newer" note. Run after upstream syncs and before each release. +argument-hint: [scope] (e.g., "session-context", "dataframe", "expr", "functions", "patterns", "pitfalls", "version-notes", "all") +--- + ---- -name: audit-skill-md -description: Audit the user-facing skill at skills/datafusion_python/SKILL.md against the current public Python API. Find new APIs that should be documented, stale mentions of removed/renamed APIs, examples that drifted from current idiomatic style, and places that need a "requires datafusion-python NN or newer" note. Run after upstream syncs and before each release. -argument-hint: [scope] (e.g., "session-context", "dataframe", "expr", "functions", "patterns", "pitfalls", "version-notes", "all") ---- - # Audit `skills/datafusion_python/SKILL.md` You are auditing the user-facing skill at diff --git a/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md index a3d82a670..828f227d8 100644 --- a/.ai/skills/check-upstream/SKILL.md +++ b/.ai/skills/check-upstream/SKILL.md @@ -1,3 +1,9 @@ +--- +name: check-upstream +description: Check if upstream Apache DataFusion features (functions, DataFrame ops, SessionContext methods, FFI types) are exposed in this Python project. Use when adding missing functions, auditing API coverage, or ensuring parity with upstream. +argument-hint: [area] (e.g., "scalar functions", "aggregate functions", "window functions", "dataframe", "session context", "ffi types", "all") +--- + ---- -name: check-upstream -description: Check if upstream Apache DataFusion features (functions, DataFrame ops, SessionContext methods, FFI types) are exposed in this Python project. Use when adding missing functions, auditing API coverage, or ensuring parity with upstream. -argument-hint: [area] (e.g., "scalar functions", "aggregate functions", "window functions", "dataframe", "session context", "ffi types", "all") ---- - # Check Upstream DataFusion Feature Coverage You are auditing the datafusion-python project to find features from the upstream Apache DataFusion Rust library that are **not yet exposed** in this Python binding project. Your goal is to identify gaps and, if asked, implement the missing bindings. diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md new file mode 100644 index 000000000..468216034 --- /dev/null +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -0,0 +1,190 @@ +--- +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)" +--- + + + +# FFI Capsule Protocol + +`datafusion-python` shares Rust objects with extension libraries through +PyCapsules. Every hook is a dunder method named `__datafusion___` 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> { ... } +``` + +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>` 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 session keeps one `Arc` for life + +`FFI_TaskContextProvider` holds its provider **weakly**, and every codec handed +to a foreign object carries one. A registered catalog provider upgrades that +handle on every `supports_filters_pushdown` and every `scan`. The handle is +bound to an `Arc` *allocation*, so anything that replaces the +allocation orphans every handle bound to the old one: +`TaskContextProvider went out of scope over FFI boundary`. + +So mutate `SessionState` in place — `*self.ctx.state_ref().write() = ...`, the +way `add_physical_optimizer_rule` and `set_session_query_planner` both do — +rather than deriving a replacement `SessionContext`. Carry the session id +across the rewrite; `SessionStateBuilder::new_from_existing` drops it and +`build` mints a fresh one, which desyncs `session_id()` from every +`TaskContext` the session hands out. + +Do not try to repair it after the fact: + +- **You cannot rebind what you cannot reach.** A codec embedded in a registered + `FFI_CatalogProvider`, and in every `FFI_SchemaProvider` and + `FFI_TableProvider` minted from it, has no Python-side handle. +- **A codec must not retain its session.** Codecs are routinely handed to a + provider that is registered straight back into the session that built them, + closing `SessionContext -> catalog -> FFI provider -> FFI codec -> + SessionContext`. + +`test_registered_providers_survive_a_planner_install` in +`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` +guards this. Its `WHERE` clause is load-bearing: filter pushdown upgrades the +weak handle during logical optimization, before plan serialization could fail +first for an unrelated reason. + +`SessionContext.enable_url_table` is the one method that mints a second +allocation for a session. Its result must not outlive the receiver. + +## Rule 7 — installing a planner mutates the session, and says so + +`set_query_planner` returns `None`, matching `add_physical_optimizer_rule`. The +query planner lives in `SessionState`, so it belongs to the session and not to +a handle on it; every context sharing that session plans through it. Do not +reintroduce a `with_query_planner` that pretends otherwise — the only way to +give a handle its own planner is a fresh `Arc`, which is what +Rule 6 forbids. + +Installing a codec rebuilds the installed planner against it, and that rebuild +reaches exactly one layer. `FFI_QueryPlanner::new_with_ffi_codecs` unwraps one +`ForeignQueryPlanner`; a fallback that planner resolved at install time sits in +its library's private data with no handle on this side, and cannot re-derive +codecs itself because `FFI_QueryPlanner` holds them by value and `Session` +exposes no accessor for the host's current ones. So do not promise that install +order is free — for a layered planner it is not. The examples cannot show this: +their fallback lives in the same cdylib as its wrapper, and `datafusion-ffi` +short-circuits a same-library hop rather than serializing. A fix has to come +from upstream; tracked in +[apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762). + +The session's planner also tracks whichever handle wrote it last, so +re-installing a planner on the original handle rebinds the session back to that +handle's codecs. `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` +pins that; changing it should be deliberate. + +## 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. diff --git a/.ai/skills/make-pythonic/SKILL.md b/.ai/skills/make-pythonic/SKILL.md index 7d490ec03..24c2bb817 100644 --- a/.ai/skills/make-pythonic/SKILL.md +++ b/.ai/skills/make-pythonic/SKILL.md @@ -1,3 +1,9 @@ +--- +name: make-pythonic +description: Audit and improve datafusion-python functions to accept native Python types (int, float, str, bool) instead of requiring explicit lit() or col() wrapping. Analyzes function signatures, checks upstream Rust implementations for type constraints, and applies the appropriate coercion pattern. +argument-hint: [scope] (e.g., "string functions", "datetime functions", "array functions", "math functions", "all", or a specific function name like "split_part") +--- + ---- -name: make-pythonic -description: Audit and improve datafusion-python functions to accept native Python types (int, float, str, bool) instead of requiring explicit lit() or col() wrapping. Analyzes function signatures, checks upstream Rust implementations for type constraints, and applies the appropriate coercion pattern. -argument-hint: [scope] (e.g., "string functions", "datetime functions", "array functions", "math functions", "all", or a specific function name like "split_part") ---- - # Make Python API Functions More Pythonic You are improving the datafusion-python API to feel more natural to Python users. The goal is to allow functions to accept native Python types (int, float, str, bool, etc.) for arguments that are contextually always or typically literal values, instead of requiring users to manually wrap them in `lit()`. diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c35801b11..d7af9b663 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -186,7 +186,7 @@ jobs: manylinux: "2_28" # FFI test wheel only needs to be built once per platform; gate to abi3. - - name: Build FFI test library + - name: Build FFI provider test library if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: @@ -196,6 +196,16 @@ jobs: args: --out dist rustup-components: rust-std + - name: Build FFI query planner test library + if: matrix.python-tag == 'abi3' + uses: PyO3/maturin-action@v1 + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + working-directory: examples/datafusion-ffi-query-planner-example + args: --out dist + rustup-components: rust-std + - name: Archive wheels uses: actions/upload-artifact@v7 with: @@ -207,7 +217,9 @@ jobs: uses: actions/upload-artifact@v7 with: name: test-ffi-manylinux-x86_64 - path: examples/datafusion-ffi-example/dist/* + path: | + examples/datafusion-ffi-example/dist/* + examples/datafusion-ffi-query-planner-example/dist/* # ============================================ # Build - Linux ARM64 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 558e751c8..047b35039 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -93,11 +93,15 @@ jobs: uv venv --python "${{ steps.setup-python.outputs.python-path }}" VENV_PY="$PWD/.venv/bin/python" uv sync --python "$VENV_PY" --dev --no-install-package datafusion + # Search recursively: the FFI artifact bundles more than one + # project, so upload-artifact keeps a `/dist/` prefix + # and the wheels are not all at the top of wheels/. WHEELS=$(find wheels/ -name "*.whl") if [ -n "$WHEELS" ]; then echo "Installing wheels:" echo "$WHEELS" - uv pip install --python "$VENV_PY" wheels/*.whl + # shellcheck disable=SC2086 # intentional split on newlines + uv pip install --python "$VENV_PY" $WHEELS else echo "ERROR: No wheels found!" exit 1 @@ -121,6 +125,8 @@ jobs: run: | cd examples/datafusion-ffi-example uv run --no-project pytest python/tests/_test*.py + cd ../datafusion-ffi-query-planner-example + uv run --no-project pytest python/tests/_test*.py - name: Run tpchgen-cli to create 1 Gb dataset if: matrix.wheel-tag == 'abi3' diff --git a/AGENTS.md b/AGENTS.md index fda08b23c..327ebd643 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -35,7 +35,24 @@ Skills follow the [Agent Skills](https://agentskills.io) open standard. Each ski To discover what skills are available, list `.ai/skills/` and read each `SKILL.md`. The frontmatter `name` and `description` fields summarize the -skill's purpose. +skill's purpose. Some descriptions begin with `TRIGGER —`; those are not tasks +to run on request but conventions to read *before* writing code that meets the +stated condition. + +## FFI Capsule Protocol + +The `__datafusion_*__` capsule getters are one protocol with a settled +convention. Before adding or changing one, read +[`.ai/skills/ffi-capsule-protocol/SKILL.md`](.ai/skills/ffi-capsule-protocol/SKILL.md). + +## Documentation Sources + +Search and edit `docs/source/`. `docs/temp/` is gitignored build output that +`grep -r` will surface with stale copies of the same pages. + +Before changing a public API, check +`docs/source/user-guide/upgrade-guides.md` for how the same API family was +migrated previously. Follow the established pattern rather than inventing one. ## Pull Requests @@ -48,7 +65,10 @@ Every pull request must follow the template in 3. **What changes are included in this PR?** — Summarize the individual changes. 4. **Are there any user-facing changes?** — Note any changes visible to users (new APIs, changed behavior, new files shipped in the package, etc.). If - there are breaking changes to public APIs, add the `api change` label. + there are breaking changes to public APIs, add the `api change` label **and + add a section to `docs/source/user-guide/upgrade-guides.md`** showing the + before and after. This applies to FFI hook method signatures, which + extension libraries implement. ## Pre-commit Checks diff --git a/Cargo.lock b/Cargo.lock index 9e0c1862a..c7632732a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -792,8 +792,7 @@ dependencies = [ [[package]] name = "datafusion" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96f76f0167ed0842b29a3d1e41be3c034c0a46409a3a703cc4cc84ee8c24abf4" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-schema", @@ -846,8 +845,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d79ec3460f6ed5c58f9b3f2d873fbc77748b82653bff1b4cdaf06de33bb4e05f" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "async-trait", @@ -871,8 +869,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b48cef241e2efcfd496fe05ae4d0d5de20793451862faefe406c397a467e12d4" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "async-trait", @@ -895,8 +892,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f72810485975c258f1b4d00baab31728470676c60c5546f366ebd0d99f05ab6" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-ipc", @@ -922,8 +918,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "533c28e75dba52f41bde187d23a1cb24ab91c7c097966824fa471e67b60320ea" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "futures", "log", @@ -933,8 +928,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b00a1fa0da26f6087136a82fea7f13c76a672cbab452d4086952a7cf770a19b" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "async-compression", @@ -970,8 +964,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ad17ec881bff2ed7768b4bfe971d3efbf3473f2fd1f9d365447bccbdf908678" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-ipc", @@ -995,8 +988,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1c9587cff8163f9bfcb186e8664ba85cb327031b8ff14330307f961ea2fc196" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-avro", @@ -1014,8 +1006,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5345285b0c3eaab412e7539b706973c083bd7e5bce575de5e0a3da488d08d1d" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "async-trait", @@ -1038,8 +1029,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da02fb9324f56bd8c53f1ee2e949547425cb66f76adc6832b10d44f80a1221d2" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "async-trait", @@ -1062,8 +1052,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c0b0dc1453952952fd5c69ad1c7f6042176e69ed233011d47e07cf74ed0949e" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-schema", @@ -1095,14 +1084,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a88fd985bc0550c36f557db69543cc9d6393b1509783520b30e902f23c555da6" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" [[package]] name = "datafusion-execution" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a98f1052f91b4991f0bf2ce1e4e36dfbdcda454a956b8c8d562c7c845e8fce1d" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-buffer", @@ -1127,8 +1114,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "464625a1f0e4b9df552d894fafcc8aac953ebbc8b0fa0acdaf20975fd615040e" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-schema", @@ -1152,8 +1138,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2604994999d5aeca1d1df645ffc98bc787447aaff05dde27aad0342b48fc1fe0" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1164,8 +1149,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ada11061028faad12bf1a45f2d5fa6624b7fcf65274fbf10a71ffe46c77cd10b" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-schema", @@ -1216,11 +1200,26 @@ dependencies = [ "pyo3-log", ] +[[package]] +name = "datafusion-ffi-query-planner-example" +version = "54.0.0" +dependencies = [ + "async-trait", + "datafusion", + "datafusion-catalog", + "datafusion-common", + "datafusion-ffi", + "datafusion-python-util", + "datafusion-session", + "pyo3", + "pyo3-build-config", + "pyo3-log", +] + [[package]] name = "datafusion-functions" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "051e97533e6af53e4aa0a0667cadc886abcaf36c4a5925019c55c0aa4c218fde" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-buffer", @@ -1251,8 +1250,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d0f1bb166d3572b6ed40e1afb2faaacade962abc08c2fcf04babee74681c56b" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1272,8 +1270,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ed756770f5f98369e181d692fd5ee6b1127ffd7322caba92f3730f9f5c92333" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1284,8 +1281,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91173fdb5c0ff2a41169a8ffa1b385b8844f18728747bb0a37e35ad7d5772a4f" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-ord", @@ -1309,8 +1305,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1bcdfb286a745461b126719c32700777e83df4f17cc44db5d71ebce5731e840" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "async-trait", @@ -1325,8 +1320,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ec4b508f1f93f00038ba3e737e894ec6c775528b4369413386655ae6125f0fc" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1342,8 +1336,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b352020834140073fbf5b46ee0ceb926e5074a9d0bcae1dbd91d0586d999cde" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1352,8 +1345,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15192effab05d38cce10e92a6fb48c967b5f166b27b7195a165a72b232569c58" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "datafusion-doc", "quote", @@ -1363,8 +1355,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "854445d9f7847e1e46089cf61b8d341a64382f14484e912c83a0f23b31216896" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "chrono", @@ -1383,8 +1374,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "671558dad1d2aa253c39c0a4c52515958b99eb91abf649f4b88d5e69cc55282f" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1406,8 +1396,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffae3d78c2da80ecc829cb58536cc5aca2e99cf1365eda694fc75bfe288861e0" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1421,8 +1410,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d9092ed15e7203fbd0903215172f7c9d18f10d94cba35137f3b3836f7c46f16" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "chrono", @@ -1439,8 +1427,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9005b6cf50b57b72d476c6ed4662b04be7ca6be5320ba9127c6d0b7e4218095b" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1459,8 +1446,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5787e4fcff4adc4fce8948441103a99705018b49c8dff0720b650bd7a15da112" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "arrow-data", @@ -1496,8 +1482,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0df0504eb9028d5e01f481af3519cde3f97ab650fd50ce7b787e94b69a7a193c" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-catalog", @@ -1523,8 +1508,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b92415a2442964f180d39cdcd8ff1edd99f9260c1499ba80a396499c2154d11" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1534,8 +1518,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62e4c0bd6af4fcabdbe201ee86fbeef2ac423a44e1d8fda56d994c9e2d1d3ad2" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -1545,8 +1528,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e651c8df0b90daed6a7be5921ec0ee379e6909705f063eeff70fd4e35010e4c" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "datafusion-common", @@ -1597,7 +1579,6 @@ dependencies = [ "arrow", "datafusion", "datafusion-ffi", - "datafusion-proto", "prost", "pyo3", "tokio", @@ -1606,8 +1587,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb56667ee38217efab19b895d9a936052cfb47ed438a19663351bdc42a6214a1" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow-schema", "async-trait", @@ -1621,8 +1601,7 @@ dependencies = [ [[package]] name = "datafusion-spark" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "992dd0b954f24cb576cbeee3555c3083b9cf7a5a5f2448ea25f7a437d444163f" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "bigdecimal", @@ -1651,8 +1630,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c29067cb9d32f8e603c45e15d61ea18f1069f96ceafeceb4e18466b8e5b31d9" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "arrow", "bigdecimal", @@ -1671,8 +1649,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8c026ddbbc33c34d0fac95a9620367c805ad6a7174cac52b57f360746c3e282" +source = "git+https://github.com/apache/datafusion?rev=1a944f43983188712aa1934d034f3087d82172e1#1a944f43983188712aa1934d034f3087d82172e1" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 9896e7421..0fabd5437 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,7 +27,12 @@ edition = "2024" rust-version = "1.88" [workspace] -members = ["crates/core", "crates/util", "examples/datafusion-ffi-example"] +members = [ + "crates/core", + "crates/util", + "examples/datafusion-ffi-example", + "examples/datafusion-ffi-query-planner-example", +] resolver = "3" [workspace.dependencies] @@ -50,6 +55,7 @@ datafusion-functions-aggregate = { version = "55.0.0" } datafusion-functions-window = { version = "55.0.0" } datafusion-spark = { version = "55.0.0" } datafusion-expr = { version = "55.0.0" } +datafusion-session = { version = "55.0.0" } prost = "0.14.3" serde_json = "1" uuid = { version = "1.23" } @@ -71,4 +77,24 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. +# +# When these are removed, raise the DataFusion requirement above to 55.1.0. +# `PySessionContext::set_session_query_planner` rebuilds an installed foreign +# query planner against the session's current codecs, and relies on +# `FFI_QueryPlanner::new_with_ffi_codecs` unwrapping a `ForeignQueryPlanner` and +# replacing its codecs. That replacement is a silent no-op before 55.1.0 +# (apache/datafusion#24722), so a 55.0.0 build would compile and then keep +# decoding with whichever codecs the planner was installed with -- installing a +# codec after `set_query_planner` would not take effect. [patch.crates-io] +datafusion = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-common = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } +datafusion-session = { git = "https://github.com/apache/datafusion", rev = "1a944f43983188712aa1934d034f3087d82172e1" } diff --git a/crates/core/src/catalog.rs b/crates/core/src/catalog.rs index 8ad49b098..dbc956611 100644 --- a/crates/core/src/catalog.rs +++ b/crates/core/src/catalog.rs @@ -30,7 +30,8 @@ use datafusion_ffi::catalog_provider::FFI_CatalogProvider; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::schema_provider::FFI_SchemaProvider; use datafusion_python_util::{ - create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, wait_for_future, + CapsuleGetterArg, call_capsule_getter, create_logical_extension_capsule, + ffi_logical_codec_from_pycapsule, wait_for_future, }; use pyo3::IntoPyObjectExt; use pyo3::exceptions::PyKeyError; @@ -626,9 +627,11 @@ fn extract_catalog_provider_from_pyobj( if catalog_provider.hasattr("__datafusion_catalog_provider__")? { let py = catalog_provider.py(); let codec_capsule = create_logical_extension_capsule(py, codec)?; - catalog_provider = catalog_provider - .getattr("__datafusion_catalog_provider__")? - .call1((codec_capsule,))?; + catalog_provider = call_capsule_getter( + catalog_provider, + "__datafusion_catalog_provider__", + CapsuleGetterArg::LogicalCodec(&codec_capsule), + )?; } let provider = if let Ok(capsule) = catalog_provider.cast::() { @@ -658,9 +661,11 @@ fn extract_schema_provider_from_pyobj( if schema_provider.hasattr("__datafusion_schema_provider__")? { let py = schema_provider.py(); let codec_capsule = create_logical_extension_capsule(py, codec)?; - schema_provider = schema_provider - .getattr("__datafusion_schema_provider__")? - .call1((codec_capsule,))?; + schema_provider = call_capsule_getter( + schema_provider, + "__datafusion_schema_provider__", + CapsuleGetterArg::LogicalCodec(&codec_capsule), + )?; } let provider = if let Ok(capsule) = schema_provider.cast::() { @@ -689,7 +694,7 @@ fn extract_logical_extension_codec( Some(obj) => obj, None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, }; - ffi_logical_codec_from_pycapsule(obj).map(Arc::new) + ffi_logical_codec_from_pycapsule(obj, None).map(Arc::new) } pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 26853e69f..94942a2d2 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -233,6 +233,14 @@ fn strip_wire_header<'a>( /// Sitting at the top of the session's logical codec stack means /// every serializer that reads `session.logical_codec()` automatically /// picks up Python-aware encoding for free. +/// +/// A codec deliberately does **not** retain the session it was built from. +/// Codecs are routinely handed to a provider that is then registered back into +/// that same session, so retaining here would close a cycle: +/// `SessionContext -> catalog -> FFI provider -> FFI codec -> here`. Keeping +/// the weak `FFI_TaskContextProvider` valid is instead a matter of never +/// replacing the session's `Arc`; see +/// `PySessionContext::set_session_query_planner`. #[derive(Debug)] pub struct PythonLogicalCodec { inner: Arc, @@ -443,6 +451,9 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio /// would round-trip at the logical level but break at the physical /// level. Both layers reuse the shared payload framing /// ([`PY_SCALAR_UDF_FAMILY`] et al.) so the wire format is identical. +/// +/// Like [`PythonLogicalCodec`], this does not retain the session it was built +/// from; see that type for why. #[derive(Debug)] pub struct PythonPhysicalCodec { inner: Arc, diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 7bbeed2f1..75bfed601 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -36,7 +36,7 @@ use datafusion::datasource::listing::{ }; use datafusion::datasource::{MemTable, TableProvider}; use datafusion::execution::context::{ - DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext, + DataFilePaths, QueryPlanner, SQLOptions, SessionConfig, SessionContext, TaskContext, }; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, UnboundedMemoryPool}; @@ -53,14 +53,16 @@ use datafusion_ffi::config::extension_options::FFI_ExtensionOptions; use datafusion_ffi::execution::FFI_TaskContextProvider; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ - create_logical_extension_capsule, create_physical_extension_capsule, - ffi_logical_codec_from_pycapsule, get_global_ctx, get_tokio_runtime, - physical_codec_from_pycapsule, physical_optimizer_rule_from_pycapsule, spawn_future, - wait_for_future, + CapsuleGetterArg, call_capsule_getter, create_logical_extension_capsule, + create_physical_extension_capsule, create_query_planner_capsule, + ffi_logical_codec_from_pycapsule, ffi_physical_codec_from_pycapsule, + ffi_query_planner_from_pycapsule, get_global_ctx, get_tokio_runtime, + physical_optimizer_rule_from_pycapsule, spawn_future, validate_pycapsule, wait_for_future, }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; @@ -199,9 +201,20 @@ impl PySessionConfig { "Expected extension object to define __datafusion_extension_options__()", )); } - let capsule = extension.call_method0("__datafusion_extension_options__")?; + // Routed through `call_capsule_getter` like every other getter in the + // family, even though this one takes no argument and so has no arity + // error to diagnose. The point is that the grep in the FFI capsule + // protocol skill turns up no exceptions. + let capsule = call_capsule_getter( + extension, + "__datafusion_extension_options__", + CapsuleGetterArg::None, + )?; let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_extension_options")?; + // No `check_ffi_version` here: `FFI_ExtensionOptions` carries no + // version field, so there is nothing to compare. let extension: NonNull = capsule .pointer_checked(Some(c"datafusion_extension_options"))? .cast(); @@ -408,6 +421,11 @@ impl PySessionContext { } pub fn enable_url_table(&self) -> PyResult { + // Pre-existing caveat, unrelated to query planners: this is the one + // method that mints a second `Arc` for a session. Any + // weak `FFI_TaskContextProvider` handed out by the receiver stays bound + // to the receiver, so the returned context must not outlive it. See + // `set_session_query_planner` for why everything else mutates in place. Ok(PySessionContext { ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()), logical_codec: Arc::clone(&self.logical_codec), @@ -719,9 +737,11 @@ impl PySessionContext { let py = factory.py(); let ffi = self.ffi_logical_codec(); let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; - factory = factory - .getattr("__datafusion_table_provider_factory__")? - .call1((codec_capsule,))?; + factory = call_capsule_getter( + factory, + "__datafusion_table_provider_factory__", + CapsuleGetterArg::LogicalCodec(&codec_capsule), + )?; } let factory: Arc = @@ -754,9 +774,11 @@ impl PySessionContext { let py = provider.py(); let ffi = self.ffi_logical_codec(); let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; - provider = provider - .getattr("__datafusion_catalog_provider_list__")? - .call1((codec_capsule,))?; + provider = call_capsule_getter( + provider, + "__datafusion_catalog_provider_list__", + CapsuleGetterArg::LogicalCodec(&codec_capsule), + )?; } let provider = if let Ok(capsule) = provider.cast::() { @@ -790,9 +812,11 @@ impl PySessionContext { let py = provider.py(); let ffi = self.ffi_logical_codec(); let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; - provider = provider - .getattr("__datafusion_catalog_provider__")? - .call1((codec_capsule,))?; + provider = call_capsule_getter( + provider, + "__datafusion_catalog_provider__", + CapsuleGetterArg::LogicalCodec(&codec_capsule), + )?; } let provider = if let Ok(capsule) = provider.cast::() { @@ -1204,13 +1228,34 @@ impl PySessionContext { let rule = physical_optimizer_rule_from_pycapsule(&rule)?; let state_ref = self.ctx.state_ref(); let mut guard = state_ref.write(); + // Rebuilding through the builder mints a fresh session id, but this + // mutates the caller's own session rather than deriving a new one, so + // the id has to survive. See `set_session_query_planner` for why losing + // it leaves `session_id()` disagreeing with every `TaskContext` the + // session hands out. let new_state = SessionStateBuilder::new_from_existing(guard.clone()) + .with_session_id(guard.session_id().to_string()) .with_physical_optimizer_rule(rule) .build(); *guard = new_state; Ok(()) } + /// Install a foreign query planner on this session. + /// + /// A session holds exactly one planner, so this replaces whichever one is + /// installed. It mutates the session rather than deriving a new one, the + /// same way `add_physical_optimizer_rule` does; see + /// [`Self::set_session_query_planner`] for why that matters to FFI handles. + pub fn set_query_planner( + slf: &Bound<'_, Self>, + planner: Bound<'_, PyAny>, + ) -> PyDataFusionResult<()> { + let planner = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?; + slf.borrow().set_session_query_planner(Some(planner)); + Ok(()) + } + pub fn table_provider(&self, name: &str, py: Python) -> PyResult { let provider = wait_for_future(py, self.ctx.table_provider(name)) // Outer error: runtime/async failure @@ -1377,52 +1422,120 @@ impl PySessionContext { PyCapsule::new_with_value(py, ffi_ctx_provider, cr"datafusion_task_context_provider") } + /// `session` exists so this matches the protocol an extension library + /// implements, where the argument is how the library reaches the session + /// it is being installed on. A session already is one, so it is ignored. + #[pyo3(signature = (session=None))] pub fn __datafusion_logical_extension_codec__<'py>( &self, py: Python<'py>, + session: Option>, ) -> PyResult> { - let ffi = self.ffi_logical_codec(); - create_logical_extension_capsule(py, ffi.as_ref()) + let _ = session; + create_logical_extension_capsule(py, self.ffi_logical_codec().as_ref()) } - pub fn with_logical_extension_codec<'py>( + /// See [`Self::__datafusion_logical_extension_codec__`] for `session`. + #[pyo3(signature = (session=None))] + pub fn __datafusion_query_planner__<'py>( &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + // An already-foreign planner is re-exported as its original handle + // rather than gaining another layer, because `new_with_ffi_codecs` + // unwraps a `ForeignQueryPlanner`. It still adopts the codecs supplied + // here, so a consumer that wraps this capsule decodes our plans with + // our codecs. + let planner = Arc::clone(self.ctx.state().query_planner()); + let ffi = FFI_QueryPlanner::new_with_ffi_codecs( + planner, + Self::ffi_logical_codec_for(&self.ctx, &self.logical_codec), + Self::ffi_physical_codec_for(&self.ctx, &self.physical_codec), + ); + create_query_planner_capsule(py, &ffi) + } + + pub fn with_logical_extension_codec<'py>( + slf: &Bound<'py, Self>, codec: Bound<'py, PyAny>, ) -> PyDataFusionResult { - let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?; + let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); - let logical_codec = Arc::new(PythonLogicalCodec::new(inner)); - Ok(Self { - ctx: Arc::clone(&self.ctx), + let this = slf.borrow(); + // Carry the receiver's inlining setting over. `PythonLogicalCodec::new` + // defaults it to on, so building the replacement without this would + // silently re-enable inline Python UDF encoding on a context that had + // opted out with `with_python_udf_inlining(enabled=False)`. + let logical_codec = Arc::new( + PythonLogicalCodec::new(inner) + .with_python_udf_inlining(this.logical_codec.python_udf_inlining()), + ); + let derived = Self { + ctx: Arc::clone(&this.ctx), logical_codec, - physical_codec: Arc::clone(&self.physical_codec), - }) + physical_codec: Arc::clone(&this.physical_codec), + }; + // The session is shared, and an installed foreign planner holds the + // codecs it was built with, so it has to be rebuilt against this one. + derived.set_session_query_planner(None); + Ok(derived) } + /// See [`Self::__datafusion_logical_extension_codec__`] for `session`. + #[pyo3(signature = (session=None))] pub fn __datafusion_physical_extension_codec__<'py>( &self, py: Python<'py>, + session: Option>, ) -> PyResult> { - let ffi = self.ffi_physical_codec(); - create_physical_extension_capsule(py, ffi.as_ref()) + let _ = session; + create_physical_extension_capsule(py, self.ffi_physical_codec().as_ref()) } pub fn with_physical_extension_codec<'py>( - &self, + slf: &Bound<'py, Self>, codec: Bound<'py, PyAny>, ) -> PyDataFusionResult { - let inner = physical_codec_from_pycapsule(&codec)?; - let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); + let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; + let inner: Arc = (&inner_ffi).into(); - Ok(Self { - ctx: Arc::clone(&self.ctx), - logical_codec: Arc::clone(&self.logical_codec), + let this = slf.borrow(); + // See `with_logical_extension_codec` for why the flag is carried over. + let physical_codec = Arc::new( + PythonPhysicalCodec::new(inner) + .with_python_udf_inlining(this.physical_codec.python_udf_inlining()), + ); + let derived = Self { + ctx: Arc::clone(&this.ctx), + logical_codec: Arc::clone(&this.logical_codec), physical_codec, - }) + }; + // See `with_logical_extension_codec`. + derived.set_session_query_planner(None); + Ok(derived) } pub fn with_python_udf_inlining(&self, enabled: bool) -> Self { + // Rebinding the session's planner is a side effect on state shared with + // every other handle, so do not pay it for a call that changes nothing. + // A defensive `with_python_udf_inlining(enabled=True)` on a context that + // already inlines would otherwise rebind the session's planner to this + // handle's codecs, and callers routinely discard the result. Returning + // the codecs as-is is observationally equivalent to the rebuild below, + // which wraps the same inner codec in a fresh `Python*Codec`. + if self.logical_codec.python_udf_inlining() == enabled + && self.physical_codec.python_udf_inlining() == enabled + { + return Self { + ctx: Arc::clone(&self.ctx), + logical_codec: Arc::clone(&self.logical_codec), + physical_codec: Arc::clone(&self.physical_codec), + }; + } + let logical_codec = Arc::new( PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner())) .with_python_udf_inlining(enabled), @@ -1431,15 +1544,72 @@ impl PySessionContext { PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) .with_python_udf_inlining(enabled), ); - Self { + let derived = Self { ctx: Arc::clone(&self.ctx), logical_codec, physical_codec, - } + }; + // See `with_logical_extension_codec`. + derived.set_session_query_planner(None); + derived } } impl PySessionContext { + /// Write the session's query planner, in place. + /// + /// Pass `Some(planner)` to install one, or `None` to rebuild whichever + /// foreign planner is already installed against this context's current + /// codecs. A foreign planner carries the FFI codecs it encodes and decodes + /// with, so replacing a codec means handing the planner new ones. A planner + /// this library owns carries no codecs and needs neither. + /// + /// This mutates `SessionState` through `state_ref()` rather than deriving a + /// new `SessionContext`, and that is the whole point. Every codec handed to + /// a foreign object carries an `FFI_TaskContextProvider`, which holds its + /// provider *weakly* and upgrades it on each callback — a registered + /// catalog provider does so on every `supports_filters_pushdown` and every + /// `scan`. Those handles are bound to one `Arc` allocation. + /// Deriving a replacement would orphan every handle bound to the old one + /// and leave callbacks failing with `TaskContextProvider went out of scope + /// over FFI boundary`, and only the codecs this struct happens to hold + /// could be moved across; one already embedded in a registered + /// `FFI_CatalogProvider` is unreachable. Keeping the single allocation for + /// the life of the session avoids the problem rather than repairing it. + fn set_session_query_planner(&self, planner: Option) { + let planner = planner.or_else(|| { + let state = self.ctx.state(); + let installed: &dyn std::any::Any = state.query_planner().as_ref(); + installed + .downcast_ref::() + .map(|planner| planner.0.clone()) + }); + let Some(planner) = planner else { + return; + }; + + let inner: Arc = (&planner).into(); + let planner: Arc = (&FFI_QueryPlanner::new_with_ffi_codecs( + inner, + Self::ffi_logical_codec_for(&self.ctx, &self.logical_codec), + Self::ffi_physical_codec_for(&self.ctx, &self.physical_codec), + )) + .into(); + + let state_ref = self.ctx.state_ref(); + let mut guard = state_ref.write(); + // `with_session_id` is load-bearing, not redundant. + // `SessionStateBuilder::new_from_existing` drops the id and `build` + // mints a fresh one, while `SessionContext` cached the original in a + // field of its own. Without this the session would report one id from + // `session_id()` and another from every `TaskContext` it hands out. + // Same reason `add_physical_optimizer_rule` carries it over. + *guard = SessionStateBuilder::new_from_existing(guard.clone()) + .with_session_id(guard.session_id().to_string()) + .with_query_planner(planner) + .build(); + } + async fn _table(&self, name: &str) -> datafusion::common::Result { self.ctx.table(name).await } @@ -1501,27 +1671,36 @@ impl PySessionContext { /// Used at every site that exports the codec across an FFI boundary /// (capsule getters, Rust wrappers for Python-defined providers, etc.). pub(crate) fn ffi_logical_codec(&self) -> Arc { - let inner: Arc = - Arc::clone(&self.logical_codec) as Arc; + Arc::new(Self::ffi_logical_codec_for(&self.ctx, &self.logical_codec)) + } + + fn ffi_logical_codec_for( + ctx: &Arc, + codec: &Arc, + ) -> FFI_LogicalExtensionCodec { + let codec: Arc = + Arc::clone(codec) as Arc; let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::clone(&self.ctx) as Arc; - Arc::new(FFI_LogicalExtensionCodec::new( - inner, - Some(runtime), - &ctx_provider, - )) + let ctx_provider = Arc::clone(ctx) as Arc; + FFI_LogicalExtensionCodec::new(codec, Some(runtime), &ctx_provider) + } + + fn ffi_physical_codec_for( + ctx: &Arc, + codec: &Arc, + ) -> FFI_PhysicalExtensionCodec { + let codec: Arc = + Arc::clone(codec) as Arc; + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = Arc::clone(ctx) as Arc; + FFI_PhysicalExtensionCodec::new(codec, Some(runtime), &ctx_provider) } /// Build an FFI-wrapped clone of the session's physical codec on demand. pub(crate) fn ffi_physical_codec(&self) -> Arc { - let inner: Arc = - Arc::clone(&self.physical_codec) as Arc; - let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::clone(&self.ctx) as Arc; - Arc::new(FFI_PhysicalExtensionCodec::new( - inner, - Some(runtime), - &ctx_provider, + Arc::new(Self::ffi_physical_codec_for( + &self.ctx, + &self.physical_codec, )) } } diff --git a/crates/core/src/udaf.rs b/crates/core/src/udaf.rs index caf7b97bc..6a2675193 100644 --- a/crates/core/src/udaf.rs +++ b/crates/core/src/udaf.rs @@ -28,7 +28,7 @@ use datafusion::logical_expr::{ Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, Signature, Volatility, }; use datafusion_ffi::udaf::FFI_AggregateUDF; -use datafusion_python_util::parse_volatility; +use datafusion_python_util::{CapsuleGetterArg, call_capsule_getter, parse_volatility}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyTuple}; @@ -365,7 +365,11 @@ impl PyAggregateUDF { } if func.hasattr("__datafusion_aggregate_udf__")? { - let capsule = func.getattr("__datafusion_aggregate_udf__")?.call0()?; + let capsule = call_capsule_getter( + func.clone(), + "__datafusion_aggregate_udf__", + CapsuleGetterArg::None, + )?; let capsule = capsule.cast::().map_err(py_datafusion_err)?; let function = aggregate_udf_from_capsule(capsule)?; return Ok(Self { function }); diff --git a/crates/core/src/udf.rs b/crates/core/src/udf.rs index 2006401db..6376c81a8 100644 --- a/crates/core/src/udf.rs +++ b/crates/core/src/udf.rs @@ -31,7 +31,7 @@ use datafusion::logical_expr::{ Volatility, }; use datafusion_ffi::udf::FFI_ScalarUDF; -use datafusion_python_util::parse_volatility; +use datafusion_python_util::{CapsuleGetterArg, call_capsule_getter, parse_volatility}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyTuple}; @@ -248,7 +248,11 @@ impl PyScalarUDF { #[staticmethod] pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { if func.hasattr("__datafusion_scalar_udf__")? { - let capsule = func.getattr("__datafusion_scalar_udf__")?.call0()?; + let capsule = call_capsule_getter( + func.clone(), + "__datafusion_scalar_udf__", + CapsuleGetterArg::None, + )?; let capsule = capsule.cast::().map_err(to_datafusion_err)?; let data: NonNull = capsule .pointer_checked(Some(c"datafusion_scalar_udf"))? diff --git a/crates/core/src/udtf.rs b/crates/core/src/udtf.rs index cffa0c12a..51ea8fa4f 100644 --- a/crates/core/src/udtf.rs +++ b/crates/core/src/udtf.rs @@ -24,10 +24,10 @@ use datafusion::execution::context::SessionContext; use datafusion::execution::session_state::SessionState; use datafusion::logical_expr::Expr; use datafusion_ffi::udtf::FFI_TableFunction; +use datafusion_python_util::call_capsule_getter; use pyo3::IntoPyObjectExt; -use pyo3::exceptions::{PyImportError, PyTypeError}; use pyo3::prelude::*; -use pyo3::types::{PyCapsule, PyDict, PyTuple, PyType}; +use pyo3::types::{PyCapsule, PyDict, PyTuple}; use crate::context::PySessionContext; use crate::errors::{py_datafusion_err, to_datafusion_err}; @@ -76,15 +76,11 @@ impl PyTableFunction { Some(session) => session, None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, }; - let capsule = func - .getattr("__datafusion_table_function__")? - .call1((session,)).map_err(|err| { - if err.get_type(py).is(PyType::new::(py)) { - PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table functions. Either downgrade DataFusion or upgrade your function library.") - } else { - err - } - })?; + let capsule = call_capsule_getter( + func.clone(), + "__datafusion_table_function__", + Some(&session), + )?; let capsule = capsule.cast::()?; let data: NonNull = capsule .pointer_checked(Some(c"datafusion_table_function"))? diff --git a/crates/core/src/udwf.rs b/crates/core/src/udwf.rs index ebec8f3bd..8935c9ba8 100644 --- a/crates/core/src/udwf.rs +++ b/crates/core/src/udwf.rs @@ -30,7 +30,7 @@ use datafusion::logical_expr::{ }; use datafusion::scalar::ScalarValue; use datafusion_ffi::udwf::FFI_WindowUDF; -use datafusion_python_util::parse_volatility; +use datafusion_python_util::{CapsuleGetterArg, call_capsule_getter, parse_volatility}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyList, PyTuple}; @@ -262,11 +262,8 @@ impl PyWindowUDF { #[staticmethod] pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { - let capsule = if func.hasattr("__datafusion_window_udf__")? { - func.getattr("__datafusion_window_udf__")?.call0()? - } else { - func - }; + let capsule = + call_capsule_getter(func, "__datafusion_window_udf__", CapsuleGetterArg::None)?; let capsule = capsule.cast::().map_err(to_datafusion_err)?; let data: NonNull = capsule diff --git a/crates/util/Cargo.toml b/crates/util/Cargo.toml index c23667b0f..00d5946a5 100644 --- a/crates/util/Cargo.toml +++ b/crates/util/Cargo.toml @@ -30,6 +30,5 @@ tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] } pyo3 = { workspace = true } datafusion = { workspace = true } datafusion-ffi = { workspace = true } -datafusion-proto = { workspace = true } arrow = { workspace = true } prost = { workspace = true } diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 9327d7f2f..5b31f7708 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -29,8 +29,8 @@ use datafusion_ffi::execution::FFI_TaskContextProvider; use datafusion_ffi::physical_optimizer::FFI_PhysicalOptimizerRule; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; use datafusion_ffi::table_provider::FFI_TableProvider; -use datafusion_proto::physical_plan::PhysicalExtensionCodec; use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyType}; @@ -156,6 +156,17 @@ pub fn parse_volatility(value: &str) -> PyDataFusionResult { }) } +/// Check that `capsule` carries the name `name`, and say so usefully if not. +/// +/// This looks redundant next to `capsule.pointer_checked(Some(name))`, which +/// also rejects a mismatched name, and it is not. `pointer_checked` bottoms out +/// in CPython's `PyCapsule_GetPointer`, whose error is the fixed string +/// `PyCapsule_GetPointer called with incorrect name` — it names neither what +/// was expected nor what was found. Handing an extension author that message +/// tells them nothing about which capsule they got wrong. +/// +/// So call this first at every extraction site. `pointer_checked` still gets +/// the name, because it is what actually guards the cast. pub fn validate_pycapsule(capsule: &Bound, name: &str) -> PyResult<()> { let capsule_name = capsule.name()?; if capsule_name.is_none() { @@ -175,28 +186,65 @@ pub fn validate_pycapsule(capsule: &Bound, name: &str) -> PyResult<() Ok(()) } +/// Reject an FFI struct built against a different major version of +/// `datafusion-ffi`. +/// +/// `found` comes from the struct's own `version` function pointer, which +/// reports the major version of the library that produced it. +/// +/// This is a diagnostic, not a soundness guarantee. Reading `version` out of +/// the struct already assumes the local field layout, and `version` is not the +/// first field on any of these types, so a sufficiently different layout can +/// fault before this ever runs. What it buys is a clear error for the case that +/// actually happens -- an extension library compiled against a different +/// DataFusion -- rather than undefined behaviour on first use, which is what +/// `datafusion_ffi::version` exists for. +/// +/// Not every FFI type carries a version. `FFI_TaskContextProvider`, +/// `FFI_TableProviderFactory`, and `FFI_ExtensionOptions` have no such field, +/// so their importers cannot check and are not expected to. +/// +/// # If the FFI ABI stabilizes +/// +/// Exact equality is the right test only while `datafusion_ffi::version` +/// tracks the DataFusion crate's semver major, which it does today +/// (`env!("CARGO_PKG_VERSION")`, `.major`). That number therefore moves on +/// every major release whether or not the ABI actually changed. +/// +/// Should a version span become compatible, **this function body is the only +/// thing to change** -- callers pass a `found` value and no policy. Relaxing it +/// at a call site instead would reintroduce the split this helper exists to +/// remove. +/// +/// The likelier fix is upstream, not here: if the ABI is stable but `version` +/// still follows the crate major, upstream's own compatibility marker is wrong +/// for every consumer, not just this one. Prefer waiting for +/// `datafusion_ffi::version` to reflect the real ABI over inventing a range +/// policy locally. +pub fn check_ffi_version(kind: &str, found: u64) -> PyResult<()> { + let expected = datafusion_ffi::version(); + if found != expected { + return Err(PyImportError::new_err(format!( + "Incompatible DataFusion {kind} major version {found}; expected {expected}. \ + Rebuild the library providing this object against a matching DataFusion." + ))); + } + Ok(()) +} + pub fn table_provider_from_pycapsule<'py>( mut obj: Bound<'py, PyAny>, session: Bound<'py, PyAny>, ) -> PyResult>> { - if obj.hasattr("__datafusion_table_provider__")? { - obj = obj - .getattr("__datafusion_table_provider__")? - .call1((session,)).map_err(|err| { - let py = obj.py(); - if err.get_type(py).is(PyType::new::(py)) { - PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table providers. Either downgrade DataFusion or upgrade your function library.") - } else { - err - } - })?; - } + obj = call_capsule_getter(obj, "__datafusion_table_provider__", Some(&session))?; if let Ok(capsule) = obj.cast::() { + validate_pycapsule(capsule, "datafusion_table_provider")?; let data: NonNull = capsule .pointer_checked(Some(c"datafusion_table_provider"))? .cast(); let provider = unsafe { data.as_ref() }; + check_ffi_version("table provider", unsafe { (provider.version)() })?; let provider: Arc = provider.into(); Ok(Some(provider)) @@ -214,23 +262,192 @@ pub fn create_logical_extension_capsule<'py>( PyCapsule::new_with_value(py, codec, cr"datafusion_logical_extension_codec") } -pub fn ffi_logical_codec_from_pycapsule(obj: Bound) -> PyResult { - let attr_name = "__datafusion_logical_extension_codec__"; - let capsule = if obj.hasattr(attr_name)? { - obj.getattr(attr_name)?.call0()? - } else { - obj +/// The single positional argument datafusion-python hands a capsule getter. +/// +/// Carrying the argument and its description together is what lets one +/// diagnostic serve the whole family: the getters do not all take a session, +/// and telling the author of a catalog provider that their method "must accept +/// the SessionContext" would send them to fix the wrong signature. +#[derive(Clone, Copy)] +pub enum CapsuleGetterArg<'a, 'py> { + /// The getter takes no arguments, so it cannot refuse one. A `TypeError` + /// from one of these is always the getter's own and is never rewritten. + None, + /// The `SessionContext` the object is being installed on. + Session(&'a Bound<'py, PyAny>), + /// The host's logical extension codec, as a + /// `datafusion_logical_extension_codec` capsule. + LogicalCodec(&'a Bound<'py, PyAny>), +} + +impl<'a, 'py> CapsuleGetterArg<'a, 'py> { + fn value(self) -> Option<&'a Bound<'py, PyAny>> { + match self { + Self::None => None, + Self::Session(value) | Self::LogicalCodec(value) => Some(value), + } + } + + fn expected(self) -> &'static str { + match self { + Self::None => "", + Self::Session(_) => "the SessionContext it is being installed on", + Self::LogicalCodec(_) => "the host's logical extension codec", + } + } +} + +impl<'a, 'py> From>> for CapsuleGetterArg<'a, 'py> { + fn from(session: Option<&'a Bound<'py, PyAny>>) -> Self { + session.map_or(Self::None, Self::Session) + } +} + +/// Calls `obj.____(arg)`, or `obj.____()` for +/// [`CapsuleGetterArg::None`]. Returns `obj` untouched when it has no such +/// attribute, so a raw capsule passes straight through. +/// +/// **Every** capsule getter goes through here rather than calling `getattr` +/// directly, so that the mapping from a refused argument to a useful error +/// lives in one place. Three importers previously each had their own copy and +/// each missed later corrections to it; the getters that take no argument route +/// through as well, so the rule has no exceptions to remember. +pub fn call_capsule_getter<'a, 'py: 'a>( + obj: Bound<'py, PyAny>, + attr_name: &str, + arg: impl Into>, +) -> PyResult> { + if !obj.hasattr(attr_name)? { + return Ok(obj); + } + + let arg = arg.into(); + let getter = obj.getattr(attr_name)?; + let result = match arg.value() { + Some(value) => getter.call1((value,)), + None => getter.call0(), }; + result.map_err(|err| { + let py = obj.py(); + if arg.value().is_none() || !err.get_type(py).is(PyType::new::(py)) { + return err; + } + + // Not every `TypeError` here means the getter refused the argument. One + // raised *inside* a correctly-signed getter would otherwise be reported + // as a version mismatch, sending an extension author to upgrade a + // library that is already correct. + // + // The two are distinguishable: an arity mismatch is raised by the call + // machinery before the getter's frame exists, so nothing unwinds and no + // traceback is attached. An error from the body unwinds that frame and + // carries one. + // + // This holds because the call originates in Rust. `call1` pushes no + // Python caller frame, so an arity error has nothing to unwind and + // arrives bare. If a Python-level shim is ever interposed between the + // host and the getter, its frame would supply a traceback and this + // check would quietly stop firing -- the diagnostic would degrade to + // the bare `TypeError` it exists to replace, with no test failing. + if err.traceback(py).is_some() { + return err; + } + + let import_err = PyImportError::new_err(format!( + "Incompatible libraries. `{attr_name}` must accept {}. \ + Upgrade the library providing this object.", + arg.expected() + )); + // Keep the original reachable as `__cause__` rather than discarding it. + import_err.set_cause(py, Some(err)); + import_err + }) +} + +pub fn ffi_logical_codec_from_pycapsule<'py>( + obj: Bound<'py, PyAny>, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult { + let capsule = call_capsule_getter(obj, "__datafusion_logical_extension_codec__", session)?; + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_logical_extension_codec")?; let data: NonNull = capsule .pointer_checked(Some(c"datafusion_logical_extension_codec"))? .cast(); let codec = unsafe { data.as_ref() }; + check_ffi_version("logical extension codec", unsafe { (codec.version)() })?; + + Ok(codec.clone()) +} + +pub fn ffi_physical_codec_from_pycapsule<'py>( + obj: Bound<'py, PyAny>, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult { + let capsule = call_capsule_getter(obj, "__datafusion_physical_extension_codec__", session)?; + + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_physical_extension_codec")?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_physical_extension_codec"))? + .cast(); + let codec = unsafe { data.as_ref() }; + check_ffi_version("physical extension codec", unsafe { (codec.version)() })?; Ok(codec.clone()) } +/// Extracts the `FFI_TaskContextProvider` a session exposes. +/// +/// An extension library exporting a codec needs one for the decode callbacks +/// its codec will receive. Taking the host's means those callbacks resolve +/// names against the session that is actually running the query, and removes +/// any need for the library to construct a `SessionContext` of its own. +pub fn ffi_task_context_provider_from_pycapsule( + session: &Bound, +) -> PyResult { + let capsule = call_capsule_getter( + session.clone(), + "__datafusion_task_context_provider__", + CapsuleGetterArg::None, + )?; + + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_task_context_provider")?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_task_context_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + + Ok(provider.clone()) +} + +pub fn create_query_planner_capsule<'py>( + py: Python<'py>, + planner: &FFI_QueryPlanner, +) -> PyResult> { + PyCapsule::new_with_value(py, planner.clone(), cr"datafusion_query_planner") +} + +pub fn ffi_query_planner_from_pycapsule<'py>( + obj: &Bound<'py, PyAny>, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult { + let capsule = call_capsule_getter(obj.clone(), "__datafusion_query_planner__", session)?; + + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_query_planner")?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_query_planner"))? + .cast(); + let planner = unsafe { data.as_ref() }; + check_ffi_version("query planner", unsafe { (planner.version)() })?; + + Ok(planner.clone()) +} + pub fn create_physical_extension_capsule<'py>( py: Python<'py>, codec: &FFI_PhysicalExtensionCodec, @@ -247,6 +464,11 @@ pub fn create_physical_extension_capsule<'py>( /// Use this when `Arc<$output_type>: From<&$ffi_type>` (infallible /// conversion). For fallible conversions use [`try_from_pycapsule!`] /// instead. +/// +/// The generated extractor does not check the FFI major version, because not +/// every FFI type carries one. If `$ffi_type` has a `version` field, call +/// [`check_ffi_version`] on it yourself, as the hand-written extractors in this +/// crate do. #[macro_export] macro_rules! from_pycapsule { ($fn_name:ident, $capsule_name:literal, $ffi_type:ty, $output_type:ty) => { @@ -256,10 +478,11 @@ macro_rules! from_pycapsule { use $crate::pyo3::prelude::*; use $crate::pyo3::types::PyCapsule; - let mut obj = obj.clone(); - if obj.hasattr(concat!("__", $capsule_name, "__"))? { - obj = obj.getattr(concat!("__", $capsule_name, "__"))?.call0()?; - } + let obj = $crate::call_capsule_getter( + obj.clone(), + concat!("__", $capsule_name, "__"), + $crate::CapsuleGetterArg::None, + )?; let capsule = obj.cast::().map_err(|_| { $crate::errors::py_datafusion_err(concat!( "Invalid ", @@ -293,10 +516,11 @@ macro_rules! try_from_pycapsule { use $crate::pyo3::prelude::*; use $crate::pyo3::types::PyCapsule; - let mut obj = obj.clone(); - if obj.hasattr(concat!("__", $capsule_name, "__"))? { - obj = obj.getattr(concat!("__", $capsule_name, "__"))?.call0()?; - } + let obj = $crate::call_capsule_getter( + obj.clone(), + concat!("__", $capsule_name, "__"), + $crate::CapsuleGetterArg::None, + )?; let capsule = obj.cast::().map_err(|_| { $crate::errors::py_datafusion_err(concat!( "Invalid ", @@ -326,13 +550,12 @@ macro_rules! try_from_pycapsule { #[doc(hidden)] pub use pyo3; -from_pycapsule!( - physical_codec_from_pycapsule, - "datafusion_physical_extension_codec", - FFI_PhysicalExtensionCodec, - dyn PhysicalExtensionCodec -); - +// There is deliberately no `physical_codec_from_pycapsule` here. These macros +// call the getter with no arguments, which is right for the two hooks below but +// wrong for `__datafusion_physical_extension_codec__`, which takes the session +// it is being installed on. Use `ffi_physical_codec_from_pycapsule`, which +// passes the session, and convert with `(&ffi).into()` if you need an +// `Arc`. from_pycapsule!( physical_optimizer_rule_from_pycapsule, "datafusion_physical_optimizer_rule", diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index bf65cad2a..31cd9391f 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -232,6 +232,202 @@ extension that has been written using this approach and the most thoroughly impl As we continue to expose more of the DataFusion features, we intend to follow this same design pattern. +## Query Planners Across Multiple Libraries + +A query can involve three independent native libraries: `datafusion-python`, a library +that owns table providers or functions, and a library that owns the query planner. The +examples use two separate extension crates so each role has a distinct shared-library +identity: + +- [`datafusion-ffi-example`] owns providers, functions, and their codecs. +- [`datafusion-ffi-query-planner-example`] owns the planner and its configuration. + +The `SessionContext` owns the codecs used for the exchange and supplies them to the +foreign planner. This lets the planner decode provider-owned objects and lets +`datafusion-python` decode the physical plan returned by the planner. The examples use +process-local tokens to demonstrate ownership; production codecs should serialize +durable metadata instead. + +The current Python API has one external logical codec and one external physical codec. +Installing another codec replaces the prior codec rather than composing a registry. +The example therefore has one external codec owner, and the planner uses built-in +physical nodes. Install the provider codecs before the planner where possible. + +The current FFI logical codec supports providers and UDFs but not arbitrary custom +`LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and +local build commands. + +### Capsule getters receive the session they are installed on + +`__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and +`__datafusion_physical_extension_codec__` all take the `SessionContext` the object is +being installed on, the same way `__datafusion_table_provider__` does: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +This exists because the FFI constructors need things an extension library does not +have. `FFI_{Logical,Physical}ExtensionCodec::new` needs a `TaskContextProvider` for the +decode callbacks the codec will receive, and `FFI_QueryPlanner::new` needs both codecs +on top of that. Taking them from the session is what keeps a library from constructing +a `SessionContext` purely to satisfy a parameter — an empty one resolves nothing, and +`FFI_TaskContextProvider` holds it weakly, so a context built inline in the getter is +already dropped by the time the capsule is used. + +A planner uses `FFI_QueryPlanner::new_with_ffi_codecs` with the two codecs it takes off +the session, and never touches a provider directly. That also matches what installation +does anyway: `set_query_planner` builds the planner against the codecs of the session +that will run the query. + +`SessionContext` accepts the argument on all three getters and ignores it, so a session +satisfies the same protocol an extension library implements. When you export the current +planner to wrap it, `ctx.__datafusion_query_planner__()` and +`ctx.__datafusion_query_planner__(ctx)` are both fine. + +### A codec decodes against the session that is running the query + +Because the provider comes from the host, a decode callback running inside an extension +library resolves names against the session running the query. A function registered with +`ctx.register_udf(...)` is visible to a foreign codec decoding a node that references it +by name, and the handle is live rather than a snapshot, so a registration made after the +codec is installed is visible too. + +This is covered in +`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`, +where the example codecs take a `require_udf_on_decode` name and resolve it out of the +task context they are handed. + +### One session, one `Arc` + +Every codec handed to a foreign object carries an `FFI_TaskContextProvider`, and that +type holds its provider **weakly**. A registered catalog provider upgrades the handle on +every `supports_filters_pushdown` and every `scan`. Those handles are bound to one +particular `Arc` allocation, not to the logical session, so anything that +replaces the allocation orphans all of them and the next query fails with +`TaskContextProvider went out of scope over FFI boundary`. + +So a `PySessionContext` keeps the `Arc` it was created with for its whole +life. Installing a query planner writes the new `SessionState` back through +`state_ref()`, exactly as `add_physical_optimizer_rule` does, rather than deriving a +replacement context. The session id is carried across that rewrite — `SessionStateBuilder` +mints a fresh one otherwise — so `session_id()` and every `TaskContext` the session hands +out keep agreeing. + +Repairing the damage instead of avoiding it does not work in general. A context can +rebuild the codecs it holds in its own fields, but a codec already embedded in a +registered `FFI_CatalogProvider` — and in every `FFI_SchemaProvider` and +`FFI_TableProvider` minted from it — is not reachable from Python at all. Nor can the +codec simply retain the session that built it: a codec handed to a provider is routinely +registered straight back into that same session, which would close the cycle +`SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and leak it. + +`SessionContext.enable_url_table` is the one exception. It clones the underlying +`SessionContext`, so the returned context has an allocation of its own and must not +outlive the receiver. + +### What a derived context shares + +`with_logical_extension_codec`, `with_physical_extension_codec`, and +`with_python_udf_inlining` return a new `SessionContext` wrapping the *same* underlying +session. Only the Python-side codec settings differ; catalogs, tables, registered +functions, and configuration are the one shared session, so a registration on either +side is visible to both. + +`set_query_planner` does not return anything. The query planner lives in `SessionState`, +so it is a property of the session rather than of a handle on it, and installing one is +visible to every context sharing that session — including ones a `with_*` call returned +earlier. Installing a codec on a session that already has a foreign planner rebuilds +that planner against the new codec for the same reason: there is one planner, and it has +to carry the codecs currently in force. This happens on the shared session, so it takes +effect even if the returned context is discarded — `ctx.with_python_udf_inlining(...)` +whose result is thrown away still leaves the session's planner carrying the codecs of +that discarded handle. A call that changes nothing is exempt: asking for the inlining +setting a context already has returns a handle without touching the session. + +The rule that falls out of this is worth stating on its own, because it is the one thing +that surprises people: + +> The session's query planner carries the codecs of the handle that most recently +> installed one. Every other path — `Expr.to_bytes(ctx)`, `ExecutionPlan.to_bytes(ctx)`, +> registering a provider — uses the codecs of the handle you call it on. + +Those can be different handles, and then one session has two codecs in effect at once: + +```python +ctx = ctx.with_logical_extension_codec(codec_a) +ctx.set_query_planner(planner) +ctx.with_logical_extension_codec(codec_b) # discarded + +Expr.to_bytes(expr, ctx) # encodes with codec_a -- ctx's own field +ctx.sql(...).collect() # plans with codec_b -- installed via the discarded handle +``` + +Chaining `ctx = ctx.with_...(...)`, as the example below does, keeps the two in step. +`test_the_planner_and_the_handle_can_hold_different_codecs` pins the divergence. + +```python +ctx = SessionContext(config) +ctx = ctx.with_logical_extension_codec(provider_logical_codec) +ctx = ctx.with_physical_extension_codec(provider_physical_codec) +ctx.set_query_planner(planner) +ctx.register_udf(my_udf) +``` + +Order is a readability preference rather than a requirement — installing a codec after a +planner rebuilds the planner against it. + +A session holds exactly one query planner. Calling `set_query_planner` again replaces the +installed planner instead of layering another one. To chain planners, have the new +planner wrap the capsule returned by `SessionContext.__datafusion_query_planner__()`, +captured before the new planner is installed, and delegate to it explicitly. + +### Rebinding a planner's codecs is one level deep + +The rebuild above swaps the codecs on the installed `ForeignQueryPlanner` handle, and +only that handle. A planner that wraps a fallback resolved that fallback when *it* was +installed, and holds the result inside its own library's private data — behind a +`create_physical_plan` function pointer, with no Python-side handle. A codec installed +afterwards therefore reaches the outer planner and not the fallback, which keeps +whichever codecs were in force when it was imported. + +Neither side can repair that: + +- **The host cannot reach it.** `FFI_QueryPlanner::new_with_ffi_codecs` unwraps exactly + one `ForeignQueryPlanner` layer. There is no deeper handle to unwrap — the same + situation as a codec embedded in a registered `FFI_CatalogProvider`. +- **The planner library cannot re-derive it.** `FFI_QueryPlanner` holds its codecs by + value, and `Session` exposes no accessor for the ones the host currently has, so + `create_physical_plan` cannot pick them up from the session it is handed. The rebuild + has to be eager, and an eager rebuild only sees the top layer. + +A fix has to come from upstream, and is tracked in +[apache/datafusion#24762](https://github.com/apache/datafusion/issues/24762). + +The stale codecs stay usable rather than dangling — they hold weak handles to the one +`Arc` that Rule 6 keeps alive — so the effect is a fallback hop +serializing with an older codec, not a failure. It is also invisible to the examples +here, which use one fallback in the same cdylib as its wrapper; `datafusion-ffi` +short-circuits a same-library hop rather than serializing, so no codec runs. A fallback +in a *different* library would serialize, and would do it with the codecs it was +imported with. + +So install the codecs before a layered planner. If a codec has to go in afterwards, +install the outer planner again *on the handle that holds the new codec* — that re-runs +its getter, which re-imports the fallback against that handle's codecs. Re-installing on +the original handle rebinds the session's planner back to the original handle's codecs +instead, which is the trap +`test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` pins. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait @@ -257,3 +453,5 @@ At the time of this writing, the FFI features are under active development. To s the latest status, we recommend reviewing the code in the [datafusion-ffi] crate. [datafusion-ffi]: https://crates.io/crates/datafusion-ffi +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/docs/source/user-guide/io/table_provider.md b/docs/source/user-guide/io/table_provider.md index 3c436ba1d..5dc2dc086 100644 --- a/docs/source/user-guide/io/table_provider.md +++ b/docs/source/user-guide/io/table_provider.md @@ -29,6 +29,11 @@ via [PyCapsule](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule). A complete example can be found in the [examples folder](https://github.com/apache/datafusion-python/tree/main/examples). +The method takes the `SessionContext` it is being registered on. Take whatever +the FFI constructor needs from that session — here the logical extension codec — +rather than building one inside your library. See the {ref}`ffi` guide for the +full capsule protocol. + ```rust #[pymethods] impl MyTableProvider { @@ -36,13 +41,13 @@ impl MyTableProvider { fn __datafusion_table_provider__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { - let name = cr"datafusion_table_provider".into(); - let provider = Arc::new(self.clone()); - let provider = FFI_TableProvider::new(provider, false, None); + let codec = ffi_logical_codec_from_pycapsule(session, None)?; + let provider = FFI_TableProvider::new_with_ffi_codec(provider, false, None, codec); - PyCapsule::new_bound(py, provider, Some(name.clone())) + PyCapsule::new_with_value(py, provider, cr"datafusion_table_provider") } } ``` diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 360e0533c..29085bc3d 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -19,6 +19,128 @@ # Upgrade Guides +## DataFusion 55.0.0 + +This release extends the change made in 52.0.0 to the remaining {ref}`ffi` hook +methods. Users who contribute their own `LogicalExtensionCodec` or +`PhysicalExtensionCodec` via FFI must update +`__datafusion_logical_extension_codec__` and +`__datafusion_physical_extension_codec__` to accept an additional +`session: Bound` parameter, and take the `TaskContextProvider` from that +session rather than constructing a `SessionContext` of their own. + +Before: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let ctx_provider: Arc = Arc::clone(&self.ctx_provider); + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +After: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +The dropped `&` on the last argument is not a typo. That parameter is +`impl Into`, so it accepts either an +`&Arc`, as before, or an `FFI_TaskContextProvider`, +which is what `ffi_task_context_provider_from_pycapsule` hands back. Both forms +compile; the argument changes because the provider now comes from the session +rather than from a field. + +A codec that keeps its own `SessionContext` still compiles, but its decode +callbacks resolve names against that empty session instead of the one running +the query, so a function registered with `SessionContext.register_udf` is not +visible to it. Taking the provider from `session` also removes a lifetime +hazard: `FFI_TaskContextProvider` holds its provider weakly, so a context +constructed inside the getter is already dropped by the time the capsule is +used. + +`SessionContext` accepts the argument on its own capsule getters and ignores +it, so existing calls such as `ctx.__datafusion_logical_extension_codec__()` +continue to work unchanged. + +New in this release, `__datafusion_query_planner__` follows the same protocol. +It receives the session and takes both extension codecs from it, so a planner +library never builds a `TaskContextProvider` at all. Install one with +`SessionContext.set_query_planner(planner)`, which mutates the session the same +way `add_physical_optimizer_rule` does and returns nothing — the query planner +lives in `SessionState`, so it belongs to the session rather than to a +particular handle on it. See the {ref}`ffi` guide for the full protocol. + +### Mismatched extension libraries now fail loudly + +Objects imported through the capsule protocol are checked against the major +version of `datafusion-ffi` this package was built with. A table provider, +extension codec, or query planner produced by a library built against a +different DataFusion major version now raises an `ImportError` naming the +version found and the one expected, instead of being used as-is. Table +providers previously performed no such check. + +This is a diagnostic rather than a soundness guarantee — reading the version +out of the struct already assumes the local field layout — but it turns the +common "extension library built against the wrong DataFusion" mistake into a +clear message rather than undefined behaviour on first use. + +`FFI_TaskContextProvider`, `FFI_TableProviderFactory`, and `FFI_ExtensionOptions` +carry no version field, so objects of those types cannot be checked. + +### Changes to the `datafusion-python-util` crate + +Extension libraries written in Rust usually depend on the +`datafusion-python-util` crate for the helpers that read these capsules. Two of +those helpers changed, because the getter they call now takes the session. + +`ffi_logical_codec_from_pycapsule` takes a second argument. Pass `Some(session)` +when importing an object from another library, so its getter receives the +session it is being installed on. Pass `None` when the object *is* a session and +you are asking it for what it holds: + +```rust +// Before +let codec = ffi_logical_codec_from_pycapsule(obj)?; + +// After +let codec = ffi_logical_codec_from_pycapsule(obj, Some(session))?; +``` + +`physical_codec_from_pycapsule` has been **removed**. It called +`__datafusion_physical_extension_codec__` with no arguments, which no longer +matches the protocol, so against an updated codec it raised a bare `TypeError` +and against an outdated one it silently produced a codec bound to the wrong +session. Use `ffi_physical_codec_from_pycapsule`, which passes the session: + +```rust +// Before +let codec: Arc = physical_codec_from_pycapsule(&obj)?; + +// After +let ffi = ffi_physical_codec_from_pycapsule(obj, Some(session))?; +let codec: Arc = (&ffi).into(); +``` + +`physical_optimizer_rule_from_pycapsule` and `task_context_from_pycapsule` are +unchanged. Their hooks take no session. + +Calling a getter that still has the old signature now raises an `ImportError` +naming the method, with the original `TypeError` retained as its `__cause__`, +rather than a bare `TypeError`. + ## DataFusion 54.0.0 The `Config` class has been removed. It was a standalone wrapper around diff --git a/examples/README.md b/examples/README.md index e0e3056d9..7bbb45dcf 100644 --- a/examples/README.md +++ b/examples/README.md @@ -49,6 +49,15 @@ Here is a direct link to the file used in the examples: - [Fan out distinct expressions to a multiprocessing pool](./multiprocessing_pickle_expr.py) - [Distribute expression evaluation across Ray actors](./ray_pickle_expr.py) +### Rust FFI Extensions + +- [Table providers, functions, and codecs](./datafusion-ffi-example/) +- [Independent query planner and planner configuration](./datafusion-ffi-query-planner-example/) + +These two crates form a three-library interoperability example with +`datafusion-python`. They are separate shared libraries so the tests exercise real FFI +type and codec boundaries rather than same-library Rust downcasts. + ### Substrait Support - [Serialize query plans using Substrait](./substrait.py) diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md new file mode 100644 index 000000000..0fa10d7f3 --- /dev/null +++ b/examples/datafusion-ffi-example/README.md @@ -0,0 +1,48 @@ + + +# DataFusion Python FFI provider example + +This crate is the **provider library** in the three-library query-planning example. It exports table providers, functions, and the logical and physical codecs needed to serialize objects owned by this library. The companion planner is in [`../datafusion-ffi-query-planner-example`](../datafusion-ffi-query-planner-example/). + +The example intentionally uses separate `cdylib` crates for these roles: + +1. **A — `datafusion-python`:** owns the `SessionContext` and executes the result. +2. **B — this crate:** owns table providers, functions, and provider execution plans. +3. **C — the planner crate:** receives the logical plan and returns a physical plan. + +Separate shared libraries guarantee distinct DataFusion library markers. This catches type-identity mistakes that a planner and provider compiled into one shared library would hide. + +## Codec behavior + +`MyLogicalExtensionCodec` serializes this example's in-memory table providers, and `MyPhysicalExtensionCodec` serializes provider-owned memory scans and opaque FFI wrappers around them. Both use documented, process-local, one-shot token registries. The registries make ownership and callback routing visible without pretending to be a portable format. They assume trusted in-process payloads and consume each token during decoding. A production provider should instead encode durable metadata from which its provider and plans can be reconstructed. + +Both codec getters take the `SessionContext` they are being installed on and pull the `TaskContextProvider` off it, so decode callbacks resolve session configuration and registered functions against the session that is running the query. Passing `require_udf_on_decode` to either constructor makes every decode call resolve a named scalar function out of that context, which is how the tests check where the registry came from. + +This example makes the provider library the sole external codec owner. Register both provider codecs before installing the planner: + +```python +ctx = ctx.with_logical_extension_codec(provider_logical_codec) +ctx = ctx.with_physical_extension_codec(provider_physical_codec) +ctx.set_query_planner(planner) +``` + +Installing a codec after the planner rebuilds the planner against it, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. The exception is a planner that wraps a fallback: the rebuild reaches the installed planner only, not the fallback inside it, so codecs-first is a requirement there. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep), which also covers why re-installing a planner rebinds the session to the codecs of whichever handle it was installed on. + +For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. diff --git a/examples/datafusion-ffi-example/pyproject.toml b/examples/datafusion-ffi-example/pyproject.toml index 7f85e9487..c51fa8a8d 100644 --- a/examples/datafusion-ffi-example/pyproject.toml +++ b/examples/datafusion-ffi-example/pyproject.toml @@ -21,7 +21,8 @@ build-backend = "maturin" [project] name = "datafusion_ffi_example" -requires-python = ">=3.9" +# Matches the abi3-py310 feature the crate builds against. +requires-python = ">=3.10" classifiers = [ "Programming Language :: Rust", "Programming Language :: Python :: Implementation :: CPython", diff --git a/examples/datafusion-ffi-example/src/catalog_provider.rs b/examples/datafusion-ffi-example/src/catalog_provider.rs index a56b5855c..75890d083 100644 --- a/examples/datafusion-ffi-example/src/catalog_provider.rs +++ b/examples/datafusion-ffi-example/src/catalog_provider.rs @@ -94,7 +94,7 @@ impl FixedSchemaProvider { ) -> PyResult> { let provider = Arc::clone(&self.inner) as Arc; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_SchemaProvider::new_with_ffi_codec(provider, None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_schema_provider") @@ -186,7 +186,7 @@ impl MyCatalogProvider { ) -> PyResult> { let provider = Arc::clone(&self.inner) as Arc; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_CatalogProvider::new_with_ffi_codec(provider, None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_catalog_provider") @@ -245,7 +245,7 @@ impl MyCatalogProviderList { ) -> PyResult> { let provider = Arc::clone(&self.inner) as Arc; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_CatalogProviderList::new_with_ffi_codec(provider, None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_catalog_provider_list") diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index eccf7b81a..3d00fdb3e 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -35,6 +35,7 @@ pub(crate) mod config; pub(crate) mod logical_extension_codec; pub(crate) mod physical_extension_codec; pub(crate) mod physical_optimizer; +pub(crate) mod required_udf; pub(crate) mod scalar_udf; pub(crate) mod table_function; pub(crate) mod table_provider; diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs index 8c3976d37..1fcaaef4c 100644 --- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs @@ -15,40 +15,91 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::collections::HashMap; +use std::fmt; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; use arrow::datatypes::SchemaRef; -use datafusion::common::{Result, TableReference}; +use datafusion::catalog::MemTable; +use datafusion::common::{DataFusionError, Result, TableReference}; use datafusion::datasource::TableProvider; -use datafusion::execution::{TaskContext, TaskContextProvider}; +use datafusion::execution::TaskContext; use datafusion::logical_expr::{Extension, LogicalPlan, ScalarUDF}; -use datafusion::prelude::SessionContext; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; -use datafusion_python_util::get_tokio_runtime; +use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio_runtime}; use pyo3::prelude::*; use pyo3::types::PyCapsule; -/// Tracks how often each `try_*_udf` entry point fires. Surface for -/// Python tests to assert the session routed UDF -/// encode/decode through this user-supplied codec rather than the -/// upstream default. +use crate::required_udf::{TaskContextProbe, resolve_required_udf}; + +const TABLE_PROVIDER_TOKEN: &[u8] = b"DFPYEXTP"; +static NEXT_TABLE_PROVIDER_ID: AtomicU64 = AtomicU64::new(1); +static TABLE_PROVIDERS: OnceLock>>> = OnceLock::new(); + +/// Hands a provider to another library in this process by token. +/// +/// Encoding inserts, decoding removes. Two consequences worth knowing before +/// copying this: +/// +/// - **Decode consumes the token.** Decoding the same encoded bytes twice +/// fails the second time with `Unknown ... table provider token`. That is +/// fine here because every plan is encoded immediately before the single +/// decode that consumes it, but it rules out anything that replays a stored +/// plan, retries a decode, or fans one encoded plan out to several readers. +/// - **An encode that is never decoded leaks.** Nothing expires entries, so a +/// plan that fails to reach its decoder keeps its provider alive for the +/// life of the process. +/// +/// Both are acceptable for an example whose job is to show that Rust type +/// identity survives a trip through two other libraries. Neither is acceptable +/// in a real codec, which should encode metadata sufficient to rebuild the +/// provider rather than parking the object here. +fn table_providers() -> &'static Mutex>> { + TABLE_PROVIDERS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn token_id(buf: &[u8], prefix: &[u8]) -> Option { + let id: [u8; 8] = buf.strip_prefix(prefix)?.try_into().ok()?; + Some(u64::from_le_bytes(id)) +} + #[derive(Debug, Default)] pub(crate) struct CallCounters { pub encode_udf: AtomicUsize, pub decode_udf: AtomicUsize, + pub encode_table_provider: AtomicUsize, + pub decode_table_provider: AtomicUsize, + pub task_ctx: TaskContextProbe, } -/// Minimal user-supplied `LogicalExtensionCodec` for integration tests. -/// Delegates everything to `DefaultLogicalExtensionCodec` and bumps -/// counters on the UDF entry points so tests can prove the wrapper -/// installed via `SessionContext.with_logical_extension_codec(...)` -/// actually gets consulted. -#[derive(Debug)] +/// Example codec for objects owned by this extension library. +/// +/// The table-provider token registry is intentionally process-local. It is a compact +/// example of preserving Rust type identity across three loaded libraries, not a +/// network serialization format. Production libraries should encode reconstructible +/// provider metadata rather than retaining objects in a global registry. +/// +/// See [`table_providers`] for the token lifecycle, which is narrower than it +/// looks: a decode consumes its token, so the same encoded plan cannot be +/// decoded twice. struct CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec, counters: Arc, + /// Scalar function every table-provider decode must resolve from the + /// `TaskContext` it is handed. See [`crate::required_udf`]. + required_udf: Option, +} + +impl fmt::Debug for CountingLogicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CountingLogicalExtensionCodec") + .field("inner", &self.inner) + .field("counters", &self.counters) + .finish_non_exhaustive() + } } impl LogicalExtensionCodec for CountingLogicalExtensionCodec { @@ -72,6 +123,21 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { schema: SchemaRef, ctx: &TaskContext, ) -> Result> { + resolve_required_udf(self.required_udf.as_deref(), ctx, &self.counters.task_ctx)?; + if let Some(id) = token_id(buf, TABLE_PROVIDER_TOKEN) { + self.counters + .decode_table_provider + .fetch_add(1, Ordering::SeqCst); + return table_providers() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .remove(&id) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "Unknown datafusion-ffi-example table provider token {id}" + )) + }); + } self.inner .try_decode_table_provider(buf, table_ref, schema, ctx) } @@ -82,6 +148,19 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { node: Arc, buf: &mut Vec, ) -> Result<()> { + if node.downcast_ref::().is_some() { + self.counters + .encode_table_provider + .fetch_add(1, Ordering::SeqCst); + let id = NEXT_TABLE_PROVIDER_ID.fetch_add(1, Ordering::SeqCst); + table_providers() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .insert(id, node); + buf.extend_from_slice(TABLE_PROVIDER_TOKEN); + buf.extend_from_slice(&id.to_le_bytes()); + return Ok(()); + } self.inner.try_encode_table_provider(table_ref, node, buf) } @@ -105,47 +184,72 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { #[derive(Clone)] pub(crate) struct MyLogicalExtensionCodec { counters: Arc, + required_udf: Option, } #[pymethods] impl MyLogicalExtensionCodec { + /// Build the codec. + /// + /// `require_udf_on_decode` names a scalar function that every table + /// provider decode must find in the `TaskContext` it is handed. Leave it + /// unset for the ordinary behaviour; set it to observe *which* session's + /// registry the FFI decode callback actually receives. #[new] - fn new() -> Self { + #[pyo3(signature = (require_udf_on_decode=None))] + fn new(require_udf_on_decode: Option) -> Self { Self { counters: Arc::new(CallCounters::default()), + required_udf: require_udf_on_decode, } } - /// Number of `try_encode_udf` invocations observed since - /// construction. + /// Number of decode calls that resolved `require_udf_on_decode`. + fn task_context_udf_resolutions(&self) -> usize { + self.counters.task_ctx.resolutions() + } + + /// Session id of the `TaskContext` the most recent decode callback ran + /// against, or `None` before any decode. + fn last_task_context_session_id(&self) -> Option { + self.counters.task_ctx.last_session_id() + } + fn encode_udf_calls(&self) -> usize { self.counters.encode_udf.load(Ordering::SeqCst) } - /// Number of `try_decode_udf` invocations observed. fn decode_udf_calls(&self) -> usize { self.counters.decode_udf.load(Ordering::SeqCst) } - /// Capsule entry point consumed by - /// `datafusion_python_util::ffi_logical_codec_from_pycapsule`. - /// datafusion-python invokes this with no arguments when the user - /// calls `ctx.with_logical_extension_codec(my_codec)`. The codec - /// owns its own bare `SessionContext` as a TaskContextProvider — - /// good enough for tests that only exercise UDF encode/decode. + fn table_provider_encode_calls(&self) -> usize { + self.counters.encode_table_provider.load(Ordering::SeqCst) + } + + fn table_provider_decode_calls(&self) -> usize { + self.counters.decode_table_provider.load(Ordering::SeqCst) + } + + /// Export the codec, bound to the session it is being installed on. + /// + /// `session` supplies the `TaskContextProvider` the FFI decode callbacks + /// resolve, so this library never constructs a `SessionContext` and the + /// callbacks see the registry of the session running the query. fn __datafusion_logical_extension_codec__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { let inner: Arc = Arc::new(CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec {}, counters: Arc::clone(&self.counters), + required_udf: self.required_udf.clone(), }); let runtime = get_tokio_runtime().handle().clone(); - let bare_session: Arc = Arc::new(SessionContext::new()); - let ctx_provider = bare_session as Arc; - let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), ctx_provider); PyCapsule::new_with_value(py, ffi, cr"datafusion_logical_extension_codec") } diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index 35ef77f6b..f9e96382e 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -15,36 +15,79 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::collections::HashMap; +use std::fmt; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; -use datafusion::common::Result; -use datafusion::execution::{TaskContext, TaskContextProvider}; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; use datafusion::logical_expr::ScalarUDF; use datafusion::physical_plan::ExecutionPlan; -use datafusion::prelude::SessionContext; +use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, }; -use datafusion_python_util::get_tokio_runtime; +use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio_runtime}; use pyo3::prelude::*; use pyo3::types::PyCapsule; +use crate::required_udf::{TaskContextProbe, resolve_required_udf}; + +const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; +static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); +static EXECUTION_PLANS: OnceLock>>> = OnceLock::new(); + +/// Execution-plan counterpart of the logical codec's provider registry, with +/// the same lifecycle: encoding inserts, decoding removes, so a decode +/// consumes its token and an encode that is never decoded leaks. See +/// [`crate::logical_extension_codec`] for why that is acceptable here and not +/// in a real codec. +fn execution_plans() -> &'static Mutex>> { + EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn token_id(buf: &[u8]) -> Option { + let id: [u8; 8] = buf.strip_prefix(EXECUTION_PLAN_TOKEN)?.try_into().ok()?; + Some(u64::from_le_bytes(id)) +} + #[derive(Debug, Default)] pub(crate) struct PhysicalCallCounters { pub encode_udf: AtomicUsize, pub decode_udf: AtomicUsize, + pub encode_execution_plan: AtomicUsize, + pub decode_execution_plan: AtomicUsize, + pub task_ctx: TaskContextProbe, } -/// Mirror of [`super::logical_extension_codec::CountingLogicalExtensionCodec`] -/// for the physical layer. Delegates to `DefaultPhysicalExtensionCodec` -/// and bumps counters on UDF encode/decode so tests can prove the -/// session routed through a user-supplied physical codec. -#[derive(Debug)] +/// Physical companion to the logical example codec. +/// +/// Provider-owned memory scan plans use a same-process token registry so the +/// owning cdylib can restore their concrete Rust type after the plan travels +/// through the independent query-planner and datafusion-python libraries. +/// +/// See [`execution_plans`] for the token lifecycle, which is narrower than it +/// looks: a decode consumes its token, so the same encoded plan cannot be +/// decoded twice. struct CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec, counters: Arc, + /// Scalar function every decode call must resolve from the `TaskContext` + /// it is handed. See [`crate::required_udf`]. + required_udf: Option, +} + +impl fmt::Debug for CountingPhysicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CountingPhysicalExtensionCodec") + .field("inner", &self.inner) + .field("counters", &self.counters) + .finish_non_exhaustive() + } } impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { @@ -55,6 +98,21 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { + resolve_required_udf(self.required_udf.as_deref(), ctx, &self.counters.task_ctx)?; + if let Some(id) = token_id(buf) { + self.counters + .decode_execution_plan + .fetch_add(1, Ordering::SeqCst); + return execution_plans() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .remove(&id) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "Unknown datafusion-ffi-example execution plan token {id}" + )) + }); + } self.inner.try_decode(buf, inputs, ctx, proto_converter) } @@ -64,6 +122,22 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { + // The provider owns DataSourceExec. A ForeignExecutionPlan can wrap a + // host-added execution decorator around that scan; retaining the opaque + // wrapper preserves its original library identity without downcasting it. + if node.is::() || node.is::() { + self.counters + .encode_execution_plan + .fetch_add(1, Ordering::SeqCst); + let id = NEXT_EXECUTION_PLAN_ID.fetch_add(1, Ordering::SeqCst); + execution_plans() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .insert(id, node); + buf.extend_from_slice(EXECUTION_PLAN_TOKEN); + buf.extend_from_slice(&id.to_le_bytes()); + return Ok(()); + } self.inner.try_encode(node, buf, proto_converter) } @@ -87,17 +161,37 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { #[derive(Clone)] pub(crate) struct MyPhysicalExtensionCodec { counters: Arc, + required_udf: Option, } #[pymethods] impl MyPhysicalExtensionCodec { + /// Build the codec. + /// + /// `require_udf_on_decode` names a scalar function that every decode call + /// must find in the `TaskContext` it is handed. Leave it unset for the + /// ordinary behaviour; set it to observe *which* session's registry the + /// FFI decode callback actually receives. #[new] - fn new() -> Self { + #[pyo3(signature = (require_udf_on_decode=None))] + fn new(require_udf_on_decode: Option) -> Self { Self { counters: Arc::new(PhysicalCallCounters::default()), + required_udf: require_udf_on_decode, } } + /// Number of decode calls that resolved `require_udf_on_decode`. + fn task_context_udf_resolutions(&self) -> usize { + self.counters.task_ctx.resolutions() + } + + /// Session id of the `TaskContext` the most recent decode callback ran + /// against, or `None` before any decode. + fn last_task_context_session_id(&self) -> Option { + self.counters.task_ctx.last_session_id() + } + fn encode_udf_calls(&self) -> usize { self.counters.encode_udf.load(Ordering::SeqCst) } @@ -106,20 +200,33 @@ impl MyPhysicalExtensionCodec { self.counters.decode_udf.load(Ordering::SeqCst) } + fn execution_plan_encode_calls(&self) -> usize { + self.counters.encode_execution_plan.load(Ordering::SeqCst) + } + + fn execution_plan_decode_calls(&self) -> usize { + self.counters.decode_execution_plan.load(Ordering::SeqCst) + } + + /// Export the codec, bound to the session it is being installed on. + /// + /// See [`crate::logical_extension_codec::MyLogicalExtensionCodec`] for why + /// `session` is taken rather than a context this library invents. fn __datafusion_physical_extension_codec__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { let inner: Arc = Arc::new(CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec {}, counters: Arc::clone(&self.counters), + required_udf: self.required_udf.clone(), }); let runtime = get_tokio_runtime().handle().clone(); - let bare_session: Arc = Arc::new(SessionContext::new()); - let ctx_provider = bare_session as Arc; - let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") } diff --git a/examples/datafusion-ffi-example/src/required_udf.rs b/examples/datafusion-ffi-example/src/required_udf.rs new file mode 100644 index 000000000..a21362d7f --- /dev/null +++ b/examples/datafusion-ffi-example/src/required_udf.rs @@ -0,0 +1,102 @@ +// 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. + +//! Support for exercising the `TaskContext` a codec is handed when it decodes. +//! +//! A codec exported over FFI carries a `TaskContextProvider`, and the decode +//! callbacks in `datafusion-ffi` resolve it to a `TaskContext` before calling +//! into the codec. Nothing in the example codecs read anything out of that +//! context, so which session it belongs to was untestable: the token +//! registries they use are keyed by an integer and ignore the registry. +//! +//! The codecs can now be asked to resolve a named scalar function from the +//! context they are given on every decode, which makes the answer observable. +//! Because the codecs take their provider from the session they are installed +//! on, a function registered on the host with `register_udf` resolves. + +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::execution::TaskContext; +use datafusion_common::error::Result as DataFusionResult; +use datafusion_common::plan_err; + +/// What the codecs record about the `TaskContext` their decode callbacks run +/// against. +#[derive(Debug, Default)] +pub(crate) struct TaskContextProbe { + resolutions: AtomicUsize, + last_session_id: Mutex>, +} + +impl TaskContextProbe { + /// Successful `require_udf_on_decode` lookups since construction. + pub(crate) fn resolutions(&self) -> usize { + self.resolutions.load(Ordering::SeqCst) + } + + /// Session id of the most recent decode callback, or `None` if the codec + /// has not been asked to decode anything yet. + /// + /// Recorded on every decode, so a test can tell *which* session the + /// callback was bound to rather than only that some session resolved a + /// name. A `SessionContext` that derives a fork must keep reporting the id + /// it reports from `session_id()`. + pub(crate) fn last_session_id(&self) -> Option { + self.last_session_id + .lock() + .expect("task context probe mutex poisoned") + .clone() + } +} + +/// Resolves `required` against `ctx`, the context the decode callback was given. +/// +/// Records `ctx`'s session id either way. `Ok(())` when nothing was requested. +/// Otherwise the name must be present in the context's scalar function +/// registry, and `probe` counts each success so a test can tell a resolved +/// lookup from a skipped one. +pub(crate) fn resolve_required_udf( + required: Option<&str>, + ctx: &TaskContext, + probe: &TaskContextProbe, +) -> DataFusionResult<()> { + // Unconditional: the session id is worth observing even when the caller + // asked for no function. + *probe + .last_session_id + .lock() + .expect("task context probe mutex poisoned") = Some(ctx.session_id().to_string()); + + let Some(name) = required else { + return Ok(()); + }; + + if ctx.scalar_functions().contains_key(name) { + probe.resolutions.fetch_add(1, Ordering::SeqCst); + return Ok(()); + } + + // A fresh SessionContext still carries every built-in, so report the count + // rather than the whole registry. + plan_err!( + "datafusion-ffi-example: decode could not resolve scalar function '{name}' \ + in the task context it was handed (session '{}', {} scalar functions registered)", + ctx.session_id(), + ctx.scalar_functions().len() + ) +} diff --git a/examples/datafusion-ffi-example/src/table_function.rs b/examples/datafusion-ffi-example/src/table_function.rs index 55543cb59..e653aeab1 100644 --- a/examples/datafusion-ffi-example/src/table_function.rs +++ b/examples/datafusion-ffi-example/src/table_function.rs @@ -48,7 +48,7 @@ impl MyTableFunction { session: Bound, ) -> PyResult> { let func = self.clone(); - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_table_function") diff --git a/examples/datafusion-ffi-example/src/table_provider.rs b/examples/datafusion-ffi-example/src/table_provider.rs index 5756e6d02..ef6430e29 100644 --- a/examples/datafusion-ffi-example/src/table_provider.rs +++ b/examples/datafusion-ffi-example/src/table_provider.rs @@ -103,7 +103,7 @@ impl MyTableProvider { .create_table() .map_err(|e: DataFusionError| PyRuntimeError::new_err(e.to_string()))?; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_TableProvider::new_with_ffi_codec(Arc::new(provider), false, None, codec); diff --git a/examples/datafusion-ffi-example/src/table_provider_factory.rs b/examples/datafusion-ffi-example/src/table_provider_factory.rs index 71dfd73ca..df0845119 100644 --- a/examples/datafusion-ffi-example/src/table_provider_factory.rs +++ b/examples/datafusion-ffi-example/src/table_provider_factory.rs @@ -77,7 +77,7 @@ impl MyTableProviderFactory { py: Python<'py>, codec: Bound, ) -> PyResult> { - let codec = ffi_logical_codec_from_pycapsule(codec)?; + let codec = ffi_logical_codec_from_pycapsule(codec, None)?; let factory = Arc::clone(&self.inner) as Arc; let factory = FFI_TableProviderFactory::new_with_ffi_codec(factory, None, codec); diff --git a/examples/datafusion-ffi-query-planner-example/Cargo.toml b/examples/datafusion-ffi-query-planner-example/Cargo.toml new file mode 100644 index 000000000..4d02c69f1 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -0,0 +1,49 @@ +# 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. + +[package] +name = "datafusion-ffi-query-planner-example" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description.workspace = true +homepage.workspace = true +repository.workspace = true +publish = false + +[dependencies] +datafusion = { workspace = true } +datafusion-catalog = { workspace = true, default-features = false } +datafusion-common = { workspace = true, default-features = false } +datafusion-ffi = { workspace = true } +datafusion-session = { workspace = true } +async-trait = { workspace = true } +datafusion-python-util.workspace = true +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "abi3-py310", +] } +pyo3-log = { workspace = true } + +[build-dependencies] +pyo3-build-config = { workspace = true } + +[lib] +name = "datafusion_ffi_query_planner_example" +crate-type = ["cdylib", "rlib"] diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md new file mode 100644 index 000000000..72f96bb8e --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -0,0 +1,60 @@ + + +# DataFusion Python FFI query planner example + +This crate is an independent query-planner Python extension. Together with [`../datafusion-ffi-example`](../datafusion-ffi-example/) it demonstrates a real three-library plan exchange: + +- **A — `datafusion-python`:** owns the session and final execution. +- **B — `datafusion-ffi-example`:** owns a table provider, UDF, and provider codecs. +- **C — this crate:** owns the query planner and its custom configuration. + +Two extension crates are used rather than placing the planner in the provider crate. Loading distinct `cdylib` images gives each library a distinct DataFusion marker and proves that foreign sessions, providers, and plans survive the actual ABI boundary. + +## Running the example + +From the repository root, build and install all three extensions, then run the +integration tests: + +```bash +maturin develop --uv +uv run maturin develop --manifest-path examples/datafusion-ffi-example/Cargo.toml +uv run maturin develop \ + --manifest-path examples/datafusion-ffi-query-planner-example/Cargo.toml +uv run pytest \ + examples/datafusion-ffi-query-planner-example/python/tests/_test*.py +``` + +The integration test follows this setup: + +```python +config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) +ctx = SessionContext(config) +ctx = ctx.with_logical_extension_codec(provider_logical_codec) +ctx = ctx.with_physical_extension_codec(provider_physical_codec) +ctx.register_table("numbers", provider) +ctx.register_udf(provider_udf) +ctx.set_query_planner(MyQueryPlanner()) +``` + +`MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. + +The provider's codec pair is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against it, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](../../docs/source/contributor-guide/ffi.md#rebinding-a-planners-codecs-is-one-level-deep). + +For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. diff --git a/examples/datafusion-ffi-query-planner-example/build.rs b/examples/datafusion-ffi-query-planner-example/build.rs new file mode 100644 index 000000000..4878d8b0e --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/build.rs @@ -0,0 +1,20 @@ +// 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. + +fn main() { + pyo3_build_config::add_extension_module_link_args(); +} diff --git a/examples/datafusion-ffi-query-planner-example/pyproject.toml b/examples/datafusion-ffi-query-planner-example/pyproject.toml new file mode 100644 index 000000000..9e34b4cd4 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/pyproject.toml @@ -0,0 +1,32 @@ +# 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. + +[build-system] +requires = ["maturin>=1.6,<2.0"] +build-backend = "maturin" + +[project] +name = "datafusion_ffi_query_planner_example" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] +dynamic = ["version"] + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py new file mode 100644 index 000000000..d046f67a6 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -0,0 +1,669 @@ +# 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. + +from __future__ import annotations + +import gc + +import pyarrow as pa +import pytest +from datafusion import Expr, SessionConfig, SessionContext, col, udf +from datafusion_ffi_example import ( + IsNullUDF, + MyCatalogProvider, + MyLogicalExtensionCodec, + MyPhysicalExtensionCodec, + MyPhysicalOptimizerRule, + MyTableProvider, +) +from datafusion_ffi_query_planner_example import MyPlannerConfig, MyQueryPlanner + + +def configured_context(max_rows: int): + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows)) + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + return ctx, logical_codec, physical_codec + + +HOST_ONLY_UDF = "my_custom_is_null" +"""Scalar function registered on the host session and nowhere else.""" + +UNREGISTERED_UDF = "not_registered_anywhere" + + +def probe_context( + *, + logical_requires: str | None = None, + physical_requires: str | None = None, + max_rows: int = 3, +): + """Three-library context whose codecs read the task context they are given. + + ``require_udf_on_decode`` makes each codec resolve a scalar function from + the ``TaskContext`` handed to its FFI decode callback, which is otherwise + unobservable: the example codecs restore objects from a token registry and + never look at the registry they are passed. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows)) + logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=logical_requires) + physical_codec = MyPhysicalExtensionCodec(require_udf_on_decode=physical_requires) + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + ctx.set_query_planner(MyQueryPlanner()) + return ctx, logical_codec, physical_codec + + +def test_logical_codec_resolves_a_host_registered_udf(): + """``try_decode_table_provider`` sees the host session's registry. + + The codec takes its task context provider from the session it is installed + on, so a function the host registered is resolvable inside a decode + callback running in the other library. + """ + ctx, logical_codec, _physical_codec = probe_context(logical_requires=HOST_ONLY_UDF) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert logical_codec.table_provider_decode_calls() > 0 + assert logical_codec.task_context_udf_resolutions() > 0 + + +def test_physical_codec_resolves_a_host_registered_udf(): + """``try_decode`` sees the host session's registry, as above.""" + ctx, _logical_codec, physical_codec = probe_context(physical_requires=HOST_ONLY_UDF) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert physical_codec.execution_plan_decode_calls() > 0 + assert physical_codec.task_context_udf_resolutions() > 0 + + +def test_codec_still_reports_a_name_registered_nowhere(): + """Negative control: resolution really is a lookup, not an unconditional pass.""" + ctx, _logical_codec, _physical_codec = probe_context( + logical_requires=UNREGISTERED_UDF + ) + + with pytest.raises( + Exception, match=rf"could not resolve scalar function '{UNREGISTERED_UDF}'" + ): + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + + +def test_codec_sees_a_udf_registered_after_it_was_installed(): + """The provider is a live handle to the session, not a snapshot of it.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF) + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + # Registered after the codec was installed and bound to this session. + ctx.register_udf(udf(IsNullUDF())) + ctx.set_query_planner(MyQueryPlanner()) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert logical_codec.task_context_udf_resolutions() > 0 + + +def test_codec_sees_a_udf_registered_after_the_planner(): + """Installing a planner does not detach the codec from the session. + + The codec is bound to the session before the planner is installed and the + function is registered afterwards. Installing writes through + ``state_ref()`` rather than deriving a new ``SessionContext``, so the + codec's task context provider still points at the one live session and + sees the later registration. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF) + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.set_query_planner(MyQueryPlanner()) + # Registered after the planner, on the same session the codec is bound to. + ctx.register_udf(udf(IsNullUDF())) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert logical_codec.table_provider_decode_calls() > 0 + assert logical_codec.task_context_udf_resolutions() > 0 + + +@pytest.mark.parametrize("codecs_first", [True, False]) +def test_registered_providers_survive_a_planner_install(codecs_first: bool): + """Neither write path may orphan a previously registered provider. + + A foreign catalog provider is handed a codec carrying a *weak* + ``FFI_TaskContextProvider`` pointing at the session it was registered on, + and upgrades it on every ``supports_filters_pushdown`` and every ``scan``. + That codec lives inside the registered ``FFI_CatalogProvider``, so nothing + on the Python side can reach it to rebind it. Deriving a replacement + ``SessionContext`` would drop the allocation those handles point at, and + the next query would fail with ``TaskContextProvider went out of scope over + FFI boundary``. Installing in place keeps the one allocation alive. + + Both parameters exercise that, because both write ``SessionState``: + ``set_query_planner`` installs the planner, and installing a codec on a + session that already has one rebuilds that planner against the new codec. + + The ``WHERE`` clause is load-bearing -- it forces filter pushdown, which + upgrades the weak handle during logical optimization. It is also why both + codecs have to be installed: without them the query fails at plan + serialization with ``LogicalExtensionCodec is not provided``, which would + mask a dangling handle rather than expose it. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=10)) + ctx = SessionContext(config) + ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider()) + + def install_codecs(ctx): + ctx = ctx.with_logical_extension_codec(MyLogicalExtensionCodec()) + return ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) + + if codecs_first: + ctx = install_codecs(ctx) + ctx.set_query_planner(MyQueryPlanner()) + else: + ctx.set_query_planner(MyQueryPlanner()) + ctx = install_codecs(ctx) + gc.collect() + + batches = ctx.sql( + "SELECT units FROM ffi_catalog.my_schema.my_table WHERE units > 5" + ).collect() + assert sorted(v for b in batches for v in b.column(0).to_pylist()) == [ + 7, + 10, + 20, + 30, + ] + + +def codec_context(max_rows: int = 3): + """Context with both example codecs installed and a table to scan. + + Unlike :func:`probe_context` the codecs ask for no function, so the only + thing they record is the session id of the task context they are handed. + No planner yet -- the caller installs one. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows)) + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + return ctx, logical_codec, physical_codec + + +def test_installing_a_planner_keeps_the_session_id(): + """A session and its decode callbacks must agree on the session id. + + Installing a planner rebuilds ``SessionState`` through + ``SessionStateBuilder``, which mints a fresh id unless handed one, while + ``SessionContext`` caches its id in a field of its own. Dropping the id + there leaves the session reporting one id from ``session_id()`` and a + different one from every ``TaskContext`` it gives a foreign codec. + + Asserting on ``session_id()`` alone cannot catch that: it reads the cached + copy, which stays correct either way. The codec-side id is the only + observable that moves, which is what makes this worth a test rather than a + one-line equality check. + """ + ctx, logical_codec, physical_codec = codec_context() + session_id = ctx.session_id() + + ctx.set_query_planner(MyQueryPlanner()) + assert ctx.session_id() == session_id + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + assert physical_codec.last_task_context_session_id() == session_id + + +def test_adding_a_physical_optimizer_rule_keeps_the_session_id(): + """The same guarantee for the other in-place ``SessionState`` rewrite. + + ``add_physical_optimizer_rule`` also rebuilds ``SessionState`` and writes + it back, so a regenerated id would desync a context from itself. + """ + ctx, logical_codec, physical_codec = codec_context() + ctx.set_query_planner(MyQueryPlanner()) + session_id = ctx.session_id() + + ctx.add_physical_optimizer_rule(MyPhysicalOptimizerRule()) + assert ctx.session_id() == session_id + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + assert physical_codec.last_task_context_session_id() == session_id + + +def test_replacing_a_planner_keeps_the_session_id(): + """Installing repeatedly must not drift the id. + + Each install rebuilds ``SessionState``, so an id carried over only on the + first write would still be lost by the second. + """ + ctx, logical_codec, _physical_codec = codec_context() + session_id = ctx.session_id() + + ctx.set_query_planner(MyQueryPlanner()) + ctx.set_query_planner(MyQueryPlanner()) + assert ctx.session_id() == session_id + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + + +@pytest.mark.parametrize("raw_capsule", [False, True]) +def test_three_library_query_planner(raw_capsule: bool): + """Host, provider, and planner exchange a real non-empty plan over FFI.""" + ctx, logical_codec, physical_codec = configured_context(max_rows=3) + planner = MyQueryPlanner() + exported_planner = ( + planner.__datafusion_query_planner__(ctx) if raw_capsule else planner + ) + ctx.set_query_planner(exported_planner) + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert batches[0].column(1).to_pylist() == [False, False, False] + assert planner.last_max_rows() == 3 + + ctx.sql("SET ffi_query_planner.max_rows = 2").collect() + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.last_max_rows() == 2 + + assert planner.plan_calls() >= 2 + assert planner.foreign_session_observed() + assert planner.foreign_provider_observed() + assert planner.foreign_plan_observed() + assert logical_codec.table_provider_encode_calls() > 0 + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_encode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + +def test_spawning_plan_across_three_libraries(): + """A plan that spawns Tokio tasks survives the full three-library round trip. + + ``target_partitions`` above one puts a ``RepartitionExec`` under the + aggregate, and that operator spawns tasks while it runs. This exercises the + codecs on a multi-node plan rather than the bare scan the other tests use. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=100)) + config = config.with_target_partitions(4) + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 3)) + + planner = MyQueryPlanner() + ctx.set_query_planner(planner) + + batches = ctx.sql( + 'SELECT "A" % 2 AS parity, count(*) AS n FROM numbers GROUP BY 1 ORDER BY 1' + ).collect() + counts = { + row[0]: row[1] + for batch in batches + for row in zip( + batch.column(0).to_pylist(), batch.column(1).to_pylist(), strict=True + ) + } + assert sum(counts.values()) == 6 + 7 + 8 + assert planner.plan_calls() > 0 + assert planner.foreign_provider_observed() + + +def test_planner_layers_on_the_session_planner(): + """A planner can wrap the one already installed and delegate to it. + + The capsule has to be captured before this planner is installed, because + ``__datafusion_query_planner__`` exports whichever planner is installed when + it is called. Capturing it afterwards would hand the planner a handle to + itself, and planning would recurse. + """ + ctx, logical_codec, physical_codec = configured_context(max_rows=3) + fallback = ctx.__datafusion_query_planner__() + planner = MyQueryPlanner(fallback=fallback) + # The capsule's FFI codecs hold weak handles to this session. Installing in + # place keeps that session alive, so the capsule stays usable; deriving a + # replacement here would fail with "TaskContextProvider went out of scope + # over FFI boundary". + ctx.set_query_planner(planner) + gc.collect() + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert planner.plan_calls() > 0 + assert planner.used_fallback() + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + +def test_a_planner_can_fall_back_to_another_planner_library(): + """A fallback may be another foreign planner, not only a session. + + The fallback is imported when this planner is installed rather than when + it is constructed, so its own getter receives the session. Importing it at + construction time would mean calling that getter with no session, which + only a ``SessionContext`` or a raw capsule tolerates -- and layering on + another planner is the case a distributed engine actually needs. + """ + ctx, logical_codec, physical_codec = configured_context(max_rows=3) + inner = MyQueryPlanner() + outer = MyQueryPlanner(fallback=inner) + ctx.set_query_planner(outer) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert outer.plan_calls() > 0 + assert outer.used_fallback() + # The delegation reached the inner planner rather than stopping at the + # default physical planner. + assert inner.plan_calls() > 0 + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + +def test_a_session_fallback_delegates_to_its_installed_planner(): + """Passing a SessionContext delegates to whatever planner it holds.""" + ctx, _logical_codec, _physical_codec = configured_context(max_rows=3) + first = MyQueryPlanner() + ctx.set_query_planner(first) + + second = MyQueryPlanner(fallback=ctx) + ctx.set_query_planner(second) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert second.used_fallback() + assert first.plan_calls() > 0 + + +def test_observations_accumulate_across_queries(): + """A later plain query must not retract what an earlier query observed. + + The ``*_observed`` accessors answer "was this ever seen". They are written + with ``fetch_or`` rather than ``store`` so a query that touches no foreign + object cannot clear a flag an earlier one set. Written with ``store``, + ``SELECT 1`` here clears ``foreign_provider_observed``, and every other + test asserting these flags after more than one query is a coincidence away + from failing. + """ + ctx, _logical_codec, _physical_codec = configured_context(max_rows=3) + planner = MyQueryPlanner() + ctx.set_query_planner(planner) + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner.foreign_session_observed() + assert planner.foreign_provider_observed() + assert planner.foreign_plan_observed() + + # Touches no table, so this plan has no foreign provider of its own. + ctx.sql("SELECT 1").collect() + assert planner.foreign_session_observed() + assert planner.foreign_provider_observed() + assert planner.foreign_plan_observed() + + # `last_max_rows` is deliberately not cumulative; it reports the last plan. + assert planner.last_max_rows() == 3 + assert planner.plan_calls() >= 2 + + +def test_second_planner_replaces_the_first(): + """A session holds exactly one planner, so installing another replaces it.""" + ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) + first = MyQueryPlanner() + second = MyQueryPlanner() + ctx.set_query_planner(first) + ctx.set_query_planner(second) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert second.plan_calls() > 0 + assert first.plan_calls() == 0 + + +def test_the_planner_reaches_every_handle_on_the_session(): + """The planner is session state, so all handles on that session use it. + + ``set_query_planner`` writes through ``state_ref()``. A context returned by + an earlier ``with_*`` call shares that session, so it plans through the new + planner too -- there is one session and one planner, not a family of + diverging copies. + """ + ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) + # Shares the session with `ctx`; predates the planner. + sibling = ctx.with_python_udf_inlining(enabled=False) + + planner = MyQueryPlanner() + ctx.set_query_planner(planner) + + batches = sibling.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.plan_calls() > 0 + + +def test_installed_codecs_outlive_python_exporters(): + ctx, logical_codec, physical_codec = configured_context(max_rows=2) + del logical_codec, physical_codec + gc.collect() + + ctx.set_query_planner(MyQueryPlanner()) + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_provider_codecs_can_be_installed_after_planner(): + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + planner = MyQueryPlanner() + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx.set_query_planner(planner) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 4, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.last_max_rows() == 2 + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + +def test_a_discarded_derived_context_still_rebinds_the_planner(): + """Installing a codec rebinds the planner even if the handle is thrown away. + + `with_logical_extension_codec` returns a context sharing this session, and + rebuilding the installed planner against the new codec happens on that + shared session rather than on the returned handle. So the rebind outlives + the handle, and the codec below takes effect on `ctx` despite `ctx`'s own + codec field never changing. + + That is spooky enough to be worth pinning as a decision. It is also forced: + `FFI_QueryPlanner` holds its codecs by value, so a planner cannot read the + session's current codecs at plan time and the rebuild has to be eager. + + A fresh codec instance is what makes it observable -- it carries its own + counters, and the planner encodes the outbound logical plan with whichever + codec it is holding. + """ + ctx, _logical_codec, _physical_codec = configured_context(max_rows=3) + ctx.set_query_planner(MyQueryPlanner()) + + later = MyLogicalExtensionCodec() + # Deliberately discarded. The rebind still lands on the shared session. + ctx.with_logical_extension_codec(later) + gc.collect() + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert later.table_provider_encode_calls() > 0 + + +def test_the_planner_and_the_handle_can_hold_different_codecs(): + """One session, two codecs in effect, depending on the path taken. + + The planner is session state and carries whichever codecs installed it + last -- here, ones that arrived through a handle that was discarded. + Everything else on a context uses that context's own codec field, which + the discarded handle never touched. So `Expr.to_bytes(ctx)` and + `ctx.sql(...)` encode with different codecs on the same `ctx`. + + Chaining ``ctx = ctx.with_...(...)`` keeps the two in step; this pins what + happens when they are allowed to diverge. + + Inlining has to be off for the assertion to say anything: with it on, a + Python UDF is encoded inline by ``PythonLogicalCodec`` and never reaches + the installed codec's ``try_encode_udf``. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + handle_codec = MyLogicalExtensionCodec() + ctx = SessionContext(config).with_python_udf_inlining(enabled=False) + ctx = ctx.with_logical_extension_codec(handle_codec) + ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.set_query_planner(MyQueryPlanner()) + + planner_codec = MyLogicalExtensionCodec() + # Discarded, but the planner keeps its codec. + ctx.with_logical_extension_codec(planner_codec) + gc.collect() + + identity = udf( + lambda arr: arr, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="identity_i64", + ) + ctx.register_udf(identity) + Expr.to_bytes(identity(col("A")), ctx) + + # Serializing through `ctx` uses `ctx`'s own codec field. + assert handle_codec.encode_udf_calls() > 0 + assert planner_codec.encode_udf_calls() == 0 + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + + # Planning through the same `ctx` uses the codec the planner was rebound to. + assert planner_codec.table_provider_encode_calls() > 0 + assert handle_codec.table_provider_encode_calls() == 0 + + +def test_an_unchanged_inlining_setting_leaves_the_planner_alone(): + """A no-op toggle must not rebind the session's planner. + + ``with_python_udf_inlining`` rebuilds the handle's codecs and rebinds the + session's planner to them. Asking for the setting a context already has + changes nothing, so it must not pay that side effect. + + Observable only once the planner is holding some *other* handle's codec: + without the guard, a defensive no-op toggle on `ctx` drags the planner back + onto `ctx`'s codec and silently undoes the install below. The rebuilt + codecs otherwise wrap the same inner codec, so nothing else distinguishes + the two paths. + """ + ctx, handle_codec, _physical_codec = codec_context() + ctx.set_query_planner(MyQueryPlanner()) + + planner_codec = MyLogicalExtensionCodec() + ctx.with_logical_extension_codec(planner_codec) # discarded; planner keeps it + gc.collect() + + # The default is on, so this asks for what `ctx` already has. + ctx.with_python_udf_inlining(enabled=True) + gc.collect() + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner_codec.table_provider_encode_calls() > 0 + assert handle_codec.table_provider_encode_calls() == 0 + + +def test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs(): + """A planner is built against the codecs of the handle it is installed from. + + The sequel to the test above, and the trap it sets up. Once a discarded + derived handle has rebound the session's planner to its codec, installing + the same planner again from the *original* handle rebuilds it against that + handle's codec instead -- which never changed. The session's planner tracks + whichever handle wrote it last, not the newest codec installed anywhere. + + So "re-install the planner after installing a codec" only repairs anything + when it is done from the handle holding the new codec. + """ + ctx, original_logical, _physical_codec = codec_context() + planner = MyQueryPlanner() + ctx.set_query_planner(planner) + + later = MyLogicalExtensionCodec() + # Deliberately discarded, exactly as in the test above. + ctx.with_logical_extension_codec(later) + gc.collect() + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert later.table_provider_encode_calls() > 0 + assert original_logical.table_provider_encode_calls() == 0 + + # `ctx`'s own codec field never changed, so this rebuilds the planner + # against `original_logical` and drops `later` from the session's planner. + ctx.set_query_planner(planner) + encodes_by_later = later.table_provider_encode_calls() + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert original_logical.table_provider_encode_calls() > 0 + assert later.table_provider_encode_calls() == encodes_by_later + + +def test_query_planner_requires_provider_codec(): + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx = SessionContext(config) + ctx.register_table("numbers", MyTableProvider(1, 3, 1)) + ctx.set_query_planner(MyQueryPlanner()) + + with pytest.raises(Exception, match=r"LogicalExtensionCodec|TableProvider"): + ctx.sql('SELECT "A" FROM numbers').collect() + + +@pytest.mark.parametrize("max_rows", ["0", "oops"]) +def test_query_planner_rejects_invalid_config(max_rows: str): + ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) + ctx.set_query_planner(MyQueryPlanner()) + + with pytest.raises(Exception, match=r"max_rows|Invalid value"): + ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/conftest.py b/examples/datafusion-ffi-query-planner-example/python/tests/conftest.py new file mode 100644 index 000000000..68f8057af --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/python/tests/conftest.py @@ -0,0 +1,42 @@ +# 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. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from collections.abc import Generator + from typing import Any + + +class _FailOnWarning(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if record.levelno >= logging.WARNING: + err = f"Unexpected log warning from '{record.name}': {self.format(record)}" + raise AssertionError(err) + + +@pytest.fixture(autouse=True) +def fail_on_log_warnings() -> Generator[None, Any, None]: + handler = _FailOnWarning() + logging.root.addHandler(handler) + yield + logging.root.removeHandler(handler) diff --git a/examples/datafusion-ffi-query-planner-example/src/config.rs b/examples/datafusion-ffi-query-planner-example/src/config.rs new file mode 100644 index 000000000..ecfa4b943 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/config.rs @@ -0,0 +1,112 @@ +// 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. + +use std::any::Any; + +use datafusion_common::config::{ + ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit, +}; +use datafusion_common::{DataFusionError, config_err}; +use datafusion_ffi::config::extension_options::FFI_ExtensionOptions; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +#[pyclass( + from_py_object, + name = "MyPlannerConfig", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Clone, Debug)] +pub(crate) struct MyPlannerConfig { + pub max_rows: usize, +} + +#[pymethods] +impl MyPlannerConfig { + #[new] + #[pyo3(signature = (max_rows=10))] + fn new(max_rows: usize) -> Self { + Self { max_rows } + } + + fn __datafusion_extension_options__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let mut config = FFI_ExtensionOptions::default(); + config + .add_config(self) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + PyCapsule::new_with_value(py, config, cr"datafusion_extension_options") + } +} + +impl Default for MyPlannerConfig { + fn default() -> Self { + Self { max_rows: 10 } + } +} + +impl ConfigExtension for MyPlannerConfig { + const PREFIX: &'static str = "ffi_query_planner"; +} + +impl ExtensionOptions for MyPlannerConfig { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> { + ConfigField::set(self, key, value) + } + + fn entries(&self) -> Vec { + vec![ConfigEntry { + key: "max_rows".to_owned(), + value: Some(self.max_rows.to_string()), + description: "Maximum rows returned by the example query planner", + }] + } +} + +impl ConfigField for MyPlannerConfig { + fn visit(&self, visitor: &mut V, _key: &str, _description: &'static str) { + self.max_rows.visit( + visitor, + "max_rows", + "Maximum rows returned by the example query planner", + ); + } + + fn set(&mut self, key: &str, value: &str) -> Result<(), DataFusionError> { + let (key, rem) = key.split_once('.').unwrap_or((key, "")); + match key { + "max_rows" => self.max_rows.set(rem, value), + _ => config_err!("Config value '{key}' not found on MyPlannerConfig"), + } + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs new file mode 100644 index 000000000..c505c1ce7 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -0,0 +1,32 @@ +// 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. + +use pyo3::prelude::*; + +use crate::config::MyPlannerConfig; +use crate::planner::MyQueryPlanner; + +mod config; +mod planner; + +#[pymodule] +fn datafusion_ffi_query_planner_example(m: &Bound<'_, PyModule>) -> PyResult<()> { + pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs new file mode 100644 index 000000000..67262e39c --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -0,0 +1,311 @@ +// 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. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use async_trait::async_trait; +use datafusion::common::DataFusionError; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::limit::GlobalLimitExec; +use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; +use datafusion_catalog::default_table_source::source_as_provider; +use datafusion_ffi::config::ExtensionOptionsFFIProvider; +use datafusion_ffi::execution_plan::ForeignExecutionPlan; +use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_ffi::session::ForeignSession; +use datafusion_ffi::table_provider::ForeignTableProvider; +use datafusion_python_util::{ + ffi_logical_codec_from_pycapsule, ffi_physical_codec_from_pycapsule, + ffi_query_planner_from_pycapsule, +}; +use datafusion_session::{QueryPlanner, Session}; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +use crate::config::MyPlannerConfig; + +/// What the planner saw, accumulated across every call rather than reset each +/// time. +/// +/// Two kinds of field here, and mixing them up is easy. `last_max_rows` +/// reports the most recent value, as its name says. Everything else is +/// cumulative: a count, or a "did this ever happen" flag written with +/// `fetch_or` so a later plan cannot retract an earlier observation. Tests +/// assert after running more than one query, so a flag that only described the +/// most recent plan would be answering a different question than the one its +/// accessor name asks. +#[derive(Default)] +struct PlannerObservations { + plan_calls: AtomicUsize, + last_max_rows: AtomicUsize, + foreign_session: AtomicBool, + foreign_provider: AtomicBool, + foreign_plan: AtomicBool, + /// Only ever set to `true`, so it is already cumulative. + used_fallback: AtomicBool, +} + +impl fmt::Debug for PlannerObservations { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PlannerObservations") + .field("plan_calls", &self.plan_calls) + .field("last_max_rows", &self.last_max_rows) + .finish_non_exhaustive() + } +} + +fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool { + if let LogicalPlan::TableScan(scan) = plan + && let Ok(provider) = source_as_provider(&scan.source) + && provider.downcast_ref::().is_some() + { + return true; + } + plan.inputs() + .iter() + .any(|input| logical_plan_has_foreign_provider(input)) +} + +fn physical_plan_has_foreign_plan(plan: &Arc) -> bool { + plan.is::() + || plan + .children() + .iter() + .any(|child| physical_plan_has_foreign_plan(child)) +} + +/// The row limit as the host spells it, where `MyPlannerConfig` is registered as +/// an ordinary config extension under its own `ConfigExtension::PREFIX`. +const MAX_ROWS_KEY: &str = "ffi_query_planner.max_rows"; + +/// The same setting as it appears once the session has crossed the FFI +/// boundary. Rebuilding a `ConfigOptions` on this side parks every foreign +/// extension inside a single `FFI_ExtensionOptions`, which is itself a config +/// extension namespaced under `datafusion_ffi`, so `ConfigOptions::entries` +/// reports the key with both prefixes. +const FFI_MAX_ROWS_KEY: &str = "datafusion_ffi.ffi_query_planner.max_rows"; + +fn planner_config(session: &dyn Session) -> datafusion::common::Result { + let options = session.config_options(); + + // Prefer the raw entry. `local_or_ffi_extension` discards a value it cannot + // parse and hands back `MyPlannerConfig::default()`, which would quietly turn + // a typo into a different row limit instead of reporting it. + let config = match options + .entries() + .into_iter() + .find(|entry| entry.key == MAX_ROWS_KEY || entry.key == FFI_MAX_ROWS_KEY) + { + Some(entry) => { + let value = entry.value.ok_or_else(|| { + DataFusionError::Configuration(format!("{} must have a value", entry.key)) + })?; + let max_rows = value.parse::().map_err(|err| { + DataFusionError::Configuration(format!( + "Invalid value '{value}' for {}: {err}", + entry.key + )) + })?; + MyPlannerConfig { max_rows } + } + None => options + .local_or_ffi_extension::() + .unwrap_or_default(), + }; + + // Validate after both paths so the fallback cannot smuggle in a limit that + // the direct path rejects. + if config.max_rows == 0 { + return Err(DataFusionError::Configuration(format!( + "{MAX_ROWS_KEY} must be greater than zero" + ))); + } + + Ok(config) +} + +#[derive(Debug)] +struct DistributedQueryPlanner { + observations: Arc, + /// Planner to hand the work to instead of planning here. + /// + /// This is how a real planner layers on top of an existing one. The capsule + /// must be captured from the session *before* this planner is installed: + /// `SessionContext.__datafusion_query_planner__` exports whatever planner + /// is installed at the time it is called, so capturing it afterwards would + /// hand this planner a handle to itself. + /// + /// Note that `Session::create_physical_plan` cannot be used for this. It + /// dispatches through the session's installed query planner, so calling it + /// from inside that planner recurses until the stack overflows. + fallback: Option>, +} + +#[async_trait] +impl QueryPlanner for DistributedQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> datafusion::common::Result> { + self.observations.plan_calls.fetch_add(1, Ordering::SeqCst); + // `fetch_or`, not `store`: these answer "was this ever seen", so a + // later plan that happens not to touch a foreign object must not + // retract what an earlier one observed. A bare `SELECT 1` after a + // scan of a foreign provider would otherwise clear the flag. + self.observations + .foreign_session + .fetch_or(session.as_any().is::(), Ordering::SeqCst); + self.observations.foreign_provider.fetch_or( + logical_plan_has_foreign_provider(logical_plan), + Ordering::SeqCst, + ); + + let config = planner_config(session)?; + self.observations + .last_max_rows + .store(config.max_rows, Ordering::SeqCst); + + let plan = match self.fallback.as_ref() { + Some(fallback) => { + self.observations + .used_fallback + .store(true, Ordering::SeqCst); + fallback.create_physical_plan(logical_plan, session).await? + } + None => { + DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, session) + .await? + } + }; + self.observations + .foreign_plan + .fetch_or(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst); + + Ok(Arc::new(GlobalLimitExec::new( + plan, + 0, + Some(config.max_rows), + ))) + } +} + +#[pyclass( + from_py_object, + name = "MyQueryPlanner", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Debug, Default, Clone)] +pub(crate) struct MyQueryPlanner { + observations: Arc, + /// Held as the Python object rather than an imported planner, and resolved + /// in `__datafusion_query_planner__` where a session is in hand. + /// + /// Importing it here would mean calling its getter with no session, which + /// only a `SessionContext` or a raw capsule accepts. Another foreign + /// planner -- the case that matters, since layering is the whole point of + /// a fallback -- implements the same protocol this type does and requires + /// the argument. + fallback: Option>>, +} + +#[pymethods] +impl MyQueryPlanner { + /// Build a planner, optionally layered on top of an existing one. + /// + /// `fallback` takes anything exporting `__datafusion_query_planner__`: + /// another planner library, a `SessionContext`, or a raw capsule. It is + /// imported when this planner is installed, not here, so that the session + /// can be handed to its getter. + /// + /// Passing a `SessionContext` delegates to whichever planner that context + /// holds at install time. If you instead capture a capsule with + /// `ctx.__datafusion_query_planner__()`, capture it *before* installing + /// this planner on that context, or the capsule will describe this planner + /// and planning will recurse. + #[new] + #[pyo3(signature = (fallback=None))] + fn new(fallback: Option>) -> Self { + Self { + fallback: fallback.map(|obj| Arc::new(obj.unbind())), + ..Self::default() + } + } + + fn used_fallback(&self) -> bool { + self.observations.used_fallback.load(Ordering::SeqCst) + } + + fn plan_calls(&self) -> usize { + self.observations.plan_calls.load(Ordering::SeqCst) + } + + fn last_max_rows(&self) -> usize { + self.observations.last_max_rows.load(Ordering::SeqCst) + } + + fn foreign_session_observed(&self) -> bool { + self.observations.foreign_session.load(Ordering::SeqCst) + } + + fn foreign_provider_observed(&self) -> bool { + self.observations.foreign_provider.load(Ordering::SeqCst) + } + + fn foreign_plan_observed(&self) -> bool { + self.observations.foreign_plan.load(Ordering::SeqCst) + } + + /// Export the planner, bound to the session it is being installed on. + /// + /// The codecs come off `session` rather than being built here. They carry + /// the host's `TaskContextProvider`, so this library never constructs a + /// `SessionContext`, and `with_query_planner` would rebind them to the + /// running session anyway. + fn __datafusion_query_planner__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + // Resolved here rather than in `new` so the fallback's own getter + // receives the session, which is what the protocol requires of every + // implementation other than a `SessionContext`. + let fallback = self + .fallback + .as_ref() + .map(|planner| { + ffi_query_planner_from_pycapsule(planner.bind(py), Some(&session)) + .map(|ffi| -> Arc { (&ffi).into() }) + }) + .transpose()?; + + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + fallback, + }); + let logical_codec = ffi_logical_codec_from_pycapsule(session.clone(), None)?; + let physical_codec = ffi_physical_codec_from_pycapsule(session, None)?; + let ffi = FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec, physical_codec); + PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") + } +} diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 94b2bb1c6..2f2cc6119 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -145,6 +145,18 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 +class QueryPlannerExportable(Protocol): + """Type hint for object that has a __datafusion_query_planner__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically + produced by a separate compiled extension. ``session`` is the + :py:class:`SessionContext` the planner is being installed on; take the + extension codecs from it rather than building your own. + """ + + def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 + + class SessionConfig: """Session configuration options.""" @@ -1759,6 +1771,51 @@ def add_physical_optimizer_rule( """ self.ctx.add_physical_optimizer_rule(rule) + def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> None: + """Install a custom query planner on this session. + + The planner is imported through its ``__datafusion_query_planner__`` + PyCapsule and installed on this context, in the same way + :meth:`~SessionContext.add_physical_optimizer_rule` installs a rule. + The query planner is part of the session state, so it applies to this + context and to every context sharing its session — including ones + already returned by + :meth:`~SessionContext.with_logical_extension_codec` and friends. + + A session holds exactly one planner, so calling this again replaces the + previous one rather than layering. To chain planners, have the new + planner wrap the capsule from + :meth:`~SessionContext.__datafusion_query_planner__`, captured + *before* the new planner is installed. + + Install any extension codecs before a layered planner. Installing a + codec afterwards rebuilds the installed planner against it, but not the + fallback inside it, which keeps the codecs it was imported with. Note + also that the planner is built against the codecs of the context this + method is called on, so installing the same planner again on a different + handle rebinds the session's planner to *that* handle's codecs. + + Args: + planner: Object exposing ``__datafusion_query_planner__`` (see + :class:`QueryPlannerExportable`) or a raw + ``datafusion_query_planner`` PyCapsule. + + Examples: + >>> from my_extension import DistributedQueryPlanner # doctest: +SKIP + >>> ctx = SessionContext() + >>> ctx.set_query_planner(DistributedQueryPlanner()) # doctest: +SKIP + >>> ctx.sql("SELECT * FROM remote_table").collect() # doctest: +SKIP + + Layer a planner on top of the one already installed by capturing + the existing planner first: + + >>> fallback = ctx.__datafusion_query_planner__() # doctest: +SKIP + >>> ctx.set_query_planner( + ... DistributedQueryPlanner(fallback=fallback) + ... ) # doctest: +SKIP + """ + self.ctx.set_query_planner(planner) + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. @@ -2178,9 +2235,22 @@ def __datafusion_task_context_provider__(self) -> Any: """Access the PyCapsule FFI_TaskContextProvider.""" return self.ctx.__datafusion_task_context_provider__() - def __datafusion_logical_extension_codec__(self) -> Any: - """Access the PyCapsule FFI_LogicalExtensionCodec.""" - return self.ctx.__datafusion_logical_extension_codec__() + def __datafusion_logical_extension_codec__(self, session: Any = None) -> Any: + """Access the PyCapsule FFI_LogicalExtensionCodec. + + ``session`` is accepted so a context satisfies the same protocol an + extension library implements, where the argument is how the library + reaches the session it is being installed on. A context already is one, + so the argument is ignored. + """ + return self.ctx.__datafusion_logical_extension_codec__(session) + + def __datafusion_query_planner__(self, session: Any = None) -> Any: + """Access the ``FFI_QueryPlanner`` PyCapsule for the current planner. + + See :meth:`__datafusion_logical_extension_codec__` for ``session``. + """ + return self.ctx.__datafusion_query_planner__(session) def with_logical_extension_codec( self, codec: LogicalExtensionCodecExportable | _PyCapsule @@ -2190,15 +2260,25 @@ def with_logical_extension_codec( Only FFI codecs are supported. Pass any object implementing ``__datafusion_logical_extension_codec__`` (see :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`). + + The returned context shares its session state with the original, so a + later registration on either is visible to both. If a custom query + planner is installed, it is rebuilt against the new codec on the shared + session, so the original context plans with the new codec too. This + happens on the shared session, so it takes effect even if the returned + context is discarded. """ new_internal = self.ctx.with_logical_extension_codec(codec) new = SessionContext.__new__(SessionContext) new.ctx = new_internal return new - def __datafusion_physical_extension_codec__(self) -> Any: - """Access the PyCapsule FFI_PhysicalExtensionCodec.""" - return self.ctx.__datafusion_physical_extension_codec__() + def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: + """Access the PyCapsule FFI_PhysicalExtensionCodec. + + See :meth:`__datafusion_logical_extension_codec__` for ``session``. + """ + return self.ctx.__datafusion_physical_extension_codec__(session) def with_physical_extension_codec( self, codec: PhysicalExtensionCodecExportable | _PyCapsule @@ -2208,6 +2288,13 @@ def with_physical_extension_codec( Only FFI codecs are supported. Pass any object implementing ``__datafusion_physical_extension_codec__`` (see :py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`). + + The returned context shares its session state with the original, so a + later registration on either is visible to both. If a custom query + planner is installed, it is rebuilt against the new codec on the shared + session, so the original context plans with the new codec too. This + happens on the shared session, so it takes effect even if the returned + context is discarded. """ new_internal = self.ctx.with_physical_extension_codec(codec) new = SessionContext.__new__(SessionContext) @@ -2250,7 +2337,13 @@ def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext: regardless of the toggle. Returns a new :class:`SessionContext` with the toggle applied; - the original session is unchanged. + the original context's own codec settings are unchanged. The + returned context shares its session state with the original, so + a later registration on either is visible to both. If a custom + query planner is installed, it is rebuilt against the new codecs + on the shared session, so the original context plans with them + too. This happens on the shared session, so it takes effect even + if the returned context is discarded. Examples: >>> import pyarrow as pa diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index 394c682ae..43b53e469 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -114,15 +114,28 @@ def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]: class LogicalExtensionCodecExportable(Protocol): - """Type hint for objects exposing ``__datafusion_logical_extension_codec__``.""" + """Type hint for objects exposing ``__datafusion_logical_extension_codec__``. - def __datafusion_logical_extension_codec__(self) -> object: ... # noqa: D105 + ``session`` is the :py:class:`~datafusion.context.SessionContext` the codec + is being installed on. Take the task context provider from it rather than + building a session of your own, so the decode callbacks resolve names + against the session that runs the query. + """ + + def __datafusion_logical_extension_codec__( # noqa: D105 + self, session: Any + ) -> object: ... class PhysicalExtensionCodecExportable(Protocol): - """Type hint for objects exposing ``__datafusion_physical_extension_codec__``.""" + """Type hint for objects exposing ``__datafusion_physical_extension_codec__``. + + See :py:class:`LogicalExtensionCodecExportable` for ``session``. + """ - def __datafusion_physical_extension_codec__(self) -> object: ... # noqa: D105 + def __datafusion_physical_extension_codec__( # noqa: D105 + self, session: Any + ) -> object: ... class ScalarUDF: diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 7d038c7a5..3c95835af 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import ctypes import datetime as dt import gzip import pathlib @@ -731,6 +732,153 @@ def test_remove_optimizer_rule(ctx): assert ctx.remove_optimizer_rule("nonexistent_rule") is False +def test_set_query_planner_rejects_wrong_capsule(ctx): + with pytest.raises(ValueError, match="datafusion_query_planner"): + ctx.set_query_planner(ctx.__datafusion_task_context_provider__()) + + +def test_with_extension_rejects_wrong_capsule(ctx): + """The extension options hook names the capsule it was handed. + + Like the rest of the capsule family, this reports which capsule turned up + rather than CPython's fixed "called with incorrect name" string. + """ + + class WrongCapsule: + def __datafusion_extension_options__(self): + return ctx.__datafusion_task_context_provider__() + + with pytest.raises(ValueError, match="datafusion_extension_options"): + SessionConfig().with_extension(WrongCapsule()) + + +def test_pre_55_codec_signature_reports_an_upgrade(ctx): + """A getter that refuses the session is named, not left as a bare TypeError. + + Extension libraries implement these getters, so the pre-55.0.0 signature + is what an out-of-date one still has. The original error stays reachable + as ``__cause__`` rather than being replaced outright. + """ + + class PreSessionCodec: + def __datafusion_logical_extension_codec__(self): + msg = "should never be called" + raise AssertionError(msg) + + with pytest.raises(ImportError, match="__datafusion_logical_extension_codec__"): + ctx.with_logical_extension_codec(PreSessionCodec()) + + with pytest.raises(ImportError) as excinfo: + ctx.with_logical_extension_codec(PreSessionCodec()) + assert isinstance(excinfo.value.__cause__, TypeError) + assert "positional argument" in str(excinfo.value.__cause__) + + +def test_type_error_inside_a_getter_is_not_reported_as_an_upgrade(ctx): + """A correctly-signed getter's own TypeError must survive unchanged. + + Only the call machinery's arity error means the library is out of date. + Rewriting every TypeError would send an author debugging their own getter + off to upgrade a library that is already correct. + """ + + class RaisesTypeError: + def __datafusion_logical_extension_codec__(self, session): + msg = "bad cast inside the getter" + raise TypeError(msg) + + with pytest.raises(TypeError, match="bad cast inside the getter"): + ctx.with_logical_extension_codec(RaisesTypeError()) + + +def test_non_type_errors_from_a_getter_propagate(ctx): + """Anything that is not a TypeError was never a signature problem.""" + + class RaisesValueError: + def __datafusion_logical_extension_codec__(self, session): + msg = "something else entirely" + raise ValueError(msg) + + with pytest.raises(ValueError, match="something else entirely"): + ctx.with_logical_extension_codec(RaisesValueError()) + + +def test_set_query_planner_capsule(ctx): + capsule = ctx.__datafusion_query_planner__() + get_name = ctypes.pythonapi.PyCapsule_GetName + get_name.argtypes = [ctypes.py_object] + get_name.restype = ctypes.c_char_p + assert get_name(capsule) == b"datafusion_query_planner" + + ctx.register_record_batches( + "query_planner_test", + [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]], + ) + ctx.set_query_planner(capsule) + assert ctx.table_exist("query_planner_test") + batches = ctx.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_installing_a_planner_leaves_the_session_intact(ctx): + """The planner is written into the existing session, not a copy of it. + + Registrations made before the install are still visible afterwards, and + ones made after are visible too -- there is a single session throughout, + so neither the catalogs nor the function registry are snapshotted. + """ + ctx.register_record_batches( + "registered_before", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + before = udf( + lambda arr: arr, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="registered_before", + ) + ctx.register_udf(before) + + ctx.set_query_planner(ctx.__datafusion_query_planner__()) + + ctx.register_record_batches( + "registered_after", + [[pa.RecordBatch.from_pydict({"value": [2]})]], + ) + after = udf( + lambda arr: arr, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="registered_after", + ) + ctx.register_udf(after) + + assert ctx.table_exist("registered_before") + assert ctx.table_exist("registered_after") + assert ctx.sql("SELECT registered_before(1)").collect() + assert ctx.sql("SELECT registered_after(1)").collect() + + +def test_contexts_sharing_a_session_share_the_planner(ctx): + """A context derived before the install still plans through the planner. + + ``with_python_udf_inlining`` returns a handle on the same session, and the + query planner lives in that session's state. + """ + sibling = ctx.with_python_udf_inlining(enabled=False) + ctx.register_record_batches( + "shared_planner_test", + [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]], + ) + + ctx.set_query_planner(ctx.__datafusion_query_planner__()) + + assert sibling.table_exist("shared_planner_test") + assert sibling.session_id() == ctx.session_id() + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) @@ -1149,3 +1297,68 @@ def test_read_csv_with_options(tmp_path, as_read, global_ctx): read_csv_with_options_inner( tmp_path, csv_content, options, expected, as_read, global_ctx ) + + +def test_pre_52_table_provider_signature_reports_an_upgrade(ctx): + """The table provider hook reports an upgrade the same way codecs do. + + The 52.0.0 signature change added the session argument. This path had its + own copy of the error mapping and so missed later corrections to it. + """ + + class PreSessionProvider: + def __datafusion_table_provider__(self): + msg = "should never be called" + raise AssertionError(msg) + + with pytest.raises(ImportError, match="__datafusion_table_provider__") as excinfo: + ctx.register_table("old_sig", PreSessionProvider()) + assert isinstance(excinfo.value.__cause__, TypeError) + + +def test_catalog_provider_getter_arity_error_names_the_codec(ctx): + """A getter refusing the codec capsule is diagnosed, not left as a TypeError. + + `__datafusion_catalog_provider__` is handed the host's logical extension + codec, not the session. It used to call `getattr(...).call1(...)` directly + and so produced a bare `TypeError` for an out-of-date library; it now routes + through `call_capsule_getter` like the rest of the family. + + The message has to name the codec rather than the SessionContext, or it + would send the author to change the wrong parameter. + """ + + class PreCodecCatalogProvider: + def __datafusion_catalog_provider__(self): + msg = "should never be called" + raise AssertionError(msg) + + with pytest.raises(ImportError, match="__datafusion_catalog_provider__") as excinfo: + ctx.register_catalog_provider("old_sig", PreCodecCatalogProvider()) + assert "logical extension codec" in str(excinfo.value) + assert "SessionContext" not in str(excinfo.value) + assert isinstance(excinfo.value.__cause__, TypeError) + + +def test_type_error_inside_a_catalog_provider_getter_propagates(ctx): + """A correctly-signed catalog getter's own TypeError survives unchanged.""" + + class RaisesTypeError: + def __datafusion_catalog_provider__(self, codec): + msg = "bad cast inside the getter" + raise TypeError(msg) + + with pytest.raises(TypeError, match="bad cast inside the getter"): + ctx.register_catalog_provider("raises", RaisesTypeError()) + + +def test_type_error_inside_a_table_provider_getter_propagates(ctx): + """A correctly-signed provider getter's own TypeError survives unchanged.""" + + class RaisesTypeError: + def __datafusion_table_provider__(self, session): + msg = "bad cast inside the getter" + raise TypeError(msg) + + with pytest.raises(TypeError, match="bad cast inside the getter"): + ctx.register_table("raises", RaisesTypeError()) diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py index 588caa21a..451f5d215 100644 --- a/python/tests/test_pickle_expr.py +++ b/python/tests/test_pickle_expr.py @@ -387,6 +387,45 @@ def test_toggle_off_then_on_restores_inline_encoding(self): decoded = Expr.from_bytes(blob_toggled, ctx=SessionContext()) assert "double" in decoded.canonical_name() + def test_installing_a_codec_preserves_strict_mode(self): + """Installing an extension codec must not re-enable inlining. + + `with_{logical,physical}_extension_codec` builds a replacement + `Python{Logical,Physical}Codec` around the imported one, and the + constructor defaults inlining to on. A context that opted out has to + keep its setting, or a codec install silently starts shipping + cloudpickled callables again. + + The session's own capsule getters are a convenient stand-in for a real + extension codec here: they return a genuine FFI codec, so the import + path under test is the same one an extension library exercises. + + This covers the logical codec. The physical one is checked in + ``test_plans.py``, since it takes an ``ExecutionPlan`` to observe. + """ + strict = SessionContext().with_python_udf_inlining(enabled=False) + e = self._build_double_udf()(col("a")) + assert b"DFPYUDF" not in e.to_bytes(strict) + + with_logical = strict.with_logical_extension_codec( + strict.__datafusion_logical_extension_codec__() + ) + assert b"DFPYUDF" not in e.to_bytes(with_logical) + + def test_installing_a_codec_preserves_inlining_when_enabled(self): + """The converse: the default stays on across a codec install. + + Guards against 'fixing' the above by hard-coding inlining off. + """ + ctx = SessionContext() + e = self._build_double_udf()(col("a")) + assert b"DFPYUDF" in e.to_bytes(ctx) + + installed = ctx.with_logical_extension_codec( + ctx.__datafusion_logical_extension_codec__() + ) + assert b"DFPYUDF" in e.to_bytes(installed) + def test_strict_roundtrip_via_registry(self): """When both sender and receiver disable inlining, the UDF travels by name only and the receiver resolves it from its diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 11e709f6b..0145d123e 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -17,6 +17,7 @@ import datetime +import pyarrow as pa import pytest from datafusion import ( ExecutionPlan, @@ -24,6 +25,8 @@ Metric, MetricsSet, SessionContext, + col, + udf, ) @@ -92,6 +95,46 @@ def test_session_with_logical_extension_codec_roundtrip(ctx, df) -> None: assert df.collect() == df_round_trip.collect() +def test_installing_a_physical_codec_preserves_strict_mode() -> None: + """Installing a physical extension codec must not re-enable inlining. + + `with_physical_extension_codec` builds a replacement `PythonPhysicalCodec` + around the imported one, and the constructor defaults inlining to on. A + context that opted out via `with_python_udf_inlining(enabled=False)` has to + keep its setting, or installing a codec silently starts embedding + cloudpickled callables in serialized execution plans. + + The logical counterpart lives in `test_pickle_expr.py`; this one needs an + `ExecutionPlan` because only the physical codec encodes it. `DFPYUDF` is + the scalar Python-UDF family prefix, shared by both layers; see + `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs. + """ + identity = udf( + lambda arr: arr, + [pa.string()], + pa.string(), + volatility="immutable", + name="identity_str", + ) + + def plan_bytes(ctx: SessionContext) -> bytes: + df = ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select( + identity(col("c1")) + ) + return df.execution_plan().to_bytes(ctx) + + # The inlining default is on, so the strict blob is what has to differ. + assert b"DFPYUDF" in plan_bytes(SessionContext()) + + strict = SessionContext().with_python_udf_inlining(enabled=False) + assert b"DFPYUDF" not in plan_bytes(strict) + + installed = strict.with_physical_extension_codec( + strict.__datafusion_physical_extension_codec__() + ) + assert b"DFPYUDF" not in plan_bytes(installed) + + def test_session_codec_capsule_getters(ctx) -> None: """SessionContext exposes both logical and physical codec capsules.""" logical = ctx.ctx.__datafusion_logical_extension_codec__() diff --git a/python/tests/test_udtf.py b/python/tests/test_udtf.py index dcb2bacc3..aa0599ffa 100644 --- a/python/tests/test_udtf.py +++ b/python/tests/test_udtf.py @@ -233,3 +233,33 @@ class FakeFFITableFunction: with pytest.raises(TypeError, match="FFI-exported table functions"): TableFunction("fake_ffi", fake, with_session=True) + + +def test_pre_52_table_function_signature_reports_an_upgrade() -> None: + """A getter that refuses the session is named, not left as a bare TypeError. + + The 52.0.0 signature change added the session argument. An out-of-date + library still has the old one, and the original error stays reachable as + ``__cause__``. + """ + + class PreSessionTableFunction: + def __datafusion_table_function__(self): + msg = "should never be called" + raise AssertionError(msg) + + with pytest.raises(ImportError, match="__datafusion_table_function__") as excinfo: + TableFunction("old_sig", PreSessionTableFunction(), None) + assert isinstance(excinfo.value.__cause__, TypeError) + + +def test_type_error_inside_a_table_function_getter_propagates() -> None: + """A correctly-signed getter's own TypeError must survive unchanged.""" + + class RaisesTypeError: + def __datafusion_table_function__(self, session): + msg = "bad cast inside the getter" + raise TypeError(msg) + + with pytest.raises(TypeError, match="bad cast inside the getter"): + TableFunction("raises", RaisesTypeError(), None)