Restore lance compress benchmarks

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

Author: AdamGS

.github/workflows/ci.yml

@@ -530,7 +530,7 @@ jobs:
       - name: Rust Tests (Windows)
         if: matrix.os == 'windows-x64'
         run: |
-          cargo nextest run --locked --workspace --all-features --no-fail-fast --exclude vortex-bench --exclude vortex-python --exclude vortex-duckdb --exclude vortex-fuzz --exclude duckdb-bench --exclude lance-bench --exclude datafusion-bench
+          cargo nextest run --locked --workspace --all-features --no-fail-fast --exclude vortex-bench --exclude vortex-python --exclude vortex-duckdb --exclude vortex-fuzz --exclude duckdb-bench --exclude lance-bench --exclude datafusion-bench --exclude random-access-bench --exclude compress-bench
       - name: Rust Tests (Other)
         if: matrix.os != 'windows-x64'
         run: cargo nextest run --locked --workspace --all-features --no-fail-fast --exclude vortex-bench --exclude vortex-duckdb

Cargo.lock

@@ -24,6 +24,7 @@ clap = { workspace = true, features = ["derive"] }
 futures = { workspace = true }
 indicatif = { workspace = true }
 itertools = { workspace = true }
+lance-bench = { path = "../lance-bench", optional = true }
 parquet = { workspace = true }
 regex = { workspace = true }
 serde = { workspace = true }
@@ -32,5 +33,8 @@ tracing = { workspace = true }
 vortex = { workspace = true }
 vortex-bench = { workspace = true }
 
+[features]
+lance = ["dep:lance-bench"]
+
 [lints]
 workspace = true

benchmarks/compress-bench/Cargo.toml

@@ -8,7 +8,6 @@ use ::vortex::array::arrays::recursive_list_from_list_view;
 use arrow_array::RecordBatch;
 use arrow_schema::Schema;
 
-pub mod bench;
 pub mod parquet;
 pub mod vortex;
 

benchmarks/compress-bench/src/lib.rs

@@ -5,26 +5,22 @@ 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;
-use vortex::array::IntoArray;
-use vortex::array::arrays::ChunkedArray;
-use vortex::array::arrays::ChunkedVTable;
-use vortex::array::builders::builder_with_capacity;
 use vortex::utils::aliases::hash_map::HashMap;
 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::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;
@@ -95,6 +91,8 @@ fn get_compressor(format: Format) -> Box<dyn Compressor> {
     match format {
         Format::OnDiskVortex => Box::new(VortexCompressor),
         Format::Parquet => Box::new(ParquetCompressor::new()),
+        #[cfg(feature = "lance")]
+        Format::Lance => todo!(),
         _ => unimplemented!("Compress bench not implemented for {format}"),
     }
 }
@@ -202,19 +200,8 @@ async fn run_benchmark_for_dataset(
     let bench_name = dataset_handle.name();
     tracing::info!("Running {bench_name} benchmark");
 
-    let vx_array = dataset_handle.to_vortex_array().await?;
-    let uncompressed = ChunkedArray::from_iter(
-        vx_array
-            .as_::<ChunkedVTable>()
-            .chunks()
-            .iter()
-            .map(|chunk| {
-                let mut builder = builder_with_capacity(chunk.dtype(), chunk.len());
-                chunk.append_to_builder(builder.as_mut());
-                builder.finish()
-            }),
-    )
-    .into_array();
+    // Get the parquet file path for this dataset
+    let parquet_path = dataset_handle.to_parquet_path().await?;
 
     let mut ratios = Vec::new();
     let mut timings = Vec::new();
@@ -229,7 +216,7 @@ async fn run_benchmark_for_dataset(
                 CompressOp::Compress => {
                     let result = benchmark_compress(
                         compressor.as_ref(),
-                        &uncompressed,
+                        &parquet_path,
                         iterations,
                         bench_name,
                     )
@@ -242,7 +229,7 @@ async fn run_benchmark_for_dataset(
                 CompressOp::Decompress => {
                     let result = benchmark_decompress(
                         compressor.as_ref(),
-                        &uncompressed,
+                        &parquet_path,
                         iterations,
                         bench_name,
                     )

benchmarks/compress-bench/src/main.rs

@@ -1,7 +1,9 @@
 // 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;
@@ -15,12 +17,8 @@ use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use parquet::basic::Compression;
 use parquet::basic::ZstdLevel;
 use parquet::file::properties::WriterProperties;
-use vortex::array::Array;
-use vortex::array::arrays::ChunkedVTable;
 use vortex_bench::Format;
-
-use crate::bench::Compressor;
-use crate::chunked_to_vec_record_batch;
+use vortex_bench::compress::Compressor;
 
 /// Compressor implementation for Parquet format with ZSTD compression.
 pub struct ParquetCompressor {
@@ -51,19 +49,35 @@ impl Compressor for ParquetCompressor {
         Format::Parquet
     }
 
-    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)?;
+    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 = builder.schema().clone();
+        let reader = builder.build()?;
+        let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>()?;
 
+        // Compress with our compression settings
         let mut buf = Vec::new();
         let start = Instant::now();
-        parquet_compress_write(batches, schema, self.compression, &mut buf)?;
+        let size = parquet_compress_write(batches, schema, self.compression, &mut buf)?;
         let elapsed = start.elapsed();
-        Ok((Bytes::from(buf), elapsed))
+        Ok((size as u64, elapsed))
     }
 
-    async fn decompress(&self, data: Bytes) -> anyhow::Result<usize> {
-        parquet_decompress_read(data)
+    async fn decompress(&self, parquet_path: &Path) -> anyhow::Result<usize> {
+        // First compress to get the bytes we'll decompress
+        let file = File::open(parquet_path)?;
+        let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
+        let schema = builder.schema().clone();
+        let reader = builder.build()?;
+        let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>()?;
+
+        let mut buf = Vec::new();
+        parquet_compress_write(batches, schema, self.compression, &mut buf)?;
+
+        // Now decompress
+        parquet_decompress_read(Bytes::from(buf))
     }
 }
 

benchmarks/compress-bench/src/parquet.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::io::Cursor;
+use std::path::Path;
 use std::sync::Arc;
 use std::time::Duration;
 use std::time::Instant;
@@ -11,13 +12,12 @@ use async_trait::async_trait;
 use bytes::Bytes;
 use futures::StreamExt;
 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::bench::Compressor;
+use vortex_bench::compress::Compressor;
+use vortex_bench::conversions::parquet_to_vortex;
 
 /// Compressor implementation for Vortex format.
 pub struct VortexCompressor;
@@ -28,15 +28,34 @@ impl Compressor for VortexCompressor {
         Format::OnDiskVortex
     }
 
-    async fn compress(&self, array: &dyn Array) -> Result<(Bytes, Duration)> {
+    async fn compress(&self, parquet_path: &Path) -> Result<(u64, Duration)> {
+        // Read the parquet file as an array stream
+        let array_stream = parquet_to_vortex(parquet_path.to_path_buf())?;
+
         let mut buf = Vec::new();
         let start = Instant::now();
-        vortex_compress_write(array, &mut buf).await?;
+        let mut cursor = Cursor::new(&mut buf);
+        SESSION
+            .write_options()
+            .write(&mut cursor, array_stream)
+            .await?;
         let elapsed = start.elapsed();
-        Ok((Bytes::from(buf), elapsed))
+
+        Ok((buf.len() as u64, elapsed))
     }
 
-    async fn decompress(&self, data: Bytes) -> Result<usize> {
+    async fn decompress(&self, parquet_path: &Path) -> Result<usize> {
+        // First compress to get the bytes we'll decompress
+        let array_stream = parquet_to_vortex(parquet_path.to_path_buf())?;
+        let mut buf = Vec::new();
+        let mut cursor = Cursor::new(&mut buf);
+        SESSION
+            .write_options()
+            .write(&mut cursor, array_stream)
+            .await?;
+
+        // Now decompress
+        let data = Bytes::from(buf);
         let scan = SESSION.open_options().open_buffer(data)?.scan()?;
         let schema = Arc::new(scan.dtype()?.to_arrow_schema()?);
 
@@ -50,13 +69,3 @@ impl Compressor for VortexCompressor {
         Ok(nbytes)
     }
 }
-
-#[inline(never)]
-pub async fn vortex_compress_write(array: &dyn Array, buf: &mut Vec<u8>) -> Result<u64> {
-    let mut cursor = Cursor::new(buf);
-    SESSION
-        .write_options()
-        .write(&mut cursor, array.to_array_stream())
-        .await?;
-    Ok(cursor.position())
-}

benchmarks/compress-bench/src/vortex.rs

@@ -20,7 +20,6 @@ lance-encoding = { version = "0.39.0" }
 
 anyhow = { workspace = true }
 arrow-cast = { workspace = true }
-arrow-schema = { workspace = true }
 async-trait = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
 futures = { workspace = true }