Add a worked example of three extension libraries in one distributed query - #1721
Draft
timsaucer wants to merge 24 commits into
Draft
Add a worked example of three extension libraries in one distributed query#1721timsaucer wants to merge 24 commits into
timsaucer wants to merge 24 commits into
Conversation
This was referenced Sep 9, 2026
timsaucer
force-pushed
the
feat/plan-partitioning-and-errors
branch
from
September 10, 2026 16:39
39304d7 to
7123e7d
Compare
First of three libraries for the multi-library distributed example in #1719, and the reference implementation the extension guide's `extension_codec_durable_metadata` section currently lacks. Every other example codec in this repository parks the live object in a process-global `HashMap` and encodes an integer token into it. The guide says plainly that this is a demonstration and not a pattern, then has nothing to point at that does it properly. This codec is that: it writes the file paths and sizes, the projection, the row limit, and the schema, so decoding needs nothing at all from the encoding process. The provider scans a directory of Parquet files and reports one output partition per file. That is the reason it exists rather than `register_parquet`: it fixes the mapping from partition index to file, so an engine can hand partition `i` to a worker and know which bytes that worker will read. Paths are sorted, because directory iteration order is unspecified and a worker that disagreed with the driver about which file is partition 3 would produce wrong answers silently rather than fail. `PartitionedParquetExec` is a leaf on purpose. A node with children hands them to the framework to encode with the host's codec, which is correct but means the interesting part of a codec — what it writes down — belongs to someone else. Everything this node needs to run is in the node, so that is what goes on the wire. It reuses `DataSourceExec` to do the actual reading; the point is to own the description of the scan across a process boundary, not to reimplement Parquet. Wire format is `DFXSTOR1 | json_len: u32 | json | arrow ipc schema`. JSON for the scalar fields because someone debugging a worker can read it, Arrow IPC for the schema because it is the only encoding that round-trips every Arrow type. The magic carries a version the codec refuses to guess at. The codec claims by downcasting to its own concrete type and hands anything else to the default codec, whose error is the chain's "not mine" signal. Tests pin one fact that makes the narrow claim obviously right: an extension codec is only ever consulted for nodes with no native encoding, so the only nodes that reach it are ones some library owns — a broad claim can only steal from a peer, never pick up slack. Ten tests, the load-bearing one being a genuinely separate interpreter spawned through `sys.executable` that builds its own session, checks the codec id it expects is installed, decodes a plan written by another process, and executes all three partitions. A token registry cannot pass that test, which is the point of writing it first. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Second of three libraries for #1719, and the one carrying the mixed-workflow case: it exposes no `__datafusion_session_components__`, so callers register its three functions and install its two codecs by hand. That is not an artificial handicap. `SessionExtensionComponents` carries codec fields only, so a function library has nowhere to put its functions — the rename in 5a1bfeb noted that UDF and provider fields will join later. Until they do, this is what a function library actually looks like, and the example should show what that costs rather than pretend every dependency has caught up. A test asserts the shape rather than describing it: `with_extensions` rejects this object, naming the hook it lacks. The functions are `dfx_net_revenue` (the TPC-H revenue expression), `dfx_weighted_avg`, and `dfx_revenue_rank`. The aggregate is written out rather than delegating to a built-in because its state is the point: two running sums, which is what lets DataFusion compute a partial aggregate per partition and merge the results. An aggregate that could only be evaluated over its whole input at once would give a different answer once split, which is exactly what a worker does to it. A test pins that by checking the plan really is `mode=Partial` and the answer is still right. Both codecs are name-only — `try_encode_*` writes nothing and `try_decode_*` rebuilds from `name`. Three worker tests, each a separate interpreter, pin what that buys, and they disagree with each other in the useful way: codec installed, nothing registered -> works, decode_calls == 1 functions registered, no codec -> works, decode_calls == 0 neither -> fails, naming dfx_net_revenue So installing the codec is an *alternative* to registering the functions, not an addition to it. The middle case is the trap worth knowing: on the driver, where the functions are registered, the registry is tried first and the codec is never consulted — so a codec that was broken or missing looks fine right up until a worker needs it. `try_decode_*` checks the name before the buffer, in that order. An empty encoding leaves `fun_definition` unset and carries no codec id, so it is the one path where a payload is offered to every installed codec in turn; a codec that trusted `buf` first would answer for names it does not own. There are two codecs because there are two plan layers and a library cannot know which one its callers will serialize — an engine shipping physical plans exercises only the physical one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Third library for #1719, and the one that makes the other two do something. Rust owns the query planner, the stage node, its codec, and a config extension; Python owns the session factory, the driver, and the worker entry point. A real engine needs both, so this crate is a mixed maturin package rather than a pure extension module. The split point is the partial aggregate. DataFusion already breaks a GROUP BY into a partial pass per input partition and a final pass that merges them, so the partial passes are independent by construction and only their output has to come back. Wrapping that subtree in a `ShuffleStageExec` is the whole rewrite. The planner plans against `LocalOptimizerSession`, which borrows the foreign session but owns the stock optimizer rule list. Without it the rules run back across FFI and hand the library `ForeignExecutionPlan` wrappers, which cannot be serialized (that is G1) and cannot be rewritten either — an engine cannot split a subtree it holds only an opaque handle to. This was validated as a spike before any of it was built. One node does both halves of the shuffle. `execute(i)` reads the file for partition `i` if it exists and otherwise computes its child and writes it on the way past, so the same node is the thing a worker runs and the thing the driver reads, and nothing has to rewrite the plan in between. A query with no workers still gets the right answer, having done the work itself. The shuffle directory travels inside the node and therefore inside its encoding, so a worker and a driver cannot disagree about where results go. `session.py` is the piece the whole example exists to motivate. There is no way to snapshot a SessionContext and restore it elsewhere — SessionConfig is write-only from Python, and `df_settings` is readable but lists `datafusion.runtime.*` keys that have no namespace to set them back into — so worker parity cannot be automated. It has to be built the same way twice from data small enough to put in a message, which is what `SessionSpec` is. Both sides call `build_session`; anything a query depends on that is not in the spec is a bug waiting for a worker to find it. Two findings this turned up, both now documented in the code: `dfx_storage` needed a *logical* codec, not just a physical one. Its scan node is physical, so a physical codec looks sufficient — but installing any FFI query planner means the session hands that planner the logical plan as protobuf, and a logical plan holds its tables as `Arc<dyn TableProvider>`. With no `try_encode_table_provider` the session fails at `execution_plan()` with "Error serializing custom table", before anything is distributed. The payload is the directory, since everything else the provider holds is read back from it. A foreign node does not print its own name. The host shows `FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1`, so the driver's tree walk has to match on containment; an anchored match works in a single-library test and fails the moment a real extension is involved. Verified end to end: three Parquet files, three worker processes, each producing one partition of partial aggregate, driver merging them to the same answer the single-process path gives. Corrupting one shuffle file breaks the driver's query, which is how we know the workers did the work rather than the driver quietly recomputing it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Seventeen tests for #1719, every one running real worker processes, plus `run_tpch.py` for the same thing against the generated TPC-H data. The four queries: a Q1-shaped distributed aggregate; the same with `dfx_udfs`' Rust scalar and aggregate functions; an inline Python UDF; and the storage library's provider read on the workers. Each compares the distributed answer against the single-process answer through the same session factory, because disagreement there is the only reliable signal that a split is wrong. Tests use a small hand-checked fixture rather than the real dataset. `tpchgen-cli` writes one file per table, so SF-1 `lineitem` is a single 220 MB file — one partition, and nothing to fan out. `run_tpch.py` re-shards it first, which is a fair illustration of the actual constraint: an engine can only spread work as widely as the data is split. Three things this pass turned up. **An empty `shuffle_dir` was writing files into the working directory.** A registered config extension always *has* an entry, so an unset directory arrives as `Some("")` rather than `None`, and the planner treated it as configured. The stage node's paths were then relative to wherever the process happened to be, so `run_local` scattered `stage-1-part-*.arrow` next to the caller and later queries read another query's leftovers back out of them — which is how six tests failed with "Batch has 3 columns but BatchCoalescer expects 5". Four of those files had already been committed by the previous change; they are deleted here. **cloudpickle captures a module attribute as the module, not as its parent.** I expected `pa.compute` inside a UDF to fail on a worker, since `import pyarrow` does not bind `pyarrow.compute` and nothing loads it transitively. It does not fail: cloudpickle resolves the attribute and stores an import of `pyarrow.compute` itself, so the worker imports the submodule on load. The real trap is a *function* with a resolvable `module.qualname` — the same callable is 1106 bytes pickled from `__main__` and 34 bytes from an importable module, because the second is a pointer. A helper at test-module scope therefore reaches the worker as `ModuleNotFoundError: No module named '_test_three_libraries'`, with `traceback: None` and nothing naming a UDF, a plan, or serialization. Both halves are pinned as tests. **An FFI query planner encodes its own output on every query.** It returns proto bytes rather than a plan handle, so both libraries' codecs show one encode apiece straight after `execution_plan()`, before the driver has asked for any bytes. Worth knowing before reading an encode counter as "this is what shipping cost". `run_tpch.py` compares floats with a tolerance rather than for equality: splitting a `sum` across partitions changes the order the additions happen in, and floating point addition is not associative, so the low bits of `sum_charge` differ legitimately between the two runs. Any distributed engine has this property, and someone diffing two runs should not conclude the split is broken. Verified: 400k rows of real `lineitem` across four worker processes, using the custom provider, its custom scan node, the engine's stage node, and both Rust functions, agreeing with the single-process result to 1e-6 relative. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three maturin builds and three test invocations for #1719, plus the documentation change that keeping three example trees requires. The plan had been to retire `datafusion-ffi-query-planner-example` and fold it into the new engine. That is now off, on evidence from building the engine: roughly fifteen of its forty-seven tests cover planner *layering*, and `dfx_engine`'s planner structurally cannot delegate to a `fallback`. Delegating hands physical planning back to the host, which returns opaque `ForeignExecutionPlan` nodes the engine can neither serialize nor split — a stage-splitting planner has to plan for itself. So the new example has nothing for those tests to nest, and deleting the crate would delete real coverage of the most subtle part of #1679's contract. Three trees then, with distinct jobs, which the guide now states up front rather than leaving a reader to infer: `examples/distributed` is the worked example and the place to start; `datafusion-ffi-example` is the capsule-protocol test bed, one of every hook exercised hard; and `datafusion-ffi-query-planner-example` is the planner-composition test bed. The guide's "three roles in a query" section described only the latter two. Two stale claims fixed while in there. `examples/README.md` linked three `sql-on-*.py` files that do not exist. The planner example's README said its planner "owns no serializable types of its own and deliberately uses only built-in physical nodes", which stopped being true when `DistributedExec` was added — and the sentence mattered, because owning a node is exactly why that library ships its codec and planner as one bundle. The `actionlint` pre-commit hook needs Docker and could not run here; the workflow files are otherwise lint-clean and parse as YAML. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… time The documentation half of #1719. Each section here exists because building the example ran into the thing it describes. **`distributing-work/query-engines.md` gains the worker-parity checklist.** This was the gap I most expected to find and did: there is no way to snapshot a SessionContext and restore it elsewhere — SessionConfig is write-only from Python, and `df_settings` is readable but lists `datafusion.runtime.*` keys with no namespace to set them back into — so parity has to be built the same way twice, and nothing said what "the same" covers. Nine items, each of which the example gets wrong somewhere on purpose to show the failure. **`extension-guide/query-planners.md` gains "plan against your own optimizer rules".** A planner returns protobuf rather than a plan handle, so every query serializes its output. Physical planning applies `session.physical_optimizers()`, which over FFI are the *host's* rules, so each one hands the library back a `ForeignExecutionPlan` — and a stock `CooperativeExec` wrapped that way has no reachable `try_to_proto`. A perfectly serializable node becomes unserializable by having crossed a boundary. It is also opaque to `downcast_ref`, so a planner that means to rewrite the plan cannot see what it was given. Wrapping the session with a locally-owned rule list fixes both, and the section says when to do that instead of delegating to a fallback: a planner does one or the other. **`extension-guide/codecs.md` gains "a table provider needs a logical codec".** A provider library reasonably concludes a physical codec is enough, since its scan is a physical node. It is enough until someone installs a query planner, which receives the logical plan — holding the provider as an `Arc<dyn TableProvider>` — and then the session fails while planning with "Error serializing custom table". Found by shipping the storage library without one. `extension_codec_durable_metadata` also now points at a codec that does encode durable metadata, which it previously could not: it described what to do, said the in-repo examples deliberately do not do it, and left the reader with no implementation to read. **`distributing-work/expressions.md` sharpens the UDF-portability rule.** The existing text said imports are captured by reference, which is true but not the useful distinction. What decides it is whether cloudpickle can resolve the name to an importable `module.qualname`: a module attribute like `pyarrow.compute` is stored as an import of that submodule and works, while a *function* in one of your modules becomes a pointer and requires your code installed on the worker. The same callable is around 1 kB from `__main__` and around 30 bytes from a package, so moving a helper into one silently changes what ships. Now a table, with the failure signature: a bare `ModuleNotFoundError` raised during plan decode, naming neither UDFs nor serialization. Also: a README for the example that says what it is not, and two checklist items — ship a logical codec with a provider, and decode in a different process in at least one test, since a token-registry codec passes every in-process round trip. Every `{ref}` added here resolves; checked by extracting defined labels and references across the docs tree. One dangling reference exists in `aggregations.md` (`spark-functions`) and predates this work. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
timsaucer
force-pushed
the
feat/distributed-extensions-example
branch
from
September 10, 2026 16:42
9b02f74 to
467fef1
Compare
The five example libraries built in the Linux x86_64 job are test fixtures, not release artifacts. The `test-ffi-manylinux-x86_64` artifact they feed is consumed only by test.yml, which installs them on a runner of the same kind. Building them in the manylinux container therefore bought nothing, and cost a container start and an in-container rustup install apiece. They now build on the host with `container: off`, which also brings them under the Swatinem/rust-cache the job already sets up. Dropping the explicit `target` matters as much as dropping the container: naming a triple puts the artifacts under `target/<triple>/debug`, where the priming build below cannot reach them. maturin runs one `cargo rustc` per library, and cargo resolves features per package, so each invocation could rebuild the shared datafusion crates under its own feature union -- measured at 78 seconds for a single transition between two of these libraries, on an otherwise warm target directory. A single `cargo build` spanning all five resolves that union once. Measured afterwards, the five per-package builds reuse every dependency and only compile their own leaf crate: about 6 seconds in total. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`insert_stage` wrapped every topmost partial aggregate but gave each one
the same `STAGE_ID`. Two stages sharing an id exchange results through the
same `stage-1-part-N.arrow` paths, and a union drives its branches
concurrently, so they race. A `UNION ALL` of two `GROUP BY`s fails
outright:
Internal error: dfx_engine: publishing .../stage-1-part-1.arrow:
No such file or directory (os error 2)
Three defects, one symptom:
- Stage ids are now allocated from a counter threaded through
`insert_stage`, not a constant. Threaded rather than global because the
driver plans the same query twice -- once to ship the stages, once to
read their results back -- and only per-call numbering has both calls
agree on which stage is which. `stages_inserted` counts every stage
instead of once per planning call.
- The temporary file a writer builds a partition in is unique to that
writer rather than derived from the final name. Two writers of one
partition were interleaving their batches into a single file and then
one rename found it already consumed, which is why the failure surfaced
as a rename error. The rename that is supposed to make publishing
atomic was the thing that broke.
- The driver ships every stage. `find_stages` returns all of them in
pre-order and `run_distributed` dispatches one worker per
`(stage_id, partition)`; leaving one behind had the driver compute it
locally while the plan it shipped claimed otherwise.
`DistributedResult.partitions`/`worker_rows` become `tasks`/`task_rows`
keyed by both, since partition 0 of two stages is two pieces of work.
Python cannot read a stage id back off a plan -- the FFI wrapper's display
replaces `ShuffleStageExec: stage=2` -- so the numbering is a convention
shared with Rust, and `_internal.stage_id(index)` is the one place the
arithmetic is written down.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`expected_codec_ids` said its ids were "read from the libraries rather than written out here", then wrote out `dfx_udfs.physical.v1`. Every other literal id in the example is in a test pinning that id, where writing it is the point; this was the only copy on a live path. A stale copy fails badly rather than obviously: on a version bump `build_session` would report "session codec ids [...] do not match the expected [...]" against a session that was in fact correct, which is the class of misdirected failure the check exists to prevent. `dfx_udfs` ships no bundle, so it has no bundle class to hang a `physical_codec_id()` on, and its id comes off a throwaway codec. `__datafusion_codec_id__` is a property of the codec object -- where the protocol puts it -- and the two static methods are a convenience their bundles need only because a bundle is not itself a codec. Adding one to `CodecObservations` for symmetry would have erased the asymmetry this library exists to show. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three things stood between a reader and a passing test run. `uv pip install ... ../..` pointed at `examples/`, which has no `pyproject.toml`; the repository root is `../../..`. And `uv run maturin develop` fails with `Failed to spawn: maturin` unless maturin is in the venv, which nothing put there. Both were invisible to anyone who already had a working environment. The third was not in the README at all. None of the example projects declared `[tool.pytest.ini_options]`, so pytest's rootdir search walked up to the repository root and ran these suites under the root's pytest settings and its `conftest.py` -- which exists to inject a doctest namespace for `datafusion` and imports numpy. A reader following the README got `ModuleNotFoundError: No module named 'numpy'` at collection, from a package the example has no use for. It worked in CI only because the venv there is the root project's, synced with `--dev`. Each example now caps the search with its own config, which also drops the `Unknown config option: asyncio_mode` warnings every run was emitting. `python_files` picks up the underscore convention the test files already follow, so `uv run pytest` needs no path argument. No tracked example file contains a Python doctest, so nothing was relying on the root's `--doctest-modules`. Verified by running the instructions verbatim in a shell with neither the repository's venv on `PATH` nor `VIRTUAL_ENV` set: 18 passed, rootdir the example directory. The two older example crates have the same latent rootdir problem and the same one-block fix; left alone here. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`test_a_session_missing_a_library_is_rejected_at_build` was named for the check and described it in its docstring, but only built a bare `SessionContext` and asserted it carried no codec ids. `build_session` was never called, so the `RuntimeError` it exists to raise had no test. Neutering the guard leaves that test passing, which is the proof. The branch cannot be reached through the function's signature: which libraries get installed is written into `build_session`, not taken from the spec, so no argument can make the session come out wrong. The new test patches `expected_codec_ids` instead and says why -- the check guards this module against being edited inconsistently, a library added to one list and not the other, rather than anything a caller passes. The mismatched peer case is already covered by `test_a_worker_whose_codecs_disagree_refuses_the_plan`, which compares a worker's session against the driver's envelope. The old test keeps its (weaker, true) assertion under a name that says what it checks: a `SessionContext` carries none of the three libraries until `build_session` puts them there. Verified by mutation: with `if installed != expected` forced false the new test fails with "DID NOT RAISE RuntimeError" and the renamed one still passes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The module doc said "See the 'Known gaps' section of the extension guide". No such section exists, and `git log -S` finds none on any branch under docs/ -- it was never a valid pointer rather than a rename casualty. The content it wanted is the section this PR adds to the query-planner guide, "Plan against your own optimizer rules", which names `LocalOptimizerSession` as its worked example. Linking to it makes the pair mutually discoverable instead of one-way. Rust cannot use a Sphinx `:ref:`, so this is the rendered URL, matching how the query-planner example's README links into the guide. Verified against a fresh `sphinx-build`: the page contains `section id="plan-against-your-own-optimizer-rules"`. The two other guide anchors this PR's READMEs link to resolve as well, and the build reports no undefined labels. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The comment justified the panic empirically -- "Physical planning never calls this; verified in the spike" -- which reads as an observation that could stop holding, and invites the obvious "fix" of keeping a cloned `TableOptions` and lending out `&mut` to that. That fix would be worse than the panic. This type borrows the session it wraps and has no table options of its own, so a caller could set a Parquet option, watch it apply to nothing, and have no way to notice. Silently dropping a mutation is not an improvement on refusing one. The real reason the panic is unreachable is stronger than the comment claimed, and is type-level: the method takes `&mut self`, `create_physical_plan` takes `&dyn Session`, nothing in DataFusion asks for a `&mut dyn Session`, and `planner.rs` binds the wrapper immutably. Probed with the compiler rather than by reading -- adding `local.table_options_mut()` at the call site fails with E0596, "cannot borrow `local` as mutable". The trait declares the method with no default impl, so it has to be written either way. Behaviour unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`planner_host_optimizer_rules` and `distributed_worker_parity` were defined and never referenced. Both rendered fine, so nothing was broken -- they were just unreachable from anywhere a reader would be standing. The planner one goes on the checklist, next to the `fallback` items, which needed it: "Your planner hook wraps `fallback` and delegates to it" was unqualified, and the new section says a planner that rewrites the plan structurally cannot delegate. The two items now name each other instead of contradicting each other. Worker parity gets two, because it has two audiences. Users arrive at it from `expressions.md`, whose portability section covers two of the nine things a worker has to reproduce and did not say there were seven more. Engine authors arrive from the checklist: the guide already says "most engines handle several of them for you -- check which", which is only checkable if engines say so. Verified against a fresh `sphinx-build`: all three render as links to `#planner-host-optimizer-rules` and `#distributed-worker-parity`, both ids present on their pages, link text taken from the section headings, and no undefined labels anywhere in the build. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`try_decode_table_provider` ignored its `schema` argument and re-read the
first Parquet footer instead. That argument is not advisory: a serialized
`CustomScan` carries the table schema, and `from_proto` resolves the
scan's projection from column *names* to indices against it --
let column_indices = columns.columns.iter()
.map(|name| schema.index_of(name))
-- while `TableScanBuilder::build` then applies those indices to the
schema the decoded provider reports:
let schema = source.schema();
... schema.fields()[*i]
Unchecked. So a column added to the directory ahead of the others since
the plan was written selects the wrong one, and a column removed indexes
off the end and panics inside DataFusion. Taking the schema the plan
carries makes both unreachable, and skips a footer read the caller had
already done.
`try_new` keeps reading the footer for the registration path, where
nobody has supplied a schema yet.
Also exposes `provider_encode_calls` / `provider_decode_calls`. Both
counters were incremented in Rust and never reachable from Python, so the
logical half of this codec had no assertion anywhere -- only the physical
`encode_calls` / `decode_calls` were exposed.
Two tests: one drives the logical round trip through `LogicalPlan.to_bytes`
and asserts each half ran, and one drifts the directory between encode and
decode and asserts the decoded scan still reports the plan's schema.
Confirmed by mutation -- restoring the footer read makes the second fail
with the drifted `[label, sensor_id, reading]`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`try_decode` parsed the projection with `filter_map`, so an element that would not convert was skipped and the decode succeeded with a shorter list. The indices are positional, so that is a different query rather than a degraded one: lose one and it reads the wrong columns, lose all of them and it reads none -- with a well-formed plan either way, and nothing downstream able to tell. Demonstrated on the real wire format before fixing it. Editing a payload's `"projection":[0,1]` to the same-width `"projection":[1e1]` -- valid JSON, but a float, so `as_u64` declines it -- decoded without complaint and produced batches whose schema was `[]` instead of `["sensor_id", "reading"]`. It now fails with "projection index 10.0 is not a column number". Both the projection and the limit are parsed fallibly, a non-array projection is rejected rather than treated as absent, and `usize::try_from` replaces `as usize` so a 64-bit index cannot truncate on a 32-bit worker. Absent and null still both mean "every column". The test edits the payload in place at constant length, because the JSON sits in a length-delimited protobuf field and a resized payload would corrupt the framing rather than exercise the codec. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Every payload error in both codecs was `internal_datafusion_err!`, so a
worker handed a truncated or version-skewed payload got:
Internal error: dfx_storage: projection index 10.0 is not a column
number.
This issue was likely caused by a bug in DataFusion's code. Please
help us to resolve this by filing a bug report in our issue tracker
The advice is wrong and the class is wrong. Upstream documents `Internal`
as "due to bugs in DataFusion", says "a user should not be able to trigger
internal errors under normal circumstances by feeding in malformed
queries, bad data, etc.", and adds that I/O errors "do NOT fall under this
category". A codec reads bytes written by another process, so nearly
everything it can hit is bad data.
Reclassified to `Execution`: payload framing and parsing in both codecs,
and the shuffle file I/O in `stage.rs`. `Internal` survives in the three
places it is the correct class, each now saying why -- `with_new_children`
and `execute(partition)`, whose arguments come from an optimizer rule
rather than a payload and so cannot be reached by any input, and encoding
an object this process already holds in memory.
The class does survive FFI, which was worth checking: `df_result!` wraps
the error as `DataFusionError::Ffi` carrying the inner error's `Display`,
so the message a Python caller sees goes from "FFI error: Internal error:
... file a bug report" to "FFI error: Execution error: ...".
A test pins it, because nothing else would notice a regression: the
malformed-projection case now asserts the message names no bug report.
The storage codec's module doc records the rule, since two classes now
coexist in one file.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`__datafusion_session_planner__` ran `ffi_query_planner_from_pycapsule` over its `fallback` and then threw the result away with `let _ = fallback;`. The hook reads as though the fallback matters and then says it does not, which is the wrong thing to copy from an example. Nothing is lost by dropping it. I had assumed the conversion was at least serving as validation, and it is not: `_export_query_planner` already runs the same function over whatever the previous hook returned -- its doc says that is deliberate, "so a malformed planner surfaces at the hook that produced it rather than at the final install". So the second conversion repeated a getter call, a capsule check and an ABI check to produce a value with no reader. `DistributedQueryPlanner::fallback` goes too. It was only ever constructed `None`, which made the `Some(fallback)` arm of `create_physical_plan` unreachable, and unreachable code in an example is a liability -- it reads as a supported path and cannot rot loudly. What that arm documented is now prose on the struct: a planner delegates or rewrites, never both, and if you write the layering kind then hold the fallback and call it directly rather than going through `Session::create_physical_plan`, which dispatches through the installed planner and recurses until the stack overflows. `datafusion-ffi-query-planner-example` remains the crate that demonstrates layering for real, so the redundant illustration here is no loss. No behaviour change: the fallback was already unused. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Six unrelated tidy-ups, no behaviour change to any query.
`driver.py` reads a worker's row count off the last line of its stdout
rather than parsing the whole buffer. A worker shares stdout with
everything loaded into it, so one stray print turns `json.loads` into a
failure a long way from its cause; an unreadable or absent report now says
which stage and partition produced it.
`dfx_storage` exports `BundledLogicalCodec`. It was constructed with
`module = "dfx_storage"` but never added to the module, so that claim was
untrue and the physical half was exported while the logical half was not.
`write_partition` writes the file with the schema the stream declares
instead of the one the child's stream reports. They agree for any
well-behaved child, and pinning it to the declaration makes a
disagreement loud: `StreamWriter::write` rejects a batch whose schema
differs, so a child contradicting its own `schema()` fails there rather
than publishing a file readers were told to expect something else from.
`run_tpch.py` raises instead of asserting. That comparison is the only
thing making the script a check rather than a demo, and `python -O` drops
an `assert`, which would leave it printing a table it never verified. The
tolerance rationale moves to the new function's docstring, where it was
otherwise duplicated.
The CI artifact is `test-example-wheels-x86_64`. It stopped being
manylinux when these builds moved to the host, and it carries five
projects rather than the two the old name and step titles implied.
`dfx_engine` documents why it declares no dependency on the sibling
libraries. I tried declaring them -- `dfx_engine.session` imports both, so
`import dfx_engine` fails without them -- and it breaks the build outright,
because neither is published:
Because dfx-storage was not found in the package registry and your
project depends on dfx-storage, we can conclude that your project's
requirements are unsatisfiable.
So the note records the dead end next to the field someone will otherwise
fill in again. The requirement itself stays in the README, beside the
install command that satisfies it.
Verified: all three wheels build, install together the way test.yml does,
and their suites pass 13 / 11 / 19. Both workflow files parse as YAML;
actionlint could not run locally, as it needs Docker.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A stage node reads partition `i` if its file exists and computes it otherwise, which is what lets one node be both halves of the exchange. The cost is that the filesystem is the state, and stage ids restart at `FIRST_STAGE_ID` for every plan -- so a second query pointed at the same shuffle directory finds the first one's files and reads them, having computed nothing. Nothing downstream can catch it. When the two plans' stage schemas differ it surfaces as an unrelated-looking schema error; when they agree -- the same query over data that has since changed -- the previous answer arrives silently. `run_distributed` now checks before building a session. The glob lives next to `partition_path` in `stage.rs` and is exported like it, so the question "does this directory hold stage output?" is asked with the same naming convention the node answers it with. The check is scoped to stage output rather than to an empty directory, because the driver writes its own encoded plans and task envelopes there too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`write_partition` pins what it writes to the node's declared schema so a child that contradicts its own `schema()` fails at the writer. The read path did the opposite: it took `reader.schema()` off the file and handed it to the stream unexamined. Reaching that code means a file was found, and finding one says nothing about who wrote it -- an older build, or a query whose stage happened to be numbered the same, both leave something readable behind. Adopting its schema pushes the disagreement up to whichever operator first uses the batches, where the message no longer names the file; where the field lists differ only in type, there may be no error at all. Compare fields and refuse a mismatch, naming the file and both field lists. Fields rather than the whole schema: a name, a type or a nullability that disagrees changes how a batch is read, and schema-level metadata does not. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`values_to_arrays` infers its length from the arguments it is given, and when they are all scalars there is nothing to infer from, so it returns one row. `dfx_net_revenue` then produced a single value for a batch of a hundred, into a column the rest of the plan sizes at a hundred. Use `to_array(args.number_rows)` per argument instead, which is the pattern `crates/core/src/udf.rs` already follows. Iterating to `number_rows` rather than to whichever argument was examined first also makes the output length a stated contract. No test: an all-literal call is constant, so `SimplifyExpressions` folds it before execution and the function only ever sees a one-row batch that agrees with `number_rows`. There is no reachable path from Python, and this repository does not run `cargo test`. The invisibility is the reason to get the shape right in a crate written to be copied. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The comment claimed these steps "only need to run once", but `wheel-tag == 'abi3'` matches five of the six matrix entries, not one. Anyone reading it to decide whether adding a suite here is expensive got the wrong answer by a factor of five. Running five times is correct and worth stating as intent rather than leaving as an accident: the example wheels are abi3, so the interpreter underneath them is the only thing that varies between those entries, and these suites are what exercise the capsule protocol from Python. The cost is four seconds for the two FFI examples, measured, against a job that takes four minutes. Renamed to match, since the step now runs the distributed example too. actionlint not run: it needs a Docker daemon, which is unavailable here. The file parses as YAML and the step list is unchanged apart from the rename. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three small things the shuffle-directory work left behind. `run_distributed` wrote each stage's encoded plan and each worker's task envelope beside the partition files. Harmless, because the stage node builds exact paths rather than globbing, but it meant "does this directory hold stage output?" -- the question `require_empty_shuffle` now asks -- had to be answered by pattern rather than by looking. They move to a `tasks/` subdirectory, so the invariant is that everything directly under the shuffle directory is stage output. The guard stays scoped to the stage node's own naming regardless: a caller pointing it at a directory of their own should get an answer about the files that would actually be read. `.gitignore` gains the shuffle output pattern. Every supported path writes these under a temporary directory, so one in the tree means a relative shuffle_dir reached a stage node -- which is how three of them were committed on this branch once already. `compare` in run_tpch.py did `abs(lhs - rhs)` before testing for None, so a null in a float column raised TypeError and reported a genuine null-handling difference between the two runs as a crash in the checker. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #1719.
Stacked on #1720. This PR's base is
feat/plan-partitioning-and-errors, so the diff shown here is only the example.Rationale for this change
#1678 made extension codecs composable and #1679 added
with_extensions, so a library can ship codecs and a planner as one atomic bundle. Both are well covered by unit tests. What the repository has never had is an example of the thing that machinery is for: several independently compiled libraries cooperating on one query whose plan actually leaves the process.The second goal was to find out what breaks. That turned out to be the more valuable half — see the findings below, most of which were found by building something the wrong way rather than by reading code.
What changes are included in this PR?
Three new crates under
examples/distributed/, and the documentation the exercise showed was missing.dfx_udfs— a scalar function, an aggregate, a window function, and the two name-only codecs that make them portable. It deliberately exposes no__datafusion_session_components__, so it is installed by hand: two codec installs and three registrations in place of one call. That is not a strawman.SessionExtensionComponentscarries codec fields only, so a library contributing functions has nowhere to put them today, and mixed setups are the normal case. A test asserts the shape rather than describing it —with_extensionsrejects the object, naming the hook it lacks.Three worker tests pin what a codec buys, and they disagree usefully: codec installed and nothing registered works (
decode_calls == 1); functions registered and no codec works (decode_calls == 0); neither fails namingdfx_net_revenue. So installing the codec is an alternative to registering the functions. The middle case is the trap — on the driver, where functions are always registered, the registry answers first and a broken codec looks fine.dfx_storage— a Parquet-directory table provider reporting one output partition per file, its own leaf scan node, and the repository's first codec that encodes durable metadata. Wire format isDFXSTOR1 | json_len:u32 | json | arrow ipc schema: JSON for the scalar fields because someone debugging a worker can read it, Arrow IPC for the schema because it is the only encoding that round-trips every Arrow type. Ten tests, the load-bearing one being a separate interpreter that builds its own session, checks the codec id it expects is installed, decodes a plan written elsewhere, and executes all three partitions. A token registry cannot pass that test, which is why it was written first.dfx_engine— a toy distributed engine in the two halves a real one has. Rust owns the query planner, the stage node, its codec, and a config extension; Python owns the session factory, the driver, and the worker entry point. It splits at the partial aggregate, because DataFusion has already split there for its own reasons and the partial passes are independent by construction.One node does both halves of the shuffle:
execute(i)reads the file for partitioniif it exists and otherwise computes its child and writes it on the way past. So the same node is the thing a worker runs and the thing the driver reads, nothing has to rewrite the plan in between, and a query run with no workers still gets the right answer.session.pyis the piece the whole example exists to motivate, and is worth reading first.Seventeen integration tests plus
run_tpch.py. Verified end to end: 400k rows of real TPC-Hlineitemacross four worker processes — using the custom provider, its custom scan node, the engine's stage node, and both Rust functions — agreeing with the single-process result to 1e-6 relative.What the exercise found
Recorded in #1719 as G1–G13. The ones that changed the design:
A stock node that crosses FFI cannot be serialized on the planner return path.
FFI_QueryPlannerreturns protobuf rather than a plan handle, so every query serializes the plan. Physical planning appliessession.physical_optimizers(), which over FFI are the host's rules, soEnsureCooperativehands the library back aForeignExecutionPlanwrapping the host'sCooperativeExec— which has no reachabletry_to_proto. A perfectly serializable node becomes unserializable by having crossed a boundary, and is opaque todowncast_refbesides, so a planner that means to rewrite the plan cannot see what it was given. The engine plans against a session that owns the stock rule list locally, which fixes both; this was spiked before anything else was built.The greedy codec claim in
datafusion-ffi-exampleis a symptom of that, not sloppiness. Its physical codec claimsnode.is::<ForeignExecutionPlan>(), which the extension guide tells authors never to do. I tried narrowing it and reverted: 31 of the 51 tests in the query-planner example fail, every one on theCooperativeExecnode above. #1720 documents why the arm exists rather than removing it.A table provider needs a logical codec, not just a physical one. A provider library reasonably concludes otherwise, since its scan is a physical node. But an installed query planner receives the logical plan, which holds tables as
Arc<dyn TableProvider>, so the session fails while planning with "Error serializing custom table" — before anything is distributed. Found by shippingdfx_storagewithout one.cloudpickle captures a module attribute as the module. I expected
pa.computein a UDF to fail on a worker; it does not, because cloudpickle resolves the attribute and stores an import ofpyarrow.compute. The real trap is a function with a resolvablemodule.qualname: the same callable is ~1 kB pickled from__main__and ~30 bytes from a package, because the second is a pointer. Both halves are pinned as tests.An empty config value is not an absent one. A registered config extension always has an entry, so an unset
shuffle_dirarrives asSome(""). The planner treated that as configured and wrote shuffle files relative to the process's working directory, where later queries read another query's leftovers back out of them.The plan changed twice, on evidence
I had planned to retire
datafusion-ffi-query-planner-example. That is off: ~15 of its tests cover planner layering, and a stage-splitting planner structurally cannot delegate to afallback— delegating hands planning back to the host and returns opaque nodes it can neither serialize nor split. Deleting the crate would delete real coverage of the most subtle part of #1679's contract. There are three example trees now, and the guide says which is which:examples/distributedis the worked example,datafusion-ffi-exampleis the capsule-protocol test bed,datafusion-ffi-query-planner-exampleis the planner-composition test bed.I had also planned to narrow the greedy codec claim, covered above.
Are there any user-facing changes?
No API changes. New files ship in the repository but not in the
datafusionwheel.Documentation:
user-guide/distributing-work/query-engines.md— replaces a 🚧 placeholder with the worker-parity checklist: nine things that have to match between a driver and a worker, and why none of it can be automated (there is no way to snapshot aSessionContext, anddf_settingslists keys that cannot be set back).extension-guide/query-planners.md— a new section on planning against your own optimizer rules, with the error you get if you do not.extension-guide/codecs.md— a new section on why a provider needs a logical codec, andextension_codec_durable_metadatanow points at a codec that actually does it, which it previously could not.user-guide/distributing-work/expressions.md— sharpens the UDF-portability rule from "imports are by reference" to what actually decides it.extension-guide/checklist.md— two items: ship a logical codec with a provider; decode in a different process in at least one test.examples/README.md, which linked three files that do not exist, and the planner example's README, which claimed its planner owns no serializable types — untrue sinceDistributedExecwas added, and the claim mattered.CI builds three more wheels and runs three more test suites, gated to
abi3like the existing ones.Verified:
pytest python/(1443 passed, 12 skipped) and all five example suites (55, 51, 10, 11, 17).pre-commit run --all-filesclean exceptactionlint, which needs Docker and could not run locally; the workflow files parse as YAML.🤖 Generated with Claude Code