Compare Vortex GPU decompression against a cuDF Parquet read

.github/workflows/pr-bench-gpu-compress.yml

@@ -32,6 +32,24 @@ jobs:
         with:
           repo-token: ${{ secrets.GITHUB_TOKEN }}
           enable-sccache: "true"
+      - name: Install DuckDB
+        # The Public BI datasets are converted from CSV to Parquet through the DuckDB CLI, as
+        # in pr-bench-compress.yml. Without it those datasets cannot build their fixture.
+        run: |
+          wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.5/duckdb_cli-linux-amd64.zip | funzip > duckdb
+          chmod +x duckdb
+          echo "$PWD" >> "$GITHUB_PATH"
+      - name: Install uv
+        uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63  # 0.18.6
+      - name: Install cuDF
+        # The GPU Parquet number is a cuDF `read_parquet`. cuDF ships prebuilt manylinux wheels
+        # on NVIDIA's index, so it stays a runtime dependency and never enters the Rust build.
+        run: |
+          uv venv --python 3.12 .venv-cudf
+          uv pip install --python .venv-cudf \
+            --extra-index-url https://pypi.nvidia.com \
+            cudf-cu12 pandas pyarrow
+          echo "$PWD/.venv-cudf/bin" >> "$GITHUB_PATH"
       - uses: ./.github/actions/system-info
       - name: Display NVIDIA GPU details
         run: |

benchmarks/compress-bench/Cargo.toml

@@ -27,6 +27,8 @@ itertools = { workspace = true }
 lance-bench = { path = "../lance-bench", optional = true }
 parquet = { workspace = true }
 regex = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
 tempfile = { workspace = true, optional = true }
 tokio = { workspace = true, features = ["full"] }
 tracing = { workspace = true }
@@ -45,7 +47,6 @@ name = "compress-bench"
 test = false
 
 [lib]
-test = false
 
 [lints]
 workspace = true

benchmarks/compress-bench/README.md

@@ -15,13 +15,100 @@ See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--format
 cargo run -p compress-bench --profile release_debug
 ```
 
-GPU decompression is opt-in and runs only the existing benchmark names allow-listed in
-`src/main.rs`:
+## GPU decompression
+
+`--gpu-decompress` is opt-in, requires the `cuda` feature, and restricts the suite to the
+GPU dataset list in `src/main.rs`. It measures decompression only, for two backends:
+
+- **Vortex** — the file is written with CUDA-compatible BtrBlocks encodings only
+  (`only_cuda_compatible`) and a CUDA flat layout, then decoded on the device all the way to
+  canonical arrays.
+- **Parquet** — the file is rewritten with GPU-friendly writer settings (see below) and read
+  back with [cuDF](https://github.com/rapidsai/cudf)'s `read_parquet`, which performs the
+  whole read on the device: page header decode, codec decompression, dictionary/RLE/plain
+  decoding and column assembly.
+
+Both sides therefore decode all the way to device-resident arrays, which is what makes the
+`vortex:parquet-<codec> gpu ratio decompress time` metric a like-for-like comparison.
 
 ```bash
 cargo run -p compress-bench --profile release_debug \
   --features cuda,unstable_encodings -- --gpu-decompress
+
+# pick the Parquet page codec the GPU file is written with (default: snappy)
+cargo run -p compress-bench --profile release_debug \
+  --features cuda,unstable_encodings -- --gpu-decompress --gpu-parquet-codec zstd
+```
+
+### cuDF
+
+cuDF is reached through its prebuilt `cudf-cu12` wheel, so it is a runtime dependency of the
+benchmark and never enters the Rust build:
+
+```bash
+uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarrow
+```
+
+`scripts/cudf-parquet-read.py` performs and times the read. Timing is taken inside that
+script, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up
+read runs first for the same reason.
+
+Both backends read a warm file by default. cuDF runs an untimed warm-up read before the timed
+one, so its timed read hits the page cache; the Vortex reader therefore does **not** use direct
+I/O by default, because `O_DIRECT` would bypass the page cache and compare a Vortex read of the
+disk against a cuDF read of RAM. `--gpu-direct-io` turns it back on to measure storage bandwidth
+instead — a different question, and the resulting ratio is not a decode comparison.
+
+The remaining asymmetry is the transfer path: the Vortex reader uses pinned buffers, while cuDF
+does its own host read and host-to-device copy.
+
+### GPU-friendly Parquet writer settings
+
+Set in `src/gpu_writer.rs`:
+
+| Setting | Value | Why |
+| --- | --- | --- |
+| writer version | `PARQUET_1_0` | v1 pages compress the whole page body; v2 pages put uncompressed levels ahead of the compressed values in the same body. |
+| compression | Snappy (default) or Zstd | Snappy is the Parquet default and has the higher device-side throughput. |
+| dictionary | enabled | Keeps the decompressed payload small; the encoding GPU Parquet readers decode fastest. |
+| data page size | 1 MiB | Large enough to amortize per-page setup, small enough to keep every SM fed. Matches the page size cuDF targets. |
+| data page row limit | 1,000,000 | The 20k-row default caps narrow columns' pages far below 1 MiB. |
+| statistics | chunk-level | Page statistics only inflate the headers a reader has to walk. |
+| row group size | 1,048,576 rows | Shared with the Vortex side as `GPU_ROW_GROUP_SIZE` — see below. |
+
+### Matching physical partitions
+
+A Parquet row group and a Vortex chunk are the same thing for this comparison: the unit the
+reader plans and dispatches over. Both formats are pinned to `GPU_ROW_GROUP_SIZE`
+(1,048,576 rows, Parquet's `DEFAULT_MAX_ROW_GROUP_ROW_COUNT`).
+
+Without this the two are not comparable. Parquet reads ~1M-row row groups, while the Vortex
+side inherits the Arrow reader's ~8K-row batches — each of which becomes its own chunk, its own
+compressed blocks and its own kernel launches, so a single dispatch turns into hundreds.
+
+Setting the Arrow reader's batch size alone is not enough: the reader also breaks at the source
+file's row group boundaries, so short batches survive. `parquet_to_vortex_chunks_with_batch_size`
+therefore concatenates the source batches and re-slices them on exact boundaries. Those batches
+are written straight through as root chunks via `ChunkedLayoutStrategy`, and read back with
+`SplitBy::RowCount(GPU_ROW_GROUP_SIZE)` so a scan batch is one whole partition.
+
+### Correctness
+
+`--gpu-verify` cross-checks device output against the CPU decoders on every iteration:
+
+- Parquet: the cuDF-read frame is compared against a CPU Parquet read of the same file.
+- Vortex: each GPU-decoded field is copied back and compared against the same field decoded
+  on the CPU, through Arrow with a pinned target type.
+
+Verification runs inline, so timings from a verifying run are not comparable to a plain one —
+run it as its own pass:
+
+```bash
+cargo run -p compress-bench --profile release_debug \
+  --features cuda,unstable_encodings -- --gpu-decompress --gpu-verify --iterations 1
 ```
 
-On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure
-storage bandwidth rather than page-cache hits.
+Any `--gpu-decompress` run reports on every dataset rather than stopping at the first failure, so
+one run shows which datasets decode correctly on the GPU and which do not. The timing tables are
+rendered before the failure summary, so a dataset the GPU cannot decode still leaves the rest of
+the matrix with numbers — the process exits non-zero either way.

benchmarks/compress-bench/src/gpu_parquet.rs

@@ -0,0 +1,140 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! GPU Parquet decompression backend, timed through cuDF.
+//!
+//! cuDF's `read_parquet` performs the whole read on the device — page header decode,
+//! codec decompression, dictionary/RLE/plain decoding and column assembly — which makes it
+//! the like-for-like opponent for the Vortex GPU backend, which likewise decodes all the way
+//! to canonical arrays on device.
+//!
+//! cuDF is reached through its prebuilt `cudf-cu12` wheel rather than by linking libcudf, so
+//! it stays a runtime dependency of this benchmark and never enters the Rust build. The
+//! measurement is taken inside [`CUDF_SCRIPT`], so interpreter start, `import cudf` and CUDA
+//! context creation are excluded; only the reads themselves are timed.
+
+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::Compressor;
+
+use crate::gpu_writer::GpuCodec;
+use crate::gpu_writer::gpu_writer_properties;
+
+/// Repo-relative path of the script that performs and times the cuDF read.
+const CUDF_SCRIPT: &str = "scripts/cudf-parquet-read.py";
+
+/// Parquet compressor whose decompression measurement is a full cuDF GPU read.
+pub struct GpuParquetCompressor {
+    codec: GpuCodec,
+    verify: bool,
+}
+
+/// What the cuDF script reports back.
+#[derive(Debug, Deserialize)]
+struct CudfReadReport {
+    /// Fastest timed read, in nanoseconds.
+    min_ns: u64,
+    rows: u64,
+    columns: u64,
+}
+
+impl GpuParquetCompressor {
+    /// Create a backend that writes pages with `codec` and times cuDF reading them back.
+    ///
+    /// When `verify` is set, the GPU read is cross-checked against a CPU Parquet read of the
+    /// same file before the measurement is reported.
+    pub fn new(codec: GpuCodec, verify: bool) -> Self {
+        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<_>, _>>()?;
+
+        let output = NamedTempFile::new()?;
+        let mut writer = ArrowWriter::try_new(
+            output.reopen()?,
+            schema,
+            Some(gpu_writer_properties(self.codec)),
+        )?;
+        for batch in batches {
+            writer.write(&batch)?;
+        }
+        writer.flush()?;
+        let size = writer.bytes_written() as u64;
+        writer.close()?;
+        Ok((output, size))
+    }
+}
+
+#[async_trait]
+impl Compressor for GpuParquetCompressor {
+    fn format(&self) -> Format {
+        Format::Parquet
+    }
+
+    async fn compress(&self, parquet_path: &Path) -> Result<(u64, Duration)> {
+        let start = Instant::now();
+        let (_file, size) = self.write_gpu_parquet(parquet_path)?;
+        Ok((size, start.elapsed()))
+    }
+
+    async fn decompress(&self, parquet_path: &Path) -> Result<Duration> {
+        let (gpu_file, _) = self.write_gpu_parquet(parquet_path)?;
+        let report = run_cudf_read(gpu_file.path(), self.verify)?;
+
+        ensure!(
+            report.rows > 0 && report.columns > 0,
+            "cuDF read {} rows and {} columns, expected a non-empty table",
+            report.rows,
+            report.columns
+        );
+
+        Ok(Duration::from_nanos(report.min_ns))
+    }
+}
+
+/// Runs the cuDF read script and returns the timing it measured.
+fn run_cudf_read(path: &Path, verify: bool) -> Result<CudfReadReport> {
+    let mut command = Command::new("python3");
+    command.arg(CUDF_SCRIPT).arg(path);
+    if verify {
+        command.arg("--verify");
+    }
+
+    let output = command.output().with_context(|| {
+        format!("failed to run {CUDF_SCRIPT}; is cudf-cu12 installed on this host?")
+    })?;
+
+    if !output.status.success() {
+        bail!(
+            "{CUDF_SCRIPT} exited with {}:\n{}",
+            output.status,
+            String::from_utf8_lossy(&output.stderr).trim()
+        );
+    }
+
+    serde_json::from_slice(&output.stdout).with_context(|| {
+        format!(
+            "could not parse the report from {CUDF_SCRIPT}: {}",
+            String::from_utf8_lossy(&output.stdout).trim()
+        )
+    })
+}