From 2288cfc3500785af8de58b72aa1fd99c3803d826 Mon Sep 17 00:00:00 2001 From: Bruce Ritchie Date: Sun, 30 Aug 2026 17:10:43 -0400 Subject: [PATCH] Prepare SQL runner for native benchmark migration: - Preserve the last result-producing query's row count when trailing DDL is executed, keeping fixed-iteration output compatible with native benchmarks. - Add collision-safe Criterion namespaces to separate results for different formats, scale factors, and configurations. Share environment parsing between the direct Criterion harness and benchmark runner so both entry points support the same behaviour. --- benchmarks/benches/sql.rs | 62 ++----- benchmarks/src/bin/benchmark_runner.rs | 219 ++++++++++++++++++++++++- benchmarks/src/sql_benchmark.rs | 31 +++- benchmarks/src/sql_benchmark_runner.rs | 195 ++++++++++++++++++++++ 4 files changed, 444 insertions(+), 63 deletions(-) diff --git a/benchmarks/benches/sql.rs b/benchmarks/benches/sql.rs index 9240a19470db..466b4b03e86b 100644 --- a/benchmarks/benches/sql.rs +++ b/benchmarks/benches/sql.rs @@ -21,13 +21,11 @@ //! `.benchmark` files. Run them with `benchmarks/bench.sh` or directly with //! Cargo, for example: `BENCH_NAME=tpch cargo bench --bench sql`. -use clap::Parser; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_benchmarks::sql_benchmark_runner::{ - BenchmarkFilter, SqlRunConfig, default_criterion_replacements, - default_sql_benchmark_directory, run_criterion_benchmarks_impl, + criterion_harness_config_from_env, default_sql_benchmark_directory, + run_criterion_benchmarks_impl_with_namespace, }; -use datafusion_benchmarks::util::CommonOpt; use datafusion_common::instant::Instant; #[cfg(feature = "snmalloc")] @@ -40,61 +38,21 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; #[global_allocator] static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; -#[derive(Debug, Parser)] -#[command(ignore_errors = true)] -struct EnvParser { - #[command(flatten)] - options: CommonOpt, - - #[arg( - env = "BENCH_PERSIST_RESULTS", - long = "persist_results", - default_value = "false", - action = clap::ArgAction::SetTrue - )] - persist_results: bool, - - #[arg( - env = "BENCH_VALIDATE", - long = "validate_results", - default_value = "false", - action = clap::ArgAction::SetTrue - )] - validate: bool, - - #[arg(env = "BENCH_NAME")] - name: Option, - - #[arg(env = "BENCH_SUBGROUP")] - subgroup: Option, - - #[arg(env = "BENCH_QUERY")] - query: Option, -} - pub fn sql(c: &mut Criterion) { env_logger::init(); let start = Instant::now(); - let args = EnvParser::parse(); - let config = SqlRunConfig { - common: args.options, - filter: BenchmarkFilter { - name: args.name, - subgroup: args.subgroup, - query: args.query, - }, - replacements: default_criterion_replacements(), - query_filename: None, - persist_results: args.persist_results, - validate_results: args.validate, - output: None, - }; + let (config, criterion_namespace) = criterion_harness_config_from_env(); println!("Loading benchmarks..."); - run_criterion_benchmarks_impl(&default_sql_benchmark_directory(), &config, c) - .unwrap_or_else(|err| panic!("failed to run SQL benchmarks: {err:?}")); + run_criterion_benchmarks_impl_with_namespace( + &default_sql_benchmark_directory(), + &config, + criterion_namespace.as_deref(), + c, + ) + .unwrap_or_else(|err| panic!("failed to run SQL benchmarks: {err:?}")); println!( "Completed benchmarks in {} ms ...", diff --git a/benchmarks/src/bin/benchmark_runner.rs b/benchmarks/src/bin/benchmark_runner.rs index fc980a63bf26..0ff442b51b84 100644 --- a/benchmarks/src/bin/benchmark_runner.rs +++ b/benchmarks/src/bin/benchmark_runner.rs @@ -28,7 +28,7 @@ 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, + prepare_benchmark, run_criterion_benchmarks_impl_with_namespace, }; use datafusion_benchmarks::sql_benchmark_suite::{ ReservedOptions, SuiteExample, SuiteMetadata, discover_suites, @@ -69,10 +69,28 @@ enum CliAction { Criterion { config: SqlRunConfig, save_baseline: Option, + criterion_namespace: Option, }, DryRun(DryRunOutput), } +fn parse_criterion_namespace(value: &str) -> std::result::Result { + if value.is_empty() + || !value.chars().all(|character| { + character.is_ascii_lowercase() + || character.is_ascii_digit() + || "_-".contains(character) + }) + { + return Err( + "namespace must be nonempty and contain only lowercase ASCII letters, digits, '_', or '-'" + .to_string(), + ); + } + + Ok(value.to_string()) +} + #[derive(Debug, Serialize)] struct ResolvedSuiteValue { value: String, @@ -180,6 +198,15 @@ struct Cli { )] save_baseline: Option, + #[arg( + long = "criterion-namespace", + env = "BENCH_NAMESPACE", + value_name = "NAMESPACE", + value_parser = parse_criterion_namespace, + help = "Append a safe namespace to Criterion benchmark groups" + )] + criterion_namespace: Option, + #[arg(short = 'p', long = "path", value_name = "PATH")] path: Option, @@ -418,9 +445,6 @@ fn cli_action_from_matches( if cli.dry_run && cli.benchmark.is_none() { return Err(exec_datafusion_err!("--dry-run requires a benchmark suite")); } - if cli.list || cli.benchmark.is_none() { - return Ok(CliAction::List); - } if cli.criterion && cli.output.is_some() { return Err(exec_datafusion_err!( "--output cannot be used with --criterion" @@ -431,6 +455,14 @@ fn cli_action_from_matches( "--save-baseline cannot be used without --criterion" )); } + if !cli.criterion && cli.criterion_namespace.is_some() { + return Err(exec_datafusion_err!( + "--criterion-namespace 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")); } @@ -579,6 +611,7 @@ fn cli_action_from_matches( Ok(CliAction::Criterion { config, save_baseline: cli.save_baseline, + criterion_namespace: cli.criterion_namespace, }) } else { Ok(CliAction::Simple(config)) @@ -596,6 +629,7 @@ async fn run_cli_action(action: CliAction, benchmark_dir: &Path) -> Result { if config.output.is_some() { return Err(exec_datafusion_err!( @@ -609,6 +643,7 @@ async fn run_cli_action(action: CliAction, benchmark_dir: &Path) -> Result, + criterion_namespace: Option<&str>, ) -> Result<()> { let mut criterion = Criterion::default() .sample_size(10) @@ -665,7 +701,12 @@ fn run_criterion_benchmarks( criterion = criterion.save_baseline(save_baseline.to_string()); } - run_criterion_benchmarks_impl(benchmark_dir, config, &mut criterion)?; + run_criterion_benchmarks_impl_with_namespace( + benchmark_dir, + config, + criterion_namespace, + &mut criterion, + )?; criterion.final_summary(); Ok(()) @@ -766,6 +807,7 @@ fn criterion_like_styles() -> clap::builder::Styles { #[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, }; @@ -776,6 +818,8 @@ mod tests { 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 { @@ -1564,6 +1608,7 @@ description = "Run query one against CSV data." let CliAction::Criterion { config, save_baseline, + criterion_namespace, } = action else { panic!("expected criterion runner"); @@ -1571,6 +1616,112 @@ description = "Run query one against CSV data." 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")); + 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) + ); + } + + #[test] + fn cli_rejects_criterion_namespace_without_criterion() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion-namespace", + "parquet", + ], + temp.path(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("--criterion-namespace")); + assert!(err.to_string().contains("--criterion")); + } + + #[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] @@ -1774,6 +1925,64 @@ description = "Run query one against CSV data." ); } + #[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(); diff --git a/benchmarks/src/sql_benchmark.rs b/benchmarks/src/sql_benchmark.rs index 24db7e0a0fb2..0c0c9375e428 100644 --- a/benchmarks/src/sql_benchmark.rs +++ b/benchmarks/src/sql_benchmark.rs @@ -246,12 +246,8 @@ impl SqlBenchmark { let result_schema = Arc::new(df.schema().as_arrow().clone()); let mut batches = df.collect().await?; - let trimmed = query.trim_start(); - // save the output for select/with queries - if starts_with_ignore_ascii_case(trimmed, "select") - || starts_with_ignore_ascii_case(trimmed, "with") - { + if is_result_statement(query) { if batches.is_empty() { batches.push(RecordBatch::new_empty(result_schema)); } @@ -273,9 +269,13 @@ impl SqlBenchmark { self.group, self.subgroup ); - result_count = self + let row_count = self .execute_sql_without_result_buffering(query, ctx) .await?; + + if is_result_statement(query) { + result_count = row_count; + } } } } @@ -1309,6 +1309,12 @@ fn starts_with_ignore_ascii_case(input: &str, prefix: &str) -> bool { .is_some_and(|value| value.eq_ignore_ascii_case(prefix)) } +fn is_result_statement(statement: &str) -> bool { + let statement = statement.trim_start(); + starts_with_ignore_ascii_case(statement, "select") + || starts_with_ignore_ascii_case(statement, "with") +} + fn split_query_statements(sql: &str) -> impl Iterator { sql.split("\n\n") .flat_map(|query| { @@ -1688,6 +1694,19 @@ mod tests { assert_eq!(row_count, 3); } + #[tokio::test] + async fn streaming_run_reports_last_select_row_count_after_ddl() { + let contents = "name Q15\n\nrun\nCREATE VIEW v AS SELECT 1 AS value;\nSELECT * FROM v;\nDROP VIEW v;\n"; + let mut benchmark = parse_benchmark(contents).await.unwrap(); + let ctx = SessionContext::new(); + + benchmark.initialize(&ctx).await.unwrap(); + let row_count = benchmark.run(&ctx, false).await.unwrap(); + + assert_eq!(row_count, 1); + assert!(ctx.table("v").await.is_err(), "trailing DDL should execute"); + } + #[tokio::test] async fn run_returns_row_count_when_saving_results() { let contents = "name Q01\n\nrun\nSELECT * FROM (VALUES (1), (2)) AS t(v)\n"; diff --git a/benchmarks/src/sql_benchmark_runner.rs b/benchmarks/src/sql_benchmark_runner.rs index 550c1d383f51..175a9de996f3 100644 --- a/benchmarks/src/sql_benchmark_runner.rs +++ b/benchmarks/src/sql_benchmark_runner.rs @@ -20,6 +20,7 @@ use crate::sql_benchmark::SqlBenchmark; use crate::util::{CommonOpt, print_memory_stats}; +use clap::Parser; use criterion::{Criterion, SamplingMode}; use datafusion::error::Result; use datafusion::prelude::SessionContext; @@ -31,6 +32,8 @@ use std::panic::{AssertUnwindSafe, catch_unwind}; use std::path::{Path, PathBuf}; use tokio::runtime::Runtime; +const CRITERION_MAX_DIRECTORY_NAME_LEN: usize = 64; + #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct BenchmarkFilter { pub name: Option, @@ -49,12 +52,81 @@ pub struct SqlRunConfig { pub output: Option, } +#[derive(Debug, Parser)] +#[command(ignore_errors = true)] +struct CriterionHarnessEnv { + #[command(flatten)] + options: CommonOpt, + + #[arg( + env = "BENCH_PERSIST_RESULTS", + long = "persist_results", + default_value = "false", + action = clap::ArgAction::SetTrue + )] + persist_results: bool, + + #[arg( + env = "BENCH_VALIDATE", + long = "validate_results", + default_value = "false", + action = clap::ArgAction::SetTrue + )] + validate: bool, + + #[arg(env = "BENCH_NAME")] + name: Option, + + #[arg(env = "BENCH_SUBGROUP")] + subgroup: Option, + + #[arg(env = "BENCH_QUERY")] + query: Option, + + #[arg(env = "BENCH_NAMESPACE")] + criterion_namespace: Option, +} + +/// Builds the direct Criterion harness configuration from its `BENCH_*` +/// environment variables. +pub fn criterion_harness_config_from_env() -> (SqlRunConfig, Option) { + let args = CriterionHarnessEnv::parse(); + let config = SqlRunConfig { + common: args.options, + filter: BenchmarkFilter { + name: args.name, + subgroup: args.subgroup, + query: args.query, + }, + replacements: default_criterion_replacements(), + query_filename: None, + persist_results: args.persist_results, + validate_results: args.validate, + output: None, + }; + + (config, args.criterion_namespace) +} + /// Runs the selected SQL benchmarks through a caller-provided Criterion instance. pub fn run_criterion_benchmarks_impl( benchmark_dir: &Path, config: &SqlRunConfig, criterion: &mut Criterion, ) -> Result<()> { + run_criterion_benchmarks_impl_with_namespace(benchmark_dir, config, None, criterion) +} + +/// Runs the selected SQL benchmarks through a caller-provided Criterion instance, +/// optionally appending a safe invocation namespace to each benchmark group. +pub fn run_criterion_benchmarks_impl_with_namespace( + benchmark_dir: &Path, + config: &SqlRunConfig, + namespace: Option<&str>, + criterion: &mut Criterion, +) -> Result<()> { + validate_criterion_namespace(namespace)?; + let rt = make_tokio_runtime()?; let listing_ctx = make_ctx(&config.common)?; let all_benchmarks = rt.block_on(load_benchmark_definitions_for_query( @@ -68,7 +140,13 @@ pub fn run_criterion_benchmarks_impl( ensure_selection(&config.filter, &all_benchmarks, &selected)?; + let mut named_benchmarks = Vec::with_capacity(selected.len()); for (group_name, benchmarks) in selected { + named_benchmarks + .push((criterion_group_name(&group_name, namespace)?, benchmarks)); + } + + for (group_name, benchmarks) in named_benchmarks { let mut group = criterion.benchmark_group(group_name); group.sample_size(10); @@ -89,6 +167,43 @@ pub fn run_criterion_benchmarks_impl( Ok(()) } +fn validate_criterion_namespace(namespace: Option<&str>) -> Result<()> { + let Some(namespace) = namespace else { + return Ok(()); + }; + + if namespace.is_empty() + || !namespace.chars().all(|character| { + character.is_ascii_lowercase() + || character.is_ascii_digit() + || "_-".contains(character) + }) + { + return Err(exec_datafusion_err!( + "criterion namespace must be nonempty and contain only lowercase ASCII letters, digits, '_', or '-'" + )); + } + + Ok(()) +} + +fn criterion_group_name(group_name: &str, namespace: Option<&str>) -> Result { + validate_criterion_namespace(namespace)?; + + let Some(namespace) = namespace else { + return Ok(group_name.to_string()); + }; + + let group_name = format!("{group_name}__{namespace}"); + if group_name.len() > CRITERION_MAX_DIRECTORY_NAME_LEN { + return Err(exec_datafusion_err!( + "criterion group with namespace must not exceed {CRITERION_MAX_DIRECTORY_NAME_LEN} bytes" + )); + } + + Ok(group_name) +} + /// Runs one benchmark case inside Criterion and converts benchmark panics to errors. fn run_criterion_benchmark( rt: &Runtime, @@ -766,4 +881,84 @@ mod tests { assert_eq!(benchmark.group(), "tpch"); assert_eq!(criterion_function_name(&benchmark), "Q01_sf1"); } + + #[test] + fn criterion_group_names_include_safe_namespaces() { + assert_eq!(criterion_group_name("tpch", None).unwrap(), "tpch"); + assert_eq!( + criterion_group_name("tpch", Some("parquet-sf1")).unwrap(), + "tpch__parquet-sf1" + ); + assert_eq!( + criterion_group_name("tpch", Some("memory_sf1")).unwrap(), + "tpch__memory_sf1" + ); + } + + #[test] + fn criterion_group_names_reject_unsafe_namespaces() { + for namespace in ["", "csv/sf1", "csv sf1", "csv.sf1", "parquét"] { + let error = criterion_group_name("tpch", Some(namespace)).unwrap_err(); + + assert!(error.to_string().contains("namespace"), "{error}"); + } + } + + #[test] + fn criterion_group_names_reject_windows_case_collisions() { + let error = criterion_group_name("tpch", Some("Parquet")).unwrap_err(); + + assert!(error.to_string().contains("lowercase"), "{error}"); + assert_eq!( + criterion_group_name("tpch", Some("parquet")).unwrap(), + "tpch__parquet" + ); + } + + #[test] + fn criterion_group_names_reject_components_criterion_would_truncate() { + let group_name = "g".repeat(55); + + assert_eq!( + criterion_group_name(&group_name, Some("1234567")) + .unwrap() + .len(), + 64 + ); + + let first = criterion_group_name(&group_name, Some("12345678")); + let second = criterion_group_name(&group_name, Some("12345679")); + + assert!(first.unwrap_err().to_string().contains("64 bytes")); + assert!(second.unwrap_err().to_string().contains("64 bytes")); + } + + #[test] + fn criterion_harness_reads_namespace_from_env_in_subprocess() { + const CHILD_ENV: &str = "DATAFUSION_CRITERION_HARNESS_ENV_TEST_CHILD"; + + if std::env::var_os(CHILD_ENV).is_some() { + let (_, namespace) = criterion_harness_config_from_env(); + + assert_eq!(namespace.as_deref(), Some("parquet_sf1")); + return; + } + + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "sql_benchmark_runner::tests::criterion_harness_reads_namespace_from_env_in_subprocess", + "--nocapture", + ]) + .env(CHILD_ENV, "1") + .env("BENCH_NAMESPACE", "parquet_sf1") + .output() + .unwrap(); + + assert!( + output.status.success(), + "child failed:\n{}", + String::from_utf8_lossy(&output.stderr) + ); + } }