perf: wrap download_many call sites with idempotent_dir_async

Use the .success marker pattern at all download_many call sites so that
a completed batch of downloads is skipped in full on re-runs:

- clickbench Partitioned: wraps the 100-shard parquet_dir download
- public_bi download_bzips: wraps the csv_bzip2 directory download
- vector_dataset download: wraps the train-shard directory download

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-bench/src/clickbench/data.rs

@@ -22,6 +22,7 @@ use crate::Format;
 pub use crate::conversions::convert_parquet_directory_to_vortex;
 use crate::datasets::data_downloads::download_data;
 use crate::datasets::data_downloads::download_many;
+use crate::idempotent_dir_async;
 
 pub static HITS_SCHEMA: LazyLock<Schema> = LazyLock::new(|| {
     use DataType::*;
@@ -195,12 +196,16 @@ impl Flavor {
                 // to add that info, see https://github.com/ClickHouse/ClickBench/issues/7.
                 info!("Downloading 100 ClickBench parquet shards");
                 let parquet_dir = basepath.join(Format::Parquet.name());
-                let downloads = (0_u32..100).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)
-                });
-                download_many(downloads).await?;
+                let parquet_dir2 = parquet_dir.clone();
+                idempotent_dir_async(&parquet_dir, || async move {
+                    let downloads = (0_u32..100).map(|idx| {
+                        let output_path = parquet_dir2.join(format!("hits_{idx}.parquet"));
+                        let url = format!("https://pub-3ba949c0f0354ac18db1f0f14f0a2c52.r2.dev/clickbench/parquet_many/hits_{idx}.parquet");
+                        (output_path, url)
+                    });
+                    download_many(downloads).await.map(|_| ())
+                })
+                .await?;
             }
         }
         Ok(())

vortex-bench/src/public_bi.rs

@@ -45,6 +45,7 @@ use crate::datasets::Dataset;
 use crate::datasets::data_downloads::decompress_bz2;
 use crate::datasets::data_downloads::download_many;
 use crate::idempotent_async;
+use crate::idempotent_dir_async;
 use crate::workspace_root;
 
 pub static PBI_DATASETS: LazyLock<PBIDatasets> = LazyLock::new(|| {
@@ -289,13 +290,21 @@ pub struct PBIData {
 
 impl PBIData {
     async fn download_bzips(&self) -> anyhow::Result<()> {
-        let downloads = self.tables.iter().map(|table| {
-            (
-                self.get_file_path(&table.name, FileType::CsvBzip2),
-                table.data_url.as_str().to_owned(),
-            )
-        });
-        download_many(downloads).await?;
+        let bzip_dir = self.base_path.join(FileType::CsvBzip2.name());
+        let downloads: Vec<(PathBuf, String)> = self
+            .tables
+            .iter()
+            .map(|table| {
+                (
+                    self.get_file_path(&table.name, FileType::CsvBzip2),
+                    table.data_url.as_str().to_owned(),
+                )
+            })
+            .collect();
+        idempotent_dir_async(&bzip_dir, || async move {
+            download_many(downloads).await.map(|_| ())
+        })
+        .await?;
         Ok(())
     }
 

vortex-bench/src/vector_dataset/download.rs

@@ -20,6 +20,7 @@ use anyhow::Result;
 
 use crate::datasets::data_downloads::download_data;
 use crate::datasets::data_downloads::download_many;
+use crate::idempotent_dir_async;
 use crate::vector_dataset::catalog::VectorDataset;
 use crate::vector_dataset::layout::LayoutSpec;
 use crate::vector_dataset::layout::TrainLayout;
@@ -102,14 +103,17 @@ pub async fn download(ds: VectorDataset, layout: TrainLayout) -> Result<DatasetP
     let train_targets = paths::train_files(ds, layout, spec.num_files());
     debug_assert_eq!(urls.len(), train_targets.len());
 
+    let train_dir = paths::train_dir(ds, layout);
     let train_downloads: Vec<(PathBuf, String)> = train_targets
         .iter()
         .cloned()
         .zip(urls.into_iter())
         .collect();
-    download_many(train_downloads)
-        .await
-        .with_context(|| format!("download train shards for {}", ds.name()))?;
+    idempotent_dir_async(&train_dir, || async move {
+        download_many(train_downloads).await.map(|_| ())
+    })
+    .await
+    .with_context(|| format!("download train shards for {}", ds.name()))?;
 
     let test = download_data(paths::test_path(ds, layout), &test_url(ds))
         .await