Document what a worker has to reproduce, and three findings that cost 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>

Author: timsaucer

docs/source/extension-guide/checklist.md

@@ -56,6 +56,14 @@ publish. Each links to the page that explains it.
 - [ ] **You round-trip a plan in a test and assert *your* codec did the work.**
       Both being installed does not mean your node reached you.
       → {ref}`extension_codec_order`
+- [ ] **You ship a logical codec too, if you contribute a table provider.** A
+      physical codec is not enough: an installed query planner receives the
+      logical plan, which holds your provider, and the session fails to plan
+      without one. → {ref}`extension_codec_provider_logical`
+- [ ] **You decode in a *different process* in at least one test.** A codec
+      that parks the object in a process-global map passes every in-process
+      round trip and fails the first real one.
+      → {ref}`extension_codec_durable_metadata`
 
 ## Bundles and planners
 
@@ -94,6 +102,8 @@ publish. Each links to the page that explains it.
       process-local token. The examples in this repository use tokens to make
       ownership observable; that is a demonstration, not a pattern.
       → {ref}`extension_codec_durable_metadata`
-- [ ] **You have integration tests across a real FFI boundary.** The two
-      example crates in this repository are the pattern: build the cdylib,
-      install the wheel, then exercise it from Python.
+- [ ] **You have integration tests across a real FFI boundary.** The example
+      trees in this repository are the pattern: build the cdylib, install the
+      wheel, then exercise it from Python. `examples/distributed` additionally
+      spawns worker processes, which is the only way to catch a codec that
+      only works in the process that wrote it.

docs/source/extension-guide/codecs.md

@@ -58,15 +58,47 @@ Your payload has to be enough to rebuild the object somewhere your process is
 not. Write the metadata a fresh instance can be constructed from — a path, a
 connection string, a schema, the options the object was created with.
 
-The example codecs in this repository do not do this, and it is worth knowing
-before copying them. They keep a process-local `HashMap` of live providers and
-encode an integer token into it: encoding inserts, decoding removes. That makes
-Rust type identity observable across three separately loaded libraries in one
-test, which is what the examples exist to show. It also means a decode consumes
-its token, so the same bytes cannot be decoded twice, one encoded plan cannot
-fan out to several readers, and a plan that never reaches a decoder keeps its
-provider alive for the life of the process. A real codec has none of those
-properties because it does not park the object anywhere.
+Two of the example codecs in this repository do not do this, and it is worth
+knowing before copying them. `datafusion-ffi-example` and
+`datafusion-ffi-query-planner-example` keep a process-local `HashMap` of live
+providers and encode an integer token into it: encoding inserts, decoding
+removes. That makes Rust type identity observable across three separately
+loaded libraries in one test, which is what those examples exist to show. It
+also means a decode consumes its token, so the same bytes cannot be decoded
+twice, one encoded plan cannot fan out to several readers, and a plan that
+never reaches a decoder keeps its provider alive for the life of the process.
+A real codec has none of those properties because it does not park the object
+anywhere.
+
+For one that does it properly, read
+[`examples/distributed/storage-library`](https://github.com/apache/datafusion-python/tree/main/examples/distributed/storage-library).
+Its payload is the file paths, the projection, the row limit, and the schema —
+enough to rebuild the scan from nothing — and its tests decode a plan in a
+separate interpreter that never registered the table.
+
+(extension_codec_provider_logical)=
+
+## A table provider needs a *logical* codec
+
+A provider library can reasonably conclude it needs only a physical codec: its
+scan is a physical node, so that is where its own type appears. That holds
+right up until someone installs a query planner.
+
+An FFI query planner is handed the **logical** plan, as protobuf. A logical
+plan holds its tables as `Arc<dyn TableProvider>`, and the default codec's
+`try_encode_table_provider` is unimplemented. So a session with your provider
+and any engine installed fails while planning, before anything is executed,
+with:
+
+```text
+Error serializing custom table ... caused by
+Execution error: No installed extension codec handled a table provider
+```
+
+Implement `try_encode_table_provider` and `try_decode_table_provider`, and
+contribute the logical codec alongside the physical one. The payload can be
+small — the storage library writes just the directory, because everything else
+it holds is read back from there — but it has to exist.
 
 (extension_codec_ids)=
 

docs/source/extension-guide/query-planners.md

@@ -38,6 +38,53 @@ against the codecs of the session that will run the query.
 `MyQueryPlanner` in [`datafusion-ffi-query-planner-example`] is the worked
 implementation.
 
+(planner_host_optimizer_rules)=
+
+## Plan against your own optimizer rules
+
+Your planner returns its plan as **protobuf**, not as a handle. Every query
+therefore serializes what you produce, and anything in it that cannot be
+encoded is your problem rather than a distant one.
+
+That matters because of where physical optimization runs. Physical planning
+applies `session.physical_optimizers()`, and when the session arrived over FFI
+those rules are the *host's* — so each one runs back across the boundary and
+hands you a `ForeignExecutionPlan` wrapping the result. `EnsureCooperative` is
+on by default and will do exactly this. A stock `CooperativeExec` produced that
+way has no reachable `try_to_proto`, so a node that is perfectly serializable
+in the process that made it becomes unserializable in yours:
+
+```text
+Internal error: Unsupported plan and extension codec failed with
+[This feature is not implemented: PhysicalExtensionCodec is not provided].
+Plan: ForeignExecutionPlan { name: "CooperativeExec", ... }
+```
+
+A foreign node is also opaque to `downcast_ref`, so a planner that means to
+*rewrite* the plan — inserting stages, say — cannot inspect what it was given.
+
+Both problems go away if the rules run on your side. Wrap the session you were
+handed in one that delegates everything except `physical_optimizers()`, and
+return the stock rule set from there:
+
+```rust
+let local = LocalOptimizerSession::new(session);   // owns PhysicalOptimizer::default().rules
+DefaultPhysicalPlanner::default()
+    .create_physical_plan(logical_plan, &local)
+    .await?
+```
+
+`LocalOptimizerSession` in
+[`examples/distributed/engine-library`](https://github.com/apache/datafusion-python/tree/main/examples/distributed/engine-library)
+is about twenty delegating methods and one override.
+
+Delegating to a `fallback` avoids the problem differently, by not planning at
+all: the plan comes back from whoever you delegated to, already concrete. That
+is the right choice for a planner that only layers behaviour on another, and
+the wrong one for a planner that needs to rewrite the result — you cannot
+rewrite a subtree you hold an opaque handle to. A planner does one or the
+other.
+
 ## One planner per session
 
 A session holds exactly one query planner. Calling `set_query_planner` again

docs/source/user-guide/distributing-work/expressions.md

@@ -119,14 +119,30 @@ requirements on the worker environment:
   stamps the sender's `(major, minor)`; mismatches raise a clear
   error naming both versions. Align the Python version on driver and
   workers.
-- **Imported modules must be importable on the worker.** cloudpickle
-  captures the callable *by value* (bytecode and closure cells travel
-  whole), but names resolved through `import` are captured *by
-  reference* — module path only. A UDF doing
-  `from mylib import transform` requires `mylib` installed on the
-  worker. Same applies to bound methods of imported classes.
-  Self-contained UDFs (no imports beyond what the worker already has,
-  e.g. `pyarrow`) avoid this entirely.
+- **Anything the callable names must be reachable on the worker.**
+  cloudpickle captures the function's own body *by value* — bytecode and
+  closure cells travel whole — but every global it refers to is captured *by
+  reference* if cloudpickle can resolve it to an importable
+  `module.qualname`. The worker then imports it by that path.
+
+  So the rule is not "imports are bad", it is **whether the name has an
+  importable home**:
+
+  | The callable refers to | Travels as | Worker needs |
+  | --- | --- | --- |
+  | a nested or `__main__`-level function | the function itself | nothing |
+  | a module, including a submodule like `pyarrow.compute` | an import of that module | the module installed |
+  | a function in an importable module of yours | a pointer to `yourmod.helper` | **your code installed** |
+
+  The third row is the one that surprises people, and the size difference
+  makes it concrete: one small function is around 1 kB pickled from
+  `__main__` and around 30 bytes from an importable module, because the
+  second is only a pointer. Moving a helper out of a script and into a
+  package silently changes what gets shipped.
+
+  It fails on the worker as a bare
+  `ModuleNotFoundError: No module named 'yourmod'`, raised while the plan is
+  being decoded, with nothing in the message about UDFs or serialization.
 
 ## Registering shared UDFs on workers
 

docs/source/user-guide/distributing-work/query-engines.md

@@ -71,11 +71,64 @@ If you install more than one library, pass them in one
 {ref}`user_guide_extensions` for the details of installing extension
 libraries, and {ref}`ffi` if you want to write an engine yourself.
 
+(distributed_worker_parity)=
+
+## What a worker has to reproduce
+
+An engine ships your plan to a process that has never seen your session. That
+process has to be able to rebuild everything the plan refers to, and there is
+**no way to snapshot a {py:class}`~datafusion.SessionContext` and restore it
+somewhere else**: {py:class}`~datafusion.SessionConfig` is write-only from
+Python, and while `information_schema.df_settings` can be read back, it lists
+`datafusion.runtime.*` keys that have no configuration namespace to set them
+into again.
+
+So parity is not automatic. It is something you build the same way twice, and
+these are the things that have to match. Most engines handle several of them
+for you — check which.
+
+- **Codec ids.** A plan records which codec wrote each payload, and decoding
+  routes on that id. Pin ids with `__datafusion_codec_id__` rather than
+  letting them default to a class's import path, and compare
+  {py:meth}`~datafusion.SessionContext.physical_extension_codec_ids` on both
+  sides before shipping anything. See {ref}`extension_codec_ids`.
+- **Functions the plan names.** A function resolves either from the receiving
+  session's registry or from a codec. Either is enough; neither is automatic.
+  A Python UDF is the exception — it travels inside the plan.
+- **Object stores**, registered *before* the plan is decoded rather than
+  before it is executed. Decoding a Parquet scan resolves its store.
+- **Config extensions**, installed before any namespaced key is set. Setting a
+  key in a namespace that has not been declared is an error, not a no-op.
+- **The Python minor version**, if any inline Python UDF is involved.
+  Cloudpickle payloads are stamped with the sender's version and refuse to
+  load on another. Launching workers with `sys.executable` makes this true by
+  construction; a hardcoded `python` does not.
+- **The `cloudpickle` version**, which is *not* stamped. Cross-version loading
+  usually works and is not guaranteed. Pin it.
+- **`target_partitions`**, if a worker re-plans anything. Left to default it
+  follows the core count, so two differently-sized machines disagree.
+
+Two more that are about lifetime rather than configuration:
+
+- **One session per worker, alive for the whole process.** An FFI codec
+  resolves names against the session captured when its bundle was installed,
+  and that reference is weak — see {ref}`extension_sessions`.
+- **Pass a context to `to_bytes`.** {py:meth}`~datafusion.ExecutionPlan.to_bytes`
+  takes an optional context, and without one it uses an empty codec chain that
+  cannot encode any library's nodes. The argument being optional makes this
+  easy to miss, because the failure only appears once an extension node is in
+  the plan.
+
 ## Available engines
 
 Query-level distribution is being built upstream. Neither project
 below is usable from datafusion-python yet; both sections will fill
-in as the integrations land.
+in as the integrations land. In the meantime the repository contains a
+worked example you can read and run:
+[`examples/distributed`](https://github.com/apache/datafusion-python/tree/main/examples/distributed)
+splits a query across worker processes using three separate extension
+libraries, and is written to make each of the requirements above visible —
+including the ways they fail.
 
 ### datafusion-distributed
 

examples/distributed/README.md

@@ -0,0 +1,147 @@
+<!---
+  Licensed to the Apache Software Foundation (ASF) under one
+  or more contributor license agreements.  See the NOTICE file
+  distributed with this work for additional information
+  regarding copyright ownership.  The ASF licenses this file
+  to you under the Apache License, Version 2.0 (the
+  "License"); you may not use this file except in compliance
+  with the License.  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing,
+  software distributed under the License is distributed on an
+  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  KIND, either express or implied.  See the License for the
+  specific language governing permissions and limitations
+  under the License.
+-->
+
+# Three libraries, one distributed query
+
+A worked example of what `datafusion-python`'s extension protocol is *for*:
+several independently compiled libraries, none of which knows about the
+others, cooperating on a single query whose work runs in separate operating
+system processes.
+
+Everything here is real. The workers are separate interpreters. The plan they
+run was serialized by the driver and decoded by them. If you break the
+serialization, the tests fail.
+
+## The three libraries
+
+| Crate | Owns | Installed with |
+| --- | --- | --- |
+| `udf-library` (`dfx_udfs`) | a scalar function, an aggregate, a window function | **by hand** — `register_udf` plus two `with_*_extension_codec` calls |
+| `storage-library` (`dfx_storage`) | a Parquet table provider and its own scan node | `with_extensions` |
+| `engine-library` (`dfx_engine`) | a query planner, a stage node, and the driver/worker machinery | `with_extensions` |
+
+One of them is deliberately old-fashioned. `dfx_udfs` exposes no
+`__datafusion_session_components__`, so it cannot be installed as a bundle and
+its caller has to do five things in the right order instead of one. That is
+not a strawman: `SessionExtensionComponents` carries codec fields only, so a
+library that contributes *functions* has nowhere to put them today. Mixed
+setups are the normal case, and this example shows what one costs.
+
+## Running it
+
+```console
+$ cd examples/distributed/engine-library
+$ uv venv && uv pip install pytest pyarrow ../.. ../storage-library ../udf-library
+$ uv run maturin develop
+$ uv run pytest python/tests/_test*.py
+```
+
+Against the real TPC-H data — generate it as
+[`examples/tpch`](../tpch/README.md) describes, then:
+
+```console
+$ uv run python ../run_tpch.py --partitions 4
+```
+
+## What actually happens
+
+The engine's planner splits the plan at the partial aggregate, which is where
+DataFusion has already split it for its own reasons: a `GROUP BY` becomes a
+partial pass per input partition and a final pass that merges them, and the
+partial passes are independent by construction.
+
+```
+SortPreservingMergeExec
+  ProjectionExec
+    AggregateExec: mode=FinalPartitioned          <- driver merges
+      RepartitionExec: Hash([l_returnflag], 2)
+        FFI_ExecutionPlan: ShuffleStageExec       <- shipped to workers
+          AggregateExec: mode=Partial             <- one worker per partition
+            FFI_ExecutionPlan: PartitionedParquetExec
+```
+
+The driver serializes the `ShuffleStageExec` subtree, starts one worker per
+partition, and waits. Each worker rebuilds an equivalent session, decodes the
+plan, runs *its* partition, and writes the result to an Arrow IPC file. The
+driver then runs the whole query itself — and the stage node, finding the
+files already there, streams them instead of recomputing.
+
+One node does both halves of that exchange, which is why nothing has to
+rewrite the plan in between. It also means a query run with no workers at all
+still gets the right answer; it just does the work itself.
+
+## The four things worth reading
+
+**`engine-library/python/dfx_engine/session.py`** is the point of the whole
+example. There is no way to snapshot a `SessionContext` and restore it
+elsewhere, so worker parity cannot be automated — it has to be *built the same
+way twice*, from data small enough to put in a message. Both the driver and
+every worker call one `build_session`. Anything a query depends on that is not
+in the `SessionSpec` is a bug waiting for a worker to find it.
+
+**`storage-library/src/codec.rs`** is the repository's only codec that encodes
+durable metadata. The others park the live object in a process-global map and
+encode an integer token, which is fine for making Rust type identity
+observable in a test and useless the moment the bytes leave the process. This
+one writes the file paths, the projection, and the schema, so the same bytes
+decode twice, decode on ten workers, and decode tomorrow.
+
+**`udf-library/python/tests/_test_udfs.py`** shows that installing a
+library's codec is an *alternative* to registering its functions, not an
+addition. Three workers, three configurations:
+
+| worker has | result |
+| --- | --- |
+| the codec, no registrations | works; the codec rebuilds each function from its name |
+| the registrations, no codec | works; the registry answers first and the codec is never consulted |
+| neither | fails, naming `dfx_net_revenue` |
+
+The middle row is the trap. On the driver, where the functions are always
+registered, a broken or missing codec looks completely fine.
+
+**`engine-library/python/tests/_test_three_libraries.py`** runs the queries,
+and pins the failure modes next to the successes — including a Python UDF that
+works on the driver and fails on the worker.
+
+## Things this example is not
+
+It writes shuffle results to local files, so "distributed" means several
+processes on one machine. Adding a network is a transport change and would not
+alter anything above it.
+
+It holds one partition of results in memory before writing, because an Arrow
+IPC stream needs its schema up front. A production engine would stream to the
+file and track completion separately.
+
+It has one stage. A real engine chains them, and the interesting problems —
+scheduling, retries, straggler handling, memory limits — all live in the part
+this example replaces with `subprocess.Popen` and a `for` loop.
+
+It is slower than running the query in one process. Four processes on one
+laptop cannot beat one process that skips a round trip through Arrow IPC
+files. The comparison the tests make is *agreement*, not speed.
+
+## Further reading
+
+- [Distributed query engines](https://datafusion.apache.org/python/user-guide/distributing-work/query-engines.html)
+  — using an engine, and the checklist for what a worker has to reproduce.
+- [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html)
+  — writing a library like these.
+- [Encode metadata, not a handle to a live object](https://datafusion.apache.org/python/extension-guide/codecs.html#encode-metadata-not-a-handle-to-a-live-object)
+  — what a codec should put on the wire, and why.