diff --git a/.github/workflows/develop-bench.yml b/.github/workflows/develop-bench.yml index 49d2e131cb3..203d34011d3 100644 --- a/.github/workflows/develop-bench.yml +++ b/.github/workflows/develop-bench.yml @@ -156,7 +156,7 @@ jobs: VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" run: | - python3 scripts/compress-split.py --formats parquet,lance,vortex --emit-ingest-records + python3 scripts/compress-split.py --formats arrow-ipc,parquet,lance,vortex --emit-ingest-records - name: Run ${{ matrix.benchmark.name }} benchmark (per-column-encoder) if: matrix.benchmark.id == 'string-bench' diff --git a/Cargo.lock b/Cargo.lock index 43862759c38..78d5cb92381 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1568,6 +1568,7 @@ version = "0.1.0" dependencies = [ "anyhow", "arrow-array 58.4.0", + "arrow-ipc 58.4.0", "arrow-schema 58.4.0", "async-trait", "bytes", @@ -9684,6 +9685,7 @@ version = "0.1.0" dependencies = [ "anyhow", "arrow-array 58.4.0", + "arrow-ipc 58.4.0", "arrow-schema 58.4.0", "arrow-select 58.4.0", "async-trait", @@ -9718,6 +9720,7 @@ dependencies = [ "sysinfo", "tabled", "target-lexicon", + "tempfile", "tokio", "tokio-stream", "tokio-util", diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 7b1dadc3209..b044dfe1263 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -17,6 +17,7 @@ publish = false [dependencies] anyhow = { workspace = true } arrow-array = { workspace = true } +arrow-ipc = { workspace = true } arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 9f4b0d46624..58c913856b4 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -1,7 +1,8 @@ # Compression benchmark -Measures compression and decompression throughput, plus resulting file sizes, for Vortex -versus Parquet (and optionally Lance) across a range of datasets: NYC taxi data, several +Measures compression and decompression throughput, plus resulting file sizes, for Vortex, +Parquet, uncompressed Arrow IPC, and optionally Lance. +The suite covers NYC taxi data and several [Public BI](https://github.com/cwida/public_bi_benchmark) tables (Arade, Bimbo, CMSprovider, Euro2016, Food, HashTags), TPC-H `l_comment` variants, and synthetic nested data. This is the workload behind the `Compression` PR comment. diff --git a/benchmarks/compress-bench/src/arrow.rs b/benchmarks/compress-bench/src/arrow.rs new file mode 100644 index 00000000000..442b2c52d7d --- /dev/null +++ b/benchmarks/compress-bench/src/arrow.rs @@ -0,0 +1,84 @@ +// 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 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::Compressor; +use vortex_bench::compress::read_projection; + +/// Uncompressed Arrow IPC file baseline. +pub struct ArrowIpcCompressor; + +#[async_trait] +impl Compressor for ArrowIpcCompressor { + fn format(&self) -> Format { + 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::, _>>()?; + + let mut buf = Vec::new(); + let start = Instant::now(); + arrow_file_write(&mut buf, &schema, &batches)?; + let elapsed = start.elapsed(); + Ok((buf.len() as u64, 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::, _>>()?; + + let mut buf = Vec::new(); + arrow_file_write(&mut buf, &schema, &batches)?; + + let start = Instant::now(); + arrow_file_read(Bytes::from(buf), schema.fields().len())?; + Ok(start.elapsed()) + } +} + +#[inline(never)] +fn arrow_file_write( + buf: &mut Vec, + schema: &Schema, + batches: &[RecordBatch], +) -> anyhow::Result<()> { + let mut writer = FileWriter::try_new(buf, schema)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + Ok(()) +} + +#[inline(never)] +fn arrow_file_read(buf: Bytes, root_columns: usize) -> anyhow::Result { + let cursor = Cursor::new(buf); + let projection = read_projection(root_columns).map(<[usize]>::to_vec); + let reader = FileReader::try_new(cursor, projection)?; + + let mut nbytes = 0; + for batch in reader { + nbytes += batch?.get_array_memory_size(); + } + Ok(nbytes) +} diff --git a/benchmarks/compress-bench/src/lib.rs b/benchmarks/compress-bench/src/lib.rs index 25347300667..636a8498912 100644 --- a/benchmarks/compress-bench/src/lib.rs +++ b/benchmarks/compress-bench/src/lib.rs @@ -3,6 +3,7 @@ #[cfg(feature = "lance")] pub use lance_bench::compress::LanceCompressor; +pub mod arrow; pub mod gpu; pub mod parquet; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index d8ef76afc93..589ec5854ae 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -10,10 +10,12 @@ use anyhow::Context; use clap::Parser; #[cfg(feature = "lance")] use compress_bench::LanceCompressor; +use compress_bench::arrow::ArrowIpcCompressor; use compress_bench::gpu::GpuCodec; use compress_bench::gpu::GpuOptions; use compress_bench::gpu::compressor as gpu_compressor; use compress_bench::parquet::ParquetCompressor; +use compress_bench::parquet::arrow_uncompressed_size; use compress_bench::vortex::VortexCompressor; use futures::FutureExt; use indicatif::ProgressBar; @@ -58,7 +60,7 @@ struct Args { long, value_delimiter = ',', value_enum, - default_values_t = vec![Format::Parquet, Format::OnDiskVortex] + default_values_t = vec![Format::ArrowIpc, Format::Parquet, Format::OnDiskVortex] )] formats: Vec, #[arg(short, long, default_value_t = 5)] @@ -188,6 +190,7 @@ fn get_compressor(format: Format, mode: BenchMode) -> Box { } match format { + Format::ArrowIpc => Box::new(ArrowIpcCompressor), Format::OnDiskVortex => Box::new(VortexCompressor), Format::Parquet => Box::new(ParquetCompressor::new()), #[cfg(feature = "lance")] @@ -217,7 +220,11 @@ async fn run_compress( output_path: Option, ingest_output: Option, ) -> anyhow::Result<()> { - let targets = formats + let timing_targets = formats + .iter() + .map(|f| Target::new(Engine::default(), *f)) + .collect_vec(); + let size_targets = formats .iter() .map(|f| Target::new(Engine::default(), *f)) .collect_vec(); @@ -360,15 +367,15 @@ async fn run_compress( // publishes the numbers for the datasets that did decode. match display_format { DisplayFormat::Table => { - render_table(&mut writer, measurements.timings, &targets)?; + render_table(&mut writer, measurements.timings, &timing_targets)?; render_table( &mut writer, - measurements.ratios, - &if formats.contains(&Format::OnDiskVortex) { - vec![Target::new(Engine::default(), Format::OnDiskVortex)] - } else { - vec![] - }, + measurements + .ratios + .into_iter() + .filter(|measurement| measurement.unit == "bytes") + .collect(), + &size_targets, )?; } DisplayFormat::GhJson => { @@ -421,6 +428,11 @@ async fn run_benchmark_for_dataset( // Get the parquet file path for this dataset let parquet_path = dataset_handle.to_parquet_path().await?; + let uncompressed_size = ops + .contains(&CompressOp::Compress) + .then(|| arrow_uncompressed_size(&parquet_path)) + .transpose() + .with_context(|| format!("measuring Arrow memory size for {bench_name}"))?; let mut ratios = Vec::new(); let mut timings = Vec::new(); @@ -460,6 +472,7 @@ async fn run_benchmark_for_dataset( v3_variant, *format, result.compressed_size, + uncompressed_size.context("compression size requires Arrow memory size")?, )); ratios.extend(result.ratios); timings.push(result.timing); diff --git a/benchmarks/compress-bench/src/parquet.rs b/benchmarks/compress-bench/src/parquet.rs index cab4f753948..e2904de2911 100644 --- a/benchmarks/compress-bench/src/parquet.rs +++ b/benchmarks/compress-bench/src/parquet.rs @@ -45,6 +45,23 @@ impl Default for ParquetCompressor { } } +/// Return the Arrow memory size after decoding the input Parquet file. +pub fn arrow_uncompressed_size(parquet_path: &Path) -> anyhow::Result { + let file = File::open(parquet_path)?; + let reader = ParquetRecordBatchReaderBuilder::try_new(file)?.build()?; + let mut total = 0u64; + + for batch in reader { + let batch = batch?; + let batch_size = u64::try_from(batch.get_array_memory_size())?; + total = total + .checked_add(batch_size) + .ok_or_else(|| anyhow::anyhow!("Arrow memory size exceeds u64"))?; + } + + Ok(total) +} + #[async_trait] impl Compressor for ParquetCompressor { fn format(&self) -> Format { diff --git a/benchmarks/datafusion-bench/src/lib.rs b/benchmarks/datafusion-bench/src/lib.rs index aeea953ba3d..4b8656af35a 100644 --- a/benchmarks/datafusion-bench/src/lib.rs +++ b/benchmarks/datafusion-bench/src/lib.rs @@ -112,7 +112,7 @@ pub fn format_to_df_format(format: Format) -> Arc { Format::OnDiskVortex | Format::VortexCompact | Format::VortexSpatialNative => Arc::new( VortexFormat::new_with_options(SESSION.clone(), vortex_table_options()), ), - Format::OnDiskDuckDB | Format::Lance => { + Format::ArrowIpc | Format::OnDiskDuckDB | Format::Lance => { unimplemented!("Format {format} cannot be turned into a DataFusion `FileFormat`") } } diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index 14949c85fcb..ae18cb240b4 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -11,8 +11,8 @@ Two access patterns are generated with a fixed seed (see [`src/main.rs`](./src/m simulating lookups with no locality. Each pattern runs over four datasets (`taxi`, `feature-vectors`, `nested-lists`, -`nested-structs`) in Parquet, Lance, and Vortex, both with a cached open file handle and -reopening the file per lookup. CI drives the full matrix via +`nested-structs`) in Arrow IPC, Parquet, Lance, and Vortex. Each format uses a cached open file +handle and a per-lookup reopen mode. CI drives the full matrix via [`scripts/random-access-split.py`](../../scripts/random-access-split.py). ## Running locally diff --git a/benchmarks/random-access-bench/src/lib.rs b/benchmarks/random-access-bench/src/lib.rs index 11888d34c43..664bbd5216c 100644 --- a/benchmarks/random-access-bench/src/lib.rs +++ b/benchmarks/random-access-bench/src/lib.rs @@ -20,6 +20,7 @@ use vortex_bench::create_output_writer; use vortex_bench::display::DisplayFormat; use vortex_bench::display::print_measurements_json; use vortex_bench::measurements::TimingMeasurement; +use vortex_bench::random_access::ArrowIpcRandomAccessor; use vortex_bench::random_access::BenchDataset; use vortex_bench::random_access::ParquetRandomAccessor; use vortex_bench::random_access::RandomAccessor; @@ -240,6 +241,10 @@ async fn open_accessor( format.ext() ); match format { + Format::ArrowIpc => { + let path = dataset.path(format).await?; + Ok(Box::new(ArrowIpcRandomAccessor::open(path, name)?)) + } Format::OnDiskVortex | Format::VortexCompact => { let path = dataset.path(format).await?; Ok(Box::new( diff --git a/benchmarks/random-access-bench/src/main.rs b/benchmarks/random-access-bench/src/main.rs index 11ef81edd3c..b7f41533c79 100644 --- a/benchmarks/random-access-bench/src/main.rs +++ b/benchmarks/random-access-bench/src/main.rs @@ -49,7 +49,7 @@ struct Args { long, value_delimiter = ',', value_parser = Format::parse_allowed, - default_values = ["parquet", "vortex", "lance"] + default_values = ["arrow-ipc", "parquet", "vortex", "lance"] )] formats: Vec, /// Time limit in seconds for each benchmark target (e.g., 10 for 10 seconds). diff --git a/scripts/compress-split.py b/scripts/compress-split.py index ecfdb262246..e216ff52bcf 100755 --- a/scripts/compress-split.py +++ b/scripts/compress-split.py @@ -83,7 +83,7 @@ def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--formats", - default="parquet,vortex", + default="arrow-ipc,parquet,vortex", help="comma-separated formats to forward to compress-bench", ) parser.add_argument( diff --git a/scripts/post-ingest.py b/scripts/post-ingest.py index 43137e0b8d3..24aea2f1d5c 100755 --- a/scripts/post-ingest.py +++ b/scripts/post-ingest.py @@ -36,7 +36,7 @@ # MUST equal `benchmarks-website/web/lib/schema-version.ts::SCHEMA_VERSION`. # Bumping this is a coordinated change across the website contract, v3.rs, and # this script. -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 def parse_args() -> argparse.Namespace: @@ -187,7 +187,7 @@ def build_commit(sha: str, repo_url: str, git_dir: Path | None) -> dict: frozenset({"dataset_variant", "env_triple"}), ), "compression_size": ( - frozenset({"commit_sha", "dataset", "format", "value_bytes"}), + frozenset({"commit_sha", "dataset", "format", "value_bytes", "uncompressed_bytes"}), frozenset({"dataset_variant"}), ), "random_access_time": ( @@ -388,6 +388,7 @@ def _memory_quartet_consistent(r: dict) -> bool: ("dataset_variant", "opt_str"), ("format", "str"), ("value_bytes", "i64"), + ("uncompressed_bytes", "i64"), ), "random_access_time": ( ("commit_sha", "str"), @@ -581,11 +582,12 @@ def _insert_compression_size(conn, mid_mod, r: dict) -> bool: """ INSERT INTO compression_sizes ( measurement_id, commit_sha, dataset, dataset_variant, - format, value_bytes - ) VALUES (%s, %s, %s, %s, %s, %s) + format, value_bytes, uncompressed_bytes + ) VALUES (%s, %s, %s, %s, %s, %s, %s) ON CONFLICT (measurement_id) DO UPDATE SET - commit_sha = excluded.commit_sha, - value_bytes = excluded.value_bytes + commit_sha = excluded.commit_sha, + value_bytes = excluded.value_bytes, + uncompressed_bytes = excluded.uncompressed_bytes RETURNING (xmax = 0) AS inserted """, ( @@ -595,6 +597,7 @@ def _insert_compression_size(conn, mid_mod, r: dict) -> bool: r.get("dataset_variant"), r["format"], r["value_bytes"], + r["uncompressed_bytes"], ), ) diff --git a/scripts/random-access-split.py b/scripts/random-access-split.py index f8db23c2298..a397d6e518d 100755 --- a/scripts/random-access-split.py +++ b/scripts/random-access-split.py @@ -18,7 +18,7 @@ PARTS_DIR = Path("parts") DATASETS = ["taxi", "feature-vectors", "nested-lists", "nested-structs"] -FORMATS = ["parquet", "lance", "vortex"] +FORMATS = ["arrow-ipc", "parquet", "lance", "vortex"] PATTERNS = ["correlated", "uniform"] OPEN_MODES = ["cached", "reopen"] diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 965a80c6259..5e1298b3411 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -33,6 +33,7 @@ vortex-tensor = { workspace = true } # TODO(connor): In the future, this might b anyhow = { workspace = true } arrow-array = { workspace = true } +arrow-ipc = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } async-trait = { workspace = true } @@ -84,6 +85,7 @@ 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/datasets/feature_vectors.rs b/vortex-bench/src/datasets/feature_vectors.rs index e09f29f49ec..173989d6e83 100644 --- a/vortex-bench/src/datasets/feature_vectors.rs +++ b/vortex-bench/src/datasets/feature_vectors.rs @@ -25,6 +25,7 @@ use crate::conversions::write_parquet_as_vortex; use crate::idempotent_async; use crate::random_access::BenchDataset; use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "feature_vectors"; @@ -46,6 +47,10 @@ impl BenchDataset for FeatureVectorsData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = feature_vectors_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => feature_vectors_vortex().await, Format::VortexCompact => feature_vectors_vortex_compact().await, Format::Parquet => feature_vectors_parquet().await, diff --git a/vortex-bench/src/datasets/nested_lists.rs b/vortex-bench/src/datasets/nested_lists.rs index 3c89ed8eb2b..92e8ae3c3c3 100644 --- a/vortex-bench/src/datasets/nested_lists.rs +++ b/vortex-bench/src/datasets/nested_lists.rs @@ -24,6 +24,7 @@ use crate::conversions::write_parquet_as_vortex; use crate::idempotent_async; use crate::random_access::BenchDataset; use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "nested_lists"; @@ -45,6 +46,10 @@ impl BenchDataset for NestedListsData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = nested_lists_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => nested_lists_vortex().await, Format::VortexCompact => nested_lists_vortex_compact().await, Format::Parquet => nested_lists_parquet().await, diff --git a/vortex-bench/src/datasets/nested_structs.rs b/vortex-bench/src/datasets/nested_structs.rs index c4a522184d3..4f0903661c7 100644 --- a/vortex-bench/src/datasets/nested_structs.rs +++ b/vortex-bench/src/datasets/nested_structs.rs @@ -25,6 +25,7 @@ use crate::conversions::write_parquet_as_vortex; use crate::idempotent_async; use crate::random_access::BenchDataset; use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "nested_structs"; @@ -46,6 +47,10 @@ impl BenchDataset for NestedStructsData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = nested_structs_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => nested_structs_vortex().await, Format::VortexCompact => nested_structs_vortex_compact().await, Format::Parquet => nested_structs_parquet().await, diff --git a/vortex-bench/src/datasets/taxi_data.rs b/vortex-bench/src/datasets/taxi_data.rs index 919a0fa4dff..e592d84aec9 100644 --- a/vortex-bench/src/datasets/taxi_data.rs +++ b/vortex-bench/src/datasets/taxi_data.rs @@ -23,6 +23,8 @@ use crate::datasets::Dataset; use crate::datasets::data_downloads::download_data; use crate::idempotent_async; use crate::random_access::BenchDataset; +use crate::random_access::data_path; +use crate::random_access::parquet_to_arrow_file; /// Dataset identifier used for data path generation. pub const DATASET: &str = "taxi"; @@ -59,6 +61,10 @@ impl BenchDataset for TaxiData { async fn path(&self, format: Format) -> Result { match format { + Format::ArrowIpc => { + let parquet_path = taxi_data_parquet().await?; + parquet_to_arrow_file(parquet_path, data_path(DATASET, Format::ArrowIpc)) + } Format::OnDiskVortex => taxi_data_vortex().await, Format::VortexCompact => taxi_data_vortex_compact().await, Format::Parquet => taxi_data_parquet().await, diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 286d9d41bcc..d68cdc43e93 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -139,6 +139,8 @@ impl Display for Target { #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, ValueEnum, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub enum Format { + #[clap(name = "arrow-ipc")] + ArrowIpc, #[clap(name = "csv")] Csv, #[clap(name = "parquet")] @@ -167,10 +169,15 @@ impl Display for Format { } /// Allowed formats for benchmark CLI arguments. -pub const ALLOWED_FORMATS: &[Format] = &[Format::Parquet, Format::OnDiskVortex, Format::Lance]; +pub const ALLOWED_FORMATS: &[Format] = &[ + Format::ArrowIpc, + Format::Parquet, + Format::OnDiskVortex, + Format::Lance, +]; impl Format { - /// Clap value parser that only accepts parquet, vortex, and lance. + /// Clap value parser that only accepts formats supported by random-access benchmarks. pub fn parse_allowed(s: &str) -> Result { let format = Format::from_str(s, true)?; if ALLOWED_FORMATS.contains(&format) { @@ -186,6 +193,7 @@ impl Format { pub fn name(&self) -> &'static str { match self { + Format::ArrowIpc => "arrow-ipc", Format::Csv => "csv", Format::Parquet => "parquet", Format::OnDiskVortex => "vortex-file-compressed", @@ -198,6 +206,7 @@ impl Format { pub fn ext(&self) -> &'static str { match self { + Format::ArrowIpc => "arrow", Format::Csv => "csv", Format::Parquet => "parquet", Format::OnDiskVortex => "vortex", diff --git a/vortex-bench/src/measurements.rs b/vortex-bench/src/measurements.rs index d6d4ad85b32..9833d2b8bba 100644 --- a/vortex-bench/src/measurements.rs +++ b/vortex-bench/src/measurements.rs @@ -351,11 +351,12 @@ pub struct CompressionTimingMeasurement { impl ToJson for CompressionTimingMeasurement { fn to_json(&self) -> serde_json::Value { let (name, engine) = match self.format { + Format::ArrowIpc => (format!("arrow-ipc {}", self.name), Engine::Vortex), Format::OnDiskVortex => (self.name.to_string(), Engine::Vortex), Format::Parquet => (format!("parquet_rs-zstd {}", self.name), Engine::Vortex), Format::Lance => (format!("lance {}", self.name), Engine::Vortex), _ => vortex_panic!( - "CompressionTimingMeasurement only supports vortex, lance, and parquet formats" + "CompressionTimingMeasurement only supports arrow-ipc, vortex, lance, and parquet formats" ), }; @@ -524,3 +525,23 @@ pub struct MemoryMeasurementJson { pub target: Target, pub env_triple: TripleJson, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compression_timing_json_supports_arrow_ipc() { + let measurement = CompressionTimingMeasurement { + name: "compress time/demo".to_string(), + format: Format::ArrowIpc, + time: Duration::from_nanos(42), + }; + + let json = measurement.to_json(); + + assert_eq!(json["name"], "arrow-ipc compress time/demo"); + assert_eq!(json["target"]["format"], "arrow-ipc"); + assert_eq!(json["value"], 42.0); + } +} diff --git a/vortex-bench/src/random_access/mod.rs b/vortex-bench/src/random_access/mod.rs index 3b9f95f5eaf..26659b3e91c 100644 --- a/vortex-bench/src/random_access/mod.rs +++ b/vortex-bench/src/random_access/mod.rs @@ -1,21 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fs::File; use std::path::PathBuf; +use std::sync::Arc; use anyhow::Result; use arrow_array::RecordBatch; +use arrow_ipc::writer::FileWriter; use async_trait::async_trait; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use vortex::array::ArrayRef; use crate::Format; +use crate::idempotent; pub mod take; // Re-export implementations +pub use take::ArrowIpcRandomAccessor; pub use take::ParquetRandomAccessor; pub use take::VortexRandomAccessor; +pub const ARROW_ROW_OFFSETS_METADATA_KEY: &str = "vortex.row_offsets"; + /// Generate the data path for a random-access benchmark dataset file. /// /// Returns a path like `random_access/{dataset}/{dataset}.{ext}` @@ -28,6 +36,42 @@ pub fn data_path(dataset: &str, format: Format) -> String { } } +/// Convert a Parquet input into an uncompressed Arrow IPC file. +pub fn parquet_to_arrow_file(parquet_path: PathBuf, arrow_path: String) -> Result { + idempotent(&arrow_path, |output_path| { + let input = File::open(parquet_path)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(input)?; + let schema = Arc::clone(builder.schema()); + let reader = builder.build()?; + let output = File::create(output_path)?; + let mut writer = FileWriter::try_new(output, schema.as_ref())?; + let mut row_offsets = vec![0_u64]; + + for batch in reader { + let batch = batch?; + let next_offset = row_offsets + .last() + .copied() + .unwrap_or_default() + .checked_add(u64::try_from(batch.num_rows())?) + .ok_or_else(|| anyhow::anyhow!("Arrow row count overflow"))?; + writer.write(&batch)?; + row_offsets.push(next_offset); + } + + writer.write_metadata( + ARROW_ROW_OFFSETS_METADATA_KEY, + row_offsets + .iter() + .map(u64::to_string) + .collect::>() + .join(","), + ); + writer.finish()?; + Ok(()) + }) +} + /// Trait for a benchmark dataset that knows how to prepare data files. #[async_trait] pub trait BenchDataset: Send + Sync { diff --git a/vortex-bench/src/random_access/take.rs b/vortex-bench/src/random_access/take.rs index fa941d07362..26cd93a9a78 100644 --- a/vortex-bench/src/random_access/take.rs +++ b/vortex-bench/src/random_access/take.rs @@ -1,17 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::BTreeMap; +use std::fs::File as StdFile; use std::iter::once; use std::path::PathBuf; use std::sync::Arc; +use anyhow::Context; use arrow_array::PrimitiveArray; use arrow_array::types::Int64Type; +use arrow_ipc::reader::FileReader; use arrow_select::concat::concat_batches; use arrow_select::take::take_record_batch; use async_trait::async_trait; use futures::stream; use itertools::Itertools; +use parking_lot::Mutex; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::arrow_reader::ArrowReaderMetadata; use parquet::arrow::arrow_reader::ArrowReaderOptions; @@ -30,9 +35,93 @@ use vortex::utils::aliases::hash_map::HashMap; use crate::Format; use crate::SESSION; +use crate::random_access::ARROW_ROW_OFFSETS_METADATA_KEY; use crate::random_access::RandomAccessor; use crate::random_access::RandomAccessorRet; +/// Random accessor for uncompressed Arrow IPC files. +pub struct ArrowIpcRandomAccessor { + name: String, + reader: Mutex>, + row_offsets: Vec, + schema: arrow_schema::SchemaRef, +} + +impl ArrowIpcRandomAccessor { + pub fn open(path: PathBuf, name: impl Into) -> anyhow::Result { + let reader = FileReader::try_new(StdFile::open(path)?, None)?; + let row_offsets = reader + .custom_metadata() + .get(ARROW_ROW_OFFSETS_METADATA_KEY) + .context("Arrow IPC file is missing row-offset metadata")? + .split(',') + .map(str::parse) + .collect::, _>>()?; + anyhow::ensure!( + row_offsets.len() == reader.num_batches() + 1, + "Arrow IPC row-offset metadata does not match the record batches" + ); + anyhow::ensure!( + row_offsets.first() == Some(&0), + "Arrow IPC row offsets must start at zero" + ); + anyhow::ensure!( + row_offsets.windows(2).all(|window| window[0] <= window[1]), + "Arrow IPC row offsets must be sorted" + ); + let schema = reader.schema(); + Ok(Self { + name: name.into(), + reader: Mutex::new(reader), + row_offsets, + schema, + }) + } +} + +#[async_trait] +impl RandomAccessor for ArrowIpcRandomAccessor { + fn format(&self) -> Format { + Format::ArrowIpc + } + + fn name(&self) -> &str { + &self.name + } + + async fn take(&self, indices: &[u64]) -> anyhow::Result { + let mut by_batch = BTreeMap::>::new(); + for &index in indices { + let batch_idx = self.row_offsets.partition_point(|offset| *offset <= index); + anyhow::ensure!( + batch_idx > 0 && batch_idx < self.row_offsets.len(), + "Arrow row index {index} is out of bounds" + ); + let batch_idx = batch_idx - 1; + by_batch + .entry(batch_idx) + .or_default() + .push(i64::try_from(index - self.row_offsets[batch_idx])?); + } + + let mut reader = self.reader.lock(); + let mut batches = Vec::with_capacity(by_batch.len()); + for (batch_idx, local_indices) in by_batch { + reader.set_index(batch_idx)?; + let batch = reader + .next() + .context("Arrow IPC record batch is missing")??; + let indices = PrimitiveArray::::from(local_indices); + batches.push(take_record_batch(&batch, &indices)?); + } + + Ok(RandomAccessorRet::RecordBatch(concat_batches( + &self.schema, + &batches, + )?)) + } +} + /// Random accessor for Vortex format files. /// /// The file handle is opened at construction time and reused across `take()` calls. @@ -190,3 +279,43 @@ impl RandomAccessor for ParquetRandomAccessor { Ok(RandomAccessorRet::RecordBatch(result)) } } + +#[cfg(test)] +mod tests { + use arrow_array::Int64Array; + use arrow_array::RecordBatch; + use arrow_ipc::writer::FileWriter; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + + use super::*; + + #[tokio::test] + async fn arrow_ipc_random_accessor_takes_rows_across_record_batches() -> anyhow::Result<()> { + let file = tempfile::NamedTempFile::new()?; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + { + let mut writer = FileWriter::try_new(file.reopen()?, schema.as_ref())?; + writer.write(&RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![0, 1]))], + )?)?; + writer.write(&RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![2, 3, 4]))], + )?)?; + writer.write_metadata(ARROW_ROW_OFFSETS_METADATA_KEY, "0,2,5"); + writer.finish()?; + } + + let accessor = ArrowIpcRandomAccessor::open(file.path().to_path_buf(), "arrow-ipc")?; + let RandomAccessorRet::RecordBatch(actual) = accessor.take(&[1, 3, 4]).await? else { + anyhow::bail!("Arrow accessor returned a Vortex array") + }; + let expected = + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1, 3, 4]))])?; + assert_eq!(actual, expected); + Ok(()) + } +} diff --git a/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap b/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap index 3a4ed1aa8a5..c026a0dd287 100644 --- a/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap +++ b/vortex-bench/src/snapshots/vortex_bench__v3__tests__snapshot_compression_size.snap @@ -7,5 +7,6 @@ expression: render(&record) "commit_sha": "", "dataset": "taxi", "format": "lance", - "value_bytes": 12345678 + "value_bytes": 12345678, + "uncompressed_bytes": 98765432 } diff --git a/vortex-bench/src/v3.rs b/vortex-bench/src/v3.rs index 735f0f4dc88..1fcaf90b990 100644 --- a/vortex-bench/src/v3.rs +++ b/vortex-bench/src/v3.rs @@ -188,6 +188,8 @@ pub struct CompressionSizeRecord { pub format: String, /// Size in bytes. pub value_bytes: u64, + /// Arrow memory size after decoding the source Parquet file. + pub uncompressed_bytes: u64, } /// A single take-time timing from `random-access-bench`. @@ -401,6 +403,7 @@ pub fn compression_size_record( dataset_variant: Option<&str>, format: Format, value_bytes: u64, + uncompressed_bytes: u64, ) -> V3Record { V3Record::CompressionSize(CompressionSizeRecord { commit_sha: GIT_COMMIT_ID.clone(), @@ -408,6 +411,7 @@ pub fn compression_size_record( dataset_variant: dataset_variant.map(str::to_string), format: format.name().to_string(), value_bytes, + uncompressed_bytes, }) } @@ -613,7 +617,7 @@ mod tests { #[test] fn snapshot_compression_size() -> anyhow::Result<()> { - let record = compression_size_record("taxi", None, Format::Lance, 12_345_678); + let record = compression_size_record("taxi", None, Format::Lance, 12_345_678, 98_765_432); assert_snapshot!(render(&record)?); Ok(()) } @@ -754,7 +758,7 @@ mod tests { }; assert_eq!(time.dataset, "tpc-h l_comment chunked"); - let record = compression_size_record("CMSprovider", None, Format::OnDiskVortex, 42); + let record = compression_size_record("CMSprovider", None, Format::OnDiskVortex, 42, 420); let V3Record::CompressionSize(size) = &record else { panic!("expected CompressionSize variant, got {record:?}"); }; @@ -798,7 +802,7 @@ mod tests { #[test] fn jsonl_round_trips_one_record_per_line() -> anyhow::Result<()> { - let record = compression_size_record("taxi", None, Format::Parquet, 100); + let record = compression_size_record("taxi", None, Format::Parquet, 100, 1_000); let mut buf: Vec = Vec::new(); write_jsonl(&mut buf, &[record.clone(), record])?; let s = String::from_utf8(buf)?;