blob: af128886fb5262f04d15cdd4ed216c2aae838696 [file] [view]
<!---
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.
-->
# 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 preferred setup uses `SessionContext.with_extensions` with extension bundles:
```python
config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
ctx = SessionContext(config).with_extensions(provider_bundle, MyPlannerExtension())
ctx.register_table("numbers", provider)
ctx.register_udf(provider_udf)
```
`MyPlannerExtension` implements both extension hooks. `__datafusion_session_components__`
receives the session it is being installed on, binds fresh codecs to that session's
task-context provider, and returns them as `SessionExtensionComponents`.
`__datafusion_session_planner__` then runs in the host's second phase, after every
bundle's codecs are installed, and builds a planner that delegates to the `fallback` it
is handed — so several libraries that each ship a planner nest instead of displacing one
another, and no planner is left carrying a chain that has since grown.
Its codecs are handed over as `BundledLogicalCodec` and `BundledPhysicalCodec` rather
than as bare capsules. `with_extensions` requires an object, because a codec's wire id
is read off the object it arrives as and a capsule has no type to read one from. Each
wrapper declares `__datafusion_codec_id__`, so the id belongs to this library and does
not change when the bundle is nested inside an application's own bundle.
The integration tests also cover the low-level chaining setup:
```python
config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
ctx = SessionContext(config)
ctx = ctx.with_logical_extension_codec(provider_logical_codec)
ctx = ctx.with_physical_extension_codec(provider_physical_codec)
ctx.register_table("numbers", provider)
ctx.register_udf(provider_udf)
ctx.set_query_planner(MyQueryPlanner())
```
`MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit.
The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](https://datafusion.apache.org/python/extension-guide/query-planners.html#install-codecs-before-a-layered-planner).
For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see the [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html).