Add a worked example of three extension libraries in one distributed query

.github/workflows/build.yml

@@ -185,41 +185,93 @@ jobs:
           features: "protoc,substrait"
           manylinux: "2_28"
 
+      # The example libraries below are test fixtures, not release artifacts:
+      # the `test-example-wheels-x86_64` artifact they feed is consumed only by
+      # test.yml, which installs them on a runner like this one. They therefore
+      # build on the host (`container: off`) rather than in the manylinux
+      # container, which drops a container start and an in-container rustup
+      # install per library and lets the Swatinem/rust-cache above cover them.
+      #
+      # None of them pass `target`, so their artifacts land in `target/debug`
+      # alongside the priming build's. Naming a target would move them under
+      # `target/<triple>/debug` and the priming build would be wasted.
+      #
+      # maturin runs one `cargo rustc` per library and cargo resolves features
+      # per package, so five separate invocations each rebuild the shared
+      # datafusion crates under their own feature union -- measured at over a
+      # minute for a single transition between two of these libraries. One
+      # build spanning all five resolves that union once; afterwards each
+      # maturin step only compiles its own leaf crate.
+      - name: Prime example crate dependencies
+        if: matrix.python-tag == 'abi3'
+        run: >
+          cargo build --lib
+          -p datafusion-ffi-example
+          -p datafusion-ffi-query-planner-example
+          -p dfx-storage
+          -p dfx-udfs
+          -p dfx-engine
+
       # FFI test wheel only needs to be built once per platform; gate to abi3.
       - name: Build FFI provider test library
         if: matrix.python-tag == 'abi3'
         uses: PyO3/maturin-action@v1
         with:
-          target: x86_64-unknown-linux-gnu
-          manylinux: "2_28"
+          container: "off"
           working-directory: examples/datafusion-ffi-example
           args: --out dist
-          rustup-components: rust-std
 
       - name: Build FFI query planner test library
         if: matrix.python-tag == 'abi3'
         uses: PyO3/maturin-action@v1
         with:
-          target: x86_64-unknown-linux-gnu
-          manylinux: "2_28"
+          container: "off"
           working-directory: examples/datafusion-ffi-query-planner-example
           args: --out dist
-          rustup-components: rust-std
+
+      # The three libraries of the distributed example. Listed in dependency
+      # order for readability only; they are independent cdylibs.
+      - name: Build distributed example storage library
+        if: matrix.python-tag == 'abi3'
+        uses: PyO3/maturin-action@v1
+        with:
+          container: "off"
+          working-directory: examples/distributed/storage-library
+          args: --out dist
+
+      - name: Build distributed example UDF library
+        if: matrix.python-tag == 'abi3'
+        uses: PyO3/maturin-action@v1
+        with:
+          container: "off"
+          working-directory: examples/distributed/udf-library
+          args: --out dist
+
+      - name: Build distributed example engine library
+        if: matrix.python-tag == 'abi3'
+        uses: PyO3/maturin-action@v1
+        with:
+          container: "off"
+          working-directory: examples/distributed/engine-library
+          args: --out dist
 
       - name: Archive wheels
         uses: actions/upload-artifact@v7
         with:
           name: dist-manylinux-x86_64-${{ matrix.python-tag }}
           path: dist/*
 
-      - name: Archive FFI test wheel
+      - name: Archive example test wheels
         if: matrix.python-tag == 'abi3'
         uses: actions/upload-artifact@v7
         with:
-          name: test-ffi-manylinux-x86_64
+          name: test-example-wheels-x86_64
           path: |
             examples/datafusion-ffi-example/dist/*
             examples/datafusion-ffi-query-planner-example/dist/*
+            examples/distributed/storage-library/dist/*
+            examples/distributed/udf-library/dist/*
+            examples/distributed/engine-library/dist/*
 
   # ============================================
   # Build - Linux ARM64

.github/workflows/test.yml

@@ -73,11 +73,11 @@ jobs:
           path: wheels/
 
       # FFI test wheel only built once (under the abi3 matrix entry in build.yml).
-      - name: Download pre-built FFI test wheel
+      - name: Download pre-built example test wheels
         if: matrix.wheel-tag == 'abi3'
         uses: actions/download-artifact@v8
         with:
-          name: test-ffi-manylinux-x86_64
+          name: test-example-wheels-x86_64
           path: wheels/
 
       - name: Install from pre-built wheels
@@ -93,9 +93,9 @@ jobs:
           uv venv --python "${{ steps.setup-python.outputs.python-path }}"
           VENV_PY="$PWD/.venv/bin/python"
           uv sync --python "$VENV_PY" --dev --no-install-package datafusion
-          # Search recursively: the FFI artifact bundles more than one
-          # project, so upload-artifact keeps a `<project>/dist/` prefix
-          # and the wheels are not all at the top of wheels/.
+          # Search recursively: the example artifact bundles five projects,
+          # so upload-artifact keeps a `<project>/dist/` prefix and the
+          # wheels are not all at the top of wheels/.
           WHEELS=$(find wheels/ -name "*.whl")
           if [ -n "$WHEELS" ]; then
             echo "Installing wheels:"
@@ -119,14 +119,29 @@ jobs:
           # free-threaded build and re-pick the system 3.12 (see install step).
           uv run --python "$PWD/.venv/bin/python" --no-project pytest -v --import-mode=importlib
 
-      # FFI + TPC-H examples only need to run once; gate to abi3 entries.
-      - name: FFI unit tests
+      # Gated to the abi3 entries because the free-threaded build has no
+      # example wheels to test against -- not to run once. This is five of
+      # the six matrix entries, and running the suites against 3.10 through
+      # 3.14 is the point: the example wheels are abi3, so the interpreter
+      # underneath them is the only thing that varies, and these suites are
+      # what exercise the capsule protocol from Python. Four seconds for the
+      # two FFI examples as of this writing.
+      - name: Example library tests
         if: matrix.wheel-tag == 'abi3'
         run: |
           cd examples/datafusion-ffi-example
           uv run --no-project pytest python/tests/_test*.py
           cd ../datafusion-ffi-query-planner-example
           uv run --no-project pytest python/tests/_test*.py
+          # The distributed example. Its tests spawn worker processes with
+          # `sys.executable`, so they need the same interpreter the wheels
+          # were installed into -- which `uv run` gives them.
+          cd ../distributed/storage-library
+          uv run --no-project pytest python/tests/_test*.py
+          cd ../udf-library
+          uv run --no-project pytest python/tests/_test*.py
+          cd ../engine-library
+          uv run --no-project pytest python/tests/_test*.py
 
       - name: Run tpchgen-cli to create 1 Gb dataset
         if: matrix.wheel-tag == 'abi3'

.gitignore

@@ -36,3 +36,12 @@ docs/mdbook/book
 
 .pyo3_build_config
 
+examples/distributed/*/.venv/
+# Left behind by the `uv` commands in examples/distributed/README.md. The
+# example projects are not locked -- only the root project is.
+examples/distributed/*/uv.lock
+# Shuffle output. Every supported path writes these under a temporary
+# directory, so one appearing in the tree means a relative shuffle_dir
+# reached a stage node and the files landed in the working directory
+# instead -- which is how three of them were committed once already.
+examples/distributed/**/stage-*-part-*.arrow

Cargo.toml

@@ -32,6 +32,9 @@ members = [
   "crates/util",
   "examples/datafusion-ffi-example",
   "examples/datafusion-ffi-query-planner-example",
+  "examples/distributed/storage-library",
+  "examples/distributed/udf-library",
+  "examples/distributed/engine-library",
 ]
 resolver = "3"
 

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
 
@@ -70,8 +78,15 @@ publish. Each links to the page that explains it.
       `with_extensions` refuses a capsule, because there would be nothing to
       name the codec by. → {ref}`extension_bundles_codecs_are_objects`
 - [ ] **Your planner hook wraps `fallback` and delegates to it.** Ignoring it
-      replaces every layer beneath you, which is legal but not composable.
+      replaces every layer beneath you, which is legal but not composable —
+      unless your planner rewrites the plan, as in the next item.
       → {ref}`extension_bundles`
+- [ ] **Your planner plans for itself if it rewrites the plan**, leaving
+      `fallback` unused. The two are exclusive: delegating hands planning back
+      to the host and returns nodes you can neither downcast nor split.
+      Planning for yourself then means supplying your *own* optimizer rules,
+      because a session that arrived over FFI carries the host's.
+      → {ref}`planner_host_optimizer_rules`
 - [ ] **Your planner hook returns `None`, not `fallback`, when it has nothing
       to contribute.** Returning `fallback` installs the session's own planner
       as a foreign one and adds an FFI hop that was not there.
@@ -94,6 +109,13 @@ 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.
+- [ ] **If you ship an engine, say which worker-parity items you handle** and
+      which you leave to your users. A `SessionContext` cannot be snapshotted
+      and restored elsewhere, so every one of them is somebody's job, and your
+      users cannot tell whose from the outside.
+      → {ref}`distributed_worker_parity`

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/index.md

@@ -59,11 +59,27 @@ this section only makes sense once they are distinct in your head:
   the codecs that serialize them.
 - **A planner library** — owns a query planner and the configuration it needs.
 
-The worked examples in this repository use two separate crates,
-[`datafusion-ffi-example`] and [`datafusion-ffi-query-planner-example`], so
-each role has a distinct shared-library identity. A real library may play more
-than one role; keeping them separate in the examples is what makes the
-boundaries observable.
+A real library may play more than one role. Keeping them in separate crates is
+what makes the boundaries observable, because each one is then a distinct
+shared library and the FFI conversions are real rather than same-library
+downcasts.
+
+The examples in this repository are three trees, and it is worth knowing which
+one answers your question:
+
+- [`examples/distributed`] is the **worked example**, and the place to start.
+  Three libraries — functions, tables, and an engine — cooperate on one query
+  whose leaf stage runs in separate worker processes. It is also the only
+  example whose plans genuinely leave the process, so it is where the codecs
+  encode durable metadata rather than tokens.
+- [`datafusion-ffi-example`] is the **capsule-protocol test bed**: one of every
+  hook, exercised hard. Read it to see the shape of a getter, not to see a
+  library designed well.
+- [`datafusion-ffi-query-planner-example`] is the **planner-composition test
+  bed**: what happens when more than one library contributes a query planner,
+  and how `fallback` nests them. The distributed example cannot cover this —
+  a planner that rewrites the plan into stages has to plan for itself, so it
+  has no use for a fallback.
 
 The session owns the codecs used for the exchange and supplies them to the
 foreign planner. That is what lets the planner decode provider-owned objects,
@@ -134,3 +150,4 @@ checklist
 
 [`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example
 [`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example
+[`examples/distributed`]: https://github.com/apache/datafusion-python/tree/main/examples/distributed