From eac735c7fa37c5e4c2808fec79e2141afdd60097 Mon Sep 17 00:00:00 2001
From: Mikhail Kot <mikhail@spiraldb.com>
Date: Wed, 12 Aug 2026 17:01:34 +0100
Subject: [PATCH] CI: run each benchmark query in a separate process

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>
---
 .agents/skills/bench-performance/SKILL.md     |  29 ++--
 bench-orchestrator/bench_orchestrator/cli.py  | 157 ++++++++++++++----
 .../bench_orchestrator/runner/executor.py     |  29 +++-
 bench-orchestrator/tests/test_cli.py          |  13 +-
 bench-orchestrator/tests/test_executor.py     |  20 ++-
 benchmarks/datafusion-bench/src/main.rs       |  18 +-
 benchmarks/duckdb-bench/src/main.rs           |  18 +-
 7 files changed, 224 insertions(+), 60 deletions(-)

diff --git a/.agents/skills/bench-performance/SKILL.md b/.agents/skills/bench-performance/SKILL.md
index c42a1271ddb..38af905d65b 100644
--- a/.agents/skills/bench-performance/SKILL.md
+++ b/.agents/skills/bench-performance/SKILL.md
@@ -110,10 +110,13 @@ assuming a flag exists.
 
 Source: `benchmarks/datafusion-bench/src/main.rs`.
 
+Runs each selected query exactly once per process.
+To repeat a query manually, loop the command.
+
 Supported diagnostics:
 
 - `--formats parquet,vortex,vortex-compact,lance,arrow`;
-- `--queries 6`, `--exclude-queries 1,2`, `--iterations N`, `--display-format gh-json`;
+- `--queries 6`, `--exclude-queries 1,2`, `--print-queries`, `--display-format gh-json`;
 - `--hide-progress-bar`, `-o /private/tmp/out.jsonl`, `--ingest-jsonl /private/tmp/out.ingest.jsonl`;
 - `--verbose`, `--tracing`, `--track-memory`, `--runner NAME`, `--opt key=value`;
 - `--explain` prints query plans instead of timing;
@@ -127,7 +130,7 @@ Examples:
 ```bash
 FEATURE_TOGGLE=1 RUST_LOG=vortex_datafusion=debug,vortex_layout=debug,datafusion=warn \
   target/release_debug/datafusion-bench tpch \
-  --display-format gh-json --iterations 5 --hide-progress-bar \
+  --display-format gh-json --hide-progress-bar \
   --formats <format> --queries <query> --show-metrics \
   -o /private/tmp/<label>.jsonl
 ```
@@ -141,14 +144,16 @@ FEATURE_TOGGLE=1 target/release_debug/datafusion-bench tpch \
 
 Source: `benchmarks/duckdb-bench/src/main.rs`.
 
+Runs each selected query exactly once per process.
+To repeat a query manually, loop the command.
+
 Supported diagnostics:
 
 - `--formats parquet,vortex,vortex-compact,duckdb`;
 - `--delete-duckdb-database` rebuilds the per-format DuckDB database;
 - `--threads N` sets DuckDB's `threads` config;
-- `--reuse` keeps one DuckDB connection across iterations, useful with Samply to keep work on the
-  same threads;
-- common flags: `--queries`, `--exclude-queries`, `--iterations`, `--display-format`,
+- `--reuse` keeps one DuckDB connection open, useful with Samply to keep work on the same threads;
+- common flags: `--queries`, `--exclude-queries`, `--print-queries`, `--display-format`,
   `--hide-progress-bar`, `-o`, `--ingest-jsonl`, `--track-memory`, `--verbose`, `--tracing`,
   `--runner`, `--opt`, `--explain`.
 
@@ -157,7 +162,7 @@ Example:
 ```bash
 RUST_LOG=duckdb_bench=trace,vortex_duckdb=debug,info \
   target/release_debug/duckdb-bench tpch \
-  --display-format gh-json --iterations 5 --hide-progress-bar \
+  --display-format gh-json --hide-progress-bar \
   --formats <baseline-format>,<candidate-format> --queries <query> --threads 8 --reuse \
   -o /private/tmp/<label>.jsonl
 ```
@@ -180,7 +185,7 @@ Example:
 ```bash
 RUST_LOG=lance_bench=debug,datafusion=warn \
   target/release_debug/lance-bench tpch \
-  --display-format gh-json --iterations 5 --hide-progress-bar \
+  --display-format gh-json --hide-progress-bar \
   --queries <query> \
   -o /private/tmp/<label>.jsonl
 ```
@@ -234,7 +239,7 @@ When investigating stream scheduling, enable the relevant flow trace and summari
 ```bash
 <FLOW_TRACE_ENV>=1 RUST_LOG=<flow-target>=debug,datafusion=warn \
   target/<profile-dir>/datafusion-bench clickbench \
-  --display-format gh-json --iterations 1 --hide-progress-bar \
+  --display-format gh-json --hide-progress-bar \
   --formats vortex --queries <query> \
   -o /private/tmp/<label>.jsonl > /private/tmp/<label>.log 2>&1
 
@@ -286,13 +291,13 @@ For Vortex/DataFusion scan I/O, prefer `--show-metrics` before OS tracing:
 
 ```bash
 target/release_debug/datafusion-bench <benchmark> \
-  --display-format gh-json --iterations 1 --hide-progress-bar \
+  --display-format gh-json --hide-progress-bar \
   --formats <format> --queries <query> --show-metrics \
   -o /private/tmp/<baseline-label>.jsonl \
   > /private/tmp/<baseline-label>.metrics.txt 2>&1
 
 FEATURE_TOGGLE=1 target/release_debug/datafusion-bench <benchmark> \
-  --display-format gh-json --iterations 1 --hide-progress-bar \
+  --display-format gh-json --hide-progress-bar \
   --formats <format> --queries <query> --show-metrics \
   -o /private/tmp/<candidate-label>.jsonl \
   > /private/tmp/<candidate-label>.metrics.txt 2>&1
@@ -350,7 +355,7 @@ profile contains only the target engine:
 FEATURE_TOGGLE=1 samply record --save-only --unstable-presymbolicate --rate 1000 \
   --output /private/tmp/<label>.profile.json.gz \
   -- target/release_debug/datafusion-bench <benchmark> \
-    --display-format gh-json --iterations 500 --hide-progress-bar \
+    --display-format gh-json --hide-progress-bar \
     --formats <format> --queries <query>
 ```
 
@@ -365,7 +370,7 @@ DuckDB profiling usually needs `--reuse`:
 samply record --save-only --unstable-presymbolicate --rate 1000 \
   --output /private/tmp/<label>.profile.json.gz \
   -- target/release_debug/duckdb-bench <benchmark> \
-    --display-format gh-json --iterations 500 --hide-progress-bar \
+    --display-format gh-json --hide-progress-bar \
     --formats <format> --queries <query> --reuse
 ```
 
diff --git a/bench-orchestrator/bench_orchestrator/cli.py b/bench-orchestrator/bench_orchestrator/cli.py
index c3daeec6363..92c1d0ee222 100644
--- a/bench-orchestrator/bench_orchestrator/cli.py
+++ b/bench-orchestrator/bench_orchestrator/cli.py
@@ -189,6 +189,69 @@ def write_result_line(line: str, store_writer, compatibility_file) -> None:
         compatibility_file.flush()
 
 
+def _median_ns(values: list[int]) -> int:
+    ordered = sorted(values)
+    count = len(ordered)
+    mid = count // 2
+    if count % 2 == 1:
+        return ordered[mid]
+    return (ordered[mid - 1] + ordered[mid]) // 2
+
+
+def merge_gh_json_runs(runs: list[list[str]]) -> list[str]:
+    """Collapse per-process gh-json timing for one query/format into a single record"""
+    timing: dict | None = None
+    others: list[str] = []
+    for lines in runs:
+        for line in lines:
+            line = line.strip()
+            if not line.startswith("{"):
+                continue
+            record = json.loads(line)
+            if record.get("all_runtimes") is None:
+                others.append(line)
+            elif timing is None:
+                timing = record
+            else:
+                timing["all_runtimes"].extend(record["all_runtimes"])
+
+    if timing is None:
+        return others
+    timing["value"] = _median_ns(timing["all_runtimes"])
+    return [json.dumps(timing), *others]
+
+
+def merge_ingest_records(paths: list[Path]) -> list[dict]:
+    """Collapse the per-process ingest records for one query/format into a single record."""
+    timing: dict | None = None
+    others: list[dict] = []
+    for path in paths:
+        with path.open(encoding="utf-8") as handle:
+            for line in handle:
+                line = line.strip()
+                if not line:
+                    continue
+                record = json.loads(line)
+                if record.get("all_runtimes_ns") is None:
+                    others.append(record)
+                elif timing is None:
+                    timing = record
+                else:
+                    timing["all_runtimes_ns"].extend(record["all_runtimes_ns"])
+
+    if timing is None:
+        return others
+    timing["value_ns"] = _median_ns(timing["all_runtimes_ns"])
+    return [timing, *others]
+
+
+def write_ingest_records(path: Path, records: list[dict]) -> None:
+    """Write merged ingest records as JSONL."""
+    with path.open("w", encoding="utf-8") as handle:
+        for record in records:
+            handle.write(json.dumps(record) + "\n")
+
+
 @app.command("prepare-data")
 def prepare_data(
     benchmark: Annotated[Benchmark, typer.Argument(help="Benchmark suite to prepare data for")],
@@ -379,37 +442,73 @@ def run(
             run_idx = 0
             for backend, backend_targets in backend_groups.items():
                 executor = BenchmarkExecutor(binary_paths[backend], backend, verbose=verbose)
+                per_query = backend in {Engine.DUCKDB, Engine.DATAFUSION}
                 for target in backend_targets:
-                    part_ingest_output = backend_ingest_output_path(ingest_temp_dir, run_idx, backend)
-                    run_idx += 1
-
-                    drop_os_caches()
-
                     try:
-                        results = executor.run(
-                            benchmark=benchmark,
-                            formats=[target.format],
-                            queries=query_list,
-                            exclude_queries=exclude_list,
-                            iterations=iterations,
-                            options=bench_opts,
-                            track_memory=track_memory,
-                            samply=samply,
-                            sample_rate=sample_rate,
-                            tracing=tracing,
-                            runner=runner,
-                            ingest_output=part_ingest_output,
-                            on_result=lambda line, store_writer=ctx.write_raw_json, compatibility=compatibility_file: (
-                                write_result_line(
-                                    line,
-                                    store_writer,
-                                    compatibility,
-                                )
-                            ),
-                        )
-                        if part_ingest_output is not None:
-                            ingest_output_parts.append(part_ingest_output)
-                        console.print(f"[green]{target}: {len(results)} results[/green]")
+                        if per_query:
+                            query_ids = executor.list_queries(benchmark, query_list, exclude_list)
+                            for query_id in query_ids:
+                                drop_os_caches()
+
+                                gh_runs: list[list[str]] = []
+                                ingest_parts: list[Path] = []
+                                for _ in range(iterations):
+                                    part = backend_ingest_output_path(ingest_temp_dir, run_idx, backend)
+                                    run_idx += 1
+                                    gh_runs.append(
+                                        executor.run(
+                                            benchmark=benchmark,
+                                            formats=[target.format],
+                                            queries=[query_id],
+                                            iterations=1,
+                                            options=bench_opts,
+                                            track_memory=track_memory,
+                                            samply=samply,
+                                            sample_rate=sample_rate,
+                                            tracing=tracing,
+                                            runner=runner,
+                                            ingest_output=part,
+                                        )
+                                    )
+                                    if part is not None:
+                                        ingest_parts.append(part)
+
+                                for line in merge_gh_json_runs(gh_runs):
+                                    write_result_line(line, ctx.write_raw_json, compatibility_file)
+
+                                if ingest_temp_dir is not None and ingest_parts:
+                                    merged_part = ingest_temp_dir / f"merged-{run_idx:04d}-{backend.value}.jsonl"
+                                    write_ingest_records(merged_part, merge_ingest_records(ingest_parts))
+                                    ingest_output_parts.append(merged_part)
+
+                            console.print(f"[green]{target}: {len(query_ids)} queries[/green]")
+                        else:
+                            part_ingest_output = backend_ingest_output_path(ingest_temp_dir, run_idx, backend)
+                            run_idx += 1
+
+                            drop_os_caches()
+
+                            def stream(line: str) -> None:
+                                write_result_line(line, ctx.write_raw_json, compatibility_file)
+
+                            results = executor.run(
+                                benchmark=benchmark,
+                                formats=[target.format],
+                                queries=query_list,
+                                exclude_queries=exclude_list,
+                                iterations=iterations,
+                                options=bench_opts,
+                                track_memory=track_memory,
+                                samply=samply,
+                                sample_rate=sample_rate,
+                                tracing=tracing,
+                                runner=runner,
+                                ingest_output=part_ingest_output,
+                                on_result=stream,
+                            )
+                            if part_ingest_output is not None:
+                                ingest_output_parts.append(part_ingest_output)
+                            console.print(f"[green]{target}: {len(results)} results[/green]")
                     except RuntimeError as exc:
                         ctx.metadata.partial = True
                         if strict_failures:
diff --git a/bench-orchestrator/bench_orchestrator/runner/executor.py b/bench-orchestrator/bench_orchestrator/runner/executor.py
index c5ccaf650ed..3b7940cda90 100644
--- a/bench-orchestrator/bench_orchestrator/runner/executor.py
+++ b/bench-orchestrator/bench_orchestrator/runner/executor.py
@@ -48,13 +48,14 @@ def build_command(
             benchmark.value,
             "--display-format",
             "gh-json",
-            "--iterations",
-            str(iterations),
             "--hide-progress-bar",
         ]
 
         if self.backend in {Engine.DATAFUSION, Engine.DUCKDB}:
             cmd.extend(["--formats", ",".join(fmt.value for fmt in formats)])
+        else:
+            # datafusion-bench and duckdb-bench run single query per process
+            cmd.extend(["--iterations", str(iterations)])
         if self.backend == Engine.DUCKDB:
             cmd.append("--delete-duckdb-database")
 
@@ -88,6 +89,30 @@ def build_command(
 
         return cmd
 
+    def list_queries(
+        self,
+        benchmark: Benchmark,
+        queries: list[int] | None = None,
+        exclude_queries: list[int] | None = None,
+    ) -> list[int]:
+        """Return query indices this benchmark selects"""
+        cmd = [str(self.binary_path), benchmark.value, "--print-queries"]
+        if queries:
+            cmd.extend(["--queries", ",".join(map(str, queries))])
+        if exclude_queries:
+            cmd.extend(["--exclude-queries", ",".join(map(str, exclude_queries))])
+
+        if self.verbose:
+            console.print(f"[dim]$ {' '.join(cmd)}[/dim]")
+
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(
+                f"Failed to list queries for {self.backend.value} {benchmark.value}: {result.stderr.strip()}"
+            )
+
+        return [int(line) for line in result.stdout.split() if line.strip()]
+
     def run(
         self,
         benchmark: Benchmark,
diff --git a/bench-orchestrator/tests/test_cli.py b/bench-orchestrator/tests/test_cli.py
index f8aee5e4f16..3be66fd753a 100644
--- a/bench-orchestrator/tests/test_cli.py
+++ b/bench-orchestrator/tests/test_cli.py
@@ -74,9 +74,9 @@ def test_run_writes_compatibility_results_output(tmp_path, monkeypatch) -> None:
     monkeypatch.setattr(cli_module, "ResultStore", lambda: run_store)
     monkeypatch.setattr(cli_module.BenchmarkBuilder, "get_binary_path", lambda self, backend: binary_path)
     monkeypatch.setattr(cli_module, "drop_os_caches", lambda: None)
+    monkeypatch.setattr(BenchmarkExecutor, "list_queries", lambda self, *args, **kwargs: [1])
 
     def fake_run(self, **kwargs):
-        kwargs["on_result"](sample_line)
         return [sample_line]
 
     monkeypatch.setattr(BenchmarkExecutor, "run", fake_run)
@@ -89,6 +89,8 @@ def fake_run(self, **kwargs):
             "--targets-json",
             '[{"engine":"datafusion","format":"parquet"}]',
             "--no-build",
+            "--iterations",
+            "1",
             "--output",
             str(output_path),
         ],
@@ -121,6 +123,7 @@ def test_run_combines_ingest_output_per_backend(tmp_path, monkeypatch) -> None:
     monkeypatch.setattr(cli_module, "ResultStore", lambda: run_store)
     monkeypatch.setattr(cli_module.BenchmarkBuilder, "get_binary_path", lambda self, backend: binary_paths[backend])
     monkeypatch.setattr(cli_module, "drop_os_caches", lambda: None)
+    monkeypatch.setattr(BenchmarkExecutor, "list_queries", lambda self, *args, **kwargs: [1])
 
     seen_backend_paths = []
 
@@ -128,7 +131,8 @@ def fake_run(self, **kwargs):
         backend_output = kwargs["ingest_output"]
         assert backend_output is not None
         assert backend_output != output_path
-        backend_output.write_text(f"{self.backend.value}-ingest\n", encoding="utf-8")
+        record = {"kind": "query_measurement", "engine": self.backend.value, "value_ns": 10, "all_runtimes_ns": [10]}
+        backend_output.write_text(json.dumps(record) + "\n", encoding="utf-8")
         seen_backend_paths.append(backend_output)
         return []
 
@@ -142,12 +146,15 @@ def fake_run(self, **kwargs):
             "--targets-json",
             '[{"engine":"datafusion","format":"parquet"},{"engine":"duckdb","format":"parquet"}]',
             "--no-build",
+            "--iterations",
+            "1",
             "--ingest-jsonl",
             str(output_path),
         ],
     )
 
     assert result.exit_code == 0
-    assert output_path.read_text(encoding="utf-8") == "datafusion-ingest\nduckdb-ingest\n"
+    engines = [json.loads(line)["engine"] for line in output_path.read_text(encoding="utf-8").splitlines()]
+    assert engines == ["datafusion", "duckdb"]
     assert len(seen_backend_paths) == 2
     assert seen_backend_paths[0] != seen_backend_paths[1]
diff --git a/bench-orchestrator/tests/test_executor.py b/bench-orchestrator/tests/test_executor.py
index 33f0f20b132..90579c77c89 100644
--- a/bench-orchestrator/tests/test_executor.py
+++ b/bench-orchestrator/tests/test_executor.py
@@ -19,13 +19,13 @@ def test_build_command_adds_duckdb_cleanup_flag() -> None:
         options={"scale-factor": "1.0"},
     )
 
-    assert cmd[:5] == [
+    assert cmd[:4] == [
         "/tmp/duckdb-bench",
         "tpch",
         "--display-format",
         "gh-json",
-        "--iterations",
     ]
+    assert "--iterations" not in cmd
     assert "--formats" in cmd
     assert "parquet,vortex" in cmd
     assert "--delete-duckdb-database" in cmd
@@ -33,6 +33,22 @@ def test_build_command_adds_duckdb_cleanup_flag() -> None:
     assert "scale-factor=1.0" in cmd
 
 
+def test_build_command_omits_iterations_for_datafusion_backend() -> None:
+    executor = BenchmarkExecutor(Path("/tmp/datafusion-bench"), Engine.DATAFUSION)
+
+    cmd = executor.build_command(benchmark=Benchmark.TPCH, formats=[Format.PARQUET])
+
+    assert "--iterations" not in cmd
+
+
+def test_build_command_keeps_iterations_for_lance_backend() -> None:
+    executor = BenchmarkExecutor(Path("/tmp/lance-bench"), Engine.LANCE)
+
+    cmd = executor.build_command(benchmark=Benchmark.TPCH, formats=[Format.LANCE], iterations=7)
+
+    assert cmd[cmd.index("--iterations") + 1] == "7"
+
+
 def test_build_command_serializes_vortex_spatial_native_format() -> None:
     executor = BenchmarkExecutor(Path("/tmp/duckdb-bench"), Engine.DUCKDB)
 
diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs
index 49a9bf92f58..1ac205795fb 100644
--- a/benchmarks/datafusion-bench/src/main.rs
+++ b/benchmarks/datafusion-bench/src/main.rs
@@ -58,9 +58,6 @@ struct Args {
     #[arg(value_enum)]
     benchmark: BenchmarkArg,
 
-    #[arg(short, long, default_value_t = 5)]
-    iterations: usize,
-
     #[arg(short, long)]
     threads: Option<usize>,
 
@@ -110,6 +107,10 @@ struct Args {
     #[arg(long, default_value_t = false)]
     explain: bool,
 
+    /// Print the selected query indices, one per line, and exit
+    #[arg(long, default_value_t = false)]
+    print_queries: bool,
+
     #[arg(long, value_delimiter = ',', value_parser = value_parser!(Format))]
     formats: Vec<Format>,
 
@@ -133,6 +134,13 @@ async fn main() -> anyhow::Result<()> {
         args.exclude_queries.as_ref(),
     );
 
+    if args.print_queries {
+        for (query_idx, _) in &filtered_queries {
+            println!("{query_idx}");
+        }
+        return Ok(());
+    }
+
     require_prepared_data(&*benchmark, &args.formats)?;
 
     let benchmark_name = benchmark.dataset().to_string();
@@ -156,9 +164,7 @@ async fn main() -> anyhow::Result<()> {
     let mode = if args.explain {
         BenchmarkMode::Explain
     } else {
-        BenchmarkMode::Run {
-            iterations: args.iterations,
-        }
+        BenchmarkMode::Run { iterations: 1 }
     };
 
     runner
diff --git a/benchmarks/duckdb-bench/src/main.rs b/benchmarks/duckdb-bench/src/main.rs
index c887ce621ba..8cf11a671a0 100644
--- a/benchmarks/duckdb-bench/src/main.rs
+++ b/benchmarks/duckdb-bench/src/main.rs
@@ -38,9 +38,6 @@ struct Args {
     #[arg(value_enum)]
     benchmark: BenchmarkArg,
 
-    #[arg(short, long, default_value_t = 5)]
-    iterations: usize,
-
     #[arg(short, long)]
     threads: Option<usize>,
 
@@ -88,6 +85,10 @@ struct Args {
     #[arg(long, default_value_t = false)]
     explain: bool,
 
+    /// Print the selected query indices, one per line, and exit
+    #[arg(long, default_value_t = false)]
+    print_queries: bool,
+
     #[arg(
         long,
         default_value_t = false,
@@ -111,6 +112,13 @@ fn main() -> anyhow::Result<()> {
         args.exclude_queries.as_ref(),
     );
 
+    if args.print_queries {
+        for (query_idx, _) in &filtered_queries {
+            println!("{query_idx}");
+        }
+        return Ok(());
+    }
+
     if args.formats.is_empty() {
         anyhow::bail!("provide a format with --formats");
     }
@@ -173,9 +181,7 @@ fn main() -> anyhow::Result<()> {
     let mode = if args.explain {
         BenchmarkMode::Explain
     } else {
-        BenchmarkMode::Run {
-            iterations: args.iterations,
-        }
+        BenchmarkMode::Run { iterations: 1 }
     };
 
     runner.run_all(