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..fc20b2b8e --- /dev/null +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -0,0 +1,146 @@ +--- +name: ffi-capsule-protocol +description: "TRIGGER — read before adding, changing, or reviewing any __datafusion_*__ capsule getter, any FFI_* export that asks for a TaskContextProvider or an extension codec, or any code that calls FFI_QueryPlanner::new / FFI_TableProvider::new / FFI_{Logical,Physical}ExtensionCodec::new. These methods are one protocol with a settled convention. Do not design it fresh; do not construct a SessionContext inside an extension library." +argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", or omit to review the whole family)" +--- + + + +# 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 fork must rebind the codecs it carries + +Installing a foreign query planner **forks** the session +(`PySessionContext::derived_parts`), because installing one writes to +`SessionState` and the receiver must not be modified. A foreign codec holds an +`FFI_TaskContextProvider` pointing at the session it was installed on, so the +fork rebinds each one onto itself via +`PySessionContext::rebound_{logical,physical}_codec`. Skip that and decode +callbacks answer from the pre-fork registry. + +Rebinding relies on `FFI_{Logical,Physical}ExtensionCodec::new` adopting the +provider on the already-foreign path, which needs DataFusion 55.1.0 or newer +(apache/datafusion#24722). It clones the handle before overwriting, so the +receiver keeps its own binding — assert that, not just that the fork works. +A codec this library owns round-trips unchanged, so the rebind is safe to apply +unconditionally. + +## Where the truth is + +- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. +- `docs/source/user-guide/upgrade-guides.md` — every past migration. +- `examples/datafusion-ffi-example/src/` — provider, catalog, function, codec + getters, all in current form. +- `examples/datafusion-ffi-query-planner-example/src/planner.rs` — planner + getter. +- `examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` + — `require_udf_on_decode` proves which session a decode callback resolves + against. Extend these when touching the protocol. 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..d433d4be2 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" [[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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" 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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-schema", @@ -1216,11 +1200,27 @@ 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-proto", + "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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-buffer", @@ -1251,8 +1251,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1272,8 +1271,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1284,8 +1282,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-ord", @@ -1309,8 +1306,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "async-trait", @@ -1325,8 +1321,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1342,8 +1337,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1352,8 +1346,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "datafusion-doc", "quote", @@ -1363,8 +1356,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "chrono", @@ -1383,8 +1375,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1406,8 +1397,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1421,8 +1411,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "chrono", @@ -1439,8 +1428,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1459,8 +1447,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-data", @@ -1496,8 +1483,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-catalog", @@ -1523,8 +1509,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1534,8 +1519,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -1545,8 +1529,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1606,8 +1589,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow-schema", "async-trait", @@ -1621,8 +1603,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "bigdecimal", @@ -1651,8 +1632,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "bigdecimal", @@ -1671,8 +1651,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/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 9896e7421..7c86cfd61 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,20 @@ 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. The +# FFI codec rebinding in `PySessionContext::derived_parts` is a silent no-op +# before that release (apache/datafusion#24722), so a 55.0.0 build would compile +# and then resolve decode callbacks against the pre-fork session. [patch.crates-io] +datafusion = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-proto = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-common = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-spark = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-expr = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-session = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } diff --git a/crates/core/src/catalog.rs b/crates/core/src/catalog.rs index 8ad49b098..ff170c257 100644 --- a/crates/core/src/catalog.rs +++ b/crates/core/src/catalog.rs @@ -689,7 +689,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..1c6658939 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -103,6 +103,7 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; use datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, @@ -233,10 +234,31 @@ 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. -#[derive(Debug)] pub struct PythonLogicalCodec { inner: Arc, python_udf_inlining: bool, + /// Keeps the exporting session alive for a codec handed across the FFI + /// boundary. + /// + /// `FFI_TaskContextProvider` stores its provider in a `Weak`, so an exported + /// codec stops working the moment the object that produced it goes out of + /// scope. Retaining the session in the inner codec survives both that and + /// `clone`, which clones the inner codec's `Arc` and so carries this along. + /// + /// Set this only on codecs that are leaving this library. A codec installed + /// *in* a session must not hold one, or the session would own itself: + /// `SessionContext -> SessionState -> query planner -> FFI codec -> here`. + exported_session: Option>, +} + +impl std::fmt::Debug for PythonLogicalCodec { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PythonLogicalCodec") + .field("inner", &self.inner) + .field("python_udf_inlining", &self.python_udf_inlining) + .finish_non_exhaustive() + } } impl PythonLogicalCodec { @@ -244,9 +266,19 @@ impl PythonLogicalCodec { Self { inner, python_udf_inlining: true, + exported_session: None, } } + /// Retain `ctx` so this codec keeps working after the exporting + /// `SessionContext` goes out of scope. Only for codecs being exported over + /// FFI; see the `exported_session` field for why installed codecs must not + /// use this. + pub fn with_exported_session(mut self, ctx: Arc) -> Self { + self.exported_session = Some(ctx); + self + } + pub fn inner(&self) -> &Arc { &self.inner } @@ -443,10 +475,31 @@ 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. -#[derive(Debug)] pub struct PythonPhysicalCodec { inner: Arc, python_udf_inlining: bool, + /// Keeps the exporting session alive for a codec handed across the FFI + /// boundary. + /// + /// `FFI_TaskContextProvider` stores its provider in a `Weak`, so an exported + /// codec stops working the moment the object that produced it goes out of + /// scope. Retaining the session in the inner codec survives both that and + /// `clone`, which clones the inner codec's `Arc` and so carries this along. + /// + /// Set this only on codecs that are leaving this library. A codec installed + /// *in* a session must not hold one, or the session would own itself: + /// `SessionContext -> SessionState -> query planner -> FFI codec -> here`. + exported_session: Option>, +} + +impl std::fmt::Debug for PythonPhysicalCodec { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PythonPhysicalCodec") + .field("inner", &self.inner) + .field("python_udf_inlining", &self.python_udf_inlining) + .finish_non_exhaustive() + } } impl PythonPhysicalCodec { @@ -454,9 +507,19 @@ impl PythonPhysicalCodec { Self { inner, python_udf_inlining: true, + exported_session: None, } } + /// Retain `ctx` so this codec keeps working after the exporting + /// `SessionContext` goes out of scope. Only for codecs being exported over + /// FFI; see the `exported_session` field for why installed codecs must not + /// use this. + pub fn with_exported_session(mut self, ctx: Arc) -> Self { + self.exported_session = Some(ctx); + self + } + pub fn inner(&self) -> &Arc { &self.inner } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 7bbeed2f1..865ebd586 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,15 @@ 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, + 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, wait_for_future, }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; @@ -1204,13 +1205,38 @@ 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 `derived_parts` 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(()) } + pub fn with_query_planner( + slf: &Bound<'_, Self>, + planner: Bound<'_, PyAny>, + ) -> PyDataFusionResult { + let planner = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?; + let this = slf.borrow(); + let (ctx, logical_codec, physical_codec) = this.derived_parts( + Arc::clone(&this.logical_codec), + Arc::clone(&this.physical_codec), + Some(planner), + ); + + Ok(Self { + ctx, + logical_codec, + physical_codec, + }) + } + 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,47 +1403,84 @@ 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.exported_ffi_logical_codec()) } - 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.exported_ffi_logical_codec(), + self.exported_ffi_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)); + let this = slf.borrow(); + let (ctx, logical_codec, physical_codec) = + this.derived_parts(logical_codec, Arc::clone(&this.physical_codec), None); Ok(Self { - ctx: Arc::clone(&self.ctx), + ctx, logical_codec, - physical_codec: Arc::clone(&self.physical_codec), + physical_codec, }) } + /// 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.exported_ffi_physical_codec()) } 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 inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; + let inner: Arc = (&inner_ffi).into(); let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); + let this = slf.borrow(); + let (ctx, logical_codec, physical_codec) = + this.derived_parts(Arc::clone(&this.logical_codec), physical_codec, None); Ok(Self { - ctx: Arc::clone(&self.ctx), - logical_codec: Arc::clone(&self.logical_codec), + ctx, + logical_codec, physical_codec, }) } @@ -1431,8 +1494,10 @@ impl PySessionContext { PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) .with_python_udf_inlining(enabled), ); + let (ctx, logical_codec, physical_codec) = + self.derived_parts(logical_codec, physical_codec, None); Self { - ctx: Arc::clone(&self.ctx), + ctx, logical_codec, physical_codec, } @@ -1440,6 +1505,119 @@ impl PySessionContext { } impl PySessionContext { + /// Return the pieces a derived `PySessionContext` should hold, binding any + /// foreign query planner and foreign codecs to the session that will run + /// the query. + /// + /// Pass `Some(planner)` to install one, or `None` to rebind whichever + /// planner the session already holds. + /// + /// With no foreign planner in play there is nothing to rebind, so the + /// existing context and codecs are shared unchanged. A foreign planner does + /// have to be rebound, and that forks the session, because installing it + /// writes to `SessionState` and the receiver must not be modified. The + /// codecs are built from the fork and the fork's state is then overwritten + /// in place, since rebuilding the context afterwards would leave them + /// pointing at a session that is no longer used for planning. + /// + /// A fork is not a deep copy. `SessionState` keeps its catalog list behind + /// an `Arc`, so catalogs and tables stay shared with the original session, + /// while registered functions and the configuration are snapshotted at the + /// time of the call. The session id is explicitly carried over, so the fork + /// and its `TaskContext`s all report the id the original session had. + fn derived_parts( + &self, + logical_codec: Arc, + physical_codec: Arc, + planner: Option, + ) -> ( + Arc, + Arc, + Arc, + ) { + let state = self.ctx.state(); + + // When the caller is only replacing codecs, recover the handle behind + // the installed planner so it can be rebound below. A planner this + // library owns is not a `ForeignQueryPlanner` and needs no rebinding, + // because it does not carry codecs of its own. + let planner = planner.or_else(|| { + 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 (Arc::clone(&self.ctx), logical_codec, physical_codec); + }; + + let ctx = Arc::new(SessionContext::new_with_state(state)); + let logical_codec = Arc::new(Self::rebound_logical_codec(&ctx, &logical_codec)); + let physical_codec = Arc::new(Self::rebound_physical_codec(&ctx, &physical_codec)); + + let inner: Arc = (&planner).into(); + let planner: Arc = (&FFI_QueryPlanner::new_with_ffi_codecs( + inner, + Self::ffi_logical_codec_for(&ctx, &logical_codec), + Self::ffi_physical_codec_for(&ctx, &physical_codec), + )) + .into(); + // `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 at `new_with_state`. Without this the fork would + // report one id from `session_id()` and another from every + // `TaskContext` it hands out. + let state = SessionStateBuilder::new_from_existing(ctx.state()) + .with_session_id(ctx.session_id()) + .with_query_planner(planner) + .build(); + *ctx.state_ref().write() = state; + + (ctx, logical_codec, physical_codec) + } + + /// Rebind a foreign logical codec to `ctx`. + /// + /// A codec imported from another library holds an `FFI_TaskContextProvider` + /// pointing at whichever session it was installed on. Its decode callbacks + /// resolve names against that session, so a fork has to move them onto the + /// fork or they keep answering from the pre-fork registry. + /// + /// `FFI_LogicalExtensionCodec::new` adopts the provider supplied here when + /// the codec is already foreign, returning a *clone* of the handle, so the + /// context this codec was derived from keeps its own binding. A codec this + /// library owns has no stored provider, and round-trips back to the same + /// `Arc` unchanged, so this is safe to apply unconditionally. + fn rebound_logical_codec( + ctx: &Arc, + codec: &Arc, + ) -> PythonLogicalCodec { + let ctx_provider = Arc::clone(ctx) as Arc; + let runtime = get_tokio_runtime().handle().clone(); + let rebound = + FFI_LogicalExtensionCodec::new(Arc::clone(codec.inner()), Some(runtime), &ctx_provider); + PythonLogicalCodec::new((&rebound).into()) + .with_python_udf_inlining(codec.python_udf_inlining()) + } + + /// Physical counterpart of [`Self::rebound_logical_codec`]. + fn rebound_physical_codec( + ctx: &Arc, + codec: &Arc, + ) -> PythonPhysicalCodec { + let ctx_provider = Arc::clone(ctx) as Arc; + let runtime = get_tokio_runtime().handle().clone(); + let rebound = FFI_PhysicalExtensionCodec::new( + Arc::clone(codec.inner()), + Some(runtime), + &ctx_provider, + ); + PythonPhysicalCodec::new((&rebound).into()) + .with_python_udf_inlining(codec.python_udf_inlining()) + } + async fn _table(&self, name: &str) -> datafusion::common::Result { self.ctx.table(name).await } @@ -1501,28 +1679,58 @@ 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) } - /// 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; + 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(&self.ctx) as Arc; - Arc::new(FFI_PhysicalExtensionCodec::new( - inner, - Some(runtime), - &ctx_provider, - )) + let ctx_provider = Arc::clone(ctx) as Arc; + FFI_PhysicalExtensionCodec::new(codec, Some(runtime), &ctx_provider) + } + + /// Build an FFI-wrapped logical codec for handing out in a PyCapsule. + /// + /// Same as [`Self::ffi_logical_codec`] except the inner codec retains this + /// session. The FFI task-context handle is weak, so without that a capsule + /// stops working as soon as the exporting `SessionContext` goes out of + /// scope, which the natural `ctx = ctx.with_query_planner(planner)` does. + /// + /// Only for codecs leaving this library. The plain builder is still correct + /// for codecs attached to objects that end up back inside this session, + /// which would otherwise make the session own itself. + fn exported_ffi_logical_codec(&self) -> FFI_LogicalExtensionCodec { + let codec = Arc::new( + PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner())) + .with_python_udf_inlining(self.logical_codec.python_udf_inlining()) + .with_exported_session(Arc::clone(&self.ctx)), + ); + Self::ffi_logical_codec_for(&self.ctx, &codec) + } + + /// Physical companion to [`Self::exported_ffi_logical_codec`]. + fn exported_ffi_physical_codec(&self) -> FFI_PhysicalExtensionCodec { + let codec = Arc::new( + PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) + .with_python_udf_inlining(self.physical_codec.python_udf_inlining()) + .with_exported_session(Arc::clone(&self.ctx)), + ); + Self::ffi_physical_codec_for(&self.ctx, &codec) } } 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/util/src/lib.rs b/crates/util/src/lib.rs index 9327d7f2f..741a9d04c 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}; @@ -175,28 +175,64 @@ 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::() { 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 +250,143 @@ 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 +/// Calls `obj.____(session)`, or `obj.____()` when no +/// session is supplied. +/// +/// The session is how an exporting library obtains the codecs and task context +/// of the session it is being installed on, instead of inventing one of its +/// own. `None` is for the reverse direction, where `obj` *is* a session and is +/// being asked for what it holds. +/// Every capsule getter must go 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. +pub fn call_capsule_getter<'py>( + obj: Bound<'py, PyAny>, + attr_name: &str, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult> { + if !obj.hasattr(attr_name)? { + return Ok(obj); + } + + let getter = obj.getattr(attr_name)?; + let result = match session { + Some(session) => getter.call1((session,)), + None => getter.call0(), }; + result.map_err(|err| { + let py = obj.py(); + if session.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. + if err.traceback(py).is_some() { + return err; + } + + let import_err = PyImportError::new_err(format!( + "Incompatible libraries. `{attr_name}` must accept the SessionContext it \ + is being installed on. Upgrade the library providing this object." + )); + // 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::()?; 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__", + 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 +403,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) => { @@ -326,13 +487,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..f107cfbe2 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -232,6 +232,138 @@ 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: `with_query_planner` rebinds a foreign planner to 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 survives a fork. Installing a foreign query planner forks the session, and the fork +rebinds every foreign codec it carries onto the new session, so a function registered +after the fork is still visible to them. `FFI_LogicalExtensionCodec::new` adopts the +provider supplied to it when the codec is already foreign, returning a clone of the +handle, so the context the fork was derived from keeps its own binding and continues to +resolve against its own session. + +That behaviour requires DataFusion 55.1.0 or newer. Before it, those constructors +silently discarded the provider, a foreign codec could not be rebound, and a fork left +it resolving against the pre-fork session +([apache/datafusion#24722](https://github.com/apache/datafusion/issues/24722)). + +All three properties are 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. + +### What a derived context shares + +`with_query_planner`, `with_logical_extension_codec`, `with_physical_extension_codec`, +and `with_python_udf_inlining` all return a new `SessionContext` rather than mutating +the receiver. How much the two contexts then share depends on whether a foreign query +planner is involved. + +Without one, the derived context wraps the *same* underlying session, so a registration +on either side is visible to both. + +`with_query_planner` is different, and so is any codec change made on a session that +already has a foreign planner installed. A foreign planner holds the FFI codecs it was +built with, so changing the codecs means rebuilding the planner against the context +that will actually run the query. That forks the session state, and the two halves of +the fork behave differently: + +- **Shared.** Catalogs, schemas, and tables. `SessionState` holds its catalog list + behind an `Arc`, so a table registered on either context is visible to both. The + runtime environment is shared for the same reason. +- **Copied.** Registered scalar, aggregate, and window functions, table functions, the + session configuration, and the analyzer and optimizer rule lists. These are + snapshotted when the derived context is created, so a UDF registered on the original + context afterwards is not visible to the derived one, and a `SET` applied to one does + not reach the other. + +The session id is carried over to the fork, so both contexts report the same id, and so +does every `TaskContext` either one hands to a foreign codec. A library that identifies a +session by its id — to correlate host-side and worker-side state, or to assert which +session a decode callback was bound to — keeps working across a fork. + +Register functions before deriving, or register them directly on the derived context: + +```python +ctx = SessionContext(config) +ctx = ctx.with_logical_extension_codec(provider_logical_codec) +ctx = ctx.with_physical_extension_codec(provider_physical_codec) +ctx = ctx.with_query_planner(planner) +ctx.register_udf(my_udf) # registered on the context that will run the query +``` + +A session holds exactly one query planner. Calling `with_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__()` +and delegate to it explicitly. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait @@ -257,3 +389,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..1ded21921 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -19,6 +19,108 @@ # 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. See the {ref}`ffi` guide +for the full protocol. + +### 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..f283d1610 --- /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 = ctx.with_query_planner(planner) +``` + +Derived contexts also rebind an installed planner when codecs change, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. + +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..263f034b8 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -0,0 +1,50 @@ +# 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-proto = { 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..af4d05b1c --- /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 = ctx.with_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 the derived context is created 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; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit. + +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..9a5d78c44 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -0,0 +1,500 @@ +# 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 pytest +from datafusion import SessionConfig, SessionContext, udf +from datafusion_ffi_example import ( + IsNullUDF, + 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 = ctx.with_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 = ctx.with_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_follows_the_session_across_a_planner_fork(): + """Installing a planner forks the session, and the codec moves with it. + + The codec is bound to the session before the fork and the function is + registered on the fork afterwards, so resolving it proves the codec's task + context provider was rebound to the forked session rather than left + pointing at the one it was installed on. + """ + 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 = ctx.with_query_planner(MyQueryPlanner()) + # Registered on the fork, after the codec was installed on its parent. + 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 + + +def test_rebinding_a_fork_leaves_the_receiver_bound_to_its_own_session(): + """Rebinding the fork's codec must not disturb the context it came from. + + ``FFI_LogicalExtensionCodec::new`` clones the handle before adopting the + new provider, so each context keeps its own binding. Both contexts here + have a planner, so both exercise the codec; the function is registered only + on the second, and only the second can resolve 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)) + + first = ctx.with_query_planner(MyQueryPlanner()) + second = first.with_query_planner(MyQueryPlanner()) + second.register_udf(udf(IsNullUDF())) + + batches = second.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + + with pytest.raises( + Exception, match=rf"could not resolve scalar function '{HOST_ONLY_UDF}'" + ): + first.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + + +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, since that is what forks. + """ + 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_a_fork_decodes_against_the_session_id_it_reports(): + """A fork and its decode callbacks must agree on the session id. + + ``with_query_planner`` forks the session state by rebuilding it through + ``SessionStateBuilder``, which mints a fresh id unless handed one, and + ``SessionContext`` caches its id in a field of its own. Dropping the id + there leaves the fork 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() + + fork = ctx.with_query_planner(MyQueryPlanner()) + assert fork.session_id() == session_id + + fork.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(): + """Mutating a session in place must not move the id it decodes against. + + ``add_physical_optimizer_rule`` rebuilds ``SessionState`` and writes it + back into the caller's own session rather than deriving a new one, so a + regenerated id would desync a context from itself with no fork to explain + it. + """ + ctx, logical_codec, physical_codec = codec_context() + fork = ctx.with_query_planner(MyQueryPlanner()) + session_id = fork.session_id() + + fork.add_physical_optimizer_rule(MyPhysicalOptimizerRule()) + assert fork.session_id() == session_id + + fork.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_rebinding_a_fork_does_not_move_the_parent_session_id(): + """Each context in a fork chain decodes against its own id. + + Two forks off one parent share the parent's id, so a test that only + compared against the parent would pass even if rebinding leaked one + context's provider into the other. Registering a function on just one fork + is what distinguishes them. + """ + ctx, logical_codec, _physical_codec = codec_context() + session_id = ctx.session_id() + + first = ctx.with_query_planner(MyQueryPlanner()) + second = first.with_query_planner(MyQueryPlanner()) + + assert first.session_id() == session_id + assert second.session_id() == session_id + + first.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + + second.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 = ctx.with_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 = ctx.with_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) + # Rebinding `ctx` drops the context that produced the capsule. The exported + # codecs retain it, so the capsule stays usable. Without that the FFI + # task-context handle is weak and planning fails with "TaskContextProvider + # went out of scope over FFI boundary". + ctx = ctx.with_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 = ctx.with_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 = ctx.with_query_planner(first) + + second = MyQueryPlanner(fallback=ctx) + ctx = ctx.with_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 = ctx.with_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 = ctx.with_query_planner(first).with_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_planner_is_not_installed_on_the_original_context(): + """``with_query_planner`` returns a fork; the receiver keeps its planner.""" + ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) + planner = MyQueryPlanner() + derived = ctx.with_query_planner(planner) + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner.plan_calls() == 0 + + derived.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + 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 = ctx.with_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).with_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_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 = ctx.with_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 = ctx.with_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/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..6e0560f15 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,50 @@ def add_physical_optimizer_rule( """ self.ctx.add_physical_optimizer_rule(rule) + def with_query_planner( + self, planner: QueryPlannerExportable | _PyCapsule + ) -> SessionContext: + """Create a new session context with a custom query planner. + + The planner is imported through its ``__datafusion_query_planner__`` + PyCapsule. The returned context carries over the current session state + and the logical and physical extension codec settings. Codec changes + made on a derived context are rebound to the planner before planning. + + 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__`. + + .. note:: Derived contexts share catalogs, not registrations + The returned context is a fork. Catalogs, schemas, and tables stay + shared with the original context, but registered functions and + configuration are copied at the time of the call. A UDF registered + on the original context afterwards is **not** visible here, while a + table registered on either context is visible to both. Register + functions before deriving, or register them on the derived context. + + Args: + planner: Object exposing ``__datafusion_query_planner__`` (see + :class:`QueryPlannerExportable`) or a raw + ``datafusion_query_planner`` PyCapsule. + + Returns: + A new context that uses the specified query planner. + + Examples: + >>> from my_extension import DistributedQueryPlanner # doctest: +SKIP + >>> ctx = SessionContext() + >>> planner = DistributedQueryPlanner() # doctest: +SKIP + >>> planner_ctx = ctx.with_query_planner(planner) # doctest: +SKIP + >>> query = planner_ctx.sql("SELECT * FROM remote_table") # doctest: +SKIP + >>> query.collect() # doctest: +SKIP + """ + new_internal = self.ctx.with_query_planner(planner) + new = SessionContext.__new__(SessionContext) + new.ctx = new_internal + return new + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. @@ -2178,9 +2234,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 +2259,24 @@ 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. The exception is a + session with a custom query planner installed: that planner has to be + rebound to the new codec, which forks the state. See + :meth:`~SessionContext.with_query_planner` for what a fork shares. """ 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 +2286,12 @@ 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. The exception is a + session with a custom query planner installed: that planner has to be + rebound to the new codec, which forks the state. See + :meth:`~SessionContext.with_query_planner` for what a fork shares. """ new_internal = self.ctx.with_physical_extension_codec(codec) new = SessionContext.__new__(SessionContext) @@ -2250,7 +2334,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 session is unchanged. The returned context shares + its session state with the original, so a later registration on + either is visible to both. The exception is a session with a + custom query planner installed: that planner has to be rebound + to the new codecs, which forks the state. See + :meth:`~SessionContext.with_query_planner` for what a fork + shares. 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..a8d5a1161 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,131 @@ def test_remove_optimizer_rule(ctx): assert ctx.remove_optimizer_rule("nonexistent_rule") is False +def test_with_query_planner_rejects_wrong_capsule(ctx): + with pytest.raises(ValueError, match="datafusion_query_planner"): + ctx.with_query_planner(ctx.__datafusion_task_context_provider__()) + + +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_with_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]})]], + ) + planner_context = ctx.with_query_planner(capsule) + assert planner_context.table_exist("query_planner_test") + batches = planner_context.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_derived_context_shares_catalogs(ctx): + """Catalogs live behind an Arc, so tables cross the fork in both directions.""" + derived = ctx.with_query_planner(ctx.__datafusion_query_planner__()) + + ctx.register_record_batches( + "registered_on_parent", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + derived.register_record_batches( + "registered_on_derived", + [[pa.RecordBatch.from_pydict({"value": [2]})]], + ) + + assert derived.table_exist("registered_on_parent") + assert ctx.table_exist("registered_on_derived") + + +def test_derived_context_snapshots_functions(ctx): + """Function registries are copied at fork time, unlike catalogs. + + A UDF registered on the parent before the fork is carried over; one + registered afterwards is not. Guards the caveat documented on + ``SessionContext.with_query_planner``. + """ + before = udf( + lambda arr: arr, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="registered_before_fork", + ) + ctx.register_udf(before) + + derived = ctx.with_query_planner(ctx.__datafusion_query_planner__()) + + after = udf( + lambda arr: arr, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="registered_after_fork", + ) + ctx.register_udf(after) + + assert derived.sql("SELECT registered_before_fork(1)").collect() + with pytest.raises(Exception, match="registered_after_fork"): + derived.sql("SELECT registered_after_fork(1)").collect() + + # The parent is unaffected by the fork. + assert ctx.sql("SELECT registered_after_fork(1)").collect() + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) @@ -1149,3 +1275,32 @@ 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_provider("old_sig", PreSessionProvider()) + assert isinstance(excinfo.value.__cause__, TypeError) + + +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_provider("raises", RaisesTypeError()) 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)