Move code around for compression

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

.github/workflows/sql-benchmarks.yml

@@ -146,7 +146,7 @@ jobs:
           # Build options string if scale_factor is set
           opts=""
           if [ -n "${{ matrix.scale_factor }}" ]; then
-            opts="--opt scale_factor=${{ matrix.scale_factor }}"
+            opts="--opt scale-factor=${{ matrix.scale_factor }}"
           fi
 
           # Generate all data formats with a single command
@@ -198,7 +198,7 @@ jobs:
           # Build options string if scale_factor is set
           opts=""
           if [ -n "${{ matrix.scale_factor }}" ]; then
-            opts="--opt scale_factor=${{ matrix.scale_factor }}"
+            opts="--opt scale-factor=${{ matrix.scale_factor }}"
           fi
 
           touch results.json
@@ -252,9 +252,9 @@ jobs:
           df_formats=$(echo "${{ matrix.targets }}" | tr ',' '\n' | grep '^datafusion:' | sed 's/datafusion://' | tr '\n' ',' | sed 's/,$//')
 
           # Build options string if scale_factor is set
-          opts="--opt remote_data_dir=${{ matrix.remote_storage }}"
+          opts="--opt remote-data-dir=${{ matrix.remote_storage }}"
           if [ -n "${{ matrix.scale_factor }}" ]; then
-            opts="--opt scale_factor=${{ matrix.scale_factor }} ${opts}"
+            opts="--opt scale-factor=${{ matrix.scale_factor }} ${opts}"
           fi
 
           touch results.json

Cargo.lock

@@ -16,10 +16,17 @@ publish = false
 
 [dependencies]
 anyhow = { workspace = true }
+arrow-array = { workspace = true }
+arrow-schema = { workspace = true }
+async-trait = { workspace = true }
+bytes = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
+futures = { workspace = true }
 indicatif = { workspace = true }
 itertools = { workspace = true }
+parquet = { workspace = true }
 regex = { workspace = true }
+serde = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 tracing = { workspace = true }
 vortex = { workspace = true }

benchmarks/compress-bench/Cargo.toml

@@ -13,14 +13,9 @@ use clap::ValueEnum;
 use serde::Serialize;
 use vortex::array::Array;
 use vortex::utils::aliases::hash_map::HashMap;
-
-use crate::Format;
-use crate::measurements::CompressionTimingMeasurement;
-use crate::measurements::CustomUnitMeasurement;
-
-// ============================================================================
-// Measurement types
-// ============================================================================
+use vortex_bench::Format;
+use vortex_bench::measurements::CompressionTimingMeasurement;
+use vortex_bench::measurements::CustomUnitMeasurement;
 
 #[derive(Default)]
 pub struct CompressMeasurements {
@@ -79,10 +74,6 @@ pub struct DecompressResult {
     pub timing: CompressionTimingMeasurement,
 }
 
-// ============================================================================
-// Compressor trait
-// ============================================================================
-
 /// Trait for format-specific compression/decompression operations.
 ///
 /// Implementations handle the actual compression logic for a specific format
@@ -94,16 +85,12 @@ pub trait Compressor: Send + Sync {
     fn format(&self) -> Format;
 
     /// Compress the array data, returning the compressed bytes.
-    async fn compress(&self, array: &dyn Array) -> Result<Bytes>;
+    async fn compress(&self, array: &dyn Array) -> Result<(Bytes, Duration)>;
 
     /// Decompress the data, returning the decompressed size in bytes.
     async fn decompress(&self, data: Bytes) -> Result<usize>;
 }
 
-// ============================================================================
-// Benchmark functions
-// ============================================================================
-
 /// Run a compression benchmark for the given compressor.
 ///
 /// Executes compression `iterations` times and returns timing statistics.
@@ -118,9 +105,7 @@ pub async fn benchmark_compress(
     let mut compressed_size = 0u64;
 
     for _ in 0..iterations {
-        let start = Instant::now();
-        let compressed = compressor.compress(uncompressed).await?;
-        let elapsed = start.elapsed();
+        let (compressed, elapsed) = compressor.compress(uncompressed).await?;
 
         compressed_size = compressed.len() as u64;
         fastest = fastest.min(elapsed);
@@ -160,7 +145,7 @@ pub async fn benchmark_decompress(
     let format = compressor.format();
 
     // First compress to get the bytes we'll decompress
-    let compressed = compressor.compress(uncompressed).await?;
+    let (compressed, _) = compressor.compress(uncompressed).await?;
 
     let mut fastest = Duration::MAX;
 

vortex-bench/src/compress/bench.rs benchmarks/compress-bench/src/bench.rsvortex-bench/src/compress/bench.rs renamed to benchmarks/compress-bench/src/bench.rs

@@ -12,19 +12,9 @@ pub mod bench;
 pub mod parquet;
 pub mod vortex;
 
-// Re-export the trait and compressor implementations for convenience
-pub use bench::CompressMeasurements;
-pub use bench::CompressOp;
-pub use bench::CompressResult;
-pub use bench::Compressor;
-pub use bench::DecompressResult;
-pub use bench::benchmark_compress;
-pub use bench::benchmark_decompress;
-pub use bench::calculate_ratios;
-pub use parquet::ParquetCompressor;
-pub use vortex::VortexCompressor;
-
-pub fn chunked_to_vec_record_batch(chunked: ChunkedArray) -> (Vec<RecordBatch>, Arc<Schema>) {
+pub fn chunked_to_vec_record_batch(
+    chunked: ChunkedArray,
+) -> anyhow::Result<(Vec<RecordBatch>, Arc<Schema>)> {
     let chunks_vec = chunked.chunks();
     assert!(!chunks_vec.is_empty(), "empty chunks");
 
@@ -34,10 +24,10 @@ pub fn chunked_to_vec_record_batch(chunked: ChunkedArray) -> (Vec<RecordBatch>,
             // TODO(connor)[ListView]: The rust Parquet implementation does not support writing
             // `ListView` to Parquet files yet.
             let converted_array = recursive_list_from_list_view(array.clone());
-            RecordBatch::try_from(converted_array.as_ref()).unwrap()
+            Ok(RecordBatch::try_from(converted_array.as_ref())?)
         })
-        .collect::<Vec<_>>();
+        .collect::<anyhow::Result<Vec<_>>>()?;
 
     let schema = batches[0].schema();
-    (batches, schema)
+    Ok((batches, schema))
 }

vortex-bench/src/compress/mod.rs benchmarks/compress-bench/src/lib.rsvortex-bench/src/compress/mod.rs renamed to benchmarks/compress-bench/src/lib.rs

@@ -5,6 +5,14 @@ use std::path::PathBuf;
 use std::time::Duration;
 
 use clap::Parser;
+use compress_bench::bench::CompressMeasurements;
+use compress_bench::bench::CompressOp;
+use compress_bench::bench::Compressor;
+use compress_bench::bench::benchmark_compress;
+use compress_bench::bench::benchmark_decompress;
+use compress_bench::bench::calculate_ratios;
+use compress_bench::parquet::ParquetCompressor;
+use compress_bench::vortex::VortexCompressor;
 use indicatif::ProgressBar;
 use itertools::Itertools;
 use regex::Regex;
@@ -17,14 +25,6 @@ use vortex_bench::BenchmarkOutput;
 use vortex_bench::Engine;
 use vortex_bench::Format;
 use vortex_bench::Target;
-use vortex_bench::compress::CompressMeasurements;
-use vortex_bench::compress::CompressOp;
-use vortex_bench::compress::Compressor;
-use vortex_bench::compress::ParquetCompressor;
-use vortex_bench::compress::VortexCompressor;
-use vortex_bench::compress::benchmark_compress;
-use vortex_bench::compress::benchmark_decompress;
-use vortex_bench::compress::calculate_ratios;
 use vortex_bench::datasets::Dataset;
 use vortex_bench::datasets::struct_list_of_ints::StructListOfInts;
 use vortex_bench::datasets::taxi_data::TaxiData;

benchmarks/compress-bench/src/main.rs

@@ -3,8 +3,9 @@
 
 use std::io::Cursor;
 use std::sync::Arc;
+use std::time::Duration;
+use std::time::Instant;
 
-use anyhow::Result;
 use arrow_array::RecordBatch;
 use arrow_schema::Schema;
 use async_trait::async_trait;
@@ -16,10 +17,10 @@ use parquet::basic::ZstdLevel;
 use parquet::file::properties::WriterProperties;
 use vortex::array::Array;
 use vortex::array::arrays::ChunkedVTable;
+use vortex_bench::Format;
 
-use crate::Format;
-use crate::compress::bench::Compressor;
-use crate::compress::chunked_to_vec_record_batch;
+use crate::bench::Compressor;
+use crate::chunked_to_vec_record_batch;
 
 /// Compressor implementation for Parquet format with ZSTD compression.
 pub struct ParquetCompressor {
@@ -50,17 +51,19 @@ impl Compressor for ParquetCompressor {
         Format::Parquet
     }
 
-    async fn compress(&self, array: &dyn Array) -> Result<Bytes> {
+    async fn compress(&self, array: &dyn Array) -> anyhow::Result<(Bytes, Duration)> {
         let chunked = array.as_::<ChunkedVTable>().clone();
-        let (batches, schema) = chunked_to_vec_record_batch(chunked);
+        let (batches, schema) = chunked_to_vec_record_batch(chunked)?;
 
         let mut buf = Vec::new();
-        parquet_compress_write(batches, schema, self.compression, &mut buf);
-        Ok(Bytes::from(buf))
+        let start = Instant::now();
+        parquet_compress_write(batches, schema, self.compression, &mut buf)?;
+        let elapsed = start.elapsed();
+        Ok((Bytes::from(buf), elapsed))
     }
 
-    async fn decompress(&self, data: Bytes) -> Result<usize> {
-        Ok(parquet_decompress_read(data))
+    async fn decompress(&self, data: Bytes) -> anyhow::Result<usize> {
+        parquet_decompress_read(data)
     }
 }
 
@@ -70,28 +73,29 @@ pub fn parquet_compress_write(
     schema: Arc<Schema>,
     compression: Compression,
     buf: &mut Vec<u8>,
-) -> usize {
+) -> anyhow::Result<usize> {
     let mut buf = Cursor::new(buf);
     let writer_properties = WriterProperties::builder()
         .set_compression(compression)
         .build();
-    let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(writer_properties)).unwrap();
+    let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(writer_properties))?;
     for batch in batches {
-        writer.write(&batch).unwrap();
+        writer.write(&batch)?;
     }
-    writer.flush().unwrap();
+    writer.flush()?;
     let n_bytes = writer.bytes_written();
-    writer.close().unwrap();
-    n_bytes
+    writer.close()?;
+    Ok(n_bytes)
 }
 
 #[inline(never)]
-pub fn parquet_decompress_read(buf: Bytes) -> usize {
-    let builder = ParquetRecordBatchReaderBuilder::try_new(buf).unwrap();
-    let reader = builder.build().unwrap();
+pub fn parquet_decompress_read(buf: Bytes) -> anyhow::Result<usize> {
+    let builder = ParquetRecordBatchReaderBuilder::try_new(buf)?;
+    let reader = builder.build()?;
     let mut nbytes = 0;
     for batch in reader {
-        nbytes += batch.unwrap().get_array_memory_size()
+        nbytes += batch?.get_array_memory_size()
     }
-    nbytes
+
+    Ok(nbytes)
 }

vortex-bench/src/compress/parquet.rs benchmarks/compress-bench/src/parquet.rsvortex-bench/src/compress/parquet.rs renamed to benchmarks/compress-bench/src/parquet.rs

@@ -3,6 +3,8 @@
 
 use std::io::Cursor;
 use std::sync::Arc;
+use std::time::Duration;
+use std::time::Instant;
 
 use anyhow::Result;
 use async_trait::async_trait;
@@ -12,10 +14,10 @@ use futures::pin_mut;
 use vortex::array::Array;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::WriteOptionsSessionExt;
+use vortex_bench::Format;
+use vortex_bench::SESSION;
 
-use crate::Format;
-use crate::SESSION;
-use crate::compress::bench::Compressor;
+use crate::bench::Compressor;
 
 /// Compressor implementation for Vortex format.
 pub struct VortexCompressor;
@@ -26,10 +28,12 @@ impl Compressor for VortexCompressor {
         Format::OnDiskVortex
     }
 
-    async fn compress(&self, array: &dyn Array) -> Result<Bytes> {
+    async fn compress(&self, array: &dyn Array) -> Result<(Bytes, Duration)> {
         let mut buf = Vec::new();
+        let start = Instant::now();
         vortex_compress_write(array, &mut buf).await?;
-        Ok(Bytes::from(buf))
+        let elapsed = start.elapsed();
+        Ok((Bytes::from(buf), elapsed))
     }
 
     async fn decompress(&self, data: Bytes) -> Result<usize> {