Clear the small stuff off the distributed example

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>

Author: timsaucer

.github/workflows/build.yml

@@ -186,7 +186,7 @@ jobs:
           manylinux: "2_28"
 
       # The example libraries below are test fixtures, not release artifacts:
-      # the `test-ffi-manylinux-x86_64` artifact they feed is consumed only by
+      # 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
@@ -261,11 +261,11 @@ jobs:
           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/*

.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:"

examples/distributed/engine-library/pyproject.toml

@@ -27,6 +27,18 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
 ]
 dynamic = ["version"]
+# No `dependencies`, deliberately, though `dfx_engine.session` imports
+# `dfx_storage` and `dfx_udfs` at module scope and `import dfx_engine` fails
+# without them. Declaring them does not work: neither is published, so
+# resolution fails before a wheel is built at all --
+#
+#   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.
+#
+# which would break the build rather than document the requirement. The
+# requirement is in the README instead, next to the install command that
+# satisfies it.
 
 [tool.maturin]
 features = ["pyo3/extension-module"]

examples/distributed/engine-library/python/dfx_engine/driver.py

@@ -142,6 +142,30 @@ def _dispatch(
     )
 
 
+def _report(task: tuple[int, int], stdout: str) -> int:
+    """Read a worker's row count off its last line of output.
+
+    The last line, not the whole stream: a worker's stdout is shared with
+    everything loaded into it, and one stray `print` from a library -- or a
+    warning some future dependency decides to write there -- would turn
+    `json.loads` on the whole buffer into a confusing failure a long way from
+    its cause.
+    """
+    stage_id, partition = task
+    lines = [line for line in stdout.splitlines() if line.strip()]
+    if not lines:
+        message = f"stage {stage_id} partition {partition} printed no report"
+        raise RuntimeError(message)
+    try:
+        return json.loads(lines[-1])["rows"]
+    except (ValueError, KeyError) as err:
+        message = (
+            f"stage {stage_id} partition {partition} printed an unreadable "
+            f"report {lines[-1]!r}"
+        )
+        raise RuntimeError(message) from err
+
+
 def run_distributed(
     sql: str, spec: SessionSpec, extra_udfs: list[ScalarUDF] | None = None
 ) -> DistributedResult:
@@ -216,7 +240,7 @@ def run_distributed(
             stage_id, partition = task
             failures.append(f"stage {stage_id} partition {partition} failed:\n{stderr}")
             continue
-        task_rows[task] = json.loads(stdout)["rows"]
+        task_rows[task] = _report(task, stdout)
 
     if failures:
         raise RuntimeError("\n".join(failures))

examples/distributed/engine-library/src/stage.rs

@@ -125,6 +125,14 @@ impl ShuffleStageExec {
         let final_path = partition_path(&self.shuffle_dir, self.stage_id, partition);
         let temp_path = temp_partition_path(&self.shuffle_dir, self.stage_id, partition);
         let shuffle_dir = self.shuffle_dir.clone();
+        // The schema the file is written with is the one this stream declares,
+        // not the one the child's stream happens to report. They agree for any
+        // well-behaved child, and pinning it to the declaration is what makes
+        // a disagreement loud: `StreamWriter::write` rejects a batch whose
+        // schema differs, so a child that contradicts its own `schema()` fails
+        // here instead of publishing a file readers were told to expect
+        // something else from.
+        let written_schema = Arc::clone(&schema);
 
         let collected = async move {
             let mut stream = input.execute(partition, context)?;
@@ -139,7 +147,7 @@ impl ShuffleStageExec {
                 let file = fs::File::create(&temp_path).map_err(|err| {
                     exec_datafusion_err!("dfx_engine: creating {}: {err}", temp_path.display())
                 })?;
-                let mut writer = StreamWriter::try_new(file, stream.schema().as_ref())
+                let mut writer = StreamWriter::try_new(file, written_schema.as_ref())
                     .map_err(|err| exec_datafusion_err!("dfx_engine: ipc writer: {err}"))?;
                 for batch in &batches {
                     writer

examples/distributed/run_tpch.py

@@ -92,6 +92,39 @@ def reshard(
     return written
 
 
+def compare(table: pa.Table, reference: pa.Table) -> None:
+    """Raise unless `table` matches `reference`, floats to 1e-6 relative.
+
+    Floats get a tolerance rather than equality. Splitting a `sum` across
+    partitions changes the order the additions happen in, and floating point
+    addition is not associative, so the last bits of `sum_charge` legitimately
+    differ between the two runs. Every distributed engine has this property;
+    it is worth knowing before someone diffs two runs and concludes the split
+    is broken.
+    """
+    if table.column_names != reference.column_names:
+        message = (
+            f"column names differ: {table.column_names} vs {reference.column_names}"
+        )
+        raise ValueError(message)
+
+    for name in table.column_names:
+        got = table.column(name).to_pylist()
+        want = reference.column(name).to_pylist()
+        if len(got) != len(want):
+            message = f"{name}: {len(got)} rows distributed, {len(want)} local"
+            raise ValueError(message)
+        for lhs, rhs in zip(got, want, strict=True):
+            close = (
+                abs(lhs - rhs) <= 1e-6 * max(1.0, abs(rhs))
+                if isinstance(lhs, float)
+                else lhs == rhs
+            )
+            if not close:
+                message = f"{name}: {lhs!r} distributed, {rhs!r} local"
+                raise ValueError(message)
+
+
 def main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
@@ -150,24 +183,11 @@ def main(argv: list[str] | None = None) -> int:
         table = pa.Table.from_batches(result.batches)
         reference = pa.Table.from_batches(local)
 
-        # Compared with a tolerance, not for equality. Splitting a `sum` across
-        # partitions changes the order the additions happen in, and floating
-        # point addition is not associative -- so the last bits of `sum_charge`
-        # legitimately differ between the two runs. Any distributed engine has
-        # this property; it is worth knowing before someone diffs two runs and
-        # concludes the split is broken.
-        assert table.column_names == reference.column_names
-        for name in table.column_names:
-            got, want = (
-                table.column(name).to_pylist(),
-                reference.column(name).to_pylist(),
-            )
-            assert len(got) == len(want), name
-            for lhs, rhs in zip(got, want, strict=True):
-                if isinstance(lhs, float):
-                    assert abs(lhs - rhs) <= 1e-6 * max(1.0, abs(rhs)), (name, lhs, rhs)
-                else:
-                    assert lhs == rhs, (name, lhs, rhs)
+        # Raises rather than asserts. This comparison is the only thing that
+        # makes the script a check rather than a demo, and `python -O` removes
+        # an `assert` -- which would leave it printing a table it never
+        # verified.
+        compare(table, reference)
 
         print("\nsame answer both ways (floats to within 1e-6 relative):\n")
         names = table.column_names

examples/distributed/storage-library/src/lib.rs

@@ -22,7 +22,7 @@
 
 use pyo3::prelude::*;
 
-use crate::extension::{BundledPhysicalCodec, DfxStorageExtension};
+use crate::extension::{BundledLogicalCodec, BundledPhysicalCodec, DfxStorageExtension};
 use crate::table_provider::PyPartitionedParquetTable;
 
 mod codec;
@@ -33,6 +33,9 @@ mod table_provider;
 #[pymodule]
 fn dfx_storage(m: &Bound<'_, PyModule>) -> PyResult<()> {
     pyo3_log::init();
+    // Both bundled codecs, so that `module = "dfx_storage"` on each is true
+    // and a caller inspecting a session's codecs sees a type it can look up.
+    m.add_class::<BundledLogicalCodec>()?;
     m.add_class::<BundledPhysicalCodec>()?;
     m.add_class::<DfxStorageExtension>()?;
     m.add_class::<PyPartitionedParquetTable>()?;