Fix statpopgen

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

Cargo.lock

@@ -196,6 +196,7 @@ termtree = { version = "0.5" }
 thiserror = "2.0.3"
 tokio = { version = "1.47" }
 tokio-stream = "0.1.17"
+tokio-util = "0.7.17"
 tpchgen = { version = "2" }
 tpchgen-arrow = { version = "2" }
 tracing = { version = "0.1.41", default-features = false }

Cargo.toml

@@ -199,7 +199,7 @@ impl DuckClient {
                 .unwrap_or_else(|| format!("*.{}", extension));
 
             let table_path = format!("{base_dir}{pattern}");
-            dbg!(&table_path);
+
             commands.push_str(&format!(
                 "CREATE {} IF NOT EXISTS {} AS SELECT * FROM read_{}('{}');\n",
                 object.to_str(),
@@ -211,8 +211,6 @@ impl DuckClient {
 
         trace!("Executing table registration commands: {commands}");
 
-        dbg!(&commands);
-
         self.execute_query(&commands)?;
 
         Ok(())

benchmarks/ddb-bench/src/lib.rs

@@ -40,14 +40,15 @@ parking_lot = { workspace = true }
 parquet = { workspace = true, features = ["async"] }
 rand = { workspace = true }
 regex = { workspace = true }
-reqwest = { workspace = true }
+reqwest = { workspace = true, features = ["stream"] }
 serde = { workspace = true }
 serde_json = { workspace = true }
 sysinfo = { workspace = true }
 tabled = { workspace = true, features = ["std"] }
 target-lexicon = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 tokio-stream = { workspace = true }
+tokio-util = { workspace = true }
 tpchgen = { workspace = true }
 tpchgen-arrow = { workspace = true }
 tracing = { workspace = true, features = ["max_level_debug"] }

vortex-bench/Cargo.toml

@@ -11,8 +11,9 @@ use parquet::arrow::AsyncArrowWriter;
 use reqwest::Client;
 use tokio::fs::File;
 use tokio::io::BufReader;
+use tokio_stream::StreamExt;
+use tokio_util::io::StreamReader;
 use tracing::info;
-use vortex::io::IoBuf;
 
 use super::StatPopGenBenchmark;
 use crate::idempotent_async;
@@ -30,6 +31,7 @@ impl StatPopGenBenchmark {
             StatPopGenBenchmark::FILE_NAME
         );
         let parquet_output_path = self.parquet_path()?;
+
         idempotent_async(
             &parquet_output_path,
             async |parquet_output_path| -> Result<()> {
@@ -48,10 +50,14 @@ impl StatPopGenBenchmark {
                     .error_for_status()
                     .context("reqwest bad status")?;
 
-                let body = response.bytes().await?;
+                // The file is BIG so we only want to download a part of it
+                let byte_stream = response
+                    .bytes_stream()
+                    .map(|x| x.map_err(std::io::Error::other));
+                let stream_reader = StreamReader::new(byte_stream);
 
                 // Wrap in BGZF reader
-                let buf_reader = BufReader::new(body.as_slice());
+                let buf_reader = BufReader::new(stream_reader);
                 let mut bgzf_reader = noodles_bgzf::r#async::io::Reader::new(buf_reader);
 
                 // Read and parse VCF header

vortex-bench/src/statpopgen/download_vcf.rs

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::fs;
-use std::path::Path;
 use std::path::PathBuf;
 
 use anyhow::Result;
@@ -14,6 +13,7 @@ use crate::BenchmarkDataset;
 use crate::Format;
 use crate::IdempotentPath;
 use crate::TableSpec;
+use crate::workspace_root;
 
 /// Statistical population genetics benchmark implementation.
 ///
@@ -75,8 +75,10 @@ impl StatPopGenBenchmark {
             }
         };
 
-        let data_url = Url::from_directory_path("statpopgen".to_data_path())
-            .map_err(|_| anyhow::anyhow!("bad data path?"))?;
+        let data_path = "statspopgen".to_data_path().join(format!("{n_rows}/"));
+
+        let data_url =
+            Url::from_directory_path(data_path).map_err(|_| anyhow::anyhow!("bad data path?"))?;
 
         Ok(Self {
             data_url,
@@ -92,7 +94,7 @@ impl StatPopGenBenchmark {
     /// The path follows the pattern: `{data_url}/{n_rows}/parquet/gnomad.genomes.v3.1.2.hgdp_tgp.chr21.parquet`
     pub fn parquet_path(&self) -> Result<PathBuf> {
         self.data_url
-            .join(&(self.n_rows.to_string() + "/parquet/"))?
+            .join("parquet/")?
             .join(&format!("{}.parquet", StatPopGenBenchmark::FILE_NAME))?
             .to_file_path()
             .map_err(|_| anyhow::anyhow!("Failed to convert data URL to filesystem path - ensure data_url uses 'file://' scheme"))
@@ -104,7 +106,7 @@ impl StatPopGenBenchmark {
     /// The path follows the pattern: `{data_url}/{n_rows}/vortex-file-compressed/gnomad.genomes.v3.1.2.hgdp_tgp.chr21.vortex`
     pub fn vortex_path(&self) -> Result<PathBuf> {
         self.data_url
-            .join(&(self.n_rows.to_string() + "/vortex-file-compressed/"))?
+            .join("vortex-file-compressed/")?
             .join(&format!("{}.vortex", StatPopGenBenchmark::FILE_NAME))?
             .to_file_path()
             .map_err(|_| anyhow::anyhow!("Failed to convert data URL to filesystem path - ensure data_url uses 'file://' scheme"))
@@ -116,7 +118,7 @@ impl StatPopGenBenchmark {
     /// The path follows the pattern: `{data_url}/{n_rows}/vortex-compact/{StatPopGenBenchmark::FILE_NAME}.vortex`
     pub fn vortex_compact_path(&self) -> Result<PathBuf> {
         self.data_url
-            .join(&(self.n_rows.to_string() + "/vortex-compact/"))?
+            .join("vortex-compact/")?
             .join(&format!("{}.vortex", StatPopGenBenchmark::FILE_NAME))?
             .to_file_path()
             .map_err(|_| anyhow::anyhow!("Failed to convert data URL to filesystem path - ensure data_url uses 'file://' scheme"))
@@ -126,7 +128,8 @@ impl StatPopGenBenchmark {
 #[async_trait::async_trait]
 impl Benchmark for StatPopGenBenchmark {
     fn queries(&self) -> Result<Vec<(usize, String)>> {
-        let queries_file = Path::new(env!("CARGO_MANIFEST_DIR"))
+        let queries_file = workspace_root()
+            .join("vortex-bench")
             .join("statpopgen")
             .with_extension("sql");
         let contents = fs::read_to_string(queries_file)?;