Add the queries and the cross-library integration tests

Seventeen tests for #1719, every one running real worker processes, plus
`run_tpch.py` for the same thing against the generated TPC-H data.

The four queries: a Q1-shaped distributed aggregate; the same with `dfx_udfs`'
Rust scalar and aggregate functions; an inline Python UDF; and the storage
library's provider read on the workers. Each compares the distributed answer
against the single-process answer through the same session factory, because
disagreement there is the only reliable signal that a split is wrong.

Tests use a small hand-checked fixture rather than the real dataset.
`tpchgen-cli` writes one file per table, so SF-1 `lineitem` is a single 220 MB
file — one partition, and nothing to fan out. `run_tpch.py` re-shards it
first, which is a fair illustration of the actual constraint: an engine can
only spread work as widely as the data is split.

Three things this pass turned up.

**An empty `shuffle_dir` was writing files into the working directory.** A
registered config extension always *has* an entry, so an unset directory
arrives as `Some("")` rather than `None`, and the planner treated it as
configured. The stage node's paths were then relative to wherever the process
happened to be, so `run_local` scattered `stage-1-part-*.arrow` next to the
caller and later queries read another query's leftovers back out of them —
which is how six tests failed with "Batch has 3 columns but BatchCoalescer
expects 5". Four of those files had already been committed by the previous
change; they are deleted here.

**cloudpickle captures a module attribute as the module, not as its parent.**
I expected `pa.compute` inside a UDF to fail on a worker, since `import
pyarrow` does not bind `pyarrow.compute` and nothing loads it transitively.
It does not fail: cloudpickle resolves the attribute and stores an import of
`pyarrow.compute` itself, so the worker imports the submodule on load. The
real trap is a *function* with a resolvable `module.qualname` — the same
callable is 1106 bytes pickled from `__main__` and 34 bytes from an importable
module, because the second is a pointer. A helper at test-module scope
therefore reaches the worker as `ModuleNotFoundError: No module named
'_test_three_libraries'`, with `traceback: None` and nothing naming a UDF, a
plan, or serialization. Both halves are pinned as tests.

**An FFI query planner encodes its own output on every query.** It returns
proto bytes rather than a plan handle, so both libraries' codecs show one
encode apiece straight after `execution_plan()`, before the driver has asked
for any bytes. Worth knowing before reading an encode counter as "this is what
shipping cost".

`run_tpch.py` compares floats with a tolerance rather than for equality:
splitting a `sum` across partitions changes the order the additions happen in,
and floating point addition is not associative, so the low bits of
`sum_charge` differ legitimately between the two runs. Any distributed engine
has this property, and someone diffing two runs should not conclude the split
is broken.

Verified: 400k rows of real `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.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: timsaucer

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

@@ -47,6 +47,7 @@
     import pyarrow as pa
     from datafusion import DataFrame, SessionContext
     from datafusion.plan import ExecutionPlan
+    from datafusion.user_defined import ScalarUDF
 
 __all__ = ["DistributedResult", "find_stage", "run_distributed"]
 
@@ -116,17 +117,27 @@ def _dispatch(
     )
 
 
-def run_distributed(sql: str, spec: SessionSpec) -> DistributedResult:
+def run_distributed(
+    sql: str, spec: SessionSpec, extra_udfs: list[ScalarUDF] | None = None
+) -> DistributedResult:
     """Run `sql`, executing its leaf stage in one worker process per partition.
 
     Requires ``spec.shuffle_dir``: without it the planner inserts no stage and
     there is nothing to distribute.
+
+    ``extra_udfs`` are registered on the driver only. They have to be here for
+    the query to *plan*, but not on the worker: a Python UDF is cloudpickled
+    into the plan and travels by value, unlike the Rust functions in
+    :func:`~dfx_engine.session.build_session`, which travel by name and so
+    have to exist on both sides.
     """
     if not spec.shuffle_dir:
         message = "run_distributed needs a shuffle_dir; build_session got none"
         raise ValueError(message)
 
     ctx, engine, _storage = build_session(spec)
+    for function in extra_udfs or []:
+        ctx.register_udf(function)
     plan = ctx.sql(sql).execution_plan()
 
     stage = find_stage(plan)
@@ -184,11 +195,14 @@ def run_distributed(sql: str, spec: SessionSpec) -> DistributedResult:
     return DistributedResult(batches, partitions, worker_rows)
 
 
-def run_local(sql: str, spec: SessionSpec) -> list[pa.RecordBatch]:
+def run_local(
+    sql: str, spec: SessionSpec, extra_udfs: list[ScalarUDF] | None = None
+) -> list[pa.RecordBatch]:
     """Run `sql` in this process, for comparison.
 
     Uses the same session factory with no shuffle directory, so the only
-    difference from :func:`run_distributed` is where the work happened.
+    difference from :func:`run_distributed` is where the work happened. Any
+    disagreement between the two is a bug in the split.
     """
     ctx, _engine, _storage = build_session(
         SessionSpec(
@@ -197,6 +211,8 @@ def run_local(sql: str, spec: SessionSpec) -> list[pa.RecordBatch]:
             target_partitions=spec.target_partitions,
         )
     )
+    for function in extra_udfs or []:
+        ctx.register_udf(function)
     return ctx.sql(sql).collect()