From 4e792542ac7df4638d1dd5b1598c6a969f5b3b13 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 14:07:40 +0000 Subject: [PATCH] perf(bench): generate TPC-H Parquet in parallel parts on a shared work queue TPC-H and TPC-DS Parquet generation ran one thread per table, so the largest table (lineitem, store_sales) set the time while most cores sat idle. Add write_parquet_tables: every part of every table is a unit on one work queue served by one worker per core, largest tables first. Each part is generated and encoded independently into its own file, and the worker that finishes a table's last part copies the parts' row groups into the final file in order, without re-encoding, keeping the embedded Arrow schema and page indexes. TPC-H tables are split with tpchgen's part/num_parts, about 4M rows per part and at most one part per core. TPC-DS tables run on the same queue as single units, because the released tpcdsgen-arrow cannot yet start generating partway through a table. TPC-H SF=10 on 4 cores: 80s -> 30s, with every table's rows, values and Arrow schema identical to before; lineitem has 60 row groups instead of 58. TPC-DS SF=10 output is unchanged (71s either way). Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01M2QNZPywc2ed6MkXoqjBdV --- vortex-bench/src/tpcds/tpcdsgen.rs | 43 ++--- vortex-bench/src/tpch/tpchgen.rs | 180 +++++++++++++++---- vortex-bench/src/utils/parquet.rs | 273 +++++++++++++++++++++++++++++ 3 files changed, 443 insertions(+), 53 deletions(-) diff --git a/vortex-bench/src/tpcds/tpcdsgen.rs b/vortex-bench/src/tpcds/tpcdsgen.rs index 9b7512088f1..5053036361d 100644 --- a/vortex-bench/src/tpcds/tpcdsgen.rs +++ b/vortex-bench/src/tpcds/tpcdsgen.rs @@ -17,8 +17,6 @@ use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Schema; use arrow_schema::SchemaRef; -use futures::StreamExt; -use futures::stream; use parquet::arrow::ArrowWriter; use tpcdsgen::config::CompatMode; use tpcdsgen::config::Session; @@ -52,8 +50,9 @@ use tracing::info; use vortex::utils::parallelism::get_available_parallelism; use crate::Format; -use crate::utils::file::idempotent; +use crate::utils::parquet::ParquetTableJob; use crate::utils::parquet::parquet_writer_properties; +use crate::utils::parquet::write_parquet_tables; const BATCH_SIZE: usize = 8192 * 16; @@ -63,8 +62,11 @@ const ROW_GROUP_SIZE: usize = 122_880; /// Generate every TPC-DS table as a single Parquet file per table under `base_dir/parquet`. /// /// The files use the column names and types of the TPC-DS specification (and of DuckDB's -/// `dsdgen`). Each table is generated on its own blocking thread. Already generated tables are -/// skipped. +/// `dsdgen`). Tables are generated on a shared pool of one worker per core. Already generated +/// tables are skipped. +/// +/// Each table is one unit of work: the released `tpcdsgen-arrow` always starts generating at a +/// table's first row, so a table cannot yet be split into parts. pub async fn generate_tpcds(base_dir: PathBuf, scale_factor: String) -> Result { let output_dir = base_dir.join(Format::Parquet.name()); fs::create_dir_all(&output_dir)?; @@ -77,23 +79,22 @@ pub async fn generate_tpcds(base_dir: PathBuf, scale_factor: String) -> Result

Result<()> { ("lineitem", TableGenerator::LineItem), ]; + if matches!(options.format, Format::Parquet | Format::OnDiskDuckDB) { + return generate_parquet_tables(&tables, options).await; + } + const MAX_CONCURRENT_FILES: usize = 32; let all_futures = tables @@ -149,6 +156,41 @@ pub async fn generate_tpch_tables(options: TpchGenOptions) -> Result<()> { Ok(()) } +/// Generate every table as Parquet, with all tables' parts sharing one pool of one worker per core. +async fn generate_parquet_tables( + tables: &[(&str, TableGenerator)], + options: TpchGenOptions, +) -> Result<()> { + let output_dir = options.output_dir.join(Format::Parquet.name()); + let scale_factor: f64 = options.scale_factor.parse()?; + let options = Arc::new(options); + let jobs = tables + .iter() + .map(|&(table_name, generator)| { + let num_parts = generator.num_parquet_parts(scale_factor); + info!( + scale_factor = options.scale_factor, + table = table_name, + num_parts, + "Generating TPC-H table as parquet", + ); + let options = Arc::clone(&options); + ParquetTableJob { + path: output_dir.join(format!("{table_name}.{}", Format::Parquet.ext())), + num_parts, + write_part: Box::new(move |part, path| { + write_parquet_part( + path, + create_batch_iterator(generator, &options, part + 1, num_parts)?, + ) + }), + } + }) + .collect(); + let workers = get_available_parallelism().unwrap_or(1); + tokio::task::spawn_blocking(move || write_parquet_tables(jobs, workers)).await? +} + #[derive(Debug, Clone, Copy)] enum TableGenerator { Nation, @@ -161,6 +203,32 @@ enum TableGenerator { LineItem, } +/// Rows per Parquet part. Each part is generated and encoded on its own thread. +const ROWS_PER_PART: f64 = (4 * ROW_GROUP_SIZE) as f64; + +impl TableGenerator { + /// Approximate row count at scale factor 1, from the TPC-H specification. + fn rows_per_scale_factor(self) -> f64 { + match self { + Self::Nation | Self::Region => 0.0, + Self::Part => 200_000.0, + Self::Supplier => 10_000.0, + Self::Customer => 150_000.0, + Self::PartSupp => 800_000.0, + Self::Orders => 1_500_000.0, + Self::LineItem => 6_000_000.0, + } + } + + /// Number of parts to generate this table's Parquet file in, at most one per core. + // `as` saturates on overflow, NaN and negatives, and the clamp bounds the result. + #[expect(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + fn num_parquet_parts(self, scale_factor: f64) -> usize { + let parts = (self.rows_per_scale_factor() * scale_factor / ROWS_PER_PART).ceil(); + (parts as usize).clamp(1, get_available_parallelism().unwrap_or(1)) + } +} + /// Generate one file for a specific table in streaming fashion fn generate_table_file( table_name: String, @@ -168,7 +236,6 @@ fn generate_table_file( options: TpchGenOptions, ) -> Result>> { let write_format = match options.format { - Format::Parquet | Format::OnDiskDuckDB => Format::Parquet, Format::OnDiskVortex => Format::OnDiskVortex, Format::VortexCompact => Format::VortexCompact, f => { @@ -190,14 +257,13 @@ fn generate_table_file( of.to_string_lossy() ); - let iter = create_batch_iterator(generator, &options)?; + let iter = create_batch_iterator(generator, &options, 1, 1)?; // 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, @@ -226,14 +292,16 @@ fn generate_table_file( Ok(future) } -/// Create a batch iterator for one full table +/// Create a batch iterator for part `part` (one-based) of `num_parts` of a table. fn create_batch_iterator( generator: TableGenerator, options: &TpchGenOptions, + part: usize, + num_parts: usize, ) -> Result> { let scale_factor = options.scale_factor.parse::()?; - let part = 1; - let num_parts = 1; + let part = i32::try_from(part)?; + let num_parts = i32::try_from(num_parts)?; let batch_size = options.batch_size; let iterator: Box = match generator { @@ -297,32 +365,21 @@ trait FileWriter { 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 = parquet_writer_properties() - .set_max_row_group_row_count(Some(ROW_GROUP_SIZE)) - .build(); - let writer = AsyncArrowWriter::try_new(file, schema, Some(props))?; - Ok(Self { writer }) - } -} - -#[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)) - } - - async fn finalize(self: Box) -> Result<()> { - self.writer.close().await?; - Ok(()) +/// Write one TPC-H table part as a standalone Parquet file. +fn write_parquet_part(path: &Path, batches: Box) -> Result<()> { + let props = parquet_writer_properties() + .set_max_row_group_row_count(Some(ROW_GROUP_SIZE)) + .build(); + let mut writer = ArrowWriter::try_new( + fs::File::create(path)?, + Arc::clone(batches.schema()), + Some(props), + )?; + for batch in batches { + writer.write(&batch)?; } + writer.close()?; + Ok(()) } /// Vortex writer for streaming TPC-H data @@ -389,3 +446,62 @@ impl FileWriter for VortexWriter { Ok(()) } } + +#[cfg(test)] +mod tests { + use arrow_select::concat::concat_batches; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use rstest::rstest; + + use super::*; + + fn read_parquet(path: &Path) -> Result<(SchemaRef, usize, RecordBatch)> { + let builder = ParquetRecordBatchReaderBuilder::try_new(fs::File::open(path)?)?; + let schema = Arc::clone(builder.schema()); + let row_groups = builder.metadata().num_row_groups(); + let batches = builder.build()?.collect::, _>>()?; + Ok(( + Arc::clone(&schema), + row_groups, + concat_batches(&schema, &batches)?, + )) + } + + #[rstest] + #[case::customer(TableGenerator::Customer)] + #[case::partsupp(TableGenerator::PartSupp)] + #[case::orders(TableGenerator::Orders)] + #[case::lineitem(TableGenerator::LineItem)] + fn parquet_parts_match_single_writer(#[case] generator: TableGenerator) -> Result<()> { + let dir = tempfile::tempdir()?; + let options = TpchGenOptions::new("0.1".to_string(), dir.path()).with_batch_size(10_000); + + let single = dir.path().join("single.parquet"); + write_parquet_part(&single, create_batch_iterator(generator, &options, 1, 1)?)?; + + let num_parts = 7; + let parts = dir.path().join("parts.parquet"); + let part_options = options; + let job = ParquetTableJob { + path: parts.clone(), + num_parts, + write_part: Box::new(move |part, part_path| { + write_parquet_part( + part_path, + create_batch_iterator(generator, &part_options, part + 1, num_parts)?, + ) + }), + }; + write_parquet_tables(vec![job], 3)?; + + let (single_schema, _, single_rows) = read_parquet(&single)?; + let (parts_schema, parts_row_groups, parts_rows) = read_parquet(&parts)?; + // The schema comparison includes the embedded Arrow schema metadata. + assert_eq!(parts_schema, single_schema); + assert_eq!(parts_rows, single_rows); + assert_eq!(parts_row_groups, num_parts); + // Only the two output files remain: the part files are cleaned up. + assert_eq!(fs::read_dir(dir.path())?.count(), 2); + Ok(()) + } +} diff --git a/vortex-bench/src/utils/parquet.rs b/vortex-bench/src/utils/parquet.rs index 3fa62e42e26..f3364f148f4 100644 --- a/vortex-bench/src/utils/parquet.rs +++ b/vortex-bench/src/utils/parquet.rs @@ -3,10 +3,25 @@ //! Shared Parquet writer settings for generated benchmark data. +use std::fs; +use std::fs::File; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use anyhow::anyhow; +use parking_lot::Mutex; use parquet::basic::Compression; use parquet::basic::ZstdLevel; +use parquet::column::writer::ColumnCloseResult; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::metadata::ParquetMetaDataReader; use parquet::file::properties::WriterProperties; use parquet::file::properties::WriterPropertiesBuilder; +use parquet::file::writer::SerializedFileWriter; /// Columns whose dictionary outgrows this fall back to plain encoding. The parquet-rs default /// (1 MiB) keeps dictionaries for unique keys, which compress far worse than plain values. A small @@ -27,3 +42,261 @@ pub fn parquet_writer_properties() -> WriterPropertiesBuilder { // zstd too little context to compress repeating columns. .set_data_page_row_count_limit(usize::MAX) } + +/// Writes part `index` (zero-based) of a table as a standalone Parquet file at the given path. +pub type WriteParquetPart = Box anyhow::Result<()> + Send + Sync>; + +/// One Parquet file to generate, split into parts that can be written independently. +pub struct ParquetTableJob { + /// The final file. Jobs whose file already exists are skipped. + pub path: PathBuf, + /// Number of parts. Concatenated in index order, they make up the whole table. + pub num_parts: usize, + pub write_part: WriteParquetPart, +} + +/// Generate Parquet tables on one shared pool of `workers` threads. +/// +/// Every part of every table is a unit on a single work queue, so one large table spreads over all +/// workers and small tables fill in around it. Tables with the most parts are queued first. The +/// worker that finishes a table's last part copies the parts' row groups into the final file, in +/// part order and without re-encoding. Each part ends with its own final row group, so a table has +/// up to `num_parts - 1` more row groups than a single writer would produce. +pub fn write_parquet_tables(mut jobs: Vec, workers: usize) -> anyhow::Result<()> { + jobs.retain(|job| !job.path.exists()); + jobs.sort_by_key(|job| std::cmp::Reverse(job.num_parts)); + for job in &jobs { + if let Some(parent) = job.path.parent() { + fs::create_dir_all(parent)?; + } + } + + let queue: Vec<(usize, usize)> = jobs + .iter() + .enumerate() + .flat_map(|(job_idx, job)| (0..job.num_parts.max(1)).map(move |part| (job_idx, part))) + .collect(); + let remaining: Vec = jobs + .iter() + .map(|job| AtomicUsize::new(job.num_parts.max(1))) + .collect(); + let next = AtomicUsize::new(0); + let failed = AtomicBool::new(false); + let first_error = Mutex::new(None); + + let run_unit = |job_idx: usize, part: usize| -> anyhow::Result<()> { + let job = &jobs[job_idx]; + (job.write_part)(part, &part_path(&job.path, part))?; + if remaining[job_idx].fetch_sub(1, Ordering::AcqRel) == 1 { + finish_table(job)?; + } + Ok(()) + }; + + std::thread::scope(|scope| { + for _ in 0..workers.clamp(1, queue.len().max(1)) { + scope.spawn(|| { + while !failed.load(Ordering::Acquire) { + let Some(&(job_idx, part)) = queue.get(next.fetch_add(1, Ordering::AcqRel)) + else { + break; + }; + if let Err(err) = run_unit(job_idx, part) { + failed.store(true, Ordering::Release); + first_error.lock().get_or_insert(err); + } + } + }); + } + }); + + match first_error.into_inner() { + Some(err) => { + // Best effort: leave no part files behind for a later run to trip over. + for job in &jobs { + for part in 0..job.num_parts.max(1) { + drop(fs::remove_file(part_path(&job.path, part))); + } + } + Err(err) + } + None => Ok(()), + } +} + +fn part_path(path: &Path, part: usize) -> PathBuf { + PathBuf::from(format!("{}.part{part}", path.display())) +} + +/// Combine a finished table's parts into its final file, then remove the parts. +fn finish_table(job: &ParquetTableJob) -> anyhow::Result<()> { + let parts: Vec = (0..job.num_parts.max(1)) + .map(|part| part_path(&job.path, part)) + .collect(); + if let [single] = parts.as_slice() { + fs::rename(single, &job.path)?; + return Ok(()); + } + // Write next to the destination and rename, so an interrupted run never leaves a partial file + // that a later run would skip. + let tmp = PathBuf::from(format!("{}.tmp", job.path.display())); + concat_parquet_files(&parts, &tmp)?; + fs::rename(&tmp, &job.path)?; + for part in &parts { + fs::remove_file(part)?; + } + Ok(()) +} + +/// Concatenate Parquet files with identical schemas by copying their encoded column chunks. +/// +/// The first input's key-value metadata, including the embedded Arrow schema, is carried over. +fn concat_parquet_files(inputs: &[PathBuf], output: &Path) -> anyhow::Result<()> { + let inputs = inputs + .iter() + .map(|path| { + let file = File::open(path)?; + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file)?; + Ok((file, metadata)) + }) + .collect::>>()?; + let first = &inputs + .first() + .ok_or_else(|| anyhow!("no Parquet files to concatenate"))? + .1; + for (_, metadata) in &inputs[1..] { + anyhow::ensure!( + metadata.file_metadata().schema() == first.file_metadata().schema(), + "Parquet parts have different schemas" + ); + } + + let props = WriterProperties::builder() + .set_key_value_metadata(first.file_metadata().key_value_metadata().cloned()) + .build(); + let mut writer = SerializedFileWriter::new( + File::create(output)?, + first.file_metadata().schema_descr().root_schema_ptr(), + Arc::new(props), + )?; + + for (input, metadata) in &inputs { + for (rg_idx, row_group) in metadata.row_groups().iter().enumerate() { + let column_indexes = metadata.column_index().and_then(|ci| ci.get(rg_idx)); + let offset_indexes = metadata.offset_index().and_then(|oi| oi.get(rg_idx)); + let mut rg_writer = writer.next_row_group()?; + for (col_idx, column) in row_group.columns().iter().enumerate() { + let close = ColumnCloseResult { + bytes_written: u64::try_from(column.compressed_size())?, + rows_written: u64::try_from(row_group.num_rows())?, + metadata: column.clone(), + bloom_filter: None, + column_index: column_indexes.and_then(|row| row.get(col_idx)).cloned(), + offset_index: offset_indexes.and_then(|row| row.get(col_idx)).cloned(), + }; + rg_writer.append_column(input, close)?; + } + rg_writer.close()?; + } + } + writer.close()?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use arrow_array::Int64Array; + use arrow_array::RecordBatch; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + use super::*; + + const ROWS_PER_PART: i64 = 1000; + + /// A job whose part `i` holds the values `i * ROWS_PER_PART..(i + 1) * ROWS_PER_PART`. + fn sequence_job(path: PathBuf, num_parts: usize) -> ParquetTableJob { + ParquetTableJob { + path, + num_parts, + write_part: Box::new(|part, path| { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let start = i64::try_from(part)? * ROWS_PER_PART; + let values = Int64Array::from_iter_values(start..start + ROWS_PER_PART); + let mut writer = + ArrowWriter::try_new(File::create(path)?, Arc::clone(&schema), None)?; + writer.write(&RecordBatch::try_new(schema, vec![Arc::new(values)])?)?; + writer.close()?; + Ok(()) + }), + } + } + + fn read_values(path: &Path) -> anyhow::Result> { + let mut values = Vec::new(); + for batch in ParquetRecordBatchReaderBuilder::try_new(File::open(path)?)?.build()? { + let batch = batch?; + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("expected an Int64 column"))?; + values.extend(column.values().iter().copied()); + } + Ok(values) + } + + #[test] + fn tables_share_the_queue_and_keep_part_order() -> anyhow::Result<()> { + let dir = tempfile::tempdir()?; + let existing = dir.path().join("existing.parquet"); + fs::write(&existing, b"not regenerated")?; + let jobs = vec![ + sequence_job(dir.path().join("one.parquet"), 1), + sequence_job(dir.path().join("many.parquet"), 9), + sequence_job(existing.clone(), 4), + ]; + + write_parquet_tables(jobs, 4)?; + + assert_eq!( + read_values(&dir.path().join("one.parquet"))?, + (0..ROWS_PER_PART).collect::>() + ); + assert_eq!( + read_values(&dir.path().join("many.parquet"))?, + (0..9 * ROWS_PER_PART).collect::>() + ); + assert_eq!(fs::read(&existing)?, b"not regenerated"); + // Part and temporary files are gone. + assert_eq!(fs::read_dir(dir.path())?.count(), 3); + Ok(()) + } + + #[test] + fn a_failed_part_fails_the_run() -> anyhow::Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("broken.parquet"); + let job = ParquetTableJob { + path: path.clone(), + num_parts: 3, + write_part: Box::new(|part, _| { + anyhow::ensure!(part != 1, "part 1 failed"); + Ok(()) + }), + }; + + let err = write_parquet_tables(vec![job], 2) + .err() + .ok_or_else(|| anyhow!("expected an error"))?; + assert!(err.to_string().contains("part 1 failed")); + assert!(!path.exists()); + assert_eq!(fs::read_dir(dir.path())?.count(), 0); + Ok(()) + } +}