Let tpchgen-cli shard the data instead of resharding it

`tpchgen-cli` has `--parts N`, which writes a table as N Parquet files into a
directory named for it. That is exactly the shape the storage library's table
provider wants -- it scans `*.parquet` under a directory, one partition per
file -- so `run_tpch.py` can generate what it needs in one call instead of
reading a 220 MB single-file `lineitem` and slicing it into shards itself.

Drops `reshard`, the column list it needed, and the `--data`/`--rows`
arguments, along with the prerequisite that a reader generate a 1 GB dataset
by hand before the script will run at all. `--scale` replaces `--rows` as the
size knob, defaulting to a tenth of scale factor 1.

The CI step installs `tpchgen-cli` itself rather than inheriting it from the
TPC-H dataset step, since it no longer reads that dataset.

Author: timsaucer

.github/workflows/test.yml

@@ -164,11 +164,15 @@ jobs:
           cd examples/tpch
           uv run --no-project pytest _tests.py
 
-      # The script the distributed example's README advertises, against the
-      # dataset generated above -- so the exact entry point a reader will
-      # copy-paste cannot rot while its test suite stays green. The row cap
-      # keeps the reshard and the worker round trip to seconds; correctness
-      # on the full shape is the pytest suites' job, this is the script.
+      # 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 run --no-project python examples/distributed/run_tpch.py --rows 400000
+        run: |
+          uv pip install tpchgen-cli
+          uv run --no-project python examples/distributed/run_tpch.py --scale 0.1

examples/distributed/README.md

@@ -66,13 +66,19 @@ Each library's tests can be run from its own directory the same way, without
 the sibling installs. `storage-library` and `udf-library` need only
 `pytest maturin ../../..`.
 
-Against the real TPC-H data — generate it as
-[`examples/tpch`](../tpch/README.md) describes, then:
+Against real TPC-H data, which the script generates for itself:
 
 ```console
+$ uv pip install tpchgen-cli
 $ uv run python ../run_tpch.py --partitions 4
 ```
 
+`tpchgen-cli` shards natively, so `--partitions 4` asks it for four Parquet
+files and tells the engine to use four partitions — the two are the same
+number because a distributed engine can only spread work as widely as the data
+is split. `--scale` sets how much data; it defaults to a tenth of TPC-H scale
+factor 1, and `--scale 1` is the full ~6M row `lineitem`.
+
 ## What actually happens
 
 The engine's planner splits the plan at the partial aggregate, which is where

examples/distributed/run_tpch.py

@@ -17,37 +17,31 @@
 
 """Run TPC-H Q1 across worker processes, and compare against one process.
 
+    uv pip install tpchgen-cli
     python examples/distributed/run_tpch.py --partitions 4
 
-Needs the TPC-H data the repository's other examples use::
-
-    mkdir -p examples/tpch/data && cd examples/tpch/data
-    uv pip install tpchgen-cli && uv run --no-project tpchgen-cli -s 1 --format=parquet
-
-`tpchgen-cli` writes one file per table, so `lineitem.parquet` is a single
-220 MB file -- one partition, and nothing to fan out. This script re-shards
-the columns Q1 needs into `--partitions` files first, which is also a fair
-illustration of the real constraint: a distributed engine can only spread work
-as widely as the data is split.
+Generates its own `lineitem` with `tpchgen-cli`, which shards natively: one
+`tpchgen-cli parquet --parts N` call writes N Parquet files. That is also a
+fair illustration of the real constraint -- a distributed engine can only
+spread work as widely as the data is split -- so the number of files is the
+same `--partitions` the engine is told to use.
 """
 
 from __future__ import annotations
 
 import argparse
 import pathlib
 import shutil
-import sys
+import subprocess
 import tempfile
 import time
 
 import pyarrow as pa
-import pyarrow.parquet as pq
 from dfx_engine.driver import run_distributed, run_local
 from dfx_engine.session import SessionSpec
 
-# Q1 without the `l_shipdate` filter and the `avg` columns, so the shard below
-# stays small. The shape that matters is unchanged: group by two low-cardinality
-# columns, aggregate, order.
+# Q1 without the `l_shipdate` filter and the `avg` columns. The shape that
+# matters is unchanged: group by two low-cardinality columns, aggregate, order.
 Q1 = """
 select l_returnflag,
        l_linestatus,
@@ -61,35 +55,37 @@
 order by l_returnflag, l_linestatus
 """
 
-COLUMNS = [
-    "l_returnflag",
-    "l_linestatus",
-    "l_quantity",
-    "l_extendedprice",
-    "l_discount",
-    "l_tax",
-]
-
-
-def reshard(
-    source: pathlib.Path, into: pathlib.Path, partitions: int, rows: int
-) -> int:
-    """Write the first `rows` rows of `source` as `partitions` Parquet files."""
-    into.mkdir(parents=True, exist_ok=True)
-    table = pq.read_table(source, columns=COLUMNS)
-    if rows:
-        table = table.slice(0, rows)
-
-    per_file = max(1, table.num_rows // partitions)
-    written = 0
-    for index in range(partitions):
-        offset = index * per_file
-        length = table.num_rows - offset if index == partitions - 1 else per_file
-        if length <= 0:
-            break
-        pq.write_table(table.slice(offset, length), into / f"part-{index}.parquet")
-        written += 1
-    return written
+
+def generate(into: pathlib.Path, partitions: int, scale: float) -> pathlib.Path:
+    """Write `lineitem` as `partitions` Parquet files, and return their directory.
+
+    `tpchgen-cli` puts a sharded table in a subdirectory named for it, so the
+    directory this returns is `into/lineitem` -- which is what the storage
+    library's table provider wants, since it scans `*.parquet` under a
+    directory and makes one partition per file.
+    """
+    executable = shutil.which("tpchgen-cli")
+    if executable is None:
+        message = (
+            "tpchgen-cli not found on PATH; install it with "
+            "`uv pip install tpchgen-cli`"
+        )
+        raise RuntimeError(message)
+
+    subprocess.run(  # noqa: S603
+        [
+            executable,
+            "parquet",
+            f"--scale-factor={scale}",
+            "--tables=lineitem",
+            f"--parts={partitions}",
+            f"--output-dir={into}",
+            "--no-progress",
+            "--quiet",
+        ],
+        check=True,
+    )
+    return into / "lineitem"
 
 
 def compare(table: pa.Table, reference: pa.Table) -> None:
@@ -131,37 +127,20 @@ def compare(table: pa.Table, reference: pa.Table) -> None:
 
 def main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument(
-        "--data",
-        type=pathlib.Path,
-        default=pathlib.Path(__file__).resolve().parents[1]
-        / "tpch"
-        / "data"
-        / "lineitem.parquet",
-    )
     parser.add_argument("--partitions", type=int, default=4)
     parser.add_argument(
-        "--rows",
-        type=int,
-        default=2_000_000,
-        help="rows to use; 0 for all of them (SF 1 lineitem is ~6M)",
+        "--scale",
+        type=float,
+        default=0.1,
+        help="TPC-H scale factor; 1 is the full ~6M row lineitem",
     )
     args = parser.parse_args(argv)
 
-    if not args.data.exists():
-        sys.stderr.write(
-            f"{args.data} not found. Generate it with:\n"
-            "  mkdir -p examples/tpch/data && cd examples/tpch/data\n"
-            "  uv pip install tpchgen-cli\n"
-            "  uv run --no-project tpchgen-cli -s 1 --format=parquet\n"
-        )
-        return 2
-
     workspace = pathlib.Path(tempfile.mkdtemp(prefix="dfx-tpch-"))
     try:
-        data = workspace / "lineitem"
-        count = reshard(args.data, data, args.partitions, args.rows)
-        print(f"resharded into {count} file(s) under {data}")
+        data = generate(workspace, args.partitions, args.scale)
+        count = len(list(data.glob("*.parquet")))
+        print(f"generated {count} file(s) under {data}")
 
         spec = SessionSpec(
             tables={"lineitem": str(data)},