diff --git a/datafusion/sqllogictest/README.md b/datafusion/sqllogictest/README.md index f0a54cf978fbf..76c28b5ecd2ec 100644 --- a/datafusion/sqllogictest/README.md +++ b/datafusion/sqllogictest/README.md @@ -184,6 +184,56 @@ EXPLAIN ANALYZE SELECT * FROM generate_series(100); Plan with Metrics LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=0, end=100, batch_size=8192], metrics=[output_rows=101, elapsed_compute=, output_bytes=] ``` +## Cookbook: Sweeping config with `configMatrix` + +Runs the same `.slt` once per combination of config values. Each directive is a comment: + +```text +# configMatrix: =,[,...] +``` + +- Repeat the directive to nest keys. Values are the cartesian product. +- Whitespace-trimmed and deduped; repeated keys merge value lists. +- Unknown key or invalid value fails fast, naming the file, key, and value. +- Each override is applied like `SET = `, so `datafusion.runtime.*` + keys (e.g. `datafusion.runtime.memory_limit`) and config-dependent UDFs behave + exactly as they would for an in-file `SET`. +- Supported by the default runner and by `--substrait-round-trip`. `--complete` + rejects a file that declares directives, since it would overwrite the file + with the output of a single combination. The Postgres runner ignores them. +- Every combination runs even if an earlier one fails. Each failure is prefixed + with the combination that produced it, and a file with more than one failing + combination reports all of them. + +Nested example (2 × 2 = 4 runs): + +```text +# configMatrix: datafusion.execution.parquet.coerce_int96=ms,us +# configMatrix: datafusion.execution.parquet.coerce_int96_tz=UTC,America/New_York + +statement ok +CREATE EXTERNAL TABLE int96_from_spark +STORED AS PARQUET +LOCATION '../../parquet-testing/data/int96_from_spark.parquet'; + +query I +select count(*) from int96_from_spark +---- +6 +``` + +Failure output: + +```text +[configMatrix: datafusion.execution.parquet.coerce_int96=ms, datafusion.execution.parquet.coerce_int96_tz=UTC] +caused by +External error: 1 errors in file .../parquet_int96_matrix.slt +``` + +When more than one combination fails, all of them are reported together under a +`N configMatrix combinations failed:` header, each still prefixed with its own +combination. + # Reference ## Running tests: Validation Mode diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs index 81b0c41d6a82a..42a324d310d1a 100644 --- a/datafusion/sqllogictest/bin/sqllogictests.rs +++ b/datafusion/sqllogictest/bin/sqllogictests.rs @@ -23,9 +23,9 @@ use datafusion::common::{DataFusionError, Result, exec_datafusion_err, exec_err} use datafusion_sqllogictest::DataFusionSubstraitRoundTrip; use datafusion_sqllogictest::TestFile; use datafusion_sqllogictest::{ - CurrentlyExecutingSqlTracker, DataFusion, Filter, TestContext, df_value_validator, - read_dir_recursive, setup_scratch_dir, should_skip_file, should_skip_record, - value_normalizer, + CurrentlyExecutingSqlTracker, DFColumnType, DataFusion, Filter, TestContext, + df_value_validator, read_dir_recursive, run_each_configuration, setup_scratch_dir, + should_skip_file, should_skip_record, test_configurations, value_normalizer, }; use futures::stream::StreamExt; use indicatif::{ @@ -443,33 +443,57 @@ async fn run_test_file_substrait_round_trip( path, relative_path, } = test_file; - let Some(test_ctx) = TestContext::try_new_for_test_file(&relative_path).await else { - info!("Skipping: {}", path.display()); - return Ok(()); - }; - setup_scratch_dir(&relative_path)?; - - let count: u64 = get_record_count(&path, "DatafusionSubstraitRoundTrip".to_string()); - let pb = mp.add(ProgressBar::new(count)); - - pb.set_style(mp_style); - pb.set_message(relative_path.display().to_string()); - let mut runner = sqllogictest::Runner::new(|| async { - Ok(DataFusionSubstraitRoundTrip::new( - test_ctx.session_ctx().clone(), - relative_path.clone(), - pb.clone(), + // Parsed once and replayed for every configuration. + let records = parse_records(&path)?; + let count = count_records(&records, "DatafusionSubstraitRoundTrip"); + let configurations = test_configurations(&path)?; + + // Reborrow so each combination's future moves in `Copy` references, keeping + // the closure callable per combination. + let path = &path; + let relative_path = &relative_path; + let records = &records; + let mp = ∓ + let mp_style = &mp_style; + let currently_executing_sql_tracker = ¤tly_executing_sql_tracker; + + // Run once per configMatrix combination. + run_each_configuration(configurations, |test_configuration| async move { + let Some((test_ctx, pb)) = setup_combination( + relative_path, + test_configuration.settings(), + mp, + mp_style, + count, ) - .with_currently_executing_sql_tracker(currently_executing_sql_tracker.clone())) - }); - runner.add_label("DatafusionSubstraitRoundTrip"); - runner.with_column_validator(strict_column_validator); - runner.with_normalizer(value_normalizer); - runner.with_validator(validator); - let res = run_file_in_runner(path, &mut runner, filters, colored_output).await; - pb.finish_and_clear(); - res + .await? + else { + info!("Skipping: {}", path.display()); + return Ok(()); + }; + + let mut runner = sqllogictest::Runner::new(|| async { + Ok(DataFusionSubstraitRoundTrip::new( + test_ctx.session_ctx().clone(), + relative_path.clone(), + pb.clone(), + ) + .with_currently_executing_sql_tracker( + currently_executing_sql_tracker.clone(), + )) + }); + runner.add_label("DatafusionSubstraitRoundTrip"); + runner.with_column_validator(strict_column_validator); + runner.with_normalizer(value_normalizer); + runner.with_validator(validator); + let result = + run_file_in_runner(path, records, &mut runner, filters, colored_output).await; + pb.finish_and_clear(); + + result + }) + .await } #[cfg(not(feature = "substrait"))] @@ -502,64 +526,87 @@ async fn run_test_file( path, relative_path, } = test_file; - let Some(test_ctx) = TestContext::try_new_for_test_file(&relative_path).await else { - info!("Skipping: {}", path.display()); - return Ok(()); - }; - setup_scratch_dir(&relative_path)?; - - let count: u64 = get_record_count(&path, "Datafusion".to_string()); - let pb = mp.add(ProgressBar::new(count)); - - pb.set_style(mp_style); - pb.set_message(relative_path.display().to_string()); - // If DataFusion configuration has changed during test file runs, errors will be - // pushed to this vec. - // HACK: managed externally because `sqllogictest` is an external dependency, and - // it doesn't have an API to directly access the inner runner. - let config_change_errors = Arc::new(Mutex::new(Vec::new())); - let mut runner = sqllogictest::Runner::new(|| async { - Ok(DataFusion::new( - test_ctx.session_ctx().clone(), - relative_path.clone(), - pb.clone(), + // Parsed once and replayed for every configuration. + let records = parse_records(&path)?; + let count = count_records(&records, "Datafusion"); + let configurations = test_configurations(&path)?; + + // Reborrow so each combination's future moves in `Copy` references, keeping + // the closure callable per combination. + let path = &path; + let relative_path = &relative_path; + let records = &records; + let mp = ∓ + let mp_style = &mp_style; + let currently_executing_sql_tracker = ¤tly_executing_sql_tracker; + + // Run once per configMatrix combination. + run_each_configuration(configurations, |test_configuration| async move { + let Some((test_ctx, pb)) = setup_combination( + relative_path, + test_configuration.settings(), + mp, + mp_style, + count, ) - .with_currently_executing_sql_tracker(currently_executing_sql_tracker.clone()) - .with_config_change_errors(Arc::clone(&config_change_errors))) - }); - runner.add_label("Datafusion"); - runner.with_column_validator(strict_column_validator); - runner.with_normalizer(value_normalizer); - runner.with_validator(validator); - let result = run_file_in_runner(path, &mut runner, filters, colored_output).await; - pb.finish_and_clear(); - - result?; + .await? + else { + info!("Skipping: {}", path.display()); + return Ok(()); + }; - // If there was no correctness error, check that the config is unchanged. - runner.shutdown_async().await; - config_change_result(&config_change_errors) + // If DataFusion configuration has changed during test file runs, errors will be + // pushed to this vec. + // HACK: managed externally because `sqllogictest` is an external dependency, and + // it doesn't have an API to directly access the inner runner. + let config_change_errors = Arc::new(Mutex::new(Vec::new())); + let mut runner = sqllogictest::Runner::new(|| async { + Ok(DataFusion::new( + test_ctx.session_ctx().clone(), + relative_path.clone(), + pb.clone(), + ) + .with_currently_executing_sql_tracker(currently_executing_sql_tracker.clone()) + .with_config_change_errors(Arc::clone(&config_change_errors))) + }); + runner.add_label("Datafusion"); + runner.with_column_validator(strict_column_validator); + runner.with_normalizer(value_normalizer); + runner.with_validator(validator); + let result = + run_file_in_runner(path, records, &mut runner, filters, colored_output).await; + pb.finish_and_clear(); + + runner.shutdown_async().await; + + // A correctness failure takes precedence; otherwise surface any config + // the file left modified. + result.and_then(|()| config_change_result(&config_change_errors)) + }) + .await } -async fn run_file_in_runner>( - path: PathBuf, +async fn run_file_in_runner( + path: &Path, + records: &[Record], runner: &mut sqllogictest::Runner, filters: &[Filter], colored_output: bool, -) -> Result<()> { - let path = path.canonicalize()?; - let records = - parse_file(&path).map_err(|e| DataFusionError::External(Box::new(e)))?; +) -> Result<()> +where + D: AsyncDB, + M: MakeConnection, +{ let mut errs = vec![]; - for record in records.into_iter() { + for record in records { if let Record::Halt { .. } = record { break; } - if should_skip_record::(&record, filters) { + if should_skip_record::(record, filters) { continue; } - if let Err(err) = runner.run_async(record).await { + if let Err(err) = runner.run_async(record.clone()).await { if colored_output { errs.push(format!("{}", err.display(true))); } else { @@ -569,6 +616,7 @@ async fn run_file_in_runner>( } if !errs.is_empty() { + let path = path.canonicalize()?; let mut msg = format!("{} errors in file {}\n\n", errs.len(), path.display()); for (i, err) in errs.iter().enumerate() { if i >= ERRS_PER_FILE_LIMIT { @@ -587,30 +635,54 @@ async fn run_file_in_runner>( Ok(()) } -#[expect(clippy::needless_pass_by_value)] -fn get_record_count(path: &PathBuf, label: String) -> u64 { - let records: Vec::ColumnType>> = - parse_file(path).unwrap(); - let mut count: u64 = 0; - - for rec in &records { - match rec { - Record::Query { conditions, .. } | Record::Statement { conditions, .. } - if conditions.is_empty() - || !conditions.contains(&Condition::SkipIf { - label: label.clone(), - }) - || conditions.contains(&Condition::OnlyIf { - label: label.clone(), - }) => - { - count += 1; +fn parse_records(path: &Path) -> Result>> { + parse_file(path).map_err(|e| DataFusionError::External(Box::new(e))) +} + +fn count_records(records: &[Record], label: &str) -> u64 { + let skip_if = Condition::SkipIf { + label: label.to_string(), + }; + let only_if = Condition::OnlyIf { + label: label.to_string(), + }; + records + .iter() + .filter(|rec| match rec { + Record::Query { conditions, .. } | Record::Statement { conditions, .. } => { + conditions.is_empty() + || !conditions.contains(&skip_if) + || conditions.contains(&only_if) } - _ => {} - } - } + _ => false, + }) + .count() as u64 +} + +/// Per-combination setup shared by the default and Substrait runners: build the +/// test context, apply the configMatrix overrides, and create the progress bar. +/// Returns `Ok(None)` when the file should be skipped (unsupported feature); the +/// caller logs the skip. +async fn setup_combination( + relative_path: &Path, + settings: &[(String, String)], + mp: &MultiProgress, + mp_style: &ProgressStyle, + count: u64, +) -> Result> { + let Some(test_ctx) = TestContext::try_new_for_test_file(relative_path).await else { + return Ok(None); + }; + setup_scratch_dir(relative_path)?; + // Before the engine is built: it snapshots config to detect drift. + test_ctx + .apply_config_overrides(settings, relative_path) + .await?; - count + let pb = mp.add(ProgressBar::new(count)); + pb.set_style(mp_style.clone()); + pb.set_message(relative_path.display().to_string()); + Ok(Some((test_ctx, pb))) } #[cfg(feature = "postgres")] @@ -629,7 +701,8 @@ async fn run_test_file_with_postgres( } = test_file; setup_scratch_dir(&relative_path)?; - let count: u64 = get_record_count(&path, "postgresql".to_string()); + let records = parse_records(&path)?; + let count = count_records(&records, "postgresql"); let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); @@ -646,7 +719,7 @@ async fn run_test_file_with_postgres( runner.with_column_validator(strict_column_validator); runner.with_normalizer(value_normalizer); runner.with_validator(validator); - let result = run_file_in_runner(path, &mut runner, filters, false).await; + let result = run_file_in_runner(&path, &records, &mut runner, filters, false).await; pb.finish_and_clear(); result } @@ -682,13 +755,25 @@ async fn run_complete_file( info!("Using complete mode to complete: {}", path.display()); + // `update_test_file` rewrites the file from a single run, so it cannot hold + // one expected-output set per configuration. + if test_configurations(&path)?.iter().any(|c| !c.is_empty()) { + return exec_err!( + "Cannot use --complete on {}: it declares `# configMatrix:` \ + directives, and completion would overwrite the file with the \ + output of a single configuration", + relative_path.display() + ); + } + let Some(test_ctx) = TestContext::try_new_for_test_file(&relative_path).await else { info!("Skipping: {}", path.display()); return Ok(()); }; setup_scratch_dir(&relative_path)?; - let count: u64 = get_record_count(&path, "Datafusion".to_string()); + // `update_test_file` re-parses the file itself, so only the count is needed. + let count = count_records(&parse_records(&path)?, "Datafusion"); let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); @@ -744,7 +829,7 @@ async fn run_complete_file_with_postgres( ); setup_scratch_dir(&relative_path)?; - let count: u64 = get_record_count(&path, "postgresql".to_string()); + let count = count_records(&parse_records(&path)?, "postgresql"); let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); diff --git a/datafusion/sqllogictest/src/config_matrix.rs b/datafusion/sqllogictest/src/config_matrix.rs new file mode 100644 index 0000000000000..4e5881c7ae23e --- /dev/null +++ b/datafusion/sqllogictest/src/config_matrix.rs @@ -0,0 +1,549 @@ +// 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. + +//! `# configMatrix:` directives sweep DataFusion config across repeated runs of +//! one SLT file: +//! +//! ```text +//! # configMatrix: =,[,...] +//! ``` +//! +//! This module validates directives, dedups their values, and expands them into +//! the [`TestConfiguration`]s the runner applies to a session. Repeated +//! directives combine as a cartesian product. + +use std::fmt; +use std::fs; +use std::future::Future; +use std::path::Path; + +use datafusion::common::{DataFusionError, Result, exec_datafusion_err}; +use itertools::Itertools; + +/// Compared case-insensitively; this spelling is the documented one. +const DIRECTIVE_MARKER: &str = "configMatrix:"; + +/// Config values to apply before a single run of an SLT file. +/// +/// Empty when the file declared no directives, which means "run once, +/// unmodified". +#[derive(Debug, Clone)] +pub struct TestConfiguration(Vec<(String, String)>); + +impl TestConfiguration { + /// The `key = value` pairs to set on the session. + pub fn settings(&self) -> &[(String, String)] { + &self.0 + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// Prefix a failure with the configuration that produced it, so a sweep + /// names the values that broke. No-op when no directives were declared. + fn attribute_failure(&self, result: Result) -> Result { + if self.is_empty() { + return result; + } + result.map_err(|e| e.context(self.to_string())) + } +} + +impl fmt::Display for TestConfiguration { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "[configMatrix: {}]", + self.0.iter().map(|(k, v)| format!("{k}={v}")).join(", ") + ) + } +} + +/// Validate, dedup, and expand the directives declared by `path`. +/// +/// Never empty: a file without directives yields one empty +/// [`TestConfiguration`], so callers keep a single loop shape. +pub fn test_configurations(path: &Path) -> Result> { + let content = fs::read_to_string(path).map_err(|e| { + exec_datafusion_err!( + "Failed to read {} for configMatrix parsing: {e}", + path.display() + ) + })?; + parse_configurations(&content, path) +} + +/// Run `run_one` once per configuration, continuing past failures so one broken +/// combination never hides the others. Each failure is prefixed with the +/// configuration that produced it (see `TestConfiguration::attribute_failure`). +/// +/// Returns `Ok(())` if all passed, the lone failure verbatim if exactly one +/// failed (so a matrix-free file keeps its original error), or an aggregate +/// naming every failing combination. +/// +/// `run_one` takes the configuration by value so its future holds no +/// higher-ranked borrow and stays `Send` for the runner's spawned task. +pub async fn run_each_configuration( + configurations: Vec, + mut run_one: F, +) -> Result<()> +where + F: FnMut(TestConfiguration) -> Fut, + Fut: Future>, +{ + let mut failures = Vec::new(); + for configuration in configurations { + // Keep a copy to attribute the failure after `run_one` consumes it. + let label = configuration.clone(); + if let Err(err) = label.attribute_failure(run_one(configuration).await) { + failures.push(err); + } + } + combine_configuration_failures(failures) +} + +/// Collapse per-configuration failures into a single [`Result`], returned by +/// [`run_each_configuration`]. +fn combine_configuration_failures(mut failures: Vec) -> Result<()> { + match failures.len() { + 0 => Ok(()), + 1 => Err(failures.pop().unwrap()), + n => { + let combined = failures.iter().join("\n\n"); + Err(DataFusionError::External( + format!("{n} configMatrix combinations failed:\n\n{combined}").into(), + )) + } + } +} + +/// `path` is used only for error messages. +fn parse_configurations(content: &str, path: &Path) -> Result> { + // Dimensions borrow `content`; values are owned only in the product below. + let mut dimensions: Vec<(&str, Vec<&str>)> = Vec::new(); + + for (idx, line) in content.lines().enumerate() { + let Some(directive) = strip_directive_prefix(line) else { + continue; + }; + let invalid = |detail: String| { + DataFusionError::Configuration(format!( + "Invalid configMatrix directive in {}:{}: {detail}", + path.display(), + idx + 1 + )) + }; + + let (key, values) = directive.split_once('=').ok_or_else(|| { + invalid(format!( + "expected `# configMatrix: =,[,...]`, got `{directive}`" + )) + })?; + + let key = key.trim(); + if key.is_empty() { + return Err(invalid("missing config key".to_string())); + } + + let values: Vec<&str> = values + .split(',') + .map(str::trim) + .filter(|value| !value.is_empty()) + .collect(); + if values.is_empty() { + return Err(invalid(format!("no values provided for `{key}`"))); + } + + // A repeated key unions its values instead of adding a dimension. + match dimensions.iter_mut().find(|(seen, _)| *seen == key) { + Some((_, merged)) => merged.extend(values), + None => dimensions.push((key, values)), + } + } + + Ok(dimensions + .into_iter() + .map(|(key, values)| { + // Dedup keeps first-seen order. + values + .into_iter() + .unique() + .map(|value| (key.to_string(), value.to_string())) + .collect_vec() + }) + .multi_cartesian_product() + .map(TestConfiguration) + .collect()) +} + +/// Text after `configMatrix:` when `line` is a directive comment. Tolerates +/// `##` banners and any marker casing. +fn strip_directive_prefix(line: &str) -> Option<&str> { + let after_hash = line.trim_start().strip_prefix('#')?; + let after_hash = after_hash.trim_start_matches('#').trim_start(); + + let (marker, rest) = after_hash.split_at_checked(DIRECTIVE_MARKER.len())?; + if !marker.eq_ignore_ascii_case(DIRECTIVE_MARKER) { + return None; + } + Some(rest.trim()) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::common::exec_err; + use std::io::Write; + + /// Settings of every configuration parsed from `content`. + fn parse(content: &str) -> Vec> { + parse_configurations(content, Path::new("test.slt")) + .unwrap() + .into_iter() + .map(|test_configuration| test_configuration.0) + .collect() + } + + fn test_configuration(settings: &[(&str, &str)]) -> TestConfiguration { + TestConfiguration( + settings + .iter() + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect(), + ) + } + + #[test] + fn parses_no_directives() { + // Zero dimensions expand to exactly one empty configuration. + let content = "# ordinary comment\nquery I\nSELECT 1\n----\n1\n"; + assert_eq!(parse(content), vec![Vec::new()]); + } + + #[test] + fn parses_single_directive() { + let configurations = parse( + "# configMatrix: datafusion.optimizer.enable_piecewise_merge_join=true,false\n", + ); + assert_eq!(configurations.len(), 2); + assert_eq!( + configurations[0], + vec![( + "datafusion.optimizer.enable_piecewise_merge_join".to_string(), + "true".to_string(), + )] + ); + assert_eq!( + configurations[1], + vec![( + "datafusion.optimizer.enable_piecewise_merge_join".to_string(), + "false".to_string(), + )] + ); + } + + #[test] + fn parses_case_insensitive_marker() { + let configurations = + parse("## CONFIGMATRIX: datafusion.execution.batch_size=1024,4096\n"); + assert_eq!(configurations.len(), 2); + assert_eq!(configurations[0][0].0, "datafusion.execution.batch_size"); + } + + #[test] + fn parses_multiple_directives_and_computes_cartesian_product() { + let configurations = parse("# configMatrix: a=1,2\n# configMatrix: b=x,y,z\n"); + assert_eq!(configurations.len(), 6); + assert_eq!( + configurations[0], + vec![ + ("a".to_string(), "1".to_string()), + ("b".to_string(), "x".to_string()), + ] + ); + assert_eq!( + configurations[5], + vec![ + ("a".to_string(), "2".to_string()), + ("b".to_string(), "z".to_string()), + ] + ); + } + + #[test] + fn rejects_missing_equals() { + let err = parse_configurations( + "# configMatrix: bogus_no_equals\n", + Path::new("test.slt"), + ) + .unwrap_err(); + assert!(err.to_string().contains("=")); + } + + #[test] + fn rejects_missing_key() { + let err = + parse_configurations("# configMatrix: =true,false\n", Path::new("test.slt")) + .unwrap_err(); + assert!(err.to_string().contains("missing config key")); + } + + #[test] + fn rejects_empty_values() { + let err = + parse_configurations("# configMatrix: some.key=\n", Path::new("test.slt")) + .unwrap_err(); + assert!(err.to_string().contains("no values provided")); + } + + #[test] + fn rejects_empty_values_on_a_repeated_key() { + let err = parse_configurations( + "# configMatrix: k=1,2\n# configMatrix: k=\n", + Path::new("test.slt"), + ) + .unwrap_err(); + assert!(err.to_string().contains("no values provided")); + } + + #[test] + fn deduplicates_repeated_values_preserving_first_order() { + let configurations = + parse("# configMatrix: some.key=false,true,false,true,false\n"); + assert_eq!(configurations.len(), 2); + assert_eq!(configurations[0][0].1, "false"); + assert_eq!(configurations[1][0].1, "true"); + } + + #[test] + fn strips_whitespace_around_values() { + let configurations = parse("# configMatrix: k= 1024 ,\t2048 , \t4096\n"); + let values: Vec<&str> = configurations.iter().map(|c| c[0].1.as_str()).collect(); + assert_eq!(values, vec!["1024", "2048", "4096"]); + } + + #[test] + fn strips_whitespace_around_key_and_around_equals() { + let configurations = parse( + "# configMatrix: datafusion.execution.batch_size = 1024 , 2048 \n", + ); + assert_eq!(configurations.len(), 2); + assert_eq!(configurations[0][0].0, "datafusion.execution.batch_size"); + } + + #[test] + fn dedup_runs_after_whitespace_is_stripped() { + // `1024`, ` 1024`, and `\t1024 ` all normalize to the same value. + let configurations = parse("# configMatrix: k= 1024 , 1024,\t1024 \n"); + assert_eq!(configurations.len(), 1); + assert_eq!(configurations[0][0].1, "1024"); + } + + #[test] + fn trailing_and_repeated_commas_are_ignored() { + let configurations = parse("# configMatrix: k=1024,,2048,\n"); + let values: Vec<&str> = configurations.iter().map(|c| c[0].1.as_str()).collect(); + assert_eq!(values, vec!["1024", "2048"]); + } + + #[test] + fn merges_repeated_directives_for_same_key() { + let configurations = parse( + "# configMatrix: datafusion.execution.batch_size=1024,2048\n\ + # configMatrix: datafusion.execution.batch_size=1024,2048\n", + ); + assert_eq!(configurations.len(), 2); + } + + #[test] + fn merges_repeated_directives_unioning_values() { + let configurations = + parse("# configMatrix: k=1024,2048\n# configMatrix: k=2048,4096\n"); + let values: Vec<&str> = configurations.iter().map(|c| c[0].1.as_str()).collect(); + assert_eq!(values, vec!["1024", "2048", "4096"]); + } + + #[test] + fn nested_matrices_expand_as_cartesian_product() { + let configurations = parse( + "# configMatrix: datafusion.execution.batch_size=1024,2048\n\ + # configMatrix: datafusion.execution.param1=true,false\n", + ); + assert_eq!(configurations.len(), 4); + for configuration in &configurations { + assert_eq!(configuration.len(), 2); + assert_eq!(configuration[0].0, "datafusion.execution.batch_size"); + assert_eq!(configuration[1].0, "datafusion.execution.param1"); + } + assert_eq!(configurations[0][0].1, "1024"); + assert_eq!(configurations[0][1].1, "true"); + assert_eq!(configurations[3][0].1, "2048"); + assert_eq!(configurations[3][1].1, "false"); + } + + #[test] + fn ignores_non_comment_lines_with_marker_text() { + assert_eq!(parse("SELECT 'configMatrix: foo=bar';\n"), vec![Vec::new()]); + } + + #[test] + fn displays_key_value_pairs() { + assert_eq!( + test_configuration(&[("a", "1"), ("b", "x")]).to_string(), + "[configMatrix: a=1, b=x]" + ); + } + + #[test] + fn test_configurations_yields_one_empty_entry_when_no_matrix() { + let mut file = tempfile::NamedTempFile::new().unwrap(); + file.write_all(b"# just a comment\nquery I\nSELECT 1\n----\n1\n") + .unwrap(); + let configurations = test_configurations(file.path()).unwrap(); + assert_eq!(configurations.len(), 1); + assert!(configurations[0].is_empty()); + } + + #[test] + fn test_configurations_yields_every_combination() { + let mut file = tempfile::NamedTempFile::new().unwrap(); + file.write_all( + b"# configMatrix: a=1,2\n\ + # configMatrix: b=x,y\n\ + query I\nSELECT 1\n----\n1\n", + ) + .unwrap(); + let configurations = test_configurations(file.path()).unwrap(); + assert_eq!(configurations.len(), 4); + assert!(configurations.iter().all(|c| !c.is_empty())); + } + + #[test] + fn attribute_failure_is_noop_without_directives() { + let err: Result<()> = Err(exec_datafusion_err!("boom")); + assert_eq!( + test_configuration(&[]) + .attribute_failure(err) + .unwrap_err() + .to_string(), + exec_datafusion_err!("boom").to_string() + ); + } + + #[test] + fn attribute_failure_names_the_configuration() { + let err: Result<()> = Err(exec_datafusion_err!("boom")); + let msg = test_configuration(&[("k", "1")]) + .attribute_failure(err) + .unwrap_err() + .to_string(); + assert!(msg.starts_with("[configMatrix: k=1]"), "got {msg}"); + assert!(msg.contains("boom"), "got {msg}"); + } + + #[tokio::test] + async fn run_each_configuration_runs_every_combination_past_failures() { + use std::sync::{Arc, Mutex}; + + // Drives the exact helper both runner paths use: every configuration + // must run and be attributed even though an earlier one fails. + let configurations = vec![ + test_configuration(&[("k", "1")]), + test_configuration(&[("k", "2")]), + test_configuration(&[("k", "3")]), + ]; + let seen = Arc::new(Mutex::new(Vec::new())); + let recorder = Arc::clone(&seen); + + let result = run_each_configuration(configurations, |configuration| { + let recorder = Arc::clone(&recorder); + async move { + let value = configuration.settings()[0].1.clone(); + recorder.lock().unwrap().push(value.clone()); + // The middle combination fails; the last must still run. + if value == "2" { + exec_err!("boom") + } else { + Ok(()) + } + } + }) + .await; + + // All three ran even though the second failed. + assert_eq!(*seen.lock().unwrap(), vec!["1", "2", "3"]); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("[configMatrix: k=2]"), "got {msg}"); + assert!(msg.contains("boom"), "got {msg}"); + } + + #[tokio::test] + async fn run_each_configuration_is_ok_when_all_pass() { + let configurations = vec![ + test_configuration(&[("k", "1")]), + test_configuration(&[("k", "2")]), + ]; + let result = run_each_configuration(configurations, |_| async { Ok(()) }).await; + assert!(result.is_ok()); + } + + #[test] + fn combine_configuration_failures_is_ok_when_empty() { + assert!(combine_configuration_failures(vec![]).is_ok()); + } + + #[test] + fn combine_configuration_failures_returns_a_lone_failure_verbatim() { + // A file without a matrix expands to one empty configuration, so its + // failure must be reported exactly as before (no combination prefix, + // no aggregate header). + let boom: Result<()> = exec_err!("boom"); + let failure = test_configuration(&[]).attribute_failure(boom).unwrap_err(); + let msg = combine_configuration_failures(vec![failure]) + .unwrap_err() + .to_string(); + assert_eq!(msg, exec_datafusion_err!("boom").to_string()); + } + + #[test] + fn combine_configuration_failures_aggregates_and_attributes_all_failures() { + let boom_1: Result<()> = exec_err!("boom-1"); + let boom_2: Result<()> = exec_err!("boom-2"); + let failures = vec![ + test_configuration(&[("k", "1")]) + .attribute_failure(boom_1) + .unwrap_err(), + test_configuration(&[("k", "2")]) + .attribute_failure(boom_2) + .unwrap_err(), + ]; + + let msg = combine_configuration_failures(failures) + .unwrap_err() + .to_string(); + + assert!( + msg.contains("2 configMatrix combinations failed"), + "got {msg}" + ); + assert!(msg.contains("[configMatrix: k=1]"), "got {msg}"); + assert!(msg.contains("boom-1"), "got {msg}"); + assert!(msg.contains("[configMatrix: k=2]"), "got {msg}"); + assert!(msg.contains("boom-2"), "got {msg}"); + } +} diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs index 6b6c40365f855..926b1fb1e871b 100644 --- a/datafusion/sqllogictest/src/lib.rs +++ b/datafusion/sqllogictest/src/lib.rs @@ -26,6 +26,7 @@ //! DataFusion sqllogictest driver +mod config_matrix; mod engines; mod test_file; @@ -46,6 +47,7 @@ mod filters; mod test_context; mod util; +pub use config_matrix::{TestConfiguration, run_each_configuration, test_configurations}; pub use filters::*; pub use test_context::TestContext; pub use test_file::TestFile; diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 9b97d3f59dac4..1990fae0d3214 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -232,6 +232,35 @@ impl TestContext { pub fn session_ctx(&self) -> &SessionContext { &self.ctx } + + /// Apply `key = value` config overrides to the `SessionContext` in place, + /// used by the runner to sweep `# configMatrix:` directives. + /// + /// Each override runs as `SET = ''`, the same path an in-file + /// `SET` takes, so it accepts any key `SET` does: `datafusion.runtime.*` keys + /// reach the runtime environment and config-dependent UDFs are refreshed. + /// Values are single-quoted (embedded quotes doubled) so strings like + /// timezones and `100M` parse as literals. + /// + /// `origin` labels error messages, typically the test file path. + pub async fn apply_config_overrides( + &self, + overrides: &[(String, String)], + origin: &Path, + ) -> Result<()> { + for (key, value) in overrides { + // Single-quote as a string literal, doubling embedded quotes. + let escaped = value.replace('\'', "''"); + let sql = format!("SET {key} = '{escaped}'"); + self.ctx.sql(&sql).await.map_err(|e| { + e.context(format!( + "configMatrix in {}: failed to set `{key}` = `{value}`", + origin.display() + )) + })?; + } + Ok(()) + } } // ============================================================================== @@ -798,3 +827,77 @@ fn register_conflicting_metadata_tables(ctx: &SessionContext) { RecordBatch::try_new(Arc::new(schema_right), vec![Arc::new(data_right)]).unwrap(); ctx.register_batch("smaller_table", batch_right).unwrap(); } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::execution::memory_pool::MemoryConsumer; + + /// Non-runtime keys land on `ConfigOptions`, same as a plain `SET`. + #[tokio::test] + async fn apply_config_overrides_sets_config_options() { + let test_ctx = TestContext::new(SessionContext::new()); + test_ctx + .apply_config_overrides( + &[( + "datafusion.execution.batch_size".to_string(), + "1234".to_string(), + )], + Path::new("test.slt"), + ) + .await + .unwrap(); + + assert_eq!( + test_ctx + .session_ctx() + .state() + .config() + .options() + .execution + .batch_size + .to_string(), + "1234" + ); + } + + /// `datafusion.runtime.*` keys reach the runtime environment, not just + /// `ConfigOptions` (a direct write would reject `runtime` keys). + #[tokio::test] + async fn apply_config_overrides_routes_runtime_keys() { + let test_ctx = TestContext::new(SessionContext::new()); + test_ctx + .apply_config_overrides( + &[( + "datafusion.runtime.memory_limit".to_string(), + "100M".to_string(), + )], + Path::new("test.slt"), + ) + .await + .unwrap(); + + // The override took effect: the pool now caps reservations at 100M. + let pool = Arc::clone(&test_ctx.session_ctx().runtime_env().memory_pool); + let reservation = MemoryConsumer::new("test").register(&pool); + assert!(reservation.try_grow(50 * 1024 * 1024).is_ok()); + assert!(reservation.try_grow(100 * 1024 * 1024).is_err()); + } + + /// An unknown key still fails fast, naming the originating file. + #[tokio::test] + async fn apply_config_overrides_reports_invalid_key() { + let test_ctx = TestContext::new(SessionContext::new()); + let err = test_ctx + .apply_config_overrides( + &[("datafusion.does.not.exist".to_string(), "1".to_string())], + Path::new("bad.slt"), + ) + .await + .unwrap_err() + .to_string(); + + assert!(err.contains("bad.slt"), "got {err}"); + assert!(err.contains("datafusion.does.not.exist"), "got {err}"); + } +}