diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index c9bd75682a1..e825fcdc614 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -235,6 +235,7 @@ jobs: uv run --no-project scripts/capture-file-sizes.py \ vortex-bench/data \ --benchmark ${{ matrix.subcommand }} \ + --dataset-dir ${{ matrix.dataset_dir }} \ --commit ${{ inputs.mode == 'pr' && github.event.pull_request.head.sha || github.sha }} \ -o sizes.json cat sizes.json >> results.json diff --git a/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py b/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py index d0aefcc580f..5c649f3aeac 100644 --- a/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py +++ b/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py @@ -108,6 +108,7 @@ id="clickbench-nvme", benchmark=Benchmark.CLICKBENCH, name="Clickbench on NVME", + dataset_dir="clickbench_partitioned", runs={ "pr": DEFAULT, "pr-compact": COMPACT, diff --git a/bench-orchestrator/bench_orchestrator/ci_matrix/model.py b/bench-orchestrator/bench_orchestrator/ci_matrix/model.py index 587f7f36da4..8889f46116e 100644 --- a/bench-orchestrator/bench_orchestrator/ci_matrix/model.py +++ b/bench-orchestrator/bench_orchestrator/ci_matrix/model.py @@ -39,6 +39,12 @@ class BenchmarkCase: iterations: int | None = None local_dir: str | None = None remote_key: str | None = None + dataset_dir: str | None = None + + @property + def data_dir_name(self) -> str: + """Return the directory under ``vortex-bench/data`` holding this benchmark's data.""" + return self.dataset_dir or self.subcommand @property def subcommand(self) -> str: diff --git a/bench-orchestrator/bench_orchestrator/ci_matrix/render.py b/bench-orchestrator/bench_orchestrator/ci_matrix/render.py index 90aa7047e4f..5a6f462fed3 100644 --- a/bench-orchestrator/bench_orchestrator/ci_matrix/render.py +++ b/bench-orchestrator/bench_orchestrator/ci_matrix/render.py @@ -42,6 +42,7 @@ def _matrix_entry(benchmark: BenchmarkCase, coverage: Coverage) -> dict[str, obj "id": benchmark.id, "subcommand": benchmark.subcommand, "name": benchmark.name, + "dataset_dir": benchmark.data_dir_name, "targets": [target.to_dict() for target in coverage.targets], "data_formats": [fmt.value for fmt in _data_formats(coverage)], } diff --git a/bench-orchestrator/tests/test_matrix.py b/bench-orchestrator/tests/test_matrix.py index a1b202a68e5..3e5ebcc1d9f 100644 --- a/bench-orchestrator/tests/test_matrix.py +++ b/bench-orchestrator/tests/test_matrix.py @@ -168,6 +168,13 @@ def test_resolver_rejects_duplicate_ids() -> None: _ = render_matrix("develop", catalog) +def test_matrix_entries_name_their_data_directory() -> None: + entries = {entry["id"]: entry for entry in render_matrix("develop", CATALOG)} + + assert entries["clickbench-nvme"]["dataset_dir"] == "clickbench_partitioned" + assert entries["tpch-nvme"]["dataset_dir"] == "tpch" + + def test_matrix_command() -> None: result = runner.invoke(cli_module.app, ["matrix", "develop"]) assert result.exit_code == 0 diff --git a/scripts/capture-file-sizes.py b/scripts/capture-file-sizes.py index 60bd87ebc06..6a72fcb10ca 100644 --- a/scripts/capture-file-sizes.py +++ b/scripts/capture-file-sizes.py @@ -19,6 +19,9 @@ def main() -> None: parser = argparse.ArgumentParser(description="Capture file sizes from benchmark data directories") parser.add_argument("data_dir", help="Data directory (e.g., vortex-bench/data)") parser.add_argument("--benchmark", required=True, help="Benchmark name (e.g., clickbench)") + parser.add_argument( + "--dataset-dir", required=True, help="The benchmark's directory under data_dir (e.g., clickbench_partitioned)" + ) parser.add_argument("--commit", required=True, help="Commit SHA") parser.add_argument("-o", "--output", required=True, help="Output JSONL file path") args = parser.parse_args() @@ -28,16 +31,9 @@ def main() -> None: print(f"Data directory not found: {data_dir}", file=sys.stderr) sys.exit(1) - # Find benchmark directories matching the name (handles flavors like clickbench_partitioned) - # Also handles exact match (e.g., tpch) - benchmark_dirs = [ - d - for d in data_dir.iterdir() - if d.is_dir() and (d.name == args.benchmark or d.name.startswith(f"{args.benchmark}_")) - ] - - if not benchmark_dirs: - print(f"No benchmark directories found matching: {args.benchmark}", file=sys.stderr) + benchmark_dir = data_dir / args.dataset_dir + if not benchmark_dir.is_dir(): + print(f"Benchmark directory not found: {benchmark_dir}", file=sys.stderr) sys.exit(1) # Formats to capture (vortex formats only, not parquet/duckdb) @@ -49,42 +45,41 @@ def main() -> None: # Walk subdirectories looking for format directories # Handle both direct format dirs (clickbench_partitioned/vortex-file-compressed/) # and scale factor subdirs (tpch/1.0/vortex-file-compressed/) - for benchmark_dir in benchmark_dirs: - for format_dir in benchmark_dir.rglob("*"): - if not format_dir.is_dir(): - continue - - format_name = format_dir.name - if format_name not in formats_to_capture: + for format_dir in benchmark_dir.rglob("*"): + if not format_dir.is_dir(): + continue + + format_name = format_dir.name + if format_name not in formats_to_capture: + continue + + # Extract scale factor from path (e.g., "1.0" for tpch/1.0/vortex-file-compressed) + # Default to "1.0" if no intermediate directory (e.g., clickbench) + path_between = format_dir.relative_to(benchmark_dir).parent + scale_factor = str(path_between) if str(path_between) != "." else "1.0" + + # Capture all files in this format directory + for file_path in format_dir.rglob("*"): + if not file_path.is_file(): continue - # Extract scale factor from path (e.g., "1.0" for tpch/1.0/vortex-file-compressed) - # Default to "1.0" if no intermediate directory (e.g., clickbench) - path_between = format_dir.relative_to(benchmark_dir).parent - scale_factor = str(path_between) if str(path_between) != "." else "1.0" - - # Capture all files in this format directory - for file_path in format_dir.rglob("*"): - if not file_path.is_file(): - continue - - size_bytes = file_path.stat().st_size - relative_path = file_path.relative_to(format_dir) - - records.append( - { - "metric": "file_size", - "unit": "bytes", - "value": size_bytes, - "commit_id": args.commit, - "file_size": { - "benchmark": args.benchmark, - "scale_factor": scale_factor, - "format": format_name, - "file": str(relative_path), - }, - } - ) + size_bytes = file_path.stat().st_size + relative_path = file_path.relative_to(format_dir) + + records.append( + { + "metric": "file_size", + "unit": "bytes", + "value": size_bytes, + "commit_id": args.commit, + "file_size": { + "benchmark": args.benchmark, + "scale_factor": scale_factor, + "format": format_name, + "file": str(relative_path), + }, + } + ) # Sort for deterministic output records.sort( diff --git a/scripts/tests/test_benchmark_reporting.py b/scripts/tests/test_benchmark_reporting.py index 55e4a673b2d..87a8f1370e9 100644 --- a/scripts/tests/test_benchmark_reporting.py +++ b/scripts/tests/test_benchmark_reporting.py @@ -701,6 +701,8 @@ def test_capture_file_sizes_emits_shared_benchmark_rows(tmp_path: Path) -> None: str(data_dir), "--benchmark", "tpch", + "--dataset-dir", + "tpch", "--commit", "deadbeef", "-o", @@ -787,3 +789,32 @@ def test_summary_reports_hot_alone_when_no_row_recorded_its_runs(tmp_path: Path) assert "**Vortex (geomean)**: hot 0.500x ✅" in summary assert "**Parquet (geomean)**: hot 1.000x ➖" in summary assert "cold" not in summary + + +def test_capture_file_sizes_reads_only_the_named_directory(tmp_path: Path) -> None: + data_dir = tmp_path / "data" + for name, size in (("clickbench_partitioned", 42), ("clickbench_partitioned_10", 7)): + format_dir = data_dir / name / "vortex-file-compressed" + format_dir.mkdir(parents=True) + (format_dir / "hits_0.vortex").write_bytes(b"x" * size) + output_path = tmp_path / "sizes.jsonl" + + subprocess.run( + [ + sys.executable, + str(CAPTURE_SCRIPT), + str(data_dir), + "--benchmark", + "clickbench", + "--dataset-dir", + "clickbench_partitioned", + "--commit", + "deadbeef", + "-o", + str(output_path), + ], + check=True, + ) + + records = [json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines()] + assert [record["value"] for record in records] == [42]