vortex-bench: run ClickBench on its first N files

vortex-bench/sql/clickbench.md

@@ -28,3 +28,6 @@ pushdown and zone-map-style pruning on a subset of the queries. See
 vx-bench run clickbench --engine datafusion,duckdb --format parquet,vortex
 vx-bench run clickbench-sorted --engine datafusion,duckdb --format parquet,vortex
 ```
+
+`--opt files=N` runs the partitioned dataset's first N files (of 100, about 1M rows each)
+from its own data directory. The known per-query row counts are checked only at 100 files.

vortex-bench/src/clickbench/benchmark.rs

@@ -18,27 +18,38 @@ use crate::utils::file::resolve_data_url;
 /// ClickBench benchmark implementation
 pub struct ClickBenchBenchmark {
     pub flavor: Flavor,
+    /// How many of the partitioned dataset's files to use, from 1 to [`CLICKBENCH_FILES`].
+    pub files: u32,
     pub queries_file: Option<String>,
     pub data_url: Url,
 }
 
 impl ClickBenchBenchmark {
     pub fn new(
         flavor: Flavor,
+        files: u32,
         queries_file: Option<String>,
         use_remote_data_dir: Option<String>,
     ) -> Result<Self> {
-        let url = Self::create_data_url(use_remote_data_dir.as_deref(), flavor)?;
+        anyhow::ensure!(
+            (1..=CLICKBENCH_FILES).contains(&files),
+            "ClickBench files must be between 1 and {CLICKBENCH_FILES}, got {files}"
+        );
+        anyhow::ensure!(
+            flavor == Flavor::Partitioned || files == CLICKBENCH_FILES,
+            "ClickBench files applies only to the partitioned flavor"
+        );
+        let url = resolve_data_url(
+            use_remote_data_dir.as_deref(),
+            &clickbench_data_dir(flavor, files),
+        )?;
         Ok(Self {
             flavor,
+            files,
             queries_file,
             data_url: url,
         })
     }
-
-    fn create_data_url(remote_data_dir: Option<&str>, flavor: Flavor) -> Result<Url> {
-        resolve_data_url(remote_data_dir, &format!("clickbench_{flavor}"))
-    }
 }
 
 /// ClickBench sorted by event time, with shard filenames shuffled to exercise sort pushdown.
@@ -89,19 +100,21 @@ impl Benchmark for ClickBenchBenchmark {
             return Ok(());
         }
 
-        let basepath = clickbench_flavor(self.flavor).to_data_path();
-        self.flavor.download(basepath).await?;
+        let basepath = clickbench_data_dir(self.flavor, self.files).to_data_path();
+        self.flavor.download(basepath, self.files).await?;
 
         Ok(())
     }
 
     fn expected_row_counts(&self) -> Option<Vec<usize>> {
-        Some(clickbench_expected_row_counts())
+        // The known row counts hold only for the complete dataset.
+        (self.files == CLICKBENCH_FILES).then(clickbench_expected_row_counts)
     }
 
     fn dataset(&self) -> BenchmarkDataset {
         BenchmarkDataset::ClickBench {
             flavor: self.flavor,
+            files: (self.files != CLICKBENCH_FILES).then_some(self.files),
         }
     }
 
@@ -110,7 +123,7 @@ impl Benchmark for ClickBenchBenchmark {
     }
 
     fn dataset_display(&self) -> String {
-        format!("clickbench_{}", self.flavor)
+        clickbench_data_dir(self.flavor, self.files)
     }
 
     fn data_url(&self) -> &Url {
@@ -168,6 +181,36 @@ impl Benchmark for ClickBenchSortedBenchmark {
     }
 }
 
-fn clickbench_flavor(flavor: Flavor) -> String {
-    format!("clickbench_{flavor}")
+/// The data directory for the first `files` files, e.g. `clickbench_partitioned_10`. The full
+/// dataset stays `clickbench_partitioned`, not `clickbench_partitioned_100`, for backcompat.
+fn clickbench_data_dir(flavor: Flavor, files: u32) -> String {
+    if files == CLICKBENCH_FILES {
+        format!("clickbench_{flavor}")
+    } else {
+        format!("clickbench_{flavor}_{files}")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn subsets_skip_the_full_dataset_row_counts() -> Result<()> {
+        let subset = ClickBenchBenchmark::new(Flavor::Partitioned, 10, None, None)?;
+        assert!(subset.expected_row_counts().is_none());
+        assert!(
+            subset
+                .data_url
+                .path()
+                .ends_with("/clickbench_partitioned_10/")
+        );
+
+        let full = ClickBenchBenchmark::new(Flavor::Partitioned, CLICKBENCH_FILES, None, None)?;
+        assert!(full.expected_row_counts().is_some());
+        assert!(full.data_url.path().ends_with("/clickbench_partitioned/"));
+
+        assert!(ClickBenchBenchmark::new(Flavor::Single, 10, None, None).is_err());
+        Ok(())
+    }
 }

vortex-bench/src/clickbench/data.rs

@@ -37,6 +37,8 @@ use crate::utils::file::temp_download_filepath;
 /// Benchmark and local data directory name for ClickBench sorted by event time.
 pub const CLICKBENCH_SORTED_NAME: &str = "clickbench-sorted";
 const CLICKBENCH_PARTITIONED_NAME: &str = "clickbench_partitioned";
+/// The number of files in the partitioned ClickBench dataset.
+pub const CLICKBENCH_FILES: u32 = 100;
 const SORTED_SHARD_COUNT: usize = 100;
 const SORTED_SHARD_COUNT_U64: u64 = 100;
 const SORTED_SHARD_FILENAME_WIDTH: usize = 3;
@@ -212,7 +214,8 @@ impl Display for Flavor {
 
 impl Flavor {
     // TODO(joe): move these elsewhere.
-    pub async fn download(&self, basepath: impl AsRef<Path>) -> anyhow::Result<()> {
+    /// Download the dataset; the partitioned flavor fetches only its first `files` files.
+    pub async fn download(&self, basepath: impl AsRef<Path>, files: u32) -> anyhow::Result<()> {
         let basepath = basepath.as_ref();
         match self {
             Flavor::Single => {
@@ -224,9 +227,9 @@ impl Flavor {
             Flavor::Partitioned => {
                 // The clickbench-provided file is missing some higher-level type info, so we reprocess it
                 // to add that info, see https://github.com/ClickHouse/ClickBench/issues/7.
-                info!("Downloading 100 ClickBench parquet shards");
+                info!("Downloading {files} ClickBench parquet shards");
                 let parquet_dir = basepath.join(Format::Parquet.name());
-                let downloads = (0_u32..100).map(|idx| {
+                let downloads = (0..files).map(|idx| {
                     let output_path = parquet_dir.join(format!("hits_{idx}.parquet"));
                     let url = format!("https://pub-3ba949c0f0354ac18db1f0f14f0a2c52.r2.dev/clickbench/parquet_many/hits_{idx}.parquet");
                     (output_path, url)
@@ -246,7 +249,9 @@ impl Flavor {
 pub async fn generate_sorted_clickbench(basepath: impl AsRef<Path>) -> anyhow::Result<()> {
     let basepath = basepath.as_ref();
     let source_base = CLICKBENCH_PARTITIONED_NAME.to_data_path();
-    Flavor::Partitioned.download(&source_base).await?;
+    Flavor::Partitioned
+        .download(&source_base, CLICKBENCH_FILES)
+        .await?;
 
     let source_parquet_dir = source_base.join(Format::Parquet.name());
     let output_parquet_dir = basepath.join(Format::Parquet.name());

vortex-bench/src/datasets/mod.rs

@@ -67,7 +67,12 @@ pub enum BenchmarkDataset {
     #[serde(rename = "tpcds")]
     TpcDS { scale_factor: String },
     #[serde(rename = "clickbench")]
-    ClickBench { flavor: Flavor },
+    ClickBench {
+        flavor: Flavor,
+        /// How many partitioned files ran, when fewer than the complete dataset.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        files: Option<u32>,
+    },
     #[serde(rename = "clickbench-sorted")]
     ClickBenchSorted,
     #[serde(rename = "public-bi")]
@@ -114,9 +119,12 @@ impl Display for BenchmarkDataset {
             BenchmarkDataset::Appian => write!(f, "appian"),
             BenchmarkDataset::TpcH { scale_factor } => write!(f, "tpch(sf={scale_factor})"),
             BenchmarkDataset::TpcDS { scale_factor } => write!(f, "tpcds(sf={scale_factor})"),
-            BenchmarkDataset::ClickBench { flavor, .. } => match flavor {
-                Flavor::Partitioned => write!(f, "clickbench-partitioned"),
-                Flavor::Single => write!(f, "clickbench-single"),
+            BenchmarkDataset::ClickBench { flavor, files } => match (flavor, files) {
+                (Flavor::Partitioned, Some(files)) => {
+                    write!(f, "clickbench-partitioned(files={files})")
+                }
+                (Flavor::Partitioned, None) => write!(f, "clickbench-partitioned"),
+                (Flavor::Single, _) => write!(f, "clickbench-single"),
             },
             BenchmarkDataset::ClickBenchSorted => write!(f, "clickbench-sorted"),
             BenchmarkDataset::PublicBi { name } => write!(f, "public-bi({name})"),

vortex-bench/src/lib.rs

@@ -12,6 +12,7 @@ use std::sync::LazyLock;
 use anyhow::bail;
 use appian::AppianBenchmark;
 use clap::ValueEnum;
+use clickbench::CLICKBENCH_FILES;
 use clickbench::ClickBenchBenchmark;
 use clickbench::ClickBenchSortedBenchmark;
 use clickbench::Flavor;
@@ -387,8 +388,9 @@ pub fn create_benchmark(b: BenchmarkArg, opts: &Opts) -> anyhow::Result<Box<dyn
         }
         BenchmarkArg::ClickBench => {
             let flavor = opts.get_as::<Flavor>("flavor").unwrap_or_default();
+            let files = opts.get_as::<u32>("files").unwrap_or(CLICKBENCH_FILES);
             let remote_data_dir = opts.get_as::<String>(REMOTE_DATA_KEY);
-            let benchmark = ClickBenchBenchmark::new(flavor, None, remote_data_dir)?;
+            let benchmark = ClickBenchBenchmark::new(flavor, files, None, remote_data_dir)?;
             Ok(Box::new(benchmark) as _)
         }
         BenchmarkArg::ClickBenchSorted => {

vortex-bench/src/v3.rs

@@ -276,7 +276,7 @@ fn canonical_tpc_scale_factor(scale_factor: &str) -> String {
 /// |---|---|---|---|---|
 /// | `TpcH { scale_factor }`     | `tpch`         | `None`              | TPC SF as string (`"1"`, `"10"`, `"100"`, `"1000"`) | Run through `canonical_tpc_scale_factor` so `"1.0"` and `"1"` collapse. |
 /// | `TpcDS { scale_factor }`    | `tpcds`        | `None`              | TPC SF as string                                    | Same canonicalization as TPC-H. |
-/// | `ClickBench { flavor: _ }`  | `clickbench`   | `None`              | `None`                                              | Migrate path drops flavor; live emitter matches so historical and live merge. |
+/// | `ClickBench { files, .. }`  | `clickbench`   | `files=N` or `None` | `None`                                              | Migrate path drops flavor; live emitter matches so historical and live merge. Only subsets set `files`. |
 /// | `ClickBenchSorted`          | `clickbench-sorted` | `None`          | `None`                                              | New live-only suite; keep separate from unsorted ClickBench history. |
 /// | `StatPopGen { n_rows: _ }`  | `statpopgen`   | `None`              | `None`                                              | Migrate path carries no SF for this suite; live drops it for the same reason. |
 /// | `PolarSignals { n_rows: _ }`| `polarsignals` | `None`              | `None`                                              | Same as StatPopGen. |
@@ -303,7 +303,11 @@ pub fn benchmark_dataset_dims(d: &BenchmarkDataset) -> (String, Option<String>,
         // v2 record names did not encode flavor, so the live emitter does the
         // same to keep historical and live records in one `clickbench` group.
         // Flavor is fixed per CI matrix entry and recoverable from there.
-        BenchmarkDataset::ClickBench { .. } => ("clickbench".to_string(), None, None),
+        BenchmarkDataset::ClickBench { files, .. } => (
+            "clickbench".to_string(),
+            files.map(|files| format!("files={files}")),
+            None,
+        ),
         BenchmarkDataset::ClickBenchSorted => ("clickbench-sorted".to_string(), None, None),
         BenchmarkDataset::PublicBi { name } => ("public-bi".to_string(), Some(name.clone()), None),
         // StatPopGen / PolarSignals: the migrate path (v2 → v3 backfill) does
@@ -608,6 +612,7 @@ mod tests {
             target: Target::new(Engine::DuckDB, Format::Parquet),
             benchmark_dataset: BenchmarkDataset::ClickBench {
                 flavor: Flavor::Partitioned,
+                files: None,
             },
             benchmark_runner: "ci-runner".to_string(),
             storage: "s3".to_string(),
@@ -718,12 +723,14 @@ mod tests {
             (
                 BenchmarkDataset::ClickBench {
                     flavor: Flavor::Partitioned,
+                    files: None,
                 },
                 "clickbench",
             ),
             (
                 BenchmarkDataset::ClickBench {
                     flavor: Flavor::Single,
+                    files: None,
                 },
                 "clickbench",
             ),
@@ -747,6 +754,32 @@ mod tests {
         }
     }
 
+    #[test]
+    fn clickbench_subsets_are_distinct_from_the_complete_dataset() -> anyhow::Result<()> {
+        let subset = BenchmarkDataset::ClickBench {
+            flavor: Flavor::Partitioned,
+            files: Some(10),
+        };
+        assert_eq!(
+            benchmark_dataset_dims(&subset),
+            ("clickbench".to_string(), Some("files=10".to_string()), None)
+        );
+
+        let full = BenchmarkDataset::ClickBench {
+            flavor: Flavor::Partitioned,
+            files: None,
+        };
+        assert_eq!(
+            serde_json::to_string(&full)?,
+            r#"{"clickbench":{"flavor":"Partitioned"}}"#
+        );
+        assert_eq!(
+            serde_json::from_str::<BenchmarkDataset>(r#"{"clickbench":{"flavor":"Partitioned"}}"#)?,
+            full
+        );
+        Ok(())
+    }
+
     #[test]
     fn clickbench_sorted_dims_are_distinct_from_clickbench() {
         let (dataset, variant, scale_factor) =