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

.ai/skills/ffi-capsule-protocol/SKILL.md

@@ -67,8 +67,13 @@ above: it takes a second argument, the planner assembled so far. A session has
 one planner slot, so planners compose by nesting rather than by chaining, and
 the host hands each bundle the previous layer instead of letting it capture one.
 Wrap `fallback` and delegate to it; returning a planner that ignores it discards
-every layer beneath, including one the session already had. It runs after every
-bundle's codecs are installed, so `ctx` carries the final chains.
+every layer beneath, including one the session already had. The one legitimate
+exception is a planner that *rewrites* the plan (splitting it into stages, say):
+it cannot delegate, because a fallback returns nodes it can neither downcast nor
+serialize, so it plans for itself against its own optimizer rules —
+`examples/distributed/engine-library` is the worked case, and
+`docs/source/extension-guide/query-planners.md` argues it. The hook runs after
+every bundle's codecs are installed, so `ctx` carries the final chains.
 
 That is also the only hook where it does. `__datafusion_session_components__`
 runs before anything is installed, so its `ctx` still carries the chains the

.github/workflows/build.yml

@@ -185,41 +185,105 @@ jobs:
           features: "protoc,substrait"
           manylinux: "2_28"
 
-      # FFI test wheel only needs to be built once per platform; gate to abi3.
+      # 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.
+      #
+      # "A runner like this one" is an assumption, not a guarantee: build and
+      # test are separate jobs, and during GitHub's ubuntu-latest image
+      # migrations they can land on different images. A wheel built against a
+      # newer glibc then fails to import on the test runner. If that ever
+      # bites, pin both jobs to the same explicit image -- or put these builds
+      # back in the manylinux container, which made it impossible.
+      #
+      # 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 building them one at a time makes each invocation
+      # rebuild the shared datafusion crates under its own feature union --
+      # measured at over a minute for a single transition between two of these
+      # libraries. One build spanning them all resolves that union once;
+      # afterwards each maturin step only compiles its own leaf crate.
+      #
+      # So this `-p` list has to name every library built below. One that is
+      # missing here still builds, but pays the rebuild this step exists to
+      # avoid -- and pushes the same cost onto whichever library follows it.
+      - 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
+
+      # Every build below is gated to abi3 because these wheels are abi3: one
+      # per platform serves every interpreter test.yml runs them against.
       - 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 distributed example's libraries. 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

@@ -72,12 +72,13 @@ jobs:
           name: dist-manylinux-x86_64-${{ matrix.wheel-tag }}
           path: wheels/
 
-      # FFI test wheel only built once (under the abi3 matrix entry in build.yml).
-      - name: Download pre-built FFI test wheel
+      # The example wheels are abi3 and build.yml builds them under its abi3
+      # entry only, so there is nothing to download for a free-threaded entry.
+      - 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 +94,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 more than one
+          # project, 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 +120,35 @@ 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 that is where the example wheels
+      # exist -- see the download step above -- and not in order to run these
+      # once. Every entry that has wheels runs them, which is the point rather
+      # than a cost to trim: the wheels are abi3, so the extension binaries
+      # are identical across those entries and the interpreter underneath is
+      # the only variable. That is what catches a build claiming the stable
+      # ABI without honouring it, and what exercises cloudpickle's
+      # version-sensitive bytecode handling in the distributed example's
+      # worker round trip.
+      #
+      # The free-threaded entries are the gap this leaves, and the one worth
+      # closing: nothing here exercises the capsule protocol without the GIL.
+      # It needs per-interpreter example wheels in build.yml first.
+      - 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'
@@ -141,3 +163,16 @@ jobs:
         run: |
           cd examples/tpch
           uv run --no-project pytest _tests.py
+
+      # The script the distributed example's README advertises, run exactly as
+      # a reader would copy-paste it -- so that entry point cannot rot while
+      # its test suite stays green. It generates its own `lineitem` rather than
+      # reading the dataset above, which is why `tpchgen-cli` is installed here
+      # too: the step is meant to stand on its own if the TPC-H steps move or
+      # go away. The scale factor keeps generation and the worker round trip to
+      # seconds; correctness on the full shape is the pytest suites' job.
+      - name: Run distributed TPC-H example
+        if: matrix.wheel-tag == 'abi3'
+        run: |
+          uv pip install tpchgen-cli
+          uv run --no-project python examples/distributed/run_tpch.py --scale 0.1

.gitignore

@@ -36,3 +36,14 @@ 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.
+# The trailing `*` also catches the `.<pid>-<n>.arrow.tmp` files a
+# writer stages before its rename, which a crash mid-write leaves behind.
+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/capsule-protocol.md

@@ -88,6 +88,47 @@ original `Arc` and skips the boundary entirely, and only otherwise wraps it in
 a `ForeignTableProvider`. Which one you get is an implementation detail, and
 both implement `TableProvider`.
 
+(extension_foreign_node_display)=
+
+## A foreign node prints as its wrapper
+
+Which one you get stops being invisible in one place you are likely to look:
+the plan's display text.
+
+`ForeignExecutionPlan` implements `DisplayAs` on its own behalf and never calls
+the wrapped node's `fmt_as`. What the receiver prints is the wrapper's format
+with the node's `name()` — the bare type name — embedded in it:
+
+```text
+FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1
+```
+
+not the `ShuffleStageExec: stage=1` that the node writes about itself. Two
+things follow, and both bite the same task: finding your own node in a plan the
+host built, which is how you confirm your planner or optimizer rule fired.
+
+**Match by containment.** Your node's name lands in the middle of the string,
+so `display().startswith("ShuffleStageExec")` and a `^ShuffleStageExec` regex
+are both false. Test for the name being *in* the text.
+
+**Anything `fmt_as` wrote is gone.** The wrapper reports a name and a child
+count; every parameter your node prints about itself — a stage number, a path,
+a partition count — is dropped, and no amount of parsing gets it back. Carry
+what you need inside the encoded plan, or recompute it from something the host
+can still see, such as the node's position in a walk of the tree.
+
+The trap is that neither mistake fails in a small test. The conversion
+described above hands back the original `Arc` when the provider and the
+receiver turn out to be the same shared library, and a node that never crossed
+the boundary prints as itself. So an anchored match — or a regex that reads a
+parameter out of the display — passes for a node your library both built and
+inspected, then silently stops matching once a second library is genuinely
+involved. There is no exception; a branch just stops being taken.
+
+`find_stages` in [`examples/distributed`] does both halves: a containment test,
+and a stage id recomputed from pre-order position because the display no longer
+carries it.
+
 ## Wrapping it in a capsule
 
 In order to share these FFI structures, we need to wrap them in some kind of
@@ -236,3 +277,4 @@ the session that is actually running the query — see
 library to construct a `SessionContext` of its own.
 
 [datafusion-ffi]: https://crates.io/crates/datafusion-ffi
+[`examples/distributed`]: https://github.com/apache/datafusion-python/tree/main/examples/distributed