Commit e6c7587
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>1 parent 2f528d2 commit e6c7587
47 files changed
Lines changed: 3386 additions & 282 deletions
File tree
- .ai/skills
- audit-skill-md
- check-upstream
- ffi-capsule-protocol
- make-pythonic
- .github/workflows
- crates
- core/src
- util
- src
- docs/source
- contributor-guide
- user-guide
- io
- examples
- datafusion-ffi-example
- src
- datafusion-ffi-query-planner-example
- python/tests
- src
- python
- datafusion
- tests
Some content is hidden
Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
1 | 7 | | |
2 | 8 | | |
3 | 9 | | |
| |||
17 | 23 | | |
18 | 24 | | |
19 | 25 | | |
20 | | - | |
21 | | - | |
22 | | - | |
23 | | - | |
24 | | - | |
25 | | - | |
26 | 26 | | |
27 | 27 | | |
28 | 28 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
1 | 7 | | |
2 | 8 | | |
3 | 9 | | |
| |||
17 | 23 | | |
18 | 24 | | |
19 | 25 | | |
20 | | - | |
21 | | - | |
22 | | - | |
23 | | - | |
24 | | - | |
25 | | - | |
26 | 26 | | |
27 | 27 | | |
28 | 28 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
| 7 | + | |
| 8 | + | |
| 9 | + | |
| 10 | + | |
| 11 | + | |
| 12 | + | |
| 13 | + | |
| 14 | + | |
| 15 | + | |
| 16 | + | |
| 17 | + | |
| 18 | + | |
| 19 | + | |
| 20 | + | |
| 21 | + | |
| 22 | + | |
| 23 | + | |
| 24 | + | |
| 25 | + | |
| 26 | + | |
| 27 | + | |
| 28 | + | |
| 29 | + | |
| 30 | + | |
| 31 | + | |
| 32 | + | |
| 33 | + | |
| 34 | + | |
| 35 | + | |
| 36 | + | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
| 56 | + | |
| 57 | + | |
| 58 | + | |
| 59 | + | |
| 60 | + | |
| 61 | + | |
| 62 | + | |
| 63 | + | |
| 64 | + | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
| 74 | + | |
| 75 | + | |
| 76 | + | |
| 77 | + | |
| 78 | + | |
| 79 | + | |
| 80 | + | |
| 81 | + | |
| 82 | + | |
| 83 | + | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
| 87 | + | |
| 88 | + | |
| 89 | + | |
| 90 | + | |
| 91 | + | |
| 92 | + | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
| 96 | + | |
| 97 | + | |
| 98 | + | |
| 99 | + | |
| 100 | + | |
| 101 | + | |
| 102 | + | |
| 103 | + | |
| 104 | + | |
| 105 | + | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
| 109 | + | |
| 110 | + | |
| 111 | + | |
| 112 | + | |
| 113 | + | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
| 117 | + | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
| 121 | + | |
| 122 | + | |
| 123 | + | |
| 124 | + | |
| 125 | + | |
| 126 | + | |
| 127 | + | |
| 128 | + | |
| 129 | + | |
| 130 | + | |
| 131 | + | |
| 132 | + | |
| 133 | + | |
| 134 | + | |
| 135 | + | |
| 136 | + | |
| 137 | + | |
| 138 | + | |
| 139 | + | |
| 140 | + | |
| 141 | + | |
| 142 | + | |
| 143 | + | |
| 144 | + | |
| 145 | + | |
| 146 | + | |
| 147 | + | |
| 148 | + | |
| 149 | + | |
| 150 | + | |
| 151 | + | |
| 152 | + | |
| 153 | + | |
| 154 | + | |
| 155 | + | |
| 156 | + | |
| 157 | + | |
| 158 | + | |
| 159 | + | |
| 160 | + | |
| 161 | + | |
| 162 | + | |
| 163 | + | |
| 164 | + | |
| 165 | + | |
| 166 | + | |
| 167 | + | |
| 168 | + | |
| 169 | + | |
| 170 | + | |
| 171 | + | |
| 172 | + | |
| 173 | + | |
| 174 | + | |
| 175 | + | |
| 176 | + | |
| 177 | + | |
| 178 | + | |
| 179 | + | |
| 180 | + | |
| 181 | + | |
| 182 | + | |
| 183 | + | |
| 184 | + | |
| 185 | + | |
| 186 | + | |
| 187 | + | |
| 188 | + | |
| 189 | + | |
| 190 | + | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
1 | 7 | | |
2 | 8 | | |
3 | 9 | | |
| |||
17 | 23 | | |
18 | 24 | | |
19 | 25 | | |
20 | | - | |
21 | | - | |
22 | | - | |
23 | | - | |
24 | | - | |
25 | | - | |
26 | 26 | | |
27 | 27 | | |
28 | 28 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
186 | 186 | | |
187 | 187 | | |
188 | 188 | | |
189 | | - | |
| 189 | + | |
190 | 190 | | |
191 | 191 | | |
192 | 192 | | |
| |||
196 | 196 | | |
197 | 197 | | |
198 | 198 | | |
| 199 | + | |
| 200 | + | |
| 201 | + | |
| 202 | + | |
| 203 | + | |
| 204 | + | |
| 205 | + | |
| 206 | + | |
| 207 | + | |
| 208 | + | |
199 | 209 | | |
200 | 210 | | |
201 | 211 | | |
| |||
207 | 217 | | |
208 | 218 | | |
209 | 219 | | |
210 | | - | |
| 220 | + | |
| 221 | + | |
| 222 | + | |
211 | 223 | | |
212 | 224 | | |
213 | 225 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
93 | 93 | | |
94 | 94 | | |
95 | 95 | | |
| 96 | + | |
| 97 | + | |
| 98 | + | |
96 | 99 | | |
97 | 100 | | |
98 | 101 | | |
99 | 102 | | |
100 | | - | |
| 103 | + | |
| 104 | + | |
101 | 105 | | |
102 | 106 | | |
103 | 107 | | |
| |||
121 | 125 | | |
122 | 126 | | |
123 | 127 | | |
| 128 | + | |
| 129 | + | |
124 | 130 | | |
125 | 131 | | |
126 | 132 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
35 | 35 | | |
36 | 36 | | |
37 | 37 | | |
38 | | - | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
39 | 56 | | |
40 | 57 | | |
41 | 58 | | |
| |||
48 | 65 | | |
49 | 66 | | |
50 | 67 | | |
51 | | - | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
52 | 72 | | |
53 | 73 | | |
54 | 74 | | |
| |||
0 commit comments