Add single-file clickbench benchmark (#2597)

This is the original thing we wanted. It didn't really work on our
runners at the time but hopefully as we're rebuilding some of that infra
we can go back to running it, and it'll be nice to at least run it
locally.

Author: AdamGS

bench-vortex/src/bin/clickbench.rs

@@ -1,25 +1,23 @@
-use std::fs::{self, File};
-use std::path::PathBuf;
+use std::fs::{self};
 use std::time::{Duration, Instant};
 
-use bench_vortex::clickbench::{self, HITS_SCHEMA, clickbench_queries};
+use bench_vortex::clickbench::{self, Flavor, HITS_SCHEMA, clickbench_queries};
 use bench_vortex::display::{DisplayFormat, RatioMode, print_measurements_json, render_table};
 use bench_vortex::measurements::QueryMeasurement;
 use bench_vortex::metrics::export_plan_spans;
 use bench_vortex::{
     Format, IdempotentPath as _, default_env_filter, execute_physical_plan,
-    feature_flagged_allocator, get_session_with_cache, idempotent, physical_plan,
+    feature_flagged_allocator, get_session_with_cache, physical_plan,
 };
 use clap::Parser;
 use datafusion_physical_plan::display::DisplayableExecutionPlan;
 use indicatif::ProgressBar;
 use itertools::Itertools;
-use log::{info, warn};
-use rayon::iter::{IntoParallelIterator, ParallelIterator as _};
+use log::warn;
 use tokio::runtime::Builder;
 use tracing::info_span;
 use tracing_futures::Instrument;
-use vortex::error::{VortexExpect, vortex_panic};
+use vortex::error::vortex_panic;
 
 feature_flagged_allocator!();
 
@@ -46,9 +44,11 @@ struct Args {
     emulate_object_store: bool,
     #[arg(long)]
     export_spans: bool,
+    #[arg(long, value_enum, default_value_t = Flavor::Partitioned)]
+    flavor: Flavor,
 }
 
-fn main() {
+fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
     // Capture `RUST_LOG` configuration
@@ -95,45 +95,10 @@ fn main() {
         None => Builder::new_multi_thread().enable_all().build(),
     }
     .expect("Failed building the Runtime");
-    let basepath = "clickbench".to_data_path();
+    let basepath = format!("clickbench_{}", args.flavor).to_data_path();
     let client = reqwest::blocking::Client::default();
 
-    // 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.
-    (0_u32..100).into_par_iter().for_each(|idx| {
-        let output_path = basepath.join("parquet").join(format!("hits_{idx}.parquet"));
-        idempotent(&output_path, |output_path| {
-            info!("Downloading file {idx}");
-            let url = format!("https://pub-3ba949c0f0354ac18db1f0f14f0a2c52.r2.dev/clickbench/parquet_many/hits_{idx}.parquet");
-
-
-            let make_req = || client.get(&url).send();
-            let mut output = None;
-
-            for attempt in 1..4 {
-                match make_req() {
-                    Ok(r) => {
-                          output = Some(r.error_for_status());
-                          break;
-                    },
-                    Err(e) => {
-                        warn!("Request for file {idx} timed out, retying for the {attempt} time");
-                        output = Some(Err(e));
-                    }
-                }
-
-                // Very basic backoff mechanism
-                std::thread::sleep(Duration::from_secs(attempt));
-            }
-
-            let mut response = output.vortex_expect("Must have value here")?;
-            let mut file = File::create(output_path)?;
-            response.copy_to(&mut file)?;
-
-            anyhow::Ok(PathBuf::from(output_path))
-        })
-        .unwrap();
-    });
+    args.flavor.download(&client, basepath.as_path())?;
 
     let queries = match args.queries.clone() {
         None => clickbench_queries(),
@@ -151,6 +116,7 @@ fn main() {
         let session_context = get_session_with_cache(args.emulate_object_store);
         let context = session_context.clone();
         let mut plans = Vec::new();
+
         match format {
             Format::Parquet => runtime.block_on(async {
                 clickbench::register_parquet_files(
@@ -254,4 +220,6 @@ fn main() {
         }
         DisplayFormat::GhJson => print_measurements_json(all_measurements).unwrap(),
     }
+
+    Ok(())
 }

bench-vortex/src/clickbench.rs

@@ -1,15 +1,21 @@
-use std::path::Path;
+use std::fs::File;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, LazyLock};
+use std::time::Duration;
 
 use arrow_schema::{DataType, Field, Schema, TimeUnit};
+use clap::ValueEnum;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::listing::{
     ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
 };
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use futures::{StreamExt, TryStreamExt, stream};
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use reqwest::IntoUrl;
+use reqwest::blocking::Response;
 use tokio::fs::{OpenOptions, create_dir_all};
-use tracing::info;
+use tracing::{info, warn};
 use vortex::TryIntoArray;
 use vortex::dtype::DType;
 use vortex::dtype::arrow::FromArrowType;
@@ -19,7 +25,7 @@ use vortex::layout::{LayoutRegistry, LayoutRegistryExt};
 use vortex::stream::ArrayStreamAdapter;
 use vortex_datafusion::persistent::VortexFormat;
 
-use crate::idempotent_async;
+use crate::{idempotent, idempotent_async};
 
 pub static HITS_SCHEMA: LazyLock<Schema> = LazyLock::new(|| {
     use DataType::*;
@@ -147,20 +153,36 @@ pub async fn register_vortex_files(
     schema: &Schema,
 ) -> anyhow::Result<()> {
     let vortex_dir = input_path.join("vortex");
+    let parquet_path = input_path.join("parquet");
     create_dir_all(&vortex_dir).await?;
 
-    stream::iter(0..100)
-        .map(|idx| {
-            let parquet_file_path = input_path
-                .join("parquet")
-                .join(format!("hits_{idx}.parquet"));
-            let output_path = vortex_dir.join(format!("hits_{idx}.{VORTEX_FILE_EXTENSION}"));
+    let parquet_inputs =
+        std::fs::read_dir(parquet_path.clone())?.collect::<std::io::Result<Vec<_>>>()?;
+    info!(
+        "Found {} parquet files in {}",
+        parquet_inputs.len(),
+        parquet_path.to_str().unwrap()
+    );
+
+    let iter = parquet_inputs
+        .iter()
+        .filter(|entry| entry.path().extension().is_some_and(|e| e == "parquet"));
+
+    stream::iter(iter)
+        .map(|dir_entry| {
+            let filename = {
+                let mut temp = dir_entry.path();
+                temp.set_extension("");
+                temp.file_name().unwrap().to_str().unwrap().to_string()
+            };
+            let parquet_file_path = parquet_path.join(format!("{filename}.parquet"));
+            let output_path = vortex_dir.join(format!("{filename}.{VORTEX_FILE_EXTENSION}"));
             let session = session.clone();
 
             tokio::spawn(async move {
                 let output_path = output_path.clone();
                 idempotent_async(&output_path, move |vtx_file| async move {
-                    info!("Processing file {idx}");
+                    info!("Processing file '{filename}'");
                     let record_batches = session
                         .read_parquet(
                             parquet_file_path.to_str().unwrap(),
@@ -259,3 +281,88 @@ pub fn clickbench_queries() -> Vec<(usize, String)> {
         .enumerate()
         .collect()
 }
+
+#[derive(ValueEnum, Default, Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub enum Flavor {
+    #[default]
+    Partitioned,
+    Single,
+}
+
+impl std::fmt::Display for Flavor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.to_possible_value().unwrap().get_name())
+    }
+}
+
+impl Flavor {
+    pub fn download(
+        &self,
+        client: &reqwest::blocking::Client,
+        basepath: impl AsRef<Path>,
+    ) -> anyhow::Result<()> {
+        let basepath = basepath.as_ref();
+        match self {
+            Flavor::Single => {
+                let output_path = basepath.join("parquet").join("hits.parquet");
+                idempotent(&output_path, |output_path| {
+                    info!("Downloading single clickbench file");
+                    let url = "https://pub-3ba949c0f0354ac18db1f0f14f0a2c52.r2.dev/clickbench/parquet_single/hits.parquet";
+                    let mut response = retry_get(client, url)?;
+                    let mut file = File::create(output_path)?;
+                    response.copy_to(&mut file)?;
+
+                    anyhow::Ok(())
+                })?;
+
+                Ok(())
+            }
+            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.
+                (0_u32..100).into_par_iter().for_each(|idx| {
+                    let output_path = basepath.join("parquet").join(format!("hits_{idx}.parquet"));
+                    idempotent(&output_path, |output_path| {
+                        info!("Downloading file {idx}");
+                        let url = format!("https://pub-3ba949c0f0354ac18db1f0f14f0a2c52.r2.dev/clickbench/parquet_many/hits_{idx}.parquet");
+                        let mut response = retry_get(client, url)?;
+                        let mut file = File::create(output_path)?;
+                        response.copy_to(&mut file)?;
+
+                        anyhow::Ok(PathBuf::from(output_path))
+                    })
+                    .unwrap();
+            });
+
+                Ok(())
+            }
+        }
+    }
+}
+
+fn retry_get(client: &reqwest::blocking::Client, url: impl IntoUrl) -> anyhow::Result<Response> {
+    let url = url.as_str();
+    let make_req = || client.get(url).send();
+
+    let mut output = None;
+
+    for attempt in 1..4 {
+        match make_req().and_then(|r| r.error_for_status()) {
+            Ok(r) => {
+                output = Some(r);
+                break;
+            }
+            Err(e) => {
+                warn!("Request errored with {e}, retying for the {attempt} time");
+            }
+        }
+
+        // Very basic backoff mechanism
+        std::thread::sleep(Duration::from_secs(attempt));
+    }
+
+    match output {
+        Some(v) => Ok(v),
+        None => anyhow::bail!("Exahusted retry attempts for {url}"),
+    }
+}