timsaucer opened a new pull request, #1721:
URL: https://github.com/apache/datafusion-python/pull/1721
# Which issue does this PR close?
Closes #1719.
**Stacked on #1720**, which is stacked on #1679. 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.
`SessionExtensionComponents` carries 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_extensions` rejects 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 naming
`dfx_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 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. 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
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, nothing has to rewrite the plan in between, and a query run with
no workers still gets the right answer.
`session.py` is 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-H `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.
## 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_QueryPlanner` returns protobuf rather than a plan handle, so every
query serializes the plan. Physical planning applies
`session.physical_optimizers()`, which over FFI are the *host's* rules, so
`EnsureCooperative` hands the library back a `ForeignExecutionPlan` wrapping
the host's `CooperativeExec` — which has no reachable `try_to_proto`. A
perfectly serializable node becomes unserializable by having crossed a
boundary, and is opaque to `downcast_ref` besides, 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-example` is a symptom of that,
not sloppiness.** Its physical codec claims
`node.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 the `CooperativeExec` node 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
shipping `dfx_storage` without one.
**cloudpickle captures a module attribute as the module.** I expected
`pa.compute` in a UDF to fail on a worker; it does not, because cloudpickle
resolves the attribute and stores an import of `pyarrow.compute`. The real trap
is a *function* with a resolvable `module.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_dir` arrives as `Some("")`.
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 a `fallback` — 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/distributed` is the worked example,
`datafusion-ffi-example` is the capsule-protocol test bed,
`datafusion-ffi-query-planner-example` is 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 `datafusion`
wheel.
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 a `SessionContext`, and `df_settings` lists 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, and `extension_codec_durable_metadata` now 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.
- Fixes `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 since `DistributedExec` was added, and the claim
mattered.
CI builds three more wheels and runs three more test suites, gated to `abi3`
like the existing ones.
Verified: `pytest python/` (1443 passed, 12 skipped) and all five example
suites (55, 51, 10, 11, 17). `pre-commit run --all-files` clean except
`actionlint`, which needs Docker and could not run locally; the workflow files
parse as YAML.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
