diff --git a/datafusion/core/benches/sql_planner.rs b/datafusion/core/benches/sql_planner.rs index de81ed021b427..48770541ca9b6 100644 --- a/datafusion/core/benches/sql_planner.rs +++ b/datafusion/core/benches/sql_planner.rs @@ -85,6 +85,55 @@ fn physical_plan(ctx: &SessionContext, rt: &Runtime, sql: &str) { })); } +/// Read the `q{N}.sql` query files from `dir`, in ascending numeric order. +/// +/// The files are discovered by listing the directory rather than by walking a +/// hardcoded range, so newly added queries are picked up without changing this +/// code. Entries that do not match `q{N}.sql` (such as the `sorted_data` +/// directory alongside the ClickBench queries) are ignored. +/// +/// Panics if `dir` contains no queries, or if the numbering is not contiguous +/// starting at `q0.sql`. A gap almost certainly means a query file was lost, and +/// silently benchmarking a subset would hide that. +fn read_numbered_queries(dir: &str) -> Vec { + let entries = std::fs::read_dir(dir) + .unwrap_or_else(|e| panic!("Failed to read query directory '{dir}': {e}")); + + let mut numbers = entries + .map(|entry| { + entry.unwrap_or_else(|e| panic!("Failed to read entry in '{dir}': {e}")) + }) + .filter_map(|entry| { + let file_name = entry.file_name(); + let stem = file_name + .to_str()? + .strip_prefix('q')? + .strip_suffix(".sql")?; + stem.parse::().ok() + }) + .collect::>(); + numbers.sort_unstable(); + + assert!( + !numbers.is_empty(), + "No `q{{N}}.sql` query files found in '{dir}'" + ); + assert_eq!( + numbers, + (0..numbers.len()).collect::>(), + "Query files in '{dir}' are not numbered contiguously from q0.sql" + ); + + numbers + .into_iter() + .map(|q| { + let path = format!("{dir}q{q}.sql"); + std::fs::read_to_string(&path) + .unwrap_or_else(|e| panic!("Failed to read query file '{path}': {e}")) + }) + .collect() +} + /// Create schema with the specified number of columns fn create_schema(column_prefix: &str, num_columns: usize) -> Schema { let fields: Fields = (0..num_columns) @@ -636,20 +685,13 @@ fn criterion_benchmark(c: &mut Criterion) { // }); // -- clickbench -- - let clickbench_queries = (0..=42) - .map(|q| { - std::fs::read_to_string(format!( - "{benchmarks_path}queries/clickbench/queries/q{q}.sql" - )) - .unwrap() - }) - .chain((0..=7).map(|q| { - std::fs::read_to_string(format!( - "{benchmarks_path}queries/clickbench/extended/q{q}.sql" - )) - .unwrap() - })) - .collect::>(); + let clickbench_queries = + read_numbered_queries(&format!("{benchmarks_path}queries/clickbench/queries/")) + .into_iter() + .chain(read_numbered_queries(&format!( + "{benchmarks_path}queries/clickbench/extended/" + ))) + .collect::>(); let clickbench_ctx = register_clickbench_hits_table(&rt);