Reuse compression/decompression results in compression benchmarks

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::<Result<Vec<_>, _>>()?;
+    async fn load(&self, parquet_path: &Path) -> anyhow::Result<Uncompressed> {
+        Uncompressed::read_arrow(parquet_path)
+    }
+
+    async fn compress(&self, input: &Uncompressed) -> anyhow::Result<Compressed> {
+        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<Duration> {
-        let file = File::open(parquet_path)?;
-        let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
-        let schema = Arc::clone(builder.schema());
-        let batches = builder.build()?.collect::<Result<Vec<_>, _>>()?;
+    async fn decompress(&self, compressed: &Compressed) -> anyhow::Result<Duration> {
+        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())
     }
 }

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<RecordBatch> = builder.build()?.collect::<Result<Vec<_>, _>>()?;
+    /// Rewrite the source data as a Parquet file with GPU-friendly writer settings.
+    fn write_gpu_parquet(&self, input: &Uncompressed) -> Result<Compressed> {
+        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> {
+        Uncompressed::read_arrow(parquet_path)
+    }
+
+    async fn compress(&self, input: &Uncompressed) -> Result<Compressed> {
+        self.write_gpu_parquet(input)
     }
 
-    async fn decompress(&self, parquet_path: &Path) -> Result<Duration> {
-        let (gpu_file, _) = self.write_gpu_parquet(parquet_path)?;
+    async fn decompress(&self, compressed: &Compressed) -> Result<Duration> {
+        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!(

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<Duration> {
-        register_cuda_layout(&SESSION);
-
+    async fn load(&self, parquet_path: &Path) -> Result<Uncompressed> {
         // 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<Compressed> {
+        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<Duration> {
+        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.

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<Compressed> = 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<u64> = 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,
                     )