From 9ee00c39c8a34757861de7426504d93498547ff3 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 28 Jul 2026 17:17:11 +0100 Subject: [PATCH 1/5] Try and replace the rust usage with a somewhat hacky call to uvx Signed-off-by: Adam Gutglick --- Cargo.lock | 18 - Cargo.toml | 3 - vortex-bench/Cargo.toml | 2 - vortex-bench/sql/tpch/README.md | 4 +- vortex-bench/src/conversions.rs | 20 +- vortex-bench/src/datasets/tpch_l_comment.rs | 13 +- vortex-bench/src/tpch/benchmark.rs | 3 +- vortex-bench/src/tpch/tpchgen.rs | 502 +++++------------- vortex-sqllogictest/slt/tpch/generate_data.sh | 2 +- vortex-test/compat-gen/Cargo.toml | 3 - .../src/fixtures/arrays/datasets/tpch.rs | 115 +++- 11 files changed, 252 insertions(+), 433 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 51188781c9a..ed82d929e33 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8938,20 +8938,6 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" -[[package]] -name = "tpchgen" -version = "2.0.2" -source = "git+https://github.com/clflushopt/tpchgen-rs.git?rev=438e9c2dbc25b2fff82c0efc08b3f13b5707874f#438e9c2dbc25b2fff82c0efc08b3f13b5707874f" - -[[package]] -name = "tpchgen-arrow" -version = "2.0.2" -source = "git+https://github.com/clflushopt/tpchgen-rs.git?rev=438e9c2dbc25b2fff82c0efc08b3f13b5707874f#438e9c2dbc25b2fff82c0efc08b3f13b5707874f" -dependencies = [ - "arrow 58.4.0", - "tpchgen", -] - [[package]] name = "tracing" version = "0.1.44" @@ -9431,8 +9417,6 @@ dependencies = [ "tokio", "tokio-stream", "tokio-util", - "tpchgen", - "tpchgen-arrow", "tracing", "tracing-perfetto", "tracing-subscriber", @@ -9523,8 +9507,6 @@ dependencies = [ "sha2", "tempfile", "tokio", - "tpchgen", - "tpchgen-arrow", "vortex", "vortex-array", "vortex-arrow", diff --git a/Cargo.toml b/Cargo.toml index a9d8bfefb85..6a2d306b631 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -267,9 +267,6 @@ tokio = { version = "1.52" } tokio-stream = "0.1.17" tokio-util = "0.7.17" vortex-array-macros = { version = "0.1.0", path = "./vortex-array-macros" } -# Pull these into non-public crates to support DF 58 -tpchgen = { version = "2.0.2", git = "https://github.com/clflushopt/tpchgen-rs.git", rev = "438e9c2dbc25b2fff82c0efc08b3f13b5707874f" } -tpchgen-arrow = { version = "2.0.2", git = "https://github.com/clflushopt/tpchgen-rs.git", rev = "438e9c2dbc25b2fff82c0efc08b3f13b5707874f" } tracing = { version = "0.1.41", default-features = false } tracing-perfetto = "0.1.5" tracing-subscriber = "0.3" diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 6dbe8ffea98..509e8bf8edf 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -68,8 +68,6 @@ 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 } tracing-perfetto = { workspace = true } tracing-subscriber = { workspace = true, features = [ diff --git a/vortex-bench/sql/tpch/README.md b/vortex-bench/sql/tpch/README.md index aa5a16fbab2..76d4cffac69 100644 --- a/vortex-bench/sql/tpch/README.md +++ b/vortex-bench/sql/tpch/README.md @@ -6,8 +6,8 @@ It stresses scan throughput, filter pushdown, joins, and aggregations over mostl and date columns. The queries in this directory (`q1.sql` ... `q22.sql`) are the standard TPC-H queries. -Data is generated deterministically by [`tpchgen`](../../src/tpch/tpchgen.rs); the benchmark -harness lives in [`src/tpch`](../../src/tpch). +Data is generated deterministically as Parquet via [`tpchgen-cli`](../../src/tpch/tpchgen.rs); +the benchmark harness lives in [`src/tpch`](../../src/tpch). ## CI variants diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 92064c25785..35b1a454351 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -235,26 +235,28 @@ pub async fn convert_parquet_directory_to_vortex( let parquet_path = input_path.join(Format::Parquet.name()); create_dir_all(&vortex_dir).await?; - let parquet_inputs = fs::read_dir(&parquet_path)?.collect::>>()?; + let parquet_inputs = fs::read_dir(&parquet_path)? + .map(|entry| entry.map(|entry| entry.path())) + .filter(|entry| { + entry + .as_ref() + .is_ok_and(|path| path.extension().is_some_and(|e| e == "parquet")) + }) + .collect::>>()?; trace!( "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")); - let concurrency = calculate_concurrency(); - futures::stream::iter(iter) - .map(|dir_entry| { + futures::stream::iter(parquet_inputs) + .map(|parquet_file_path| { let filename = { - let mut temp = dir_entry.path(); + let mut temp = parquet_file_path.clone(); 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}.{}", format.ext())); tokio::spawn( diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index ebe6ff2e96a..e1a51497494 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -18,9 +18,11 @@ use vortex::expr::col; use vortex::expr::pack; use vortex::file::OpenOptionsSessionExt; +use crate::CompactionStrategy; use crate::Format; use crate::IdempotentPath; use crate::SESSION; +use crate::conversions::convert_parquet_directory_to_vortex; use crate::datasets::Dataset; use crate::tpch::tpchgen::TpchGenOptions; use crate::tpch::tpchgen::generate_tpch_tables; @@ -58,13 +60,12 @@ impl Dataset for TPCHLCommentChunked { let scale_factor_dir = base_path.join("1.0"); let data_dir = scale_factor_dir.join(Format::OnDiskVortex.name()); - // Generate TPC-H CSV data if it doesn't exist if !data_dir.exists() { - // Use blocking call like TPC-H benchmark does - let options = TpchGenOptions::new("1.0".to_string(), scale_factor_dir) - .with_format(Format::OnDiskVortex); - - futures::executor::block_on(generate_tpch_tables(options))?; + let options = TpchGenOptions::new("1.0".to_string(), scale_factor_dir.clone()) + .with_format(Format::Parquet); + generate_tpch_tables(options).await?; + convert_parquet_directory_to_vortex(&scale_factor_dir, CompactionStrategy::Default) + .await?; } let mut chunks: Vec = vec![]; diff --git a/vortex-bench/src/tpch/benchmark.rs b/vortex-bench/src/tpch/benchmark.rs index 4dfdaf1207d..cea63a17f11 100644 --- a/vortex-bench/src/tpch/benchmark.rs +++ b/vortex-bench/src/tpch/benchmark.rs @@ -93,8 +93,7 @@ impl Benchmark for TpcHBenchmark { .to_file_path() .map_err(|_| anyhow::anyhow!("Invalid file URL: {}", self.data_url.as_str()))?; - let options = TpchGenOptions::new(self.scale_factor.clone(), base_data_dir) - .with_max_file_size_mb(Some(600)); + let options = TpchGenOptions::new(self.scale_factor.clone(), base_data_dir); tpchgen::generate_tpch_tables(options).await?; diff --git a/vortex-bench/src/tpch/tpchgen.rs b/vortex-bench/src/tpch/tpchgen.rs index ba2646b365e..24c7620c056 100644 --- a/vortex-bench/src/tpch/tpchgen.rs +++ b/vortex-bench/src/tpch/tpchgen.rs @@ -1,63 +1,35 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::env; use std::fs; use std::path::Path; use std::path::PathBuf; -use std::sync::Arc; +use anyhow::Context; use anyhow::Result; -use anyhow::anyhow; -use arrow_array::RecordBatch; -use arrow_schema::SchemaRef; -use futures::FutureExt; -use futures::StreamExt; -use futures::future::BoxFuture; -use parquet::arrow::AsyncArrowWriter; -use parquet::basic::Compression; -use parquet::file::properties::WriterProperties; -use tokio::fs::File as TokioFile; -use tokio::sync::Semaphore; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; -use tokio_stream::wrappers::UnboundedReceiverStream; -use tpchgen::generators::CustomerGenerator; -use tpchgen::generators::LineItemGenerator; -use tpchgen::generators::NationGenerator; -use tpchgen::generators::OrderGenerator; -use tpchgen::generators::PartGenerator; -use tpchgen::generators::PartSuppGenerator; -use tpchgen::generators::RegionGenerator; -use tpchgen::generators::SupplierGenerator; -use tpchgen_arrow::RecordBatchIterator; +use tokio::process::Command; use tracing::info; -use vortex::array::ArrayRef; -use vortex::array::stream::ArrayStreamAdapter; -use vortex::dtype::DType; -use vortex::error::VortexExpect; -use vortex::file::WriteOptionsSessionExt; -use vortex_arrow::FromArrowArray; -use vortex_arrow::FromArrowType; -use crate::CompactionStrategy; use crate::Format; use crate::IdempotentPath; -use crate::SESSION; -use crate::utils::file::idempotent_async; -/// Configuration for TPC-H data generation +/// Python package used to run the TPC-H generator. +/// +/// `tpchgen-rs` v3 depends on Arrow 59 through `tpchgen-arrow`. Running the CLI behind a Parquet +/// boundary keeps that Arrow dependency out of the Vortex Rust workspace. +const TPCHGEN_CLI_PACKAGE: &str = "tpchgen-cli==3.0.0"; +const TPCHGEN_CLI_BIN: &str = "tpchgen-cli"; + +/// Configuration for TPC-H data generation. #[derive(Debug, Clone)] pub struct TpchGenOptions { - /// Scale factor (0.01, 0.1, 1, 10, 100, 1000) + /// Scale factor (0.01, 0.1, 1, 10, 100, 1000). pub scale_factor: String, - /// Output directory + /// Output directory. pub output_dir: PathBuf, - /// Output format + /// Output format. pub format: Format, - /// Batch size for streaming - pub batch_size: usize, - /// The max size of uncompressed file .tbl that we should generate - pub max_file_size_mb: Option, } impl Default for TpchGenOptions { @@ -66,8 +38,6 @@ impl Default for TpchGenOptions { scale_factor: "1.0".to_string(), output_dir: "tpch".to_data_path(), format: Format::Parquet, - batch_size: 8192 * 64, - max_file_size_mb: None, } } } @@ -85,368 +55,150 @@ impl TpchGenOptions { self.format = format; self } - - pub fn with_batch_size(mut self, batch_size: usize) -> Self { - self.batch_size = batch_size; - self - } - - pub fn with_max_file_size_mb(mut self, max_file_size_mb: Option) -> Self { - self.max_file_size_mb = max_file_size_mb; - self - } } -/// Generate all TPC-H tables for a single scale factor +/// Generate all TPC-H tables for a single scale factor. +/// +/// This always generates Parquet files. Other benchmark formats are produced by reading these +/// Parquet files through the workspace Arrow version and converting them in `data-gen`. pub async fn generate_tpch_tables(options: TpchGenOptions) -> Result<()> { + if !matches!(options.format, Format::Parquet | Format::OnDiskDuckDB) { + anyhow::bail!( + "TPC-H generation only creates Parquet base data; use data-gen conversion for {}", + options.format + ); + } + fs::create_dir_all(&options.output_dir)?; + let parquet_dir = options.output_dir.join(Format::Parquet.name()); + fs::create_dir_all(&parquet_dir)?; let tables = [ - ("nation", TableGenerator::Nation), - ("region", TableGenerator::Region), - ("part", TableGenerator::Part), - ("supplier", TableGenerator::Supplier), - ("customer", TableGenerator::Customer), - ("partsupp", TableGenerator::PartSupp), - ("orders", TableGenerator::Orders), - ("lineitem", TableGenerator::LineItem), + "nation", "region", "part", "supplier", "customer", "partsupp", "orders", "lineitem", ]; - const MAX_CONCURRENT_FILES: usize = 32; - - let all_futures = tables - .iter() - .map(|(table_name, generator)| { - info!( - scale_factor = options.scale_factor, - format = %options.format, - table = table_name, - "Generating TPC-H table", - ); - let table_name = table_name.to_string(); - - generate_table_files(table_name, *generator, options.clone()) - }) - .collect::>>()? - .into_iter() - .flatten() - .collect::>>>(); - - let limiter = Arc::new(Semaphore::new(MAX_CONCURRENT_FILES)); - let (tx, rx) = mpsc::unbounded_channel(); - - for f in all_futures { - let limiter = Arc::clone(&limiter); - let tx = tx.clone(); - tokio::task::spawn(async move { - let _guard = limiter.acquire().await?; - tx.send(f.await)?; - - anyhow::Ok(()) - }); - } - - drop(tx); - - let mut rx = UnboundedReceiverStream::new(rx); - - while let Some(r) = rx.next().await { - r?; + for table_name in tables { + info!( + scale_factor = options.scale_factor, + format = %Format::Parquet, + table = table_name, + "Generating TPC-H table", + ); + generate_table_file(table_name, &options, &parquet_dir).await?; } Ok(()) } -#[derive(Debug, Clone, Copy)] -enum TableGenerator { - Nation, - Region, - Part, - Supplier, - Customer, - PartSupp, - Orders, - LineItem, -} - -impl TableGenerator { - // Returns the size in MBs of the table at scale factor 1, if the table is constant size - // return None - fn uncompressed_data_size(&self) -> Option { - match self { - TableGenerator::Nation => None, - TableGenerator::Region => None, - TableGenerator::Part => Some(113), - TableGenerator::Supplier => Some(1), - TableGenerator::Customer => Some(23), - TableGenerator::PartSupp => Some(113), - TableGenerator::Orders => Some(164), - TableGenerator::LineItem => Some(725), - } +/// Generate one Parquet file for a specific table. +async fn generate_table_file( + table_name: &str, + options: &TpchGenOptions, + parquet_dir: &Path, +) -> Result<()> { + let output_file = parquet_dir.join(format!("{table_name}_0.parquet")); + if output_file.exists() { + return Ok(()); } -} - -/// Generate files for a specific table in streaming fashion -fn generate_table_files( - table_name: String, - generator: TableGenerator, - options: TpchGenOptions, -) -> Result>>> { - let write_format = match options.format { - Format::Parquet | Format::OnDiskDuckDB => Format::Parquet, - Format::OnDiskVortex => Format::OnDiskVortex, - Format::VortexCompact => Format::VortexCompact, - f => { - anyhow::bail!("{f} format is not supported by tpchgen"); - } - }; - - let output_dir = options.output_dir.join(write_format.name()); - - fs::create_dir_all(&output_dir)?; - - // Calculate number of partitions without creating expensive iterators - let num_parts = calculate_num_parts(generator, &options)?; - - let mut futures = Vec::new(); - - for partition_idx in 0..num_parts { - let output_file = output_dir.join(format!( - "{table_name}_{partition_idx}.{}", - write_format.ext() - )); - - // Clone necessary data for the async closure - let options_clone = options.clone(); - let table_name = table_name.to_string(); - - let future = async move { - let of = output_file.clone(); - idempotent_async(output_file.to_string_lossy().as_ref(), |path| async move { - info!( - "Generating {table_name} table as {write_format}, at {}", - of.to_string_lossy() - ); - - // Create the specific iterator for this partition only when we need to generate - let iter = create_single_batch_iterator(generator, &options_clone, partition_idx)?; - // Create generator and process batches in streaming fashion - let schema = Arc::clone(iter.schema()); - - // Create writer based on format - let mut writer: Box = match write_format { - Format::Parquet => Box::new(ParquetWriter::new(path, schema).await?), - Format::OnDiskVortex => Box::new(VortexWriter::new( - path, - schema, - CompactionStrategy::Default, - )?), - Format::VortexCompact => Box::new(VortexWriter::new( - path, - schema, - CompactionStrategy::Compact, - )?), - _ => unreachable!(), - }; - - for batch in iter { - writer.write_batch(&batch).await?; - } - - writer.finalize().await?; - Ok::<(), anyhow::Error>(()) - }) - .await?; - Ok(()) - } - .boxed(); - - futures.push(future); + let scratch_dir = parquet_dir.join(format!(".{table_name}-tpchgen-{}", uuid::Uuid::new_v4())); + fs::create_dir_all(&scratch_dir)?; + + let result = async { + generate_table_with_cli(table_name, &options.scale_factor, &scratch_dir).await?; + let generated_file = generated_file_path(table_name, &scratch_dir)?; + fs::rename(&generated_file, &output_file).with_context(|| { + format!( + "failed to move generated TPC-H file from {} to {}", + generated_file.display(), + output_file.display() + ) + }) + } + .await; + + fs::remove_dir_all(&scratch_dir).ok(); + result +} + +async fn generate_table_with_cli( + table_name: &str, + scale_factor: &str, + output_dir: &Path, +) -> Result<()> { + let mut command = tpchgen_command(); + command + .arg("parquet") + .arg("-s") + .arg(scale_factor) + .arg("--tables") + .arg(table_name) + .arg("--output-dir") + .arg(output_dir) + .arg("--no-progress") + .arg("--quiet"); + + let output = command + .output() + .await + .with_context(|| format!("failed to spawn {TPCHGEN_CLI_BIN} via uvx"))?; + + if !output.status.success() { + let status = output.status; + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + anyhow::bail!( + "{TPCHGEN_CLI_BIN} failed while generating TPC-H table {table_name}: {status}\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); } - Ok(futures) + Ok(()) } -/// Calculate the number of partitions without creating expensive iterators -#[expect(clippy::cast_possible_truncation)] -fn calculate_num_parts(generator: TableGenerator, options: &TpchGenOptions) -> Result { - let scale_factor = options.scale_factor.parse::()?; - - let num_parts = if let Some((data_size, max_file_size)) = generator - .uncompressed_data_size() - .zip(options.max_file_size_mb) - { - #[expect(clippy::cast_precision_loss)] - let file_size = (data_size as f64 * scale_factor).ceil() as u64; - file_size.div_ceil(max_file_size) +fn tpchgen_command() -> Command { + if let Some(cli) = env::var_os("TPCHGEN_CLI") { + Command::new(cli) } else { - 1 - }; - - Ok(num_parts as usize) -} - -/// Create a single batch iterator for a specific partition -#[expect(clippy::cast_possible_truncation)] -fn create_single_batch_iterator( - generator: TableGenerator, - options: &TpchGenOptions, - partition_idx: usize, -) -> Result> { - let scale_factor = options.scale_factor.parse::()?; - let num_parts = calculate_num_parts(generator, options)? as i32; - let part = (partition_idx + 1) as i32; // 1-indexed - let batch_size = options.batch_size; - - let iterator: Box = match generator { - TableGenerator::Nation => Box::new( - tpchgen_arrow::NationArrow::new(NationGenerator::new(scale_factor, part, num_parts)) - .with_batch_size(batch_size), - ), - TableGenerator::Region => Box::new( - tpchgen_arrow::RegionArrow::new(RegionGenerator::new(scale_factor, part, num_parts)) - .with_batch_size(batch_size), - ), - TableGenerator::Part => Box::new( - tpchgen_arrow::PartArrow::new(PartGenerator::new(scale_factor, part, num_parts)) - .with_batch_size(batch_size), - ), - TableGenerator::Supplier => Box::new( - tpchgen_arrow::SupplierArrow::new(SupplierGenerator::new( - scale_factor, - part, - num_parts, - )) - .with_batch_size(batch_size), - ), - TableGenerator::Customer => Box::new( - tpchgen_arrow::CustomerArrow::new(CustomerGenerator::new( - scale_factor, - part, - num_parts, - )) - .with_batch_size(batch_size), - ), - TableGenerator::PartSupp => Box::new( - tpchgen_arrow::PartSuppArrow::new(PartSuppGenerator::new( - scale_factor, - part, - num_parts, - )) - .with_batch_size(batch_size), - ), - TableGenerator::Orders => Box::new( - tpchgen_arrow::OrderArrow::new(OrderGenerator::new(scale_factor, part, num_parts)) - .with_batch_size(batch_size), - ), - TableGenerator::LineItem => Box::new( - tpchgen_arrow::LineItemArrow::new(LineItemGenerator::new( - scale_factor, - part, - num_parts, - )) - .with_batch_size(batch_size), - ), - }; - - Ok(iterator) -} - -/// Common interface for file writers -#[async_trait::async_trait] -trait FileWriter { - async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()>; - async fn finalize(self: Box) -> Result<()>; -} - -/// Parquet writer for streaming TPC-H data -struct ParquetWriter { - writer: AsyncArrowWriter, -} - -impl ParquetWriter { - async fn new(path: PathBuf, schema: SchemaRef) -> Result { - let file = TokioFile::create(&path).await?; - let props = WriterProperties::builder() - .set_compression(Compression::SNAPPY) - .set_bloom_filter_enabled(true) - .build(); - let writer = AsyncArrowWriter::try_new(file, schema, Some(props))?; - Ok(Self { writer }) + let mut command = Command::new("uvx"); + command.args(["--from", TPCHGEN_CLI_PACKAGE, TPCHGEN_CLI_BIN]); + command } } -#[async_trait::async_trait] -impl FileWriter for ParquetWriter { - async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { - self.writer.write(batch).await.map_err(|e| anyhow!(e)) +fn generated_file_path(table_name: &str, scratch_dir: &Path) -> Result { + let single_file = scratch_dir.join(format!("{table_name}.parquet")); + if single_file.exists() { + return Ok(single_file); } - async fn finalize(self: Box) -> Result<()> { - self.writer.close().await?; - Ok(()) - } + anyhow::bail!( + "expected {TPCHGEN_CLI_BIN} to write {}", + single_file.display() + ); } -/// Vortex writer for streaming TPC-H data -struct VortexWriter { - sender: Option>>, - write_task: Option>>, -} - -impl VortexWriter { - fn new( - path: PathBuf, - schema: SchemaRef, - compaction_strategy: CompactionStrategy, - ) -> Result { - // Increase buffer size to avoid backpressure issues - let (sender, receiver) = mpsc::channel(2); - let dtype = DType::from_arrow(schema); - let file_path = path; - let write_task = Some(tokio::spawn(async move { - let stream = ArrayStreamAdapter::new(dtype, ReceiverStream::new(receiver)); +#[cfg(test)] +mod tests { + use super::*; - let mut file = TokioFile::create(&file_path).await?; - compaction_strategy - .apply_options(SESSION.write_options()) - .write(&mut file, stream) - .await - .map_err(|e| anyhow!("Vortex write failed: {}", e))?; + #[test] + fn generated_file_path_matches_cli_layout() -> Result<()> { + let scratch_dir = test_scratch_dir()?; + let expected = scratch_dir.join("nation.parquet"); + fs::write(&expected, [])?; - Ok(()) - })); + assert_eq!(generated_file_path("nation", &scratch_dir)?, expected); - Ok(Self { - sender: Some(sender), - write_task, - }) - } -} - -#[async_trait::async_trait] -impl FileWriter for VortexWriter { - async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { - let array = ArrayRef::from_arrow(batch, false)?; - self.sender - .as_ref() - .vortex_expect("sender closed early") - .send(Ok(array)) - .await - .map_err(|_| anyhow!("Failed to send array to write task")) + fs::remove_dir_all(scratch_dir)?; + Ok(()) } - async fn finalize(mut self: Box) -> Result<()> { - // Close the sender to signal end of stream - self.sender.take(); - - // Wait for write task to complete - if let Some(task) = self.write_task.take() { - task.await - .map_err(|e| anyhow!("Write task failed: {}", e))??; - } - - Ok(()) + fn test_scratch_dir() -> Result { + let path = std::env::temp_dir().join(format!( + "vortex-bench-tpchgen-test-{}", + uuid::Uuid::new_v4() + )); + fs::create_dir_all(&path)?; + Ok(path) } } diff --git a/vortex-sqllogictest/slt/tpch/generate_data.sh b/vortex-sqllogictest/slt/tpch/generate_data.sh index 42b40d22a05..a0b1a20f2aa 100755 --- a/vortex-sqllogictest/slt/tpch/generate_data.sh +++ b/vortex-sqllogictest/slt/tpch/generate_data.sh @@ -17,7 +17,7 @@ mkdir -p "${DATA_DIR}" # 1. Generate TPC-H data at scale factor 0.1 echo "Generating TPC-H data (SF=0.1)..." -uvx tpchgen-cli -s 0.1 --format=parquet --output-dir "${DATA_DIR}/" +uvx --from tpchgen-cli==3.0.0 tpchgen-cli parquet -s 0.1 --output-dir "${DATA_DIR}/" --no-progress # 2. Convert each parquet file to Vortex format for f in "${DATA_DIR}"/*.parquet; do diff --git a/vortex-test/compat-gen/Cargo.toml b/vortex-test/compat-gen/Cargo.toml index 4a62aca3671..c03a1fda047 100644 --- a/vortex-test/compat-gen/Cargo.toml +++ b/vortex-test/compat-gen/Cargo.toml @@ -29,10 +29,7 @@ vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-session = { workspace = true } -# TPC-H generation arrow-array = { workspace = true } -tpchgen = { workspace = true } -tpchgen-arrow = { workspace = true } # ClickBench parquet reading + writing arrow-select = { workspace = true } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs index 59e97a7e587..23f4a6e66ac 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/tpch.rs @@ -1,22 +1,117 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::env; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::process::Command; + use arrow_array::RecordBatch; -use tpchgen::generators::LineItemGenerator; -use tpchgen::generators::OrderGenerator; -use tpchgen_arrow::RecordBatchIterator; +use bytes::Bytes; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_arrow::FromArrowArray; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::fixtures::DatasetFixture; -const SCALE_FACTOR: f64 = 0.01; +const SCALE_FACTOR: &str = "0.01"; +const BATCH_SIZE: usize = 65_536; +const TPCHGEN_CLI_PACKAGE: &str = "tpchgen-cli==3.0.0"; +const TPCHGEN_CLI_BIN: &str = "tpchgen-cli"; + +fn cached_tpch_parquet(table: &str) -> VortexResult { + let crate_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let data_dir = crate_dir.join("data").join("tpch"); + let dest = data_dir.join(format!("{table}.parquet")); + + if dest.exists() { + return Ok(dest); + } + + fs::create_dir_all(&data_dir).map_err(|e| vortex_err!("failed to create data dir: {e}"))?; + + let scratch = tempfile::Builder::new() + .prefix(&format!("{table}-")) + .tempdir_in(&data_dir) + .map_err(|e| vortex_err!("failed to create TPC-H scratch dir: {e}"))?; + + generate_tpch_parquet(table, scratch.path())?; + + let generated = scratch.path().join(format!("{table}.parquet")); + if !generated.exists() { + return Err(vortex_err!( + "expected {TPCHGEN_CLI_BIN} to write {}", + generated.display() + )); + } + + fs::rename(&generated, &dest).map_err(|e| { + vortex_err!( + "failed to move generated TPC-H parquet from {} to {}: {e}", + generated.display(), + dest.display() + ) + })?; + + Ok(dest) +} + +fn generate_tpch_parquet(table: &str, output_dir: &Path) -> VortexResult<()> { + let output = tpchgen_command() + .arg("parquet") + .arg("-s") + .arg(SCALE_FACTOR) + .arg("--tables") + .arg(table) + .arg("--output-dir") + .arg(output_dir) + .arg("--no-progress") + .arg("--quiet") + .output() + .map_err(|e| vortex_err!("failed to spawn {TPCHGEN_CLI_BIN} via uvx: {e}"))?; + + if !output.status.success() { + let status = output.status; + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(vortex_err!( + "{TPCHGEN_CLI_BIN} failed while generating TPC-H table {table}: {status}\nstdout:\n{stdout}\nstderr:\n{stderr}" + )); + } + + Ok(()) +} + +fn tpchgen_command() -> Command { + if let Some(cli) = env::var_os("TPCHGEN_CLI") { + Command::new(cli) + } else { + let mut command = Command::new("uvx"); + command.args(["--from", TPCHGEN_CLI_PACKAGE, TPCHGEN_CLI_BIN]); + command + } +} + +fn collect_parquet_as_vortex(path: PathBuf) -> VortexResult { + let file_bytes = fs::read(&path) + .map_err(|e| vortex_err!("failed to read cached parquet at {}: {e}", path.display()))?; + let bytes = Bytes::from(file_bytes); + + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes) + .map_err(|e| vortex_err!("failed to open parquet: {e}"))? + .with_batch_size(BATCH_SIZE) + .build() + .map_err(|e| vortex_err!("failed to build parquet reader: {e}"))?; + + let batches: Vec = reader + .collect::, _>>() + .map_err(|e| vortex_err!("failed to read parquet batches: {e}"))?; -fn collect_batches_as_vortex(iter: impl RecordBatchIterator) -> VortexResult { - let batches: Vec = iter.collect(); Ok(ChunkedArray::from_iter( batches .into_iter() @@ -38,9 +133,7 @@ impl DatasetFixture for TpchLineitemFixture { } fn build(&self) -> VortexResult { - let generator = LineItemGenerator::new(SCALE_FACTOR, 1, 1); - let arrow_iter = tpchgen_arrow::LineItemArrow::new(generator).with_batch_size(65_536); - collect_batches_as_vortex(arrow_iter) + collect_parquet_as_vortex(cached_tpch_parquet("lineitem")?) } } @@ -56,9 +149,7 @@ impl DatasetFixture for TpchOrdersFixture { } fn build(&self) -> VortexResult { - let generator = OrderGenerator::new(SCALE_FACTOR, 1, 1); - let arrow_iter = tpchgen_arrow::OrderArrow::new(generator).with_batch_size(65_536); - collect_batches_as_vortex(arrow_iter) + collect_parquet_as_vortex(cached_tpch_parquet("orders")?) } } From 8ab1a94aec860d94c830c42c794da95473ba6454 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 28 Jul 2026 17:20:46 +0100 Subject: [PATCH 2/5] fmt cargo.toml Signed-off-by: Adam Gutglick --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 6a2d306b631..1e3bc072926 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -266,12 +266,12 @@ thiserror = "2.0.3" tokio = { version = "1.52" } tokio-stream = "0.1.17" tokio-util = "0.7.17" -vortex-array-macros = { version = "0.1.0", path = "./vortex-array-macros" } tracing = { version = "0.1.41", default-features = false } tracing-perfetto = "0.1.5" tracing-subscriber = "0.3" url = "2.5.7" uuid = { version = "1.23", features = ["js"] } +vortex-array-macros = { version = "0.1.0", path = "./vortex-array-macros" } wasm-bindgen-futures = "0.4.54" wkb = "0.9.2" xshell = "0.2.6" From 45bfc69c7f3516d6ab8e531deea54da9e4f302e3 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 28 Jul 2026 17:26:54 +0100 Subject: [PATCH 3/5] fix Signed-off-by: Adam Gutglick --- vortex-bench/src/tpch/tpchgen.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-bench/src/tpch/tpchgen.rs b/vortex-bench/src/tpch/tpchgen.rs index 24c7620c056..896cf433b47 100644 --- a/vortex-bench/src/tpch/tpchgen.rs +++ b/vortex-bench/src/tpch/tpchgen.rs @@ -194,7 +194,7 @@ mod tests { } fn test_scratch_dir() -> Result { - let path = std::env::temp_dir().join(format!( + let path = env::temp_dir().join(format!( "vortex-bench-tpchgen-test-{}", uuid::Uuid::new_v4() )); From bd136c2d7d0e26ecf34293c02cc0b55bf46885a8 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 28 Jul 2026 18:44:15 +0100 Subject: [PATCH 4/5] fix Signed-off-by: Adam Gutglick --- .github/workflows/ci.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5fa6a1714a1..04099c19ca1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -407,6 +407,8 @@ jobs: - uses: ./.github/actions/setup-prebuild with: enable-sccache: "true" + - name: Install uv + uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 - name: Rust Tests (Windows) run: | cargo nextest run --cargo-profile ci --locked --workspace --all-features --no-fail-fast ` From cd229bea685fd28048a4a6ff53b1f70bf65632c4 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 28 Jul 2026 18:56:06 +0100 Subject: [PATCH 5/5] fix Signed-off-by: Adam Gutglick --- .github/workflows/musl.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/musl.yml b/.github/workflows/musl.yml index c331dbc2859..cfee6f93444 100644 --- a/.github/workflows/musl.yml +++ b/.github/workflows/musl.yml @@ -54,6 +54,9 @@ jobs: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + - name: Install uv + uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 + - name: Install nextest shell: bash # Prebuilt static musl nextest binary; building it from source would