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.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f8db548..b752f2b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml
@@ -164,11 +164,15 @@ 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
diff --git a/examples/distributed/README.md b/examples/distributed/README.md index c9bb23c..78b79e8 100644 --- a/examples/distributed/README.md +++ b/examples/distributed/README.md
@@ -66,13 +66,19 @@ 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
diff --git a/examples/distributed/run_tpch.py b/examples/distributed/run_tpch.py index dae88db..4dff6f5 100644 --- a/examples/distributed/run_tpch.py +++ b/examples/distributed/run_tpch.py
@@ -17,18 +17,14 @@ """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 @@ -36,18 +32,16 @@ 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 generate(into: pathlib.Path, partitions: int, scale: float) -> pathlib.Path: + """Write `lineitem` as `partitions` Parquet files, and return their directory. -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) + `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) - 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 + 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 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)},