From 1c7a4be48053c64d6be1a39cc2f217b9427610a8 Mon Sep 17 00:00:00 2001 From: Bruce Ritchie Date: Sun, 19 Jul 2026 08:28:33 -0400 Subject: [PATCH] Add support for running benchmarks with command line arguments, --help and --dry-run --- Cargo.lock | 52 +- benchmarks/Cargo.toml | 3 +- benchmarks/benches/sql.rs | 6 +- benchmarks/sql_benchmarks/README.md | 119 +- .../clickbench/clickbench.suite | 25 + .../clickbench_extended.suite | 21 + .../clickbench_sorted/clickbench_sorted.suite | 28 + benchmarks/sql_benchmarks/h2o/h2o.suite | 33 + benchmarks/sql_benchmarks/hj/hj.suite | 21 + benchmarks/sql_benchmarks/imdb/imdb.suite | 26 + benchmarks/sql_benchmarks/nlj/nlj.suite | 11 + .../predicate_eval/predicate_eval.suite | 33 +- .../push_down_topk/push_down_topk.suite | 21 + benchmarks/sql_benchmarks/smj/smj.suite | 11 + .../sql_benchmarks/sort_tpch/sort_tpch.suite | 28 + benchmarks/sql_benchmarks/tpcds/tpcds.suite | 21 + benchmarks/sql_benchmarks/tpch/tpch.suite | 47 + .../wide_schema/wide_schema.suite | 14 + benchmarks/src/bin/benchmark_runner.rs | 1458 +++++++++++++++-- benchmarks/src/lib.rs | 1 + benchmarks/src/sql_benchmark_runner.rs | 225 ++- benchmarks/src/sql_benchmark_suite.rs | 849 ++++++++++ 22 files changed, 2874 insertions(+), 179 deletions(-) create mode 100644 benchmarks/sql_benchmarks/clickbench/clickbench.suite create mode 100644 benchmarks/sql_benchmarks/clickbench_extended/clickbench_extended.suite create mode 100644 benchmarks/sql_benchmarks/clickbench_sorted/clickbench_sorted.suite create mode 100644 benchmarks/sql_benchmarks/h2o/h2o.suite create mode 100644 benchmarks/sql_benchmarks/hj/hj.suite create mode 100644 benchmarks/sql_benchmarks/imdb/imdb.suite create mode 100644 benchmarks/sql_benchmarks/nlj/nlj.suite create mode 100644 benchmarks/sql_benchmarks/push_down_topk/push_down_topk.suite create mode 100644 benchmarks/sql_benchmarks/smj/smj.suite create mode 100644 benchmarks/sql_benchmarks/sort_tpch/sort_tpch.suite create mode 100644 benchmarks/sql_benchmarks/tpcds/tpcds.suite create mode 100644 benchmarks/sql_benchmarks/tpch/tpch.suite create mode 100644 benchmarks/sql_benchmarks/wide_schema/wide_schema.suite create mode 100644 benchmarks/src/sql_benchmark_suite.rs diff --git a/Cargo.lock b/Cargo.lock index 6f776e5e965f9..074995061b79d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1797,6 +1797,7 @@ dependencies = [ "tempfile", "tokio", "tokio-util", + "toml", ] [[package]] @@ -5610,6 +5611,15 @@ dependencies = [ "syn", ] +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + [[package]] name = "serde_tokenstream" version = "0.2.3" @@ -6326,6 +6336,30 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "0.9.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" +dependencies = [ + "indexmap 2.14.0", + "serde_core", + "serde_spanned", + "toml_datetime 0.7.5+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 0.7.15", +] + +[[package]] +name = "toml_datetime" +version = "0.7.5+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +dependencies = [ + "serde_core", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -6342,9 +6376,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" dependencies = [ "indexmap 2.14.0", - "toml_datetime", + "toml_datetime 1.1.1+spec-1.1.0", "toml_parser", - "winnow", + "winnow 1.0.2", ] [[package]] @@ -6353,9 +6387,15 @@ version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" dependencies = [ - "winnow", + "winnow 1.0.2", ] +[[package]] +name = "toml_writer" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" + [[package]] name = "tonic" version = "0.14.6" @@ -7273,6 +7313,12 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" + [[package]] name = "winnow" version = "1.0.2" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 5dae70761f9a7..a8dcd704efeda 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -43,7 +43,7 @@ mimalloc_extended = ["libmimalloc-sys/extended"] arrow = { workspace = true } async-trait = "0.1" bytes = { workspace = true } -clap = { version = "4.6.0", features = ["derive", "env"] } +clap = { version = "4.6.0", features = ["derive", "env", "string"] } criterion = { workspace = true, features = ["html_reports"] } datafusion = { workspace = true, default-features = true } datafusion-common = { workspace = true, default-features = true } @@ -62,6 +62,7 @@ serde_json = { workspace = true } snmalloc-rs = { version = "0.7", optional = true } tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] } tokio-util = { version = "0.7.17" } +toml = "0.9.8" [dev-dependencies] datafusion-proto = { workspace = true, features = ["parquet"] } diff --git a/benchmarks/benches/sql.rs b/benchmarks/benches/sql.rs index 83351b8205ddc..9240a19470db9 100644 --- a/benchmarks/benches/sql.rs +++ b/benchmarks/benches/sql.rs @@ -24,8 +24,8 @@ use clap::Parser; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_benchmarks::sql_benchmark_runner::{ - BenchmarkFilter, SqlRunConfig, default_sql_benchmark_directory, - run_criterion_benchmarks_impl, + BenchmarkFilter, SqlRunConfig, default_criterion_replacements, + default_sql_benchmark_directory, run_criterion_benchmarks_impl, }; use datafusion_benchmarks::util::CommonOpt; use datafusion_common::instant::Instant; @@ -84,6 +84,8 @@ pub fn sql(c: &mut Criterion) { subgroup: args.subgroup, query: args.query, }, + replacements: default_criterion_replacements(), + query_filename: None, persist_results: args.persist_results, validate_results: args.validate, output: None, diff --git a/benchmarks/sql_benchmarks/README.md b/benchmarks/sql_benchmarks/README.md index f92baf6e73bbf..dfb09e0a3a4a2 100644 --- a/benchmarks/sql_benchmarks/README.md +++ b/benchmarks/sql_benchmarks/README.md @@ -43,24 +43,96 @@ in the community: | `tpcds` | TPC‑DS queries | | `tpch` | TPC‑H queries | | `wide_schema` | Small-projection queries on a wide (1024-col, 256-file) synthetic dataset; runs `wide` + `narrow` subgroups for comparison | -| `predicate_eval` | Conjunctive (AND) filter-evaluation micro-benchmarks; each subgroup is a different predicate pattern, to test how an adaptive predicate-ordering system behaves across them ([#11262](https://github.com/apache/datafusion/issues/11262)). Subgroups (`BENCH_SUBGROUP`): `costsel`, `cost`, `selectivity`, `cardinality`, `width`, `scale`, `neutral`, `correlation`, `drift`. Toggle a system under test with its native `DATAFUSION_*` env var | +| `predicate_eval` | Conjunctive (AND) filter-evaluation micro-benchmarks; each subgroup is a different predicate pattern, to test how an adaptive predicate-ordering system behaves across them ([#11262](https://github.com/apache/datafusion/issues/11262)). Subgroups (`--subgroup`): `costsel`, `cost`, `selectivity`, `cardinality`, `width`, `scale`, `neutral`, `correlation`, `drift`. Configure the system under test through its DataFusion settings. | # Running Benchmarks -The easiest way to run a benchmark is to use the `bench.sh` shell script (up one level from this document) -as it takes care of configuring any required environment variables and can populate any required data files. -However, it is possible to directly run a sql benchmark using the `cargo bench` command. For example: +Use `benchmark_runner` to run SQL benchmarks. It reads each suite's `.suite` +file and exposes the suite's configuration as command-line options. Use the +`bench.sh` shell script one level above this directory to download or generate +required data files. ```shell -BENCH_NAME=tpch cargo bench --bench sql +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch ``` +## SQL benchmark runner + +The `benchmark_runner` binary discovers suites from this directory and exposes +suite-specific options alongside the common benchmark options. The suite name +must come before all options. + +```bash +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- --list +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --help +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --query 15 --format csv +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- clickbench --partitioning partitioned --dry-run +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --query 1 --result-mode persist +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --query 1 --result-mode validate +``` + +Use `--path PATH` or `-p PATH` to override `DATA_DIR` for a suite that declares +that path replacement. Suites without a `DATA_DIR` replacement reject the +option. Suite-specific values follow this precedence: command-line option, +environment variable, then the default in the suite metadata. + +`--dry-run` prints the resolved suite, filters, run mode, common options, +suite-specific values, and path replacements as JSON. It reports source metadata +for suite options and path replacements. It validates the command but does not +load benchmark definitions, create a session, read datasets, execute SQL, or +write benchmark results. + +Use `--result-mode persist` to save query results or `--result-mode validate` to +compare them with saved results. The default, `--result-mode none`, does neither. +For compatibility with direct Criterion runs, the runner also reads +`BENCH_PERSIST_RESULTS` and `BENCH_VALIDATE`. Persistence takes precedence when +both variables are `true`. An explicit `--result-mode` overrides both variables. + +### Suite metadata + +Each discoverable suite has one TOML metadata file named +`/.suite`. The runner accepts these top-level fields: + +| Field | Required | Description | +|-------|----------|-------------| +| `description` | Yes | Non-empty text shown by `--list` and suite help. | +| `query_pattern` | No | Relative benchmark filename pattern. It must contain exactly one `{QUERY_ID}` or `{QUERY_ID_PADDED}` placeholder and defaults to `q{QUERY_ID_PADDED}.benchmark`. | +| `path_replacements` | No | Map of replacement names to paths. Relative paths resolve from the suite directory. `DATA_DIR` enables `--path/-p`. | +| `options` | No | Array of suite-specific option tables described below. | +| `examples` | No | Array of `command` and `description` pairs appended to suite help. Both values must contain text. | + +Each `[[options]]` table has these fields: + +| Field | Required | Description | +|-------|----------|-------------| +| `name` | Yes | Long option name without `--`; use lowercase ASCII letters, digits, and hyphens. | +| `short` | No | One ASCII letter or digit without `-`. | +| `env` | Yes | Environment variable that supplies the option value. | +| `default` | Yes | Value used when neither the command line nor the environment supplies one. | +| `values` | No | Accepted values. Omit the field to allow any value. Include `"..."` to allow the listed values plus any other value. Without `"..."`, the list is closed. | +| `help` | Yes | Non-empty text shown in suite help. | + +Option names, short names, and environment keys must be unique within a suite. +An option environment key cannot also appear in `path_replacements`. Suite +options cannot reuse the runner's global names: `help` (`-h`), `query` (`-q`), +`subgroup`, `iterations` (`-i`), `partitions` (`-n`), `batch-size` (`-s`), +`mem-pool-type`, `memory-limit`, `sort-spill-reservation-bytes`, `debug` (`-d`), +`simulate-latency`, `criterion`, `list`, `output` (`-o`), `save-baseline`, +`path` (`-p`), `result-mode`, or `dry-run`. + # Benchmark configuration -Sql benchmarks are configured via environment variables. Cargo's bench command and -[criterion](https://github.com/criterion-rs/criterion.rs) (the underlying benchmark framework) have an unfortunate -limitation in that custom command arguments cannot be passed into a benchmark. The alternative is to use environment -variables to pass in arguments which is what is used here. +`benchmark_runner` is the preferred interface for configuring and running SQL +benchmarks. Run ` --help` to see the common and suite-specific options: + +```shell +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- h2o --help +``` + +The runner maps suite options to the environment variables below for +compatibility with benchmark files and direct Criterion runs. Direct +`cargo bench --bench sql` invocations cannot accept custom arguments, so they +still use environment variables. The SQL benchmarking tool uses the following environment variables: @@ -76,10 +148,10 @@ The SQL benchmarking tool uses the following environment variables: | MEM_POOL_TYPE | The memory pool type to use, should be one of "fair" or "greedy". | | MEMORY_LIMIT | Memory limit (e.g. '100M', '1.5G'). If not specified, run all pre-defined memory limits for given query if there's any, otherwise run with no memory limit. | -Example – Run the H2O window benchmarks on the 'small' sized CSV data files: +Example: run the H2O window benchmarks on the small CSV data files: -``` bash -BENCH_NAME=h2o BENCH_SUBGROUP=window H2O_BENCH_SIZE=small H20_FILE_TYPE=csv cargo bench --bench sql +```shell +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- h2o --subgroup window --size small --format csv ``` Some benchmarks use custom environment variables as outlined below: @@ -101,21 +173,22 @@ Some benchmarks use custom environment variables as outlined below: ## How it works -SQL benchmarks are run via cargo's bench command using [criterion](https://docs.rs/criterion/latest/criterion/) -for running and gathering statistics of each sql being benchmarked. +The runner executes SQL benchmarks with its basic runner by default. Pass +`--criterion` to gather statistics with +[Criterion](https://docs.rs/criterion/latest/criterion/). Each individual benchmark is represented by a `.benchmark` file that contains a number of directives instructing the tool on how to load data, run initializations, run assertions, run the benchmark, optionally persist and validate results, and finally run any cleanup if required. -Variables are supported in two forms: +Benchmark files support replacement variables in two forms: -* string substitution based on environment variables (with default values if unset): \${ENV_VAR} and +* string substitution with an optional default: \${ENV_VAR} and \${ENV_VAR:-default}. -* if / else based on whether an environment variable is true or not +* if / else based on whether a replacement value is true or not (\${ENV_VAR:-default|true value|false value}). In this form only the value `true` (case-insensitive) selects the - true branch; any other set value selects the false branch. If ENV_VAR is unset, the valud of `default` is used to -* select the branch. + true branch; any other supplied value selects the false branch. If the value is absent, the parser uses `default` to + select the branch. Comments in files are supported with lines starting with # or --. @@ -157,8 +230,8 @@ The above showcases the use of defaults for variables: `${NAME:-default}` The name of the benchmark. This will be used as part of the display name used by criterion.

Example:
name Q${QUERY_NUMBER_PADDED}
-The `name` directive also makes the value available to benchmark-file replacements as `BENCH_NAME`. This is separate -from the `BENCH_NAME` environment variable used to select which benchmark group to run. +The `name` directive also makes the value available to benchmark-file replacements as `BENCH_NAME`. This value is +separate from the suite name passed to `benchmark_runner`. @@ -221,8 +294,8 @@ The run directive called during execution of the benchmark. If a path to a file the run directive that path will be parsed and any sql statements in that file will be executed during the benchmark run. If no path is specified the next line is required to be the sql statement to execute.

Multiple statements are allowed within a single run directive, however a benchmark file may contain only one run directive. When -running with `BENCH_PERSIST_RESULTS` or `BENCH_VALIDATE`, only the last `SELECT` or `WITH` statement from that run -directive will be used for comparison.

The run directive (including any following sql statement) must be +when persisting or validating results, only the last `SELECT` or `WITH` statement from that run directive will be used +for comparison.

The run directive (including any following sql statement) must be followed by a blank line.

Example:
run sql_benchmarks/imdb/queries/${QUERY_NUMBER_PADDED}.sql
diff --git a/benchmarks/sql_benchmarks/clickbench/clickbench.suite b/benchmarks/sql_benchmarks/clickbench/clickbench.suite new file mode 100644 index 0000000000000..74d8a5cc2ae56 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/clickbench.suite @@ -0,0 +1,25 @@ +description = "ClickBench analytics queries over the hits dataset" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "partitioning" +env = "CLICKBENCH_TYPE" +default = "single" +values = ["single", "partitioned"] +help = "Selects the single-file or partitioned ClickBench dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench" +description = "Run all ClickBench queries against the single-file dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench --query 7" +description = "Run ClickBench query 7." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench --partitioning partitioned" +description = "Run all ClickBench queries against the partitioned dataset." diff --git a/benchmarks/sql_benchmarks/clickbench_extended/clickbench_extended.suite b/benchmarks/sql_benchmarks/clickbench_extended/clickbench_extended.suite new file mode 100644 index 0000000000000..dfca00b4a03db --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/clickbench_extended.suite @@ -0,0 +1,21 @@ +description = "Extended ClickBench queries over the hits dataset" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "partitioning" +env = "CLICKBENCH_TYPE" +default = "single" +values = ["single", "partitioned"] +help = "Selects the single-file or partitioned ClickBench dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_extended" +description = "Run all extended ClickBench queries against the single-file dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_extended --query 4 --partitioning partitioned" +description = "Run extended ClickBench query 4 against the partitioned dataset." diff --git a/benchmarks/sql_benchmarks/clickbench_sorted/clickbench_sorted.suite b/benchmarks/sql_benchmarks/clickbench_sorted/clickbench_sorted.suite new file mode 100644 index 0000000000000..5c8a0909e3f55 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_sorted/clickbench_sorted.suite @@ -0,0 +1,28 @@ +description = "ClickBench query over a pre-sorted hits dataset" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "sort-column" +env = "SORTED_BY" +default = "EventTime" +values = ["EventTime", "..."] +help = "Selects the column used to sort the ClickBench data." + +[[options]] +name = "sort-order" +env = "SORTED_ORDER" +default = "ASC" +values = ["ASC", "DESC"] +help = "Selects the sort direction for the ClickBench data." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_sorted" +description = "Run the sorted ClickBench query ordered by EventTime ascending." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_sorted --sort-column UserID --sort-order DESC" +description = "Run the query over data sorted by UserID descending." diff --git a/benchmarks/sql_benchmarks/h2o/h2o.suite b/benchmarks/sql_benchmarks/h2o/h2o.suite new file mode 100644 index 0000000000000..27d83285ba026 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/h2o.suite @@ -0,0 +1,33 @@ +description = "H2O group-by, join, and window SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "size" +env = "H2O_BENCH_SIZE" +default = "small" +values = ["small", "medium", "big"] +help = "Selects the H2O dataset size." + +[[options]] +name = "format" +short = "f" +env = "H2O_FILE_TYPE" +default = "csv" +values = ["csv", "parquet"] +help = "Selects the H2O data format." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o" +description = "Run all H2O queries with the small CSV datasets." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o --query 3 --subgroup window" +description = "Run H2O window query 3." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o --subgroup join --size medium -f parquet" +description = "Run the H2O join queries with the medium Parquet dataset." diff --git a/benchmarks/sql_benchmarks/hj/hj.suite b/benchmarks/sql_benchmarks/hj/hj.suite new file mode 100644 index 0000000000000..31ff12a046c59 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/hj.suite @@ -0,0 +1,21 @@ +description = "Hash join SQL benchmarks derived from TPC-H" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "..."] +help = "Selects the TPC-H scale factor used by the hash join benchmarks." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- hj" +description = "Run all hash join queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- hj --query 16 --scale-factor 10" +description = "Run hash join query 16 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/imdb/imdb.suite b/benchmarks/sql_benchmarks/imdb/imdb.suite new file mode 100644 index 0000000000000..7422b06bbc345 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/imdb.suite @@ -0,0 +1,26 @@ +description = "Join Order Benchmark queries over the IMDb dataset" + +query_pattern = "{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "format" +short = "f" +env = "IMDB_FILE_TYPE" +default = "parquet" +values = ["parquet", "csv"] +help = "Selects the IMDb data format." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- imdb" +description = "Run all IMDb queries against Parquet data." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- imdb --query 01a" +description = "Run IMDb query 01a." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- imdb --query 01a -f csv" +description = "Run IMDb query 01a against CSV data." diff --git a/benchmarks/sql_benchmarks/nlj/nlj.suite b/benchmarks/sql_benchmarks/nlj/nlj.suite new file mode 100644 index 0000000000000..21b4cb298cd8e --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/nlj.suite @@ -0,0 +1,11 @@ +description = "Nested-loop join SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- nlj" +description = "Run all nested-loop join queries." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- nlj --query 7" +description = "Run nested-loop join query 7." diff --git a/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.suite b/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.suite index aba11e06ff166..af1a326cd8c51 100644 --- a/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.suite +++ b/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.suite @@ -1,2 +1,31 @@ -name = "predicate_eval" -description = "Micro-benchmarks for conjunctive (AND) filter evaluation. Each subgroup exercises a different predicate pattern (per-predicate cost, selectivity, conjunct count, string-column width, row count, correlation, selectivity drift, plus an order-neutral control) so the suite can show how an adaptive predicate-ordering system behaves across them -- the kind of change these benchmarks are meant to help drive, e.g. https://github.com/apache/datafusion/issues/11262. By default it measures DataFusion's built-in left-deep AND short-circuit and sets no engine config of its own; toggle a system under test with its native DATAFUSION_* env var (the harness reads SessionConfig::from_env), e.g. DATAFUSION_EXECUTION_ADAPTIVE_FILTER_REORDERING=true. Subgroups (BENCH_SUBGROUP): costsel, cost, selectivity, cardinality, width, scale, neutral, correlation, drift. Size synthetic data with PRED_ROWS and string-column width with PRED_FILL." +description = "Conjunctive filter evaluation micro-benchmarks covering predicate cost, selectivity, cardinality, width, scale, correlation, and drift" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "rows" +short = "r" +env = "PRED_ROWS" +default = "1000000" +values = ["1000000", "..."] +help = "Sets the number of rows in generated predicate-evaluation datasets." + +[[options]] +name = "fill" +short = "f" +env = "PRED_FILL" +default = "30" +values = ["2", "30", "170", "..."] +help = "Sets the filler width for generated string columns." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- predicate_eval" +description = "Run all predicate-evaluation subgroups with default data sizes." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- predicate_eval --query 20 --subgroup selectivity" +description = "Run selectivity query 20." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- predicate_eval --subgroup width -r 500000 -f 170" +description = "Run the width subgroup with 500,000 extra-wide rows." diff --git a/benchmarks/sql_benchmarks/push_down_topk/push_down_topk.suite b/benchmarks/sql_benchmarks/push_down_topk/push_down_topk.suite new file mode 100644 index 0000000000000..a70139c7669ca --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/push_down_topk.suite @@ -0,0 +1,21 @@ +description = "TopK pushdown benchmarks for ORDER BY LIMIT over TPC-H joins" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "..."] +help = "Selects the TPC-H scale factor used by the TopK benchmarks." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- push_down_topk" +description = "Run all TopK pushdown queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- push_down_topk --query 3 --scale-factor 10" +description = "Run TopK pushdown query 3 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/smj/smj.suite b/benchmarks/sql_benchmarks/smj/smj.suite new file mode 100644 index 0000000000000..44db22ffe20f5 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/smj.suite @@ -0,0 +1,11 @@ +description = "Sort-merge join SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- smj" +description = "Run all sort-merge join queries." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- smj --query 12" +description = "Run sort-merge join query 12." diff --git a/benchmarks/sql_benchmarks/sort_tpch/sort_tpch.suite b/benchmarks/sql_benchmarks/sort_tpch/sort_tpch.suite new file mode 100644 index 0000000000000..38ee9c132b284 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/sort_tpch.suite @@ -0,0 +1,28 @@ +description = "Sorting benchmarks over the TPC-H lineitem table" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "..."] +help = "Selects the TPC-H scale factor." + +[[options]] +name = "sorted" +env = "BENCH_SORTED" +default = "false" +values = ["false", "true"] +help = "Controls whether the lineitem table is loaded in l_orderkey order." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- sort_tpch" +description = "Run all TPC-H sorting queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- sort_tpch --query 4 --sorted true" +description = "Run sorting query 4 over pre-sorted lineitem data." diff --git a/benchmarks/sql_benchmarks/tpcds/tpcds.suite b/benchmarks/sql_benchmarks/tpcds/tpcds.suite new file mode 100644 index 0000000000000..7261c3d4dfc6d --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/tpcds.suite @@ -0,0 +1,21 @@ +description = "TPC-DS SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "100", "..."] +help = "Selects the TPC-DS scale factor." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpcds" +description = "Run all TPC-DS queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpcds --query 42 --scale-factor 10" +description = "Run TPC-DS query 42 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/tpch/tpch.suite b/benchmarks/sql_benchmarks/tpch/tpch.suite new file mode 100644 index 0000000000000..0330cc0f32584 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpch/tpch.suite @@ -0,0 +1,47 @@ +description = "TPC-H SQL benchmarks" + +# Query patterns control how numeric QUERY_ID values map to .benchmark files +# during discovery and command resolution. Use exactly one query-id token: +# - {QUERY_ID_PADDED}: two-digit ids, such as q01.benchmark +# - {QUERY_ID}: unpadded ids, such as query-1.benchmark +# If omitted, this defaults to q{QUERY_ID_PADDED}.benchmark. +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +# Path replacements define path-like variables used while parsing benchmark +# files. Relative paths are resolved from this suite file's directory and then +# passed to SqlBenchmark's replacement mapping, so benchmark SQL can refer to +# values such as ${DATA_DIR}. For timed runs, the runner's --path/-p option +# overrides DATA_DIR. +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "format" +short = "f" +env = "TPCH_FILE_TYPE" +default = "parquet" +values = ["parquet", "csv", "mem"] +help = "Selects the TPC-H data format." + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "..."] +help = "Selects the TPC-H scale factor." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch" +description = "Run all TPC-H queries with the default parquet SF1 configuration." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch --query 15" +description = "Run TPC-H query 15 with the default parquet SF1 configuration." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch --query 15 -f csv" +description = "Run TPC-H query 15 against CSV data." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch --query 15 --scale-factor 10" +description = "Run TPC-H query 15 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/wide_schema/wide_schema.suite b/benchmarks/sql_benchmarks/wide_schema/wide_schema.suite new file mode 100644 index 0000000000000..275f15e102677 --- /dev/null +++ b/benchmarks/sql_benchmarks/wide_schema/wide_schema.suite @@ -0,0 +1,14 @@ +description = "Projection benchmarks over synthetic wide and narrow schemas" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- wide_schema" +description = "Run all wide-schema projection queries." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- wide_schema --query 2 --subgroup narrow" +description = "Run query 2 with the narrow schema." diff --git a/benchmarks/src/bin/benchmark_runner.rs b/benchmarks/src/bin/benchmark_runner.rs index 71b3b1b9e0a87..3a208b0aa8f93 100644 --- a/benchmarks/src/bin/benchmark_runner.rs +++ b/benchmarks/src/bin/benchmark_runner.rs @@ -17,22 +17,29 @@ //! DataFusion SQL benchmark runner. -use clap::{ArgAction, ArgMatches, CommandFactory, FromArgMatches, Parser}; +use clap::{ + Arg, ArgAction, ArgMatches, Command, CommandFactory, FromArgMatches, Parser, + ValueEnum, +}; use criterion::Criterion; use datafusion::error::Result; use datafusion::prelude::SessionContext; use datafusion_benchmarks::sql_benchmark::SqlBenchmark; use datafusion_benchmarks::sql_benchmark_runner::{ BenchmarkFilter, SqlRunConfig, default_sql_benchmark_directory, ensure_selection, - filter_benchmarks, finish_benchmark, format_benchmark_list, - load_benchmark_definitions, make_ctx, prepare_benchmark, - run_criterion_benchmarks_impl, sort_benchmarks, + filter_benchmarks, finish_benchmark, load_benchmark_definitions_for_query, make_ctx, + prepare_benchmark, run_criterion_benchmarks_impl, +}; +use datafusion_benchmarks::sql_benchmark_suite::{ + ReservedOptions, SuiteExample, SuiteMetadata, discover_suites, }; use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt, print_memory_stats}; use datafusion_common::instant::Instant; use datafusion_common::{DataFusionError, exec_datafusion_err}; use datafusion_common_runtime::SpawnedTask; -use std::collections::BTreeMap; +use serde::{Serialize, Serializer}; +use std::collections::{BTreeMap, BTreeSet}; +use std::ffi::{OsStr, OsString}; use std::io::IsTerminal; use std::path::Path; @@ -63,6 +70,67 @@ enum CliAction { config: SqlRunConfig, save_baseline: Option, }, + DryRun(DryRunOutput), +} + +#[derive(Debug, Serialize)] +struct ResolvedSuiteValue { + value: String, + #[serde(serialize_with = "serialize_value_source")] + source: datafusion_benchmarks::sql_benchmark_suite::ValueSource, + environment: String, +} + +#[derive(Debug, Serialize)] +struct ResolvedPathValue { + value: String, + #[serde(serialize_with = "serialize_value_source")] + source: datafusion_benchmarks::sql_benchmark_suite::ValueSource, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +enum RunMode { + Simple, + Criterion, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, ValueEnum)] +#[serde(rename_all = "snake_case")] +enum ResultMode { + #[default] + None, + Persist, + Validate, +} + +impl ResultMode { + fn config_flags(self) -> (bool, bool) { + match self { + Self::None => (false, false), + Self::Persist => (true, false), + Self::Validate => (false, true), + } + } +} + +#[derive(Debug, Serialize)] +struct DryRunCommonOptions { + iterations: usize, + partitions: Option, + batch_size: Option, +} + +#[derive(Debug, Serialize)] +struct DryRunOutput { + suite: String, + query: Option, + subgroup: Option, + mode: RunMode, + result_mode: ResultMode, + common_options: DryRunCommonOptions, + suite_options: BTreeMap, + path_replacements: BTreeMap, } #[derive(Debug, Parser)] @@ -111,13 +179,26 @@ struct Cli { help = "Save Criterion measurements to the named baseline" )] save_baseline: Option, + + #[arg(short = 'p', long = "path", value_name = "PATH")] + path: Option, + + #[arg( + long = "result-mode", + value_enum, + value_name = "MODE", + help = "Handle expected results: none, persist, or validate" + )] + result_mode: Option, + + #[arg(long = "dry-run", action = ArgAction::SetTrue)] + dry_run: bool, } /// Parses CLI arguments, runs the selected action, and prints any output. async fn run_cli() -> Result<()> { - let matches = Cli::command().get_matches(); - let action = cli_action_from_matches(&matches)?; - let output = run_cli_action(action, &default_sql_benchmark_directory()).await?; + let benchmark_dir = default_sql_benchmark_directory(); + let output = run_cli_from(std::env::args_os(), &benchmark_dir).await?; if !output.is_empty() { println!("{output}"); @@ -126,15 +207,220 @@ async fn run_cli() -> Result<()> { Ok(()) } +fn serialize_value_source( + source: &datafusion_benchmarks::sql_benchmark_suite::ValueSource, + serializer: S, +) -> std::result::Result +where + S: Serializer, +{ + let value = match source { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine => { + "command_line" + } + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Environment => { + "environment" + } + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default => "default", + }; + serializer.serialize_str(value) +} + +fn clap_display_output(error: &DataFusionError) -> Option { + let DataFusionError::External(error) = error else { + return None; + }; + let error = error.downcast_ref::()?; + matches!( + error.kind(), + clap::error::ErrorKind::DisplayHelp | clap::error::ErrorKind::DisplayVersion + ) + .then(|| error.to_string()) +} + +async fn run_cli_from(args: I, benchmark_dir: &Path) -> Result +where + I: IntoIterator, + T: Into + Clone, +{ + match parse_cli_from(args, benchmark_dir) { + Ok(action) => run_cli_action(action, benchmark_dir).await, + Err(error) => clap_display_output(&error).ok_or(error), + } +} + +fn format_examples(examples: &[SuiteExample]) -> String { + if examples.is_empty() { + return String::new(); + } + + let mut output = String::from("Examples:\n"); + for example in examples { + output.push_str(" "); + output.push_str(example.command()); + output.push_str("\n "); + output.push_str(example.description()); + output.push('\n'); + } + output +} + +fn build_cli(suite: Option<&SuiteMetadata>) -> Command { + let mut command = Cli::command(); + + if let Some(suite) = suite { + command = command.about(suite.description().to_string()); + + for option in suite.options() { + let mut arg = Arg::new(option.name().to_string()) + .long(option.name().to_string()) + .help(option.help().to_string()) + .env(option.env().to_string()) + .default_value(option.default().to_string()); + + if let Some(short) = option.short() { + arg = arg.short(short); + } + if let Some(values) = option + .values() + .filter(|values| !values.iter().any(|value| value == "...")) + { + arg = arg.value_parser(values.to_vec()); + } + + command = command.arg(arg); + } + + let examples = format_examples(suite.examples()); + + if !examples.is_empty() { + command = command.after_help(examples); + } + } + + command +} + +fn reserved_options() -> (BTreeSet, BTreeSet) { + let command = Cli::command(); + let long = command + .get_arguments() + .filter_map(|arg| arg.get_long().map(ToOwned::to_owned)) + .collect(); + let short = command + .get_arguments() + .filter_map(|arg| arg.get_short()) + .collect(); + (long, short) +} + +fn suite_metadata(benchmark_dir: &Path) -> Result> { + let (long, short) = reserved_options(); + discover_suites( + benchmark_dir, + &ReservedOptions { + long: &long, + short: &short, + }, + ) +} + +fn format_suite_list(suites: &[SuiteMetadata]) -> String { + let mut output = String::from("SQL benchmarks:\n"); + for suite in suites { + let query_word = if suite.benchmark_count() == 1 { + "query " + } else { + "queries " + }; + output.push_str(&format!( + " {:<24} {} {query_word}{}\n", + suite.name(), + suite.benchmark_count(), + suite.description() + )); + } + output.trim_end().to_string() +} + +fn locate_suite_arg(args: &[OsString]) -> Result> { + let Some(argument) = args.get(1) else { + return Ok(None); + }; + if argument == OsStr::new("--help") + || argument == OsStr::new("-h") + || argument == OsStr::new("--list") + || argument == OsStr::new("--dry-run") + { + return Ok(None); + } + let suite = argument.to_str().ok_or_else(|| { + DataFusionError::External("suite name is not valid Unicode".into()) + })?; + if suite.starts_with('-') { + return Err(exec_datafusion_err!( + "suite must be the first argument; options must follow the suite" + )); + } + Ok(Some(suite)) +} + +fn try_parse_cli_from(args: I, benchmark_dir: &Path) -> Result +where + I: IntoIterator, + T: Into + Clone, +{ + let args = args.into_iter().map(Into::into).collect::>(); + let suite = locate_suite_arg(&args)?; + let (long, short) = reserved_options(); + let reserved = ReservedOptions { + long: &long, + short: &short, + }; + let suite = suite + .map(|name| { + if !benchmark_dir.join(name).is_dir() { + let available = discover_suites(benchmark_dir, &reserved)?; + return Err(exec_datafusion_err!( + "unknown benchmark '{name}'\n\n{}", + format_suite_list(&available) + )); + } + SuiteMetadata::load(benchmark_dir, name, &reserved) + }) + .transpose()?; + let matches = build_cli(suite.as_ref()) + .try_get_matches_from(args) + .map_err(|error| DataFusionError::External(Box::new(error)))?; + + cli_action_from_matches(&matches, suite.as_ref()) +} + +fn parse_cli_from(args: I, benchmark_dir: &Path) -> Result +where + I: IntoIterator, + T: Into + Clone, +{ + try_parse_cli_from(args, benchmark_dir) +} + /// Converts parsed arguments into an executable action and validates mode options. -fn cli_action_from_matches(matches: &ArgMatches) -> Result { +fn cli_action_from_matches( + matches: &ArgMatches, + suite: Option<&SuiteMetadata>, +) -> Result { let cli = Cli::from_arg_matches(matches) .map_err(|e| DataFusionError::External(Box::new(e)))?; + if cli.dry_run && cli.list { + return Err(exec_datafusion_err!("--list cannot be used with --dry-run")); + } + if cli.dry_run && cli.benchmark.is_none() { + return Err(exec_datafusion_err!("--dry-run requires a benchmark suite")); + } if cli.list || cli.benchmark.is_none() { return Ok(CliAction::List); } - if cli.criterion && cli.output.is_some() { return Err(exec_datafusion_err!( "--output cannot be used with --criterion" @@ -145,6 +431,9 @@ fn cli_action_from_matches(matches: &ArgMatches) -> Result { "--save-baseline cannot be used without --criterion" )); } + if !cli.criterion && cli.common.iterations == 0 { + return Err(exec_datafusion_err!("iterations must be greater than zero")); + } // we need to know if iterations was set on the command line, not the default value let iterations_from_cli = matches.value_source("iterations") @@ -155,21 +444,136 @@ fn cli_action_from_matches(matches: &ArgMatches) -> Result { "--iterations cannot be used with --criterion" )); } - if !cli.criterion && cli.common.iterations == 0 { - return Err(exec_datafusion_err!("iterations must be greater than zero")); + + let suite = + suite.ok_or_else(|| exec_datafusion_err!("benchmark suite is required"))?; + + if cli.path.is_some() && !suite.path_replacements().contains_key("DATA_DIR") { + return Err(exec_datafusion_err!( + "--path cannot be used because suite '{}' does not declare DATA_DIR", + suite.name() + )); } - let config = SqlRunConfig { + let result_mode = resolve_result_mode(cli.result_mode)?; + let (persist_results, validate_results) = result_mode.config_flags(); + let mut config = SqlRunConfig { common: cli.common, filter: BenchmarkFilter { name: cli.benchmark, subgroup: cli.subgroup, query: cli.query, }, - persist_results: false, - validate_results: false, + replacements: Default::default(), + query_filename: None, + persist_results, + validate_results, output: cli.output, }; + let suite_options: BTreeMap = suite + .options() + .iter() + .map(|option| { + let value = matches + .get_one::(option.name()) + .expect("suite options always have defaults") + .clone(); + let source = match matches.value_source(option.name()) { + Some(clap::parser::ValueSource::CommandLine) => { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine + } + Some(clap::parser::ValueSource::EnvVariable) => { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Environment + } + Some(clap::parser::ValueSource::DefaultValue) => { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default + } + vs => unreachable!("unexpected suite option source: {vs:?}"), + }; + ( + option.name().to_string(), + ResolvedSuiteValue { + value, + source, + environment: option.env().to_string(), + }, + ) + }) + .collect(); + let path_replacements: BTreeMap = suite + .path_replacements() + .iter() + .map(|(key, default)| { + let (value, source) = if key == "DATA_DIR" { + cli.path.as_ref().map_or_else( + || { + ( + default.display().to_string(), + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default, + ) + }, + |path| { + ( + path.display().to_string(), + datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine, + ) + }, + ) + } else { + ( + default.display().to_string(), + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default, + ) + }; + ( + key.to_ascii_lowercase(), + ResolvedPathValue { value, source }, + ) + }) + .collect(); + + config.replacements = suite_options + .values() + .map(|resolved| { + ( + resolved.environment.to_ascii_lowercase(), + resolved.value.clone(), + ) + }) + .chain( + path_replacements + .iter() + .map(|(key, resolved)| (key.clone(), resolved.value.clone())), + ) + .collect(); + config.query_filename = config + .filter + .query + .as_deref() + .map(|query| suite.query_filename(query)) + .transpose()?; + + if cli.dry_run { + let mode = if cli.criterion { + RunMode::Criterion + } else { + RunMode::Simple + }; + return Ok(CliAction::DryRun(DryRunOutput { + suite: suite.name().to_string(), + query: config.filter.query.clone(), + subgroup: config.filter.subgroup.clone(), + mode, + result_mode, + common_options: DryRunCommonOptions { + iterations: config.common.iterations, + partitions: config.common.partitions, + batch_size: config.common.batch_size, + }, + suite_options, + path_replacements, + })); + } if cli.criterion { Ok(CliAction::Criterion { @@ -184,13 +588,7 @@ fn cli_action_from_matches(matches: &ArgMatches) -> Result { /// Executes a parsed CLI action and returns any text that should be printed. async fn run_cli_action(action: CliAction, benchmark_dir: &Path) -> Result { match action { - CliAction::List => { - let ctx = SessionContext::new(); - let benchmarks = - load_benchmarks(&BenchmarkFilter::default(), &ctx, benchmark_dir).await?; - - Ok(format_benchmark_list(&benchmarks)) - } + CliAction::List => Ok(format_suite_list(&suite_metadata(benchmark_dir)?)), CliAction::Simple(config) => { run_simple_benchmarks(benchmark_dir, config).await?; Ok(String::new()) @@ -218,21 +616,39 @@ async fn run_cli_action(action: CliAction, benchmark_dir: &Path) -> Result serde_json::to_string_pretty(&output) + .map_err(|error| DataFusionError::External(Box::new(error))), } } -/// Loads benchmark definitions, applies CLI-style filters, and sorts each group. -pub async fn load_benchmarks( - filter: &BenchmarkFilter, - ctx: &SessionContext, - benchmark_dir: &Path, -) -> Result>> { - let benches = load_benchmark_definitions(filter, ctx, benchmark_dir).await?; - let mut benches = filter_benchmarks(filter, benches); +fn resolve_result_mode(explicit: Option) -> Result { + if let Some(mode) = explicit { + return Ok(mode); + } - sort_benchmarks(&mut benches); + let persist = parse_compat_bool("BENCH_PERSIST_RESULTS")?; + let validate = parse_compat_bool("BENCH_VALIDATE")?; - Ok(benches) + Ok(if persist { + ResultMode::Persist + } else if validate { + ResultMode::Validate + } else { + ResultMode::None + }) +} + +fn parse_compat_bool(name: &str) -> Result { + let Some(value) = std::env::var_os(name) else { + return Ok(false); + }; + let value = value + .into_string() + .map_err(|_| exec_datafusion_err!("{name} contains invalid UTF-8"))?; + + value.parse::().map_err(|_| { + exec_datafusion_err!("invalid value '{value}' for {name}; expected true or false") + }) } /// Builds the default Criterion runner and optionally records a named baseline. @@ -265,8 +681,14 @@ pub async fn run_simple_benchmarks( } let listing_ctx = make_ctx(&config.common)?; - let all_benchmarks = - load_benchmark_definitions(&config.filter, &listing_ctx, benchmark_dir).await?; + let all_benchmarks = load_benchmark_definitions_for_query( + &config.filter, + &listing_ctx, + benchmark_dir, + &config.replacements, + config.query_filename.as_deref(), + ) + .await?; let selected = filter_benchmarks(&config.filter, all_benchmarks.clone()); let mut run = BenchmarkRun::new(); @@ -341,9 +763,97 @@ fn criterion_like_styles() -> clap::builder::Styles { #[cfg(test)] mod tests { use super::*; - use datafusion_benchmarks::sql_benchmark_runner::unknown_benchmark_error; + use datafusion_benchmarks::sql_benchmark_runner::{ + load_benchmark_definitions, sort_benchmarks, unknown_benchmark_error, + }; + use datafusion_benchmarks::sql_benchmark_suite::ValueSource; + use std::collections::HashMap; + use std::ffi::OsString; use std::fs; use std::path::{Path, PathBuf}; + use std::sync::{Mutex, MutexGuard}; + + static ENV_MUTEX: Mutex<()> = Mutex::new(()); + + struct ScopedEnv { + previous: Vec<(&'static str, Option)>, + _lock: MutexGuard<'static, ()>, + } + + impl ScopedEnv { + fn set(name: &'static str, value: impl Into) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let previous = std::env::var_os(name); + // SAFETY: ENV_MUTEX serializes changes made through ScopedEnv in this + // test module; it does not synchronize environment access elsewhere. + unsafe { std::env::set_var(name, value.into()) }; + Self { + previous: vec![(name, previous)], + _lock: lock, + } + } + + fn remove(name: &'static str) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let previous = std::env::var_os(name); + // SAFETY: ENV_MUTEX serializes changes made through ScopedEnv in this + // test module; it does not synchronize environment access elsewhere. + unsafe { std::env::remove_var(name) }; + Self { + previous: vec![(name, previous)], + _lock: lock, + } + } + + fn set_many(changes: [(&'static str, Option<&str>); N]) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let mut previous = Vec::with_capacity(N); + for (name, value) in changes { + previous.push((name, std::env::var_os(name))); + // SAFETY: this guard holds ENV_MUTEX until it restores all entries. + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + } + Self { + previous, + _lock: lock, + } + } + } + + impl Drop for ScopedEnv { + fn drop(&mut self) { + // SAFETY: this guard holds ENV_MUTEX until after all entries are restored. + unsafe { + for (name, previous) in self.previous.drain(..).rev() { + match previous { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + } + } + } + + /// Loads benchmark definitions, applies CLI-style filters, and sorts each group. + async fn load_benchmarks( + filter: &BenchmarkFilter, + ctx: &SessionContext, + benchmark_dir: &Path, + ) -> Result>> { + let benches = + load_benchmark_definitions(filter, ctx, benchmark_dir, &Default::default()) + .await?; + let mut benches = filter_benchmarks(filter, benches); + + sort_benchmarks(&mut benches); + + Ok(benches) + } fn write_benchmark(root: &Path, relative_path: &str, contents: &str) -> PathBuf { let path = root.join(relative_path); @@ -354,6 +864,13 @@ mod tests { path } + fn write_suite(root: &Path, name: &str, description: &str) -> PathBuf { + let path = root.join(name).join(format!("{name}.suite")); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(&path, format!("description = {description:?}\n")).unwrap(); + path + } + fn common(iterations: usize) -> CommonOpt { CommonOpt { iterations, @@ -367,88 +884,678 @@ mod tests { } } - fn parse_cli_from(args: I) -> Result + async fn run_cli_with_dir(args: I, benchmark_dir: &Path) -> Result where I: IntoIterator, - T: Into + Clone, + T: Into + Clone, { - let matches = Cli::command() - .try_get_matches_from(args) - .map_err(|e| DataFusionError::External(Box::new(e)))?; + run_cli_from(args, benchmark_dir).await + } - cli_action_from_matches(&matches) + fn suite_root() -> tempfile::TempDir { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + fs::write( + temp.path().join("alpha/alpha.suite"), + r#"description = "Alpha benchmark" + +[path_replacements] +DATA_DIR = "data" + +[[options]] +name = "format" +short = "f" +env = "ALPHA_FORMAT" +default = "parquet" +values = ["parquet", "csv"] +help = "Alpha input format" + +[[examples]] +command = "benchmark_runner alpha -q 1 -f csv" +description = "Run query one against CSV data." +"#, + ) + .unwrap(); + temp } - async fn run_cli_with_dir(args: I, benchmark_dir: &Path) -> Result - where - I: IntoIterator, - T: Into + Clone, - { - run_cli_action(parse_cli_from(args)?, benchmark_dir).await + #[test] + fn suite_help_contains_metadata() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let error = + try_parse_cli_from(["benchmark_runner", "alpha", "--help"], temp.path()) + .unwrap_err(); + let help = error.to_string(); + + assert!(help.contains("Alpha benchmark"), "{help}"); + assert!(help.contains("--format"), "{help}"); + assert!(help.contains("ALPHA_FORMAT"), "{help}"); + assert!( + help.contains("benchmark_runner alpha -q 1 -f csv"), + "{help}" + ); + } + + #[test] + fn accepts_interleaved_named_options() { + let temp = suite_root(); + let action = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--partitions", + "2", + "--format", + "csv", + "--query", + "5", + ], + temp.path(), + ) + .unwrap(); + let CliAction::Simple(config) = action else { + panic!("expected simple run") + }; + + assert_eq!(config.common.partitions, Some(2)); + assert_eq!(config.filter.query.as_deref(), Some("5")); + assert_eq!(config.query_filename.as_deref(), Some("q05.benchmark")); + assert_eq!(config.replacements["alpha_format"], "csv"); + assert_eq!( + config.replacements["data_dir"], + temp.path().join("alpha/data").display().to_string() + ); + } + + #[test] + fn dry_run_uses_suite_default() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let action = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "parquet"); + assert_eq!(output.suite_options["format"].source, ValueSource::Default); + } + + #[test] + fn result_mode_defaults_to_none() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let CliAction::DryRun(output) = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap() + else { + panic!("expected dry run"); + }; + assert_eq!(output.result_mode, ResultMode::None); + } + + #[tokio::test] + async fn result_mode_persist_writes_expected_results() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let result_path = temp.path().join("expected.csv"); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + &format!( + "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", + result_path.display() + ), + ); + + run_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "1", + "--result-mode", + "persist", + ], + temp.path(), + ) + .await + .unwrap(); + + let persisted = fs::read_to_string(result_path).unwrap(); + assert!(persisted.contains("value"), "{persisted}"); + assert!(persisted.contains('1'), "{persisted}"); + } + + #[tokio::test] + async fn result_mode_validate_accepts_expected_results() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let result_path = temp.path().join("expected.csv"); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + &format!( + "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", + result_path.display() + ), + ); + fs::write(&result_path, "value\n1\n").unwrap(); + + run_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "1", + "--result-mode", + "validate", + ], + temp.path(), + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn result_mode_validate_reports_mismatched_results() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let result_path = temp.path().join("expected.csv"); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + &format!( + "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", + result_path.display() + ), + ); + fs::write(&result_path, "value\n2\n").unwrap(); + + let error = run_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "1", + "--result-mode", + "validate", + ], + temp.path(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("expected value"), "{error}"); + } + + #[test] + fn explicit_result_modes_populate_config() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", Some("invalid")), + ("BENCH_VALIDATE", Some("invalid")), + ]); + let temp = suite_root(); + for (value, expected, persist, validate) in [ + ("none", ResultMode::None, false, false), + ("persist", ResultMode::Persist, true, false), + ("validate", ResultMode::Validate, false, true), + ] { + let CliAction::DryRun(output) = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--result-mode", + value, + "--dry-run", + ], + temp.path(), + ) + .unwrap() else { + panic!("expected dry run"); + }; + assert_eq!(output.result_mode, expected); + + let CliAction::Simple(config) = parse_cli_from( + ["benchmark_runner", "alpha", "--result-mode", value], + temp.path(), + ) + .unwrap() else { + panic!("expected simple run"); + }; + assert_eq!(config.persist_results, persist); + assert_eq!(config.validate_results, validate); + } + } + + #[test] + fn compatibility_environment_resolves_result_mode() { + for (persist, validate, expected) in [ + (Some("true"), None, ResultMode::Persist), + (None, Some("true"), ResultMode::Validate), + (Some("true"), Some("true"), ResultMode::Persist), + (Some("false"), Some("false"), ResultMode::None), + ] { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", persist), + ("BENCH_VALIDATE", validate), + ]); + let temp = suite_root(); + let CliAction::DryRun(output) = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap() + else { + panic!("expected dry run"); + }; + assert_eq!(output.result_mode, expected); + } + } + + #[test] + fn invalid_result_mode_environment_is_rejected_without_cli_override() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", Some("invalid")), + ("BENCH_VALIDATE", Some("false")), + ]); + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap_err(); + assert!( + error.to_string().contains("BENCH_PERSIST_RESULTS"), + "{error}" + ); + } + + #[test] + fn invalid_result_mode_cli_value_lists_allowed_values() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let error = parse_cli_from( + ["benchmark_runner", "alpha", "--result-mode", "invalid"], + temp.path(), + ) + .unwrap_err(); + let message = error.to_string(); + for allowed in ["none", "persist", "validate"] { + assert!(message.contains(allowed), "{message}"); + } + } + + #[test] + fn environment_beats_suite_default() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "csv"); + let temp = suite_root(); + let action = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "csv"); + assert_eq!( + output.suite_options["format"].source, + ValueSource::Environment + ); + } + + #[test] + fn cli_equals_syntax_beats_environment_and_default() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "parquet"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "--format=csv", "--dry-run"], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "csv"); + assert_eq!( + output.suite_options["format"].source, + ValueSource::CommandLine + ); + } + + #[test] + fn empty_environment_value_is_validated() { + let _env = ScopedEnv::set("ALPHA_FORMAT", ""); + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap_err(); + + let message = error.to_string(); + assert!(message.contains("a value is required"), "{message}"); + assert!(message.contains("parquet, csv"), "{message}"); + } + + #[cfg(unix)] + #[test] + fn non_unicode_environment_value_is_rejected_by_clap() { + use std::os::unix::ffi::OsStringExt; + + let _env = ScopedEnv::set("ALPHA_FORMAT", OsString::from_vec(vec![0xff])); + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap_err(); + + assert!(error.to_string().contains("invalid UTF-8"), "{error}"); + } + + #[test] + fn attached_short_cli_value_beats_invalid_environment() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "invalid"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "-fcsv", "--dry-run"], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "csv"); + assert_eq!( + output.suite_options["format"].source, + ValueSource::CommandLine + ); + } + + #[test] + fn invalid_suite_environment_does_not_block_help() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "invalid"); + let temp = suite_root(); + let error = + try_parse_cli_from(["benchmark_runner", "alpha", "--help"], temp.path()) + .unwrap_err(); + let help = error.to_string(); + + assert!(help.contains("Alpha benchmark"), "{help}"); + assert!(help.contains("--format"), "{help}"); + } + + #[test] + fn dry_run_rejects_list_instead_of_listing() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let error = parse_cli_from( + ["benchmark_runner", "alpha", "--list", "--dry-run"], + temp.path(), + ) + .unwrap_err(); + + assert!(error.to_string().contains("--list"), "{error}"); + assert!(error.to_string().contains("--dry-run"), "{error}"); + } + + #[test] + fn dry_run_requires_suite() { + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "--dry-run"], temp.path()).unwrap_err(); + + assert!(error.to_string().contains("--dry-run"), "{error}"); + assert!(error.to_string().contains("suite"), "{error}"); + } + + #[test] + fn dry_run_resolves_default_and_overridden_path() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let default_action = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap(); + let CliAction::DryRun(default_output) = default_action else { + panic!("expected dry run") + }; + assert_eq!( + default_output.path_replacements["data_dir"].source, + ValueSource::Default + ); + + let override_action = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--path", + "/tmp/alpha-data", + "--dry-run", + ], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(override_output) = override_action else { + panic!("expected dry run") + }; + assert_eq!( + override_output.path_replacements["data_dir"].value, + "/tmp/alpha-data" + ); + assert_eq!( + override_output.path_replacements["data_dir"].source, + ValueSource::CommandLine + ); + } + + #[test] + fn path_is_rejected_without_data_dir_replacement() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + fs::write( + temp.path().join("alpha/alpha.suite"), + fs::read_to_string(temp.path().join("alpha/alpha.suite")) + .unwrap() + .replace("[path_replacements]\nDATA_DIR = \"data\"\n\n", ""), + ) + .unwrap(); + + let error = + parse_cli_from(["benchmark_runner", "alpha", "--path", "data"], temp.path()) + .unwrap_err(); + assert!(error.to_string().contains("--path"), "{error}"); + assert!(error.to_string().contains("DATA_DIR"), "{error}"); + } + + #[test] + fn criterion_dry_run_reports_mode_and_keeps_cross_checks() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "--criterion", "--dry-run"], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + assert_eq!(output.mode, RunMode::Criterion); + + let error = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--iterations", + "2", + "--dry-run", + ], + temp.path(), + ) + .unwrap_err(); + assert!(error.to_string().contains("--iterations"), "{error}"); + } + + #[test] + fn dry_run_rejects_invalid_query_form() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let error = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "../secret", + "--dry-run", + ], + temp.path(), + ) + .unwrap_err(); + + assert!( + error.to_string().contains("invalid query identifier"), + "{error}" + ); + } + + #[test] + fn uppercase_q_dry_run_uses_same_query_filename_as_lowercase_q() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + + assert!(matches!( + parse_cli_from( + ["benchmark_runner", "alpha", "--query", "Q1", "--dry-run",], + temp.path(), + ) + .unwrap(), + CliAction::DryRun(_) + )); + + let filename = |query| { + let CliAction::Simple(config) = parse_cli_from( + ["benchmark_runner", "alpha", "--query", query], + temp.path(), + ) + .unwrap() else { + panic!("expected simple run") + }; + config.query_filename + }; + + assert_eq!(filename("Q1"), filename("q1")); + } + + #[tokio::test] + async fn dry_run_returns_deterministic_json_without_reading_benchmarks() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + fs::write( + temp.path().join("alpha/benchmarks/q01.benchmark"), + "this benchmark is intentionally invalid", + ) + .unwrap(); + + let output = run_cli_with_dir( + [ + "benchmark_runner", + "alpha", + "--query", + "7", + "--partitions", + "2", + "--dry-run", + ], + temp.path(), + ) + .await + .unwrap(); + let json: serde_json::Value = serde_json::from_str(&output).unwrap(); + + assert_eq!(json["suite"], "alpha"); + assert_eq!(json["query"], "7"); + assert_eq!(json["mode"], "simple"); + assert_eq!(json["common_options"]["partitions"], 2); + assert_eq!(json["suite_options"]["format"]["source"], "default"); } #[test] fn cli_lists_when_benchmark_is_omitted() { - let action = parse_cli_from(["benchmark_runner"]).unwrap(); + let temp = suite_root(); + let action = parse_cli_from(["benchmark_runner"], temp.path()).unwrap(); assert!(matches!(action, CliAction::List)); } #[test] fn cli_lists_with_explicit_list_flag() { - let action = parse_cli_from(["benchmark_runner", "--list"]).unwrap(); + let temp = suite_root(); + let action = parse_cli_from(["benchmark_runner", "--list"], temp.path()).unwrap(); assert!(matches!(action, CliAction::List)); } #[test] fn cli_defaults_to_basic_runner() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); let action = - parse_cli_from(["benchmark_runner", "tpch", "--query", "1"]).unwrap(); + parse_cli_from(["benchmark_runner", "alpha", "--query", "1"], temp.path()) + .unwrap(); let CliAction::Simple(config) = action else { panic!("expected basic runner"); }; - assert_eq!(config.filter.name.as_deref(), Some("tpch")); + assert_eq!(config.filter.name.as_deref(), Some("alpha")); assert_eq!(config.filter.query.as_deref(), Some("1")); } #[test] fn cli_reads_query_from_env() { - let previous = std::env::var_os("BENCH_QUERY"); - // SAFETY: This test restores BENCH_QUERY before returning and does not - // spawn threads while the environment variable is overridden. - unsafe { - std::env::set_var("BENCH_QUERY", "8"); - } - - let action = parse_cli_from(["benchmark_runner", "tpch"]); - - unsafe { - match previous { - Some(value) => std::env::set_var("BENCH_QUERY", value), - None => std::env::remove_var("BENCH_QUERY"), - } - } - + let _env = ScopedEnv::set("BENCH_QUERY", "8"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "--format", "parquet"], + temp.path(), + ); let action = action.unwrap(); let CliAction::Simple(config) = action else { panic!("expected basic runner"); }; - assert_eq!(config.filter.name.as_deref(), Some("tpch")); + assert_eq!(config.filter.name.as_deref(), Some("alpha")); assert_eq!(config.filter.query.as_deref(), Some("8")); } #[test] fn cli_accepts_criterion_runner() { - let action = parse_cli_from([ - "benchmark_runner", - "tpch", - "--criterion", - "--save-baseline", - "main", - ]) + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let action = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--save-baseline", + "main", + ], + temp.path(), + ) .unwrap(); let CliAction::Criterion { @@ -459,19 +1566,24 @@ mod tests { panic!("expected criterion runner"); }; - assert_eq!(config.filter.name.as_deref(), Some("tpch")); + assert_eq!(config.filter.name.as_deref(), Some("alpha")); assert_eq!(save_baseline.as_deref(), Some("main")); } #[test] fn cli_rejects_output_with_criterion() { - let err = parse_cli_from([ - "benchmark_runner", - "tpch", - "--criterion", - "--output", - "results.json", - ]) + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--output", + "results.json", + ], + temp.path(), + ) .unwrap_err(); assert!(err.to_string().contains("--output")); @@ -480,8 +1592,13 @@ mod tests { #[test] fn cli_rejects_save_baseline_without_criterion() { - let err = parse_cli_from(["benchmark_runner", "tpch", "--save-baseline", "main"]) - .unwrap_err(); + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + ["benchmark_runner", "alpha", "--save-baseline", "main"], + temp.path(), + ) + .unwrap_err(); assert!(err.to_string().contains("--save-baseline")); assert!(err.to_string().contains("--criterion")); @@ -489,13 +1606,18 @@ mod tests { #[test] fn cli_rejects_iterations_with_criterion() { - let err = parse_cli_from([ - "benchmark_runner", - "tpch", - "--criterion", - "--iterations", - "3", - ]) + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--iterations", + "3", + ], + temp.path(), + ) .unwrap_err(); assert!(err.to_string().contains("--iterations")); @@ -504,8 +1626,13 @@ mod tests { #[test] fn cli_rejects_zero_basic_iterations() { - let err = parse_cli_from(["benchmark_runner", "tpch", "--iterations", "0"]) - .unwrap_err(); + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + ["benchmark_runner", "alpha", "--iterations", "0"], + temp.path(), + ) + .unwrap_err(); assert!(err.to_string().contains("iterations")); } @@ -519,6 +1646,7 @@ mod tests { "alpha/benchmarks/q01.benchmark", "name Q01\n\nrun\nSELECT 1\n", ); + write_suite(temp.path(), "alpha", "Alpha workload"); let output = run_cli_with_dir(["benchmark_runner"], temp.path()) .await @@ -537,6 +1665,7 @@ mod tests { "alpha/benchmarks/q01.benchmark", "name Q01\n\nrun\nSELECT 1\n", ); + write_suite(temp.path(), "alpha", "Alpha workload"); let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) .await @@ -547,14 +1676,45 @@ mod tests { } #[tokio::test] - async fn run_cli_reports_unknown_benchmark_with_list() { - let temp = tempfile::tempdir().unwrap(); + async fn run_cli_top_level_help_is_successful_output() { + let temp = suite_root(); + let output = run_cli_with_dir(["benchmark_runner", "--help"], temp.path()) + .await + .unwrap(); - write_benchmark( + assert!(output.contains("Run DataFusion SQL benchmarks"), "{output}"); + assert!(output.contains("Usage:"), "{output}"); + } + + #[tokio::test] + async fn run_cli_suite_help_is_successful_output() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let output = + run_cli_with_dir(["benchmark_runner", "alpha", "--help"], temp.path()) + .await + .unwrap(); + + assert!(output.contains("Alpha benchmark"), "{output}"); + assert!(output.contains("--format"), "{output}"); + } + + #[tokio::test] + async fn run_cli_real_parse_error_remains_an_error() { + let temp = suite_root(); + let error = run_cli_with_dir( + ["benchmark_runner", "alpha", "--not-an-option"], temp.path(), - "alpha/benchmarks/q01.benchmark", - "name Q01\n\nrun\nSELECT 1\n", - ); + ) + .await + .unwrap_err(); + + assert!(error.to_string().contains("unexpected argument"), "{error}"); + } + + #[tokio::test] + async fn run_cli_reports_unknown_benchmark_with_list() { + let temp = suite_root(); let err = run_cli_with_dir(["benchmark_runner", "missing"], temp.path()) .await @@ -590,6 +1750,8 @@ mod tests { subgroup: None, query: Some("1".to_string()), }, + replacements: HashMap::new(), + query_filename: None, persist_results: false, validate_results: false, output: None, @@ -626,6 +1788,8 @@ mod tests { subgroup: None, query: Some("9".to_string()), }, + replacements: HashMap::new(), + query_filename: None, persist_results: false, validate_results: false, output: None, @@ -663,6 +1827,8 @@ mod tests { subgroup: Some("narrow".to_string()), query: None, }, + replacements: HashMap::new(), + query_filename: None, persist_results: false, validate_results: false, output: None, @@ -701,6 +1867,8 @@ mod tests { subgroup: None, query: Some("1".to_string()), }, + replacements: HashMap::new(), + query_filename: None, persist_results: false, validate_results: false, output: Some(output.clone()), @@ -732,6 +1900,8 @@ mod tests { subgroup: None, query: Some("1".to_string()), }, + replacements: HashMap::new(), + query_filename: None, persist_results: false, validate_results: false, output: None, @@ -819,13 +1989,8 @@ mod tests { } #[tokio::test] - async fn benchmark_replacements_default_data_dir_to_benchmarks_data() { + async fn benchmark_replacements_use_explicit_data_dir() { let temp = tempfile::tempdir().unwrap(); - let previous = std::env::var_os("DATA_DIR"); - - unsafe { - std::env::remove_var("DATA_DIR"); - } write_benchmark( temp.path(), @@ -834,22 +1999,19 @@ mod tests { ); let ctx = SessionContext::new(); - let benches = - load_benchmarks(&BenchmarkFilter::default(), &ctx, temp.path()).await; - - unsafe { - match previous { - Some(value) => std::env::set_var("DATA_DIR", value), - None => std::env::remove_var("DATA_DIR"), - } - } - - let benches = benches.unwrap(); - let expected = PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("data") .to_string_lossy() .into_owned(); + let replacements = HashMap::from([("data_dir".to_string(), expected.clone())]); + let benches = load_benchmark_definitions( + &BenchmarkFilter::default(), + &ctx, + temp.path(), + &replacements, + ) + .await + .unwrap(); assert_eq!(benches["clickbench"][0].subgroup(), expected); } @@ -929,7 +2091,7 @@ mod tests { } #[tokio::test] - async fn list_output_is_sorted_and_includes_counts() { + async fn list_output_is_sorted_and_includes_counts_and_descriptions() { let temp = tempfile::tempdir().unwrap(); write_benchmark( @@ -944,19 +2106,63 @@ mod tests { ); write_benchmark( temp.path(), - "alpha/benchmarks/q02.benchmark", + "beta/benchmarks/q02.benchmark", "name Q02\n\nrun\nSELECT 2\n", ); + write_suite(temp.path(), "alpha", "Alpha workload"); + write_suite(temp.path(), "beta", "Beta workload"); - let ctx = SessionContext::new(); - let benches = load_benchmarks(&BenchmarkFilter::default(), &ctx, temp.path()) + let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) .await .unwrap(); - let output = format_benchmark_list(&benches); - assert!(output.starts_with("SQL benchmarks:\n alpha")); - assert!(output.contains("alpha 2 queries")); - assert!(output.contains("beta 1 query")); + assert_eq!( + output, + "SQL benchmarks:\n alpha 1 query Alpha workload\n beta 2 queries Beta workload" + ); + } + + #[tokio::test] + async fn list_does_not_parse_benchmark_sql() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "not valid benchmark syntax", + ); + write_suite(temp.path(), "alpha", "Alpha workload"); + + let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) + .await + .unwrap(); + + assert_eq!( + output, + "SQL benchmarks:\n alpha 1 query Alpha workload" + ); + } + + #[tokio::test] + async fn list_malformed_metadata_names_its_file() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + let metadata_path = temp.path().join("alpha/alpha.suite"); + fs::write(&metadata_path, "not valid metadata").unwrap(); + + let error = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains(&metadata_path.display().to_string()), + "{error}" + ); } #[tokio::test] diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index 8d24d44a174e3..7d8b7044bbdd8 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -28,6 +28,7 @@ pub mod sort_pushdown; pub mod sort_tpch; pub mod sql_benchmark; pub mod sql_benchmark_runner; +pub mod sql_benchmark_suite; pub mod tpcds; pub mod tpch; pub mod util; diff --git a/benchmarks/src/sql_benchmark_runner.rs b/benchmarks/src/sql_benchmark_runner.rs index f1cb3ad2f71b4..414a4c0e91dc6 100644 --- a/benchmarks/src/sql_benchmark_runner.rs +++ b/benchmarks/src/sql_benchmark_runner.rs @@ -42,6 +42,8 @@ pub struct BenchmarkFilter { pub struct SqlRunConfig { pub common: CommonOpt, pub filter: BenchmarkFilter, + pub replacements: HashMap, + pub query_filename: Option, pub persist_results: bool, pub validate_results: bool, pub output: Option, @@ -55,10 +57,12 @@ pub fn run_criterion_benchmarks_impl( ) -> Result<()> { let rt = make_tokio_runtime()?; let listing_ctx = make_ctx(&config.common)?; - let all_benchmarks = rt.block_on(load_benchmark_definitions( + let all_benchmarks = rt.block_on(load_benchmark_definitions_for_query( &config.filter, &listing_ctx, benchmark_dir, + &config.replacements, + config.query_filename.as_deref(), ))?; let selected = filter_benchmarks(&config.filter, all_benchmarks.clone()); @@ -134,6 +138,23 @@ pub fn default_sql_benchmark_directory() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("sql_benchmarks") } +/// Replacements used by the Criterion SQL benchmark harness. +pub fn default_criterion_replacements() -> HashMap { + criterion_replacements(std::env::var("DATA_DIR").ok()) +} + +fn criterion_replacements(data_dir: Option) -> HashMap { + HashMap::from([( + "data_dir".to_string(), + data_dir.unwrap_or_else(|| { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("data") + .to_string_lossy() + .into_owned() + }), + )]) +} + fn make_tokio_runtime() -> Result { tokio::runtime::Builder::new_multi_thread() .enable_all() @@ -163,11 +184,41 @@ pub async fn load_benchmark_definitions( filter: &BenchmarkFilter, ctx: &SessionContext, benchmark_dir: &Path, + replacements: &HashMap, +) -> Result>> { + load_benchmark_definitions_for_query(filter, ctx, benchmark_dir, replacements, None) + .await +} + +/// Loads benchmark definitions, optionally limiting discovery to one filename. +pub async fn load_benchmark_definitions_for_query( + filter: &BenchmarkFilter, + ctx: &SessionContext, + benchmark_dir: &Path, + replacements: &HashMap, + query_filename: Option<&str>, ) -> Result>> { let mut benches = BTreeMap::new(); - let replacements = benchmark_replacements(filter); + let mut replacements = replacements.clone(); + let selected_suite_dir = filter + .name + .as_ref() + .map(|name| benchmark_dir.join(name.to_ascii_lowercase())) + .filter(|path| path.is_dir()); + let discovery_dir = selected_suite_dir.as_deref().unwrap_or(benchmark_dir); + if let Some(subgroup) = &filter.subgroup { + replacements.insert("bench_subgroup".to_string(), subgroup.to_string()); + } - for path in discover_benchmark_paths(benchmark_dir)? { + for path in discover_benchmark_paths(discovery_dir)? + .into_iter() + .filter(|path| { + query_filename.is_none_or(|filename| { + path.file_name() + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(filename)) + }) + }) + { let benchmark = SqlBenchmark::new_with_replacements( ctx, &path, @@ -186,25 +237,6 @@ pub async fn load_benchmark_definitions( Ok(benches) } -/// Builds template replacements from CLI values that also appear in benchmark files. -fn benchmark_replacements(filter: &BenchmarkFilter) -> HashMap { - let mut replacements = HashMap::new(); - let data_dir = std::env::var("DATA_DIR").unwrap_or_else(|_| { - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("data") - .to_string_lossy() - .into_owned() - }); - - replacements.insert("data_dir".to_string(), data_dir); - - if let Some(subgroup) = &filter.subgroup { - replacements.insert("bench_subgroup".to_string(), subgroup.to_string()); - } - - replacements -} - pub fn sort_benchmarks(benchmarks: &mut BTreeMap>) { benchmarks .values_mut() @@ -560,6 +592,155 @@ mod tests { path } + #[tokio::test] + async fn caller_replacements_reach_parser() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nload\nSELECT '${ALPHA_FORMAT}'\n\nrun\nSELECT 1\n", + ); + let replacements = + HashMap::from([("alpha_format".to_string(), "csv".to_string())]); + + let result = load_benchmark_definitions( + &BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("1".to_string()), + }, + &SessionContext::new(), + temp.path(), + &replacements, + ) + .await; + + assert!(result.is_ok()); + } + + #[tokio::test] + async fn query_filename_filters_paths_before_parsing() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q07.benchmark", + "name Q07\n\nrun\nSELECT 7\n", + ); + write_benchmark( + temp.path(), + "alpha/benchmarks/q08.benchmark", + "this is not a benchmark definition", + ); + write_benchmark( + temp.path(), + "beta/benchmarks/q07.benchmark", + "this is not a benchmark definition", + ); + + let benches = load_benchmark_definitions_for_query( + &BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("7".to_string()), + }, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("q07.benchmark"), + ) + .await + .unwrap(); + + assert_eq!(benches["alpha"].len(), 1); + assert_eq!(benches["alpha"][0].name(), "Q07"); + } + + #[test] + fn criterion_replacements_use_benchmarks_data_directory() { + let expected = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("data") + .to_string_lossy() + .into_owned(); + + assert_eq!(criterion_replacements(None)["data_dir"], expected); + } + + #[test] + fn criterion_replacements_use_explicit_data_directory() { + let replacements = criterion_replacements(Some("/custom/data".to_string())); + + assert_eq!(replacements["data_dir"], "/custom/data"); + } + + #[tokio::test] + async fn query_filename_keeps_matches_in_multiple_subgroups() { + let temp = tempfile::tempdir().unwrap(); + for subgroup in ["aggregate", "window"] { + write_benchmark( + temp.path(), + &format!("alpha/benchmarks/{subgroup}/q03.benchmark"), + &format!("name Q03\nsubgroup {subgroup}\n\nrun\nSELECT 3\n"), + ); + } + + let filter = BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("3".to_string()), + }; + let benches = load_benchmark_definitions_for_query( + &filter, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("q03.benchmark"), + ) + .await + .unwrap(); + assert_eq!(filter_benchmarks(&filter, benches)["alpha"].len(), 2); + + let filter = BenchmarkFilter { + subgroup: Some("window".to_string()), + ..filter + }; + let benches = load_benchmark_definitions_for_query( + &filter, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("q03.benchmark"), + ) + .await + .unwrap(); + assert_eq!(filter_benchmarks(&filter, benches)["alpha"].len(), 1); + } + + #[tokio::test] + async fn query_filename_accepts_alphanumeric_pattern() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "imdb/benchmarks/01a.benchmark", + "name Q01a\n\nrun\nSELECT 1\n", + ); + + let benches = load_benchmark_definitions_for_query( + &BenchmarkFilter { + name: Some("imdb".to_string()), + subgroup: None, + query: Some("1a".to_string()), + }, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("01a.benchmark"), + ) + .await + .unwrap(); + + assert_eq!(benches["imdb"][0].name(), "Q01a"); + } + #[test] fn normalizes_query_like_existing_sql_harness() { assert_eq!(normalize_query("1"), "Q01"); diff --git a/benchmarks/src/sql_benchmark_suite.rs b/benchmarks/src/sql_benchmark_suite.rs new file mode 100644 index 0000000000000..aa7a3c5d52c8e --- /dev/null +++ b/benchmarks/src/sql_benchmark_suite.rs @@ -0,0 +1,849 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Metadata parsing and validation for SQL benchmark suites. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::io; +use std::path::{Component, Path, PathBuf}; + +use datafusion_common::{DataFusionError, Result}; +use serde::Deserialize; + +const DEFAULT_QUERY_PATTERN: &str = "q{QUERY_ID_PADDED}.benchmark"; + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawSuite { + description: String, + query_pattern: Option, + #[serde(default)] + path_replacements: BTreeMap, + #[serde(default)] + options: Vec, + #[serde(default)] + examples: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawSuiteOption { + name: String, + short: Option, + env: String, + default: String, + values: Option>, + help: String, +} + +/// Validated metadata for one benchmark suite. +#[derive(Debug, Clone)] +pub struct SuiteMetadata { + name: String, + directory: PathBuf, + description: String, + query_pattern: String, + path_replacements: BTreeMap, + options: Vec, + examples: Vec, + benchmark_count: usize, +} + +/// A suite-specific command-line option. +#[derive(Debug, Clone)] +pub struct SuiteOption { + name: String, + short: Option, + env: String, + default: String, + values: Option>, + help: String, +} + +/// An example invocation from a suite metadata file. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SuiteExample { + command: String, + description: String, +} + +/// Global option names unavailable to suite-specific options. +pub struct ReservedOptions<'a> { + pub long: &'a BTreeSet, + pub short: &'a BTreeSet, +} + +/// Where a resolved option value originated. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ValueSource { + CommandLine, + Environment, + Default, +} + +/// An option value together with its origin. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResolvedValue { + pub value: String, + pub source: ValueSource, +} + +fn metadata_error(message: impl Into) -> DataFusionError { + DataFusionError::Configuration(message.into()) +} + +impl SuiteMetadata { + /// Loads and validates `//.suite`. + pub fn load(root: &Path, name: &str, reserved: &ReservedOptions) -> Result { + let directory = root.join(name); + let metadata_path = directory.join(format!("{name}.suite")); + let contents = fs::read_to_string(&metadata_path)?; + let raw: RawSuite = toml::from_str(&contents).map_err(|error| { + metadata_error(format!("{}: {error}", metadata_path.display())) + })?; + Self::from_raw(name, directory, raw, reserved) + } + + fn from_raw( + name: &str, + directory: PathBuf, + raw: RawSuite, + reserved: &ReservedOptions, + ) -> Result { + if raw.description.trim().is_empty() { + return Err(metadata_error("suite description must not be empty")); + } + + let query_pattern = raw + .query_pattern + .clone() + .unwrap_or_else(|| DEFAULT_QUERY_PATTERN.to_string()); + + validate_query_pattern(&query_pattern)?; + + let mut long_names = BTreeSet::new(); + let mut short_names = BTreeSet::new(); + let mut env_names = BTreeSet::new(); + let mut options = Vec::with_capacity(raw.options.len()); + + for option in raw.options { + Self::validate_option( + reserved, + &mut long_names, + &mut short_names, + &mut env_names, + &raw.path_replacements, + &option, + )?; + + let suite_option = SuiteOption { + name: option.name, + short: option.short.as_deref().map(parse_short).transpose()?, + env: option.env, + default: option.default, + values: option.values, + help: option.help, + }; + + if !suite_option.accepts(&suite_option.default) { + return Err(metadata_error(format!( + "default value '{}' is not accepted by option '{}'", + suite_option.default, suite_option.name + ))); + } + + options.push(suite_option); + } + + for example in &raw.examples { + if example.command.trim().is_empty() { + return Err(metadata_error("example command must not be empty")); + } + if example.description.trim().is_empty() { + return Err(metadata_error("example description must not be empty")); + } + } + + let path_replacements = raw + .path_replacements + .into_iter() + .map(|(key, path)| { + let path = PathBuf::from(path); + let path = if path.is_relative() { + directory.join(path) + } else { + path + }; + (key, path) + }) + .collect(); + let benchmark_count = count_benchmarks(&directory)?; + + Ok(Self { + name: name.to_string(), + directory, + description: raw.description, + query_pattern, + path_replacements, + options, + examples: raw.examples, + benchmark_count, + }) + } + + fn validate_option( + reserved: &ReservedOptions, + long_names: &mut BTreeSet, + short_names: &mut BTreeSet, + env_names: &mut BTreeSet, + path_replacements: &BTreeMap, + option: &RawSuiteOption, + ) -> Result<()> { + if !valid_long_name(&option.name) { + return Err(metadata_error(format!( + "invalid option name '{}'", + option.name + ))); + } + if reserved.long.contains(&option.name) || !long_names.insert(option.name.clone()) + { + return Err(metadata_error(format!( + "option name '{}' is reserved or duplicated", + option.name + ))); + } + let short = option.short.as_deref().map(parse_short).transpose()?; + if let Some(short) = short + && (reserved.short.contains(&short) || !short_names.insert(short)) + { + return Err(metadata_error(format!( + "option short name '{short}' is reserved or duplicated" + ))); + } + if !env_names.insert(option.env.clone()) { + return Err(metadata_error(format!( + "option environment key '{}' is duplicated", + option.env + ))); + } + if path_replacements.contains_key(&option.env) { + return Err(metadata_error(format!( + "environment key '{}' is used by both an option and a path replacement", + option.env + ))); + } + if option.help.trim().is_empty() { + return Err(metadata_error(format!( + "help for option '{}' must not be empty", + option.name + ))); + } + + Ok(()) + } + + pub fn name(&self) -> &str { + &self.name + } + + pub fn directory(&self) -> &Path { + &self.directory + } + + pub fn description(&self) -> &str { + &self.description + } + + pub fn query_pattern(&self) -> &str { + &self.query_pattern + } + + pub fn path_replacements(&self) -> &BTreeMap { + &self.path_replacements + } + + pub fn options(&self) -> &[SuiteOption] { + &self.options + } + + pub fn examples(&self) -> &[SuiteExample] { + &self.examples + } + + pub fn benchmark_count(&self) -> usize { + self.benchmark_count + } + + /// Formats a query identifier using this suite's query pattern. + pub fn query_filename(&self, query: &str) -> Result { + let query = query.strip_prefix(['q', 'Q']).unwrap_or(query); + let digit_count = query.bytes().take_while(u8::is_ascii_digit).count(); + + if digit_count == 0 || !query.bytes().all(|byte| byte.is_ascii_alphanumeric()) { + return Err(metadata_error(format!( + "invalid query identifier '{query}'" + ))); + } + + let (digits, suffix) = query.split_at(digit_count); + let replacement = if self.query_pattern.contains("{QUERY_ID_PADDED}") { + let digits = digits.trim_start_matches('0'); + let digits = if digits.is_empty() { "0" } else { digits }; + format!("{digits:0>2}{suffix}") + } else { + query.to_string() + }; + + Ok(self + .query_pattern + .replace("{QUERY_ID_PADDED}", &replacement) + .replace("{QUERY_ID}", &replacement)) + } +} + +impl SuiteOption { + pub fn name(&self) -> &str { + &self.name + } + + pub fn short(&self) -> Option { + self.short + } + + pub fn env(&self) -> &str { + &self.env + } + + pub fn default(&self) -> &str { + &self.default + } + + pub fn values(&self) -> Option<&[String]> { + self.values.as_deref() + } + + pub fn help(&self) -> &str { + &self.help + } + + /// Whether `value` belongs to this option's configured value set. + pub fn accepts(&self, value: &str) -> bool { + self.values.as_ref().is_none_or(|values| { + values + .iter() + .any(|allowed| allowed == value || allowed == "...") + }) + } +} + +impl SuiteExample { + pub fn command(&self) -> &str { + &self.command + } + + pub fn description(&self) -> &str { + &self.description + } +} + +/// Finds and loads suite metadata immediately below `root`, sorted by name. +pub fn discover_suites( + root: &Path, + reserved: &ReservedOptions, +) -> Result> { + let mut suites = Vec::new(); + + for entry in collect_sorted_entries(fs::read_dir(root)?)? { + if !entry.file_type()?.is_dir() { + continue; + } + + let name = entry.file_name().to_string_lossy().into_owned(); + let expected = entry.path().join(format!("{name}.suite")); + let suite_files = collect_sorted_entries(fs::read_dir(entry.path())?)? + .into_iter() + .filter(|entry| entry.path().extension().is_some_and(|ext| ext == "suite")) + .collect::>(); + + if suite_files.is_empty() { + continue; + } + if suite_files.len() != 1 || suite_files[0].path() != expected { + return Err(metadata_error(format!( + "suite metadata filename must match directory name '{name}'" + ))); + } + + suites.push(SuiteMetadata::load(root, &name, reserved)?); + } + + suites.sort_by(|left, right| left.name.cmp(&right.name)); + + Ok(suites) +} + +fn valid_long_name(name: &str) -> bool { + name.bytes() + .next() + .is_some_and(|byte| byte.is_ascii_alphanumeric()) + && name.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-' + }) +} + +fn parse_short(short: &str) -> Result { + let mut chars = short.chars(); + let value = chars.next().filter(char::is_ascii_alphanumeric); + + match (value, chars.next()) { + (Some(value), None) => Ok(value), + _ => Err(metadata_error(format!( + "invalid option short name '{short}': expected one ASCII alphanumeric character" + ))), + } +} + +fn validate_query_pattern(pattern: &str) -> Result<()> { + let path = Path::new(pattern); + if path.is_absolute() { + return Err(metadata_error("query pattern must not be absolute")); + } + if path + .components() + .any(|component| component == Component::ParentDir) + { + return Err(metadata_error( + "query pattern must not contain a parent component", + )); + } + + let placeholders = pattern.matches("{QUERY_ID}").count() + + pattern.matches("{QUERY_ID_PADDED}").count(); + if placeholders != 1 { + return Err(metadata_error( + "query pattern must contain exactly one query identifier placeholder", + )); + } + + Ok(()) +} + +fn count_benchmarks(directory: &Path) -> Result { + let mut count = 0; + for entry in collect_sorted_entries(fs::read_dir(directory)?)? { + if entry.file_type()?.is_dir() { + count += count_benchmarks(&entry.path())?; + } else if entry + .path() + .extension() + .is_some_and(|ext| ext == "benchmark") + { + count += 1; + } + } + + Ok(count) +} + +fn collect_sorted_entries( + entries: impl IntoIterator>, +) -> io::Result> { + let mut entries = entries.into_iter().collect::>>()?; + entries.sort_by_key(|entry| entry.file_name()); + + Ok(entries) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sql_benchmark_runner::default_sql_benchmark_directory; + use std::collections::BTreeSet; + use std::fs; + use std::io; + use std::path::Path; + + fn reserved() -> ReservedOptions<'static> { + let long = Box::leak(Box::new(BTreeSet::from([ + "help".to_string(), + "query".to_string(), + ]))); + let short = Box::leak(Box::new(BTreeSet::from(['h', 'q']))); + ReservedOptions { long, short } + } + + fn write_suite(root: &Path, name: &str, metadata: &str) { + let directory = root.join(name); + fs::create_dir_all(&directory).unwrap(); + fs::write(directory.join(format!("{name}.suite")), metadata).unwrap(); + } + + fn minimal(extra: &str) -> String { + format!("description = \"Benchmark\"\n{extra}") + } + + #[test] + fn loads_complete_suite() { + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + r#" +description = "Alpha benchmark" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" +[path_replacements] +DATA_DIR = "../../data" +[[options]] +name = "format" +short = "f" +env = "ALPHA_FORMAT" +default = "parquet" +values = ["parquet", "csv"] +help = "Select the file format." +[[examples]] +command = "benchmark_runner alpha -q 1 -f csv" +description = "Run query 1 against CSV." +"#, + ); + + let suite = SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap(); + assert_eq!(suite.name(), "alpha"); + assert_eq!(suite.description(), "Alpha benchmark"); + assert_eq!(suite.options()[0].short(), Some('f')); + assert!(suite.options()[0].accepts("csv")); + assert!(!suite.options()[0].accepts("json")); + assert_eq!( + suite.path_replacements()["DATA_DIR"], + temp.path().join("alpha/../../data") + ); + assert_eq!(suite.examples().len(), 1); + } + + #[test] + fn rejects_unknown_field() { + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + "description = \"Alpha\"\ndescripton = \"bad\"\n", + ); + let error = SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains("descripton")); + assert!(error.to_string().contains("alpha.suite")); + } + + #[test] + fn validates_value_sets() { + let closed = suite_option(Some(vec!["csv", "parquet"])); + assert!(closed.accepts("csv")); + assert!(!closed.accepts("json")); + assert!(suite_option(Some(vec!["1", "10", "..."])).accepts("100")); + assert!(suite_option(None).accepts("anything")); + } + + fn suite_option(values: Option>) -> SuiteOption { + SuiteOption { + name: "format".to_string(), + short: Some('f'), + env: "FORMAT".to_string(), + default: "csv".to_string(), + values: values.map(|values| values.into_iter().map(str::to_string).collect()), + help: "Format".to_string(), + } + } + + #[test] + fn rejects_invalid_metadata() { + let cases = [ + ("empty description", "description = \" \"\n", "description"), + ( + "invalid long", + &minimal( + "[[options]]\nname = \"Bad_name\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "Bad_name", + ), + ( + "long starts hyphen", + &minimal( + "[[options]]\nname = \"-bad\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "-bad", + ), + ( + "long reserved", + &minimal( + "[[options]]\nname = \"query\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "query", + ), + ( + "short long", + &minimal( + "[[options]]\nname = \"format\"\nshort = \"ff\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "ff", + ), + ( + "short invalid", + &minimal( + "[[options]]\nname = \"format\"\nshort = \"-\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "short", + ), + ( + "short reserved", + &minimal( + "[[options]]\nname = \"format\"\nshort = \"q\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "q", + ), + ( + "empty help", + &minimal( + "[[options]]\nname = \"format\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \" \"\n", + ), + "help", + ), + ( + "bad default", + &minimal( + "[[options]]\nname = \"format\"\nenv = \"ENV\"\ndefault = \"json\"\nvalues = [\"csv\"]\nhelp = \"help\"\n", + ), + "json", + ), + ( + "absolute pattern", + "description = \"Benchmark\"\nquery_pattern = \"/q{QUERY_ID}.benchmark\"\n", + "absolute", + ), + ( + "parent pattern", + "description = \"Benchmark\"\nquery_pattern = \"../q{QUERY_ID}.benchmark\"\n", + "parent", + ), + ( + "no placeholder", + "description = \"Benchmark\"\nquery_pattern = \"q.benchmark\"\n", + "placeholder", + ), + ( + "two placeholders", + "description = \"Benchmark\"\nquery_pattern = \"{QUERY_ID}-{QUERY_ID_PADDED}.benchmark\"\n", + "exactly one", + ), + ( + "empty example command", + &minimal("[[examples]]\ncommand = \" \"\ndescription = \"example\"\n"), + "command", + ), + ( + "empty example description", + &minimal("[[examples]]\ncommand = \"runner\"\ndescription = \" \"\n"), + "description", + ), + ]; + + for (name, metadata, expected) in cases { + let temp = tempfile::tempdir().unwrap(); + write_suite(temp.path(), "alpha", metadata); + let error = + SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains(expected), "{name}: {error}"); + } + } + + #[test] + fn rejects_duplicate_and_colliding_options() { + let fields = [("name", "format"), ("short", "f"), ("env", "FORMAT")]; + for (field, value) in fields { + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + &minimal(&format!( + r#" +[[options]] +name = "format" +short = "f" +env = "FORMAT" +default = "x" +help = "help" +[[options]] +name = "{name}" +short = "{short}" +env = "{env}" +default = "x" +help = "help" +"#, + name = if field == "name" { value } else { "other" }, + short = if field == "short" { value } else { "o" }, + env = if field == "env" { value } else { "OTHER" } + )), + ); + let error = + SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains(value), "{field}: {error}"); + } + + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + &minimal( + "[path_replacements]\nFORMAT = \"data\"\n[[options]]\nname = \"format\"\nenv = \"FORMAT\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + ); + let error = SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains("FORMAT")); + } + + #[test] + fn discovers_sorted_suites_and_counts_benchmarks() { + let temp = tempfile::tempdir().unwrap(); + write_suite(temp.path(), "zeta", "description = \"Zeta\"\n"); + write_suite(temp.path(), "alpha", "description = \"Alpha\"\n"); + fs::create_dir_all(temp.path().join("alpha/nested")).unwrap(); + fs::write(temp.path().join("alpha/q01.benchmark"), "").unwrap(); + fs::write(temp.path().join("alpha/nested/q02.benchmark"), "").unwrap(); + fs::write(temp.path().join("alpha/ignored.sql"), "").unwrap(); + let suites = discover_suites(temp.path(), &reserved()).unwrap(); + assert_eq!( + suites.iter().map(SuiteMetadata::name).collect::>(), + ["alpha", "zeta"] + ); + assert_eq!(suites[0].benchmark_count(), 2); + } + + #[test] + fn checked_in_suites_cover_benchmark_directories() { + let root = default_sql_benchmark_directory(); + for entry in fs::read_dir(&root).unwrap() { + let entry = entry.unwrap(); + if !entry.file_type().unwrap().is_dir() { + continue; + } + let directory = entry.path(); + let has_benchmark = count_benchmarks(&directory).unwrap() > 0; + if has_benchmark { + let name = entry.file_name().to_string_lossy().into_owned(); + assert!( + directory.join(format!("{name}.suite")).is_file(), + "benchmark directory {name} is missing {name}.suite" + ); + } + } + + let long = BTreeSet::from([ + "batch-size".to_string(), + "debug".to_string(), + "iterations".to_string(), + "output".to_string(), + "partitions".to_string(), + "path".to_string(), + "query".to_string(), + ]); + let short = BTreeSet::from(['q', 'i', 'n', 's', 'd', 'p', 'o']); + let suites = discover_suites( + &root, + &ReservedOptions { + long: &long, + short: &short, + }, + ) + .unwrap(); + let by_name = suites + .iter() + .map(|suite| (suite.name(), suite)) + .collect::>(); + + assert_eq!( + by_name["imdb"].query_filename("1a").unwrap(), + "01a.benchmark" + ); + assert_eq!( + by_name["imdb"].query_filename("01a").unwrap(), + "01a.benchmark" + ); + assert_eq!(by_name["clickbench"].options()[0].name(), "partitioning"); + assert!( + by_name["tpch"] + .options() + .iter() + .all(|option| option.short() != Some('s')) + ); + } + + #[test] + fn rejects_mismatched_suite_filename() { + let temp = tempfile::tempdir().unwrap(); + fs::create_dir_all(temp.path().join("wrong")).unwrap(); + fs::write( + temp.path().join("wrong/other.suite"), + "description = \"Wrong\"", + ) + .unwrap(); + + let error = discover_suites(temp.path(), &reserved()).unwrap_err(); + assert!(error.to_string().contains("wrong")); + } + + #[test] + fn propagates_directory_entry_errors() { + let entries = std::iter::once(Err::(io::Error::new( + io::ErrorKind::PermissionDenied, + "entry denied", + ))); + + let error = collect_sorted_entries(entries).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::PermissionDenied); + assert_eq!(error.to_string(), "entry denied"); + } + + #[test] + fn formats_query_filenames() { + let temp = tempfile::tempdir().unwrap(); + write_suite(temp.path(), "padded", "description = \"Padded\"\n"); + write_suite( + temp.path(), + "plain", + "description = \"Plain\"\nquery_pattern = \"{QUERY_ID}.benchmark\"\n", + ); + let padded = SuiteMetadata::load(temp.path(), "padded", &reserved()).unwrap(); + let plain = SuiteMetadata::load(temp.path(), "plain", &reserved()).unwrap(); + + assert_eq!(padded.query_filename("7").unwrap(), "q07.benchmark"); + assert_eq!(padded.query_filename("07").unwrap(), "q07.benchmark"); + assert_eq!( + padded.query_filename("Q1").unwrap(), + padded.query_filename("q1").unwrap() + ); + assert_eq!( + plain.query_filename("Q01a").unwrap(), + plain.query_filename("q01a").unwrap() + ); + assert_eq!( + padded.query_filename("184467440737095516160").unwrap(), + "q184467440737095516160.benchmark" + ); + assert_eq!(plain.query_filename("01a").unwrap(), "01a.benchmark"); + assert!(plain.query_filename("abc").is_err()); + assert!(plain.query_filename("1-a").is_err()); + } +}