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) + ); + } }