From f39662fcd9d7e4f4ef4d6524cdb4ad093cdbc972 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 3 Sep 2026 11:46:34 +0100 Subject: [PATCH 1/3] Reuse compression/decompression results in compression benchmarks Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + benchmarks/compress-bench/src/arrow.rs | 45 ++++--- benchmarks/compress-bench/src/gpu/parquet.rs | 51 +++++--- benchmarks/compress-bench/src/gpu/vortex.rs | 40 ++++-- benchmarks/compress-bench/src/main.rs | 32 +++-- benchmarks/compress-bench/src/parquet.rs | 49 ++++--- benchmarks/compress-bench/src/vortex.rs | 38 +++--- benchmarks/lance-bench/src/compress.rs | 73 ++++++----- vortex-bench/Cargo.toml | 3 +- vortex-bench/src/compress/mod.rs | 128 ++++++++++++++++--- 10 files changed, 310 insertions(+), 150 deletions(-) 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..1054ebbccd1 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 } @@ -59,6 +60,7 @@ regex = { workspace = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +tempfile = { workspace = true } sha2 = { workspace = true } spatialbench = { workspace = true } spatialbench-arrow = { 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..332ed529124 100644 --- a/vortex-bench/src/compress/mod.rs +++ b/vortex-bench/src/compress/mod.rs @@ -1,15 +1,27 @@ // 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::utils::aliases::hash_map::HashMap; use crate::Format; @@ -84,10 +96,83 @@ 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"), + } + } +} + +/// 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 +193,55 @@ 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. 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?; + 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 +258,7 @@ pub async fn benchmark_compress( Ok(CompressResult { time: fastest, compressed_size, + compressed, timing, all_runs, ratios, @@ -177,10 +267,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 +279,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); From fc2a30c0f23b4aa59e5972266ac372bdbb4a3888 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 3 Sep 2026 13:54:28 +0100 Subject: [PATCH 2/3] fixes Signed-off-by: Robert Kruszewski --- vortex-bench/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 1054ebbccd1..0e078f3471f 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -60,7 +60,6 @@ regex = { workspace = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } -tempfile = { workspace = true } sha2 = { workspace = true } spatialbench = { workspace = true } spatialbench-arrow = { workspace = true } @@ -68,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 } From c7361f2af11a6ff061d75357451f7eb7b9453db0 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 3 Sep 2026 22:02:45 +0000 Subject: [PATCH 3/3] Clear cached array stats between compression iterations Reusing one Vortex array across iterations let every run after the first find its statistics already cached, so the timed region shrank to just the encoding. Reset the input before each iteration so the measurement matches a fresh conversion. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019CX6MXw5FFYnqQNdV7K4ci --- vortex-bench/src/compress/mod.rs | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/vortex-bench/src/compress/mod.rs b/vortex-bench/src/compress/mod.rs index 332ed529124..d98bd6a795d 100644 --- a/vortex-bench/src/compress/mod.rs +++ b/vortex-bench/src/compress/mod.rs @@ -22,6 +22,7 @@ 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; @@ -145,6 +146,27 @@ impl Uncompressed { _ => 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. @@ -219,7 +241,8 @@ pub trait Compressor: Send + Sync { /// Run a compression benchmark for the given compressor. /// -/// Compresses the same `input` `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, input: &Uncompressed, @@ -232,6 +255,7 @@ pub async fn benchmark_compress( let mut compressed = None; for _ in 0..iterations { + input.reset(); let result = compressor.compress(input).await?; fastest = fastest.min(result.elapsed);