Skip to content

perf(bench): generate TPC-H Parquet in parallel parts on a shared work queue #10557

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
joseph-isaacs wants to merge 1 commit into
base: develop
Choose a base branch
from
Draft

perf(bench): generate TPC-H Parquet in parallel parts on a shared work queue #10557

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 22 additions & 21 deletions vortex-bench/src/tpcds/tpcdsgen.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand All @@ -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<PathBuf> {
let output_dir = base_dir.join(Format::Parquet.name());
fs::create_dir_all(&output_dir)?;
Expand All @@ -77,23 +79,22 @@ pub async fn generate_tpcds(base_dir: PathBuf, scale_factor: String) -> Result<P
.with_compat_mode(CompatMode::C)
.build()?;

let parallelism = get_available_parallelism().unwrap_or(1);
let mut tasks = stream::iter(TABLES.map(|table| {
let session = session.clone();
let path = output_dir.join(format!("{table}.{}", Format::Parquet.ext()));
tokio::task::spawn_blocking(move || {
idempotent(&path, |tmp| {
info!(table, scale_factor, "Generating TPC-DS table");
write_parquet(tmp, table_batches(table, session))
})
let jobs = TABLES
.map(|table| {
let session = session.clone();
ParquetTableJob {
path: output_dir.join(format!("{table}.{}", Format::Parquet.ext())),
num_parts: 1,
write_part: Box::new(move |_, path| {
info!(table, scale_factor, "Generating TPC-DS table");
write_parquet(path, table_batches(table, session.clone()))
}),
}
})
}))
.buffer_unordered(parallelism);

while let Some(result) = tasks.next().await {
result??;
}
drop(tasks);
.into_iter()
.collect();
let workers = get_available_parallelism().unwrap_or(1);
tokio::task::spawn_blocking(move || write_parquet_tables(jobs, workers)).await??;

Ok(output_dir)
}
Expand Down
180 changes: 148 additions & 32 deletions vortex-bench/src/tpch/tpchgen.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ use arrow_schema::SchemaRef;
use futures::FutureExt;
use futures::StreamExt;
use futures::future::BoxFuture;
use parquet::arrow::AsyncArrowWriter;
use parquet::arrow::ArrowWriter;
use tokio::fs::File as TokioFile;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
Expand All @@ -33,14 +33,17 @@ use vortex::array::ArrayRef;
use vortex::array::stream::ArrayStreamAdapter;
use vortex::error::VortexExpect;
use vortex::file::WriteOptionsSessionExt;
use vortex::utils::parallelism::get_available_parallelism;
use vortex_arrow::ArrowSessionExt;

use crate::CompactionStrategy;
use crate::Format;
use crate::IdempotentPath;
use crate::SESSION;
use crate::utils::file::idempotent_async;
use crate::utils::parquet::ParquetTableJob;
use crate::utils::parquet::parquet_writer_properties;
use crate::utils::parquet::write_parquet_tables;

/// The parquet-rs default, kept explicit so the shared writer settings do not change the
/// parallelism of existing TPC-H Parquet benchmarks.
Expand Down Expand Up @@ -105,6 +108,10 @@ pub async fn generate_tpch_tables(options: TpchGenOptions) -> 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
Expand Down Expand Up @@ -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,
Expand All @@ -161,14 +203,39 @@ 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,
generator: TableGenerator,
options: TpchGenOptions,
) -> Result<BoxFuture<'static, Result<()>>> {
let write_format = match options.format {
Format::Parquet | Format::OnDiskDuckDB => Format::Parquet,
Format::OnDiskVortex => Format::OnDiskVortex,
Format::VortexCompact => Format::VortexCompact,
f => {
Expand All @@ -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<dyn FileWriter + Send> = match write_format {
Format::Parquet => Box::new(ParquetWriter::new(path, schema).await?),
Format::OnDiskVortex => Box::new(VortexWriter::new(
path,
schema,
Expand Down Expand Up @@ -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<Box<dyn RecordBatchIterator>> {
let scale_factor = options.scale_factor.parse::<f64>()?;
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<dyn RecordBatchIterator> = match generator {
Expand Down Expand Up @@ -297,32 +365,21 @@ trait FileWriter {
async fn finalize(self: Box<Self>) -> Result<()>;
}

/// Parquet writer for streaming TPC-H data
struct ParquetWriter {
writer: AsyncArrowWriter<TokioFile>,
}

impl ParquetWriter {
async fn new(path: PathBuf, schema: SchemaRef) -> Result<Self> {
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<Self>) -> 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<dyn RecordBatchIterator>) -> 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
Expand Down Expand Up @@ -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::<Result<Vec<_>, _>>()?;
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(())
}
}
Loading
Loading