From 8216ecfeb5a72193ad6d4ec34c5c9369176edbf8 Mon Sep 17 00:00:00 2001 From: Bruce Ritchie Date: Tue, 5 May 2026 15:32:12 -0400 Subject: [PATCH] Add benchmark_runner info and query commands --- benchmarks/Cargo.toml | 2 +- benchmarks/benches/sql.rs | 9 +- benchmarks/sql_benchmarks/tpch/tpch.suite | 34 +- benchmarks/src/benchmark_runner/cli.rs | 958 ++++++++++++- benchmarks/src/benchmark_runner/mod.rs | 325 ++++- benchmarks/src/benchmark_runner/native.rs | 362 +++++ benchmarks/src/benchmark_runner/output.rs | 583 +++++++- benchmarks/src/benchmark_runner/selector.rs | 356 +++++ benchmarks/src/benchmark_runner/style.rs | 58 + benchmarks/src/benchmark_runner/suite.rs | 1360 +++++++++++++++++-- benchmarks/src/bin/benchmark_runner.rs | 5 +- benchmarks/src/sql_benchmark.rs | 211 ++- 12 files changed, 4049 insertions(+), 214 deletions(-) create mode 100644 benchmarks/src/benchmark_runner/native.rs create mode 100644 benchmarks/src/benchmark_runner/selector.rs create mode 100644 benchmarks/src/benchmark_runner/style.rs diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 97e0d901b95f9..d5f146f63edfb 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -44,7 +44,7 @@ anstream = "1.0" arrow = { workspace = true } async-trait = "0.1" bytes = { workspace = true } -clap = { version = "4.6.1", features = ["derive", "env", "color"] } +clap = { version = "4.6.1", features = ["derive", "env", "color", "string"] } criterion = { workspace = true, features = ["html_reports"] } datafusion = { workspace = true, default-features = true } datafusion-common = { workspace = true, default-features = true } diff --git a/benchmarks/benches/sql.rs b/benchmarks/benches/sql.rs index eade3194d1402..85f8fe51723b8 100644 --- a/benchmarks/benches/sql.rs +++ b/benchmarks/benches/sql.rs @@ -172,10 +172,10 @@ fn handle_run( name: &str, ) { rt.block_on(async { - benchmark + let _ = benchmark .run(ctx, args.validate) .await - .unwrap_or_else(|err| panic!("Failed to run benchmark {name}: {err:?}")) + .unwrap_or_else(|err| panic!("Failed to run benchmark {name}: {err:?}")); }); } @@ -208,7 +208,7 @@ fn handle_verify( info!("Verifying benchmark {name} results ..."); rt.block_on(async { - benchmark + let _ = benchmark .run(ctx, true) .await .unwrap_or_else(|err| panic!("Failed to run benchmark {name}: {err:?}")); @@ -281,7 +281,8 @@ async fn load_benchmarks( for path in paths { debug!("Loading benchmark from {path}"); - let benchmark = SqlBenchmark::new(ctx, &path, &*SQL_BENCHMARK_DIRECTORY).await?; + let benchmark = + SqlBenchmark::new(ctx, &path, &*SQL_BENCHMARK_DIRECTORY, None).await?; let entries = benches .entry(benchmark.group().to_string()) .or_insert(vec![]); diff --git a/benchmarks/sql_benchmarks/tpch/tpch.suite b/benchmarks/sql_benchmarks/tpch/tpch.suite index 317b32e57f45f..0a6174f449f58 100644 --- a/benchmarks/sql_benchmarks/tpch/tpch.suite +++ b/benchmarks/sql_benchmarks/tpch/tpch.suite @@ -1,9 +1,24 @@ -name = "tpch" 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." @@ -11,6 +26,23 @@ help = "Selects the TPC-H data format." [[options]] name = "scale-factor" short = "sf" +env = "BENCH_SIZE" default = "1" values = ["1", "10"] help = "Selects the TPC-H scale factor." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- run tpch" +description = "Run all TPC-H queries with the default parquet SF1 configuration." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- run tpch 15" +description = "Run TPC-H query 15 with the default parquet SF1 configuration." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- run tpch 15 -f csv" +description = "Run TPC-H query 15 against CSV data." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- run tpch 15 -sf 10" +description = "Run TPC-H query 15 at scale factor 10." diff --git a/benchmarks/src/benchmark_runner/cli.rs b/benchmarks/src/benchmark_runner/cli.rs index 30fb71b5a30cc..49b1680dbfe42 100644 --- a/benchmarks/src/benchmark_runner/cli.rs +++ b/benchmarks/src/benchmark_runner/cli.rs @@ -17,65 +17,973 @@ //! CLI construction and argument conversion for `benchmark_runner`. //! -//! This module owns the clap command tree for the initial runner surface: -//! top-level help and suite listing. +//! This module owns the clap command tree for `list`, `info`, and `query`. +//! Suite options are loaded from `.suite` files and registered as dynamic clap +//! arguments so help output can describe suite-specific flags. Short aliases +//! with more than one character, such as `-sf`, are normalized before clap +//! parses the command because clap only supports one-character short flags +//! directly. -use clap::builder::styling::{AnsiColor, Styles}; -use clap::{ArgMatches, Command}; +use crate::benchmark_runner::native::NativeBenchmarks; +use crate::benchmark_runner::output::write_help_section; +use crate::benchmark_runner::style::{HELP_STYLES, header}; +use crate::benchmark_runner::suite::{SuiteConfig, SuiteOption, SuiteRegistry}; +use crate::util::CommonOpt; +use clap::{Arg, ArgAction, ArgMatches, Args, Command}; use datafusion_common::{Result, exec_datafusion_err}; +use std::collections::{BTreeMap, BTreeSet}; +use std::ffi::{OsStr, OsString}; +use std::fmt::Write as _; -const HELP_STYLES: Styles = Styles::styled() - .header(AnsiColor::Green.on_default().bold()) - .usage(AnsiColor::Green.on_default().bold()) - .literal(AnsiColor::Cyan.on_default().bold()) - .placeholder(AnsiColor::Cyan.on_default()); +const BENCHMARK_TARGET_HEADING: &str = "SQL Benchmark Target"; +const RUNNER_OPTIONS_HEADING: &str = "Runner Options"; +const DATAFUSION_OPTIONS_HEADING: &str = "DataFusion Options"; +const SUITE_OPTIONS_HEADING: &str = "Suite Options"; +const DATAFUSION_ARG_IDS: &[&str] = &[ + "partitions", + "batch_size", + "mem_pool_type", + "memory_limit", + "sort_spill_reservation_bytes", + "debug", + "simulate_latency", +]; +const RUN_ENV_ARG_IDS: &[&str] = &[ + "iterations", + "partitions", + "batch_size", + "mem_pool_type", + "memory_limit", + "sort_spill_reservation_bytes", + "debug", + "simulate_latency", +]; +const HIDDEN_RUN_ARG_IDS: &[&str] = &["debug", "sort_spill_reservation_bytes"]; +/// Parsed top-level command selected by clap and dynamic suite metadata. #[derive(Debug)] pub enum RunnerCommand { + /// Print root command help. Help, + /// List SQL suites and visible native benchmarks. List, + /// Show SQL suite or query metadata. + Info(InspectArgs), + /// Show native benchmark help-style metadata. + InfoNative(String), + /// Print SQL parsed from one benchmark query. + Query(InspectArgs), } -/// Builds the command tree for help and suite listing. -pub fn build_cli() -> Command { +/// Suite-option visibility used while constructing selector subcommands. +#[derive(Clone, Copy)] +enum SelectorSuiteOptions<'a> { + /// Include every discovered suite option. + All, + /// Include only options from the selected suite. + Suite(&'a SuiteConfig), + /// Do not include suite options. + None, +} + +/// Selector subcommands that accept a SQL suite name. +#[derive(Clone, Copy)] +enum SelectorCommand { + /// The `info` subcommand. + Info, + /// The `query` subcommand. + Query, +} + +/// Common selector arguments for SQL inspection and execution commands. +#[derive(Debug, Clone)] +pub struct InspectArgs { + /// Suite selector provided on the command line. + pub selector: String, + /// Optional query id provided on the command line. + pub query_id: Option, + /// Suite option values supplied by the user. + pub suite_options: BTreeMap, +} + +/// Per-subcommand suite-option visibility used for help rendering. +#[derive(Clone, Copy)] +struct SelectorOptionVisibility<'a> { + info: SelectorSuiteOptions<'a>, + query: SelectorSuiteOptions<'a>, +} + +/// Constructs visibility rules for normal and help-specific command trees. +impl<'a> SelectorOptionVisibility<'a> { + /// Shows all suite options for each selector command. + fn all() -> Self { + Self { + info: SelectorSuiteOptions::All, + query: SelectorSuiteOptions::All, + } + } + + /// Shows a selected option set for one selector command's help output. + fn for_help(command: SelectorCommand, options: SelectorSuiteOptions<'a>) -> Self { + let mut visibility = Self::all(); + match command { + SelectorCommand::Info => visibility.info = options, + SelectorCommand::Query => visibility.query = options, + }; + visibility + } +} + +/// Precomputed suite option metadata used while parsing dynamic CLI options. +pub(crate) struct SuiteCliOptions { + /// All suite option long names. + names: BTreeSet, + /// All long option spellings, including the leading `--`. + long_args: BTreeSet, + /// All short alias spellings, including the leading `-`. + alias_args: BTreeSet, + /// Per-suite short alias to long option name mappings. + aliases_by_suite: BTreeMap>, +} + +/// Collects dynamic suite option metadata from the registry. +impl SuiteCliOptions { + /// Collects suite option names and aliases used by clap parsing and + /// pre-parse alias normalization. + pub(crate) fn new(registry: &SuiteRegistry) -> Self { + let mut names = BTreeSet::new(); + let mut long_args = BTreeSet::new(); + let mut alias_args = BTreeSet::new(); + let mut aliases_by_suite = BTreeMap::new(); + + for suite in registry.suites() { + let mut aliases = BTreeMap::new(); + + for option in &suite.options { + names.insert(option.name.clone()); + long_args.insert(OsString::from(format!("--{}", option.name))); + + if let Some(short) = &option.short { + let alias = OsString::from(format!("-{short}")); + alias_args.insert(alias.clone()); + aliases.insert(alias, option.name.clone()); + } + } + aliases_by_suite.insert(suite.name.clone(), aliases); + } + + Self { + names, + long_args, + alias_args, + aliases_by_suite, + } + } +} + +/// Builds the full command tree with all discovered suite options available. +pub(crate) fn build_cli(registry: &SuiteRegistry, native: &NativeBenchmarks) -> Command { + build_cli_with_selector_options(registry, native, SelectorOptionVisibility::all()) +} + +/// Builds a command tree using precomputed suite option metadata. +pub(crate) fn build_cli_for_args_with_options( + registry: &SuiteRegistry, + suite_cli_options: &SuiteCliOptions, + native: &NativeBenchmarks, + args: &[OsString], +) -> Command { + match selector_help_request(suite_cli_options, args) { + Some((command, Some(selector))) => { + let options = registry + .get(&selector) + .map(SelectorSuiteOptions::Suite) + .unwrap_or(SelectorSuiteOptions::None); + + build_cli_with_help_options(registry, native, command, options) + } + Some((command, None)) => build_cli_with_help_options( + registry, + native, + command, + SelectorSuiteOptions::None, + ), + None => build_cli(registry, native), + } +} + +/// Builds a help-oriented command tree for one selector subcommand. +fn build_cli_with_help_options( + registry: &SuiteRegistry, + native: &NativeBenchmarks, + command: SelectorCommand, + options: SelectorSuiteOptions<'_>, +) -> Command { + build_cli_with_selector_options( + registry, + native, + SelectorOptionVisibility::for_help(command, options), + ) +} + +/// Builds the clap command tree with configurable suite-option visibility per selector subcommand. +fn build_cli_with_selector_options( + registry: &SuiteRegistry, + native: &NativeBenchmarks, + visibility: SelectorOptionVisibility<'_>, +) -> Command { + let info = add_configured_suite_options( + base_selector_command("info", "Show benchmark metadata") + .override_usage( + "benchmark_runner info [QUERY_ID]\n benchmark_runner info ", + ), + registry, + visibility.info, + ); + let info = add_native_info_subcommands(info, native); + let query = add_configured_suite_options( + base_selector_command("query", "Show parsed benchmark SQL"), + registry, + visibility.query, + ); + Command::new("benchmark_runner") - .about("Inspect DataFusion SQL benchmark suites.") + .about("Inspect DataFusion benchmarks.") .styles(HELP_STYLES) .subcommand_required(false) .arg_required_else_help(false) .disable_help_subcommand(true) - .subcommand(Command::new("help").about("Print help")) - .subcommand(Command::new("list").about("List SQL benchmark suites")) + .subcommand( + Command::new("help") + .about("Print help") + .disable_help_flag(true), + ) + .subcommand( + Command::new("list") + .about("List benchmark suites") + .disable_help_flag(true), + ) + .subcommand(info) + .subcommand(query) +} + +/// Builds the shared `run` command metadata reused by SQL `info` output. +fn base_run_command() -> Command { + let run = CommonOpt::augment_args( + base_selector_command("run", "Run benchmarks").override_usage(run_usage()), + ) + .mut_arg("iterations", |arg| arg.help_heading(RUNNER_OPTIONS_HEADING)) + .arg( + Arg::new("persist-results") + .long("persist-results") + .action(ArgAction::SetTrue) + .help("Persist benchmark results before the timed run (SQL suites only.)") + .help_heading(RUNNER_OPTIONS_HEADING), + ) + .arg( + Arg::new("validate-results") + .long("validate-results") + .action(ArgAction::SetTrue) + .help("Validate benchmark results before the timed run (SQL suites only.)") + .help_heading(RUNNER_OPTIONS_HEADING), + ) + .arg( + Arg::new("criterion") + .long("criterion") + .action(ArgAction::SetTrue) + .help( + "Run selected benchmarks with Criterion instead of fixed iterations (SQL suites only.)", + ) + .help_heading(RUNNER_OPTIONS_HEADING), + ) + .arg( + Arg::new("save-baseline") + .long("save-baseline") + .value_name("BASELINE") + .help("Save Criterion results under a named baseline (SQL suites only.)") + .help_heading(RUNNER_OPTIONS_HEADING), + ) + .arg( + Arg::new("path") + .long("path") + .short('p') + .value_name("PATH") + .help("Path to benchmark data (default: benchmarks/data)") + .help_heading(RUNNER_OPTIONS_HEADING), + ) + .arg( + Arg::new("output") + .long("output") + .short('o') + .value_name("PATH") + .help("Write run timings and row counts as JSON") + .help_heading(RUNNER_OPTIONS_HEADING), + ); + let run = set_help_heading(run, DATAFUSION_ARG_IDS, DATAFUSION_OPTIONS_HEADING); + let run = hide_env_annotations(run, RUN_ENV_ARG_IDS); + + hide_help_entries(run, HIDDEN_RUN_ARG_IDS) } -/// Converts clap matches into a typed command. -pub(crate) fn command_from_matches(matches: &ArgMatches) -> Result { +/// Returns the condensed SQL-only usage string for run help. +fn run_usage() -> &'static str { + "benchmark_runner run [OPTIONS] [QUERY_ID]" +} + +/// Formats the SQL benchmark `run` help sections reused by SQL `info`. +pub(crate) fn format_sql_run_help_sections_styled( + registry: &SuiteRegistry, +) -> Result { + let run = add_configured_suite_options( + base_run_command(), + registry, + SelectorSuiteOptions::None, + ); + let mut output = String::new(); + + writeln!(output, "{}", header(format!("Usage: {}", run_usage())))?; + writeln!(output)?; + write_help_section(&mut output, &run, BENCHMARK_TARGET_HEADING)?; + write_help_section(&mut output, &run, RUNNER_OPTIONS_HEADING)?; + write_help_section(&mut output, &run, DATAFUSION_OPTIONS_HEADING)?; + + Ok(output.trim_end().to_string()) +} + +/// Converts clap matches into a typed command using precomputed suite option metadata. +pub(crate) fn command_from_matches_with_options( + suite_cli_options: &SuiteCliOptions, + native: &NativeBenchmarks, + matches: &ArgMatches, +) -> Result { match matches.subcommand() { None | Some(("help", _)) => Ok(RunnerCommand::Help), Some(("list", _)) => Ok(RunnerCommand::List), + Some(("info", submatches)) => { + if let Some((name, _native_matches)) = submatches.subcommand() { + if !native.contains(name) { + return Err(exec_datafusion_err!( + "unknown native benchmark '{name}'" + )); + } + reject_suite_options_for_native(suite_cli_options, submatches, name)?; + + return Ok(RunnerCommand::InfoNative(name.to_string())); + } + + Ok(RunnerCommand::Info(inspect_args_with_missing_selector( + suite_cli_options, + submatches, + "Missing benchmark name", + )?)) + } + Some(("query", submatches)) => Ok(RunnerCommand::Query(inspect_args( + suite_cli_options, + submatches, + )?)), Some((name, _)) => Err(exec_datafusion_err!("Unknown command '{name}'")), } } +/// Creates the common `info` and `query` selector command shape. +fn base_selector_command(name: &'static str, about: &'static str) -> Command { + Command::new(name) + .about(about) + .styles(HELP_STYLES) + .arg( + Arg::new("selector") + .value_name("SUITE") + .help("SQL benchmark suite name") + .help_heading(BENCHMARK_TARGET_HEADING) + .display_order(1), + ) + .arg( + Arg::new("query_id") + .value_name("QUERY_ID") + .help("Optional SQL query id, such as 1, 01, 15, or 00") + .help_heading(BENCHMARK_TARGET_HEADING) + .display_order(2), + ) +} + +/// Adds each distinct suite-defined option to a command. +fn add_suite_options(mut command: Command, registry: &SuiteRegistry) -> Command { + let mut seen = BTreeSet::new(); + + for suite in registry.suites() { + for option in &suite.options { + if seen.insert(option.name.clone()) { + command = add_suite_option(command, option, SUITE_OPTIONS_HEADING, false); + } + } + } + command +} + +/// Applies the requested suite-option visibility mode to a selector command. +fn add_configured_suite_options( + command: Command, + registry: &SuiteRegistry, + suite_options: SelectorSuiteOptions<'_>, +) -> Command { + match suite_options { + SelectorSuiteOptions::All => add_suite_options(command, registry), + SelectorSuiteOptions::Suite(suite) => add_suite_options_for(command, suite), + SelectorSuiteOptions::None => command, + } +} + +/// Adds each visible native benchmark as an info-only selector subcommand. +fn add_native_info_subcommands( + mut command: Command, + native: &NativeBenchmarks, +) -> Command { + for benchmark in native.visible() { + command = command.subcommand(benchmark.info_command().hide(true)); + } + + command +} + +/// Adds only one suite's options under a suite-specific help heading. +fn add_suite_options_for(mut command: Command, suite: &SuiteConfig) -> Command { + let heading = suite_options_heading(&suite.name); + + for option in &suite.options { + command = add_suite_option(command, option, heading.clone(), true); + } + + command +} + +/// Registers one suite-defined option with clap, using direct short aliases +/// when clap supports them and documenting multi-character aliases otherwise. +fn add_suite_option( + command: Command, + option: &SuiteOption, + heading: impl Into, + register_short_alias: bool, +) -> Command { + let name = option.name.clone(); + let (short, help) = match &option.short { + Some(short) if register_short_alias && short.chars().count() == 1 => { + (short.chars().next(), option.help.clone()) + } + Some(short) => (None, format!("{} Alias: -{short}.", option.help)), + None => (None, option.help.clone()), + }; + let mut arg = Arg::new(name.clone()) + .long(name) + .value_name("VALUE") + .help(help) + .help_heading(heading); + if let Some(short) = short { + arg = arg.short(short); + } + + command.arg(arg) +} + +/// Formats the suite-specific heading used for selected suite options. +fn suite_options_heading(suite_name: &str) -> String { + format!("{SUITE_OPTIONS_HEADING} ({suite_name})") +} + +/// Detects whether the provided arguments are asking for help on a selector +/// subcommand and, when possible, extracts the suite name preceding `--help`. +fn selector_help_request( + suite_cli_options: &SuiteCliOptions, + args: &[OsString], +) -> Option<(SelectorCommand, Option)> { + if !args.iter().any(|arg| is_help_arg(arg)) { + return None; + } + + selector_request(suite_cli_options, args, true) +} + +/// Scans raw CLI arguments to find a selector before clap parses them. +fn selector_request( + suite_cli_options: &SuiteCliOptions, + args: &[OsString], + stop_at_help: bool, +) -> Option<(SelectorCommand, Option)> { + let (command_position, command) = + args.iter().enumerate().find_map(|(position, arg)| { + selector_command(arg).map(|command| (position, command)) + })?; + let mut skip_next = false; + + for arg in &args[command_position + 1..] { + if skip_next { + skip_next = false; + continue; + } + + if stop_at_help && is_help_arg(arg) { + return Some((command, None)); + } + + if is_value_taking_selector_option(suite_cli_options, command, arg) { + skip_next = true; + continue; + } + + if arg.to_string_lossy().starts_with('-') { + continue; + } + + return Some((command, Some(arg.to_string_lossy().into_owned()))); + } + + Some((command, None)) +} + +/// Converts a raw argument into a selector command name when it matches. +fn selector_command(arg: &OsStr) -> Option { + match arg.to_str()? { + "info" => Some(SelectorCommand::Info), + "query" => Some(SelectorCommand::Query), + _ => None, + } +} + +/// Returns whether a raw argument requests clap help. +fn is_help_arg(arg: &OsStr) -> bool { + arg == "--help" || arg == "-h" +} + +/// Returns whether an argument is an option whose next token should be skipped +/// while looking for a suite selector before `--help`. +fn is_value_taking_selector_option( + suite_cli_options: &SuiteCliOptions, + _command: SelectorCommand, + arg: &OsStr, +) -> bool { + let arg_text = arg.to_string_lossy(); + if arg_text.contains('=') { + return false; + } + + suite_cli_options.long_args.contains(arg) + || suite_cli_options.alias_args.contains(arg) +} + +/// Rewrites suite option aliases using precomputed suite option metadata. +pub(crate) fn normalize_suite_option_aliases_with_options( + suite_cli_options: &SuiteCliOptions, + args: I, +) -> Vec +where + I: IntoIterator, + T: AsRef, +{ + let args = args + .into_iter() + .map(|arg| arg.as_ref().to_os_string()) + .collect::>(); + let aliases = selector_request(suite_cli_options, &args, false) + .and_then(|(_, selector)| selector) + .and_then(|selector| suite_cli_options.aliases_by_suite.get(&selector)); + + args.into_iter() + .map(|arg| { + aliases + .and_then(|aliases| aliases.get(&arg)) + .map(|option| OsString::from(format!("--{option}"))) + .unwrap_or(arg) + }) + .collect() +} + +/// Extracts selector, optional query id, and suite option values from a +/// selector subcommand's clap matches. +fn inspect_args( + suite_cli_options: &SuiteCliOptions, + matches: &ArgMatches, +) -> Result { + inspect_args_with_missing_selector(suite_cli_options, matches, "Missing SUITE") +} + +/// Extracts inspection args using a caller-specific missing-selector message. +fn inspect_args_with_missing_selector( + suite_cli_options: &SuiteCliOptions, + matches: &ArgMatches, + missing_selector_message: &'static str, +) -> Result { + let selector = matches + .get_one::("selector") + .cloned() + .ok_or_else(|| exec_datafusion_err!("{missing_selector_message}"))?; + let query_id = matches.get_one::("query_id").cloned(); + let suite_options = suite_options_from_matches(suite_cli_options, matches); + + Ok(InspectArgs { + selector, + query_id, + suite_options, + }) +} + +/// Rejects SQL suite options that clap accepted before native selection. +fn reject_suite_options_for_native( + suite_cli_options: &SuiteCliOptions, + matches: &ArgMatches, + native_name: &str, +) -> Result<()> { + let suite_options = suite_options_from_matches(suite_cli_options, matches); + + if let Some(option) = suite_options.keys().next() { + return Err(exec_datafusion_err!( + "suite option '--{option}' cannot be used with native benchmark '{native_name}'" + )); + } + + Ok(()) +} + +/// Collects only suite-defined option values that were present in clap matches. +fn suite_options_from_matches( + suite_cli_options: &SuiteCliOptions, + matches: &ArgMatches, +) -> BTreeMap { + let mut values = BTreeMap::new(); + + for option_name in &suite_cli_options.names { + if let Some(value) = matches.get_one::(option_name) { + values.insert(option_name.clone(), value.clone()); + } + } + + values +} + +/// Moves existing clap arguments into a shared help section. +fn set_help_heading( + mut command: Command, + arg_ids: &[&str], + heading: &'static str, +) -> Command { + for id in arg_ids { + command = command.mut_arg(id, |arg| arg.help_heading(heading)); + } + + command +} + +/// Hides environment-variable annotations from help while preserving the +/// underlying argument behavior. +fn hide_env_annotations(mut command: Command, arg_ids: &[&str]) -> Command { + for id in arg_ids { + command = command.mut_arg(id, |arg| arg.hide_env(true)); + } + + command +} + +/// Hides selected inherited options from help while keeping them parseable. +fn hide_help_entries(mut command: Command, arg_ids: &[&str]) -> Command { + for id in arg_ids { + command = command.mut_arg(id, |arg| arg.hide(true)); + } + + command +} + #[cfg(test)] mod tests { use super::*; + use crate::benchmark_runner::native::NativeBenchmarks; + use std::fs; + + fn registry() -> SuiteRegistry { + SuiteRegistry::discover( + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("sql_benchmarks"), + ) + .unwrap() + } + + fn native_registry(registry: &SuiteRegistry) -> NativeBenchmarks { + NativeBenchmarks::new(registry) + } + + fn build_cli_with_native(registry: &SuiteRegistry) -> Command { + build_cli(registry, &native_registry(registry)) + } + + fn os_args(args: I) -> Vec + where + I: IntoIterator, + T: AsRef, + { + args.into_iter() + .map(|arg| arg.as_ref().to_os_string()) + .collect() + } + + fn build_cli_for_args(registry: &SuiteRegistry, args: I) -> Command + where + I: IntoIterator, + T: AsRef, + { + let suite_cli_options = SuiteCliOptions::new(registry); + let args = os_args(args); + let native = native_registry(registry); + build_cli_for_args_with_options(registry, &suite_cli_options, &native, &args) + } + + fn command_from_matches( + registry: &SuiteRegistry, + matches: &ArgMatches, + ) -> Result { + let suite_cli_options = SuiteCliOptions::new(registry); + let native = native_registry(registry); + command_from_matches_with_options(&suite_cli_options, &native, matches) + } + + fn parse_command(registry: &SuiteRegistry, args: I) -> Result + where + I: IntoIterator, + T: AsRef, + { + let args = os_args(args); + let matches = build_cli_with_native(registry) + .try_get_matches_from(args) + .unwrap(); + + command_from_matches(registry, &matches) + } + + fn normalize_suite_option_aliases( + registry: &SuiteRegistry, + args: I, + ) -> Vec + where + I: IntoIterator, + T: AsRef, + { + let suite_cli_options = SuiteCliOptions::new(registry); + normalize_suite_option_aliases_with_options(&suite_cli_options, args) + } + + fn registry_with_suite_aliases() -> SuiteRegistry { + let tempdir = tempfile::tempdir().unwrap(); + + for (suite_name, option_name, value) in [ + ("left", "left-option", "left-value"), + ("right", "right-option", "right-value"), + ] { + let suite_dir = tempdir.path().join(suite_name); + fs::create_dir(&suite_dir).unwrap(); + fs::write( + suite_dir.join(format!("{suite_name}.suite")), + format!( + r#" +description = "{suite_name} benchmark suite" + +[[options]] +name = "{option_name}" +short = "x" +env = "SUITE_OPTION" +default = "{value}" +values = ["{value}"] +help = "Selects the suite option." +"# + ), + ) + .unwrap(); + } + + SuiteRegistry::discover(tempdir.path()).unwrap() + } + + #[test] + fn help_and_list_subcommands_reject_help_flags() { + for command in ["help", "list"] { + for help_flag in ["-h", "--help"] { + let err = build_cli_with_native(®istry()) + .try_get_matches_from(["benchmark_runner", command, help_flag]) + .unwrap_err(); + + assert_eq!(err.kind(), clap::error::ErrorKind::UnknownArgument); + } + } + } + + #[test] + fn help_and_list_subcommands_accept_no_options() { + for (command, expected) in + [("help", RunnerCommand::Help), ("list", RunnerCommand::List)] + { + let registry = registry(); + let actual = parse_command(®istry, ["benchmark_runner", command]).unwrap(); + + assert_eq!( + std::mem::discriminant(&actual), + std::mem::discriminant(&expected) + ); + } + } #[test] - fn list_rejects_unrecognized_options() { - let matches = - build_cli().try_get_matches_from(["benchmark_runner", "list", "--format"]); + fn list_rejects_suite_options_because_they_are_ignored() { + let matches = build_cli_with_native(®istry()).try_get_matches_from([ + "benchmark_runner", + "list", + "--format", + "csv", + ]); assert!(matches.is_err(), "{matches:?}"); } #[test] - fn help_mentions_list_command() { - let err = build_cli() - .try_get_matches_from(["benchmark_runner", "--help"]) - .unwrap_err(); - let help = err.to_string(); + fn help_mentions_query_id() { + let help = + build_cli_for_args(®istry(), ["benchmark_runner", "query", "--help"]) + .try_get_matches_from(["benchmark_runner", "query", "--help"]) + .unwrap_err() + .to_string(); + + assert!(help.contains("SUITE")); + assert!(help.contains("SQL benchmark suite name")); + assert!(help.contains("QUERY_ID")); + assert!(help.contains("Optional SQL query id")); + } + + #[test] + fn info_without_selector_reports_missing_benchmark_name() { + let registry = registry(); + let matches = build_cli_with_native(®istry) + .try_get_matches_from(["benchmark_runner", "info"]) + .unwrap(); + + let err = command_from_matches(®istry, &matches).unwrap_err(); + + assert!( + err.to_string() + .contains("Execution error: Missing benchmark name") + ); + } + + #[test] + fn sql_run_help_sections_use_run_command_metadata() { + let output = format_sql_run_help_sections_styled(®istry()).unwrap(); + + assert!(output.contains("Usage:")); + assert!(output.contains("SQL Benchmark Target:")); + assert!(output.contains("Runner Options:")); + assert!(output.contains("DataFusion Options:")); + assert!(output.contains("--save-baseline ")); + assert!(output.lines().any(|line| { + line.contains("--persist-results") + && line.contains( + "Persist benchmark results before the timed run (SQL suites only.)", + ) + })); + assert!(output.contains( + "Validate benchmark results before the timed run (SQL suites only.)" + )); + assert!(output.contains( + "Run selected benchmarks with Criterion instead of fixed iterations (SQL suites only.)" + )); + assert!(output.contains( + "Save Criterion results under a named baseline (SQL suites only.)" + )); + assert!(!output.contains( + "--persist-results\n Persist benchmark results before the timed run" + )); + } + + #[test] + fn info_native_selector_parses_as_native_info() { + let registry = registry(); + let command = + parse_command(®istry, ["benchmark_runner", "info", "hj"]).unwrap(); + + assert!(matches!(command, RunnerCommand::InfoNative(name) if name == "hj")); + } + + #[test] + fn info_native_selector_does_not_require_run_options() { + let registry = registry(); + let command = + parse_command(®istry, ["benchmark_runner", "info", "tpcds"]).unwrap(); + + assert!(matches!(command, RunnerCommand::InfoNative(name) if name == "tpcds")); + } + + #[test] + fn info_native_rejects_parent_suite_options() { + let registry = registry(); + let matches = build_cli_with_native(®istry) + .try_get_matches_from(["benchmark_runner", "info", "--format", "csv", "hj"]) + .unwrap(); + + let err = command_from_matches(®istry, &matches).unwrap_err(); + + assert!(err.to_string().contains("--format")); + assert!(err.to_string().contains("native benchmark")); + } + + #[test] + fn parses_suite_option_short_alias() { + let registry = registry(); + let args = normalize_suite_option_aliases( + ®istry, + [ + "benchmark_runner", + "info", + "tpch", + "15", + "-f", + "csv", + "-sf", + "10", + ], + ); + let matches = build_cli_with_native(®istry) + .try_get_matches_from(args) + .unwrap(); + let command = command_from_matches(®istry, &matches).unwrap(); + + let RunnerCommand::Info(args) = command else { + panic!("expected info command"); + }; + assert_eq!( + args.suite_options.get("format").map(String::as_str), + Some("csv") + ); + assert_eq!( + args.suite_options.get("scale-factor").map(String::as_str), + Some("10") + ); + } + + #[test] + fn normalizes_suite_option_aliases_for_selected_suite() { + for (suite_name, option_name, value) in [ + ("left", "left-option", "left-value"), + ("right", "right-option", "right-value"), + ] { + let registry = registry_with_suite_aliases(); + let args = normalize_suite_option_aliases( + ®istry, + ["benchmark_runner", "info", suite_name, "-x", value], + ); + let matches = build_cli_with_native(®istry) + .try_get_matches_from(args) + .unwrap(); + let command = command_from_matches(®istry, &matches).unwrap(); - assert!(help.contains("list")); + let RunnerCommand::Info(args) = command else { + panic!("expected info command"); + }; + assert_eq!( + args.suite_options.get(option_name).map(String::as_str), + Some(value) + ); + assert_eq!(args.suite_options.len(), 1); + } } } diff --git a/benchmarks/src/benchmark_runner/mod.rs b/benchmarks/src/benchmark_runner/mod.rs index 458e5f974c152..5bd85e3dc2af8 100644 --- a/benchmarks/src/benchmark_runner/mod.rs +++ b/benchmarks/src/benchmark_runner/mod.rs @@ -15,73 +15,229 @@ // specific language governing permissions and limitations // under the License. -//! Command-line inspection for SQL benchmark suites. +//! Command-line inspection for DataFusion SQL and native benchmarks. //! -//! This module backs the `benchmark_runner` binary. The initial command -//! surface lists discovered SQL benchmark suites from `.suite` files and -//! prints the top-level help. +//! This module backs the `benchmark_runner` binary. SQL suites and their +//! configurable options are discovered from text `.suite` files, individual +//! `.benchmark` files are parsed through [`SqlBenchmark`], and selected +//! benchmarks can be inspected as metadata or parsed SQL. Native benchmarks are +//! exposed through their existing Rust `RunOpt` metadata while they are migrated +//! to SQL suites. //! //! Common invocations: //! //! ```text //! cargo run --bin benchmark_runner -- --help //! cargo run --release --bin benchmark_runner -- list +//! cargo run --release --bin benchmark_runner -- info tpch +//! cargo run --release --bin benchmark_runner -- info tpch 15 +//! cargo run --release --bin benchmark_runner -- info hj +//! cargo run --release --bin benchmark_runner -- query tpch 15 //! ``` //! +//! In suite-based `query` commands, the final numeric argument is the benchmark +//! query id. For example, `15` targets `q15.benchmark`, and `1` or `01` targets +//! `q01.benchmark`. For `info`, the query id is optional: +//! `info tpch` reports suite metadata, while `info tpch 15` also reports +//! metadata for `q15.benchmark`. +//! //! The public entry point is [`run_cli`]. The submodules are kept private so //! the command-line flow remains the single supported API: //! -//! - `cli` builds the clap command tree and parses the selected command. -//! - `suite` loads `.suite` metadata and discovers benchmark query files. -//! - `output` formats colored `list` command output. +//! - `cli` builds the clap command tree, parses inspection flags, and handles +//! dynamic suite options from suite files. +//! - `native` describes Rust `RunOpt` benchmarks, filters entries shadowed by +//! SQL suites, and formats native benchmark metadata. +//! - `suite` loads `.suite` metadata, validates suite options, discovers +//! benchmark query files recursively, and maps chosen options to SQL +//! replacements. +//! - `selector` resolves a suite name and optional query id into the concrete +//! benchmark files to inspect. +//! - `output` formats colored `list`, `info`, and `query` command output. mod cli; +mod native; mod output; +mod selector; +mod style; mod suite; -use crate::benchmark_runner::cli::{RunnerCommand, build_cli, command_from_matches}; -use crate::benchmark_runner::output::format_suite_list_styled; +use crate::benchmark_runner::cli::{ + InspectArgs, RunnerCommand, SuiteCliOptions, build_cli_for_args_with_options, + command_from_matches_with_options, format_sql_run_help_sections_styled, + normalize_suite_option_aliases_with_options, +}; +use crate::benchmark_runner::native::{NativeBenchmark, NativeBenchmarks}; +use crate::benchmark_runner::output::{ + format_benchmark_list_styled, write_info_styled, write_queries_styled, + write_suite_info_styled, +}; +use crate::benchmark_runner::selector::{ResolvedBenchmarkTarget, ResolvedSuite}; use crate::benchmark_runner::suite::SuiteRegistry; +use crate::sql_benchmark::SqlBenchmark; +use clap::error::ErrorKind; use datafusion::error::Result; +use datafusion::prelude::SessionContext; use datafusion_common::DataFusionError; +use std::ffi::OsString; use std::io::Write; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; /// Runs the benchmark runner command-line flow for the provided argument list. /// -/// This discovers suite metadata, parses the help/list command, and dispatches -/// to the selected implementation. -pub fn run_cli(args: I) -> Result<()> +/// This discovers suite metadata, normalizes suite option aliases, builds the +/// appropriate clap command for the requested help context, and dispatches to +/// the selected `list`, `info`, or `query` implementation. +pub async fn run_cli(args: I) -> Result<()> where I: IntoIterator, - T: Clone + Into, + T: Into, { let benchmark_dir = default_benchmark_dir(); let registry = SuiteRegistry::discover(&benchmark_dir)?; - let mut cli = build_cli(); + let native = NativeBenchmarks::new(®istry); + let suite_cli_options = SuiteCliOptions::new(®istry); + let args = args.into_iter().map(Into::into).collect::>(); + let args = normalize_suite_option_aliases_with_options(&suite_cli_options, args); + + let root_help_requested = is_root_help_request(&args); + let mut cli = + build_cli_for_args_with_options(®istry, &suite_cli_options, &native, &args); let matches = match cli.try_get_matches_from_mut(args) { Ok(matches) => matches, - Err(e) if e.kind() == clap::error::ErrorKind::DisplayHelp => { - e.print()?; + Err(e) if e.kind() == ErrorKind::DisplayHelp => { + if root_help_requested { + print_root_help(&cli)?; + } else { + e.print()?; + } return Ok(()); } Err(e) => return Err(DataFusionError::External(Box::new(e))), }; - let command = command_from_matches(&matches)?; + + let command = + command_from_matches_with_options(&suite_cli_options, &native, &matches)?; match command { RunnerCommand::Help => { - cli.print_long_help()?; - println!(); + print_root_help(&cli)?; } RunnerCommand::List => { - print_styled(&format_suite_list_styled(®istry)?)?; + print_styled(&format_benchmark_list_styled(®istry, &native)?)?; + } + RunnerCommand::InfoNative(name) => { + let benchmark = native.get(&name).ok_or_else(|| { + DataFusionError::Configuration(format!( + "unknown native benchmark '{name}'" + )) + })?; + print_native_help(benchmark)?; + } + RunnerCommand::Info(args) => { + print_styled(&load_info_output(®istry, &benchmark_dir, &args).await?)?; + } + RunnerCommand::Query(args) => { + if native.contains(&args.selector) { + return Err(native_query_error(&args.selector)); + } + + let (_target, sql) = load_benchmark(®istry, &benchmark_dir, &args).await?; + + print_styled(&write_queries_styled(&sql)?)?; } } Ok(()) } +/// Returns whether the raw args request root help directly. +fn is_root_help_request(args: &[OsString]) -> bool { + matches!( + args.get(1).and_then(|arg| arg.to_str()), + Some("-h" | "--help") + ) && args.len() == 2 +} + +/// Prints root help with the explicit `help` command hidden from options. +fn print_root_help(cli: &clap::Command) -> Result<()> { + let mut cli = cli.clone().mut_arg("help", |arg| arg.hide(true)); + cli.print_long_help()?; + println!(); + Ok(()) +} + +/// Prints compact native benchmark help shared by native info selectors. +fn print_native_help(benchmark: &NativeBenchmark) -> Result<()> { + print_styled(&benchmark.info_text()?)?; + println!(); + Ok(()) +} + +/// Builds the error returned when `query` is used with a native benchmark. +fn native_query_error(selector: &str) -> DataFusionError { + DataFusionError::Configuration(format!( + "query is only supported for SQL benchmark suites; '{selector}' is a native benchmark" + )) +} + +/// Resolves a selector and parses the single benchmark it names. +/// +/// The `info` and `query` subcommands are parse-only operations: they load the +/// benchmark with suite replacement values but do not initialize benchmark SQL. +async fn load_benchmark( + registry: &SuiteRegistry, + benchmark_dir: &Path, + args: &InspectArgs, +) -> Result<(ResolvedBenchmarkTarget, SqlBenchmark)> { + let target = ResolvedBenchmarkTarget::resolve( + registry, + &args.selector, + args.query_id.as_deref(), + &args.suite_options, + )?; + let benchmark_path = target.single_benchmark_for_inspection()?.path.clone(); + let ctx = SessionContext::new(); + let sql = SqlBenchmark::new_with_replacements( + &ctx, + &benchmark_path, + benchmark_dir, + &target.replacement_values, + None, + ) + .await?; + + Ok((target, sql)) +} + +/// Formats SQL `info` output. +/// +/// A suite selector without `QUERY_ID` reports suite metadata, options, and +/// examples. A selector with `QUERY_ID` additionally parses the single +/// benchmark file and reports benchmark metadata. +async fn load_info_output( + registry: &SuiteRegistry, + benchmark_dir: &Path, + args: &InspectArgs, +) -> Result { + if args.query_id.is_none() { + let suite = registry.get(&args.selector).ok_or_else(|| { + DataFusionError::Configuration(format!( + "unknown benchmark suite '{}'", + args.selector + )) + })?; + suite.resolve_option_values(&args.suite_options)?; + let run_help_sections = format_sql_run_help_sections_styled(registry)?; + return write_suite_info_styled(&ResolvedSuite::from(suite), &run_help_sections); + } + + let (target, sql) = load_benchmark(registry, benchmark_dir, args).await?; + let run_help_sections = format_sql_run_help_sections_styled(registry)?; + + write_info_styled(target.suite.as_ref(), &sql, &run_help_sections) +} + /// Writes already styled output through `anstream` so ANSI color handling /// matches clap help output on supported terminals. fn print_styled(output: &str) -> Result<()> { @@ -96,9 +252,136 @@ fn print_styled(output: &str) -> Result<()> { /// benchmarks crate manifest directory. fn default_benchmark_dir() -> PathBuf { let repo_root_path = PathBuf::from("benchmarks/sql_benchmarks"); + if repo_root_path.exists() { repo_root_path } else { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("sql_benchmarks") } } + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + + fn strip_ansi(input: &str) -> String { + let mut output = String::new(); + let mut chars = input.chars(); + + while let Some(c) = chars.next() { + if c == '\x1b' { + for c in chars.by_ref() { + if c.is_ascii_alphabetic() { + break; + } + } + } else { + output.push(c); + } + } + + output + } + + #[test] + fn native_query_error_mentions_sql_only_query() { + let error = native_query_error("hj").to_string(); + + assert!(error.contains("query is only supported for SQL benchmark suites")); + assert!(error.contains("hj")); + } + + #[tokio::test] + async fn sql_info_without_query_id_formats_suite_info() { + let benchmark_dir = default_benchmark_dir(); + let registry = SuiteRegistry::discover(&benchmark_dir).unwrap(); + let args = InspectArgs { + selector: "tpch".to_string(), + query_id: None, + suite_options: BTreeMap::new(), + }; + + let output = load_info_output(®istry, &benchmark_dir, &args) + .await + .unwrap(); + let plain = strip_ansi(&output); + + assert!(plain.contains("Benchmark suite")); + assert!(plain.contains("suite")); + assert!(plain.contains("tpch")); + assert!( + plain.contains("Usage: benchmark_runner run [OPTIONS] [QUERY_ID]") + ); + assert!(output.contains("\u{1b}[1m\u{1b}[32mUsage:")); + assert!(plain.contains("SQL Benchmark Target:")); + assert!(output.contains("\u{1b}[1m\u{1b}[32mSQL Benchmark Target:")); + assert!(plain.contains("Runner Options:")); + assert!(output.contains("\u{1b}[1m\u{1b}[32mRunner Options:")); + assert!(plain.contains("--save-baseline ")); + assert!(plain.contains("DataFusion Options:")); + assert!(output.contains("\u{1b}[1m\u{1b}[32mDataFusion Options:")); + assert!(plain.contains("Suite Options (tpch):")); + assert!(plain.contains("--format")); + assert!(plain.contains("--scale-factor")); + assert!(plain.contains("Examples")); + assert!(!plain.contains(" value:")); + assert!(!plain.contains("q01.benchmark")); + } + + #[tokio::test] + async fn sql_info_with_query_id_includes_run_help_sections() { + let benchmark_dir = default_benchmark_dir(); + let registry = SuiteRegistry::discover(&benchmark_dir).unwrap(); + let args = InspectArgs { + selector: "tpch".to_string(), + query_id: Some("1".to_string()), + suite_options: BTreeMap::new(), + }; + + let output = load_info_output(®istry, &benchmark_dir, &args) + .await + .unwrap(); + let plain = strip_ansi(&output); + + assert!(plain.contains("Benchmark")); + assert!(plain.contains("q01.benchmark")); + assert!( + plain.contains("Usage: benchmark_runner run [OPTIONS] [QUERY_ID]") + ); + assert!(plain.contains("SQL Benchmark Target:")); + assert!(plain.contains("Runner Options:")); + assert!(plain.contains("DataFusion Options:")); + assert!(plain.contains("Suite Options (tpch):")); + assert!(plain.contains("Examples")); + assert!(!plain.contains(" value:")); + } + + #[tokio::test] + async fn sql_info_without_query_id_does_not_require_queries() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_dir = tempdir.path().join("empty"); + std::fs::create_dir(&suite_dir).unwrap(); + std::fs::write( + suite_dir.join("empty.suite"), + r#" +description = "Empty SQL benchmark suite" +"#, + ) + .unwrap(); + let registry = SuiteRegistry::discover(tempdir.path()).unwrap(); + let args = InspectArgs { + selector: "empty".to_string(), + query_id: None, + suite_options: BTreeMap::new(), + }; + + let output = load_info_output(®istry, tempdir.path(), &args) + .await + .unwrap(); + + assert!(output.contains("Benchmark suite")); + assert!(output.contains("empty")); + assert!(output.contains("Empty SQL benchmark suite")); + } +} diff --git a/benchmarks/src/benchmark_runner/native.rs b/benchmarks/src/benchmark_runner/native.rs new file mode 100644 index 0000000000000..1dddb1c07061a --- /dev/null +++ b/benchmarks/src/benchmark_runner/native.rs @@ -0,0 +1,362 @@ +// 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. + +//! Native benchmark metadata for `benchmark_runner`. +//! +//! Native benchmarks are the Rust `RunOpt`-based benchmarks historically +//! exposed through `dfbench.rs`. SQL suites take precedence, so native entries +//! whose selector matches a discovered `.suite` are hidden from this runner. + +use crate::benchmark_runner::output::write_options_section; +use crate::benchmark_runner::style::{HELP_STYLES, header}; +use crate::benchmark_runner::suite::SuiteRegistry; +use crate::{ + cancellation, clickbench, h2o, hj, imdb, nlj, smj, sort_pushdown, sort_tpch, tpcds, +}; +use clap::{Arg, ArgAction, Args, ColorChoice, Command}; +use datafusion::error::Result; +use std::collections::BTreeSet; +use std::fmt::Write as _; + +/// Metadata and clap wiring for one native Rust benchmark. +#[derive(Debug, Clone)] +pub(crate) struct NativeBenchmark { + kind: NativeBenchmarkKind, +} + +/// Builds commands and formatted help for one native benchmark. +impl NativeBenchmark { + /// Creates metadata for one native benchmark kind. + fn new(kind: NativeBenchmarkKind) -> Self { + Self { kind } + } + + /// Returns the command-line selector for this benchmark. + pub(crate) fn name(&self) -> &'static str { + self.kind.name() + } + + /// Returns the description shown by `list` and help output. + pub(crate) fn description(&self) -> &'static str { + self.kind.description() + } + + /// Builds the clap command used to render native benchmark options. + pub(crate) fn command(&self) -> Command { + configure_help( + self.kind.augment_args( + Command::new(self.name()) + .about(self.description()) + .styles(HELP_STYLES) + .color(ColorChoice::Always), + ), + ) + } + + /// Builds the hidden clap subcommand used to recognize native info selectors. + pub(crate) fn info_command(&self) -> Command { + Command::new(self.name()) + .about(self.description()) + .styles(HELP_STYLES) + .color(ColorChoice::Always) + } + + /// Renders compact native help text for the `info` command. + pub(crate) fn info_text(&self) -> Result { + let mut command = self + .command() + .bin_name(format!("benchmark_runner info {}", self.name())); + format_native_help(&mut command) + } +} + +/// Formats native benchmark help in the compact runner style. +fn format_native_help(command: &mut Command) -> Result { + let mut output = String::new(); + + if let Some(about) = command.get_about() { + writeln!(output, "{about}\n")?; + } + + writeln!( + output, + "{}", + header(command.render_usage().to_string().trim_end()) + )?; + writeln!(output)?; + write_options_section(&mut output, command)?; + + Ok(output) +} + +/// Applies benchmark-runner help behavior to a native `RunOpt` command. +fn configure_help(command: Command) -> Command { + hide_env_annotations(command).disable_help_flag(true).arg( + Arg::new("help") + .short('h') + .long("help") + .action(ArgAction::HelpShort) + .help("Print help"), + ) +} + +/// Hides environment variable annotations inherited from `RunOpt` arguments. +fn hide_env_annotations(mut command: Command) -> Command { + let arg_ids = command + .get_arguments() + .map(|arg| arg.get_id().to_string()) + .collect::>(); + + for id in arg_ids { + command = command.mut_arg(id, |arg| arg.hide_env(true)); + } + + command +} + +/// Visible native benchmark registry after excluding names shadowed by suites. +#[derive(Debug, Clone)] +pub(crate) struct NativeBenchmarks { + visible: Vec, +} + +/// Discovers and looks up native benchmarks exposed by the runner. +impl NativeBenchmarks { + /// Builds the visible native benchmark list for the current suite registry. + pub(crate) fn new(suites: &SuiteRegistry) -> Self { + let suite_names = suites + .suites() + .iter() + .map(|suite| suite.name.as_str()) + .collect::>(); + let visible = native_benchmarks() + .into_iter() + .filter(|benchmark| !suite_names.contains(benchmark.name())) + .collect(); + + Self { visible } + } + + /// Returns visible native benchmarks in display order. + pub(crate) fn visible(&self) -> &[NativeBenchmark] { + &self.visible + } + + /// Returns whether a native benchmark selector is visible. + pub(crate) fn contains(&self, name: &str) -> bool { + self.get(name).is_some() + } + + /// Finds one visible native benchmark by selector. + pub(crate) fn get(&self, name: &str) -> Option<&NativeBenchmark> { + self.visible + .iter() + .find(|benchmark| benchmark.name() == name) + } +} + +/// Enumerates the native `RunOpt` benchmarks that can be exposed. +#[derive(Debug, Clone, Copy)] +enum NativeBenchmarkKind { + Cancellation, + Clickbench, + H2o, + Hj, + Imdb, + Nlj, + Smj, + SortPushdown, + SortTpch, + Tpcds, +} + +const NATIVE_BENCHMARKS: &[NativeBenchmarkKind] = &[ + NativeBenchmarkKind::Cancellation, + NativeBenchmarkKind::Clickbench, + NativeBenchmarkKind::H2o, + NativeBenchmarkKind::Hj, + NativeBenchmarkKind::Imdb, + NativeBenchmarkKind::Nlj, + NativeBenchmarkKind::Smj, + NativeBenchmarkKind::SortPushdown, + NativeBenchmarkKind::SortTpch, + NativeBenchmarkKind::Tpcds, +]; + +/// Provides metadata and clap augmentation for native benchmark kinds. +impl NativeBenchmarkKind { + /// Returns the command-line selector for this native benchmark kind. + fn name(self) -> &'static str { + match self { + Self::Cancellation => "cancellation", + Self::Clickbench => "clickbench", + Self::H2o => "h2o", + Self::Hj => "hj", + Self::Imdb => "imdb", + Self::Nlj => "nlj", + Self::Smj => "smj", + Self::SortPushdown => "sort-pushdown", + Self::SortTpch => "sort-tpch", + Self::Tpcds => "tpcds", + } + } + + /// Returns the description shown in list and help output. + fn description(self) -> &'static str { + match self { + Self::Cancellation => "Test performance of cancelling queries", + Self::Clickbench => "Driver program to run the ClickBench benchmark", + Self::H2o => "Run the H2O benchmark", + Self::Hj => "Run the Hash Join benchmark", + Self::Imdb => "Run the Join Order Benchmark using the IMDB dataset", + Self::Nlj => "Run the Nested Loop Join benchmark", + Self::Smj => "Run the Sort Merge Join benchmark", + Self::SortPushdown => "Benchmark sort pushdown optimization", + Self::SortTpch => "Run end-to-end sort queries on a TPC-H dataset", + Self::Tpcds => "Run the TPC-DS benchmark", + } + } + + /// Adds this benchmark's `RunOpt` arguments to a clap command for display. + fn augment_args(self, command: Command) -> Command { + match self { + Self::Cancellation => augment::(command), + Self::Clickbench => augment::(command), + Self::H2o => augment::(command), + Self::Hj => augment::(command), + Self::Imdb => augment::(command), + Self::Nlj => augment::(command), + Self::Smj => augment::(command), + Self::SortPushdown => augment::(command), + Self::SortTpch => augment::(command), + Self::Tpcds => augment::(command), + } + } +} + +/// Builds metadata for every native benchmark kind before suite filtering. +fn native_benchmarks() -> Vec { + NATIVE_BENCHMARKS + .iter() + .copied() + .map(NativeBenchmark::new) + .collect() +} + +/// Adds one concrete native `RunOpt` type's clap arguments to a command. +fn augment(command: Command) -> Command +where + T: Args, +{ + T::augment_args(command) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::benchmark_runner::suite::SuiteRegistry; + use std::path::PathBuf; + + fn manifest_path(path: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) + } + + fn registry() -> SuiteRegistry { + SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap() + } + + #[test] + fn visible_native_benchmarks_include_hj() { + let suites = registry(); + let native = NativeBenchmarks::new(&suites); + + assert!(native.contains("hj")); + assert!(native.get("hj").is_some()); + } + + #[test] + fn visible_native_benchmarks_exclude_suite_names() { + let suites = registry(); + let native = NativeBenchmarks::new(&suites); + + assert!(suites.get("tpch").is_some()); + assert!(!native.contains("tpch")); + assert!(native.get("tpch").is_none()); + } + + #[test] + fn native_help_for_hj_contains_runopt_options() { + let suites = registry(); + let native = NativeBenchmarks::new(&suites); + let hj = native.get("hj").unwrap(); + let mut command = hj + .command() + .bin_name(format!("benchmark_runner info {}", hj.name())); + let help = command + .try_get_matches_from_mut(["benchmark_runner info hj", "--help"]) + .unwrap_err() + .to_string(); + + assert!(help.contains("Run the Hash Join benchmark")); + assert!(help.contains("--query")); + assert!(help.contains("--path")); + assert!(help.contains("--iterations")); + assert!(help.contains("Number of iterations of each test run [default: 3]")); + assert!(!help.contains("This micro-benchmark focuses")); + assert!(!help.contains("[env:")); + } + + #[test] + fn native_info_text_matches_condensed_help() { + let suites = registry(); + let native = NativeBenchmarks::new(&suites); + let hj = native.get("hj").unwrap(); + let info = hj.info_text().unwrap(); + + assert!(info.contains("Usage:")); + assert!(info.contains("\u{1b}[")); + assert!(info.contains("--iterations")); + assert!(info.contains("Number of iterations of each test run")); + assert!(info.lines().any(|line| { + line.contains("-i, --iterations ") + && line.contains("Number of iterations of each test run") + })); + assert!( + !info.contains("--iterations \n Number of iterations") + ); + assert!(info.contains("-d, --debug")); + assert!(!info.contains("-d, --debug ")); + assert!(!info.contains("This micro-benchmark focuses")); + assert!(!info.contains("[env:")); + } + + #[test] + fn native_info_text_preserves_required_args_in_compact_usage() { + let suites = registry(); + let native = NativeBenchmarks::new(&suites); + let tpcds = native.get("tpcds").unwrap(); + let info = tpcds.info_text().unwrap(); + + assert!(info.contains( + "Usage: benchmark_runner info tpcds [OPTIONS] --path --query_path " + )); + assert!(info.contains("-p, --path ")); + assert!(info.contains("-Q, --query_path ")); + assert!(info.contains("Path to data files")); + assert!(!info.contains("--path \n Path to data files")); + } +} diff --git a/benchmarks/src/benchmark_runner/output.rs b/benchmarks/src/benchmark_runner/output.rs index 9eb975311e29f..b69dfe95150be 100644 --- a/benchmarks/src/benchmark_runner/output.rs +++ b/benchmarks/src/benchmark_runner/output.rs @@ -17,79 +17,324 @@ //! Formatting helpers for human-readable benchmark runner output. //! -//! The runner intentionally uses the same colored style for `list` output as -//! clap uses for help text. +//! The runner intentionally uses the same colored style for `list`, `info`, +//! and `query` output as clap uses for help text. These functions only format +//! already parsed suite or benchmark metadata; they should not execute +//! benchmarks or gather runtime-derived details. -use crate::benchmark_runner::suite::{SuiteConfig, SuiteOption, SuiteRegistry}; -use clap::builder::styling::{AnsiColor, Style}; +use crate::benchmark_runner::native::{NativeBenchmark, NativeBenchmarks}; +use crate::benchmark_runner::selector::ResolvedSuite; +use crate::benchmark_runner::style::{header, label, literal, placeholder, value}; +use crate::benchmark_runner::suite::{ + SuiteConfig, SuiteExample, SuiteOption, SuiteRegistry, +}; +use crate::sql_benchmark::{QueryDirective, SqlBenchmark}; +use clap::{Arg, Command}; use datafusion_common::Result; -use std::fmt::{Display, Write as _}; +use std::fmt::Write as _; -/// Formats the `list` command output with suite summaries, query-id hints, and -/// configurable suite options. -pub fn format_suite_list_styled(registry: &SuiteRegistry) -> Result { +/// Formats the `list` command output with benchmark names and descriptions. +pub fn format_benchmark_list_styled( + registry: &SuiteRegistry, + native: &NativeBenchmarks, +) -> Result { let mut output = String::new(); for suite in registry.suites() { write_suite_list_entry(&mut output, suite)?; } + for benchmark in native.visible() { + write_native_list_entry(&mut output, benchmark)?; + } + Ok(output) } /// Writes one suite entry for the `list` command. fn write_suite_list_entry(output: &mut String, suite: &SuiteConfig) -> Result<()> { writeln!(output, "{}", header(&suite.name))?; - writeln!(output, " {}: {}", label("description"), suite.description)?; + writeln!(output, " {}", suite.description)?; + + Ok(()) +} - let queries = suite.discover_queries()?; +/// Writes one native benchmark entry for the `list` command. +fn write_native_list_entry( + output: &mut String, + benchmark: &NativeBenchmark, +) -> Result<()> { + writeln!(output, "{}", header(benchmark.name()))?; + writeln!(output, " {}", benchmark.description())?; - if let (Some(first), Some(last)) = (queries.first(), queries.last()) { - writeln!( - output, - " {}: {}-{} discovered under {} as {}", - label("query ids"), - value(first.id), - value(last.id), - suite.query_search_root().display(), - literal("qNN.benchmark") - )?; + Ok(()) +} + +/// Writes parse-time metadata for one benchmark and its containing suite. +/// +/// This intentionally avoids runtime-derived details such as row counts so the +/// `info` command can remain a cheap inspection operation. +pub fn write_info_styled( + suite: Option<&ResolvedSuite>, + benchmark: &SqlBenchmark, + run_help_sections: &str, +) -> Result { + let mut output = String::new(); + + writeln!(output, "{}", header("Benchmark"))?; + + if let Some(suite) = suite { + writeln!(output, " {}: {}", label("suite"), value(&suite.name))?; } - writeln!(output, " {}:", label("options"))?; - for option in &suite.options { - write_suite_list_option(output, option)?; + writeln!(output, " {}: {}", label("name"), value(benchmark.name()))?; + writeln!(output, " {}: {}", label("group"), value(benchmark.group()))?; + writeln!( + output, + " {}: {}", + label("subgroup"), + value(benchmark.subgroup()) + )?; + writeln!( + output, + " {}: {}", + label("file"), + benchmark.benchmark_path().display() + )?; + + if let Some(suite) = suite { + writeln!(output, " {}: {}", label("description"), suite.description)?; } - Ok(()) + if let Some(suite) = suite { + write_run_help_sections(&mut output, run_help_sections)?; + write_suite_sections(&mut output, suite)?; + } + + Ok(output) +} + +/// Writes suite-level metadata without requiring a single benchmark query file. +pub fn write_suite_info_styled( + suite: &ResolvedSuite, + run_help_sections: &str, +) -> Result { + let mut output = String::new(); + + writeln!(output, "{}", header("Benchmark suite"))?; + writeln!(output, " {}: {}", label("suite"), value(&suite.name))?; + writeln!(output, " {}: {}", label("description"), suite.description)?; + write_run_help_sections(&mut output, run_help_sections)?; + write_suite_sections(&mut output, suite)?; + + Ok(output) } -/// Writes one suite option summary for the `list` command. -fn write_suite_list_option(output: &mut String, option: &SuiteOption) -> Result<()> { - let values = option - .values +/// Writes one named clap help section in the runner's compact styled format. +pub fn write_help_section( + output: &mut String, + command: &Command, + heading: &str, +) -> Result<()> { + writeln!(output, "{}", header(format!("{heading}:")))?; + + let entries = command + .get_arguments() + .filter_map(|arg| { + if arg.is_hide_set() || arg.get_help_heading() != Some(heading) { + return None; + } + + let help = arg.get_help()?; + Some((arg_display(arg), help.to_string())) + }) + .collect::>(); + let flag_width = entries .iter() - .map(|v| { - if v == &option.default { - format!("{} ({})", value(v), label("default")) - } else { - value(v) + .map(|(flag, _help)| flag.len()) + .max() + .unwrap_or(0); + + for (flag, help) in entries { + write_help_entry(output, &flag, flag_width, &help)?; + } + + writeln!(output)?; + + Ok(()) +} + +/// Writes all visible options for a command in compact single-line form. +pub fn write_options_section(output: &mut String, command: &Command) -> Result<()> { + writeln!(output, "{}", header("Options:"))?; + + let entries = command + .get_arguments() + .filter_map(|arg| { + if arg.is_hide_set() { + return None; } + + Some((arg_display(arg), arg_help(arg)?)) }) + .collect::>(); + let flag_width = entries + .iter() + .map(|(flag, _help)| flag.len()) + .max() + .unwrap_or(0); + + for (flag, help) in entries { + write_help_entry(output, &flag, flag_width, &help)?; + } + + Ok(()) +} + +/// Builds help text for one clap argument, including defaults and values. +fn arg_help(arg: &Arg) -> Option { + let mut help = arg.get_help()?.to_string(); + + if arg.get_action().takes_values() { + let defaults = arg + .get_default_values() + .iter() + .map(|value| value.to_string_lossy()) + .collect::>(); + if !defaults.is_empty() { + write!(help, " [default: {}]", defaults.join(", ")).ok()?; + } + + let possible_values = arg + .get_possible_values() + .into_iter() + .map(|value| value.get_name().to_string()) + .collect::>(); + if !possible_values.is_empty() { + write!(help, " [possible values: {}]", possible_values.join(", ")).ok()?; + } + } + + Some(help) +} + +/// Builds the display form for one clap argument. +fn arg_display(arg: &Arg) -> String { + match arg.get_id().as_str() { + "selector" => return "".to_string(), + "query_id" => return "[QUERY_ID]".to_string(), + _ => {} + } + + let mut display = String::new(); + + if let Some(short) = arg.get_short() { + display.push('-'); + display.push(short); + } + + if let Some(long) = arg.get_long() { + if !display.is_empty() { + display.push_str(", "); + } + display.push_str("--"); + display.push_str(long); + } + + if arg.get_action().takes_values() + && let Some(value_names) = arg.get_value_names() + { + display.push(' '); + display.push_str(&format_value_names(value_names.iter())); + } + + display +} + +/// Formats one or more clap value names as angle-bracket placeholders. +fn format_value_names<'a>( + value_names: impl Iterator, +) -> String { + value_names + .map(|value| format!("<{value}>")) .collect::>() - .join(", "); + .join(" ") +} +/// Writes one aligned help entry. +fn write_help_entry( + output: &mut String, + flag: &str, + flag_width: usize, + help: &str, +) -> Result<()> { + let padding = " ".repeat(flag_width.saturating_sub(flag.len()) + 2); + + writeln!(output, " {}{padding}{help}", literal(flag))?; + + Ok(()) +} + +/// Appends preformatted run-help sections to info output. +fn write_run_help_sections(output: &mut String, run_help_sections: &str) -> Result<()> { + if !run_help_sections.is_empty() { + writeln!(output, "\n{}", run_help_sections.trim_end())?; + } + + Ok(()) +} + +/// Writes suite options and examples sections for info output. +fn write_suite_sections(output: &mut String, suite: &ResolvedSuite) -> Result<()> { + writeln!( + output, + "\n{}", + header(format!("Suite Options ({}):", suite.name)) + )?; + + for option in &suite.options { + write_suite_option_details(output, option)?; + } + + if suite.examples.is_empty() { + return Ok(()); + } + + writeln!(output, "\n{}", header("Examples"))?; + + for example in &suite.examples { + write_example_details(output, example)?; + } + + Ok(()) +} + +/// Writes one suite option for `info`, including default, replacement +/// variable, allowed values, and help text. +fn write_suite_option_details(output: &mut String, option: &SuiteOption) -> Result<()> { writeln!( output, - " {} {} {}", + " {}\n {}: {}\n {}: {}\n {}: {}\n {}: {}\n", literal(option_display(option)), - placeholder(""), - values + label("default"), + value(&option.default), + label("variable"), + placeholder(&option.env), + label("values"), + option + .values + .iter() + .map(value) + .collect::>() + .join(", "), + label("help"), + option.help )?; Ok(()) } +/// Formats one suite option with its short alias when present. fn option_display(option: &SuiteOption) -> String { match &option.short { Some(short) => format!("-{short}, --{}", option.name), @@ -97,39 +342,87 @@ fn option_display(option: &SuiteOption) -> String { } } -fn header(text: impl Display) -> String { - styled(AnsiColor::Green.on_default().bold(), text) -} +/// Writes one suite example command and its description for `info`. +fn write_example_details(output: &mut String, example: &SuiteExample) -> Result<()> { + writeln!( + output, + " {}\n {}\n", + literal(&example.command), + example.description + )?; -fn literal(text: impl Display) -> String { - styled(AnsiColor::Cyan.on_default().bold(), text) + Ok(()) } -fn placeholder(text: impl Display) -> String { - styled(AnsiColor::Cyan.on_default(), text) -} +/// Writes the SQL query text parsed from a benchmark file by directive. +/// +/// Only directives that contain queries are emitted, and each query is numbered +/// within its directive for easier inspection. +pub fn write_queries_styled(benchmark: &SqlBenchmark) -> Result { + let mut output = String::new(); -fn value(text: impl Display) -> String { - styled(AnsiColor::Green.on_default(), text) -} + for directive in [ + QueryDirective::Init, + QueryDirective::Load, + QueryDirective::Run, + QueryDirective::Cleanup, + ] { + if let Some(queries) = benchmark.queries().get(&directive) { + writeln!(output, "{}", header(directive.as_str()))?; -fn label(text: impl Display) -> String { - styled(Style::new().bold(), text) -} + for (idx, query) in queries.iter().enumerate() { + writeln!( + output, + "{}\n{}", + literal(format!("-- query {}", idx + 1)), + query.trim() + )?; + } + + output.push('\n'); + } + } -fn styled(style: Style, text: impl Display) -> String { - format!("{style}{text}{style:#}") + Ok(output) } #[cfg(test)] mod tests { use super::*; + use crate::benchmark_runner::native::NativeBenchmarks; + use datafusion::prelude::SessionContext; + use std::cell::OnceCell; + use std::collections::BTreeMap; use std::path::PathBuf; fn manifest_path(path: &str) -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) } + fn example_suite() -> ResolvedSuite { + ResolvedSuite::from(&SuiteConfig { + name: "example-suite".to_string(), + description: "Example suite description".to_string(), + query_pattern: "q{QUERY_ID_PADDED}.benchmark".to_string(), + path_replacements: BTreeMap::new(), + options: vec![], + examples: vec![], + suite_path: PathBuf::from("example.suite"), + suite_dir: PathBuf::from("."), + query_cache: OnceCell::new(), + }) + } + + fn assert_in_order(output: &str, expected: &[&str]) { + let mut offset = 0; + for expected in expected { + let Some(index) = output[offset..].find(expected) else { + panic!("expected '{expected}' after byte {offset} in:\n{output}"); + }; + offset += index + expected.len(); + } + } + fn strip_ansi(input: &str) -> String { let mut output = String::new(); let mut chars = input.chars(); @@ -148,26 +441,188 @@ mod tests { } #[test] - fn list_output_mentions_tpch_options() { + fn list_output_only_includes_benchmark_names_and_descriptions() { let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); - let output = strip_ansi(&format_suite_list_styled(®istry).unwrap()); + let native = NativeBenchmarks::new(®istry); + let output = + strip_ansi(&format_benchmark_list_styled(®istry, &native).unwrap()); - assert!(output.contains("tpch\n description: TPC-H SQL benchmarks")); - assert!(output.contains("-f, --format parquet (default), csv, mem")); - assert!(output.contains("-sf, --scale-factor 1 (default), 10")); + assert!(output.contains("tpch\n TPC-H SQL benchmarks")); + assert!(output.contains("hj\n Run the Hash Join benchmark")); + assert!(!output.contains("description:")); + assert!(!output.contains("options:")); + assert!(!output.contains("examples:")); } #[test] fn styled_list_output_includes_ansi_sequences() { let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); - let output = format_suite_list_styled(®istry).unwrap(); + let native = NativeBenchmarks::new(®istry); + let output = format_benchmark_list_styled(®istry, &native).unwrap(); assert!(output.contains("\u{1b}[")); assert!(output.contains("tpch")); - assert!(output.contains("-f")); - assert!(output.contains("--format")); - assert!(output.contains("-sf")); - assert!(output.contains("--scale-factor")); - assert!(output.contains("")); + } + + #[test] + fn list_output_excludes_native_benchmarks_shadowed_by_sql_suites() { + let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); + let native = NativeBenchmarks::new(®istry); + let output = + strip_ansi(&format_benchmark_list_styled(®istry, &native).unwrap()); + let lines = output.lines().collect::>(); + + assert!(output.contains("tpch\n TPC-H SQL benchmarks")); + assert_eq!(lines.iter().filter(|line| **line == "tpch").count(), 1); + } + + #[tokio::test] + async fn info_output_is_parse_only() { + let ctx = SessionContext::new(); + let dir = tempfile::tempdir().expect("temp dir should be created"); + let benchmark_path = dir.path().join("q01.benchmark"); + std::fs::write( + &benchmark_path, + "name Q01\ngroup example\nsubgroup default\nrun\nSELECT 1\n", + ) + .expect("benchmark file should be written"); + let benchmark = SqlBenchmark::new(&ctx, &benchmark_path, dir.path(), None) + .await + .expect("benchmark should parse"); + + let suite = example_suite(); + let output = write_info_styled( + Some(&suite), + &benchmark, + "Usage: benchmark_runner run [OPTIONS] [QUERY_ID]", + ) + .unwrap(); + + assert!(output.contains("Benchmark")); + assert_in_order( + &output, + &[ + "suite", + "example-suite", + "name", + "Q01", + "group", + "example", + "subgroup", + "default", + "file", + "q01.benchmark", + "description", + "Example suite description", + ], + ); + assert!(!output.contains("query id")); + assert!(!output.contains("Parsed benchmark")); + assert!(!output.contains("initialized")); + assert!(!output.contains("expected plan strings")); + assert!(!output.contains("query directives")); + assert!(!output.contains("assert queries")); + assert!(!output.contains("result queries")); + assert!(!output.contains("row count")); + assert!(!output.contains("rows returned")); + } + + #[tokio::test] + async fn styled_info_output_includes_ansi_sequences() { + let ctx = SessionContext::new(); + let dir = tempfile::tempdir().expect("temp dir should be created"); + let benchmark_path = dir.path().join("q01.benchmark"); + std::fs::write( + &benchmark_path, + "name Q01\ngroup example\nsubgroup default\nrun\nSELECT 1\n", + ) + .expect("benchmark file should be written"); + let benchmark = SqlBenchmark::new(&ctx, &benchmark_path, dir.path(), None) + .await + .expect("benchmark should parse"); + + let suite = example_suite(); + let output = write_info_styled( + Some(&suite), + &benchmark, + "Usage: benchmark_runner run [OPTIONS] [QUERY_ID]", + ) + .unwrap(); + + assert!(output.contains("\u{1b}[")); + assert!(output.contains("Benchmark")); + assert!(output.contains("suite")); + assert!(output.contains("name")); + assert!(output.contains("Suite Options (example-suite):")); + assert!(output.contains("group")); + assert!(output.contains("subgroup")); + assert!(output.contains("file")); + assert!(output.contains("description")); + assert!(!output.contains("query id")); + assert!(!output.contains("Parsed benchmark")); + assert!(!output.contains("initialized")); + assert!(!output.contains("expected plan strings")); + assert!(!output.contains("query directives")); + assert!(!output.contains("assert queries")); + assert!(!output.contains("result queries")); + } + + #[test] + fn suite_info_omits_empty_examples_section() { + let suite = example_suite(); + let output = write_suite_info_styled(&suite, "").unwrap(); + + assert!(output.contains("Benchmark suite")); + assert!(!output.contains("Examples")); + } + + #[tokio::test] + async fn query_output_formats_loaded_query_directives() { + let ctx = SessionContext::new(); + let dir = tempfile::tempdir().expect("temp dir should be created"); + let benchmark_path = dir.path().join("q01.benchmark"); + std::fs::write( + &benchmark_path, + "init\nCREATE TEMP TABLE t AS SELECT 1 AS value;\n\nrun\nSELECT value FROM t\n\ncleanup\nDROP TABLE t\n", + ) + .expect("benchmark file should be written"); + let benchmark = SqlBenchmark::new(&ctx, &benchmark_path, dir.path(), None) + .await + .expect("benchmark should parse"); + + let output = write_queries_styled(&benchmark).unwrap(); + + assert_in_order( + &output, + &[ + "init", + "-- query 1", + "CREATE TEMP TABLE t AS SELECT 1 AS value;", + "run", + "-- query 1", + "SELECT value FROM t", + "cleanup", + "-- query 1", + "DROP TABLE t", + ], + ); + } + + #[tokio::test] + async fn styled_query_output_includes_ansi_sequences() { + let ctx = SessionContext::new(); + let dir = tempfile::tempdir().expect("temp dir should be created"); + let benchmark_path = dir.path().join("q01.benchmark"); + std::fs::write(&benchmark_path, "run\nSELECT 1\n") + .expect("benchmark file should be written"); + let benchmark = SqlBenchmark::new(&ctx, &benchmark_path, dir.path(), None) + .await + .expect("benchmark should parse"); + + let output = write_queries_styled(&benchmark).unwrap(); + + assert!(output.contains("\u{1b}[")); + assert!(output.contains("run")); + assert!(output.contains("SELECT 1")); } } diff --git a/benchmarks/src/benchmark_runner/selector.rs b/benchmarks/src/benchmark_runner/selector.rs new file mode 100644 index 0000000000000..006990d993c80 --- /dev/null +++ b/benchmarks/src/benchmark_runner/selector.rs @@ -0,0 +1,356 @@ +// 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. + +//! Benchmark target resolution. +//! +//! Selectors name a suite, such as `tpch`. Optional query IDs are normalized +//! so values like `1`, `01`, and `001` resolve through the suite's query file +//! pattern. + +use crate::benchmark_runner::suite::{ + SuiteConfig, SuiteExample, SuiteOption, SuiteRegistry, +}; +use datafusion_common::{DataFusionError, Result}; +use std::collections::{BTreeMap, HashMap}; +use std::path::PathBuf; + +/// Concrete benchmark file selected by suite and optional query id. +#[derive(Debug, Clone)] +pub struct ResolvedBenchmark { + /// Benchmark file selected for inspection or execution. + pub path: PathBuf, + /// Numeric query id when it can be inferred from a `qNN.benchmark` file. + #[cfg_attr( + not(test), + expect(dead_code, reason = "used by the stacked run command branch") + )] + pub query_id: Option, + /// Human-readable benchmark label used in run output and JSON. + #[cfg_attr( + not(test), + expect(dead_code, reason = "used by the stacked run command branch") + )] + pub label: String, +} + +/// Suite metadata needed after selector resolution. +/// +/// This intentionally carries only display and suite-option metadata. It does +/// not include the full suite config, filesystem paths, or discovery caches. +#[derive(Debug, Clone)] +pub struct ResolvedSuite { + /// Suite selector used on the command line. + pub name: String, + /// Human-readable suite description. + pub description: String, + /// Suite-specific CLI options shown by `info`. + pub options: Vec, + /// Example commands shown by `info`. + pub examples: Vec, +} + +/// Copies display metadata from a suite config into a resolved suite value. +impl From<&SuiteConfig> for ResolvedSuite { + fn from(suite: &SuiteConfig) -> Self { + Self { + name: suite.name.clone(), + description: suite.description.clone(), + options: suite.options.clone(), + examples: suite.examples.clone(), + } + } +} + +/// Fully resolved SQL benchmark target ready for inspection or execution. +#[derive(Debug, Clone)] +pub struct ResolvedBenchmarkTarget { + /// Containing suite metadata, when the target came from or is inside a + /// discovered suite. + pub suite: Option, + /// Final suite option values after applying defaults and CLI overrides. + #[cfg_attr( + not(test), + expect(dead_code, reason = "used by the stacked run command branch") + )] + pub option_values: BTreeMap, + /// SQL replacement values passed to `SqlBenchmark`. + pub replacement_values: HashMap, + /// Concrete benchmark files selected for inspection or execution. + pub benchmarks: Vec, +} + +/// Resolves suite selectors, query ids, and suite option values. +impl ResolvedBenchmarkTarget { + /// Resolves a user selector into the concrete benchmark files to inspect or run. + /// + /// Suite selectors apply validated suite options and optional query-id + /// filtering. + pub fn resolve( + registry: &SuiteRegistry, + selector: &str, + query_id: Option<&str>, + suite_options: &BTreeMap, + ) -> Result { + let suite = registry.get(selector).ok_or_else(|| { + DataFusionError::Configuration(format!( + "unknown benchmark suite '{selector}'" + )) + })?; + let option_values = suite.resolve_option_values(suite_options)?; + let replacement_values = suite.option_replacements(&option_values); + let queries = suite.discover_queries()?; + let benchmarks = if let Some(query_id) = query_id { + let id = parse_query_id(query_id)?; + let query = queries + .into_iter() + .find(|query| query.id == id) + .ok_or_else(|| { + DataFusionError::Configuration(format!( + "suite '{}' has no QUERY_ID {}", + suite.name, id + )) + })?; + vec![ResolvedBenchmark { + path: query.path, + query_id: Some(id), + label: format!("{}/q{id:02}", suite.name), + }] + } else { + queries + .into_iter() + .map(|query| ResolvedBenchmark { + label: format!("{}/q{:02}", suite.name, query.id), + path: query.path, + query_id: Some(query.id), + }) + .collect() + }; + + if benchmarks.is_empty() { + return Err(DataFusionError::Configuration(format!( + "suite '{}' has no benchmark files", + suite.name + ))); + } + + Ok(Self { + suite: Some(ResolvedSuite::from(suite)), + option_values, + replacement_values, + benchmarks, + }) + } + + /// Returns the single resolved benchmark required by parse-only inspection + /// commands. + pub fn single_benchmark_for_inspection(&self) -> Result<&ResolvedBenchmark> { + match self.benchmarks.as_slice() { + [benchmark] => Ok(benchmark), + _ => Err(DataFusionError::Configuration( + "query requires exactly one benchmark; provide QUERY_ID".to_string(), + )), + } + } +} + +/// Parses a user-provided query id while accepting zero-padded forms. +fn parse_query_id(input: &str) -> Result { + let input = input.trim(); + + input.parse::().map_err(|_| { + DataFusionError::Configuration( + "QUERY_ID must be a numeric query id such as 1, 01, or 15".to_string(), + ) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::benchmark_runner::suite::SuiteRegistry; + use std::collections::BTreeMap; + use std::path::PathBuf; + + fn manifest_path(path: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) + } + + fn registry() -> SuiteRegistry { + SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap() + } + + #[test] + fn parses_query_ids() { + assert_eq!(parse_query_id("1").unwrap(), 1); + assert_eq!(parse_query_id("01").unwrap(), 1); + assert_eq!(parse_query_id("001").unwrap(), 1); + assert_eq!(parse_query_id("15").unwrap(), 15); + + let err = parse_query_id("").unwrap_err().to_string(); + assert!(err.contains("QUERY_ID")); + + let err = parse_query_id("q15").unwrap_err().to_string(); + assert!(err.contains("QUERY_ID")); + + assert_eq!(parse_query_id("0").unwrap(), 0); + assert_eq!(parse_query_id("00").unwrap(), 0); + assert_eq!(parse_query_id("000").unwrap(), 0); + } + + #[test] + fn resolves_tpch_query_id_to_q15() { + let registry = registry(); + let resolved = ResolvedBenchmarkTarget::resolve( + ®istry, + "tpch", + Some("15"), + &BTreeMap::new(), + ) + .unwrap(); + + assert_eq!(resolved.benchmarks.len(), 1); + assert!(resolved.benchmarks[0].path.ends_with("q15.benchmark")); + assert_eq!(resolved.benchmarks[0].query_id, Some(15)); + assert_eq!(resolved.benchmarks[0].label, "tpch/q15"); + } + + #[test] + fn selected_suite_options_override_defaults() { + let registry = registry(); + let mut options = BTreeMap::new(); + options.insert("format".to_string(), "csv".to_string()); + options.insert("scale-factor".to_string(), "10".to_string()); + + let resolved = + ResolvedBenchmarkTarget::resolve(®istry, "tpch", Some("1"), &options) + .unwrap(); + + assert_eq!( + resolved.replacement_values.get("TPCH_FILE_TYPE").unwrap(), + "csv" + ); + assert_eq!(resolved.replacement_values.get("BENCH_SIZE").unwrap(), "10"); + assert_eq!( + resolved.option_values.get("format").map(String::as_str), + Some("csv") + ); + } + + #[test] + fn suite_without_query_id_resolves_all_queries_in_numeric_order() { + let registry = registry(); + let resolved = + ResolvedBenchmarkTarget::resolve(®istry, "tpch", None, &BTreeMap::new()) + .unwrap(); + + let ids = resolved + .benchmarks + .iter() + .map(|benchmark| benchmark.query_id.unwrap()) + .collect::>(); + assert_eq!(ids.first(), Some(&1)); + assert_eq!(ids.last(), Some(&22)); + assert_eq!( + resolved.replacement_values.get("TPCH_FILE_TYPE").unwrap(), + "parquet" + ); + assert_eq!(resolved.replacement_values.get("BENCH_SIZE").unwrap(), "1"); + } + + #[test] + fn query_id_resolution_uses_suite_query_pattern() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_dir = tempdir.path().join("custom"); + std::fs::create_dir(&suite_dir).unwrap(); + std::fs::write(suite_dir.join("query-7.benchmark"), "run\nSELECT 7\n").unwrap(); + std::fs::write( + suite_dir.join("custom.suite"), + r#" +description = "Custom query pattern" +query_pattern = "query-{QUERY_ID}.benchmark" +"#, + ) + .unwrap(); + let registry = SuiteRegistry::discover(tempdir.path()).unwrap(); + + let resolved = ResolvedBenchmarkTarget::resolve( + ®istry, + "custom", + Some("7"), + &BTreeMap::new(), + ) + .unwrap(); + + assert_eq!(resolved.benchmarks.len(), 1); + assert_eq!(resolved.benchmarks[0].query_id, Some(7)); + assert!(resolved.benchmarks[0].path.ends_with("query-7.benchmark")); + } + + #[test] + fn query_id_resolution_uses_discovered_numeric_id() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_dir = tempdir.path().join("custom"); + std::fs::create_dir(&suite_dir).unwrap(); + std::fs::write(suite_dir.join("q001.benchmark"), "run\nSELECT 1\n").unwrap(); + std::fs::write( + suite_dir.join("custom.suite"), + r#" +description = "Custom padded query id" +"#, + ) + .unwrap(); + let registry = SuiteRegistry::discover(tempdir.path()).unwrap(); + + let resolved = ResolvedBenchmarkTarget::resolve( + ®istry, + "custom", + Some("1"), + &BTreeMap::new(), + ) + .unwrap(); + + assert_eq!(resolved.benchmarks.len(), 1); + assert_eq!(resolved.benchmarks[0].query_id, Some(1)); + assert!(resolved.benchmarks[0].path.ends_with("q001.benchmark")); + } + + #[test] + fn inspection_requires_single_benchmark() { + let registry = registry(); + let resolved = + ResolvedBenchmarkTarget::resolve(®istry, "tpch", None, &BTreeMap::new()) + .unwrap(); + + let err = resolved.single_benchmark_for_inspection().unwrap_err(); + assert!(err.to_string().contains("QUERY_ID")); + } + + #[test] + fn benchmark_file_path_is_rejected() { + let registry = registry(); + let path = manifest_path("sql_benchmarks/tpch/benchmarks/q15.benchmark"); + let err = ResolvedBenchmarkTarget::resolve( + ®istry, + path.to_str().unwrap(), + None, + &BTreeMap::new(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("unknown benchmark suite")); + } +} diff --git a/benchmarks/src/benchmark_runner/style.rs b/benchmarks/src/benchmark_runner/style.rs new file mode 100644 index 0000000000000..d8e46cc972e50 --- /dev/null +++ b/benchmarks/src/benchmark_runner/style.rs @@ -0,0 +1,58 @@ +// 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. + +//! Shared styling for benchmark runner clap help and formatted output. + +use clap::builder::styling::{AnsiColor, Style, Styles}; +use std::fmt::Display; + +/// Shared clap color palette for benchmark runner help output. +pub(crate) const HELP_STYLES: Styles = Styles::styled() + .header(AnsiColor::Green.on_default().bold()) + .usage(AnsiColor::Green.on_default().bold()) + .literal(AnsiColor::Cyan.on_default().bold()) + .placeholder(AnsiColor::Cyan.on_default()); + +/// Formats section headings with the runner's heading color. +pub(crate) fn header(text: impl Display) -> String { + styled(AnsiColor::Green.on_default().bold(), text) +} + +/// Formats command-line flags, commands, and other literal text. +pub(crate) fn literal(text: impl Display) -> String { + styled(AnsiColor::Cyan.on_default().bold(), text) +} + +/// Formats placeholder names such as SQL replacement variables. +pub(crate) fn placeholder(text: impl Display) -> String { + styled(AnsiColor::Cyan.on_default(), text) +} + +/// Formats selected values in informational output. +pub(crate) fn value(text: impl Display) -> String { + styled(AnsiColor::Green.on_default(), text) +} + +/// Formats field labels in informational output. +pub(crate) fn label(text: impl Display) -> String { + styled(Style::new().bold(), text) +} + +/// Applies one clap style and resets it after the text. +fn styled(style: Style, text: impl Display) -> String { + format!("{style}{text}{style:#}") +} diff --git a/benchmarks/src/benchmark_runner/suite.rs b/benchmarks/src/benchmark_runner/suite.rs index fe5b3339b1643..a78264138ae92 100644 --- a/benchmarks/src/benchmark_runner/suite.rs +++ b/benchmarks/src/benchmark_runner/suite.rs @@ -18,20 +18,46 @@ //! Suite-file loading and validation. //! //! A suite is described by a text `.suite` file that declares which options -//! the runner should display. Query discovery recursively scans from the suite +//! the runner should expose and how selected option values map to SQL +//! replacement variables. Query discovery recursively scans from the suite //! file's directory for `qNN.benchmark` files. Discovered queries are cached -//! lazily because they are reused by listing during a single CLI run. +//! lazily because they are reused by listing and target resolution during a +//! single CLI run. use datafusion_common::{DataFusionError, Result}; use serde::{Deserialize, Serialize}; use std::cell::OnceCell; -use std::collections::HashSet; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fs::{self, DirEntry}; use std::path::{Path, PathBuf}; +const DEFAULT_QUERY_PATTERN: &str = "q{QUERY_ID_PADDED}.benchmark"; +const QUERY_ID_TOKEN: &str = "{QUERY_ID}"; +const QUERY_ID_PADDED_TOKEN: &str = "{QUERY_ID_PADDED}"; +const RESERVED_CLI_OPTION_NAMES: &[&str] = &[ + "batch-size", + "criterion", + "debug", + "help", + "iterations", + "mem-pool-type", + "memory-limit", + "output", + "partitions", + "persist-results", + "query_id", + "save-baseline", + "selector", + "simulate-latency", + "sort-spill-reservation-bytes", + "validate-results", +]; +const RESERVED_CLI_SHORT_ALIASES: &[&str] = &["d", "h", "i", "n", "o", "p", "s"]; + +/// Discovered query benchmark file belonging to one suite. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SuiteQuery { - /// Numeric query id parsed from a `qNN.benchmark` file. + /// Numeric query id parsed from the configured query file pattern. pub id: usize, /// File name as it appears on disk, for example `q01.benchmark`. pub file_name: String, @@ -46,13 +72,23 @@ pub struct SuiteQuery { /// caching during a single runner invocation. #[derive(Debug, Deserialize, Serialize)] pub struct SuiteConfig { - /// Suite selector used on the command line, such as `tpch`. + /// Suite selector inferred from the `.suite` file name, such as `tpch`. + #[serde(skip)] pub name: String, - /// Human-readable suite description shown by `list`. + /// Human-readable suite description shown by `list` and `info`. pub description: String, - /// Suite-specific options shown by `list`. + /// Query file naming pattern. Defaults to `q{QUERY_ID_PADDED}.benchmark`. + #[serde(default = "default_query_pattern")] + pub query_pattern: String, + /// Path-like SQL replacement values resolved relative to the suite file. + #[serde(default)] + pub path_replacements: BTreeMap, + /// Suite-specific CLI options and their SQL replacement variables. #[serde(default)] pub options: Vec, + /// Example commands shown by `list` and `info`. + #[serde(default)] + pub examples: Vec, /// Path to the `.suite` file that produced this config. #[serde(skip)] pub suite_path: PathBuf, @@ -64,12 +100,36 @@ pub struct SuiteConfig { pub(crate) query_cache: OnceCell>, } +/// Raw TOML shape accepted from a `.suite` file. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct SuiteFileConfig { + /// Human-readable suite description shown by list and info. + description: String, + /// Query file naming pattern used for discovery and query-id lookup. + #[serde(default = "default_query_pattern")] + query_pattern: String, + /// SQL replacement paths resolved relative to the suite file. + #[serde(default)] + path_replacements: BTreeMap, + /// Suite-defined command-line options and their replacement variables. + #[serde(default)] + options: Vec, + /// Example commands shown by suite info output. + #[serde(default)] + examples: Vec, +} + +/// Clones suite metadata while intentionally resetting the lazy query cache. impl Clone for SuiteConfig { fn clone(&self) -> Self { Self { name: self.name.clone(), description: self.description.clone(), + query_pattern: self.query_pattern.clone(), + path_replacements: self.path_replacements.clone(), options: self.options.clone(), + examples: self.examples.clone(), suite_path: self.suite_path.clone(), suite_dir: self.suite_dir.clone(), query_cache: OnceCell::new(), @@ -77,14 +137,18 @@ impl Clone for SuiteConfig { } } +/// Suite-defined command-line option and its SQL replacement behavior. #[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] pub struct SuiteOption { - /// Long option name, without the leading `--`. + /// Long CLI option name, without the leading `--`. pub name: String, - /// Optional short alias, without the leading `-`. + /// Optional short CLI alias, without the leading `-`. #[serde(default)] pub short: Option, - /// Default option value. + /// SQL replacement variable set when this option is selected. + pub env: String, + /// Default option value used when the CLI does not provide one. pub default: String, /// Allowed option values. pub values: Vec, @@ -92,41 +156,70 @@ pub struct SuiteOption { pub help: String, } +/// Example command configured in a `.suite` file. +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct SuiteExample { + /// Example command to display. + pub command: String, + /// Explanation of what the example command runs. + pub description: String, +} + /// Discovered suite metadata, sorted by suite name. #[derive(Debug, Clone)] pub struct SuiteRegistry { suites: Vec, } +/// Loads, validates, and resolves metadata for one suite file. impl SuiteConfig { /// Loads, parses, and validates one `.suite` file. /// /// The suite file path and containing directory are stored on the returned - /// config so later discovery can be resolved relative to the suite file. + /// config so later discovery and path replacements can be resolved relative + /// to the suite file. pub fn from_file(path: impl AsRef) -> Result { let suite_path = path.as_ref().to_path_buf(); + let suite_name = suite_name_from_path(&suite_path)?; let contents = fs::read_to_string(&suite_path).map_err(|e| { DataFusionError::External( format!("failed to read suite file {}: {e}", suite_path.display()).into(), ) })?; - let mut suite: Self = toml::from_str(&contents).map_err(|e| { + let raw: SuiteFileConfig = toml::from_str(&contents).map_err(|e| { DataFusionError::External( format!("failed to parse suite file {}: {e}", suite_path.display()) .into(), ) })?; - suite.suite_dir = suite_path + let suite_dir = suite_path .parent() .map(Path::to_path_buf) .unwrap_or_else(|| PathBuf::from(".")); - suite.suite_path = suite_path; + let suite = Self { + name: suite_name, + description: raw.description, + query_pattern: raw.query_pattern, + path_replacements: raw.path_replacements, + options: raw.options, + examples: raw.examples, + suite_path, + suite_dir, + query_cache: OnceCell::new(), + }; + suite.validate()?; Ok(suite) } + /// Returns a configured suite option by command-line name. + pub fn option(&self, name: &str) -> Option<&SuiteOption> { + self.options.iter().find(|option| option.name == name) + } + /// Returns the directory recursively searched for query benchmark files. pub fn query_search_root(&self) -> &Path { &self.suite_dir @@ -135,8 +228,8 @@ impl SuiteConfig { /// Discovers and caches the suite's benchmark query files. /// /// Query files are found by recursively scanning from the suite directory, - /// accepting only `qNN.benchmark` files, sorting by numeric query id, and - /// rejecting duplicate ids. + /// accepting only files matching this suite's `query_pattern`, sorting by + /// numeric query id, and rejecting duplicate ids. pub fn discover_queries(&self) -> Result> { if let Some(queries) = self.query_cache.get() { return Ok(queries.clone()); @@ -150,8 +243,9 @@ impl SuiteConfig { /// Performs uncached query discovery and duplicate-id validation. fn scan_queries(&self) -> Result> { let mut queries = Vec::new(); + let query_pattern = QueryPattern::new(&self.query_pattern)?; - self.scan_query_dir(self.query_search_root(), &mut queries)?; + self.scan_query_dir(self.query_search_root(), &query_pattern, &mut queries)?; queries.sort_by(|left, right| { left.id .cmp(&right.id) @@ -178,7 +272,12 @@ impl SuiteConfig { /// Recursively scans a directory and appends valid query benchmark files to /// the provided collection. - fn scan_query_dir(&self, dir: &Path, queries: &mut Vec) -> Result<()> { + fn scan_query_dir( + &self, + dir: &Path, + query_pattern: &QueryPattern, + queries: &mut Vec, + ) -> Result<()> { let mut entries = read_dir_entries(dir, "benchmark query directory")?; entries.sort_by_key(|entry| entry.file_name()); @@ -195,7 +294,7 @@ impl SuiteConfig { })?; if file_type.is_dir() { - self.scan_query_dir(&path, queries)?; + self.scan_query_dir(&path, query_pattern, queries)?; continue; } @@ -207,7 +306,7 @@ impl SuiteConfig { } let file_name = entry.file_name().to_string_lossy().into_owned(); - if let Some(id) = parse_query_file_name(&file_name) { + if let Some(id) = query_pattern.parse_file_name(&file_name) { queries.push(SuiteQuery { id, file_name, @@ -219,26 +318,106 @@ impl SuiteConfig { Ok(()) } + /// Merges supplied suite option values with defaults and validates that all + /// values are allowed for this suite. + pub fn resolve_option_values( + &self, + supplied: &BTreeMap, + ) -> Result> { + for key in supplied.keys() { + if self.option(key).is_none() { + return Err(DataFusionError::Configuration(format!( + "unknown option '--{}' for suite '{}'", + key, self.name + ))); + } + } + + let mut values = BTreeMap::new(); + for option in &self.options { + let value = supplied + .get(&option.name) + .map(String::as_str) + .unwrap_or(option.default.as_str()); + + if !option.values.iter().any(|allowed| allowed == value) { + return Err(DataFusionError::Configuration(format!( + "invalid value '{}' for --{}; expected one of: {}", + value, + option.name, + option.values.join(", ") + ))); + } + + values.insert(option.name.clone(), value.to_string()); + } + + Ok(values) + } + + /// Converts selected suite option values and path replacements into the + /// replacement map passed to `SqlBenchmark`. + /// + /// Path replacement values are resolved relative to the suite file unless + /// they are already absolute. + pub fn option_replacements( + &self, + values: &BTreeMap, + ) -> HashMap { + let mut replacements = + HashMap::with_capacity(self.path_replacements.len() + self.options.len()); + + for (key, value) in &self.path_replacements { + let value = if value.is_absolute() { + value.to_string_lossy().into_owned() + } else { + self.suite_dir.join(value).to_string_lossy().into_owned() + }; + + replacements.insert(key.clone(), value); + } + + replacements.extend(self.options.iter().filter_map(|option| { + values + .get(&option.name) + .map(|value| (option.env.clone(), value.clone())) + })); + + replacements + } + /// Validates suite metadata that cannot be enforced by TOML deserialization. fn validate(&self) -> Result<()> { self.validate_suite_fields()?; + self.validate_path_replacements()?; self.validate_options() } /// Validates required suite-level fields. fn validate_suite_fields(&self) -> Result<()> { - if self.name.trim().is_empty() { - return Err(DataFusionError::Configuration( - "suite name cannot be empty".to_string(), - )); + QueryPattern::new(&self.query_pattern)?; + + Ok(()) + } + + /// Validates static path replacement declarations. + fn validate_path_replacements(&self) -> Result<()> { + for key in self.path_replacements.keys() { + if key.trim().is_empty() { + return Err(DataFusionError::Configuration( + "path replacement name cannot be empty".to_string(), + )); + } } Ok(()) } - /// Validates suite-defined option declarations. + /// Validates suite-defined CLI option declarations. fn validate_options(&self) -> Result<()> { let mut option_names = HashSet::new(); + let mut short_aliases = HashSet::new(); + for option in &self.options { if !option_names.insert(option.name.as_str()) { return Err(DataFusionError::Configuration(format!( @@ -248,14 +427,142 @@ impl SuiteConfig { } option.validate()?; + + if let Some(short) = &option.short + && !short_aliases.insert(short.as_str()) + { + return Err(DataFusionError::Configuration(format!( + "duplicate short alias '{short}'" + ))); + } } Ok(()) } } +/// Parsed suite query filename pattern used for discovery and lookup. +#[derive(Debug)] +struct QueryPattern<'a> { + /// File name prefix before the query-id token. + prefix: &'a str, + /// File name suffix after the query-id token. + suffix: &'a str, + /// Whether file names use padded or unpadded query ids. + query_id_format: QueryIdFormat, +} + +/// Query id formatting mode selected by the suite query pattern token. +#[derive(Debug, Clone, Copy)] +enum QueryIdFormat { + /// Formats query ids as at least two digits. + Padded, + /// Formats query ids without padding. + Unpadded, +} + +/// Parses query filename patterns and applies them in both directions. +impl<'a> QueryPattern<'a> { + /// Creates a validated pattern from one `.suite` `query_pattern` value. + fn new(pattern: &'a str) -> Result { + if pattern.trim().is_empty() { + return Err(DataFusionError::Configuration( + "query_pattern cannot be empty".to_string(), + )); + } + + if !pattern.ends_with(".benchmark") { + return Err(DataFusionError::Configuration( + "query_pattern must name .benchmark files".to_string(), + )); + } + + let padded_count = pattern.matches(QUERY_ID_PADDED_TOKEN).count(); + let unpadded_count = pattern.matches(QUERY_ID_TOKEN).count(); + + if padded_count + unpadded_count != 1 { + return Err(DataFusionError::Configuration(format!( + "query_pattern must contain exactly one {QUERY_ID_TOKEN} or {QUERY_ID_PADDED_TOKEN} token" + ))); + } + + if padded_count == 1 { + let (prefix, suffix) = + pattern.split_once(QUERY_ID_PADDED_TOKEN).ok_or_else(|| { + DataFusionError::Configuration(format!( + "query_pattern must contain {QUERY_ID_PADDED_TOKEN}" + )) + })?; + Ok(Self { + prefix, + suffix, + query_id_format: QueryIdFormat::Padded, + }) + } else { + let (prefix, suffix) = + pattern.split_once(QUERY_ID_TOKEN).ok_or_else(|| { + DataFusionError::Configuration(format!( + "query_pattern must contain {QUERY_ID_TOKEN}" + )) + })?; + Ok(Self { + prefix, + suffix, + query_id_format: QueryIdFormat::Unpadded, + }) + } + } + + /// Extracts a query id from a file name when it matches this pattern. + fn parse_file_name(&self, file_name: &str) -> Option { + let query_id = file_name + .strip_prefix(self.prefix)? + .strip_suffix(self.suffix)?; + + if query_id.is_empty() || !query_id.chars().all(|c| c.is_ascii_digit()) { + return None; + } + + if matches!(self.query_id_format, QueryIdFormat::Padded) && query_id.len() < 2 { + return None; + } + + query_id.parse().ok() + } +} + +/// Returns the default padded `qNN.benchmark` query filename pattern. +fn default_query_pattern() -> String { + DEFAULT_QUERY_PATTERN.to_string() +} + +/// Infers and validates a suite selector from a `.suite` file path. +fn suite_name_from_path(path: &Path) -> Result { + let file_stem = path.file_stem().ok_or_else(|| { + DataFusionError::Configuration(format!( + "suite file {} has no file name stem", + path.display() + )) + })?; + let name = file_stem.to_str().ok_or_else(|| { + DataFusionError::Configuration(format!( + "suite file name {} is not valid UTF-8", + path.display() + )) + })?; + + if !is_valid_suite_name(name) { + return Err(DataFusionError::Configuration(format!( + "invalid suite file name '{name}'; expected lowercase ASCII letters, digits, hyphens, and underscores" + ))); + } + + Ok(name.to_string()) +} + +/// Validates suite-defined CLI option metadata. impl SuiteOption { - /// Validates one suite-defined option. + /// Validates one suite-defined CLI option. fn validate(&self) -> Result<()> { if self.name.trim().is_empty() { return Err(DataFusionError::Configuration( @@ -265,13 +572,27 @@ impl SuiteOption { if !is_valid_cli_option_name(&self.name) { return Err(DataFusionError::Configuration(format!( - "invalid option name '{}'; expected lowercase ASCII letters, digits, and hyphens", + "invalid CLI option name '{}'; expected lowercase ASCII letters, digits, and hyphens", + self.name + ))); + } + + if RESERVED_CLI_OPTION_NAMES.contains(&self.name.as_str()) { + return Err(DataFusionError::Configuration(format!( + "reserved CLI option name '{}'", self.name ))); } self.validate_short_alias()?; + if self.env.trim().is_empty() { + return Err(DataFusionError::Configuration(format!( + "env for option '{}' cannot be empty", + self.name + ))); + } + if self.help.trim().is_empty() { return Err(DataFusionError::Configuration(format!( "help for option '{}' cannot be empty", @@ -306,7 +627,7 @@ impl SuiteOption { Ok(()) } - /// Validates the optional short alias for one suite-defined option. + /// Validates the optional short alias for one suite-defined CLI option. fn validate_short_alias(&self) -> Result<()> { let Some(short) = &self.short else { return Ok(()); @@ -326,10 +647,20 @@ impl SuiteOption { ))); } + if short.chars().count() == 1 + && RESERVED_CLI_SHORT_ALIASES.contains(&short.as_str()) + { + return Err(DataFusionError::Configuration(format!( + "reserved short alias '{short}' for option '{}'", + self.name + ))); + } + Ok(()) } } +/// Discovers suite files and provides lookup by selector name. impl SuiteRegistry { /// Discovers all suite files below the SQL benchmark root. /// @@ -360,6 +691,7 @@ impl SuiteRegistry { let mut suite_files = Vec::new(); let suite_dir = entry.path(); + for suite_entry in read_dir_entries(&suite_dir, "benchmark suite directory")? { let path = suite_entry.path(); @@ -389,112 +721,972 @@ impl SuiteRegistry { Ok(Self { suites }) } - /// Returns discovered suites sorted by selector name. + /// Returns a suite by its command-line selector name. + pub fn get(&self, name: &str) -> Option<&SuiteConfig> { + self.suites.iter().find(|suite| suite.name == name) + } + + /// Returns discovered suites in display order. pub fn suites(&self) -> &[SuiteConfig] { &self.suites } } -fn parse_query_file_name(file_name: &str) -> Option { - let query_id = file_name.strip_prefix('q')?.strip_suffix(".benchmark")?; +/// Reads directory entries and wraps filesystem errors with runner-specific context. +fn read_dir_entries(path: &Path, description: &str) -> Result> { + let entries = fs::read_dir(path).map_err(|e| { + DataFusionError::External( + format!("failed to read {description} {}: {e}", path.display()).into(), + ) + })?; - if query_id.len() < 2 || !query_id.chars().all(|c| c.is_ascii_digit()) { - return None; - } + collect_dir_entries(path, description, entries) +} - query_id.parse().ok() +/// Collects directory entries while preserving contextual filesystem errors. +fn collect_dir_entries( + path: &Path, + description: &str, + entries: impl IntoIterator>, +) -> Result> { + entries + .into_iter() + .map(|entry| { + entry.map_err(|e| { + DataFusionError::External( + format!("failed to read {description} entry {}: {e}", path.display()) + .into(), + ) + }) + }) + .collect() } +/// Checks the restricted option-name syntax accepted for suite-defined CLI flags and aliases. fn is_valid_cli_option_name(name: &str) -> bool { - name.chars() - .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-') -} + let mut previous_was_hyphen = false; + for c in name.chars() { + let is_valid = c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-'; + if !is_valid || (c == '-' && previous_was_hyphen) { + return false; + } + previous_was_hyphen = c == '-'; + } -fn read_dir_entries(dir: &Path, label: &str) -> Result> { - let entries = fs::read_dir(dir).map_err(|e| { - DataFusionError::External( - format!("failed to read {label} {}: {e}", dir.display()).into(), - ) - })?; + !name.starts_with('-') && !name.ends_with('-') +} - entries - .collect::>>() - .map_err(|e| DataFusionError::External(Box::new(e))) +/// Checks the restricted filename syntax accepted for suite selectors. +fn is_valid_suite_name(name: &str) -> bool { + !name.is_empty() + && name + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_') } #[cfg(test)] mod tests { use super::*; use std::fs; + use std::io; fn manifest_path(path: &str) -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) } + fn suite_option_with_short(short: Option<&str>) -> SuiteOption { + SuiteOption { + name: "format".to_string(), + short: short.map(str::to_string), + env: "DATA_FORMAT".to_string(), + default: "parquet".to_string(), + values: vec!["parquet".to_string()], + help: "Selects the data format.".to_string(), + } + } + + fn suite_with_format_option() -> SuiteConfig { + SuiteConfig { + name: "example".to_string(), + description: "Example benchmark suite".to_string(), + query_pattern: DEFAULT_QUERY_PATTERN.to_string(), + path_replacements: BTreeMap::new(), + options: vec![suite_option_with_short(None)], + examples: vec![], + suite_path: PathBuf::from("example.suite"), + suite_dir: PathBuf::from("."), + query_cache: OnceCell::new(), + } + } + + fn suite_file_error_message(suite_contents: impl AsRef) -> String { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("invalid.suite"); + fs::write(&suite_path, suite_contents.as_ref()).unwrap(); + + SuiteConfig::from_file(&suite_path).unwrap_err().to_string() + } + + fn assert_suite_file_error_contains( + suite_contents: impl AsRef, + expected_message: &str, + ) { + let message = suite_file_error_message(suite_contents); + + assert!( + message.contains(expected_message), + "expected error containing {expected_message:?}, got {message:?}" + ); + } + #[test] - fn discovers_tpch_suite_file() { - let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); + fn read_dir_entries_reports_directory_open_errors() { + let tempdir = tempfile::tempdir().unwrap(); + let missing_dir = tempdir.path().join("missing"); + + let err = read_dir_entries(&missing_dir, "test directory") + .expect_err("read should fail"); + let message = err.to_string(); + + assert!( + message.contains("failed to read test directory"), + "{message}" + ); + assert!( + message.contains(&missing_dir.display().to_string()), + "{message}" + ); + } - assert_eq!(registry.suites().len(), 1); - assert_eq!(registry.suites()[0].name, "tpch"); - assert_eq!(registry.suites()[0].options.len(), 2); + #[test] + fn read_dir_entries_reports_directory_entry_errors() { + let tempdir = tempfile::tempdir().unwrap(); + let err = collect_dir_entries( + tempdir.path(), + "test directory", + std::iter::once(Err::(io::Error::other("entry failed"))), + ) + .expect_err("read should fail"); + let message = err.to_string(); + + assert!( + message.contains("failed to read test directory entry"), + "{message}" + ); + assert!( + message.contains(&tempdir.path().display().to_string()), + "{message}" + ); + assert!(message.contains("entry failed"), "{message}"); } #[test] - fn discovers_query_ids_in_numeric_order() { - let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); - let suite = ®istry.suites()[0]; + fn from_file_reports_missing_file_errors() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("missing.suite"); + + let err = SuiteConfig::from_file(&suite_path).unwrap_err(); + let message = err.to_string(); + + assert!(message.contains("failed to read suite file"), "{message}"); + assert!( + message.contains(&suite_path.display().to_string()), + "{message}" + ); + } + + #[test] + fn from_file_reports_toml_parse_errors() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("invalid.suite"); + fs::write(&suite_path, "name = [").unwrap(); + + let err = SuiteConfig::from_file(&suite_path).unwrap_err(); + let message = err.to_string(); + + assert!(message.contains("failed to parse suite file"), "{message}"); + assert!( + message.contains(&suite_path.display().to_string()), + "{message}" + ); + } + + #[test] + fn query_pattern_new_rejects_missing_query_id_token() { + let err = QueryPattern::new("query.benchmark").unwrap_err(); + + assert!( + err.to_string() + .contains("query_pattern must contain exactly one"), + "{err}" + ); + } + + #[test] + fn query_pattern_new_rejects_multiple_query_id_tokens() { + let err = + QueryPattern::new("q{QUERY_ID}-{QUERY_ID_PADDED}.benchmark").unwrap_err(); + + assert!( + err.to_string() + .contains("query_pattern must contain exactly one"), + "{err}" + ); + } + + #[test] + fn resolve_option_values_rejects_unknown_option() { + let suite = suite_with_format_option(); + let supplied = BTreeMap::from([("unknown".to_string(), "parquet".to_string())]); + + let err = suite.resolve_option_values(&supplied).unwrap_err(); + let message = err.to_string(); + + assert!(message.contains("unknown option '--unknown'"), "{message}"); + assert!(message.contains("suite 'example'"), "{message}"); + } + + #[test] + fn resolve_option_values_rejects_invalid_value() { + let suite = suite_with_format_option(); + let supplied = BTreeMap::from([("format".to_string(), "csv".to_string())]); + + let err = suite.resolve_option_values(&supplied).unwrap_err(); + let message = err.to_string(); + + assert!( + message.contains("invalid value 'csv' for --format"), + "{message}" + ); + assert!(message.contains("expected one of: parquet"), "{message}"); + } + + #[test] + fn parses_suite_file_with_multiple_options() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + fs::write( + &suite_path, + r#" +description = "Example benchmark suite" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "format" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet", "csv"] +help = "Selects the data format." + +[[options]] +name = "scale-factor" +env = "SCALE_FACTOR" +default = "1" +values = ["1", "10"] +help = "Selects the scale factor." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- run example" +description = "Run the example suite." +"#, + ) + .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + + assert_eq!(suite.name, "example"); + assert_eq!(suite.description, "Example benchmark suite"); + assert_eq!(suite.query_pattern, "q{QUERY_ID_PADDED}.benchmark"); + assert_eq!(suite.suite_path, suite_path); + assert_eq!(suite.suite_dir, tempdir.path()); + assert_eq!(suite.query_search_root(), tempdir.path()); + assert_eq!(suite.options.len(), 2); + assert_eq!(suite.option("format").unwrap().default, "parquet"); + assert_eq!(suite.option("scale-factor").unwrap().values, ["1", "10"]); + assert!(suite.option("missing").is_none()); + assert_eq!(suite.examples.len(), 1); + assert_eq!( + suite.examples[0].command, + "cargo run --release --bin benchmark_runner -- run example" + ); + } + + #[test] + fn suite_name_is_inferred_from_file_stem() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example_suite.suite"); + fs::write( + &suite_path, + r#" +description = "Example benchmark suite" +"#, + ) + .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + + assert_eq!(suite.name, "example_suite"); + } + + #[test] + fn validation_rejects_top_level_suite_name_field() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + fs::write( + &suite_path, + r#" +name = "example" +description = "Example benchmark suite" +"#, + ) + .unwrap(); + + let err = SuiteConfig::from_file(&suite_path).unwrap_err(); + + assert!(err.to_string().contains("unknown field")); + } + + #[test] + fn path_replacements_are_resolved_relative_to_suite_dir() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + fs::write( + &suite_path, + r#" +description = "Example benchmark suite" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../data" +"#, + ) + .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + let replacements = suite.option_replacements(&BTreeMap::new()); + let expected_data_dir = tempdir + .path() + .join("../data") + .to_string_lossy() + .into_owned(); + + assert_eq!( + replacements.get("DATA_DIR").map(String::as_str), + Some(expected_data_dir.as_str()) + ); + } + + #[test] + fn suite_file_without_query_pattern_uses_default_pattern() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + fs::write( + &suite_path, + r#" +description = "Example benchmark suite" +"#, + ) + .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + + assert_eq!(suite.query_pattern, "q{QUERY_ID_PADDED}.benchmark"); + } + + #[test] + fn checked_in_tpch_suite_has_expected_options() { + let suite = + SuiteConfig::from_file(manifest_path("sql_benchmarks/tpch/tpch.suite")) + .unwrap(); + + assert_eq!(suite.name, "tpch"); + assert_eq!( + suite.query_search_root(), + manifest_path("sql_benchmarks/tpch") + ); + assert_eq!(suite.option("format").unwrap().env, "TPCH_FILE_TYPE"); + assert_eq!(suite.option("format").unwrap().short.as_deref(), Some("f")); + assert_eq!( + suite.option("format").unwrap().values, + ["parquet", "csv", "mem"] + ); + assert_eq!(suite.option("scale-factor").unwrap().env, "BENCH_SIZE"); + assert_eq!(suite.option("scale-factor").unwrap().values, ["1", "10"]); + } + + #[test] + fn discover_queries_finds_tpch_queries_sorted_by_id() { + let suite = + SuiteConfig::from_file(manifest_path("sql_benchmarks/tpch/tpch.suite")) + .unwrap(); + let queries = suite.discover_queries().unwrap(); - let ids = queries.iter().map(|query| query.id).collect::>(); - assert_eq!(ids.first(), Some(&1)); - assert_eq!(ids.last(), Some(&22)); - assert!(!ids.contains(&0)); + let query_ids = queries.iter().map(|query| query.id).collect::>(); + assert!(query_ids.starts_with(&[1, 2, 3])); + assert!(query_ids.ends_with(&[20, 21, 22])); + + let q01 = queries.iter().find(|query| query.id == 1).unwrap(); + assert_eq!(q01.file_name, "q01.benchmark"); + assert_eq!( + q01.path, + manifest_path("sql_benchmarks/tpch/benchmarks/q01.benchmark") + ); + + let q15 = queries.iter().find(|query| query.id == 15).unwrap(); + assert_eq!(q15.file_name, "q15.benchmark"); + + let q22 = queries.iter().find(|query| query.id == 22).unwrap(); + assert_eq!(q22.file_name, "q22.benchmark"); } #[test] - fn rejects_duplicate_suite_names() { - let dir = tempfile::tempdir().unwrap(); - let one = dir.path().join("one"); - let two = dir.path().join("two"); - fs::create_dir_all(&one).unwrap(); - fs::create_dir_all(&two).unwrap(); + fn discover_queries_recursively_from_suite_dir_without_benchmark_dir() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + let nested_dir = tempdir.path().join("nested/queries"); + fs::create_dir_all(&nested_dir).unwrap(); + fs::write(nested_dir.join("q02.benchmark"), "run\nSELECT 2\n").unwrap(); + fs::write(tempdir.path().join("q01.benchmark"), "run\nSELECT 1\n").unwrap(); + fs::write(nested_dir.join("notes.benchmark"), "run\nSELECT 3\n").unwrap(); fs::write( - one.join("suite.suite"), - "name = \"dup\"\ndescription = \"one\"\n", + &suite_path, + r#" +description = "Example benchmark suite" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" +"#, ) .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + let queries = suite.discover_queries().unwrap(); + + assert_eq!( + queries + .iter() + .map(|query| (query.id, query.file_name.as_str())) + .collect::>(), + vec![(1, "q01.benchmark"), (2, "q02.benchmark")] + ); + assert_eq!(queries[0].path, tempdir.path().join("q01.benchmark")); + assert_eq!(queries[1].path, nested_dir.join("q02.benchmark")); + } + + #[test] + fn discover_queries_uses_unpadded_query_pattern() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + let nested_dir = tempdir.path().join("nested"); + fs::create_dir_all(&nested_dir).unwrap(); + fs::write(tempdir.path().join("query-1.benchmark"), "run\nSELECT 1\n").unwrap(); + fs::write(nested_dir.join("query-12.benchmark"), "run\nSELECT 12\n").unwrap(); + fs::write(tempdir.path().join("q01.benchmark"), "run\nSELECT 99\n").unwrap(); fs::write( - two.join("suite.suite"), - "name = \"dup\"\ndescription = \"two\"\n", + &suite_path, + r#" +description = "Example benchmark suite" +query_pattern = "query-{QUERY_ID}.benchmark" +"#, ) .unwrap(); - let err = SuiteRegistry::discover(dir.path()).unwrap_err(); - assert!(err.to_string().contains("duplicate suite name")); + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + let queries = suite.discover_queries().unwrap(); + + assert_eq!( + queries + .iter() + .map(|query| (query.id, query.file_name.as_str())) + .collect::>(), + vec![(1, "query-1.benchmark"), (12, "query-12.benchmark")] + ); + } + + #[test] + fn discover_queries_uses_custom_padded_query_pattern() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + fs::write(tempdir.path().join("case-03.benchmark"), "run\nSELECT 3\n").unwrap(); + fs::write(tempdir.path().join("case-3.benchmark"), "run\nSELECT 99\n").unwrap(); + fs::write( + &suite_path, + r#" +description = "Example benchmark suite" +query_pattern = "case-{QUERY_ID_PADDED}.benchmark" +"#, + ) + .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + let queries = suite.discover_queries().unwrap(); + + assert_eq!( + queries + .iter() + .map(|query| (query.id, query.file_name.as_str())) + .collect::>(), + vec![(3, "case-03.benchmark")] + ); } #[test] - fn rejects_invalid_option_metadata() { - let dir = tempfile::tempdir().unwrap(); - let suite_dir = dir.path().join("suite"); - fs::create_dir_all(&suite_dir).unwrap(); + fn discover_queries_rejects_duplicate_query_ids() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("example.suite"); + let nested_dir = tempdir.path().join("nested"); + fs::create_dir_all(&nested_dir).unwrap(); + fs::write(tempdir.path().join("q01.benchmark"), "run\nSELECT 1\n").unwrap(); + fs::write(nested_dir.join("q01.benchmark"), "run\nSELECT 1\n").unwrap(); fs::write( - suite_dir.join("bad.suite"), + &suite_path, + r#" +description = "Example benchmark suite" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" +"#, + ) + .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + let err = suite.discover_queries().unwrap_err(); + + assert!(err.to_string().contains("duplicate QUERY_ID 1")); + } + + #[test] + fn cloning_suite_config_does_not_copy_runtime_caches() { + let suite = + SuiteConfig::from_file(manifest_path("sql_benchmarks/tpch/tpch.suite")) + .unwrap(); + + let queries = suite.discover_queries().unwrap(); + assert!(!queries.is_empty()); + assert!(suite.query_cache.get().is_some()); + + let cloned = suite.clone(); + + assert!(cloned.query_cache.get().is_none()); + } + + #[test] + fn validation_rejects_duplicate_option_names() { + assert_suite_file_error_contains( + r#" +description = "Duplicate options" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "format" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet"] +help = "Selects the data format." + +[[options]] +name = "format" +env = "OTHER_DATA_FORMAT" +default = "csv" +values = ["csv"] +help = "Selects another data format." +"#, + "duplicate option name", + ); + } + + #[test] + fn validation_rejects_default_values_missing_from_values() { + let message = suite_file_error_message( + r#" +description = "Missing default" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "format" +env = "DATA_FORMAT" +default = "parquet" +values = ["csv"] +help = "Selects the data format." +"#, + ); + + assert!(message.contains("default")); + assert!(message.contains("values")); + } + + #[test] + fn validation_rejects_invalid_suite_file_name() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("Invalid.suite"); + fs::write( + &suite_path, + r#" +description = "Invalid suite file name" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" +"#, + ) + .unwrap(); + + let err = SuiteConfig::from_file(&suite_path).unwrap_err(); + assert!(err.to_string().contains("invalid suite file name")); + } + + #[test] + fn suite_file_name_allows_any_valid_characters() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_path = tempdir.path().join("-example__suite-.suite"); + fs::write( + &suite_path, + r#" +description = "Example benchmark suite" +"#, + ) + .unwrap(); + + let suite = SuiteConfig::from_file(&suite_path).unwrap(); + + assert_eq!(suite.name, "-example__suite-"); + } + + #[test] + fn validation_rejects_unknown_top_level_fields() { + assert_suite_file_error_contains( + r#" +description = "Example benchmark suite" +unexpected = "value" +"#, + "unknown field", + ); + } + + #[test] + fn validation_rejects_unknown_option_fields() { + assert_suite_file_error_contains( + r#" +description = "Example benchmark suite" + +[[options]] +name = "format" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet"] +help = "Selects the data format." +unexpected = "value" +"#, + "unknown field", + ); + } + + #[test] + fn validation_rejects_query_pattern_without_query_id_token() { + assert_suite_file_error_contains( + r#" +description = "Invalid pattern" +query_pattern = "query.benchmark" +"#, + "query_pattern", + ); + } + + #[test] + fn validation_rejects_query_pattern_with_multiple_query_id_tokens() { + assert_suite_file_error_contains( + r#" +description = "Invalid pattern" +query_pattern = "q{QUERY_ID}-{QUERY_ID_PADDED}.benchmark" +"#, + "query_pattern", + ); + } + + #[test] + fn validation_rejects_empty_query_pattern() { + assert_suite_file_error_contains( + r#" +description = "Invalid pattern" +query_pattern = "" +"#, + "query_pattern cannot be empty", + ); + } + + #[test] + fn validation_rejects_query_pattern_without_benchmark_extension() { + assert_suite_file_error_contains( + r#" +description = "Invalid pattern" +query_pattern = "q{QUERY_ID}.sql" +"#, + "query_pattern must name .benchmark files", + ); + } + + #[test] + fn validation_rejects_empty_path_replacement_name() { + assert_suite_file_error_contains( + r#" +description = "Invalid path replacement" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +"" = "../data" +"#, + "path replacement name cannot be empty", + ); + } + + #[test] + fn validation_rejects_empty_option_name() { + assert_suite_file_error_contains( + r#" +description = "Empty option name" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet"] +help = "Selects the data format." +"#, + "option name cannot be empty", + ); + } + + #[test] + fn validation_rejects_invalid_option_name() { + assert_suite_file_error_contains( + r#" +description = "Invalid option name" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "Format" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet"] +help = "Selects the data format." +"#, + "invalid CLI option name 'Format'", + ); + } + + #[test] + fn validation_rejects_empty_short_alias() { + let option = suite_option_with_short(Some(" ")); + + let err = option.validate_short_alias().unwrap_err(); + + assert!( + err.to_string() + .contains("short alias for option 'format' cannot be empty") + ); + } + + #[test] + fn validation_rejects_invalid_short_alias() { + let option = suite_option_with_short(Some("F")); + + let err = option.validate_short_alias().unwrap_err(); + let message = err.to_string(); + + assert!(message.contains("invalid short alias 'F'"), "{message}"); + assert!(message.contains("option 'format'"), "{message}"); + } + + #[test] + fn validation_rejects_duplicate_short_aliases() { + assert_suite_file_error_contains( r#" -name = "bad" -description = "bad suite" +description = "Duplicate short aliases" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" [[options]] -name = "BAD" -default = "one" -values = ["one"] -help = "bad" +name = "format" +short = "f" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet"] +help = "Selects the data format." + +[[options]] +name = "filter" +short = "f" +env = "FILTER" +default = "none" +values = ["none"] +help = "Selects the filter." +"#, + "duplicate short alias 'f'", + ); + } + + #[test] + fn validation_rejects_empty_option_env() { + assert_suite_file_error_contains( + r#" +description = "Empty option env" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "format" +env = "" +default = "parquet" +values = ["parquet"] +help = "Selects the data format." +"#, + "env for option 'format' cannot be empty", + ); + } + + #[test] + fn validation_rejects_empty_option_help() { + assert_suite_file_error_contains( + r#" +description = "Empty option help" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "format" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet"] +help = "" +"#, + "help for option 'format' cannot be empty", + ); + } + + #[test] + fn validation_rejects_empty_option_values() { + assert_suite_file_error_contains( + r#" +description = "Empty option values" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "format" +env = "DATA_FORMAT" +default = "parquet" +values = [] +help = "Selects the data format." +"#, + "values for option 'format' cannot be empty", + ); + } + + #[test] + fn validation_rejects_reserved_cli_option_names() { + assert_suite_file_error_contains( + r#" +description = "Reserved option name" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "iterations" +env = "ITERATIONS" +default = "1" +values = ["1"] +help = "Conflicts with CommonOpt." +"#, + "reserved CLI option name", + ); + } + + #[test] + fn validation_rejects_new_runner_option_names() { + for option_name in ["criterion", "save-baseline"] { + assert_suite_file_error_contains( + format!( + r#" +description = "Reserved runner option" +query_pattern = "q{{QUERY_ID_PADDED}}.benchmark" + +[[options]] +name = "{option_name}" +env = "RESERVED" +default = "value" +values = ["value"] +help = "Conflicts with benchmark_runner run options." +"# + ), + "reserved CLI option name", + ); + } + } + + #[test] + fn validation_rejects_reserved_short_aliases() { + for short in ["h", "p", "o", "i", "n", "s", "d"] { + assert_suite_file_error_contains( + format!( + r#" +description = "Reserved short alias" +query_pattern = "q{{QUERY_ID_PADDED}}.benchmark" + +[[options]] +name = "format" +short = "{short}" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet"] +help = "Conflicts with benchmark_runner short options." +"# + ), + "reserved short alias", + ); + } + } + + #[test] + fn validation_rejects_duplicate_option_values() { + assert_suite_file_error_contains( + r#" +description = "Duplicate option values" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "format" +env = "DATA_FORMAT" +default = "parquet" +values = ["parquet", "parquet"] +help = "Selects the data format." +"#, + "duplicate value", + ); + } + + #[test] + fn registry_discovery_rejects_duplicate_suite_names() { + let tempdir = tempfile::tempdir().unwrap(); + let suite_one_dir = tempdir.path().join("one"); + let suite_two_dir = tempdir.path().join("two"); + fs::create_dir(&suite_one_dir).unwrap(); + fs::create_dir(&suite_two_dir).unwrap(); + fs::write( + suite_one_dir.join("duplicate.suite"), + r#" +description = "First duplicate suite" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" +"#, + ) + .unwrap(); + fs::write( + suite_two_dir.join("duplicate.suite"), + r#" +description = "Second duplicate suite" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" "#, ) .unwrap(); - let err = SuiteRegistry::discover(dir.path()).unwrap_err(); - assert!(err.to_string().contains("invalid option name")); + let err = SuiteRegistry::discover(tempdir.path()).unwrap_err(); + assert!(err.to_string().contains("duplicate suite name")); + } + + #[test] + fn registry_discovery_finds_tpch() { + let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); + let tpch = registry.get("tpch").unwrap(); + + assert_eq!(tpch.name, "tpch"); } } diff --git a/benchmarks/src/bin/benchmark_runner.rs b/benchmarks/src/bin/benchmark_runner.rs index 831d6f96d7e22..f009e8d24cbb9 100644 --- a/benchmarks/src/bin/benchmark_runner.rs +++ b/benchmarks/src/bin/benchmark_runner.rs @@ -30,9 +30,10 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; #[global_allocator] static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; -fn main() { +#[tokio::main] +async fn main() { env_logger::init(); - if let Err(e) = run_cli(std::env::args()) { + if let Err(e) = run_cli(std::env::args()).await { eprintln!("{e}"); std::process::exit(1); } diff --git a/benchmarks/src/sql_benchmark.rs b/benchmarks/src/sql_benchmark.rs index fc6da24b8a9b2..8b829fd6acf40 100644 --- a/benchmarks/src/sql_benchmark.rs +++ b/benchmarks/src/sql_benchmark.rs @@ -74,6 +74,24 @@ impl SqlBenchmark { ctx: &SessionContext, full_path: impl AsRef, benchmark_directory: impl AsRef, + data_path: Option<&Path>, + ) -> Result { + Self::new_with_replacements( + ctx, + full_path, + benchmark_directory, + &HashMap::new(), + data_path, + ) + .await + } + + pub async fn new_with_replacements( + ctx: &SessionContext, + full_path: impl AsRef, + benchmark_directory: impl AsRef, + replacement_overrides: &HashMap, + data_path: Option<&Path>, ) -> Result { let full_path = full_path.as_ref(); let benchmark_directory = benchmark_directory.as_ref(); @@ -97,6 +115,16 @@ impl SqlBenchmark { "BENCHMARK_DIR", benchmark_directory.to_string_lossy().into_owned(), ); + for (key, value) in replacement_overrides { + insert_replacement(&mut bm.replacement_mapping, key, value.clone()); + } + if let Some(data_path) = data_path { + insert_replacement( + &mut bm.replacement_mapping, + "DATA_DIR", + data_path.to_string_lossy().into_owned(), + ); + } let path = bm.benchmark_path.clone(); bm.process_file(ctx, &path).await?; @@ -201,7 +229,11 @@ impl SqlBenchmark { /// # Errors /// Returns an error if a `run` query fails or if expected plan strings /// are not found. - pub async fn run(&mut self, ctx: &SessionContext, save_results: bool) -> Result<()> { + pub async fn run( + &mut self, + ctx: &SessionContext, + save_results: bool, + ) -> Result { let run_queries = self .queries .get(&QueryDirective::Run) @@ -270,7 +302,7 @@ impl SqlBenchmark { // Store results for verification self.last_results = Some(result); - Ok(()) + Ok(result_count) } /// Calls run and persists results to disk as a CSV file. @@ -283,7 +315,7 @@ impl SqlBenchmark { /// Returns an error if no results are available or if writing to the /// target path fails. pub async fn persist(&mut self, ctx: &SessionContext) -> Result<()> { - self.run(ctx, true).await?; + let _ = self.run(ctx, true).await?; // Check if we have result queries to persist for if self.result_queries.is_empty() { @@ -627,6 +659,10 @@ impl SqlBenchmark { &self.assert_queries } + pub fn expected_plan_strings(&self) -> &[String] { + &self.expect + } + pub fn is_loaded(&self) -> bool { self.is_loaded } @@ -655,7 +691,7 @@ impl QueryDirective { } } - fn as_str(self) -> &'static str { + pub fn as_str(self) -> &'static str { match self { Self::Load => "load", Self::Run => "run", @@ -839,8 +875,13 @@ impl BenchmarkDirective { debug!("Processing {} file: {}", splits[0], splits[1]); - let query_file = fs::read_to_string(splits[1]).map_err(|e| { - exec_datafusion_err!("Failed to read query file {}: {e}", splits[1]) + let query_file_path = + resolve_benchmark_relative_path(splits[1], &bench.replacement_mapping); + let query_file = fs::read_to_string(&query_file_path).map_err(|e| { + exec_datafusion_err!( + "Failed to read query file {}: {e}", + query_file_path.display() + ) })?; let query_file = query_file.replace("\r\n", "\n"); @@ -989,6 +1030,9 @@ impl BenchmarkDirective { } let path = process_replacements(splits[1], &bench.replacement_mapping)?; + let path = resolve_benchmark_relative_path(path, &bench.replacement_mapping) + .to_string_lossy() + .into_owned(); bench.result_queries.push(BenchmarkQuery { path: Some(path), @@ -1076,8 +1120,11 @@ impl BenchmarkDirective { )); } + let template_path = + resolve_benchmark_relative_path(splits[1], &bench.replacement_mapping); + // template: update the path to read - bench.benchmark_path = PathBuf::from(splits[1]); + bench.benchmark_path = template_path.clone(); line.clear(); @@ -1121,7 +1168,7 @@ impl BenchmarkDirective { } // restart the load from the template file - Box::pin(bench.process_file(ctx, Path::new(splits[1]))).await + Box::pin(bench.process_file(ctx, &template_path)).await } async fn process_include( @@ -1137,7 +1184,10 @@ impl BenchmarkDirective { )); } - Box::pin(bench.process_file(ctx, Path::new(splits[1]))).await + let include_path = + resolve_benchmark_relative_path(splits[1], &bench.replacement_mapping); + + Box::pin(bench.process_file(ctx, &include_path)).await } fn process_echo( @@ -1413,6 +1463,33 @@ fn lookup_replacement_value( get_env(&key.to_uppercase()) } +fn resolve_benchmark_relative_path( + path: impl AsRef, + replacement_map: &HashMap, +) -> PathBuf { + let path = path.as_ref(); + if path.is_absolute() || path.exists() { + return path.to_path_buf(); + } + + let Some(benchmark_dir) = replacement_map.get("benchmark_dir").map(PathBuf::from) + else { + return path.to_path_buf(); + }; + let Some(benchmark_dir_name) = benchmark_dir.file_name() else { + return path.to_path_buf(); + }; + + if !path.starts_with(benchmark_dir_name) { + return path.to_path_buf(); + } + + benchmark_dir + .parent() + .map(|parent| parent.join(path)) + .unwrap_or_else(|| path.to_path_buf()) +} + fn read_query_from_reader( reader: &mut BenchmarkFileReader, sql: &str, @@ -1569,7 +1646,7 @@ mod tests { async fn parse_benchmark_file(path: &Path) -> Result { let ctx = SessionContext::new(); let path_string = path.to_string_lossy().into_owned(); - SqlBenchmark::new(&ctx, &path_string, "/tmp").await + SqlBenchmark::new(&ctx, &path_string, "/tmp", None).await } async fn parse_benchmark(contents: &str) -> Result { @@ -1579,6 +1656,19 @@ mod tests { parse_benchmark_file(&path).await } + async fn parse_benchmark_with_replacements( + contents: &str, + replacements: &[(&str, &str)], + ) -> Result { + let temp_dir = tempdir().expect("failed to create benchmark test directory"); + let path = write_test_file(&temp_dir, "parser.benchmark", contents); + let ctx = SessionContext::new(); + let replacements = replacement_map(replacements); + + SqlBenchmark::new_with_replacements(&ctx, &path, "/tmp", &replacements, None) + .await + } + async fn assert_parse_error(contents: &str, expected_message: &str) { let error = parse_benchmark(contents) .await @@ -1691,6 +1781,103 @@ mod tests { assert_eq!(actual, "sf10"); } + #[tokio::test] + async fn benchmark_parsing_uses_explicit_replacements_without_env_mutation() { + let benchmark = parse_benchmark_with_replacements( + "name explicit replacements\nrun\nSELECT '${FILE_TYPE}'", + &[("FILE_TYPE", "parquet")], + ) + .await + .expect("benchmark should parse"); + + assert_eq!( + benchmark.queries()[&QueryDirective::Run], + ["SELECT 'parquet'"] + ); + } + + #[tokio::test] + async fn benchmark_parsing_uses_explicit_data_path_replacement() { + let temp_dir = tempdir().expect("failed to create benchmark test directory"); + let path = write_test_file( + &temp_dir, + "parser.benchmark", + "name explicit data path\nrun\nSELECT '${DATA_DIR:-data}'", + ); + let ctx = SessionContext::new(); + let data_path = temp_dir.path().join("custom-data"); + + let benchmark = SqlBenchmark::new(&ctx, &path, "/tmp", Some(data_path.as_path())) + .await + .expect("benchmark should parse"); + + assert_eq!( + benchmark.replacement_mapping().get("data_dir"), + Some(&data_path.to_string_lossy().into_owned()) + ); + assert_eq!( + benchmark.queries()[&QueryDirective::Run], + [format!("SELECT '{}'", data_path.display())] + ); + } + + #[tokio::test] + async fn benchmark_data_path_overrides_replacement_map_data_dir() { + let temp_dir = tempdir().expect("failed to create benchmark test directory"); + let path = write_test_file( + &temp_dir, + "parser.benchmark", + "name explicit data path\nrun\nSELECT '${DATA_DIR}'", + ); + let ctx = SessionContext::new(); + let replacements = replacement_map(&[("DATA_DIR", "suite-data")]); + let data_path = temp_dir.path().join("cli-data"); + + let benchmark = SqlBenchmark::new_with_replacements( + &ctx, + &path, + "/tmp", + &replacements, + Some(data_path.as_path()), + ) + .await + .expect("benchmark should parse"); + + assert_eq!( + benchmark.replacement_mapping().get("data_dir"), + Some(&data_path.to_string_lossy().into_owned()) + ); + assert_eq!( + benchmark.queries()[&QueryDirective::Run], + [format!("SELECT '{}'", data_path.display())] + ); + } + + #[tokio::test] + async fn run_returns_row_count() { + let ctx = SessionContext::new(); + let mut benchmark = parse_benchmark("run\nSELECT 1\n") + .await + .expect("benchmark should parse"); + + assert_eq!(benchmark.run(&ctx, false).await.unwrap(), 1); + assert_eq!(benchmark.run(&ctx, true).await.unwrap(), 1); + } + + #[tokio::test] + async fn exposes_expected_plan_strings() { + let benchmark = + parse_benchmark("expect_plan PlaceholderRowExec\nrun\nSELECT 1\n") + .await + .expect("benchmark should parse"); + + assert_eq!( + benchmark.expected_plan_strings(), + &["PlaceholderRowExec".to_string()] + ); + assert_eq!(QueryDirective::Run.as_str(), "run"); + } + #[test] fn process_replacements_uses_default_for_missing_variable() { let replacements = HashMap::new(); @@ -2264,7 +2451,7 @@ NULL|(empty) let ctx = SessionContext::new(); let path_string = benchmark_path.to_string_lossy().into_owned(); - let mut benchmark = SqlBenchmark::new(&ctx, &path_string, "/tmp") + let mut benchmark = SqlBenchmark::new(&ctx, &path_string, "/tmp", None) .await .expect("benchmark should parse"); @@ -2577,7 +2764,7 @@ NULL|(empty) let ctx = SessionContext::new(); let benchmark_path_string = benchmark_path.to_string_lossy().into_owned(); - let result = SqlBenchmark::new(&ctx, &benchmark_path_string, "/tmp").await; + let result = SqlBenchmark::new(&ctx, &benchmark_path_string, "/tmp", None).await; let error = result.expect_err("benchmark parsing should fail"); let message = error.to_string();