diff --git a/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py b/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py index d3af5100a1d..d0aefcc580f 100644 --- a/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py +++ b/bench-orchestrator/bench_orchestrator/ci_matrix/catalog.py @@ -9,10 +9,10 @@ from dataclasses import replace -from ..config import Benchmark, Format +from ..config import Benchmark, Engine, Format from .model import BenchmarkCase, Catalog, Coverage, Storage from .render import merge_coverage -from .targets import df, duck +from .targets import df, duck, targets PRESETS = { "develop": "Every regular SQL benchmark at full target coverage.", @@ -78,6 +78,31 @@ # Concrete benchmark cases +# TPC-DS SF=100 runs one engine/format target per job, each on its own machine, so a single slow +# target cannot push the others past the job timeout and lose their results. +TPCDS_SF100_TARGETS = ( + (Engine.DATAFUSION, Format.PARQUET, "DataFusion", "Parquet"), + (Engine.DATAFUSION, Format.VORTEX, "DataFusion", "Vortex"), + (Engine.DUCKDB, Format.PARQUET, "DuckDB", "Parquet"), + (Engine.DUCKDB, Format.VORTEX, "DuckDB", "Vortex"), +) +TPCDS_SF100_CASES = tuple( + BenchmarkCase( + id=f"tpcds-{storage.value}-100-{engine.value}-{fmt.value}", + benchmark=Benchmark.TPCDS, + name=f"TPC-DS SF=100 on {storage.value.upper()} ({engine_name}, {fmt_name})", + storage=storage, + scale_factor=100.0, + iterations=3, + local_dir="vortex-bench/data/tpcds/100.0" if storage is Storage.S3 else None, + # A distinct key per job: each job clears and re-uploads its own copy of the data. + remote_key=f"tpcds/100.0-{engine.value}-{fmt.value}" if storage is Storage.S3 else None, + runs={"nightly": Coverage(targets(engine, fmt))}, + ) + for storage in (Storage.NVME, Storage.S3) + for engine, fmt, engine_name, fmt_name in TPCDS_SF100_TARGETS +) + BENCHMARKS = ( BenchmarkCase( id="clickbench-nvme", @@ -193,15 +218,6 @@ scale_factor=10.0, runs={"nightly": DEFAULT}, ), - BenchmarkCase( - id="tpcds-nvme-100", - benchmark=Benchmark.TPCDS, - name="TPC-DS SF=100 on NVME", - scale_factor=100.0, - # 99 queries across four targets: fewer passes keep the job inside its timeout. - iterations=3, - runs={"nightly": DEFAULT}, - ), BenchmarkCase( id="tpcds-s3-10", benchmark=Benchmark.TPCDS, @@ -212,17 +228,7 @@ remote_key="tpcds/10.0", runs={"nightly": DEFAULT}, ), - BenchmarkCase( - id="tpcds-s3-100", - benchmark=Benchmark.TPCDS, - name="TPC-DS SF=100 on S3", - storage=Storage.S3, - scale_factor=100.0, - iterations=3, - local_dir="vortex-bench/data/tpcds/100.0", - remote_key="tpcds/100.0", - runs={"nightly": DEFAULT}, - ), + *TPCDS_SF100_CASES, BenchmarkCase( id="statpopgen", benchmark=Benchmark.STATPOPGEN, diff --git a/bench-orchestrator/tests/test_matrix.py b/bench-orchestrator/tests/test_matrix.py index 9ce0b0dd83f..a1b202a68e5 100644 --- a/bench-orchestrator/tests/test_matrix.py +++ b/bench-orchestrator/tests/test_matrix.py @@ -59,9 +59,12 @@ "tpch-nvme", "tpch-s3", "tpcds-nvme-10", - "tpcds-nvme-100", "tpcds-s3-10", - "tpcds-s3-100", + *( + f"tpcds-{storage}-100-{target}" + for storage in ("nvme", "s3") + for target in ("datafusion-parquet", "datafusion-vortex", "duckdb-parquet", "duckdb-vortex") + ), ), } @@ -86,6 +89,17 @@ def test_matrix_presets(preset: str, expected_ids: tuple[str, ...]) -> None: assert "${{" not in json.dumps(entries) +def test_tpcds_sf100_runs_one_target_per_job() -> None: + nightly = [ + entry for entry in _entries("nightly") if str(entry["id"]).startswith("tpcds-") and "-100-" in str(entry["id"]) + ] + + assert len(nightly) == 8 + assert all(len(_targets(entry)) == 1 for entry in nightly) + s3_keys = [entry["remote_key"] for entry in nightly if "remote_key" in entry] + assert len(s3_keys) == len(set(s3_keys)) == 4 + + def test_pr_target_selection() -> None: develop = {entry["id"]: entry for entry in _entries("develop")} pr = {entry["id"]: entry for entry in _entries("pr")} diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index fda1638a63e..1398ea37835 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -12,6 +12,8 @@ use futures::StreamExt; use futures::TryStreamExt; use parquet::arrow::AsyncArrowWriter; use parquet::arrow::ParquetRecordBatchStreamBuilder; +use parquet::arrow::arrow_reader::ArrowReaderMetadata; +use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::ParquetRecordBatchStream; use parquet::file::metadata::KeyValue; use parquet::file::metadata::ParquetMetaData; @@ -177,27 +179,94 @@ pub fn parquet_to_vortex_stream( }) } +/// How many row groups of one Parquet file to decode concurrently. +/// +/// Each in-flight row group is held fully decoded, so the count is bounded by the per-stream +/// memory budget ([`MEMORY_PER_STREAM_GB`]) divided by the largest row group's uncompressed size, +/// and by the available parallelism. Files with row groups too large to decode two at a time +/// return 1 and keep the batch-at-a-time streaming path. +fn row_group_decode_concurrency(metadata: &ParquetMetaData) -> usize { + let largest_row_group = metadata + .row_groups() + .iter() + .map(|row_group| u64::try_from(row_group.total_byte_size()).unwrap_or(u64::MAX)) + .max() + .unwrap_or(0) + .max(1); + let budget = MEMORY_PER_STREAM_GB * 1024 * 1024 * 1024; + let fits = usize::try_from(budget / largest_row_group).unwrap_or(usize::MAX); + fits.min(get_available_parallelism().unwrap_or(1)) + .min(metadata.num_row_groups()) + .max(1) +} + +/// Decode a Parquet file's row groups concurrently, yielding canonical Vortex chunks in file order. +/// +/// A single [`ParquetRecordBatchStream`] decodes on one thread, so a large fact table would +/// otherwise convert at the speed of one core while the parallel compressor waits for input. +/// Each row group gets its own reader over the shared footer `metadata`, and at most +/// `concurrency` decoded row groups are held in memory at once. +fn parquet_row_groups_to_vortex_stream( + parquet_path: PathBuf, + metadata: ArrowReaderMetadata, + concurrency: usize, +) -> impl futures::Stream> { + let num_row_groups = metadata.metadata().num_row_groups(); + futures::stream::iter(0..num_row_groups) + .map(move |row_group| { + let parquet_path = parquet_path.clone(); + let metadata = metadata.clone(); + tokio::spawn(async move { + let file = File::open(parquet_path) + .await + .map_err(|e| vortex_err!(External: e))?; + let reader = ParquetRecordBatchStreamBuilder::new_with_metadata(file, metadata) + .with_row_groups(vec![row_group]) + .build() + .map_err(|e| vortex_err!(External: e))?; + parquet_to_vortex_stream(reader) + .try_collect::>() + .await + }) + }) + .buffered(concurrency) + .map(|joined| joined.map_err(|e| vortex_err!(External: e))?) + .map_ok(|chunks| futures::stream::iter(chunks.into_iter().map(Ok))) + .try_flatten() +} + /// Convert a single Parquet file to Vortex format using streaming. /// /// Streams data directly from Parquet to Vortex without loading the entire file into memory. +/// Row groups are decoded in parallel, within the per-stream memory budget. pub async fn convert_parquet_file_to_vortex( parquet_path: &Path, output_path: &Path, compaction: CompactionStrategy, ) -> anyhow::Result<()> { - let file = File::open(parquet_path).await?; - let builder = ParquetRecordBatchStreamBuilder::new(file).await?; + let mut file = File::open(parquet_path).await?; + let metadata = ArrowReaderMetadata::load_async(&mut file, ArrowReaderOptions::new()).await?; // GeoParquet geometry tagging. - let spatial_columns = geoparquet_columns(builder.metadata()); + let spatial_columns = geoparquet_columns(metadata.metadata()); let dtype = tag_spatial_dtype( SESSION .arrow() - .from_arrow_schema(builder.schema().as_ref())?, + .from_arrow_schema(metadata.schema().as_ref())?, &spatial_columns, )?; - let stream = parquet_to_vortex_stream(builder.build()?) - .map(move |chunk| chunk.and_then(|chunk| tag_spatial_array(chunk, &spatial_columns))); + let concurrency = row_group_decode_concurrency(metadata.metadata()); + let stream = if concurrency > 1 { + StreamExt::boxed(parquet_row_groups_to_vortex_stream( + parquet_path.to_path_buf(), + metadata, + concurrency, + )) + } else { + let reader = ParquetRecordBatchStreamBuilder::new_with_metadata(file, metadata).build()?; + StreamExt::boxed(parquet_to_vortex_stream(reader)) + } + .map(move |chunk| chunk.and_then(|chunk| tag_spatial_array(chunk, &spatial_columns))); let mut output_file = OpenOptions::new() .write(true)