| // 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. | |
| //! DataFusion SQL benchmark runner. | |
| use clap::{ | |
| Arg, ArgAction, ArgMatches, Command, CommandFactory, FromArgMatches, Parser, | |
| ValueEnum, | |
| }; | |
| use criterion::Criterion; | |
| use datafusion::error::Result; | |
| use datafusion::prelude::SessionContext; | |
| use datafusion_benchmarks::sql_benchmark::SqlBenchmark; | |
| use datafusion_benchmarks::sql_benchmark_runner::{ | |
| BenchmarkFilter, SqlRunConfig, default_sql_benchmark_directory, ensure_selection, | |
| filter_benchmarks, finish_benchmark, load_benchmark_definitions_for_query, make_ctx, | |
| prepare_benchmark, run_criterion_benchmarks_impl_with_namespace, | |
| validate_criterion_namespace, | |
| }; | |
| use datafusion_benchmarks::sql_benchmark_suite::{ | |
| ReservedOptions, SuiteExample, SuiteMetadata, discover_suites, | |
| }; | |
| use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt, print_memory_stats}; | |
| use datafusion_common::instant::Instant; | |
| use datafusion_common::{DataFusionError, exec_datafusion_err}; | |
| use datafusion_common_runtime::SpawnedTask; | |
| use serde::{Serialize, Serializer}; | |
| use std::collections::{BTreeMap, BTreeSet}; | |
| use std::ffi::{OsStr, OsString}; | |
| use std::fmt::Write as _; | |
| use std::io::IsTerminal; | |
| use std::path::Path; | |
| #[cfg(feature = "snmalloc")] | |
| #[global_allocator] | |
| static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; | |
| // `cargo clippy --all-features` enables both allocator features, so prefer | |
| // `snmalloc` in that case and fall back to `mimalloc` otherwise. | |
| #[cfg(all(not(feature = "snmalloc"), feature = "mimalloc"))] | |
| #[global_allocator] | |
| static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; | |
| #[tokio::main] | |
| async fn main() { | |
| env_logger::init(); | |
| if let Err(error) = run_cli().await { | |
| eprintln!("Error: {error}"); | |
| std::process::exit(1); | |
| } | |
| } | |
| #[derive(Debug)] | |
| enum CliAction { | |
| List, | |
| Simple(SqlRunConfig), | |
| Criterion { | |
| config: SqlRunConfig, | |
| save_baseline: Option<String>, | |
| criterion_namespace: Option<String>, | |
| }, | |
| DryRun(DryRunOutput), | |
| } | |
| fn parse_criterion_namespace(value: &str) -> std::result::Result<String, String> { | |
| validate_criterion_namespace(Some(value)).map_err(|error| error.to_string())?; | |
| Ok(value.to_string()) | |
| } | |
| #[derive(Debug, Serialize)] | |
| struct ResolvedSuiteValue { | |
| value: String, | |
| #[serde(serialize_with = "serialize_value_source")] | |
| source: datafusion_benchmarks::sql_benchmark_suite::ValueSource, | |
| environment: String, | |
| } | |
| #[derive(Debug, Serialize)] | |
| struct ResolvedPathValue { | |
| value: String, | |
| #[serde(serialize_with = "serialize_value_source")] | |
| source: datafusion_benchmarks::sql_benchmark_suite::ValueSource, | |
| } | |
| #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] | |
| #[serde(rename_all = "snake_case")] | |
| enum RunMode { | |
| Simple, | |
| Criterion, | |
| } | |
| #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, ValueEnum)] | |
| #[serde(rename_all = "snake_case")] | |
| enum ResultMode { | |
| #[default] | |
| None, | |
| Persist, | |
| Validate, | |
| } | |
| impl ResultMode { | |
| fn config_flags(self) -> (bool, bool) { | |
| match self { | |
| Self::None => (false, false), | |
| Self::Persist => (true, false), | |
| Self::Validate => (false, true), | |
| } | |
| } | |
| } | |
| #[derive(Debug, Serialize)] | |
| struct DryRunCommonOptions { | |
| iterations: usize, | |
| partitions: Option<usize>, | |
| batch_size: Option<usize>, | |
| } | |
| #[derive(Debug, Serialize)] | |
| struct DryRunOutput { | |
| suite: String, | |
| query: Option<String>, | |
| subgroup: Option<String>, | |
| mode: RunMode, | |
| result_mode: ResultMode, | |
| common_options: DryRunCommonOptions, | |
| suite_options: BTreeMap<String, ResolvedSuiteValue>, | |
| path_replacements: BTreeMap<String, ResolvedPathValue>, | |
| } | |
| #[derive(Debug, Parser)] | |
| #[command( | |
| name = "benchmark_runner", | |
| about = "Run DataFusion SQL benchmarks", | |
| styles = criterion_like_styles(), | |
| )] | |
| struct Cli { | |
| #[arg(value_name = "BENCHMARK", help = "SQL benchmark group to run")] | |
| benchmark: Option<String>, | |
| #[arg(short = 'q', long = "query", env = "BENCH_QUERY")] | |
| query: Option<String>, | |
| #[arg(long = "subgroup", env = "BENCH_SUBGROUP")] | |
| subgroup: Option<String>, | |
| #[command(flatten)] | |
| common: CommonOpt, | |
| #[arg( | |
| long = "criterion", | |
| action = ArgAction::SetTrue, | |
| help = "Run benchmarks with Criterion" | |
| )] | |
| criterion: bool, | |
| #[arg( | |
| long = "list", | |
| action = ArgAction::SetTrue, | |
| help = "List available SQL benchmark groups" | |
| )] | |
| list: bool, | |
| #[arg( | |
| short = 'o', | |
| long = "output", | |
| help = "Write simple runner results as JSON to this path" | |
| )] | |
| output: Option<std::path::PathBuf>, | |
| #[arg( | |
| long = "save-baseline", | |
| value_name = "BASELINE", | |
| help = "Save Criterion measurements to the named baseline" | |
| )] | |
| save_baseline: Option<String>, | |
| #[arg( | |
| long = "criterion-namespace", | |
| env = "BENCH_NAMESPACE", | |
| value_name = "NAMESPACE", | |
| value_parser = parse_criterion_namespace, | |
| help = "Append a safe namespace to benchmark groups (enables Criterion)" | |
| )] | |
| criterion_namespace: Option<String>, | |
| #[arg(short = 'p', long = "path", value_name = "PATH")] | |
| path: Option<std::path::PathBuf>, | |
| #[arg( | |
| long = "result-mode", | |
| value_enum, | |
| value_name = "MODE", | |
| help = "Handle expected results: none, persist, or validate" | |
| )] | |
| result_mode: Option<ResultMode>, | |
| #[arg(long = "dry-run", action = ArgAction::SetTrue)] | |
| dry_run: bool, | |
| } | |
| /// Parses CLI arguments, runs the selected action, and prints any output. | |
| async fn run_cli() -> Result<()> { | |
| let benchmark_dir = default_sql_benchmark_directory(); | |
| let output = run_cli_from(std::env::args_os(), &benchmark_dir).await?; | |
| if !output.is_empty() { | |
| println!("{output}"); | |
| } | |
| Ok(()) | |
| } | |
| fn serialize_value_source<S>( | |
| source: &datafusion_benchmarks::sql_benchmark_suite::ValueSource, | |
| serializer: S, | |
| ) -> std::result::Result<S::Ok, S::Error> | |
| where | |
| S: Serializer, | |
| { | |
| let value = match source { | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine => { | |
| "command_line" | |
| } | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::Environment => { | |
| "environment" | |
| } | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default => "default", | |
| }; | |
| serializer.serialize_str(value) | |
| } | |
| fn clap_display_output(error: &DataFusionError) -> Option<String> { | |
| let DataFusionError::External(error) = error else { | |
| return None; | |
| }; | |
| let error = error.downcast_ref::<clap::Error>()?; | |
| matches!( | |
| error.kind(), | |
| clap::error::ErrorKind::DisplayHelp | clap::error::ErrorKind::DisplayVersion | |
| ) | |
| .then(|| error.to_string()) | |
| } | |
| async fn run_cli_from<I, T>(args: I, benchmark_dir: &Path) -> Result<String> | |
| where | |
| I: IntoIterator<Item = T>, | |
| T: Into<OsString> + Clone, | |
| { | |
| match parse_cli_from(args, benchmark_dir) { | |
| Ok(action) => run_cli_action(action, benchmark_dir).await, | |
| Err(error) => clap_display_output(&error).ok_or(error), | |
| } | |
| } | |
| fn format_examples(examples: &[SuiteExample]) -> String { | |
| if examples.is_empty() { | |
| return String::new(); | |
| } | |
| let mut output = String::from("Examples:\n"); | |
| for example in examples { | |
| output.push_str(" "); | |
| output.push_str(example.command()); | |
| output.push_str("\n "); | |
| output.push_str(example.description()); | |
| output.push('\n'); | |
| } | |
| output | |
| } | |
| fn build_cli(suite: Option<&SuiteMetadata>) -> Command { | |
| let mut command = Cli::command(); | |
| if let Some(suite) = suite { | |
| command = command.about(suite.description().to_string()); | |
| for option in suite.options() { | |
| let mut arg = Arg::new(option.name().to_string()) | |
| .long(option.name().to_string()) | |
| .help(option.help().to_string()) | |
| .env(option.env().to_string()) | |
| .default_value(option.default().to_string()); | |
| if let Some(short) = option.short() { | |
| arg = arg.short(short); | |
| } | |
| if let Some(values) = option | |
| .values() | |
| .filter(|values| !values.iter().any(|value| value == "...")) | |
| { | |
| arg = arg.value_parser(values.to_vec()); | |
| } | |
| command = command.arg(arg); | |
| } | |
| let examples = format_examples(suite.examples()); | |
| if !examples.is_empty() { | |
| command = command.after_help(examples); | |
| } | |
| } | |
| command | |
| } | |
| fn reserved_options() -> (BTreeSet<String>, BTreeSet<char>) { | |
| let command = Cli::command(); | |
| let long = command | |
| .get_arguments() | |
| .filter_map(|arg| arg.get_long().map(ToOwned::to_owned)) | |
| .collect(); | |
| let short = command | |
| .get_arguments() | |
| .filter_map(|arg| arg.get_short()) | |
| .collect(); | |
| (long, short) | |
| } | |
| fn suite_metadata(benchmark_dir: &Path) -> Result<Vec<SuiteMetadata>> { | |
| let (long, short) = reserved_options(); | |
| discover_suites( | |
| benchmark_dir, | |
| &ReservedOptions { | |
| long: &long, | |
| short: &short, | |
| }, | |
| ) | |
| } | |
| fn format_suite_list(suites: &[SuiteMetadata]) -> String { | |
| let mut output = String::from("SQL benchmarks:\n"); | |
| for suite in suites { | |
| let query_word = if suite.benchmark_count() == 1 { | |
| "query " | |
| } else { | |
| "queries " | |
| }; | |
| writeln!( | |
| output, | |
| " {:<24} {} {query_word}{}", | |
| suite.name(), | |
| suite.benchmark_count(), | |
| suite.description() | |
| ) | |
| .ok(); | |
| } | |
| output.trim_end().to_string() | |
| } | |
| fn locate_suite_arg(args: &[OsString]) -> Result<Option<&str>> { | |
| let Some(argument) = args.get(1) else { | |
| return Ok(None); | |
| }; | |
| if argument == OsStr::new("--help") | |
| || argument == OsStr::new("-h") | |
| || argument == OsStr::new("--list") | |
| || argument == OsStr::new("--dry-run") | |
| { | |
| return Ok(None); | |
| } | |
| let suite = argument.to_str().ok_or_else(|| { | |
| DataFusionError::External("suite name is not valid Unicode".into()) | |
| })?; | |
| if suite.starts_with('-') { | |
| return Err(exec_datafusion_err!( | |
| "suite must be the first argument; options must follow the suite" | |
| )); | |
| } | |
| Ok(Some(suite)) | |
| } | |
| fn try_parse_cli_from<I, T>(args: I, benchmark_dir: &Path) -> Result<CliAction> | |
| where | |
| I: IntoIterator<Item = T>, | |
| T: Into<OsString> + Clone, | |
| { | |
| let args = args.into_iter().map(Into::into).collect::<Vec<_>>(); | |
| let suite = locate_suite_arg(&args)?; | |
| let (long, short) = reserved_options(); | |
| let reserved = ReservedOptions { | |
| long: &long, | |
| short: &short, | |
| }; | |
| let suite = suite | |
| .map(|name| { | |
| if !benchmark_dir.join(name).is_dir() { | |
| let available = discover_suites(benchmark_dir, &reserved)?; | |
| return Err(exec_datafusion_err!( | |
| "unknown benchmark '{name}'\n\n{}", | |
| format_suite_list(&available) | |
| )); | |
| } | |
| SuiteMetadata::load(benchmark_dir, name, &reserved) | |
| }) | |
| .transpose()?; | |
| let matches = build_cli(suite.as_ref()) | |
| .try_get_matches_from(args) | |
| .map_err(|error| DataFusionError::External(Box::new(error)))?; | |
| cli_action_from_matches(&matches, suite.as_ref()) | |
| } | |
| fn parse_cli_from<I, T>(args: I, benchmark_dir: &Path) -> Result<CliAction> | |
| where | |
| I: IntoIterator<Item = T>, | |
| T: Into<OsString> + Clone, | |
| { | |
| try_parse_cli_from(args, benchmark_dir) | |
| } | |
| /// Converts parsed arguments into an executable action and validates mode options. | |
| fn cli_action_from_matches( | |
| matches: &ArgMatches, | |
| suite: Option<&SuiteMetadata>, | |
| ) -> Result<CliAction> { | |
| let mut cli = Cli::from_arg_matches(matches) | |
| .map_err(|e| DataFusionError::External(Box::new(e)))?; | |
| let criterion_implied = !cli.criterion && cli.criterion_namespace.is_some(); | |
| cli.criterion |= criterion_implied; | |
| if cli.dry_run && cli.list { | |
| return Err(exec_datafusion_err!("--list cannot be used with --dry-run")); | |
| } | |
| if cli.dry_run && cli.benchmark.is_none() { | |
| return Err(exec_datafusion_err!("--dry-run requires a benchmark suite")); | |
| } | |
| if cli.criterion && cli.output.is_some() { | |
| return Err(exec_datafusion_err!( | |
| "--output cannot be used with --criterion" | |
| )); | |
| } | |
| if !cli.criterion && cli.save_baseline.is_some() { | |
| return Err(exec_datafusion_err!( | |
| "--save-baseline cannot be used without --criterion" | |
| )); | |
| } | |
| if cli.list || cli.benchmark.is_none() { | |
| return Ok(CliAction::List); | |
| } | |
| if !cli.criterion && cli.common.iterations == 0 { | |
| return Err(exec_datafusion_err!("iterations must be greater than zero")); | |
| } | |
| // we need to know if iterations was set on the command line, not the default value | |
| let iterations_from_cli = matches.value_source("iterations") | |
| == Some(clap::parser::ValueSource::CommandLine); | |
| if cli.criterion && iterations_from_cli { | |
| return Err(exec_datafusion_err!( | |
| "--iterations cannot be used with --criterion" | |
| )); | |
| } | |
| let suite = | |
| suite.ok_or_else(|| exec_datafusion_err!("benchmark suite is required"))?; | |
| if cli.path.is_some() && !suite.path_replacements().contains_key("DATA_DIR") { | |
| return Err(exec_datafusion_err!( | |
| "--path cannot be used because suite '{}' does not declare DATA_DIR", | |
| suite.name() | |
| )); | |
| } | |
| let result_mode = resolve_result_mode(cli.result_mode)?; | |
| let (persist_results, validate_results) = result_mode.config_flags(); | |
| let mut config = SqlRunConfig { | |
| common: cli.common, | |
| filter: BenchmarkFilter { | |
| name: cli.benchmark, | |
| subgroup: cli.subgroup, | |
| query: cli.query, | |
| }, | |
| replacements: Default::default(), | |
| query_filename: None, | |
| persist_results, | |
| validate_results, | |
| output: cli.output, | |
| }; | |
| let suite_options: BTreeMap<String, ResolvedSuiteValue> = suite | |
| .options() | |
| .iter() | |
| .map(|option| { | |
| let value = matches | |
| .get_one::<String>(option.name()) | |
| .expect("suite options always have defaults") | |
| .clone(); | |
| let source = match matches.value_source(option.name()) { | |
| Some(clap::parser::ValueSource::CommandLine) => { | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine | |
| } | |
| Some(clap::parser::ValueSource::EnvVariable) => { | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::Environment | |
| } | |
| Some(clap::parser::ValueSource::DefaultValue) => { | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default | |
| } | |
| vs => unreachable!("unexpected suite option source: {vs:?}"), | |
| }; | |
| ( | |
| option.name().to_string(), | |
| ResolvedSuiteValue { | |
| value, | |
| source, | |
| environment: option.env().to_string(), | |
| }, | |
| ) | |
| }) | |
| .collect(); | |
| let path_replacements: BTreeMap<String, ResolvedPathValue> = suite | |
| .path_replacements() | |
| .iter() | |
| .map(|(key, default)| { | |
| let (value, source) = if key == "DATA_DIR" { | |
| cli.path.as_ref().map_or_else( | |
| || { | |
| ( | |
| default.display().to_string(), | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default, | |
| ) | |
| }, | |
| |path| { | |
| ( | |
| path.display().to_string(), | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine, | |
| ) | |
| }, | |
| ) | |
| } else { | |
| ( | |
| default.display().to_string(), | |
| datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default, | |
| ) | |
| }; | |
| ( | |
| key.to_ascii_lowercase(), | |
| ResolvedPathValue { value, source }, | |
| ) | |
| }) | |
| .collect(); | |
| config.replacements = suite_options | |
| .values() | |
| .map(|resolved| { | |
| ( | |
| resolved.environment.to_ascii_lowercase(), | |
| resolved.value.clone(), | |
| ) | |
| }) | |
| .chain( | |
| path_replacements | |
| .iter() | |
| .map(|(key, resolved)| (key.clone(), resolved.value.clone())), | |
| ) | |
| .collect(); | |
| config.query_filename = config | |
| .filter | |
| .query | |
| .as_deref() | |
| .map(|query| suite.query_filename(query)) | |
| .transpose()?; | |
| if criterion_implied { | |
| eprintln!( | |
| "Enabling Criterion because --criterion-namespace or BENCH_NAMESPACE is set." | |
| ); | |
| } | |
| if cli.dry_run { | |
| let mode = if cli.criterion { | |
| RunMode::Criterion | |
| } else { | |
| RunMode::Simple | |
| }; | |
| return Ok(CliAction::DryRun(DryRunOutput { | |
| suite: suite.name().to_string(), | |
| query: config.filter.query.clone(), | |
| subgroup: config.filter.subgroup.clone(), | |
| mode, | |
| result_mode, | |
| common_options: DryRunCommonOptions { | |
| iterations: config.common.iterations, | |
| partitions: config.common.partitions, | |
| batch_size: config.common.batch_size, | |
| }, | |
| suite_options, | |
| path_replacements, | |
| })); | |
| } | |
| if cli.criterion { | |
| Ok(CliAction::Criterion { | |
| config, | |
| save_baseline: cli.save_baseline, | |
| criterion_namespace: cli.criterion_namespace, | |
| }) | |
| } else { | |
| Ok(CliAction::Simple(config)) | |
| } | |
| } | |
| /// Executes a parsed CLI action and returns any text that should be printed. | |
| async fn run_cli_action(action: CliAction, benchmark_dir: &Path) -> Result<String> { | |
| match action { | |
| CliAction::List => Ok(format_suite_list(&suite_metadata(benchmark_dir)?)), | |
| CliAction::Simple(config) => { | |
| run_simple_benchmarks(benchmark_dir, config).await?; | |
| Ok(String::new()) | |
| } | |
| CliAction::Criterion { | |
| config, | |
| save_baseline, | |
| criterion_namespace, | |
| } => { | |
| if config.output.is_some() { | |
| return Err(exec_datafusion_err!( | |
| "--output cannot be used with --criterion" | |
| )); | |
| } | |
| let benchmark_dir = benchmark_dir.to_path_buf(); | |
| SpawnedTask::spawn_blocking(move || { | |
| run_criterion_benchmarks( | |
| &benchmark_dir, | |
| &config, | |
| save_baseline.as_deref(), | |
| criterion_namespace.as_deref(), | |
| ) | |
| }) | |
| .await | |
| .map_err(|e| DataFusionError::External(Box::new(e)))??; | |
| Ok(String::new()) | |
| } | |
| CliAction::DryRun(output) => serde_json::to_string_pretty(&output) | |
| .map_err(|error| DataFusionError::External(Box::new(error))), | |
| } | |
| } | |
| fn resolve_result_mode(explicit: Option<ResultMode>) -> Result<ResultMode> { | |
| if let Some(mode) = explicit { | |
| return Ok(mode); | |
| } | |
| let persist = parse_compat_bool("BENCH_PERSIST_RESULTS")?; | |
| let validate = parse_compat_bool("BENCH_VALIDATE")?; | |
| Ok(if persist { | |
| ResultMode::Persist | |
| } else if validate { | |
| ResultMode::Validate | |
| } else { | |
| ResultMode::None | |
| }) | |
| } | |
| fn parse_compat_bool(name: &str) -> Result<bool> { | |
| let Some(value) = std::env::var_os(name) else { | |
| return Ok(false); | |
| }; | |
| let value = value | |
| .into_string() | |
| .map_err(|_| exec_datafusion_err!("{name} contains invalid UTF-8"))?; | |
| value.parse::<bool>().map_err(|_| { | |
| exec_datafusion_err!("invalid value '{value}' for {name}; expected true or false") | |
| }) | |
| } | |
| /// Builds the default Criterion runner and optionally records a named baseline. | |
| fn run_criterion_benchmarks( | |
| benchmark_dir: &Path, | |
| config: &SqlRunConfig, | |
| save_baseline: Option<&str>, | |
| criterion_namespace: Option<&str>, | |
| ) -> Result<()> { | |
| let mut criterion = Criterion::default() | |
| .sample_size(10) | |
| .with_output_color(std::io::stdout().is_terminal()); | |
| if let Some(save_baseline) = save_baseline { | |
| criterion = criterion.save_baseline(save_baseline.to_string()); | |
| } | |
| run_criterion_benchmarks_impl_with_namespace( | |
| benchmark_dir, | |
| config, | |
| criterion_namespace, | |
| &mut criterion, | |
| )?; | |
| criterion.final_summary(); | |
| Ok(()) | |
| } | |
| /// Runs selected benchmarks with fixed iteration counts and optional JSON output. | |
| pub async fn run_simple_benchmarks( | |
| benchmark_dir: &Path, | |
| config: SqlRunConfig, | |
| ) -> Result<()> { | |
| if config.common.iterations == 0 { | |
| return Err(exec_datafusion_err!("iterations must be greater than zero")); | |
| } | |
| let listing_ctx = make_ctx(&config.common)?; | |
| let all_benchmarks = load_benchmark_definitions_for_query( | |
| &config.filter, | |
| &listing_ctx, | |
| benchmark_dir, | |
| &config.replacements, | |
| config.query_filename.as_deref(), | |
| ) | |
| .await?; | |
| let selected = filter_benchmarks(&config.filter, all_benchmarks.clone()); | |
| let mut run = BenchmarkRun::new(); | |
| ensure_selection(&config.filter, &all_benchmarks, &selected)?; | |
| for (_group, benchmarks) in selected { | |
| for mut benchmark in benchmarks { | |
| let ctx = make_ctx(&config.common)?; | |
| let result = | |
| run_simple_benchmark(&ctx, &mut benchmark, &config, &mut run).await; | |
| let cleanup_result = benchmark.cleanup(&ctx).await; | |
| finish_benchmark(result, cleanup_result)?; | |
| } | |
| } | |
| run.maybe_write_json(config.output.as_ref())?; | |
| Ok(()) | |
| } | |
| /// Runs one benchmark case, recording each timed iteration. | |
| async fn run_simple_benchmark( | |
| ctx: &SessionContext, | |
| benchmark: &mut SqlBenchmark, | |
| config: &SqlRunConfig, | |
| run: &mut BenchmarkRun, | |
| ) -> Result<()> { | |
| prepare_benchmark(ctx, benchmark, config).await?; | |
| let case_name = benchmark_case_name(benchmark); | |
| // Each case gets its own `SessionContext`, so hand over its pool before the | |
| // case starts. | |
| run.set_memory_pool(&ctx.runtime_env().memory_pool); | |
| run.start_new_case(&case_name); | |
| for iteration in 0..config.common.iterations { | |
| let start = Instant::now(); | |
| let row_count = benchmark.run(ctx, false).await?; | |
| let elapsed = start.elapsed(); | |
| let ms = elapsed.as_secs_f64() * 1000.0; | |
| println!("{case_name} iteration {iteration}: {ms:.1} ms, {row_count} rows"); | |
| run.write_iter(elapsed, row_count); | |
| } | |
| print_memory_stats(&*ctx.runtime_env().memory_pool); | |
| Ok(()) | |
| } | |
| fn benchmark_case_name(benchmark: &SqlBenchmark) -> String { | |
| let mut name = format!("{}/{}", benchmark.group(), benchmark.name()); | |
| if !benchmark.subgroup().is_empty() { | |
| name.push('/'); | |
| name.push_str(benchmark.subgroup()); | |
| } | |
| name | |
| } | |
| fn criterion_like_styles() -> clap::builder::Styles { | |
| use clap::builder::styling::AnsiColor; | |
| clap::builder::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()) | |
| } | |
| #[cfg(test)] | |
| mod tests { | |
| use super::*; | |
| use datafusion_benchmarks::sql_benchmark_runner::run_criterion_benchmarks_impl; | |
| use datafusion_benchmarks::sql_benchmark_runner::{ | |
| load_benchmark_definitions, sort_benchmarks, unknown_benchmark_error, | |
| }; | |
| use datafusion_benchmarks::sql_benchmark_suite::ValueSource; | |
| use std::collections::HashMap; | |
| use std::ffi::OsString; | |
| use std::fs; | |
| use std::path::{Path, PathBuf}; | |
| use std::sync::{Mutex, MutexGuard}; | |
| const CLI_ENV_TEST_CHILD: &str = "DATAFUSION_BENCHMARK_RUNNER_ENV_TEST_CHILD"; | |
| static ENV_MUTEX: Mutex<()> = Mutex::new(()); | |
| struct ScopedEnv { | |
| previous: Vec<(&'static str, Option<OsString>)>, | |
| _lock: MutexGuard<'static, ()>, | |
| } | |
| impl ScopedEnv { | |
| fn set(name: &'static str, value: impl Into<OsString>) -> Self { | |
| let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); | |
| let previous = std::env::var_os(name); | |
| // SAFETY: ENV_MUTEX serializes changes made through ScopedEnv in this | |
| // test module; it does not synchronize environment access elsewhere. | |
| unsafe { std::env::set_var(name, value.into()) }; | |
| Self { | |
| previous: vec![(name, previous)], | |
| _lock: lock, | |
| } | |
| } | |
| fn remove(name: &'static str) -> Self { | |
| let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); | |
| let previous = std::env::var_os(name); | |
| // SAFETY: ENV_MUTEX serializes changes made through ScopedEnv in this | |
| // test module; it does not synchronize environment access elsewhere. | |
| unsafe { std::env::remove_var(name) }; | |
| Self { | |
| previous: vec![(name, previous)], | |
| _lock: lock, | |
| } | |
| } | |
| fn set_many<const N: usize>(changes: [(&'static str, Option<&str>); N]) -> Self { | |
| let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); | |
| let mut previous = Vec::with_capacity(N); | |
| for (name, value) in changes { | |
| previous.push((name, std::env::var_os(name))); | |
| // SAFETY: this guard holds ENV_MUTEX until it restores all entries. | |
| unsafe { | |
| match value { | |
| Some(value) => std::env::set_var(name, value), | |
| None => std::env::remove_var(name), | |
| } | |
| } | |
| } | |
| Self { | |
| previous, | |
| _lock: lock, | |
| } | |
| } | |
| } | |
| impl Drop for ScopedEnv { | |
| fn drop(&mut self) { | |
| // SAFETY: this guard holds ENV_MUTEX until after all entries are restored. | |
| unsafe { | |
| for (name, previous) in self.previous.drain(..).rev() { | |
| match previous { | |
| Some(value) => std::env::set_var(name, value), | |
| None => std::env::remove_var(name), | |
| } | |
| } | |
| } | |
| } | |
| } | |
| /// Loads benchmark definitions, applies CLI-style filters, and sorts each group. | |
| async fn load_benchmarks( | |
| filter: &BenchmarkFilter, | |
| ctx: &SessionContext, | |
| benchmark_dir: &Path, | |
| ) -> Result<BTreeMap<String, Vec<SqlBenchmark>>> { | |
| let benches = | |
| load_benchmark_definitions(filter, ctx, benchmark_dir, &Default::default()) | |
| .await?; | |
| let mut benches = filter_benchmarks(filter, benches); | |
| sort_benchmarks(&mut benches); | |
| Ok(benches) | |
| } | |
| fn write_benchmark(root: &Path, relative_path: &str, contents: &str) -> PathBuf { | |
| let path = root.join(relative_path); | |
| fs::create_dir_all(path.parent().unwrap()).unwrap(); | |
| fs::write(&path, contents).unwrap(); | |
| path | |
| } | |
| fn write_suite(root: &Path, name: &str, description: &str) -> PathBuf { | |
| let path = root.join(name).join(format!("{name}.suite")); | |
| fs::create_dir_all(path.parent().unwrap()).unwrap(); | |
| fs::write(&path, format!("description = {description:?}\n")).unwrap(); | |
| path | |
| } | |
| fn common(iterations: usize) -> CommonOpt { | |
| CommonOpt { | |
| iterations, | |
| partitions: None, | |
| batch_size: None, | |
| mem_pool_type: "fair".to_string(), | |
| memory_limit: None, | |
| sort_spill_reservation_bytes: None, | |
| debug: false, | |
| simulate_latency: false, | |
| } | |
| } | |
| async fn run_cli_with_dir<I, T>(args: I, benchmark_dir: &Path) -> Result<String> | |
| where | |
| I: IntoIterator<Item = T>, | |
| T: Into<OsString> + Clone, | |
| { | |
| run_cli_from(args, benchmark_dir).await | |
| } | |
| fn suite_root() -> tempfile::TempDir { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| fs::write( | |
| temp.path().join("alpha/alpha.suite"), | |
| r#"description = "Alpha benchmark" | |
| [path_replacements] | |
| DATA_DIR = "data" | |
| [[options]] | |
| name = "format" | |
| short = "f" | |
| env = "ALPHA_FORMAT" | |
| default = "parquet" | |
| values = ["parquet", "csv"] | |
| help = "Alpha input format" | |
| [[examples]] | |
| command = "benchmark_runner alpha -q 1 -f csv" | |
| description = "Run query one against CSV data." | |
| "#, | |
| ) | |
| .unwrap(); | |
| temp | |
| } | |
| #[test] | |
| fn suite_help_contains_metadata() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let error = | |
| try_parse_cli_from(["benchmark_runner", "alpha", "--help"], temp.path()) | |
| .unwrap_err(); | |
| let help = error.to_string(); | |
| assert!(help.contains("Alpha benchmark"), "{help}"); | |
| assert!(help.contains("--format"), "{help}"); | |
| assert!(help.contains("ALPHA_FORMAT"), "{help}"); | |
| assert!( | |
| help.contains("benchmark_runner alpha -q 1 -f csv"), | |
| "{help}" | |
| ); | |
| } | |
| #[test] | |
| fn accepts_interleaved_named_options() { | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--partitions", | |
| "2", | |
| "--format", | |
| "csv", | |
| "--query", | |
| "5", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::Simple(config) = action else { | |
| panic!("expected simple run") | |
| }; | |
| assert_eq!(config.common.partitions, Some(2)); | |
| assert_eq!(config.filter.query.as_deref(), Some("5")); | |
| assert_eq!(config.query_filename.as_deref(), Some("q05.benchmark")); | |
| assert_eq!(config.replacements["alpha_format"], "csv"); | |
| assert_eq!( | |
| config.replacements["data_dir"], | |
| temp.path().join("alpha/data").display().to_string() | |
| ); | |
| } | |
| #[test] | |
| fn dry_run_uses_suite_default() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let action = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap(); | |
| let CliAction::DryRun(output) = action else { | |
| panic!("expected dry run") | |
| }; | |
| assert_eq!(output.suite_options["format"].value, "parquet"); | |
| assert_eq!(output.suite_options["format"].source, ValueSource::Default); | |
| } | |
| #[test] | |
| fn result_mode_defaults_to_none() { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", None), | |
| ("BENCH_VALIDATE", None), | |
| ]); | |
| let temp = suite_root(); | |
| let CliAction::DryRun(output) = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap() | |
| else { | |
| panic!("expected dry run"); | |
| }; | |
| assert_eq!(output.result_mode, ResultMode::None); | |
| } | |
| #[tokio::test] | |
| async fn result_mode_persist_writes_expected_results() { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", None), | |
| ("BENCH_VALIDATE", None), | |
| ]); | |
| let temp = suite_root(); | |
| let result_path = temp.path().join("expected.csv"); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| &format!( | |
| "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", | |
| result_path.display() | |
| ), | |
| ); | |
| run_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--query", | |
| "1", | |
| "--result-mode", | |
| "persist", | |
| ], | |
| temp.path(), | |
| ) | |
| .await | |
| .unwrap(); | |
| let persisted = fs::read_to_string(result_path).unwrap(); | |
| assert!(persisted.contains("value"), "{persisted}"); | |
| assert!(persisted.contains('1'), "{persisted}"); | |
| } | |
| #[tokio::test] | |
| async fn result_mode_validate_accepts_expected_results() { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", None), | |
| ("BENCH_VALIDATE", None), | |
| ]); | |
| let temp = suite_root(); | |
| let result_path = temp.path().join("expected.csv"); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| &format!( | |
| "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", | |
| result_path.display() | |
| ), | |
| ); | |
| fs::write(&result_path, "value\n1\n").unwrap(); | |
| run_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--query", | |
| "1", | |
| "--result-mode", | |
| "validate", | |
| ], | |
| temp.path(), | |
| ) | |
| .await | |
| .unwrap(); | |
| } | |
| #[tokio::test] | |
| async fn result_mode_validate_reports_mismatched_results() { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", None), | |
| ("BENCH_VALIDATE", None), | |
| ]); | |
| let temp = suite_root(); | |
| let result_path = temp.path().join("expected.csv"); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| &format!( | |
| "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", | |
| result_path.display() | |
| ), | |
| ); | |
| fs::write(&result_path, "value\n2\n").unwrap(); | |
| let error = run_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--query", | |
| "1", | |
| "--result-mode", | |
| "validate", | |
| ], | |
| temp.path(), | |
| ) | |
| .await | |
| .unwrap_err(); | |
| assert!(error.to_string().contains("expected value"), "{error}"); | |
| } | |
| #[test] | |
| fn explicit_result_modes_populate_config() { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", Some("invalid")), | |
| ("BENCH_VALIDATE", Some("invalid")), | |
| ]); | |
| let temp = suite_root(); | |
| for (value, expected, persist, validate) in [ | |
| ("none", ResultMode::None, false, false), | |
| ("persist", ResultMode::Persist, true, false), | |
| ("validate", ResultMode::Validate, false, true), | |
| ] { | |
| let CliAction::DryRun(output) = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--result-mode", | |
| value, | |
| "--dry-run", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap() else { | |
| panic!("expected dry run"); | |
| }; | |
| assert_eq!(output.result_mode, expected); | |
| let CliAction::Simple(config) = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--result-mode", value], | |
| temp.path(), | |
| ) | |
| .unwrap() else { | |
| panic!("expected simple run"); | |
| }; | |
| assert_eq!(config.persist_results, persist); | |
| assert_eq!(config.validate_results, validate); | |
| } | |
| } | |
| #[test] | |
| fn compatibility_environment_resolves_result_mode() { | |
| for (persist, validate, expected) in [ | |
| (Some("true"), None, ResultMode::Persist), | |
| (None, Some("true"), ResultMode::Validate), | |
| (Some("true"), Some("true"), ResultMode::Persist), | |
| (Some("false"), Some("false"), ResultMode::None), | |
| ] { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", persist), | |
| ("BENCH_VALIDATE", validate), | |
| ]); | |
| let temp = suite_root(); | |
| let CliAction::DryRun(output) = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap() | |
| else { | |
| panic!("expected dry run"); | |
| }; | |
| assert_eq!(output.result_mode, expected); | |
| } | |
| } | |
| #[test] | |
| fn invalid_result_mode_environment_is_rejected_without_cli_override() { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", Some("invalid")), | |
| ("BENCH_VALIDATE", Some("false")), | |
| ]); | |
| let temp = suite_root(); | |
| let error = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap_err(); | |
| assert!( | |
| error.to_string().contains("BENCH_PERSIST_RESULTS"), | |
| "{error}" | |
| ); | |
| } | |
| #[test] | |
| fn invalid_result_mode_cli_value_lists_allowed_values() { | |
| let _env = ScopedEnv::set_many([ | |
| ("BENCH_PERSIST_RESULTS", None), | |
| ("BENCH_VALIDATE", None), | |
| ]); | |
| let temp = suite_root(); | |
| let error = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--result-mode", "invalid"], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| let message = error.to_string(); | |
| for allowed in ["none", "persist", "validate"] { | |
| assert!(message.contains(allowed), "{message}"); | |
| } | |
| } | |
| #[test] | |
| fn environment_beats_suite_default() { | |
| let _env = ScopedEnv::set("ALPHA_FORMAT", "csv"); | |
| let temp = suite_root(); | |
| let action = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap(); | |
| let CliAction::DryRun(output) = action else { | |
| panic!("expected dry run") | |
| }; | |
| assert_eq!(output.suite_options["format"].value, "csv"); | |
| assert_eq!( | |
| output.suite_options["format"].source, | |
| ValueSource::Environment | |
| ); | |
| } | |
| #[test] | |
| fn cli_equals_syntax_beats_environment_and_default() { | |
| let _env = ScopedEnv::set("ALPHA_FORMAT", "parquet"); | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--format=csv", "--dry-run"], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::DryRun(output) = action else { | |
| panic!("expected dry run") | |
| }; | |
| assert_eq!(output.suite_options["format"].value, "csv"); | |
| assert_eq!( | |
| output.suite_options["format"].source, | |
| ValueSource::CommandLine | |
| ); | |
| } | |
| #[test] | |
| fn empty_environment_value_is_validated() { | |
| let _env = ScopedEnv::set("ALPHA_FORMAT", ""); | |
| let temp = suite_root(); | |
| let error = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap_err(); | |
| let message = error.to_string(); | |
| assert!(message.contains("a value is required"), "{message}"); | |
| assert!(message.contains("parquet, csv"), "{message}"); | |
| } | |
| #[cfg(unix)] | |
| #[test] | |
| fn non_unicode_environment_value_is_rejected_by_clap() { | |
| use std::os::unix::ffi::OsStringExt; | |
| let _env = ScopedEnv::set("ALPHA_FORMAT", OsString::from_vec(vec![0xff])); | |
| let temp = suite_root(); | |
| let error = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap_err(); | |
| assert!(error.to_string().contains("invalid UTF-8"), "{error}"); | |
| } | |
| #[test] | |
| fn attached_short_cli_value_beats_invalid_environment() { | |
| let _env = ScopedEnv::set("ALPHA_FORMAT", "invalid"); | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| ["benchmark_runner", "alpha", "-fcsv", "--dry-run"], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::DryRun(output) = action else { | |
| panic!("expected dry run") | |
| }; | |
| assert_eq!(output.suite_options["format"].value, "csv"); | |
| assert_eq!( | |
| output.suite_options["format"].source, | |
| ValueSource::CommandLine | |
| ); | |
| } | |
| #[test] | |
| fn invalid_suite_environment_does_not_block_help() { | |
| let _env = ScopedEnv::set("ALPHA_FORMAT", "invalid"); | |
| let temp = suite_root(); | |
| let error = | |
| try_parse_cli_from(["benchmark_runner", "alpha", "--help"], temp.path()) | |
| .unwrap_err(); | |
| let help = error.to_string(); | |
| assert!(help.contains("Alpha benchmark"), "{help}"); | |
| assert!(help.contains("--format"), "{help}"); | |
| } | |
| #[test] | |
| fn dry_run_rejects_list_instead_of_listing() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let error = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--list", "--dry-run"], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(error.to_string().contains("--list"), "{error}"); | |
| assert!(error.to_string().contains("--dry-run"), "{error}"); | |
| } | |
| #[test] | |
| fn dry_run_requires_suite() { | |
| let temp = suite_root(); | |
| let error = | |
| parse_cli_from(["benchmark_runner", "--dry-run"], temp.path()).unwrap_err(); | |
| assert!(error.to_string().contains("--dry-run"), "{error}"); | |
| assert!(error.to_string().contains("suite"), "{error}"); | |
| } | |
| #[test] | |
| fn dry_run_resolves_default_and_overridden_path() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let default_action = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap(); | |
| let CliAction::DryRun(default_output) = default_action else { | |
| panic!("expected dry run") | |
| }; | |
| assert_eq!( | |
| default_output.path_replacements["data_dir"].source, | |
| ValueSource::Default | |
| ); | |
| let override_action = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--path", | |
| "/tmp/alpha-data", | |
| "--dry-run", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::DryRun(override_output) = override_action else { | |
| panic!("expected dry run") | |
| }; | |
| assert_eq!( | |
| override_output.path_replacements["data_dir"].value, | |
| "/tmp/alpha-data" | |
| ); | |
| assert_eq!( | |
| override_output.path_replacements["data_dir"].source, | |
| ValueSource::CommandLine | |
| ); | |
| } | |
| #[test] | |
| fn path_is_rejected_without_data_dir_replacement() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| fs::write( | |
| temp.path().join("alpha/alpha.suite"), | |
| fs::read_to_string(temp.path().join("alpha/alpha.suite")) | |
| .unwrap() | |
| .replace("[path_replacements]\nDATA_DIR = \"data\"\n\n", ""), | |
| ) | |
| .unwrap(); | |
| let error = | |
| parse_cli_from(["benchmark_runner", "alpha", "--path", "data"], temp.path()) | |
| .unwrap_err(); | |
| assert!(error.to_string().contains("--path"), "{error}"); | |
| assert!(error.to_string().contains("DATA_DIR"), "{error}"); | |
| } | |
| #[test] | |
| fn criterion_dry_run_reports_mode_and_keeps_cross_checks() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--criterion", "--dry-run"], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::DryRun(output) = action else { | |
| panic!("expected dry run") | |
| }; | |
| assert_eq!(output.mode, RunMode::Criterion); | |
| let error = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion", | |
| "--iterations", | |
| "2", | |
| "--dry-run", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(error.to_string().contains("--iterations"), "{error}"); | |
| } | |
| #[test] | |
| fn dry_run_rejects_invalid_query_form() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let error = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--query", | |
| "../secret", | |
| "--dry-run", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!( | |
| error.to_string().contains("invalid query identifier"), | |
| "{error}" | |
| ); | |
| } | |
| #[test] | |
| fn uppercase_q_dry_run_uses_same_query_filename_as_lowercase_q() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| assert!(matches!( | |
| parse_cli_from( | |
| ["benchmark_runner", "alpha", "--query", "Q1", "--dry-run",], | |
| temp.path(), | |
| ) | |
| .unwrap(), | |
| CliAction::DryRun(_) | |
| )); | |
| let filename = |query| { | |
| let CliAction::Simple(config) = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--query", query], | |
| temp.path(), | |
| ) | |
| .unwrap() else { | |
| panic!("expected simple run") | |
| }; | |
| config.query_filename | |
| }; | |
| assert_eq!(filename("Q1"), filename("q1")); | |
| } | |
| #[tokio::test] | |
| async fn dry_run_returns_deterministic_json_without_reading_benchmarks() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| fs::write( | |
| temp.path().join("alpha/benchmarks/q01.benchmark"), | |
| "this benchmark is intentionally invalid", | |
| ) | |
| .unwrap(); | |
| let output = run_cli_with_dir( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--query", | |
| "7", | |
| "--partitions", | |
| "2", | |
| "--dry-run", | |
| ], | |
| temp.path(), | |
| ) | |
| .await | |
| .unwrap(); | |
| let json: serde_json::Value = serde_json::from_str(&output).unwrap(); | |
| assert_eq!(json["suite"], "alpha"); | |
| assert_eq!(json["query"], "7"); | |
| assert_eq!(json["mode"], "simple"); | |
| assert_eq!(json["common_options"]["partitions"], 2); | |
| assert_eq!(json["suite_options"]["format"]["source"], "default"); | |
| } | |
| #[test] | |
| fn cli_lists_when_benchmark_is_omitted() { | |
| let temp = suite_root(); | |
| let action = parse_cli_from(["benchmark_runner"], temp.path()).unwrap(); | |
| assert!(matches!(action, CliAction::List)); | |
| } | |
| #[test] | |
| fn cli_lists_with_explicit_list_flag() { | |
| let temp = suite_root(); | |
| let action = parse_cli_from(["benchmark_runner", "--list"], temp.path()).unwrap(); | |
| assert!(matches!(action, CliAction::List)); | |
| } | |
| #[test] | |
| fn cli_defaults_to_basic_runner() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let action = | |
| parse_cli_from(["benchmark_runner", "alpha", "--query", "1"], temp.path()) | |
| .unwrap(); | |
| let CliAction::Simple(config) = action else { | |
| panic!("expected basic runner"); | |
| }; | |
| assert_eq!(config.filter.name.as_deref(), Some("alpha")); | |
| assert_eq!(config.filter.query.as_deref(), Some("1")); | |
| } | |
| #[test] | |
| fn cli_reads_query_from_env() { | |
| let _env = ScopedEnv::set("BENCH_QUERY", "8"); | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--format", "parquet"], | |
| temp.path(), | |
| ); | |
| let action = action.unwrap(); | |
| let CliAction::Simple(config) = action else { | |
| panic!("expected basic runner"); | |
| }; | |
| assert_eq!(config.filter.name.as_deref(), Some("alpha")); | |
| assert_eq!(config.filter.query.as_deref(), Some("8")); | |
| } | |
| #[test] | |
| fn cli_accepts_criterion_runner() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion", | |
| "--save-baseline", | |
| "main", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::Criterion { | |
| config, | |
| save_baseline, | |
| criterion_namespace, | |
| } = action | |
| else { | |
| panic!("expected criterion runner"); | |
| }; | |
| assert_eq!(config.filter.name.as_deref(), Some("alpha")); | |
| assert_eq!(save_baseline.as_deref(), Some("main")); | |
| assert_eq!(criterion_namespace, None); | |
| } | |
| #[test] | |
| fn cli_accepts_safe_criterion_namespace() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion", | |
| "--criterion-namespace", | |
| "parquet_sf1-ordered", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::Criterion { | |
| criterion_namespace, | |
| .. | |
| } = action | |
| else { | |
| panic!("expected criterion runner"); | |
| }; | |
| assert_eq!(criterion_namespace.as_deref(), Some("parquet_sf1-ordered")); | |
| } | |
| #[test] | |
| fn cli_reads_criterion_namespace_from_env() { | |
| if std::env::var(CLI_ENV_TEST_CHILD).as_deref() == Ok("namespace") { | |
| let temp = suite_root(); | |
| let action = | |
| parse_cli_from(["benchmark_runner", "alpha", "--criterion"], temp.path()) | |
| .unwrap(); | |
| let CliAction::Criterion { | |
| criterion_namespace, | |
| .. | |
| } = action | |
| else { | |
| panic!("expected criterion runner"); | |
| }; | |
| assert_eq!(criterion_namespace.as_deref(), Some("csv_sf1")); | |
| let action = | |
| parse_cli_from(["benchmark_runner", "alpha"], temp.path()).unwrap(); | |
| assert!( | |
| matches!(action, CliAction::Criterion { criterion_namespace: Some(ref namespace), .. } if namespace == "csv_sf1") | |
| ); | |
| let action = | |
| parse_cli_from(["benchmark_runner", "--list"], temp.path()).unwrap(); | |
| assert!(matches!(action, CliAction::List)); | |
| let action = | |
| parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) | |
| .unwrap(); | |
| assert!( | |
| matches!(action, CliAction::DryRun(output) if matches!(output.mode, RunMode::Criterion)) | |
| ); | |
| let action = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion-namespace", | |
| "parquet_sf1", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| assert!( | |
| matches!(action, CliAction::Criterion { criterion_namespace: Some(ref namespace), .. } if namespace == "parquet_sf1") | |
| ); | |
| for options in [["--output", "results.json"], ["--iterations", "2"]] { | |
| let error = parse_cli_from( | |
| ["benchmark_runner", "alpha"].into_iter().chain(options), | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(error.to_string().contains(options[0]), "{error}"); | |
| assert!(error.to_string().contains("--criterion"), "{error}"); | |
| } | |
| return; | |
| } | |
| let output = std::process::Command::new(std::env::current_exe().unwrap()) | |
| .args([ | |
| "--exact", | |
| "tests::cli_reads_criterion_namespace_from_env", | |
| "--nocapture", | |
| ]) | |
| .env(CLI_ENV_TEST_CHILD, "namespace") | |
| .env("BENCH_NAMESPACE", "csv_sf1") | |
| .env("ALPHA_FORMAT", "parquet") | |
| .output() | |
| .unwrap(); | |
| assert!( | |
| output.status.success(), | |
| "child failed:\n{}", | |
| String::from_utf8_lossy(&output.stderr) | |
| ); | |
| assert!( | |
| String::from_utf8_lossy(&output.stderr).contains("Enabling Criterion"), | |
| "missing inferred-mode notice on stderr" | |
| ); | |
| assert!( | |
| !String::from_utf8_lossy(&output.stdout).contains("Enabling Criterion"), | |
| "inferred-mode notice must not contaminate stdout" | |
| ); | |
| } | |
| #[test] | |
| fn cli_criterion_namespace_enables_criterion() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let action = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion-namespace", | |
| "parquet", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| let CliAction::Criterion { | |
| criterion_namespace, | |
| .. | |
| } = action | |
| else { | |
| panic!("expected criterion runner"); | |
| }; | |
| assert_eq!(criterion_namespace.as_deref(), Some("parquet")); | |
| } | |
| #[test] | |
| fn cli_criterion_namespace_preserves_mode_checks() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| for options in [["--output", "results.json"], ["--iterations", "2"]] { | |
| let error = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion-namespace", | |
| "parquet", | |
| ] | |
| .into_iter() | |
| .chain(options), | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(error.to_string().contains(options[0]), "{error}"); | |
| assert!(error.to_string().contains("--criterion"), "{error}"); | |
| } | |
| let action = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion-namespace", | |
| "parquet", | |
| "--save-baseline", | |
| "baseline", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap(); | |
| assert!( | |
| matches!(action, CliAction::Criterion { save_baseline: Some(ref name), .. } if name == "baseline") | |
| ); | |
| } | |
| #[test] | |
| fn cli_rejects_unsafe_criterion_namespaces() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| for namespace in ["", "Parquet", "csv/sf1", "csv sf1", "csv.sf1", "parquét"] { | |
| let error = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion", | |
| "--criterion-namespace", | |
| namespace, | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(error.to_string().contains("namespace"), "{error}"); | |
| } | |
| } | |
| #[test] | |
| fn cli_rejects_output_with_criterion() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let err = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion", | |
| "--output", | |
| "results.json", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(err.to_string().contains("--output")); | |
| assert!(err.to_string().contains("--criterion")); | |
| } | |
| #[test] | |
| fn cli_rejects_save_baseline_without_criterion() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let err = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--save-baseline", "main"], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(err.to_string().contains("--save-baseline")); | |
| assert!(err.to_string().contains("--criterion")); | |
| } | |
| #[test] | |
| fn cli_rejects_iterations_with_criterion() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let err = parse_cli_from( | |
| [ | |
| "benchmark_runner", | |
| "alpha", | |
| "--criterion", | |
| "--iterations", | |
| "3", | |
| ], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(err.to_string().contains("--iterations")); | |
| assert!(err.to_string().contains("--criterion")); | |
| } | |
| #[test] | |
| fn cli_rejects_zero_basic_iterations() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let err = parse_cli_from( | |
| ["benchmark_runner", "alpha", "--iterations", "0"], | |
| temp.path(), | |
| ) | |
| .unwrap_err(); | |
| assert!(err.to_string().contains("iterations")); | |
| } | |
| #[tokio::test] | |
| async fn run_cli_lists_when_no_benchmark_is_supplied() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| write_suite(temp.path(), "alpha", "Alpha workload"); | |
| let output = run_cli_with_dir(["benchmark_runner"], temp.path()) | |
| .await | |
| .unwrap(); | |
| assert!(output.contains("SQL benchmarks:")); | |
| assert!(output.contains("alpha")); | |
| } | |
| #[tokio::test] | |
| async fn run_cli_lists_with_explicit_list_flag() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| write_suite(temp.path(), "alpha", "Alpha workload"); | |
| let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) | |
| .await | |
| .unwrap(); | |
| assert!(output.contains("SQL benchmarks:")); | |
| assert!(output.contains("alpha")); | |
| } | |
| #[tokio::test] | |
| async fn run_cli_top_level_help_is_successful_output() { | |
| let temp = suite_root(); | |
| let output = run_cli_with_dir(["benchmark_runner", "--help"], temp.path()) | |
| .await | |
| .unwrap(); | |
| assert!(output.contains("Run DataFusion SQL benchmarks"), "{output}"); | |
| assert!(output.contains("Usage:"), "{output}"); | |
| } | |
| #[tokio::test] | |
| async fn run_cli_suite_help_is_successful_output() { | |
| let _env = ScopedEnv::remove("ALPHA_FORMAT"); | |
| let temp = suite_root(); | |
| let output = | |
| run_cli_with_dir(["benchmark_runner", "alpha", "--help"], temp.path()) | |
| .await | |
| .unwrap(); | |
| assert!(output.contains("Alpha benchmark"), "{output}"); | |
| assert!(output.contains("--format"), "{output}"); | |
| } | |
| #[tokio::test] | |
| async fn run_cli_real_parse_error_remains_an_error() { | |
| let temp = suite_root(); | |
| let error = run_cli_with_dir( | |
| ["benchmark_runner", "alpha", "--not-an-option"], | |
| temp.path(), | |
| ) | |
| .await | |
| .unwrap_err(); | |
| assert!(error.to_string().contains("unexpected argument"), "{error}"); | |
| } | |
| #[tokio::test] | |
| async fn run_cli_reports_unknown_benchmark_with_list() { | |
| let temp = suite_root(); | |
| let err = run_cli_with_dir(["benchmark_runner", "missing"], temp.path()) | |
| .await | |
| .unwrap_err(); | |
| let message = err.to_string(); | |
| assert!(message.contains("unknown benchmark 'missing'"), "{message}"); | |
| assert!(message.contains("alpha"), "{message}"); | |
| } | |
| #[test] | |
| fn sql_runner_debug_plans_include_execution_metrics() { | |
| const CHILD_ENV: &str = "DATAFUSION_SQL_DEBUG_TEST_MODE"; | |
| if let Ok(mode) = std::env::var(CHILD_ENV) { | |
| let temp = suite_root(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT value FROM (VALUES (3), (1), (2)) AS t(value) ORDER BY value\n", | |
| ); | |
| let mut args = vec!["benchmark_runner", "alpha", "--query", "1"]; | |
| if mode == "criterion" { | |
| args.push("--criterion"); | |
| } else { | |
| args.extend(["--iterations", "1"]); | |
| } | |
| if mode == "buffered" { | |
| args.extend(["--result-mode", "persist"]); | |
| } | |
| match parse_cli_from(args, temp.path()).unwrap() { | |
| CliAction::Simple(config) => { | |
| tokio::runtime::Runtime::new() | |
| .unwrap() | |
| .block_on(run_simple_benchmarks(temp.path(), config)) | |
| .unwrap(); | |
| } | |
| CliAction::Criterion { config, .. } => { | |
| let output = tempfile::tempdir().unwrap(); | |
| let mut criterion = Criterion::default() | |
| .output_directory(output.path()) | |
| .warm_up_time(std::time::Duration::from_millis(1)) | |
| .measurement_time(std::time::Duration::from_millis(10)) | |
| .without_plots(); | |
| run_criterion_benchmarks_impl(temp.path(), &config, &mut criterion) | |
| .unwrap(); | |
| } | |
| _ => panic!("expected a benchmark run"), | |
| } | |
| return; | |
| } | |
| for mode in ["simple", "criterion", "buffered"] { | |
| for debug in [false, true] { | |
| let output = std::process::Command::new(std::env::current_exe().unwrap()) | |
| .args([ | |
| "--exact", | |
| "tests::sql_runner_debug_plans_include_execution_metrics", | |
| "--nocapture", | |
| ]) | |
| // Other CLI tests mutate environment variables concurrently. | |
| .env_clear() | |
| // Keep subprocess coverage output in cargo-llvm-cov's directory. | |
| .envs( | |
| std::env::var_os("LLVM_PROFILE_FILE") | |
| .map(|value| ("LLVM_PROFILE_FILE", value)), | |
| ) | |
| .env(CHILD_ENV, mode) | |
| .env("DEBUG", debug.to_string()) | |
| .output() | |
| .unwrap(); | |
| let stdout = String::from_utf8(output.stdout).unwrap(); | |
| assert!( | |
| output.status.success(), | |
| "{mode}, debug={debug}: {stdout}\n{}", | |
| String::from_utf8_lossy(&output.stderr) | |
| ); | |
| for heading in [ | |
| "=== Logical plan ===", | |
| "=== Optimized logical plan ===", | |
| "=== Physical plan ===", | |
| "=== Physical plan with metrics ===", | |
| ] { | |
| assert_eq!(stdout.contains(heading), debug, "{mode}: {stdout}"); | |
| } | |
| if debug { | |
| let metrics = stdout | |
| .split("=== Physical plan with metrics ===") | |
| .nth(1) | |
| .unwrap(); | |
| assert!(metrics.contains("output_rows=3"), "{mode}: {stdout}"); | |
| } | |
| } | |
| } | |
| } | |
| #[test] | |
| fn criterion_runner_saves_named_baseline() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| let output = tempfile::tempdir().unwrap(); | |
| let mut criterion = Criterion::default() | |
| .sample_size(10) | |
| .warm_up_time(std::time::Duration::from_millis(1)) | |
| .measurement_time(std::time::Duration::from_millis(10)) | |
| .without_plots() | |
| .output_directory(output.path()) | |
| .save_baseline("acceptance".to_string()); | |
| let config = SqlRunConfig { | |
| common: common(3), | |
| filter: BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: None, | |
| query: Some("1".to_string()), | |
| }, | |
| replacements: HashMap::new(), | |
| query_filename: None, | |
| persist_results: false, | |
| validate_results: false, | |
| output: None, | |
| }; | |
| run_criterion_benchmarks_impl(temp.path(), &config, &mut criterion).unwrap(); | |
| criterion.final_summary(); | |
| assert!( | |
| output | |
| .path() | |
| .join("alpha") | |
| .join("Q01") | |
| .join("acceptance") | |
| .join("estimates.json") | |
| .exists() | |
| ); | |
| } | |
| #[test] | |
| fn criterion_namespaces_write_distinct_artifacts_across_fresh_instances() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| let output = tempfile::tempdir().unwrap(); | |
| let config = SqlRunConfig { | |
| common: common(3), | |
| filter: BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: None, | |
| query: Some("1".to_string()), | |
| }, | |
| replacements: HashMap::new(), | |
| query_filename: None, | |
| persist_results: false, | |
| validate_results: false, | |
| output: None, | |
| }; | |
| for namespace in ["parquet", "memory"] { | |
| let mut criterion = Criterion::default() | |
| .sample_size(10) | |
| .warm_up_time(std::time::Duration::from_millis(1)) | |
| .measurement_time(std::time::Duration::from_millis(10)) | |
| .without_plots() | |
| .output_directory(output.path()) | |
| .save_baseline("acceptance".to_string()); | |
| run_criterion_benchmarks_impl_with_namespace( | |
| temp.path(), | |
| &config, | |
| Some(namespace), | |
| &mut criterion, | |
| ) | |
| .unwrap(); | |
| criterion.final_summary(); | |
| } | |
| for group in ["alpha__parquet", "alpha__memory"] { | |
| assert!( | |
| output | |
| .path() | |
| .join(group) | |
| .join("Q01") | |
| .join("acceptance") | |
| .join("estimates.json") | |
| .exists(), | |
| "missing Criterion artifact for {group}" | |
| ); | |
| } | |
| } | |
| #[tokio::test] | |
| async fn simple_runner_reports_unknown_query_for_known_benchmark() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| let config = SqlRunConfig { | |
| common: common(1), | |
| filter: BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: None, | |
| query: Some("9".to_string()), | |
| }, | |
| replacements: HashMap::new(), | |
| query_filename: None, | |
| persist_results: false, | |
| validate_results: false, | |
| output: None, | |
| }; | |
| let err = run_simple_benchmarks(temp.path(), config) | |
| .await | |
| .unwrap_err(); | |
| let message = err.to_string(); | |
| assert!( | |
| message.contains("no SQL benchmark query matched benchmark 'alpha'"), | |
| "{message}" | |
| ); | |
| assert!(message.contains("query '9'"), "{message}"); | |
| assert!(message.contains("normalized: 'Q09'"), "{message}"); | |
| assert!(message.contains("Available alpha queries:"), "{message}"); | |
| assert!(message.contains("Q01"), "{message}"); | |
| assert!(!message.contains("unknown benchmark 'alpha'"), "{message}"); | |
| } | |
| #[tokio::test] | |
| async fn simple_runner_reports_unknown_subgroup_for_known_benchmark() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\nsubgroup wide\n\nrun\nSELECT 1\n", | |
| ); | |
| let config = SqlRunConfig { | |
| common: common(1), | |
| filter: BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: Some("narrow".to_string()), | |
| query: None, | |
| }, | |
| replacements: HashMap::new(), | |
| query_filename: None, | |
| persist_results: false, | |
| validate_results: false, | |
| output: None, | |
| }; | |
| let err = run_simple_benchmarks(temp.path(), config) | |
| .await | |
| .unwrap_err(); | |
| let message = err.to_string(); | |
| assert!( | |
| message.contains( | |
| "no SQL benchmark subgroup matched benchmark 'alpha' with subgroup 'narrow'" | |
| ), | |
| "{message}" | |
| ); | |
| assert!(message.contains("Available alpha subgroups:"), "{message}"); | |
| assert!(message.contains("wide"), "{message}"); | |
| assert!(!message.contains("unknown benchmark 'alpha'"), "{message}"); | |
| } | |
| #[tokio::test] | |
| async fn basic_runner_executes_iterations_and_writes_json() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| let output = temp.path().join("results.json"); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT * FROM (VALUES (1), (2)) AS t(v)\n", | |
| ); | |
| let config = SqlRunConfig { | |
| common: common(2), | |
| filter: BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: None, | |
| query: Some("1".to_string()), | |
| }, | |
| replacements: HashMap::new(), | |
| query_filename: None, | |
| persist_results: false, | |
| validate_results: false, | |
| output: Some(output.clone()), | |
| }; | |
| run_simple_benchmarks(temp.path(), config).await.unwrap(); | |
| let json = fs::read_to_string(output).unwrap(); | |
| assert!(json.contains("\"query\": \"alpha/Q01\"")); | |
| assert!(json.contains("\"row_count\": 2")); | |
| assert_eq!(json.matches("\"row_count\": 2").count(), 2); | |
| } | |
| #[tokio::test] | |
| async fn basic_runner_reports_run_and_cleanup_failures() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT * FROM missing_run_table\n\ncleanup\nSELECT * FROM missing_cleanup_table\n", | |
| ); | |
| let config = SqlRunConfig { | |
| common: common(1), | |
| filter: BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: None, | |
| query: Some("1".to_string()), | |
| }, | |
| replacements: HashMap::new(), | |
| query_filename: None, | |
| persist_results: false, | |
| validate_results: false, | |
| output: None, | |
| }; | |
| let err = run_simple_benchmarks(temp.path(), config) | |
| .await | |
| .unwrap_err(); | |
| let message = err.to_string(); | |
| assert!(message.contains("missing_run_table"), "{message}"); | |
| assert!(message.contains("cleanup also failed"), "{message}"); | |
| assert!(message.contains("missing_cleanup_table"), "{message}"); | |
| } | |
| #[tokio::test] | |
| async fn discovery_lists_groups_from_directories() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| write_benchmark( | |
| temp.path(), | |
| "beta/benchmarks/q02.benchmark", | |
| "name Q02\n\nrun\nSELECT 2\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let benches = load_benchmarks(&BenchmarkFilter::default(), &ctx, temp.path()) | |
| .await | |
| .unwrap(); | |
| assert_eq!(benches["alpha"].len(), 1); | |
| assert_eq!(benches["beta"].len(), 1); | |
| } | |
| #[tokio::test] | |
| async fn discovery_filters_benchmark_subgroup_and_query() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\nsubgroup wide\n\nrun\nSELECT 1\n", | |
| ); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q02.benchmark", | |
| "name Q02\nsubgroup narrow\n\nrun\nSELECT 2\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let filter = BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: Some("wide".to_string()), | |
| query: Some("1".to_string()), | |
| }; | |
| let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); | |
| assert_eq!(benches.len(), 1); | |
| assert_eq!(benches["alpha"].len(), 1); | |
| assert_eq!(benches["alpha"][0].name(), "Q01"); | |
| } | |
| #[tokio::test] | |
| #[expect( | |
| clippy::literal_string_with_formatting_args, | |
| reason = "The `${VAR:-default}` braces are shell-style placeholders, not format args" | |
| )] | |
| async fn cli_subgroup_filter_is_used_for_benchmark_replacements() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "wide_schema/benchmarks/q01.benchmark", | |
| "name Q01\nsubgroup ${BENCH_SUBGROUP:-wide}\n\nrun\nSELECT '${BENCH_SUBGROUP:-wide}'\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let filter = BenchmarkFilter { | |
| name: Some("wide_schema".to_string()), | |
| subgroup: Some("narrow".to_string()), | |
| query: None, | |
| }; | |
| let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); | |
| assert_eq!(benches["wide_schema"].len(), 1); | |
| assert_eq!(benches["wide_schema"][0].subgroup(), "narrow"); | |
| } | |
| #[tokio::test] | |
| #[expect( | |
| clippy::literal_string_with_formatting_args, | |
| reason = "The `${VAR:-default}` braces are shell-style placeholders, not format args" | |
| )] | |
| async fn benchmark_replacements_use_explicit_data_dir() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "clickbench/benchmarks/q01.benchmark", | |
| "name Q01\nsubgroup ${DATA_DIR:-data}\n\nrun\nSELECT 1\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let expected = PathBuf::from(env!("CARGO_MANIFEST_DIR")) | |
| .join("data") | |
| .to_string_lossy() | |
| .into_owned(); | |
| let replacements = HashMap::from([("data_dir".to_string(), expected.clone())]); | |
| let benches = load_benchmark_definitions( | |
| &BenchmarkFilter::default(), | |
| &ctx, | |
| temp.path(), | |
| &replacements, | |
| ) | |
| .await | |
| .unwrap(); | |
| assert_eq!(benches["clickbench"][0].subgroup(), expected); | |
| } | |
| #[tokio::test] | |
| async fn query_filter_matches_starts_with_when_exact_match_is_absent() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01a.benchmark", | |
| "name Q01a\n\nrun\nSELECT 1\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let filter = BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: None, | |
| query: Some("1".to_string()), | |
| }; | |
| let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); | |
| assert_eq!(benches["alpha"].len(), 1); | |
| assert_eq!(benches["alpha"][0].name(), "Q01a"); | |
| } | |
| #[tokio::test] | |
| async fn query_filter_matches_token_start_when_exact_match_is_absent() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "predicate_eval/benchmarks/costsel/q01.benchmark", | |
| "name costsel_q01_regexp_selective_last\n\nrun\nSELECT 1\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let filter = BenchmarkFilter { | |
| name: Some("predicate_eval".to_string()), | |
| subgroup: None, | |
| query: Some("1".to_string()), | |
| }; | |
| let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); | |
| assert_eq!(benches["predicate_eval"].len(), 1); | |
| assert_eq!( | |
| benches["predicate_eval"][0].name(), | |
| "costsel_q01_regexp_selective_last" | |
| ); | |
| } | |
| #[tokio::test] | |
| async fn query_filter_prefers_starts_with_match_over_token_match() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/token.benchmark", | |
| "name costsel_q01_regexp_selective_last\n\nrun\nSELECT 1\n", | |
| ); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01a.benchmark", | |
| "name Q01a\n\nrun\nSELECT 2\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let filter = BenchmarkFilter { | |
| name: Some("alpha".to_string()), | |
| subgroup: None, | |
| query: Some("1".to_string()), | |
| }; | |
| let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); | |
| assert_eq!(benches["alpha"].len(), 1); | |
| assert_eq!(benches["alpha"][0].name(), "Q01a"); | |
| } | |
| #[tokio::test] | |
| async fn list_output_is_sorted_and_includes_counts_and_descriptions() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "beta/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| write_benchmark( | |
| temp.path(), | |
| "beta/benchmarks/q02.benchmark", | |
| "name Q02\n\nrun\nSELECT 2\n", | |
| ); | |
| write_suite(temp.path(), "alpha", "Alpha workload"); | |
| write_suite(temp.path(), "beta", "Beta workload"); | |
| let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) | |
| .await | |
| .unwrap(); | |
| assert_eq!( | |
| output, | |
| "SQL benchmarks:\n alpha 1 query Alpha workload\n beta 2 queries Beta workload" | |
| ); | |
| } | |
| #[tokio::test] | |
| async fn list_does_not_parse_benchmark_sql() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "not valid benchmark syntax", | |
| ); | |
| write_suite(temp.path(), "alpha", "Alpha workload"); | |
| let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) | |
| .await | |
| .unwrap(); | |
| assert_eq!( | |
| output, | |
| "SQL benchmarks:\n alpha 1 query Alpha workload" | |
| ); | |
| } | |
| #[tokio::test] | |
| async fn list_malformed_metadata_names_its_file() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| let metadata_path = temp.path().join("alpha/alpha.suite"); | |
| fs::write(&metadata_path, "not valid metadata").unwrap(); | |
| let error = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) | |
| .await | |
| .unwrap_err(); | |
| assert!( | |
| error | |
| .to_string() | |
| .contains(&metadata_path.display().to_string()), | |
| "{error}" | |
| ); | |
| } | |
| #[tokio::test] | |
| async fn unknown_benchmark_error_includes_available_benchmarks() { | |
| let temp = tempfile::tempdir().unwrap(); | |
| write_benchmark( | |
| temp.path(), | |
| "alpha/benchmarks/q01.benchmark", | |
| "name Q01\n\nrun\nSELECT 1\n", | |
| ); | |
| let ctx = SessionContext::new(); | |
| let benches = load_benchmarks(&BenchmarkFilter::default(), &ctx, temp.path()) | |
| .await | |
| .unwrap(); | |
| let message = unknown_benchmark_error("missing", &benches).to_string(); | |
| assert!(message.contains("unknown benchmark 'missing'"), "{message}"); | |
| assert!(message.contains("alpha"), "{message}"); | |
| } | |
| } |