feat(runner, bin): support test case partitioning (#257)

Author: BugenZhao

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+## [0.28.0] - 2025-03-06
+
+* runner: Add `Partitioner` and `Runner::with_partitioner` to enable partitioning of test cases, allowing only a subset of the glob result to be executed. This can be helpful for running tests in parallel in CI.
+* bin: Add `--partition-id` and `--partition-count` to set a hash partitioning for the test cases. If users are running in Buildkite CI with `parallelism: ..` specified in the workflow file, this will be automatically configured.
+
 ## [0.27.2] - 2025-02-18
 
 * engines/bin: fix stdin to be closed properly to avoid hangs in the `external` engine.

Cargo.lock

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["sqllogictest", "sqllogictest-bin", "sqllogictest-engines", "tests"]
 
 [workspace.package]
-version = "0.27.2"
+version = "0.28.0"
 edition = "2021"
 homepage = "https://github.com/risinglightdb/sqllogictest-rs"
 keywords = ["sql", "database", "parser", "cli"]

Cargo.toml

@@ -24,8 +24,8 @@ glob = "0.3"
 itertools = "0.13"
 quick-junit = { version = "0.5" }
 rand = "0.8"
-sqllogictest = { path = "../sqllogictest", version = "0.27" }
-sqllogictest-engines = { path = "../sqllogictest-engines", version = "0.27" }
+sqllogictest = { path = "../sqllogictest", version = "0.28" }
+sqllogictest-engines = { path = "../sqllogictest-engines", version = "0.28" }
 tokio = { version = "1", features = [
     "rt",
     "rt-multi-thread",

sqllogictest-bin/Cargo.toml

@@ -1,6 +1,7 @@
 mod engines;
 
 use std::collections::{BTreeMap, BTreeSet, HashSet};
+use std::hash::{DefaultHasher, Hash, Hasher};
 use std::io::{stdout, Read, Seek, SeekFrom, Write};
 use std::path::{Path, PathBuf};
 use std::time::{Duration, Instant};
@@ -19,7 +20,7 @@ use rand::seq::SliceRandom;
 use sqllogictest::substitution::well_known;
 use sqllogictest::{
     default_column_validator, default_normalizer, default_validator, update_record_with_output,
-    AsyncDB, Injected, MakeConnection, Record, Runner,
+    AsyncDB, Injected, MakeConnection, Partitioner, Record, Runner,
 };
 use tokio_util::task::AbortOnDropHandle;
 
@@ -32,6 +33,10 @@ pub enum Color {
     Never,
 }
 
+// Env keys for partitioning.
+const PARTITION_ID_ENV_KEY: &str = "SLT_PARTITION_ID";
+const PARTITION_COUNT_ENV_KEY: &str = "SLT_PARTITION_COUNT";
+
 #[derive(Parser, Debug, Clone)]
 #[clap(about, version, author)]
 struct Opt {
@@ -112,6 +117,18 @@ struct Opt {
     /// The engine name is a label by default.
     #[clap(long = "label")]
     labels: Vec<String>,
+
+    /// Partition ID for sharding the test files. When used with `partition_count`,
+    /// divides the test files into shards based on the hash of the file path.
+    ///
+    /// Useful for running tests in parallel across multiple CI jobs. Currently
+    /// automatically configured in Buildkite.
+    #[clap(long, env = PARTITION_ID_ENV_KEY)]
+    partition_id: Option<u64>,
+
+    /// Total number of partitions for test sharding. More details in `partition_id`.
+    #[clap(long, env = PARTITION_COUNT_ENV_KEY)]
+    partition_count: Option<u64>,
 }
 
 /// Connection configuration.
@@ -138,10 +155,62 @@ impl DBConfig {
     }
 }
 
+struct HashPartitioner {
+    count: u64,
+    id: u64,
+}
+
+impl HashPartitioner {
+    fn new(count: u64, id: u64) -> Result<Self> {
+        if count == 0 {
+            bail!("partition count must be greater than zero");
+        }
+        if id >= count {
+            bail!("partition id (zero-based) must be less than count");
+        }
+        Ok(Self { count, id })
+    }
+}
+
+impl Partitioner for HashPartitioner {
+    fn matches(&self, file_name: &str) -> bool {
+        let mut hasher = DefaultHasher::new();
+        file_name.hash(&mut hasher);
+        hasher.finish() % self.count == self.id
+    }
+}
+
+#[allow(clippy::needless_return)]
+fn import_partition_config_from_ci() {
+    if std::env::var_os(PARTITION_ID_ENV_KEY).is_some()
+        || std::env::var_os(PARTITION_COUNT_ENV_KEY).is_some()
+    {
+        // Ignore if already set.
+        return;
+    }
+
+    // Buildkite
+    {
+        const ID: &str = "BUILDKITE_PARALLEL_JOB";
+        const COUNT: &str = "BUILDKITE_PARALLEL_JOB_COUNT";
+
+        if let (Some(id), Some(count)) = (std::env::var_os(ID), std::env::var_os(COUNT)) {
+            std::env::set_var(PARTITION_ID_ENV_KEY, id);
+            std::env::set_var(PARTITION_COUNT_ENV_KEY, count);
+            eprintln!("Imported partition config from Buildkite.");
+            return;
+        }
+    }
+
+    // TODO: more CI providers
+}
+
 #[tokio::main]
 pub async fn main() -> Result<()> {
     tracing_subscriber::fmt::init();
 
+    import_partition_config_from_ci();
+
     let cli = Opt::command().disable_help_flag(true).arg(
         Arg::new("help")
             .long("help")
@@ -167,6 +236,8 @@ pub async fn main() -> Result<()> {
         r#override,
         format,
         labels,
+        partition_count,
+        partition_id,
     } = Opt::from_arg_matches(&matches)
         .map_err(|err| err.exit())
         .unwrap();
@@ -205,17 +276,34 @@ pub async fn main() -> Result<()> {
         Color::Auto => {}
     }
 
+    let partitioner = if let Some(count) = partition_count {
+        let id = partition_id.context("parallel job count is specified but job id is not")?;
+        Some(HashPartitioner::new(count, id)?)
+    } else {
+        None
+    };
+
     let glob_patterns = files;
-    let mut files: Vec<PathBuf> = Vec::new();
-    for glob_pattern in glob_patterns.into_iter() {
-        let pathbufs = glob::glob(&glob_pattern).context("failed to read glob pattern")?;
-        for pathbuf in pathbufs.into_iter().try_collect::<_, Vec<_>, _>()? {
-            files.push(pathbuf)
+    let mut all_files = Vec::new();
+
+    for glob_pattern in glob_patterns {
+        let mut files: Vec<PathBuf> = glob::glob(&glob_pattern)
+            .context("failed to read glob pattern")?
+            .try_collect()?;
+
+        // Test against partitioner only if there are multiple files matched, e.g., expanded from an `*`.
+        if files.len() > 1 {
+            if let Some(partitioner) = &partitioner {
+                let len = files.len();
+                files.retain(|path| partitioner.matches(path.to_str().unwrap()));
+                let len_after = files.len();
+                eprintln!(
+                    "Running {len_after} out of {len} test cases for glob pattern \"{glob_pattern}\" based on partitioning.",
+                );
+            }
         }
-    }
 
-    if files.is_empty() {
-        bail!("no test case found");
+        all_files.extend(files);
     }
 
     let config = DBConfig {
@@ -227,7 +315,7 @@ pub async fn main() -> Result<()> {
     };
 
     if r#override || format {
-        return update_test_files(files, &engine, config, format).await;
+        return update_test_files(all_files, &engine, config, format).await;
     }
 
     let mut report = Report::new(junit.clone().unwrap_or_else(|| "sqllogictest".to_string()));
@@ -241,7 +329,7 @@ pub async fn main() -> Result<()> {
             jobs,
             keep_db_on_failure,
             &mut test_suite,
-            files,
+            all_files,
             &engine,
             config,
             &labels,
@@ -252,7 +340,7 @@ pub async fn main() -> Result<()> {
     } else {
         run_serial(
             &mut test_suite,
-            files,
+            all_files,
             &engine,
             config,
             &labels,

sqllogictest-bin/src/main.rs

@@ -20,7 +20,7 @@ postgres-types = { version = "0.2.8", features = ["derive", "with-chrono-0_4"] }
 rust_decimal = { version = "1.36.0", features = ["tokio-pg"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
-sqllogictest = { path = "../sqllogictest", version = "0.27" }
+sqllogictest = { path = "../sqllogictest", version = "0.28" }
 thiserror = "2"
 tokio = { version = "1", features = [
     "rt",

sqllogictest-engines/Cargo.toml

@@ -522,6 +522,27 @@ pub fn strict_column_validator<T: ColumnType>(actual: &Vec<T>, expected: &Vec<T>
             .any(|(actual_column, expected_column)| actual_column != expected_column)
 }
 
+/// Decide whether a test file should be run. Useful for partitioning tests into multiple
+/// parallel machines to speed up test runs.
+pub trait Partitioner: Send + Sync + 'static {
+    /// Returns true if the given file name matches the partition and should be run.
+    fn matches(&self, file_name: &str) -> bool;
+}
+
+impl<F> Partitioner for F
+where
+    F: Fn(&str) -> bool + Send + Sync + 'static,
+{
+    fn matches(&self, file_name: &str) -> bool {
+        self(file_name)
+    }
+}
+
+/// The default partitioner matches all files.
+pub fn default_partitioner(_file_name: &str) -> bool {
+    true
+}
+
 #[derive(Default)]
 pub(crate) struct RunnerLocals {
     /// The temporary directory. Test cases can use `__TEST_DIR__` to refer to this directory.
@@ -560,6 +581,7 @@ pub struct Runner<D: AsyncDB, M: MakeConnection<Conn = D>> {
     // normalizer is used to normalize the result text
     normalizer: Normalizer,
     column_type_validator: ColumnTypeValidator<D::ColumnType>,
+    partitioner: Arc<dyn Partitioner>,
     substitution_on: bool,
     sort_mode: Option<SortMode>,
     result_mode: Option<ResultMode>,
@@ -580,6 +602,7 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
             validator: default_validator,
             normalizer: default_normalizer,
             column_type_validator: default_column_validator,
+            partitioner: Arc::new(default_partitioner),
             substitution_on: false,
             sort_mode: None,
             result_mode: None,
@@ -611,6 +634,13 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
         self.column_type_validator = validator;
     }
 
+    /// Set the partitioner for the runner. Only files that match the partitioner will be run.
+    ///
+    /// This only takes effect when running tests in parallel.
+    pub fn with_partitioner(&mut self, partitioner: impl Partitioner + 'static) {
+        self.partitioner = Arc::new(partitioner);
+    }
+
     pub fn with_hash_threshold(&mut self, hash_threshold: usize) {
         self.hash_threshold = hash_threshold;
     }
@@ -1281,13 +1311,18 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
     {
         let files = glob::glob(glob).expect("failed to read glob pattern");
         let mut tasks = vec![];
-        // let conn_builder = Arc::new(conn_builder);
 
         for (idx, file) in files.enumerate() {
             // for every slt file, we create a database against table conflict
             let file = file.unwrap();
-            let db_name = file.to_str().expect("not a UTF-8 filename");
-            let db_name = db_name.replace([' ', '.', '-', '/'], "_");
+            let filename = file.to_str().expect("not a UTF-8 filename");
+
+            // Skip files that don't match the partitioner.
+            if !self.partitioner.matches(filename) {
+                continue;
+            }
+
+            let db_name = filename.replace([' ', '.', '-', '/'], "_");
 
             self.conn
                 .run_default(&format!("CREATE DATABASE {db_name};"))
@@ -1305,6 +1340,7 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                 validator: self.validator,
                 normalizer: self.normalizer,
                 column_type_validator: self.column_type_validator,
+                partitioner: self.partitioner.clone(),
                 substitution_on: self.substitution_on,
                 sort_mode: self.sort_mode,
                 result_mode: self.result_mode,