More work untangaling random access and compression

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

Author: AdamGS

Cargo.lock

@@ -0,0 +1,29 @@
+[package]
+name = "compress-bench"
+description = "Compression benchmarks for Vortex and other formats"
+authors.workspace = true
+categories.workspace = true
+edition.workspace = true
+homepage.workspace = true
+include.workspace = true
+keywords.workspace = true
+license.workspace = true
+readme.workspace = true
+repository.workspace = true
+rust-version.workspace = true
+version.workspace = true
+publish = false
+
+[dependencies]
+anyhow = { workspace = true }
+clap = { workspace = true, features = ["derive"] }
+indicatif = { workspace = true }
+itertools = { workspace = true }
+regex = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
+tracing = { workspace = true }
+vortex = { workspace = true }
+vortex-bench = { workspace = true }
+
+[lints]
+workspace = true

benchmarks/compress-bench/Cargo.toml

@@ -11,8 +11,6 @@ use clap::Parser;
 use indicatif::ProgressBar;
 use itertools::Itertools;
 use regex::Regex;
-use tokio::runtime::Runtime;
-use vortex::array::Array;
 use vortex::array::IntoArray;
 use vortex::array::arrays::ChunkedArray;
 use vortex::array::arrays::ChunkedVTable;
@@ -21,9 +19,14 @@ use vortex::utils::aliases::hash_map::HashMap;
 use vortex_bench::Engine;
 use vortex_bench::Format;
 use vortex_bench::Target;
-use vortex_bench::compress::bench as compress;
-use vortex_bench::compress::bench::CompressMeasurements;
-use vortex_bench::compress::bench::CompressOp;
+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;
@@ -33,8 +36,6 @@ use vortex_bench::display::DisplayFormat;
 use vortex_bench::display::print_measurements_json;
 use vortex_bench::display::render_table;
 use vortex_bench::downloadable_dataset::DownloadableDataset;
-use vortex_bench::measurements::CompressionTimingMeasurement;
-use vortex_bench::measurements::CustomUnitMeasurement;
 use vortex_bench::public_bi::PBI_DATASETS;
 use vortex_bench::public_bi::PBIDataset::Arade;
 use vortex_bench::public_bi::PBIDataset::Bimbo;
@@ -43,7 +44,6 @@ use vortex_bench::public_bi::PBIDataset::Euro2016;
 use vortex_bench::public_bi::PBIDataset::Food;
 use vortex_bench::public_bi::PBIDataset::HashTags;
 use vortex_bench::setup_logging_and_tracing;
-use vortex_bench::utils::new_tokio_runtime;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
@@ -58,8 +58,6 @@ struct Args {
     #[arg(short, long, default_value_t = 5)]
     iterations: usize,
     #[arg(short, long)]
-    threads: Option<usize>,
-    #[arg(short, long)]
     verbose: bool,
     #[arg(
         long,
@@ -77,26 +75,33 @@ struct Args {
     tracing: bool,
 }
 
-fn main() -> anyhow::Result<()> {
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
     setup_logging_and_tracing(args.verbose, args.tracing)?;
 
-    let runtime = new_tokio_runtime(args.threads)?;
-
-    compress(
-        runtime,
+    run_compress(
         args.iterations,
         args.datasets.map(|d| Regex::new(&d)).transpose()?,
         args.formats,
         args.ops,
         args.display_format,
         &args.output_path,
     )
+    .await
+}
+
+/// Get a compressor for the given format.
+fn get_compressor(format: Format) -> Box<dyn Compressor> {
+    match format {
+        Format::OnDiskVortex => Box::new(VortexCompressor),
+        Format::Parquet => Box::new(ParquetCompressor::new()),
+        _ => unimplemented!("Compress bench not implemented for {format}"),
+    }
 }
 
-fn compress(
-    runtime: Runtime,
+async fn run_compress(
     iterations: usize,
     datasets_filter: Option<Regex>,
     formats: Vec<Format>,
@@ -151,21 +156,15 @@ fn compress(
 
     let progress = ProgressBar::new((datasets.len() * formats.len() * ops.len()) as u64);
 
-    let measurements = datasets
-        .into_iter()
-        .map(|dataset_handle| {
-            benchmark_compress(
-                &runtime,
-                &progress,
-                &formats,
-                &ops,
-                iterations,
-                dataset_handle,
-            )
-        })
-        .try_collect::<_, Vec<_>, _>()?
-        .into_iter()
-        .collect::<CompressMeasurements>();
+    let mut measurements = vec![];
+
+    for dataset_handle in datasets.into_iter() {
+        let m = run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle)
+            .await?;
+        measurements.push(m);
+    }
+
+    let measurements = CompressMeasurements::from_iter(measurements);
 
     progress.finish();
 
@@ -195,28 +194,7 @@ fn compress(
     }
 }
 
-// Type aliases for compression and decompression function signatures.
-type CompressFn = fn(
-    &Runtime,
-    &dyn Array,
-    usize,
-    &str,
-) -> anyhow::Result<(
-    Duration,
-    u64,
-    Vec<CustomUnitMeasurement>,
-    CompressionTimingMeasurement,
-)>;
-
-type DecompressFn = fn(
-    &Runtime,
-    &dyn Array,
-    usize,
-    &str,
-) -> anyhow::Result<(Duration, CompressionTimingMeasurement)>;
-
-pub fn benchmark_compress(
-    runtime: &Runtime,
+async fn run_benchmark_for_dataset(
     progress: &ProgressBar,
     formats: &[Format],
     ops: &[CompressOp],
@@ -226,7 +204,7 @@ pub fn benchmark_compress(
     let bench_name = dataset_handle.name();
     tracing::info!("Running {bench_name} benchmark");
 
-    let vx_array = runtime.block_on(async { dataset_handle.to_vortex_array().await })?;
+    let vx_array = dataset_handle.to_vortex_array().await?;
     let uncompressed = ChunkedArray::from_iter(
         vx_array
             .as_::<ChunkedVTable>()
@@ -246,36 +224,33 @@ pub fn benchmark_compress(
     let mut compressed_sizes: HashMap<Format, u64> = HashMap::new();
 
     for format in formats {
+        let compressor = get_compressor(*format);
+
         for op in ops {
             let time = match op {
                 CompressOp::Compress => {
-                    // Select the compression function based on format.
-                    let compress_fn: CompressFn = match format {
-                        Format::OnDiskVortex => compress::benchmark_vortex_compress,
-                        Format::Parquet => compress::benchmark_parquet_compress,
-                        _ => unimplemented!("Compress bench not implemented for {format}"),
-                    };
-
-                    let (time, size, ratios_part, timing) =
-                        compress_fn(runtime, &uncompressed, iterations, bench_name)?;
-                    compressed_sizes.insert(*format, size);
-                    ratios.extend(ratios_part);
-                    timings.push(timing);
-
-                    time
+                    let result = benchmark_compress(
+                        compressor.as_ref(),
+                        &uncompressed,
+                        iterations,
+                        bench_name,
+                    )
+                    .await?;
+                    compressed_sizes.insert(*format, result.compressed_size);
+                    ratios.extend(result.ratios);
+                    timings.push(result.timing);
+                    result.time
                 }
                 CompressOp::Decompress => {
-                    // Select the decompression function based on format.
-                    let decompress_fn: DecompressFn = match format {
-                        Format::OnDiskVortex => compress::benchmark_vortex_decompress,
-                        Format::Parquet => compress::benchmark_parquet_decompress,
-                        _ => unimplemented!("Decompress bench not implemented for {format}"),
-                    };
-
-                    let (time, timing) =
-                        decompress_fn(runtime, &uncompressed, iterations, bench_name)?;
-                    timings.push(timing);
-                    time
+                    let result = benchmark_decompress(
+                        compressor.as_ref(),
+                        &uncompressed,
+                        iterations,
+                        bench_name,
+                    )
+                    .await?;
+                    timings.push(result.timing);
+                    result.time
                 }
             };
 
@@ -285,7 +260,7 @@ pub fn benchmark_compress(
     }
 
     // Calculate cross-format ratios after all measurements.
-    compress::calculate_ratios(
+    calculate_ratios(
         &measurements_map,
         &compressed_sizes,
         bench_name,

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

@@ -6,7 +6,6 @@ use std::sync::Arc;
 
 use clap::Parser;
 use clap::value_parser;
-
 use lance::datafusion::LanceTableProvider;
 use lance::dataset::Dataset;
 use lance::deps::arrow_array::RecordBatch;

benchmarks/lance-bench/src/main.rs

@@ -0,0 +1,26 @@
+[package]
+name = "random-access-bench"
+description = "Random access benchmarks for Vortex and other formats"
+authors.workspace = true
+categories.workspace = true
+edition.workspace = true
+homepage.workspace = true
+include.workspace = true
+keywords.workspace = true
+license.workspace = true
+readme.workspace = true
+repository.workspace = true
+rust-version.workspace = true
+version.workspace = true
+publish = false
+
+[dependencies]
+anyhow = { workspace = true }
+clap = { workspace = true, features = ["derive"] }
+indicatif = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
+vortex = { workspace = true }
+vortex-bench = { workspace = true }
+
+[lints]
+workspace = true