diff --git a/Cargo.lock b/Cargo.lock index 0a367df1534..b0b9f37d024 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10639,6 +10639,7 @@ dependencies = [ "arrow-schema 59.2.0", "arrow-select 59.2.0", "async-trait", + "bytes", "bzip2", "clap", "futures", diff --git a/benchmarks/compress-bench/src/arrow.rs b/benchmarks/compress-bench/src/arrow.rs index 442b2c52d7d..98175d3c94a 100644 --- a/benchmarks/compress-bench/src/arrow.rs +++ b/benchmarks/compress-bench/src/arrow.rs @@ -1,22 +1,23 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fs::File; use std::io::Cursor; use std::path::Path; -use std::sync::Arc; use std::time::Duration; use std::time::Instant; +use anyhow::bail; use arrow_array::RecordBatch; use arrow_ipc::reader::FileReader; use arrow_ipc::writer::FileWriter; use arrow_schema::Schema; use async_trait::async_trait; use bytes::Bytes; -use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex_bench::Format; +use vortex_bench::compress::Compressed; +use vortex_bench::compress::CompressedData; use vortex_bench::compress::Compressor; +use vortex_bench::compress::Uncompressed; use vortex_bench::compress::read_projection; /// Uncompressed Arrow IPC file baseline. @@ -28,30 +29,38 @@ impl Compressor for ArrowIpcCompressor { Format::ArrowIpc } - async fn compress(&self, parquet_path: &Path) -> anyhow::Result<(u64, Duration)> { - let file = File::open(parquet_path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let schema = Arc::clone(builder.schema()); - let batches = builder.build()?.collect::, _>>()?; + async fn load(&self, parquet_path: &Path) -> anyhow::Result { + Uncompressed::read_arrow(parquet_path) + } + + async fn compress(&self, input: &Uncompressed) -> anyhow::Result { + let (schema, batches) = input.arrow()?; let mut buf = Vec::new(); let start = Instant::now(); - arrow_file_write(&mut buf, &schema, &batches)?; + arrow_file_write(&mut buf, schema, batches)?; let elapsed = start.elapsed(); - Ok((buf.len() as u64, elapsed)) + + Ok(Compressed { + size: buf.len() as u64, + data: CompressedData::Bytes(Bytes::from(buf)), + elapsed, + }) } - async fn decompress(&self, parquet_path: &Path) -> anyhow::Result { - let file = File::open(parquet_path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let schema = Arc::clone(builder.schema()); - let batches = builder.build()?.collect::, _>>()?; + async fn decompress(&self, compressed: &Compressed) -> anyhow::Result { + let CompressedData::Bytes(buf) = &compressed.data else { + bail!("Arrow IPC decompression expects in-memory bytes"); + }; - let mut buf = Vec::new(); - arrow_file_write(&mut buf, &schema, &batches)?; + // The projection needs the column count; read the footer before the clock starts. + let root_columns = FileReader::try_new(Cursor::new(buf.clone()), None)? + .schema() + .fields() + .len(); let start = Instant::now(); - arrow_file_read(Bytes::from(buf), schema.fields().len())?; + arrow_file_read(buf.clone(), root_columns)?; Ok(start.elapsed()) } } diff --git a/benchmarks/compress-bench/src/gpu/parquet.rs b/benchmarks/compress-bench/src/gpu/parquet.rs index 2d01a1527eb..20c199fa331 100644 --- a/benchmarks/compress-bench/src/gpu/parquet.rs +++ b/benchmarks/compress-bench/src/gpu/parquet.rs @@ -22,19 +22,21 @@ use std::path::Path; use std::process::Command; use std::sync::Arc; use std::time::Duration; +use std::time::Instant; use anyhow::Context; use anyhow::Result; use anyhow::bail; use anyhow::ensure; -use arrow_array::RecordBatch; use async_trait::async_trait; use parquet::arrow::ArrowWriter; -use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use serde::Deserialize; use tempfile::NamedTempFile; use vortex_bench::Format; +use vortex_bench::compress::Compressed; +use vortex_bench::compress::CompressedData; use vortex_bench::compress::Compressor; +use vortex_bench::compress::Uncompressed; use crate::gpu::writer::GpuCodec; use crate::gpu::writer::gpu_writer_properties; @@ -84,25 +86,30 @@ impl GpuParquetCompressor { Self { codec, verify } } - /// Rewrite the source Parquet file with GPU-friendly writer settings. - fn write_gpu_parquet(&self, parquet_path: &Path) -> Result<(NamedTempFile, u64)> { - let builder = ParquetRecordBatchReaderBuilder::try_new(std::fs::File::open(parquet_path)?)?; - let schema = Arc::clone(builder.schema()); - let batches: Vec = builder.build()?.collect::, _>>()?; + /// Rewrite the source data as a Parquet file with GPU-friendly writer settings. + fn write_gpu_parquet(&self, input: &Uncompressed) -> Result { + let (schema, batches) = input.arrow()?; let output = NamedTempFile::new()?; + let start = Instant::now(); let mut writer = ArrowWriter::try_new( output.reopen()?, - schema, + Arc::clone(schema), Some(gpu_writer_properties(self.codec)), )?; for batch in batches { - writer.write(&batch)?; + writer.write(batch)?; } writer.flush()?; let size = writer.bytes_written() as u64; writer.close()?; - Ok((output, size)) + let elapsed = start.elapsed(); + + Ok(Compressed { + data: CompressedData::File(output), + size, + elapsed, + }) } } @@ -112,21 +119,25 @@ impl Compressor for GpuParquetCompressor { Format::Parquet } - /// Unsupported: GPU mode measures decompression only. + /// Writes the file that [`Self::decompress`] reads. /// - /// `--gpu-decompress` restricts the suite to [`CompressOp::Decompress`], so nothing calls - /// this. It used to time [`Self::write_gpu_parquet`], but that measures the host Parquet - /// writer rather than anything on the device, and the result was never rendered — so it was - /// a number nobody could read and nobody should have compared. The Vortex GPU backend - /// refuses the same way. + /// `--gpu-decompress` restricts the suite to [`CompressOp::Decompress`], so the timing this + /// returns is never published. That is deliberate: it measures the host Parquet writer + /// rather than anything on the device. /// /// [`CompressOp::Decompress`]: vortex_bench::compress::CompressOp::Decompress - async fn compress(&self, _parquet_path: &Path) -> Result<(u64, Duration)> { - bail!("GPU compress-bench only supports decompression measurements") + async fn load(&self, parquet_path: &Path) -> Result { + Uncompressed::read_arrow(parquet_path) + } + + async fn compress(&self, input: &Uncompressed) -> Result { + self.write_gpu_parquet(input) } - async fn decompress(&self, parquet_path: &Path) -> Result { - let (gpu_file, _) = self.write_gpu_parquet(parquet_path)?; + async fn decompress(&self, compressed: &Compressed) -> Result { + let CompressedData::File(gpu_file) = &compressed.data else { + bail!("GPU Parquet decompression expects a file on disk"); + }; let report = run_cudf_read(gpu_file.path(), self.verify)?; ensure!( diff --git a/benchmarks/compress-bench/src/gpu/vortex.rs b/benchmarks/compress-bench/src/gpu/vortex.rs index be1ab61868f..0a1e0797027 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/benchmarks/compress-bench/src/gpu/vortex.rs @@ -33,7 +33,10 @@ use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::benchmark_write_options; +use vortex_bench::compress::Compressed; +use vortex_bench::compress::CompressedData; use vortex_bench::compress::Compressor; +use vortex_bench::compress::Uncompressed; use vortex_bench::conversions::parquet_to_vortex_chunks_with_batch_size; use vortex_cuda::CanonicalCudaExt; use vortex_cuda::CudaExecutionCtx; @@ -71,21 +74,26 @@ impl Compressor for GpuVortexCompressor { Format::OnDiskVortex } - async fn compress(&self, _parquet_path: &Path) -> Result<(u64, Duration)> { - anyhow::bail!("GPU compress-bench only supports decompression measurements") - } - - async fn decompress(&self, parquet_path: &Path) -> Result { - register_cuda_layout(&SESSION); - + async fn load(&self, parquet_path: &Path) -> Result { // Rebatch to the same partition size the GPU Parquet file is written with. Left alone, // the Arrow reader hands back ~8K-row batches, each of which becomes its own Vortex // chunk and its own set of kernel launches. - let uncompressed = parquet_to_vortex_chunks_with_batch_size( + let chunks = parquet_to_vortex_chunks_with_batch_size( parquet_path.to_path_buf(), Some(GPU_ROW_GROUP_SIZE), ) .await?; + Ok(Uncompressed::Vortex(chunks.into_array())) + } + + /// Writes the CUDA-compatible file that [`Self::decompress`] reads. + /// + /// GPU mode never publishes this timing: `--gpu-decompress` restricts the suite to + /// decompression, and the write runs on the host anyway. + async fn compress(&self, input: &Uncompressed) -> Result { + register_cuda_layout(&SESSION); + + let array = input.vortex()?; let gpu_file = NamedTempFile::new()?; let mut output = tokio::fs::File::create(gpu_file.path()).await?; // Write those batches straight through as root chunks, so a chunk on disk is one @@ -96,13 +104,27 @@ impl Compressor for GpuVortexCompressor { .only_cuda_compatible() .build(), ))); + let start = Instant::now(); benchmark_write_options(SESSION.write_options()) .with_strategy(strategy) - .write(&mut output, uncompressed.into_array().to_array_stream()) + .write(&mut output, array.to_array_stream()) .await?; output.sync_all().await?; + let elapsed = start.elapsed(); drop(output); + Ok(Compressed { + size: gpu_file.as_file().metadata()?.len(), + data: CompressedData::File(gpu_file), + elapsed, + }) + } + + async fn decompress(&self, compressed: &Compressed) -> Result { + let CompressedData::File(gpu_file) = &compressed.data else { + bail!("GPU Vortex decompression expects a file on disk"); + }; + // Verification is a precondition on the measurement below, not a substitute for it. It // used to return its own elapsed time, which bundled a file copy, a second host scan and // every Arrow conversion into a number the table then published as a decode time. diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 589ec5854ae..d2b753e4e73 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -28,6 +28,7 @@ use vortex_bench::LogFormat; use vortex_bench::Target; use vortex_bench::compress::CompressMeasurements; use vortex_bench::compress::CompressOp; +use vortex_bench::compress::Compressed; use vortex_bench::compress::Compressor; use vortex_bench::compress::benchmark_compress; use vortex_bench::compress::benchmark_decompress; @@ -442,18 +443,21 @@ async fn run_benchmark_for_dataset( for format in formats { let compressor = get_compressor(*format, mode); + // Read the source once per format; every compression iteration starts from it. + let input = compressor + .load(&parquet_path) + .await + .with_context(|| format!("loading {bench_name} for {format}"))?; + // Compressed output shared by both ops, so decompression never compresses again. + let mut compressed: Option = None; for op in ops { let time = match op { CompressOp::Compress => { - let result = benchmark_compress( - compressor.as_ref(), - &parquet_path, - iterations, - bench_name, - ) - .await - .with_context(|| format!("compressing {bench_name} as {format}"))?; + let result = + benchmark_compress(compressor.as_ref(), &input, iterations, bench_name) + .await + .with_context(|| format!("compressing {bench_name} as {format}"))?; compressed_sizes.insert(*format, result.compressed_size); let all_runs_ns: Vec = result .all_runs @@ -476,12 +480,22 @@ async fn run_benchmark_for_dataset( )); ratios.extend(result.ratios); timings.push(result.timing); + compressed = Some(result.compressed); result.time } CompressOp::Decompress => { + let input = match &compressed { + Some(input) => input, + None => compressed.insert( + compressor + .compress(&input) + .await + .with_context(|| format!("compressing {bench_name} as {format}"))?, + ), + }; let result = benchmark_decompress( compressor.as_ref(), - &parquet_path, + input, iterations, &decompress_name, ) diff --git a/benchmarks/compress-bench/src/parquet.rs b/benchmarks/compress-bench/src/parquet.rs index e2904de2911..ddb5faa2afa 100644 --- a/benchmarks/compress-bench/src/parquet.rs +++ b/benchmarks/compress-bench/src/parquet.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; +use anyhow::bail; use arrow_array::RecordBatch; use arrow_schema::Schema; use async_trait::async_trait; @@ -19,7 +20,10 @@ use parquet::basic::Compression; use parquet::basic::ZstdLevel; use parquet::file::properties::WriterProperties; use vortex_bench::Format; +use vortex_bench::compress::Compressed; +use vortex_bench::compress::CompressedData; use vortex_bench::compress::Compressor; +use vortex_bench::compress::Uncompressed; use vortex_bench::compress::read_projection; /// Compressor implementation for Parquet format with ZSTD compression. @@ -68,45 +72,40 @@ impl Compressor for ParquetCompressor { Format::Parquet } - async fn compress(&self, parquet_path: &Path) -> anyhow::Result<(u64, Duration)> { - // Read the input parquet file - let file = File::open(parquet_path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let schema = Arc::clone(builder.schema()); - let reader = builder.build()?; - let batches: Vec = reader.collect::, _>>()?; + async fn load(&self, parquet_path: &Path) -> anyhow::Result { + Uncompressed::read_arrow(parquet_path) + } + + async fn compress(&self, input: &Uncompressed) -> anyhow::Result { + let (schema, batches) = input.arrow()?; // Compress with our compression settings let mut buf = Vec::new(); let start = Instant::now(); - let size = parquet_compress_write(batches, schema, self.compression, &mut buf)?; + let size = parquet_compress_write(batches, Arc::clone(schema), self.compression, &mut buf)?; let elapsed = start.elapsed(); - Ok((size as u64, elapsed)) - } - async fn decompress(&self, parquet_path: &Path) -> anyhow::Result { - // First compress to get the bytes we'll decompress - let file = File::open(parquet_path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let schema = Arc::clone(builder.schema()); - let reader = builder.build()?; - let batches: Vec = reader.collect::, _>>()?; - - let mut buf = Vec::new(); - parquet_compress_write(batches, schema, self.compression, &mut buf)?; + Ok(Compressed { + data: CompressedData::Bytes(Bytes::from(buf)), + size: size as u64, + elapsed, + }) + } - let buf = Bytes::from(buf); + async fn decompress(&self, compressed: &Compressed) -> anyhow::Result { + let CompressedData::Bytes(buf) = &compressed.data else { + bail!("Parquet decompression expects in-memory bytes"); + }; - // Now decompress let timer = Instant::now(); - parquet_decompress_read(buf)?; + parquet_decompress_read(buf.clone())?; Ok(timer.elapsed()) } } #[inline(never)] pub fn parquet_compress_write( - batches: Vec, + batches: &[RecordBatch], schema: Arc, compression: Compression, buf: &mut Vec, @@ -117,7 +116,7 @@ pub fn parquet_compress_write( .build(); let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(writer_properties))?; for batch in batches { - writer.write(&batch)?; + writer.write(batch)?; } writer.flush()?; let n_bytes = writer.bytes_written(); diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index a030a54b903..7dfefb8f83b 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -8,6 +8,7 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; +use anyhow::bail; use async_trait::async_trait; use bytes::Bytes; use futures::StreamExt; @@ -22,7 +23,10 @@ use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::benchmark_write_options; +use vortex_bench::compress::Compressed; +use vortex_bench::compress::CompressedData; use vortex_bench::compress::Compressor; +use vortex_bench::compress::Uncompressed; use vortex_bench::compress::read_projection; use vortex_bench::conversions::parquet_to_vortex_chunks; @@ -35,34 +39,36 @@ impl Compressor for VortexCompressor { Format::OnDiskVortex } - async fn compress(&self, parquet_path: &Path) -> Result<(u64, Duration)> { - // Read the parquet file as an array stream - let uncompressed = parquet_to_vortex_chunks(parquet_path.to_path_buf()).await?; + async fn load(&self, parquet_path: &Path) -> Result { + let chunks = parquet_to_vortex_chunks(parquet_path.to_path_buf()).await?; + Ok(Uncompressed::Vortex(chunks.into_array())) + } + + async fn compress(&self, input: &Uncompressed) -> Result { + let array = input.vortex()?; let mut buf = Vec::new(); let start = Instant::now(); let mut cursor = Cursor::new(&mut buf); benchmark_write_options(SESSION.write_options()) - .write(&mut cursor, uncompressed.into_array().to_array_stream()) + .write(&mut cursor, array.to_array_stream()) .await?; let elapsed = start.elapsed(); - Ok((buf.len() as u64, elapsed)) + Ok(Compressed { + size: buf.len() as u64, + data: CompressedData::Bytes(Bytes::from(buf)), + elapsed, + }) } - async fn decompress(&self, parquet_path: &Path) -> Result { - // First compress to get the bytes we'll decompress - let uncompressed = parquet_to_vortex_chunks(parquet_path.to_path_buf()).await?; - let mut buf = Vec::new(); - let mut cursor = Cursor::new(&mut buf); - benchmark_write_options(SESSION.write_options()) - .write(&mut cursor, uncompressed.into_array().to_array_stream()) - .await?; + async fn decompress(&self, compressed: &Compressed) -> Result { + let CompressedData::Bytes(data) = &compressed.data else { + bail!("Vortex decompression expects in-memory bytes"); + }; - // Now decompress let start = Instant::now(); - let data = Bytes::from(buf); - let mut scan = SESSION.open_options().open_buffer(data)?.scan()?; + let mut scan = SESSION.open_options().open_buffer(data.clone())?.scan()?; let source_dtype = scan.dtype()?; let root_columns = source_dtype .as_struct_fields_opt() diff --git a/benchmarks/lance-bench/src/compress.rs b/benchmarks/lance-bench/src/compress.rs index 4d8df8f24dd..d21e0462feb 100644 --- a/benchmarks/lance-bench/src/compress.rs +++ b/benchmarks/lance-bench/src/compress.rs @@ -9,17 +9,22 @@ use std::time::Duration; use std::time::Instant; use anyhow::anyhow; +use anyhow::bail; use async_trait::async_trait; use futures::StreamExt; use lance::dataset::Dataset; use lance::dataset::WriteParams; use lance::deps::arrow_array::RecordBatch; use lance::deps::arrow_array::RecordBatchIterator; +use lance::deps::arrow_schema::SchemaRef; use lance_encoding::version::LanceFileVersion; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use tempfile::TempDir; use vortex_bench::Format; +use vortex_bench::compress::Compressed; +use vortex_bench::compress::CompressedData; use vortex_bench::compress::Compressor; +use vortex_bench::compress::Uncompressed; use vortex_bench::compress::read_projection; use crate::convert::convert_utf8view_batch; @@ -84,6 +89,15 @@ pub fn calculate_lance_size(dataset_path: &Path) -> anyhow::Result { Ok(total_size) } +/// Subdirectory of the temp directory holding the Lance dataset. +const DATASET_DIR: &str = "dataset"; + +/// Source batches in Lance's Arrow version, which differs from the workspace's. +struct LanceInput { + schema: SchemaRef, + batches: Vec, +} + /// Compressor implementation for Lance format. /// /// Lance writes to the filesystem rather than in-memory buffers, so this implementation @@ -96,7 +110,7 @@ impl Compressor for LanceCompressor { Format::Lance } - async fn compress(&self, parquet_path: &Path) -> anyhow::Result<(u64, Duration)> { + async fn load(&self, parquet_path: &Path) -> anyhow::Result { // Read the input parquet file let file = File::open(parquet_path)?; let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; @@ -105,15 +119,26 @@ impl Compressor for LanceCompressor { let batches: Vec = reader.collect::, _>>()?; // Convert Utf8View columns to Utf8 (Lance doesn't support Utf8View) - let converted_batches: Vec = batches + let batches = batches .into_iter() .map(convert_utf8view_batch) .collect::>>()?; - let converted_schema = convert_utf8view_schema(&schema); + let schema = convert_utf8view_schema(&schema); + + Ok(Uncompressed::Opaque(Box::new(LanceInput { + schema, + batches, + }))) + } + + async fn compress(&self, input: &Uncompressed) -> anyhow::Result { + let LanceInput { schema, batches } = input.opaque()?; + // The Lance writer wants an owned reader; batch clones only bump buffer refcounts. + let batches = batches.clone(); // Create temp directory for Lance dataset let temp_dir = TempDir::new()?; - let dataset_path = temp_dir.path().join("dataset"); + let dataset_path = temp_dir.path().join(DATASET_DIR); fs::create_dir_all(&dataset_path)?; let start = Instant::now(); @@ -122,8 +147,7 @@ impl Compressor for LanceCompressor { let path_str = dataset_path .to_str() .ok_or_else(|| anyhow!("Failed to convert path to str"))?; - let reader_iter = - RecordBatchIterator::new(converted_batches.into_iter().map(Ok), converted_schema); + let reader_iter = RecordBatchIterator::new(batches.into_iter().map(Ok), Arc::clone(schema)); let write_params = WriteParams::with_storage_version(LanceFileVersion::V2_0); Dataset::write(reader_iter, path_str, Some(write_params)).await?; @@ -132,39 +156,22 @@ impl Compressor for LanceCompressor { // Calculate size of Lance files on disk let size = calculate_lance_size(&dataset_path)?; - Ok((size, elapsed)) + Ok(Compressed { + data: CompressedData::Dir(temp_dir), + size, + elapsed, + }) } - async fn decompress(&self, parquet_path: &Path) -> anyhow::Result { - // First compress to get the Lance dataset - let file = File::open(parquet_path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let schema = Arc::clone(builder.schema()); - let reader = builder.build()?; - let batches: Vec = reader.collect::, _>>()?; - - // Convert Utf8View columns to Utf8 (Lance doesn't support Utf8View) - let converted_batches: Vec = batches - .into_iter() - .map(convert_utf8view_batch) - .collect::>>()?; - let converted_schema = convert_utf8view_schema(&schema); - - // Create temp directory for Lance dataset - let temp_dir = TempDir::new()?; - let dataset_path = temp_dir.path().join("dataset"); - fs::create_dir_all(&dataset_path)?; - - // Write to Lance format + async fn decompress(&self, compressed: &Compressed) -> anyhow::Result { + let CompressedData::Dir(temp_dir) = &compressed.data else { + bail!("Lance decompression expects a dataset directory"); + }; + let dataset_path = temp_dir.path().join(DATASET_DIR); let path_str = dataset_path .to_str() .ok_or_else(|| anyhow!("Failed to convert path to str"))?; - let reader_iter = - RecordBatchIterator::new(converted_batches.into_iter().map(Ok), converted_schema); - let write_params = WriteParams::with_storage_version(LanceFileVersion::V2_0); - Dataset::write(reader_iter, path_str, Some(write_params)).await?; - // Now decompress let start = Instant::now(); lance_decompress_read(path_str).await?; Ok(start.elapsed()) diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 5e1298b3411..0e078f3471f 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -37,6 +37,7 @@ arrow-ipc = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } async-trait = { workspace = true } +bytes = { workspace = true } bzip2 = { workspace = true } clap = { workspace = true, features = ["derive"] } futures = { workspace = true } @@ -66,6 +67,7 @@ spatialbench-parquet = { workspace = true } sysinfo = { workspace = true } tabled = { workspace = true, features = ["std"] } target-lexicon = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["full"] } tokio-stream = { workspace = true } tokio-util = { workspace = true } @@ -85,7 +87,6 @@ wkb = { workspace = true } [dev-dependencies] insta = { workspace = true } rstest = { workspace = true } -tempfile = { workspace = true } [features] unstable_encodings = ["vortex/unstable_encodings"] diff --git a/vortex-bench/src/compress/mod.rs b/vortex-bench/src/compress/mod.rs index 6e261b871d2..d98bd6a795d 100644 --- a/vortex-bench/src/compress/mod.rs +++ b/vortex-bench/src/compress/mod.rs @@ -1,15 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; use std::borrow::Cow; use std::fmt; +use std::fs::File; use std::path::Path; +use std::sync::Arc; use std::time::Duration; +use anyhow::Context; use anyhow::Result; +use anyhow::bail; +use arrow_array::RecordBatch; +use arrow_schema::Schema; use async_trait::async_trait; +use bytes::Bytes; use clap::ValueEnum; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use serde::Serialize; +use tempfile::NamedTempFile; +use tempfile::TempDir; +use vortex::array::ArrayRef; +use vortex::expr::stats::Stat; use vortex::utils::aliases::hash_map::HashMap; use crate::Format; @@ -84,10 +97,104 @@ impl fmt::Display for CompressOp { } } +/// Source data in the form a format compresses, read from Parquet once per format. +pub enum Uncompressed { + /// Arrow record batches sharing one schema. + Arrow { + schema: Arc, + batches: Vec, + }, + /// A Vortex array. + Vortex(ArrayRef), + /// Format-private input, for backends whose Arrow crate version differs from the + /// workspace's and so cannot share [`Self::Arrow`]. + Opaque(Box), +} + +impl Uncompressed { + /// Read a Parquet file into Arrow record batches. + pub fn read_arrow(parquet_path: &Path) -> Result { + let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(parquet_path)?)?; + let schema = Arc::clone(builder.schema()); + let batches = builder.build()?.collect::, _>>()?; + + Ok(Self::Arrow { schema, batches }) + } + + /// The Arrow schema and batches, or an error if this is not Arrow data. + pub fn arrow(&self) -> Result<(&Arc, &[RecordBatch])> { + match self { + Self::Arrow { schema, batches } => Ok((schema, batches)), + _ => bail!("expected Arrow record batches"), + } + } + + /// The Vortex array, or an error if this is not Vortex data. + pub fn vortex(&self) -> Result<&ArrayRef> { + match self { + Self::Vortex(array) => Ok(array), + _ => bail!("expected a Vortex array"), + } + } + + /// The format-private input as `T`, or an error if this holds something else. + pub fn opaque(&self) -> Result<&T> { + match self { + Self::Opaque(input) => input + .downcast_ref() + .ok_or_else(|| anyhow::anyhow!("format-private input has an unexpected type")), + _ => bail!("expected format-private input"), + } + } + + /// Return the input to the state [`Self::read_arrow`] or a fresh conversion would produce. + /// + /// The Vortex writer computes statistics inside the timed region and caches them on the + /// array, so reusing one array across iterations would let every run after the first skip + /// that work. Clearing the cache keeps each iteration's measurement comparable. + pub fn reset(&self) { + if let Self::Vortex(array) = self { + clear_stats(array); + } + } +} + +/// Clear cached statistics on `array` and every array beneath it. +fn clear_stats(array: &ArrayRef) { + for stat in Stat::all() { + array.statistics().clear(stat); + } + for child in array.children_iter() { + clear_stats(child); + } +} + +/// Where a format keeps its compressed output. +/// +/// Temporary files and directories are removed on drop, so the output lives exactly as long as +/// the value does. +pub enum CompressedData { + /// In-memory file image. + Bytes(Bytes), + /// Single file on disk. + File(NamedTempFile), + /// Directory of files on disk. + Dir(TempDir), +} + +/// Output of one compression run: the compressed data, its size and the compression time. +pub struct Compressed { + pub data: CompressedData, + pub size: u64, + pub elapsed: Duration, +} + /// Result of a compression benchmark run. pub struct CompressResult { pub time: Duration, pub compressed_size: u64, + /// Output of the last iteration, kept so decompression need not compress again. + pub compressed: Compressed, pub timing: CompressionTimingMeasurement, /// Per-iteration encode wall times. Captured for v3 emission. pub all_runs: Vec, @@ -108,51 +215,57 @@ pub struct DecompressResult { /// (e.g., Vortex, Parquet, Lance). The benchmark functions use this trait /// to run timing measurements. /// -/// The input data is provided as a path to a Parquet file, which implementations -/// read and convert as needed for their target format. +/// The input data is provided as a path to a Parquet file. [`Self::load`] reads it once into +/// the form the format compresses; the timed operations never touch the file again. #[async_trait] pub trait Compressor: Send + Sync { /// The format this compressor handles. fn format(&self) -> Format; - /// Compress data from a Parquet file, returning the compressed size in bytes and elapsed time. + /// Read a Parquet file into the form [`Self::compress`] consumes. Not timed. + async fn load(&self, parquet_path: &Path) -> Result; + + /// Compress output previously produced by [`Self::load`]. /// - /// The implementation should read the Parquet file and compress it - /// to the target format. - async fn compress(&self, parquet_path: &Path) -> Result<(u64, Duration)>; + /// Only the compression itself should be timed. + async fn compress(&self, input: &Uncompressed) -> Result; - /// Decompress data from the Parquet file (after compressing), returning the decompressed size. + /// Decompress output previously produced by [`Self::compress`]. /// - /// This method first compresses the data to the target format, then decompresses it. /// The timing returned should only measure the decompression phase. /// /// Format implementations apply the fixed wide-table read projection when the input schema /// matches the projection benchmark. - async fn decompress(&self, parquet_path: &Path) -> Result; + async fn decompress(&self, compressed: &Compressed) -> Result; } /// Run a compression benchmark for the given compressor. /// -/// Executes compression `iterations` times and returns timing statistics. +/// Compresses the same `input` `iterations` times and returns timing statistics. The input is +/// [reset](Uncompressed::reset) before every iteration so none of them starts warm. pub async fn benchmark_compress( compressor: &dyn Compressor, - parquet_path: &Path, + input: &Uncompressed, iterations: usize, bench_name: &str, ) -> Result { let format = compressor.format(); let mut fastest = Duration::MAX; - let mut compressed_size = 0u64; let mut all_runs = Vec::with_capacity(iterations); + let mut compressed = None; for _ in 0..iterations { - let (size, elapsed) = compressor.compress(parquet_path).await?; + input.reset(); + let result = compressor.compress(input).await?; - compressed_size = size; - fastest = fastest.min(elapsed); - all_runs.push(elapsed); + fastest = fastest.min(result.elapsed); + all_runs.push(result.elapsed); + compressed = Some(result); } + let compressed = compressed.context("--iterations must be at least 1")?; + let compressed_size = compressed.size; + let ratios = vec![CustomUnitMeasurement { name: format!("{} size/{bench_name}", format.name()), format, @@ -169,6 +282,7 @@ pub async fn benchmark_compress( Ok(CompressResult { time: fastest, compressed_size, + compressed, timing, all_runs, ratios, @@ -177,10 +291,10 @@ pub async fn benchmark_compress( /// Run a decompression benchmark for the given compressor. /// -/// Benchmarks decompression `iterations` times. +/// Decompresses the same `compressed` output `iterations` times. pub async fn benchmark_decompress( compressor: &dyn Compressor, - parquet_path: &Path, + compressed: &Compressed, iterations: usize, bench_name: &str, ) -> Result { @@ -189,7 +303,7 @@ pub async fn benchmark_decompress( let mut all_runs = Vec::with_capacity(iterations); for _ in 0..iterations { - let elapsed = compressor.decompress(parquet_path).await?; + let elapsed = compressor.decompress(compressed).await?; fastest = fastest.min(elapsed); all_runs.push(elapsed);