Add FFI query planner support (#1677)

* Add FFI query planner support

AI Disclosure: This code was written in part by an AI agent.:

* Add three-library FFI planner example

AI Disclosure: This code was written in part by an AI agent.:

* Update FFI query planner integration

AI Disclosure: This code was written in part by an AI agent.:

* add rat

* fix: install FFI test wheels from nested artifact paths

The FFI test wheel artifact now bundles two projects, so upload-artifact
preserves a `<project>/dist/` prefix instead of placing the wheels at the
artifact root. The install step globbed `wheels/*.whl`, which no longer
matched them, so the FFI wheels were silently skipped and the FFI unit
tests failed with `ModuleNotFoundError: No module named
'datafusion_ffi_example'`.

Install the recursive `find` results instead of re-globbing.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* refactor: address review of FFI query planner support

Collapse the two duplicated planner-install blocks into a single
`ctx_with_rebound_planner`. A derived context shares the existing
`SessionContext` when there is no foreign planner to rebind, and forks
only when one is installed, since the FFI codecs capture the context
they are built against.

Document what that fork shares. Catalogs, tables, and the runtime
environment stay shared; registered functions, configuration, and the
optimizer rule lists are snapshotted. The caveat lands on all four
derivation methods and on a new contributor-guide subsection, with
tests covering both halves.

Explain why `RuntimeAwareQueryPlanner` exists at all. Upstream's
`ForeignQueryPlanner` is the consumer-side adapter that lets an
`FFI_QueryPlanner` satisfy the `QueryPlanner` trait, which is what makes
a planner from another shared library installable in a `SessionState`.
Its trait method receives only a `&LogicalPlan` and a `&dyn Session`, so
it has nowhere to obtain a runtime handle and passes `None`.

Throughout datafusion-ffi each library attaches its own runtime to the
objects it exports, so a producer-side wrapper can enter that runtime
before running its own library's code. A provider owned by another
library keeps its owner's runtime even when it travels through our
catalog, because `FFI_TableProvider::new_with_ffi_codec` unwraps a
`ForeignTableProvider` back to the original handle and discards the
runtime passed alongside it. `session_runtime` is that same rule applied
to the session: `FFI_SessionRef` is our object and every callback on it
runs our code.

It matters for what those callbacks hand back. A plan produced by our
own planner returns as `FFI_ExecutionPlan::new(plan, runtime)`, and
`execute` enters that runtime before calling into the plan; the same
holds for our physical optimizer rules and for tables we own rather than
re-export. The delegation case this type exists for is exactly that
shape. A foreign planner falling back to our planner through
`__datafusion_query_planner__` receives a plan whose execution needs our
runtime, and datafusion-python owns that runtime as a process global
while the Python thread calling in carries no ambient one.

The same reasoning is why `__datafusion_query_planner__` re-exports
through the adapter rather than unwrapping to the inner handle. A
consumer reaching us through `ForeignQueryPlanner` calls with `None`, so
the adapter is what restores our handle on the way back out. Unwrapping
would save a planning-time round trip and silently drop it.

In the planner example, match the two real spellings of the row-limit
config key exactly instead of by suffix, and validate after both lookup
paths so the fallback cannot accept `max_rows = 0`. The key appears
twice because rebuilding a `ConfigOptions` across the FFI boundary
parks every foreign extension inside a single `FFI_ExtensionOptions`,
itself namespaced under `datafusion_ffi`.

Also declare `requires-python = ">=3.10"` on the provider example to
match the `abi3-py310` feature it builds against, and link both example
READMEs to the contributor guide rather than restating its caveats.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* refactor: drop the runtime adapter and fix exported capsule lifetimes

Remove `RuntimeAwareQueryPlanner`. It existed to re-attach our Tokio
handle to the session we hand to a foreign planner, on the reasoning that
`ForeignQueryPlanner` passes `session_runtime: None`. That handle turns
out to have no reachable path: the query planner FFI exchanges serialized
bytes rather than plan handles, a provider owned by another library keeps
its own runtime because `FFI_TableProvider::new_with_ffi_codec` unwraps a
`ForeignTableProvider` back to the original handle, and we execute on our
own runtime regardless. Setting the handle to `None` left every test
passing. Codec rebinding now downcasts upstream's `ForeignQueryPlanner`
directly, which also stops `__datafusion_query_planner__` adding a second
layer, since `new_with_ffi_codecs` already unwraps that type. The
`datafusion-session` dependency is no longer needed in crates/core.

Keep the exporting session alive for codecs handed out in a PyCapsule.
`FFI_TaskContextProvider` stores its provider in a `Weak`, so a capsule
stopped working as soon as the `SessionContext` that produced it went out
of scope. That made the natural spelling of the documented fallback
pattern fail:

    fallback = ctx.__datafusion_query_planner__()
    ctx = ctx.with_query_planner(MyPlanner(fallback=fallback))

Rebinding `ctx` dropped the exporter and planning then failed with
"TaskContextProvider went out of scope over FFI boundary". Both Python
codecs gained an opt-in `exported_session`, set only by the three capsule
getters. The keep-alive lives in the inner codec because the consumer
clones the FFI handle out of the capsule and `clone` clones the inner
codec's `Arc`, so a capsule-scoped keep-alive would die too early. It is
deliberately opt-in: the same codecs are also attached to providers and
catalogs that end up back inside the session, where a strong reference
would close a `SessionContext -> SessionState -> query planner -> FFI
codec` cycle. Both structs now implement `Debug` by hand, because
`SessionContext` is not `Debug`.

Add two example tests. One drives a plan containing `RepartitionExec`,
which spawns Tokio tasks as it runs, through all three libraries, so the
codecs are exercised on a multi-node plan rather than a bare scan. The
other layers a planner on top of the session's existing planner using the
capsule captured beforehand, which is the delegation pattern upstream
prescribes; `Session::create_physical_plan` cannot be used for this,
because it dispatches through the installed planner and recurses.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* add override for datafusion version to pre-release testing of upstream fix

* remove unintentionally committed files

* Empty commit to trigger CI

* fix: keep the example planner's exported task context alive

`FFI_TaskContextProvider` downgrades the provider it is given to a
`Weak`, so building one inline in `__datafusion_query_planner__` left the
capsule carrying a provider that was already dropped by the time it
returned. Every codec callback through that capsule would have failed
with "TaskContextProvider went out of scope over FFI boundary". The
example did not notice because it ships the default codecs and no custom
extension nodes, so `try_decode` is never reached.

`MyQueryPlanner` now owns the context and hands out clones of it. The
`QueryPlanner` the capsule carries holds a reference too, so the capsule
stays usable even when the Python object that exported it is dropped
first.

Document the distinction the inline construction obscured. The
`TaskContextProvider` supplied at export time backs the exporting
library's own codec callbacks, decoding that library's nodes in its own
registry. It is unrelated to the `&dyn Session` that later arrives at
`create_physical_plan`, which belongs to the host, and it could not be
derived from that session in any case, since the codecs are built before
any session exists.

Rename `PlannerConfig` to `MyPlannerConfig` to match `MyQueryPlanner`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* test: cover which session a foreign codec decodes against

The example codecs restore objects from a process-local token registry and
never read the `TaskContext` their FFI decode callbacks are handed, so
which session that context belongs to was untestable. The token path
ignores the registry entirely, which is why an empty
`SessionContext::new()` has served as the exported provider without
anyone noticing.

Both codecs now accept `require_udf_on_decode`. When set, every decode
call resolves that scalar function out of the task context it was given
and fails with the session id if it is absent, which makes the answer
observable. Each codec registers a marker function on the context it
exports, so a name owned by the codec's library and a name owned by the
host can be told apart.

Four tests use it. The two library-local cases pass: a foreign codec
resolves against the session its own library supplied. The two
host-registered cases are `xfail(strict=True)`, because a function
registered on the host with `register_udf` is not visible to a foreign
codec's decode callback at all. A fifth pins the current error so the
failure mode stays legible. Strict xfail means the pair will announce
itself if the upstream design changes.

Document the rule this establishes, and correct the surrounding section:
`with_query_planner` rebuilds a foreign planner against the session that
will run the query, so the provider a planner library supplies is
replaced on that path. Codecs installed through
`with_logical_extension_codec` and `with_physical_extension_codec` keep
the provider their own library exported, which is the case these tests
exercise.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* feat: pass the session to the capsule getters that need it

`FFI_QueryPlanner::new` and `FFI_{Logical,Physical}ExtensionCodec::new`
ask an extension library for a `TaskContextProvider`, and a planner for
two codecs on top of that. A library has none of those. Both examples
answered with `Arc::new(SessionContext::new())`, an empty session that
resolves nothing, held weakly by `FFI_TaskContextProvider` and therefore
also a lifetime hazard.

The table provider protocol already solved this: the host calls
`__datafusion_table_provider__(session)` and the library takes what it
needs off the session. Do the same for the other three getters.
`__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`,
and `__datafusion_physical_extension_codec__` now receive the
`SessionContext` they are being installed on. A codec takes the task
context provider from it; a planner takes both codecs and uses
`new_with_ffi_codecs`, which needs no provider at all. Neither example
constructs a `SessionContext` any more.

Decode callbacks consequently resolve against the session running the
query. The two `xfail(strict=True)` tests from the previous commit now
pass unmodified: a scalar function registered on the host with
`register_udf` is visible inside a decode callback executing in another
library, for both the logical and physical codec. A negative control
keeps the check honest, and a further test covers a function registered
after the codec was installed, since the provider is a live handle rather
than a snapshot.

`PySessionContext` gains an `ancestors` list. A foreign codec is built
against the session current at the time it is installed and holds it
weakly, so installing a foreign planner afterwards — which forks — would
strand the codec once the Python name is rebound. The keep-alive lives on
`PySessionContext` rather than on the codec because nothing reachable
from a `SessionContext` reaches a `PySessionContext`, so it cannot close
a cycle. What it does not paper over is the fork itself: a function
registered after the fork is not visible to a codec bound to the session
before it, which is the existing derived-context caveat seen from the
codec's side, and is covered by a test.

`SessionContext` accepts and ignores the argument on all three getters,
so a session satisfies the same protocol a library implements and
`ctx.__datafusion_query_planner__()` keeps working for the delegation
pattern. Calling a stale getter that takes no session now reports an
incompatible-library error naming the method, matching what
`table_provider_from_pycapsule` does.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: record the FFI capsule protocol as a convention

The session-passing rule was already settled for four getters and
documented in the 52.0.0 upgrade guide, but nothing pointed an agent or a
new contributor at it before they wrote a fifth. Write it down where it
will be found.

Add the 55.0.0 upgrade guide entry this branch owes. Changing
`__datafusion_logical_extension_codec__` and
`__datafusion_physical_extension_codec__` to take a session breaks every
extension library implementing them, so it needs before/after Rust in the
same shape as the 52.0.0 entry.

Correct `user-guide/io/table_provider.md`. It still showed the pre-52.0.0
signature with no session and a `PyCapsule::new_bound` call, so the one
page a reader is most likely to find contradicted the convention.

Add `.ai/skills/ffi-capsule-protocol/`. Its description is written as a
trigger rather than a task, because the existing skills are all things to
run on request and a convention read as one would be skipped. It leads
with enumerating the family, which is the step that makes the rest
unnecessary.

Point `CLAUDE.md` at it, since that file loads unconditionally and a skill
only helps once someone goes looking. Also note that `docs/temp/` is
gitignored build output that `grep -r` surfaces with stale copies, and
require an upgrade guide section alongside the `api change` label, so a
breaking change forces a visit to the file that records the conventions.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: put skill frontmatter before the license header

Every `.ai/skills/*/SKILL.md` opened with the ASF header and only then the
YAML frontmatter, which has to be the first thing in the file. The result
was that no skill's `description` was readable: the skill listing showed
`<!---` for all of them, so the field meant to say when a skill applies
said nothing. `skills/datafusion_python/SKILL.md` already had the right
order and was the model to follow.

Move the header below the frontmatter in all four. Apache RAT still
approves each file — it looks for the license anywhere, not at the top —
verified with rat 0.13.

This matters most for the new `ffi-capsule-protocol` skill, whose
description is written as a trigger condition rather than a task name.
The existing skills are all tasks to run on request, so a convention that
has to be read *before* writing code is easy to filter out while skimming
for something to invoke. Note the distinction in the skills section of
`AGENTS.md`.

Then remove what that makes redundant. `AGENTS.md` had grown a copy of the
skill's opening grep and a summary of its central rule. Two copies of one
convention, with the more discoverable copy free to drift, is exactly the
failure this branch already fixed in `user-guide/io/table_provider.md`.
`AGENTS.md` now says only when to look and where; the skill owns the
procedure. The `docs/source` versus `docs/temp` note moves the other way,
out of the skill and into `AGENTS.md`, where it applies to everything
rather than to this one protocol.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Update temporary DF version with corrections in FFI

* feat: rebind foreign codecs when a planner install forks the session

Installing a foreign query planner writes to `SessionState`, and
`with_query_planner` must not modify its receiver, so it forks. A foreign
codec holds an `FFI_TaskContextProvider` pointing at the session it was
installed on, and until now the fork could not move it: passing a new
provider to `FFI_LogicalExtensionCodec::new` was silently discarded
whenever the codec was already foreign. The fork rebound only its own
outer wrapper, so decode callbacks in the extension library kept
answering from the pre-fork registry, and the pre-fork session had to be
retained or the weakly held provider dangled.

apache/datafusion#24722 fixes the discard; those constructors now adopt
the provider on the already-foreign path. Repoint the patch at the branch
carrying it and rebind both codecs onto the fork.

Verified the branch carries everything already pinned rather than trusting
the commit graph, which reports the two as diverged: across 3811 files the
only differences are the four constructors from the fix, and
`datafusion/ffi/src/session/mod.rs` is byte-identical, so the
`create_physical_plan` codec fix arrives as its branch-55 backport.

`ancestors` and its helpers are deleted. They existed only to keep the
pre-fork session alive for a codec that could not be moved off it, and a
codec bound to the running session needs no such anchor.

Three tests, replacing two that were weaker than they looked. One
registers a function on the fork after the codec was installed on its
parent and resolves it, which is the direct evidence the rebind happened;
it failed before this change. One installs a planner twice and asserts the
first context still cannot resolve a function registered only on the
second, covering the clone-before-adopt half — a rebind that mutated the
shared handle would pass the first test and fail this one. The third keeps
the live-handle case. The test it replaces required a name registered
nowhere, so it passed for the same reason as the negative control and
never exercised a fork at all.

Note the version floor in `Cargo.toml` rather than raising it now: the
patched branch still reports 55.0.0, so the requirement can only move to
55.1.0 when the patch section is removed. Building against 55.0.0 without
the patch would compile and silently skip the rebind.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: restore base64 0.23.1 in Cargo.lock

Regenerating the lock against the patched DataFusion fork silently
downgraded base64 from 0.23.1 to 0.23.0. Nothing requires the older
version -- neither the fork nor upstream 55.0.0 constrains it -- so this
was incidental churn from the lockfile refresh, not a resolution result.

Restores the checksum main already had and re-points the three
dependents (datafusion-common, datafusion-functions, parquet). No other
dependency moves; cargo metadata --locked still resolves cleanly.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: preserve session id across a planner fork

Installing a foreign query planner forks the session state, and the fork
was minting a new session id. SessionStateBuilder::new_from_existing
drops the id and build() replaces it with a fresh UUID, while
SessionContext had already cached the original into a field of its own
back at new_with_state. Overwriting the state in place afterwards left
the two disagreeing: session_id() returned the pre-fork id, every
TaskContext handed to a foreign codec carried a different one.

Nothing in DataFusion core keys on the session id beyond debug logging,
so this broke no in-tree behavior. It matters at the FFI boundary, where
session id equality is the idiom for "which session is this codec bound
to", and for extension libraries correlating host-side and worker-side
state. Upstream hit the same case in SessionContext::enable_url_table
and preserves the id explicitly, guarded by preserve_session_context_id.

Passing the id through the builder makes the fork, its state, and its
TaskContexts agree, which is what the derived_parts doc comment and the
FFI contributor guide already claimed.

Verified by reading the id out of a decode callback via the example
codec's require_udf_on_decode error path -- the only way to observe the
state-side id from Python -- with and without the fix.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: preserve session id in add_physical_optimizer_rule

Same drift just fixed in derived_parts, but on a path that never forks.
add_physical_optimizer_rule rebuilds SessionState through
SessionStateBuilder::new_from_existing and writes it straight back into
the caller's own session, so the fresh id build() mints replaces the one
SessionContext had already cached at construction. The session the user
is holding then reports one id from session_id() and a different one
from every TaskContext it hands out, with no derivation to explain it.

Reproduced against a foreign codec, reading the id back out of a decode
callback: identical setup differing only by an add_physical_optimizer_rule
call went from MATCH to DRIFT, and back to MATCH with the id threaded
through the builder.

This is the last new_from_existing call site in the crate.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* test: guard the session id a codec decodes against

The two session id fixes had no regression guard. A Python-level
assertion cannot provide one: session_id() reads a copy SessionContext
caches at construction, which stayed correct through both bugs. The id
that actually moved was the one inside the TaskContext handed to a
foreign codec's decode callback, which nothing exposed.

Give the example codecs a TaskContextProbe that records it. This
replaces the bare AtomicUsize the require_udf_on_decode support used, so
the counter and the session id are recorded together, and the id is
recorded on every decode rather than only when a function was requested.

Three tests, all against the codec-side id rather than session_id():
a fork agrees with its codecs, add_physical_optimizer_rule does not move
the id, and a two-deep fork chain leaves both halves on the parent's id.

Confirmed non-vacuous: with both fixes reverted all three fail and the
other 17 tests pass; with only the derived_parts fix restored, exactly
the add_physical_optimizer_rule test still fails.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: stop reporting any getter TypeError as an outdated library

call_capsule_getter rewrote every TypeError from a capsule getter into
"Incompatible libraries ... Upgrade the library providing this object",
and dropped the original. Only an arity mismatch means the library is
out of date. An extension author whose own getter raised a TypeError --
a bad cast, a wrong argument to something it called -- was told the
error was a version problem and lost the error that would have located
it.

The two are distinguishable without guessing at message text: an arity
mismatch is raised by the call machinery before the getter's frame
exists, so no frame unwinds and no traceback is attached, while an error
from the body carries one. Verified to hold for both pure-Python and
pyo3-compiled getters, which is the case that matters here since
extension libraries are compiled.

Also chains the original as __cause__ on the paths that do report an
upgrade, so the arity error stays readable.

Tests cover all three outcomes. Confirmed non-vacuous: dropping the
traceback check fails only the inside-the-getter test, dropping
set_cause fails only the upgrade test.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix!: remove physical_codec_from_pycapsule

The from_pycapsule! macros call the getter with no arguments. That is
correct for __datafusion_physical_optimizer_rule__ and
__datafusion_task_context_provider__, which take no session, but
__datafusion_physical_extension_codec__ now takes the session it is
being installed on, so this helper was the one member of the family left
speaking the old protocol.

Nothing in the tree called it, but datafusion-python-util is published
by `cargo publish --workspace`, so it was still reachable. Against an
updated codec it raised a bare TypeError, bypassing the ImportError that
names the method. Against an outdated one it succeeded and produced a
codec resolving names against the wrong session -- the silent failure
the rest of this work exists to prevent.

Removing it is a breaking change to that crate, but the crate already
breaks this release: ffi_logical_codec_from_pycapsule gained its session
parameter. A compile error pointing at the replacement beats a helper
that quietly binds to nothing.

Callers move to ffi_physical_codec_from_pycapsule, which passes the
session, plus (&ffi).into() where an Arc<dyn PhysicalExtensionCodec> is
wanted -- what crates/core already does.

Documents both helper changes in the 55.0.0 upgrade guide, which until
now covered only the __datafusion_*__ method signatures and not the Rust
helpers the same authors call.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: check the FFI major version on every importer that can

Only ffi_query_planner_from_pycapsule validated the version a capsule
reported. The codec and table provider importers dereference a foreign
struct through the same `unsafe { data.as_ref() }` and were happy to
accept one built against a different DataFusion.

Extracts the planner's inline check into check_ffi_version and applies
it to the logical codec, physical codec, and table provider importers as
well. The helper is pub so extension libraries writing their own
importers can use it.

Two things the symmetry cannot reach, both now documented where someone
would look:

FFI_TaskContextProvider, FFI_TableProviderFactory, and
FFI_ExtensionOptions carry no version field, so their importers cannot
check. The from_pycapsule!/try_from_pycapsule! macros are #[macro_export]
and generic over the FFI type, so requiring a version field there would
break downstream users holding one of those three; they stay unchecked
and their doc comment now says to call check_ffi_version directly.

This is a diagnostic, not a soundness guarantee, and the helper says so:
`version` is not the first field on any of these structs, so reading it
already assumes the local layout. It turns the realistic failure -- a
library compiled against a different DataFusion -- into a clear error
instead of undefined behaviour on first use, which is what
datafusion_ffi::version is documented to be for.

Verified all four sites are wired by inverting the comparison and
confirming each one fires from the test suites.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: say where to relax the FFI version check

Exact equality is only right while datafusion_ffi::version tracks the
crate's semver major, which it does today, so the number moves on every
major release whether or not the ABI changed. If a version span later
becomes compatible, a maintainer needs to know that this one body holds
the whole policy -- callers pass a value and no decision -- and that
relaxing it at a call site would reintroduce the split the helper was
added to remove.

Also records the likelier resolution: if the ABI is stable but version
still follows the crate major, upstream's compatibility marker is wrong
for every consumer, so the fix belongs there rather than in a local
range policy.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: route every capsule getter through call_capsule_getter

The table provider and table function importers each carried their own
copy of the TypeError-to-ImportError mapping, predating
call_capsule_getter and never folded into it. Both therefore missed the
correction it since received: they rewrote a TypeError raised inside a
correctly-signed getter into "upgrade your library", and discarded the
original.

Three copies of one mapping, two of them stale, is the reason to have
one. Both now call the shared helper, so they pick up the traceback
discrimination and the __cause__ chain, and any later correction reaches
all three by construction.

Their messages named DataFusion 52.0.0. The shared message names the
method that refused the argument instead, which points at the specific
hook rather than a release, and the upgrade guide carries the version
detail. call_capsule_getter is now pub, with a doc comment saying to use
it rather than calling getattr directly.

Tests cover both outcomes on both paths. Verified against the previous
build that they are non-vacuous: before this change the raises-inside
case produced the same misleading ImportError as the old-signature case.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: explain the dropped & in the codec migration snippet

The before and after snippets pass the task context provider
differently, by reference in one and by value in the other, with nothing
saying why. Read as a diff it looks like a typo in one of them, and a
reader correcting it would be puzzled when both versions compile.

Both are valid: the parameter is impl Into<FFI_TaskContextProvider>,
which is satisfied by &Arc<dyn TaskContextProvider> and by
FFI_TaskContextProvider itself, and the latter is what
ffi_task_context_provider_from_pycapsule returns. The argument changes
because the provider now comes from the session instead of a field,
which is the point of the migration.

The contributor guide shows only the post-migration form, so it needs no
equivalent note.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: spell out the token registry lifecycle in the example codecs

The README already describes these as one-shot registries that consume
each token during decoding, but the source comments did not, and the
source is what someone reuses the pattern from. The existing comment
warned that the registry is process-local without saying that a decode
removes its entry, which is the constraint most likely to bite.

Documents both consequences on the registry accessors, where the
mechanism lives, with a pointer from each struct doc:

- Decode consumes the token, so the same encoded bytes cannot be decoded
  twice. Fine here because every plan is encoded immediately before the
  one decode that consumes it, but it rules out replaying a stored plan,
  retrying a decode, or fanning one plan out to several readers.
- An encode that never reaches a decoder leaks for the life of the
  process. Normal operation does not: encode and decode counts balance
  exactly across repeated queries, which is what makes remove-on-decode
  the right trade here rather than a leak on every call.

Comments only.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: resolve a planner fallback when it is installed, not constructed

MyQueryPlanner::new imported its fallback immediately, with no session
to pass, so the fallback's getter was called with no arguments. That
works for a SessionContext, whose getter takes the session optionally,
and for a raw capsule, which has no getter at all. It fails for another
foreign planner, which implements the same protocol this type does and
requires the argument -- and layering on another planner is the case a
distributed engine actually needs. The docstring claimed fallback "takes
anything exporting __datafusion_query_planner__", which was not true.

Holds the Python object instead and imports it in
__datafusion_query_planner__, where the session is in hand and can be
forwarded. All three fallback kinds now work.

Deferring also removes a footgun rather than adding one. Passing a
SessionContext now delegates to whichever planner it holds at install
time, and since with_query_planner calls the getter before installing,
the context still reports its previous planner, so wrapping a context in
a planner installed on that same context does not recurse.

Arc<Py<PyAny>> rather than Py<PyAny> because pyo3 0.29 gates Py: Clone
behind the py-clone feature, and this type derives Clone. Matches how
PythonTableFunctionCallable holds its callable.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: accumulate planner observations instead of overwriting them

foreign_session, foreign_provider, and foreign_plan were written with
store, so each one described only the most recent plan. Their accessors
are named foreign_*_observed, which asks whether the thing was ever
seen, and the tests assert them after running more than one query. The
existing tests passed by luck.

Reproduced: after scanning a foreign provider and then running
SELECT 1, foreign_provider_observed goes from True back to False.

Writes them with fetch_or so a later plan cannot retract what an earlier
one observed. plan_calls already accumulated, used_fallback only ever
stores true so it was already cumulative, and last_max_rows is
deliberately last-wins as its name says. Documents that split on the
struct, since it is the kind of thing that gets "tidied" back.

Confirmed non-vacuous: with store restored, exactly the new test fails
and the other 22 pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Update rev for upstream datafusion to pre-release of 55.1.0

* refactor!: install the query planner in place instead of forking

`with_query_planner` derived a new `SessionContext` to install a planner,
on the grounds that the receiver must not be modified. That mints a fresh
`Arc<SessionContext>` allocation, and every FFI handle in play is bound to
an allocation rather than to the logical session: `FFI_TaskContextProvider`
holds its provider weakly, and a registered catalog provider upgrades that
handle on every `supports_filters_pushdown` and every `scan`.

So the natural `ctx = ctx.with_query_planner(planner)` dropped the session
a foreign catalog had been registered on, and the next query failed with
`TaskContextProvider went out of scope over FFI boundary`. Reproduced with
a `MyCatalogProvider` registered before the install and a `WHERE` clause to
force pushdown during logical optimization.

Rebinding cannot cover this. It reaches the codecs `PySessionContext` holds
in its own fields; a codec embedded in a registered `FFI_CatalogProvider` —
and in every `FFI_SchemaProvider` and `FFI_TableProvider` minted from it —
has no Python-side handle. Nor can a codec retain the session that built
it: codecs are routinely handed to a provider that is registered straight
back into that session, closing `SessionContext -> catalog -> FFI provider
-> FFI codec -> SessionContext`.

Install in place instead, writing `SessionState` back through `state_ref()`
exactly as `add_physical_optimizer_rule` already did. A session keeps one
`Arc<SessionContext>` for life, so no handle is ever orphaned and the bug
cannot occur. This deletes the fork and everything that existed to repair
it: `ancestors`, `rebound_{logical,physical}_codec`, `exported_session` on
both codecs, and the `exported_ffi_*` builders.

The query planner lives in `SessionState`, so it belongs to the session
rather than to a handle on it. `with_query_planner(planner) ->
SessionContext` therefore becomes `set_query_planner(planner) -> None`,
matching `add_physical_optimizer_rule`.

The 55.1.0 pin may no longer be needed — its stated reason in Cargo.toml is
the rebinding this removes — but that is left alone pending a check of the
rest of the PR.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* test: install codecs in the provider-survival test

The test registered a foreign catalog provider and no codecs, so it could
never finish: once the planner install stopped orphaning the provider, the
query got past filter pushdown and then failed at plan serialization with
`LogicalExtensionCodec is not provided`. That is the same unrelated failure
`test_query_planner_requires_provider_codec` already covers, and it would
mask a dangling handle rather than expose one.

Install both provider codecs, and fold the codec-install-after-planner case
in as a parameter rather than a near-duplicate test. Both orderings write
`SessionState` — one installs the planner, the other rebuilds it against a
new codec — so both exercise the path that must not replace the session's
`Arc<SessionContext>`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: keep the UDF inlining setting when installing a codec

`with_logical_extension_codec` and `with_physical_extension_codec` built
the replacement wrapper with `Python{Logical,Physical}Codec::new`, whose
constructor defaults `python_udf_inlining` to true. Installing a codec on
a context that had opted out therefore turned inlining back on without
saying so:

    ctx = SessionContext().with_python_udf_inlining(enabled=False)
    ctx = ctx.with_logical_extension_codec(codec)   # inlining silently back on

That matters beyond a stale flag. Inlining is what embeds a cloudpickled
callable in the wire format, and it is opt-out precisely because that is
not portable across interpreters and not something every deployment wants
to ship. A codec install is not a request to change it.

Carry the receiver's setting across instead. Both new tests fail on the
prior build with `DFPYUDF` reappearing in the blob, and the paired
`..._preserves_inlining_when_enabled` case pins the default-on direction so
the fix cannot degenerate into hard-coding it off.

The physical case is covered in `test_plans.py` rather than alongside the
logical one: `Expr.to_bytes` only routes through the logical codec, so an
assertion there would pass with the physical bug still present. It takes an
`ExecutionPlan` to observe.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: route the last seven capsule getters through call_capsule_getter

`call_capsule_getter` claimed every capsule getter went through it, so the
mapping from a refused argument to a diagnosable error would live in one
place. Seven sites still called `getattr(...).call0()` or `.call1(...)`
directly, so the claim was false and four of them — catalog provider,
schema provider, catalog provider list, table provider factory — still
handed an out-of-date extension library a bare `TypeError`.

Those four take the host's logical extension codec rather than the session,
which is why they could not simply be passed through as they stood: the
diagnostic would have told a catalog author their method "must accept the
SessionContext", pointing them at the wrong parameter. Carry the argument
and its description together in a `CapsuleGetterArg` so one diagnostic can
serve getters that take a session, getters that take a codec, and getters
that take nothing. `Option<&Bound<PyAny>>` still converts into it, so the
documented `*_from_pycapsule` helper signatures are unchanged.

The three zero-argument getters (scalar, aggregate, window UDF) route
through as well. Nothing can be refused there, but the rule is easier to
follow with no exceptions to remember.

Also add `validate_pycapsule` to `table_provider_from_pycapsule` and
`ffi_logical_codec_from_pycapsule`, the two extraction sites that lacked
it. This is not redundant with `pointer_checked`, despite appearances:
`pointer_checked` bottoms out in CPython's `PyCapsule_GetPointer`, whose
error is the fixed string `PyCapsule_GetPointer called with incorrect
name` and names neither the expected capsule nor the one received. Say so
in a doc comment so it does not get "simplified" away later.

Drop the unused `datafusion-proto` dependency from the query planner
example while here.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: correct the pin rationale and derived-context caveats

The Cargo.toml comment justified the pre-release pin by "the FFI codec
rebinding in `PySessionContext::derived_parts`", a symbol that no longer
exists. That rebinding was not removed, it moved into
`set_session_query_planner`, which every `with_*` method calls. It depends on
`FFI_QueryPlanner::new_with_ffi_codecs` unwrapping a `ForeignQueryPlanner` and
replacing its codecs, a swap that is a silent no-op before 55.1.0
(apache/datafusion#24722). So the pin is still required, not droppable.

The `with_*` methods rebuild the installed planner on the *shared* session, so
the rebind takes effect even when the returned context is discarded.
`with_python_udf_inlining` additionally claimed "the original session is
unchanged", which the rebuild contradicts; it is the context's own codec
settings that are unchanged.

Also:
- Note that the arity-vs-body TypeError split in `call_capsule_getter` holds
  only because the call originates in Rust. A Python-level shim between the
  host and the getter would supply a traceback and silently disable it.
- Document the new FFI major-version gate in the 55.0.0 upgrade guide. Table
  providers previously performed no such check, so a mismatched extension
  library that used to load now raises ImportError.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* test: pin the planner rebind and guard the planner example on log warnings

`test_a_discarded_derived_context_still_rebinds_the_planner` covers the
surprising half of the shared-session rebuild: a codec installed through a
context that is then thrown away still binds to the session's planner. A fresh
codec instance makes it observable, since the planner encodes the outbound
logical plan with whichever codec it holds. Verified non-vacuous -- a codec
built but never installed reports zero encode calls.

The query planner example had no conftest, so it ran without the autouse
fail-on-log-warning handler the provider example uses, despite calling
`pyo3_log::init()` for the same reason. Copied verbatim; the suite passes
under it with no allowlist needed.

Two table provider tests called the deprecated `register_table_provider`,
which is a one-line forwarder to `register_table`, so they reached the same
capsule path while emitting DeprecationWarning.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* chore: drop the stub uv.lock from the query planner example

It locked nothing -- a single editable entry for the crate itself and no
dependencies. Nothing consumed it either: CI runs both example suites with
`uv run --no-project`, and the older datafusion-ffi-example has no lock file
at all. The codespell skip list in pyproject.toml matches on a bare `uv.lock`
glob, so it still covers the repository root lock.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* docs: scope the planner codec rebind to one layer

Installing a codec rebuilds the installed planner against it, but the
rebuild reaches exactly one `ForeignQueryPlanner`. A planner that resolved
a fallback at install time keeps that fallback's codecs, and neither side
can repair it: the host has no handle past the first layer, and the planner
library cannot re-derive codecs at plan time because `FFI_QueryPlanner`
holds them by value and `Session` exposes no accessor for the host's
current ones. Tracked upstream in apache/datafusion#24762.

The examples cannot demonstrate it. Their fallback lives in the same cdylib
as its wrapper, and `From<&FFI_QueryPlanner>` short-circuits on a matching
`library_marker_id`, so a same-library hop never serializes. Measured: a
layered planner produces the same codec traffic as a flat one.

What is demonstrable is that the session's planner tracks whichever handle
wrote it last, so re-installing a planner from the original handle rebinds
the session back to that handle's codecs rather than picking up a codec
installed through a derived one. Pinned by a new test as the sequel to
`test_a_discarded_derived_context_still_rebinds_the_planner`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* fix: skip the planner rebind when inlining is unchanged

`with_python_udf_inlining` rebuilds the handle's codecs and rebinds the
session's planner to them, which is state shared with every other handle on
that session. Asking for the setting a context already has changes nothing,
so it should not pay that side effect: a defensive no-op toggle on `ctx`
otherwise drags the planner back onto `ctx`'s codecs and silently undoes a
codec installed through another handle.

Returning the existing codecs is observationally equivalent to the rebuild
otherwise -- it wraps the same inner codec in a fresh `Python*Codec` -- so
the guard is only visible through that side effect. The new test fails
without it with `assert 0 > 0`.

Also pins the divergence the rebind creates. The planner carries the codecs
of whichever handle installed it last; every other path on a context uses
that context's own codec field. Those can be different handles, and then
`Expr.to_bytes(ctx)` and `ctx.sql(...)` encode with different codecs on the
same `ctx`. Stated as a rule in the FFI guide rather than left implicit in
the description of the mechanism.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* refactor: route the last three capsule getters through the helper

`call_capsule_getter` documents itself as having no exceptions to remember,
and had three. The `from_pycapsule!` and `try_from_pycapsule!` macros each
hand-rolled the getattr/call0 prologue, and `SessionConfig::with_extension`
called `__datafusion_extension_options__` directly. All three take no
argument, so with `CapsuleGetterArg::None` the helper returns the original
error untouched and the substitution is behavior-identical.

This is what makes the claim true rather than aspirational: the grep in the
FFI capsule protocol skill now turns up no bare call sites.

`with_extension` also gains the `validate_pycapsule` check every other
extractor has, so a mismatched capsule is named instead of raising CPython's
fixed "called with incorrect name". It keeps its own `hasattr` precheck,
whose AttributeError is more useful than the helper's pass-through, and
takes no version check because `FFI_ExtensionOptions` carries no version
field.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* chore: drop the unused datafusion-proto dep from the util crate

Its only use was `physical_codec_from_pycapsule`, whose `dyn
PhysicalExtensionCodec` output type came from that crate. That helper is
gone, replaced by `ffi_physical_codec_from_pycapsule`, which returns the FFI
type and leaves the conversion to the caller. Nothing in `crates/util/src`
references `datafusion_proto` any more, and it is not re-exported, so
extension libraries depending on this crate are unaffected.

The package stays in Cargo.lock; crates/core still uses it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
47 files changed
tree: f42c93fd03251555615efa7b10b4a6a18f119ec5
  1. .ai/
  2. .claude/
  3. .github/
  4. benchmarks/
  5. ci/
  6. crates/
  7. dev/
  8. docs/
  9. examples/
  10. python/
  11. skills/
  12. .asf.yaml
  13. .dockerignore
  14. .gitignore
  15. .gitmodules
  16. .pre-commit-config.yaml
  17. AGENTS.md
  18. Cargo.lock
  19. Cargo.toml
  20. CHANGELOG.md
  21. conftest.py
  22. LICENSE.txt
  23. pyproject.toml
  24. README.md
  25. rust-toolchain.toml
  26. rustfmt.toml
  27. uv.lock
README.md

DataFusion in Python

Python test Python Release Build

This is a Python library that binds to Apache Arrow in-memory query engine DataFusion.

DataFusion's Python bindings can be used as a foundation for building new data systems in Python. Here are some examples:

  • Dask SQL uses DataFusion's Python bindings for SQL parsing, query planning, and logical plan optimizations, and then transpiles the logical plan to Dask operations for execution.
  • DataFusion Ballista is a distributed SQL query engine that extends DataFusion's Python bindings for distributed use cases.
  • DataFusion Ray is another distributed query engine that uses DataFusion's Python bindings.

Features

  • Execute queries using SQL or DataFrames against CSV, Parquet, and JSON data sources.
  • Queries are optimized using DataFusion's query optimizer.
  • Execute user-defined Python code from SQL.
  • Exchange data with Pandas and other DataFrame libraries that support PyArrow.
  • Serialize and deserialize query plans in Substrait format.
  • Experimental support for transpiling SQL queries to DataFrame calls with Polars, Pandas, and cuDF.

For tips on tuning parallelism, see Maximizing CPU Usage in the configuration guide.

Example Usage

The following example demonstrates running a SQL query against a Parquet file using DataFusion, storing the results in a Pandas DataFrame, and then plotting a chart.

The Parquet file used in this example can be downloaded from the following page:

from datafusion import SessionContext

# Create a DataFusion context
ctx = SessionContext()

# Register table with context
ctx.register_parquet('taxi', 'yellow_tripdata_2021-01.parquet')

# Execute SQL
df = ctx.sql("select passenger_count, count(*) "
             "from taxi "
             "where passenger_count is not null "
             "group by passenger_count "
             "order by passenger_count")

# convert to Pandas
pandas_df = df.to_pandas()

# create a chart
fig = pandas_df.plot(kind="bar", title="Trip Count by Number of Passengers").get_figure()
fig.savefig('chart.png')

This produces the following chart:

Chart

Registering a DataFrame as a View

You can use SessionContext's register_view method to convert a DataFrame into a view and register it with the context.

from datafusion import SessionContext, col, literal

# Create a DataFusion context
ctx = SessionContext()

# Create sample data
data = {"a": [1, 2, 3, 4, 5], "b": [10, 20, 30, 40, 50]}

# Create a DataFrame from the dictionary
df = ctx.from_pydict(data, "my_table")

# Filter the DataFrame (for example, keep rows where a > 2)
df_filtered = df.filter(col("a") > literal(2))

# Register the dataframe as a view with the context
ctx.register_view("view1", df_filtered)

# Now run a SQL query against the registered view
df_view = ctx.sql("SELECT * FROM view1")

# Collect the results
results = df_view.collect()

# Convert results to a list of dictionaries for display
result_dicts = [batch.to_pydict() for batch in results]

print(result_dicts)

This will output:

[{'a': [3, 4, 5], 'b': [30, 40, 50]}]

Configuration

It is possible to configure runtime (memory and disk settings) and configuration settings when creating a context.

runtime = (
    RuntimeEnvBuilder()
    .with_disk_manager_os()
    .with_fair_spill_pool(10000000)
)
config = (
    SessionConfig()
    .with_create_default_catalog_and_schema(True)
    .with_default_catalog_and_schema("foo", "bar")
    .with_target_partitions(8)
    .with_information_schema(True)
    .with_repartition_joins(False)
    .with_repartition_aggregations(False)
    .with_repartition_windows(False)
    .with_parquet_pruning(False)
    .set("datafusion.execution.parquet.pushdown_filters", "true")
)
ctx = SessionContext(config, runtime)

Refer to the API documentation for more information.

Printing the context will show the current configuration settings.

print(ctx)

Extensions

For information about how to extend DataFusion Python, please see the extensions page of the online documentation.

More Examples

See examples for more information.

Executing Queries with DataFusion

Running User-Defined Python Code

Substrait Support

How to install

uv

uv add datafusion

Pip

pip install datafusion
# or
python -m pip install datafusion

Conda

conda install -c conda-forge datafusion

You can verify the installation by running:

>>> import datafusion
>>> datafusion.__version__
'0.6.0'

Using DataFusion with AI coding assistants

This project ships a SKILL.md that teaches AI coding assistants how to write idiomatic DataFusion Python. It follows the Agent Skills open standard.

Preferred: npx skills add apache/datafusion-python — installs the skill in Claude Code, Cursor, Windsurf, Cline, Codex, Copilot, Gemini CLI, and other supported agents.

Manual: paste this line into your project's AGENTS.md / CLAUDE.md:

For DataFusion Python code, see https://github.com/apache/datafusion-python/blob/main/skills/datafusion_python/SKILL.md

How to develop

This assumes that you have rust and cargo installed. We use the workflow recommended by pyo3 and maturin. The Maturin tools used in this workflow can be installed either via uv or pip. Both approaches should offer the same experience. It is recommended to use uv since it has significant performance improvements over pip.

Currently for protobuf support either protobuf or cmake must be installed.

Bootstrap (uv):

By default uv will attempt to build the datafusion python package. For our development we prefer to build manually. This means that when creating your virtual environment using uv sync you need to pass in the additional --no-install-package datafusion and for uv run commands the additional parameter --no-project

# fetch this repo
git clone git@github.com:apache/datafusion-python.git
# cd to the repo root
cd datafusion-python/
# create the virtual environment
uv sync --dev --no-install-package datafusion
# activate the environment
source .venv/bin/activate

Bootstrap (pip):

# fetch this repo
git clone git@github.com:apache/datafusion-python.git
# cd to the repo root
cd datafusion-python/
# prepare development environment (used to build wheel / install in development)
python3 -m venv .venv
# activate the venv
source .venv/bin/activate
# update pip itself if necessary
python -m pip install -U pip
# install dependencies
python -m pip install -r pyproject.toml

The tests rely on test data in git submodules.

git submodule update --init

Whenever rust code changes (your changes or via git pull):

# make sure you activate the venv using "source venv/bin/activate" first
maturin develop --uv
python -m pytest

Alternatively if you are using uv you can do the following without needing to activate the virtual environment:

uv run --no-project maturin develop --uv
uv run --no-project pytest

To run the FFI tests within the examples folder, after you have built datafusion-python with the previous commands:

cd examples/datafusion-ffi-example
uv run --no-project maturin develop --uv
uv run --no-project pytest python/tests/_test_*py

Running & Installing pre-commit hooks

datafusion-python takes advantage of pre-commit to assist developers with code linting to help reduce the number of commits that ultimately fail in CI due to linter errors. Using the pre-commit hooks is optional for the developer but certainly helpful for keeping PRs clean and concise.

Our pre-commit hooks can be installed by running pre-commit install, which will install the configurations in your DATAFUSION_PYTHON_ROOT/.github directory and run each time you perform a commit, failing to complete the commit if an offending lint is found allowing you to make changes locally before pushing.

The pre-commit hooks can also be run adhoc without installing them by simply running pre-commit run --all-files.

NOTE: the current pre-commit hooks require docker, and cmake. See note on protobuf above.

Running linters without using pre-commit

There are scripts in ci/scripts for running Rust and Python linters.

./ci/scripts/python_lint.sh
./ci/scripts/rust_clippy.sh
./ci/scripts/rust_fmt.sh
./ci/scripts/rust_toml_fmt.sh

Checking Upstream DataFusion Coverage

This project includes an AI agent skill for auditing which features from the upstream Apache DataFusion Rust library are not yet exposed in these Python bindings. This is useful when adding missing functions, auditing API coverage, or ensuring parity with upstream.

The skill accepts an optional area argument:

scalar functions
aggregate functions
window functions
dataframe
session context
ffi types
all

If no argument is provided, it defaults to checking all areas. The skill will fetch the upstream DataFusion documentation, compare it against the functions and methods exposed in this project, and produce a coverage report listing what is currently exposed and what is missing.

The skill definition lives in .ai/skills/check-upstream/SKILL.md and follows the Agent Skills open standard. It can be used by any AI coding agent that supports skill discovery, or followed manually.

How to update dependencies

To change test dependencies, change the pyproject.toml and run

uv sync --dev --no-install-package datafusion