From 37af5b06bf27568475ee57648ed728718c2b83a6 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 12:02:29 +0000 Subject: [PATCH 01/16] Add a GPU Parquet decompression backend to the compression benchmark The GPU compression benchmark only measured Vortex, and only on a single dataset, so it could not say anything about how Vortex GPU decompression compares to Parquet, nor about encodings beyond FSST strings. Parquet compresses each page body independently, which is exactly the batch shape nvCOMP's device decompressors take and how cuDF's Parquet reader gets pages off the CPU. This adds a Parquet backend built on that: column chunks are staged on the device through the same pinned, direct-I/O reader the Vortex backend uses, then every page in a row group is decompressed in one batched nvCOMP launch. - vortex-nvcomp: bind the batched Snappy decompression entrypoints and the per-algorithm alignment queries, and share `DecompressBackend` between the Snappy and Zstd wrappers. - compress-bench: locate compressed page bodies by walking the per-page Thrift headers (`parquet::format::PageHeader` is deprecated and `parquet`'s own parser is crate-private), and write files with GPU-friendly settings: v1 pages, dictionary encoding, 1 MiB pages, Snappy by default. - Run both Vortex and Parquet under `--gpu-decompress`, and expand the GPU dataset set from one to nine so ALP, bit-packed, run-end, date/time-parts and null-heavy columns are covered alongside FSST strings. - Add `--gpu-verify`, which compares every GPU-decompressed page against the host codec and every GPU-decoded Vortex field against the CPU decode, and run it as a CI step before the timed benchmark. Independently of that flag, nvCOMP's per-page status and size arrays are checked on every iteration. Page decoding is not part of the Parquet measurement, so its numbers are an upper bound on a full GPU Parquet reader; the README states this. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 10 + Cargo.lock | 4 + Cargo.toml | 1 + benchmarks/compress-bench/Cargo.toml | 10 +- benchmarks/compress-bench/README.md | 73 ++- benchmarks/compress-bench/src/gpu_parquet.rs | 551 ++++++++++++++++++ benchmarks/compress-bench/src/gpu_vortex.rs | 47 +- benchmarks/compress-bench/src/lib.rs | 3 + benchmarks/compress-bench/src/main.rs | 164 ++++-- .../compress-bench/src/parquet_pages.rs | 530 +++++++++++++++++ vortex-cuda/nvcomp/build.rs | 7 + vortex-cuda/nvcomp/src/backend.rs | 52 ++ vortex-cuda/nvcomp/src/lib.rs | 4 + vortex-cuda/nvcomp/src/snappy.rs | 222 +++++++ vortex-cuda/nvcomp/src/zstd.rs | 47 +- 15 files changed, 1655 insertions(+), 70 deletions(-) create mode 100644 benchmarks/compress-bench/src/gpu_parquet.rs create mode 100644 benchmarks/compress-bench/src/parquet_pages.rs create mode 100644 vortex-cuda/nvcomp/src/backend.rs create mode 100644 vortex-cuda/nvcomp/src/snappy.rs diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 3abb5b54533..45bc59fb317 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -44,6 +44,16 @@ jobs: cargo build --locked --package compress-bench --profile release_debug --features cuda,unstable_encodings - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh + - name: Verify GPU decompression correctness + shell: bash + env: + RUST_BACKTRACE: full + FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + # Cross-checks every GPU-decompressed page and field against the CPU decoders before + # any timings are taken. Verification runs inline, so this pass is not timed. + run: | + target/release_debug/compress-bench \ + --gpu-decompress --gpu-verify --iterations 1 -d table - name: Run GPU compression benchmark shell: bash env: diff --git a/Cargo.lock b/Cargo.lock index c764cd8eab7..e7c6dd2dfdb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1566,12 +1566,15 @@ dependencies = [ "async-trait", "bytes", "clap", + "cudarc", "futures", "indicatif", "itertools 0.14.0", "lance-bench", "parquet 58.4.0", "regex", + "rstest", + "snap", "tempfile", "tokio", "tracing", @@ -1579,6 +1582,7 @@ dependencies = [ "vortex-arrow", "vortex-bench", "vortex-cuda", + "zstd", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 997b3281d25..edf730760b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -245,6 +245,7 @@ similar = "3.0.0" sketches-ddsketch = "0.4.0" smallvec = "1.15.1" smol = "2.0.2" +snap = "1.1.2" spatialbench = "0.2" spatialbench-arrow = "0.2" # spatialbench still pins arrow 56, two majors behind the workspace arrow. Until upstream diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 4046a12d42e..c20e2861b2d 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -21,12 +21,14 @@ arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } clap = { workspace = true, features = ["derive"] } +cudarc = { workspace = true, optional = true } futures = { workspace = true } indicatif = { workspace = true } itertools = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { workspace = true } +snap = { workspace = true } tempfile = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } @@ -34,9 +36,14 @@ vortex = { workspace = true } vortex-arrow = { workspace = true } vortex-bench = { workspace = true } vortex-cuda = { workspace = true, optional = true } +zstd = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +tempfile = { workspace = true } [features] -cuda = ["dep:tempfile", "dep:vortex-cuda"] +cuda = ["dep:cudarc", "dep:tempfile", "dep:vortex-cuda"] lance = ["dep:lance-bench"] unstable_encodings = ["vortex/unstable_encodings", "vortex-cuda?/unstable_encodings"] @@ -45,7 +52,6 @@ name = "compress-bench" test = false [lib] -test = false [lints] workspace = true diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index bf2d3efc1db..132ce153450 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -15,13 +15,78 @@ 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 every + page body is decompressed on the device with nvCOMP's batched Snappy or Zstd entrypoints, + the same decomposition cuDF's Parquet reader uses: column chunks are staged on the device, + then all pages go through a single batched launch. ```bash cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress + +# pick the Parquet page codec (default: snappy) +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress --gpu-parquet-codec zstd ``` -On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure -storage bandwidth rather than page-cache hits. +On Linux both backends read through the same pinned, direct-I/O (`O_DIRECT`) reader, so +repeated iterations measure storage bandwidth rather than page-cache hits. + +### What the Parquet GPU number does and does not include + +Included: column chunk I/O, the host-to-device transfer, and the batched codec launch. + +Not included: page *decoding* — the dictionary, RLE and plain decoders that turn a +decompressed page into an Arrow array — because there is no Rust GPU Parquet page decoder to +call. The Vortex backend it is compared against decodes all the way to canonical arrays, so +the Parquet figure is an upper bound on what a full GPU Parquet reader could reach, and the +`vortex:parquet- gpu ratio decompress time` metric is biased in Parquet's favour. + +Walking the per-page Thrift headers also happens on the host, once per file, outside the +measurement; a real GPU Parquet reader decodes page headers on the device. + +### GPU-friendly Parquet writer settings + +Set in `src/parquet_pages.rs`: + +| Setting | Value | Why | +| --- | --- | --- | +| writer version | `PARQUET_1_0` | v1 pages compress the whole page body, which is the unit nvCOMP decompresses. v2 pages put uncompressed levels ahead of the compressed values in the same body. | +| compression | Snappy (default) or Zstd | The two Parquet codecs nvCOMP implements. Snappy has the higher device throughput and is the Parquet default. | +| dictionary | enabled | Keeps the decompressed payload small; the encoding GPU Parquet readers decode fastest. | +| data page size | 1 MiB | Large enough to amortize per-chunk 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 that have to be walked on the host. | + +### Correctness + +`--gpu-verify` cross-checks device output against the CPU decoders on every iteration: + +- Parquet: each decompressed page is copied back and compared byte-for-byte against the host + Snappy/Zstd output for the same compressed bytes. +- 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 +``` + +Independently of `--gpu-verify`, every Parquet GPU run checks nvCOMP's per-page status and +output-size arrays after the measurement, because a batched launch reports per-page failures +in device memory rather than by failing the launch. + +Page-header scanning is covered by CPU-only unit tests in `src/parquet_pages.rs`, which +assert the located page bodies decompress to exactly the bytes `parquet`'s own page reader +produces. diff --git a/benchmarks/compress-bench/src/gpu_parquet.rs b/benchmarks/compress-bench/src/gpu_parquet.rs new file mode 100644 index 00000000000..b8e6195d784 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_parquet.rs @@ -0,0 +1,551 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! GPU Parquet decompression backend. +//! +//! Parquet's compressed unit is the page, and every page body is an independent block of +//! Snappy or Zstd. That is precisely the batch shape nvCOMP's device decompressors take, and +//! it is how cuDF's Parquet reader gets pages off the CPU: read the column chunks to the +//! device, then decompress every page in one batched launch. +//! +//! What this measures is the decompression stage of a Parquet read — column chunk I/O, +//! host-to-device transfer, and the batched codec launch. Page *decoding* (the dictionary, +//! RLE and plain decoders that turn a decompressed page into an Arrow array) is not +//! included, because there is no Rust GPU Parquet page decoder to call. The Vortex GPU +//! backend it is compared against decodes all the way to canonical arrays, so the Parquet +//! numbers here are an upper bound on what a full GPU Parquet reader could achieve, and the +//! comparison is favourable to Parquet. + +use std::fs::File; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use anyhow::anyhow; +use anyhow::ensure; +use arrow_array::RecordBatch; +use async_trait::async_trait; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::metadata::ParquetMetaDataReader; +use tempfile::NamedTempFile; +use vortex::array::buffer::BufferHandle; +use vortex::array::buffer::DeviceBuffer; +use vortex::buffer::Alignment; +use vortex::buffer::Buffer; +use vortex::error::vortex_err; +use vortex::io::VortexReadAt; +use vortex::io::session::RuntimeSessionExt; +use vortex_bench::Format; +use vortex_bench::SESSION; +use vortex_bench::compress::Compressor; +use vortex_cuda::CudaBufferExt; +use vortex_cuda::CudaDeviceBuffer; +use vortex_cuda::CudaExecutionCtx; +use vortex_cuda::CudaSession; +use vortex_cuda::CudaSessionExt; +use vortex_cuda::PooledFileReadAt; +use vortex_cuda::PooledFileReadAtOptions; +use vortex_cuda::nvcomp::AlignmentRequirements; +use vortex_cuda::nvcomp::snappy; +use vortex_cuda::nvcomp::sys; +use vortex_cuda::nvcomp::sys::nvcompStatus_t; +use vortex_cuda::nvcomp::zstd as nvcomp_zstd; + +use crate::parquet_pages::ColumnChunkPages; +use crate::parquet_pages::GpuCodec; +use crate::parquet_pages::gpu_writer_properties; +use crate::parquet_pages::scan_compressed_pages; + +/// Parquet compressor whose decompression measurement runs the page codec on the GPU. +pub struct GpuParquetCompressor { + codec: GpuCodec, + verify: bool, +} + +impl GpuParquetCompressor { + /// Create a backend that writes pages with `codec` and decompresses them with nvCOMP. + /// + /// When `verify` is set, every decompressed page is copied back and compared against the + /// host codec's output 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(File::open(parquet_path)?)?; + let schema = Arc::clone(builder.schema()); + let batches: Vec = builder.build()?.collect::, _>>()?; + + 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)) + } + + fn check_codec(&self, metadata: &ParquetMetaData) -> Result<()> { + for row_group in metadata.row_groups() { + for column in row_group.columns() { + ensure!( + self.codec.matches(column.compression()), + "column {} was written with {:?}, expected {}", + column.column_path(), + column.compression(), + self.codec.name() + ); + } + } + Ok(()) + } +} + +#[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 { + let (gpu_file, _) = self.write_gpu_parquet(parquet_path)?; + + let file = File::open(gpu_file.path())?; + let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; + self.check_codec(&metadata)?; + drop(file); + + // The page map is a property of the file, so it is resolved once rather than on every + // iteration. A GPU Parquet reader decodes page headers on the device as part of the + // read; this benchmark cannot, so the host walk is kept out of the measurement. + let file_bytes = std::fs::read(gpu_file.path())?; + let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; + let plan = DecompressPlan::build(&row_groups, self.codec)?; + // Only verification needs the host copy afterwards, to decompress the same bytes. + let file_bytes = self.verify.then_some(file_bytes); + + let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + let reader = open_reader(gpu_file.path(), &cuda_ctx)?; + + // Row groups are staged and released one at a time so device residency stays bounded + // by the largest row group rather than by the whole file, matching how the Vortex + // backend streams one batch at a time. + let mut elapsed = Duration::ZERO; + let mut checks = Vec::with_capacity(plan.row_groups.len()); + for row_group in &plan.row_groups { + let start = Instant::now(); + let device_chunks = stage_row_group(row_group, &reader).await?; + let output = + decompress_pages(row_group, &device_chunks, self.codec, &mut cuda_ctx).await?; + cuda_ctx.synchronize_stream()?; + elapsed += start.elapsed(); + + let DecompressOutput { + output, + actual_sizes, + statuses, + } = output; + drop(device_chunks); + + if self.verify { + // Verified inline, since holding every row group's decompressed output would + // defeat the point of streaming them. This is why a verifying run's timings + // are not comparable. + let file_bytes = file_bytes + .as_deref() + .ok_or_else(|| anyhow!("verification requires the host file bytes"))?; + verify_against_host(row_group, file_bytes, self.codec, output).await?; + } + checks.push((actual_sizes, statuses)); + } + + // Checked outside the measurement, but on every iteration: a batched launch reports + // per-page failures in device memory rather than by failing the launch, so without + // this a broken run would simply look fast. + for (row_group, (actual_sizes, statuses)) in plan.row_groups.iter().zip(checks) { + check_statuses(row_group, actual_sizes, statuses).await?; + } + + Ok(elapsed) + } +} + +/// Stages a row group's column chunks on the device, bounding concurrent pinned staging. +async fn stage_row_group( + row_group: &RowGroupPlan, + reader: &PooledFileReadAt, +) -> Result> { + let reads = row_group + .chunks + .iter() + .map(|chunk| reader.read_at(chunk.offset, chunk.len, Alignment::of::())) + .collect::>(); + + stream::iter(reads) + .buffered(STAGING_CONCURRENCY) + .try_collect::>() + .await + .map_err(|e| anyhow!("failed to stage Parquet column chunks on the device: {e}")) +} + +/// Column chunks staged on the device at once. Matches `vortex-cuda`'s file read concurrency. +const STAGING_CONCURRENCY: usize = 32; + +/// Opens the GPU Parquet file through the same pinned, direct-I/O reader the Vortex GPU +/// backend uses, so both formats measure storage bandwidth rather than page-cache hits. +fn open_reader(path: &Path, cuda_ctx: &CudaExecutionCtx) -> Result { + let pool = Arc::clone(SESSION.cuda_session().pinned_buffer_pool()); + let options = PooledFileReadAtOptions::default(); + #[cfg(target_os = "linux")] + let options = options.with_direct_io(); + + Ok(PooledFileReadAt::open_with_options( + path, + SESSION.handle(), + pool, + cuda_ctx.stream().clone(), + options, + )?) +} + +/// A page to decompress, addressed relative to its column chunk's device buffer. +struct PlannedPage { + chunk: usize, + offset_in_chunk: usize, + compressed_len: usize, + uncompressed_len: usize, + output_offset: usize, +} + +/// The byte range of a column chunk to stage on the device. +struct PlannedChunk { + offset: u64, + len: usize, +} + +/// Everything needed to issue one batched nvCOMP launch over a row group's pages. +struct RowGroupPlan { + chunks: Vec, + pages: Vec, + output_len: usize, + max_uncompressed: usize, +} + +/// The per-row-group work for one file. +struct DecompressPlan { + row_groups: Vec, +} + +impl DecompressPlan { + fn build(row_groups: &[Vec], codec: GpuCodec) -> Result { + let alignment = decompress_alignments(codec)?; + ensure!( + alignment.output.is_power_of_two(), + "nvcomp reported a non-power-of-two output alignment of {}", + alignment.output + ); + + let row_groups = row_groups + .iter() + .map(|chunks| RowGroupPlan::build(chunks, codec, alignment)) + .collect::>>()?; + + ensure!( + row_groups.iter().any(|plan| !plan.pages.is_empty()), + "Parquet file contains no compressed pages" + ); + + Ok(Self { row_groups }) + } +} + +impl RowGroupPlan { + fn build( + chunks: &[ColumnChunkPages], + codec: GpuCodec, + alignment: AlignmentRequirements, + ) -> Result { + let mut planned_chunks = Vec::with_capacity(chunks.len()); + let mut pages = Vec::new(); + let mut output_len = 0usize; + let mut max_uncompressed = 0usize; + + for (index, chunk) in chunks.iter().enumerate() { + planned_chunks.push(PlannedChunk { + offset: chunk.offset, + len: chunk.len, + }); + + for page in &chunk.pages { + let offset_in_chunk = usize::try_from( + u64::try_from(page.offset)? + .checked_sub(chunk.offset) + .ok_or_else(|| anyhow!("page offset precedes its column chunk"))?, + )?; + // Pages are decompressed in place from their column chunk's device buffer. + // CUDA allocations are at least 256-byte aligned, so a page's device address + // meets nvcomp's requirement exactly when its chunk-relative offset does. + ensure!( + offset_in_chunk.is_multiple_of(alignment.input), + "page at file offset {} sits {offset_in_chunk} bytes into its column chunk, \ + which does not meet nvcomp's {} byte input alignment for {}; \ + use --gpu-parquet-codec snappy", + page.offset, + alignment.input, + codec.name() + ); + + let output_offset = output_len.next_multiple_of(alignment.output); + output_len = output_offset + page.uncompressed_len; + max_uncompressed = max_uncompressed.max(page.uncompressed_len); + + pages.push(PlannedPage { + chunk: index, + offset_in_chunk, + compressed_len: page.compressed_len, + uncompressed_len: page.uncompressed_len, + output_offset, + }); + } + } + + Ok(Self { + chunks: planned_chunks, + pages, + output_len, + max_uncompressed, + }) + } +} + +fn decompress_alignments(codec: GpuCodec) -> Result { + match codec { + GpuCodec::Snappy => { + snappy::decompress_alignment_requirements(snappy::SnappyDecompressOpts::default()) + } + GpuCodec::Zstd => nvcomp_zstd::decompress_alignment_requirements( + nvcomp_zstd::ZstdDecompressOpts::default(), + ), + } + .map_err(|e| anyhow!("nvcomp alignment query failed: {e}")) +} + +/// Device-side outputs of a batched decompression launch. +struct DecompressOutput { + output: CudaSlice, + actual_sizes: CudaSlice, + statuses: CudaSlice, +} + +/// Enqueues the batched decompression of every page in `plan` onto the context's stream. +async fn decompress_pages( + plan: &RowGroupPlan, + device_chunks: &[BufferHandle], + codec: GpuCodec, + ctx: &mut CudaExecutionCtx, +) -> Result { + let num_pages = plan.pages.len(); + + let temp_size = match codec { + GpuCodec::Snappy => { + snappy::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) + } + GpuCodec::Zstd => { + nvcomp_zstd::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) + } + } + .map_err(|e| anyhow!("nvcomp temp size query failed: {e}"))?; + + let chunk_bases = device_chunks + .iter() + .map(|handle| handle.cuda_device_ptr()) + .collect::, _>>()?; + + let mut output = ctx.device_alloc::(plan.output_len)?; + // Only the allocation address is needed to build the output pointer table; the device + // write itself is tracked by the guard taken around the launch below. + let output_base = { + let (base, _) = output.device_ptr(ctx.stream()); + base + }; + + let mut compressed_ptrs = Vec::with_capacity(num_pages); + let mut compressed_sizes = Vec::with_capacity(num_pages); + let mut uncompressed_sizes = Vec::with_capacity(num_pages); + let mut output_ptrs = Vec::with_capacity(num_pages); + for page in &plan.pages { + compressed_ptrs.push(chunk_bases[page.chunk] + page.offset_in_chunk as u64); + compressed_sizes.push(page.compressed_len); + uncompressed_sizes.push(page.uncompressed_len); + output_ptrs.push(output_base + page.output_offset as u64); + } + + let (compressed_ptrs, compressed_sizes, uncompressed_sizes, output_ptrs) = futures::try_join!( + ctx.copy_to_device(compressed_ptrs)?, + ctx.copy_to_device(compressed_sizes)?, + ctx.copy_to_device(uncompressed_sizes)?, + ctx.copy_to_device(output_ptrs)? + )?; + + let mut actual_sizes: CudaSlice = ctx.device_alloc(num_pages)?; + let mut statuses: CudaSlice = ctx.device_alloc(num_pages)?; + let mut temp: CudaSlice = ctx.device_alloc(temp_size)?; + + let stream = ctx.stream(); + let compressed_ptrs_view = compressed_ptrs.cuda_view::()?; + let compressed_sizes_view = compressed_sizes.cuda_view::()?; + let uncompressed_sizes_view = uncompressed_sizes.cuda_view::()?; + let output_ptrs_view = output_ptrs.cuda_view::()?; + + let (compressed_ptrs_ptr, record_compressed_ptrs) = compressed_ptrs_view.device_ptr(stream); + let (compressed_sizes_ptr, record_compressed_sizes) = compressed_sizes_view.device_ptr(stream); + let (uncompressed_sizes_ptr, record_uncompressed_sizes) = + uncompressed_sizes_view.device_ptr(stream); + let (output_ptrs_ptr, record_output_ptrs) = output_ptrs_view.device_ptr(stream); + let (_output_ptr, record_output) = output.device_ptr_mut(stream); + let (actual_sizes_ptr, record_actual_sizes) = actual_sizes.device_ptr_mut(stream); + let (statuses_ptr, record_statuses) = statuses.device_ptr_mut(stream); + let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); + + ctx.launch_external(plan.output_len, || { + // SAFETY: every pointer is derived from a live device allocation sized by the plan, + // and each batch metadata array holds exactly `num_pages` entries. + unsafe { + match codec { + GpuCodec::Snappy => snappy::decompress_async( + compressed_ptrs_ptr as _, + compressed_sizes_ptr as _, + uncompressed_sizes_ptr as _, + actual_sizes_ptr as _, + num_pages, + temp_ptr as _, + temp_size, + output_ptrs_ptr as _, + statuses_ptr as _, + stream.cu_stream().cast(), + ), + GpuCodec::Zstd => nvcomp_zstd::decompress_async( + compressed_ptrs_ptr as _, + compressed_sizes_ptr as _, + uncompressed_sizes_ptr as _, + actual_sizes_ptr as _, + num_pages, + temp_ptr as _, + temp_size, + output_ptrs_ptr as _, + statuses_ptr as _, + stream.cu_stream().cast(), + ), + } + .map_err(|e| vortex_err!("nvcomp decompress_async failed: {}", e)) + } + })?; + + drop(( + record_compressed_ptrs, + record_compressed_sizes, + record_uncompressed_sizes, + record_output_ptrs, + record_output, + record_actual_sizes, + record_statuses, + record_temp, + )); + // The temporary workspace must outlive the launch, which the stream ordering guarantees + // only while the allocation is alive. + drop(temp); + + Ok(DecompressOutput { + output, + actual_sizes, + statuses, + }) +} + +/// Copies the per-page status and size arrays back and fails on any mismatch. +async fn check_statuses( + plan: &RowGroupPlan, + actual_sizes: CudaSlice, + statuses: CudaSlice, +) -> Result<()> { + let statuses = CudaDeviceBuffer::new(statuses) + .copy_to_host(Alignment::of::())? + .await?; + let actual_sizes = CudaDeviceBuffer::new(actual_sizes) + .copy_to_host(Alignment::of::())? + .await?; + + let statuses = Buffer::::from_byte_buffer(statuses); + let actual_sizes = Buffer::::from_byte_buffer(actual_sizes); + + for (index, page) in plan.pages.iter().enumerate() { + let status = statuses.as_slice()[index]; + ensure!( + status == sys::nvcompStatus_t_nvcompSuccess, + "page {index} failed to decompress with nvcomp status {status}" + ); + let actual = actual_sizes.as_slice()[index]; + ensure!( + actual == page.uncompressed_len, + "page {index} decompressed to {actual} bytes, expected {}", + page.uncompressed_len + ); + } + + Ok(()) +} + +/// Compares every decompressed page against the host codec's output for the same bytes. +async fn verify_against_host( + plan: &RowGroupPlan, + file_bytes: &[u8], + codec: GpuCodec, + output: CudaSlice, +) -> Result<()> { + let device_output = CudaDeviceBuffer::new(output) + .copy_to_host(Alignment::of::())? + .await?; + let device_output = device_output.as_slice(); + + for (index, page) in plan.pages.iter().enumerate() { + let chunk = &plan.chunks[page.chunk]; + let start = usize::try_from(chunk.offset)? + page.offset_in_chunk; + let compressed = &file_bytes[start..start + page.compressed_len]; + let expected = codec.decompress_host(compressed, page.uncompressed_len)?; + let actual = &device_output[page.output_offset..page.output_offset + page.uncompressed_len]; + ensure!( + actual == expected.as_slice(), + "page {index} decompressed on the GPU differs from the host codec output" + ); + } + + tracing::info!( + "verified {} GPU-decompressed {} pages against the host codec", + plan.pages.len(), + codec.name() + ); + Ok(()) +} diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 3dbb68bc7a8..0e31d78ec00 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -8,9 +8,13 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; +use anyhow::ensure; +use arrow_schema::Field; use async_trait::async_trait; use futures::StreamExt; use tempfile::NamedTempFile; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; @@ -18,10 +22,12 @@ use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; +use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::compress::Compressor; use vortex_bench::conversions::parquet_to_vortex_chunks; +use vortex_cuda::CanonicalCudaExt; use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; #[cfg(target_os = "linux")] @@ -31,7 +37,20 @@ use vortex_cuda::layout::CudaFlatLayoutStrategy; use vortex_cuda::layout::register_cuda_layout; /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. -pub struct GpuVortexCompressor; +pub struct GpuVortexCompressor { + verify: bool, +} + +impl GpuVortexCompressor { + /// Create the backend. + /// + /// When `verify` is set, each GPU-decoded field is copied back to the host and compared + /// against the same field decoded on the CPU. Verification runs inline, so timings from a + /// verifying run are not comparable to a plain one. + pub fn new(verify: bool) -> Self { + Self { verify } + } +} #[async_trait] impl Compressor for GpuVortexCompressor { @@ -75,7 +94,13 @@ impl Compressor for GpuVortexCompressor { while let Some(batch) = batches.next().await { let record = batch?.execute::(cuda_ctx.execution_ctx())?; for field in record.iter_unmasked_fields() { - black_box(field.clone().execute_cuda(&mut cuda_ctx).await?); + let decoded = field.clone().execute_cuda(&mut cuda_ctx).await?; + if self.verify { + let host = decoded.into_host().await?.into_array(); + verify_field(field, host, cuda_ctx.execution_ctx())?; + } else { + black_box(decoded); + } } } cuda_ctx.synchronize_stream()?; @@ -83,3 +108,21 @@ impl Compressor for GpuVortexCompressor { Ok(start.elapsed()) } } + +/// Fails unless a GPU-decoded field matches the same field decoded on the CPU. +fn verify_field(compressed: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { + let expected = SESSION + .arrow() + .execute_arrow(compressed.clone(), None, ctx)?; + // Pin the Arrow target type so the two sides cannot land on different but equivalent + // encodings of the same logical values. + let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable()); + let actual = SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?; + + ensure!( + expected.to_data() == actual.to_data(), + "GPU decode of a {} field does not match the CPU decode", + compressed.encoding_id() + ); + Ok(()) +} diff --git a/benchmarks/compress-bench/src/lib.rs b/benchmarks/compress-bench/src/lib.rs index 68039996605..e82f8574dc6 100644 --- a/benchmarks/compress-bench/src/lib.rs +++ b/benchmarks/compress-bench/src/lib.rs @@ -4,6 +4,9 @@ #[cfg(feature = "lance")] pub use lance_bench::compress::LanceCompressor; #[cfg(feature = "cuda")] +pub mod gpu_parquet; +#[cfg(feature = "cuda")] pub mod gpu_vortex; pub mod parquet; +pub mod parquet_pages; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 1b1603e52c8..a915112897f 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -8,8 +8,11 @@ use clap::Parser; #[cfg(feature = "lance")] use compress_bench::LanceCompressor; #[cfg(feature = "cuda")] +use compress_bench::gpu_parquet::GpuParquetCompressor; +#[cfg(feature = "cuda")] use compress_bench::gpu_vortex::GpuVortexCompressor; use compress_bench::parquet::ParquetCompressor; +use compress_bench::parquet_pages::GpuCodec; use compress_bench::vortex::VortexCompressor; use indicatif::ProgressBar; use itertools::Itertools; @@ -35,6 +38,7 @@ 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::CustomUnitMeasurement; use vortex_bench::public_bi::PBI_DATASETS; use vortex_bench::public_bi::PBIDataset::Arade; use vortex_bench::public_bi::PBIDataset::Bimbo; @@ -67,11 +71,23 @@ struct Args { ops: Vec, #[arg(long)] datasets: Option, - /// Run GPU decompression for the allow-listed benchmarks. + /// Run GPU decompression for the GPU-supported benchmarks. /// - /// This filters the suite to GPU-supported dataset names and runs only Vortex decompression. + /// Restricts the suite to datasets with verified CUDA decode support and measures + /// decompression only, for both Vortex and Parquet. #[arg(long)] gpu_decompress: bool, + /// Page codec the GPU Parquet backend writes and decompresses with. + /// + /// Snappy is the Parquet default and the codec nvCOMP decompresses fastest. + #[arg(long, value_enum, default_value_t)] + gpu_parquet_codec: GpuCodec, + /// Cross-check every GPU-decompressed result against the CPU decoder. + /// + /// Verification runs inline, so timings from a verifying run are not comparable to a + /// plain one. Intended to be run as its own pass. + #[arg(long)] + gpu_verify: bool, #[arg(short, long, default_value_t, value_enum)] display_format: DisplayFormat, #[arg(short, long)] @@ -97,8 +113,16 @@ async fn main() -> anyhow::Result<()> { anyhow::bail!("--gpu-decompress requires building compress-bench with --features cuda"); } - let (formats, ops) = if args.gpu_decompress { - (vec![Format::OnDiskVortex], vec![CompressOp::Decompress]) + let gpu = args.gpu_decompress.then_some(GpuOptions { + codec: args.gpu_parquet_codec, + verify: args.gpu_verify, + }); + + let (formats, ops) = if gpu.is_some() { + ( + vec![Format::Parquet, Format::OnDiskVortex], + vec![CompressOp::Decompress], + ) } else { (args.formats, args.ops) }; @@ -108,7 +132,7 @@ async fn main() -> anyhow::Result<()> { args.datasets.map(|d| Regex::new(&d)).transpose()?, formats, ops, - args.gpu_decompress, + gpu, args.display_format, args.output_path, args.ingest_output, @@ -116,15 +140,35 @@ async fn main() -> anyhow::Result<()> { .await } +/// Settings for the GPU decompression mode. +#[derive(Clone, Copy, Debug)] +struct GpuOptions { + /// Parquet page codec to write and decompress on the device. + codec: GpuCodec, + /// Cross-check decompressed output against the CPU decoders. + #[cfg_attr( + not(feature = "cuda"), + expect(dead_code, reason = "only the CUDA backends read this") + )] + verify: bool, +} + /// Get a compressor for the given format. -fn get_compressor(format: Format, gpu_decompress: bool) -> Box { - if gpu_decompress { +fn get_compressor(format: Format, gpu: Option) -> Box { + if let Some(gpu) = gpu { #[cfg(feature = "cuda")] + return match format { + Format::OnDiskVortex => { + Box::new(GpuVortexCompressor::new(gpu.verify)) as Box + } + Format::Parquet => Box::new(GpuParquetCompressor::new(gpu.codec, gpu.verify)), + _ => unimplemented!("GPU compress bench not implemented for {format}"), + }; + #[cfg(not(feature = "cuda"))] { - return Box::new(GpuVortexCompressor); + let _ = gpu; + unreachable!("GPU feature validation happens before selecting compressors"); } - #[cfg(not(feature = "cuda"))] - unreachable!("GPU feature validation happens before selecting compressors"); } match format { @@ -151,7 +195,7 @@ async fn run_compress( datasets_filter: Option, formats: Vec, ops: Vec, - gpu_decompress: bool, + gpu: Option, display_format: DisplayFormat, output_path: Option, ingest_output: Option, @@ -178,15 +222,26 @@ async fn run_compress( // ), ]; - // Add an existing benchmark name here only after its CUDA-compatible compression and - // decompression kernels have been verified end to end. - #[expect( - clippy::useless_vec, - reason = "this is an intentionally incremental allow-list of benchmark names" - )] - let gpu_decompress_benchmarks = vec!["TPC-H l_comment canonical"]; + // Datasets run in GPU mode. Add one only after its CUDA-compatible compression and + // decompression kernels have been verified end to end with `--gpu-verify`. Between them + // these cover FSST strings, ALP and bit-packed numerics, run-end and date/time-parts + // encodings, and columns with nulls. + // + // `StructListOfInts` is deliberately absent: its list layouts have no verified CUDA + // decode path yet. + let gpu_datasets: [&dyn Dataset; 9] = [ + &TPCHLCommentCanonical as &dyn Dataset, + &TPCHLCommentChunked, + &TaxiData, + PBI_DATASETS.get(Arade), + PBI_DATASETS.get(Bimbo), + PBI_DATASETS.get(CMSprovider), + PBI_DATASETS.get(Euro2016), + PBI_DATASETS.get(Food), + PBI_DATASETS.get(HashTags), + ]; - let datasets: Vec<&dyn Dataset> = [ + let all_datasets: Vec<&dyn Dataset> = [ &TaxiData as &dyn Dataset, PBI_DATASETS.get(Arade), PBI_DATASETS.get(Bimbo), @@ -206,10 +261,15 @@ async fn run_compress( ] .into_iter() .chain(structlistofints.iter().map(|d| d as &dyn Dataset)) + .collect(); + + let datasets: Vec<&dyn Dataset> = if gpu.is_some() { + gpu_datasets.to_vec() + } else { + all_datasets + } + .into_iter() .filter(|d| { - if gpu_decompress && !gpu_decompress_benchmarks.contains(&d.name()) { - return false; - } if let Some(filter) = datasets_filter.as_ref() { filter.is_match(d.name()) } else { @@ -226,15 +286,9 @@ async fn run_compress( let mut v3_records: Vec = Vec::new(); for dataset_handle in datasets.into_iter() { - let (m, mut records) = run_benchmark_for_dataset( - &progress, - &formats, - &ops, - iterations, - dataset_handle, - gpu_decompress, - ) - .await?; + let (m, mut records) = + run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle, gpu) + .await?; measurements.push(m); v3_records.append(&mut records); } @@ -275,7 +329,7 @@ async fn run_benchmark_for_dataset( ops: &[CompressOp], iterations: usize, dataset_handle: &dyn Dataset, - gpu_decompress: bool, + gpu: Option, ) -> anyhow::Result<(CompressMeasurements, Vec)> { let bench_name = dataset_handle.name(); let (v3_dataset, v3_variant) = dataset_handle.v3_dataset_dims(); @@ -291,7 +345,7 @@ async fn run_benchmark_for_dataset( let mut v3_records: Vec = Vec::new(); for format in formats { - let compressor = get_compressor(*format, gpu_decompress); + let compressor = get_compressor(*format, gpu); for op in ops { let time = match op { @@ -342,7 +396,7 @@ async fn run_benchmark_for_dataset( v3_records.push(v3::compression_time_record( &result.timing, v3_dataset, - if gpu_decompress { + if gpu.is_some() { Some("gpu") } else { v3_variant @@ -361,12 +415,42 @@ async fn run_benchmark_for_dataset( } // Calculate cross-format ratios after all measurements. - calculate_ratios( - &measurements_map, - &compressed_sizes, - bench_name, - &mut ratios, - ); + match gpu { + // The shared ratio labels name the CPU suite's codec, which the GPU run does not + // necessarily use, so GPU mode emits its own correctly-labelled ratio. + Some(gpu) => push_gpu_ratio(&measurements_map, gpu, bench_name, &mut ratios), + None => calculate_ratios( + &measurements_map, + &compressed_sizes, + bench_name, + &mut ratios, + ), + } Ok((CompressMeasurements { timings, ratios }, v3_records)) } + +/// Emit the Vortex-versus-Parquet decompression ratio for a GPU run. +fn push_gpu_ratio( + measurements: &HashMap<(Format, CompressOp), Duration>, + gpu: GpuOptions, + bench_name: &str, + ratios: &mut Vec, +) { + let (Some(vortex_time), Some(parquet_time)) = ( + measurements.get(&(Format::OnDiskVortex, CompressOp::Decompress)), + measurements.get(&(Format::Parquet, CompressOp::Decompress)), + ) else { + return; + }; + + ratios.push(CustomUnitMeasurement { + name: format!( + "vortex:parquet-{} gpu ratio decompress time/{bench_name}", + gpu.codec.name() + ), + format: Format::OnDiskVortex, + unit: std::borrow::Cow::from("ratio"), + value: vortex_time.as_nanos() as f64 / parquet_time.as_nanos() as f64, + }); +} diff --git a/benchmarks/compress-bench/src/parquet_pages.rs b/benchmarks/compress-bench/src/parquet_pages.rs new file mode 100644 index 00000000000..d55598cd8dd --- /dev/null +++ b/benchmarks/compress-bench/src/parquet_pages.rs @@ -0,0 +1,530 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Locating the compressed page bodies inside a Parquet file. +//! +//! Parquet compresses each page body independently with a block codec, which is exactly the +//! shape nvCOMP's batched decompression entrypoints consume: an array of independent +//! compressed chunks with known uncompressed sizes. This module finds those chunks so the +//! GPU backend can hand the whole batch to the device in one launch, the same decomposition +//! cuDF's Parquet reader uses. +//! +//! Column chunk ranges come from the file footer; page boundaries within a chunk are only +//! discoverable by walking the per-page Thrift headers, so a minimal Thrift compact-protocol +//! reader lives here. `parquet::format::PageHeader` is deprecated and scheduled for removal, +//! and `parquet`'s own page-header parser is crate-private, so neither can be used. + +use anyhow::Result; +use anyhow::bail; +use anyhow::ensure; +use clap::ValueEnum; +use parquet::basic::Compression; +use parquet::basic::ZstdLevel; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::EnabledStatistics; +use parquet::file::properties::WriterProperties; +use parquet::file::properties::WriterVersion; + +/// Target size of a data page written for GPU decompression. +/// +/// nvCOMP decompresses one chunk per page, so pages must be large enough to amortize the +/// per-chunk setup yet numerous enough to fill the device. ~1 MiB is the page size cuDF's +/// Parquet reader is tuned around. +pub const GPU_DATA_PAGE_SIZE: usize = 1024 * 1024; + +/// Row cap per data page. +/// +/// `parquet`'s default caps pages at 20k rows, which produces pages far below +/// [`GPU_DATA_PAGE_SIZE`] for narrow columns and leaves the device underfed. +const GPU_DATA_PAGE_ROW_COUNT_LIMIT: usize = 1_000_000; + +/// Parquet page codecs that nvCOMP can decompress on the device. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)] +pub enum GpuCodec { + /// The Parquet default, and the codec with the highest device-side throughput. + #[default] + Snappy, + /// Matches the codec used by the CPU Parquet benchmark, at lower device throughput. + Zstd, +} + +impl GpuCodec { + /// The Parquet compression setting for this codec. + pub fn to_parquet(self) -> Compression { + match self { + GpuCodec::Snappy => Compression::SNAPPY, + GpuCodec::Zstd => Compression::ZSTD(ZstdLevel::default()), + } + } + + /// Short lowercase name, used in measurement labels. + pub fn name(self) -> &'static str { + match self { + GpuCodec::Snappy => "snappy", + GpuCodec::Zstd => "zstd", + } + } + + /// Whether a column chunk's codec matches this one. + pub fn matches(self, compression: Compression) -> bool { + matches!( + (self, compression), + (GpuCodec::Snappy, Compression::SNAPPY) | (GpuCodec::Zstd, Compression::ZSTD(_)) + ) + } + + /// Decompress a single page body on the host, for cross-checking device output. + pub fn decompress_host(self, compressed: &[u8], uncompressed_len: usize) -> Result> { + let decompressed = match self { + GpuCodec::Snappy => snap::raw::Decoder::new().decompress_vec(compressed)?, + GpuCodec::Zstd => zstd::bulk::decompress(compressed, uncompressed_len)?, + }; + ensure!( + decompressed.len() == uncompressed_len, + "page decompressed to {} bytes, page header declared {uncompressed_len}", + decompressed.len() + ); + Ok(decompressed) + } +} + +/// Writer properties tuned for GPU decompression. +pub fn gpu_writer_properties(codec: GpuCodec) -> WriterProperties { + WriterProperties::builder() + // V1 data pages compress the entire page body, which is the unit nvCOMP decompresses. + // V2 pages place uncompressed repetition/definition levels ahead of the compressed + // values inside one page body, which the batched entrypoints cannot express. + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_compression(codec.to_parquet()) + // Dictionary encoding keeps the decompressed payload small and is the encoding GPU + // Parquet readers decode fastest. + .set_dictionary_enabled(true) + .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) + .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) + // Per-page statistics only inflate the page headers that have to be walked on the host. + .set_statistics_enabled(EnabledStatistics::Chunk) + .build() +} + +/// A compressed page body located within a Parquet file. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CompressedPage { + /// Offset of the compressed body, i.e. just past the page header. + pub offset: usize, + /// Length of the compressed body in bytes. + pub compressed_len: usize, + /// Length of the body once decompressed. + pub uncompressed_len: usize, +} + +/// The pages of one column chunk, alongside the byte range the chunk occupies. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ColumnChunkPages { + /// Offset of the column chunk within the file. + pub offset: u64, + /// Length of the column chunk in bytes. + pub len: usize, + /// Compressed pages of this chunk, in file order, with file-absolute offsets. + pub pages: Vec, +} + +/// Walks every column chunk's page headers and returns the compressed page bodies. +/// +/// The outer `Vec` is one entry per row group, which is the unit a reader can stage on the +/// device and release before moving on. Chunks and the pages within them are in file order. +pub fn scan_compressed_pages( + file_bytes: &[u8], + metadata: &ParquetMetaData, +) -> Result>> { + let mut row_groups = Vec::with_capacity(metadata.row_groups().len()); + + for row_group in metadata.row_groups() { + let mut chunks = Vec::with_capacity(row_group.columns().len()); + for column in row_group.columns() { + let (chunk_offset, chunk_len) = column.byte_range(); + let mut pages = Vec::new(); + let (start, len) = (chunk_offset, chunk_len); + let start = usize::try_from(start)?; + let end = start + .checked_add(usize::try_from(len)?) + .filter(|end| *end <= file_bytes.len()) + .ok_or_else(|| { + anyhow::anyhow!( + "column chunk range {start}..+{len} extends past the {} byte file", + file_bytes.len() + ) + })?; + + let mut pos = start; + while pos < end { + let header = read_page_header(&file_bytes[pos..end])?; + let body = pos + header.header_len; + let body_end = body + .checked_add(header.compressed_len) + .filter(|body_end| *body_end <= end) + .ok_or_else(|| { + anyhow::anyhow!( + "page body at {body} of {} bytes overruns its column chunk", + header.compressed_len + ) + })?; + + match header.page_type { + PageType::Data | PageType::Dictionary => pages.push(CompressedPage { + offset: body, + compressed_len: header.compressed_len, + uncompressed_len: header.uncompressed_len, + }), + PageType::DataV2 => bail!( + "v2 data pages are not GPU-decompressible as a single chunk; \ + write the file with WriterVersion::PARQUET_1_0" + ), + // Index pages are not part of the column data and are never written by + // `parquet`; skip over the body rather than decompressing it. + PageType::Index => {} + } + + pos = body_end; + } + + chunks.push(ColumnChunkPages { + offset: chunk_offset, + len: usize::try_from(chunk_len)?, + pages, + }); + } + + row_groups.push(chunks); + } + + Ok(row_groups) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PageType { + Data, + Index, + Dictionary, + DataV2, +} + +struct PageHeaderInfo { + header_len: usize, + page_type: PageType, + compressed_len: usize, + uncompressed_len: usize, +} + +/// Thrift compact-protocol field types. +mod ttype { + pub(super) const STOP: u8 = 0x00; + pub(super) const BOOL_TRUE: u8 = 0x01; + pub(super) const BOOL_FALSE: u8 = 0x02; + pub(super) const I8: u8 = 0x03; + pub(super) const I16: u8 = 0x04; + pub(super) const I32: u8 = 0x05; + pub(super) const I64: u8 = 0x06; + pub(super) const DOUBLE: u8 = 0x07; + pub(super) const BINARY: u8 = 0x08; + pub(super) const LIST: u8 = 0x09; + pub(super) const SET: u8 = 0x0a; + pub(super) const MAP: u8 = 0x0b; + pub(super) const STRUCT: u8 = 0x0c; + pub(super) const UUID: u8 = 0x0d; +} + +/// Guards against unbounded recursion on malformed headers. +const MAX_STRUCT_DEPTH: u32 = 32; + +/// Reads the `PageHeader` at the start of `buf`, returning its fields and encoded length. +fn read_page_header(buf: &[u8]) -> Result { + let mut reader = CompactReader { buf, pos: 0 }; + let mut page_type = None; + let mut uncompressed_len = None; + let mut compressed_len = None; + let mut last_field_id = 0i16; + + while let Some((field_id, field_type)) = reader.read_field_header(&mut last_field_id)? { + match (field_id, field_type) { + (1, ttype::I32) => page_type = Some(reader.read_i32()?), + (2, ttype::I32) => uncompressed_len = Some(reader.read_i32()?), + (3, ttype::I32) => compressed_len = Some(reader.read_i32()?), + _ => reader.skip_value(field_type, 0)?, + } + } + + let page_type = match page_type { + Some(0) => PageType::Data, + Some(1) => PageType::Index, + Some(2) => PageType::Dictionary, + Some(3) => PageType::DataV2, + Some(other) => bail!("unknown Parquet page type {other}"), + None => bail!("page header is missing its page type"), + }; + + let uncompressed_len = + uncompressed_len.ok_or_else(|| anyhow::anyhow!("page header is missing its size"))?; + let compressed_len = compressed_len + .ok_or_else(|| anyhow::anyhow!("page header is missing its compressed size"))?; + + Ok(PageHeaderInfo { + header_len: reader.pos, + page_type, + compressed_len: usize::try_from(compressed_len)?, + uncompressed_len: usize::try_from(uncompressed_len)?, + }) +} + +/// Minimal reader for the subset of the Thrift compact protocol that page headers use. +struct CompactReader<'a> { + buf: &'a [u8], + pos: usize, +} + +impl CompactReader<'_> { + fn read_u8(&mut self) -> Result { + let byte = *self + .buf + .get(self.pos) + .ok_or_else(|| anyhow::anyhow!("page header ends mid-field"))?; + self.pos += 1; + Ok(byte) + } + + fn advance(&mut self, len: usize) -> Result<()> { + let end = self + .pos + .checked_add(len) + .filter(|end| *end <= self.buf.len()) + .ok_or_else(|| anyhow::anyhow!("page header ends mid-value"))?; + self.pos = end; + Ok(()) + } + + fn read_varint(&mut self) -> Result { + let mut value = 0u64; + for shift in (0..64).step_by(7) { + let byte = self.read_u8()?; + value |= u64::from(byte & 0x7f) << shift; + if byte & 0x80 == 0 { + return Ok(value); + } + } + bail!("varint in page header is not terminated") + } + + fn read_zigzag(&mut self) -> Result { + let encoded = self.read_varint()?; + Ok(((encoded >> 1) as i64) ^ -((encoded & 1) as i64)) + } + + fn read_i32(&mut self) -> Result { + Ok(i32::try_from(self.read_zigzag()?)?) + } + + /// Reads the next field header, or `None` at the struct's STOP byte. + fn read_field_header(&mut self, last_field_id: &mut i16) -> Result> { + let header = self.read_u8()?; + if header == ttype::STOP { + return Ok(None); + } + + let field_type = header & 0x0f; + let delta = header >> 4; + let field_id = if delta == 0 { + i16::try_from(self.read_zigzag()?)? + } else { + last_field_id + .checked_add(i16::from(delta)) + .ok_or_else(|| anyhow::anyhow!("field id overflow in page header"))? + }; + *last_field_id = field_id; + + Ok(Some((field_id, field_type))) + } + + fn skip_struct(&mut self, depth: u32) -> Result<()> { + ensure!( + depth < MAX_STRUCT_DEPTH, + "page header nests structs more than {MAX_STRUCT_DEPTH} deep" + ); + let mut last_field_id = 0i16; + while let Some((_, field_type)) = self.read_field_header(&mut last_field_id)? { + self.skip_value(field_type, depth + 1)?; + } + Ok(()) + } + + /// Skips a field value. Booleans carry their value in the field type, so consume nothing. + fn skip_value(&mut self, field_type: u8, depth: u32) -> Result<()> { + match field_type { + ttype::BOOL_TRUE | ttype::BOOL_FALSE => Ok(()), + ttype::I8 => self.advance(1), + ttype::I16 | ttype::I32 | ttype::I64 => self.read_varint().map(|_| ()), + ttype::DOUBLE => self.advance(8), + ttype::UUID => self.advance(16), + ttype::BINARY => { + let len = usize::try_from(self.read_varint()?)?; + self.advance(len) + } + ttype::LIST | ttype::SET => { + let (len, element_type) = self.read_collection_header()?; + for _ in 0..len { + self.skip_element(element_type, depth + 1)?; + } + Ok(()) + } + ttype::MAP => { + let len = usize::try_from(self.read_varint()?)?; + if len > 0 { + let types = self.read_u8()?; + let (key_type, value_type) = (types >> 4, types & 0x0f); + for _ in 0..len { + self.skip_element(key_type, depth + 1)?; + self.skip_element(value_type, depth + 1)?; + } + } + Ok(()) + } + ttype::STRUCT => self.skip_struct(depth), + other => bail!("unsupported Thrift compact type {other} in page header"), + } + } + + /// Skips one collection element. Unlike fields, booleans here occupy a byte of their own. + fn skip_element(&mut self, element_type: u8, depth: u32) -> Result<()> { + match element_type { + ttype::BOOL_TRUE | ttype::BOOL_FALSE => self.advance(1), + other => self.skip_value(other, depth), + } + } + + fn read_collection_header(&mut self) -> Result<(usize, u8)> { + let header = self.read_u8()?; + let element_type = header & 0x0f; + let len = match header >> 4 { + 0x0f => usize::try_from(self.read_varint()?)?, + short_len => usize::from(short_len), + }; + Ok((len, element_type)) + } +} + +#[cfg(test)] +mod tests { + use std::fs::File; + use std::sync::Arc; + + use arrow_array::Int64Array; + use arrow_array::RecordBatch; + use arrow_array::StringArray; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::ParquetMetaDataReader; + use parquet::file::reader::FileReader; + use parquet::file::reader::SerializedFileReader; + use rstest::rstest; + + use super::*; + + fn sample_batch() -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("ints", DataType::Int64, false), + Field::new("strings", DataType::Utf8, false), + ])); + let ints = Int64Array::from_iter_values((0..50_000).map(|i| i % 977)); + let strings = + StringArray::from_iter_values((0..50_000).map(|i| format!("value-{}", i % 1_000))); + Ok(RecordBatch::try_new( + schema, + vec![Arc::new(ints), Arc::new(strings)], + )?) + } + + fn write_sample(path: &std::path::Path, codec: GpuCodec) -> Result<()> { + let batch = sample_batch()?; + let file = File::create(path)?; + let mut writer = + ArrowWriter::try_new(file, batch.schema(), Some(gpu_writer_properties(codec)))?; + writer.write(&batch)?; + writer.close()?; + Ok(()) + } + + /// The page bodies we locate must decompress to exactly the bytes `parquet` itself reads. + #[rstest] + #[case(GpuCodec::Snappy)] + #[case(GpuCodec::Zstd)] + fn scanned_pages_match_parquet_reader(#[case] codec: GpuCodec) -> Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("sample.parquet"); + write_sample(&path, codec)?; + + let file = File::open(&path)?; + let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; + let file_bytes = std::fs::read(&path)?; + let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; + let pages = row_groups + .iter() + .flatten() + .flat_map(|chunk| chunk.pages.iter()) + .collect::>(); + + let reader = SerializedFileReader::new(File::open(&path)?)?; + let mut expected = Vec::new(); + for row_group in 0..reader.metadata().num_row_groups() { + let row_group_reader = reader.get_row_group(row_group)?; + for column in 0..row_group_reader.num_columns() { + let mut page_reader = row_group_reader.get_column_page_reader(column)?; + while let Some(page) = page_reader.get_next_page()? { + expected.push(page.buffer().to_vec()); + } + } + } + + assert_eq!(pages.len(), expected.len(), "page count mismatch"); + assert!(!pages.is_empty(), "expected the sample file to have pages"); + + for (page, expected) in pages.iter().zip(expected.iter()) { + let compressed = &file_bytes[page.offset..page.offset + page.compressed_len]; + let decompressed = codec.decompress_host(compressed, page.uncompressed_len)?; + assert_eq!(&decompressed, expected); + } + + Ok(()) + } + + /// Pages must tile their column chunks exactly, with no gap left unaccounted for. + #[test] + fn scanned_pages_cover_every_column_chunk() -> Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("sample.parquet"); + write_sample(&path, GpuCodec::Snappy)?; + + let file = File::open(&path)?; + let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; + let file_bytes = std::fs::read(&path)?; + let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; + let pages = row_groups + .iter() + .flatten() + .flat_map(|chunk| chunk.pages.iter()) + .collect::>(); + + let compressed: usize = pages.iter().map(|page| page.compressed_len).sum(); + let chunk_total: i64 = metadata + .row_groups() + .iter() + .flat_map(|rg| rg.columns()) + .map(|col| col.compressed_size()) + .sum(); + + // The chunk total includes the page headers, so the page bodies must be strictly + // smaller but within a header's worth per page. + assert!(compressed < usize::try_from(chunk_total)?); + assert!(compressed > usize::try_from(chunk_total)? - pages.len() * 256); + Ok(()) + } +} diff --git a/vortex-cuda/nvcomp/build.rs b/vortex-cuda/nvcomp/build.rs index 877fc68c58d..f34507378db 100644 --- a/vortex-cuda/nvcomp/build.rs +++ b/vortex-cuda/nvcomp/build.rs @@ -90,13 +90,20 @@ fn main() { let bindings = bindgen::Builder::default() .header(include_dir.join("nvcomp.h").to_string_lossy()) .header(include_dir.join("nvcomp/zstd.h").to_string_lossy()) + .header(include_dir.join("nvcomp/snappy.h").to_string_lossy()) .clang_arg(format!("-I{}", include_dir.display())) .clang_arg(format!("-I{}", cuda_stub_dir.display())) .allowlist_type("nvcompStatus_t") + .allowlist_type("nvcompAlignmentRequirements_t") .allowlist_type("nvcompBatchedZstdDecompressOpts_t") + .allowlist_type("nvcompBatchedSnappyDecompressOpts_t") .allowlist_type("nvcompDecompressBackend_t") .allowlist_function("nvcompBatchedZstdDecompressGetTempSizeAsync") + .allowlist_function("nvcompBatchedZstdDecompressGetRequiredAlignments") .allowlist_function("nvcompBatchedZstdDecompressAsync") + .allowlist_function("nvcompBatchedSnappyDecompressGetTempSizeAsync") + .allowlist_function("nvcompBatchedSnappyDecompressGetRequiredAlignments") + .allowlist_function("nvcompBatchedSnappyDecompressAsync") .dynamic_library_name("NvcompLibrary") .dynamic_link_require_all(true) .wrap_unsafe_ops(true) diff --git a/vortex-cuda/nvcomp/src/backend.rs b/vortex-cuda/nvcomp/src/backend.rs new file mode 100644 index 00000000000..5be1ebbde4c --- /dev/null +++ b/vortex-cuda/nvcomp/src/backend.rs @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Options and metadata shared by nvcomp's batched decompression APIs. + +use crate::sys; + +/// Backend selection for nvcomp decompression. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum DecompressBackend { + /// Let nvcomp auto-select the best backend for the hardware. + #[default] + Default, + /// Use hardware decompression + Hardware, + /// Use CUDA + Cuda, +} + +impl DecompressBackend { + pub(crate) fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { + match self { + Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, + Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, + Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, + } + } +} + +/// Minimum buffer alignments required by an nvcomp algorithm. +/// +/// Buffers passed to the batched decompression entrypoints must satisfy these alignments. +/// Exceeding them (for example 16- or 32-byte alignment) may improve throughput. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AlignmentRequirements { + /// Minimum alignment of each compressed input chunk. + pub input: usize, + /// Minimum alignment of each decompressed output chunk. + pub output: usize, + /// Minimum alignment of the temporary workspace buffer. + pub temp: usize, +} + +impl From for AlignmentRequirements { + fn from(value: sys::nvcompAlignmentRequirements_t) -> Self { + Self { + input: value.input, + output: value.output, + temp: value.temp, + } + } +} diff --git a/vortex-cuda/nvcomp/src/lib.rs b/vortex-cuda/nvcomp/src/lib.rs index 3ab502baa63..859eecbea22 100644 --- a/vortex-cuda/nvcomp/src/lib.rs +++ b/vortex-cuda/nvcomp/src/lib.rs @@ -30,9 +30,13 @@ use std::sync::OnceLock; )] pub mod sys; +mod backend; mod error; +pub mod snappy; pub mod zstd; +pub use backend::AlignmentRequirements; +pub use backend::DecompressBackend; pub use error::NvcompError; /// The loaded nvcomp library instance. diff --git a/vortex-cuda/nvcomp/src/snappy.rs b/vortex-cuda/nvcomp/src/snappy.rs new file mode 100644 index 00000000000..175c19d3a13 --- /dev/null +++ b/vortex-cuda/nvcomp/src/snappy.rs @@ -0,0 +1,222 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Wrappers around nvcomp's batched Snappy decompression API. +//! +//! Snappy is the default Parquet page codec and is the fastest of the codecs nvcomp +//! implements on device, which makes it the codec of choice when feeding Parquet pages +//! to the GPU. + +use std::ffi::c_void; + +use crate::backend::AlignmentRequirements; +pub use crate::backend::DecompressBackend; +use crate::error::NvcompError; +use crate::error::check_status; +use crate::nvcomp_library; +use crate::sys; + +/// The largest compressed chunk the Snappy decompressor accepts, in bytes. +pub const MAX_COMPRESSED_CHUNK_SIZE: usize = (1 << 31) - 1; + +/// Options for batched Snappy decompression. +#[derive(Debug, Clone, Copy, Default)] +pub struct SnappyDecompressOpts { + /// Which nvcomp backend performs the decompression. + pub backend: DecompressBackend, + /// Sort chunks by size before submitting them to the hardware decompression engine. + /// + /// Only used when `backend` selects the hardware engine. + pub sort_before_hw_decompress: bool, +} + +impl SnappyDecompressOpts { + fn to_nvcomp(self) -> sys::nvcompBatchedSnappyDecompressOpts_t { + sys::nvcompBatchedSnappyDecompressOpts_t { + backend: self.backend.to_nvcomp(), + sort_before_hw_decompress: i32::from(self.sort_before_hw_decompress), + reserved: [0; 56], + } + } +} + +/// Computes required temporary buffer size for batched Snappy decompression. +/// +/// # Arguments +/// +/// * `num_chunks` - Number of compressed chunks to decompress +/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk +/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks +/// +/// # Returns +/// +/// The required size in bytes for the temporary buffer. +pub fn get_decompress_temp_size( + num_chunks: usize, + max_uncompressed_chunk_bytes: usize, + max_total_uncompressed_bytes: usize, +) -> Result { + get_decompress_temp_size_with_opts( + num_chunks, + max_uncompressed_chunk_bytes, + max_total_uncompressed_bytes, + SnappyDecompressOpts::default(), + ) +} + +/// Computes required temporary buffer size with custom options. +/// +/// # Arguments +/// +/// * `num_chunks` - Number of compressed chunks to decompress +/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk +/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks +/// * `opts` - Decompression options +/// +/// # Returns +/// +/// The required size in bytes for the temporary buffer. +pub fn get_decompress_temp_size_with_opts( + num_chunks: usize, + max_uncompressed_chunk_bytes: usize, + max_total_uncompressed_bytes: usize, + opts: SnappyDecompressOpts, +) -> Result { + let library = nvcomp_library()?; + + let mut temp_bytes: usize = 0; + + let status = unsafe { + library.nvcompBatchedSnappyDecompressGetTempSizeAsync( + num_chunks, + max_uncompressed_chunk_bytes, + opts.to_nvcomp(), + &raw mut temp_bytes, + max_total_uncompressed_bytes, + ) + }; + + check_status(status)?; + Ok(temp_bytes) +} + +/// Returns the minimum buffer alignments required by batched Snappy decompression. +pub fn decompress_alignment_requirements( + opts: SnappyDecompressOpts, +) -> Result { + let library = nvcomp_library()?; + + let mut requirements = sys::nvcompAlignmentRequirements_t { + input: 0, + output: 0, + temp: 0, + }; + + let status = unsafe { + library.nvcompBatchedSnappyDecompressGetRequiredAlignments( + opts.to_nvcomp(), + &raw mut requirements, + ) + }; + + check_status(status)?; + Ok(requirements.into()) +} + +/// Launches batched Snappy decompression asynchronously on the GPU. +/// +/// This function decompresses multiple raw Snappy blocks in parallel on the GPU. All +/// pointer arguments must point to device memory, and the operation is executed +/// asynchronously on the provided CUDA stream. +/// +/// # Arguments +/// +/// * `device_compressed_ptrs` - Device pointer to array of pointers to compressed chunks +/// * `device_compressed_bytes` - Device pointer to array of compressed chunk sizes +/// * `device_uncompressed_bytes` - Device pointer to array of expected uncompressed sizes +/// * `device_actual_uncompressed_bytes` - Device pointer to array for actual uncompressed sizes (output) +/// * `num_chunks` - Number of chunks to decompress +/// * `device_temp_ptr` - Device pointer to temporary workspace buffer +/// * `temp_bytes` - Size of temporary buffer in bytes +/// * `device_uncompressed_ptrs` - Device pointer to array of pointers to output buffers +/// * `device_statuses` - Device pointer to array for per-chunk status codes (output) +/// * `stream` - CUDA stream to execute on +/// +/// # Safety +/// +/// - All device pointers must be valid and point to properly allocated device memory +/// - `device_compressed_ptrs` must point to valid device pointers +/// - `device_uncompressed_ptrs` must point to valid device pointers +/// - Each output buffer must have at least the corresponding `device_uncompressed_bytes` size +/// - `device_temp_ptr` must have at least `temp_bytes` allocated +/// - The stream must be valid +#[expect(clippy::too_many_arguments)] +pub unsafe fn decompress_async( + device_compressed_ptrs: *const *const c_void, + device_compressed_bytes: *const usize, + device_uncompressed_bytes: *const usize, + device_actual_uncompressed_bytes: *mut usize, + num_chunks: usize, + device_temp_ptr: *mut c_void, + temp_bytes: usize, + device_uncompressed_ptrs: *const *mut c_void, + device_statuses: *mut sys::nvcompStatus_t, + stream: sys::cudaStream_t, +) -> Result<(), NvcompError> { + // SAFETY: Caller has to ensure all pointers are valid. + unsafe { + decompress_async_with_opts( + device_compressed_ptrs, + device_compressed_bytes, + device_uncompressed_bytes, + device_actual_uncompressed_bytes, + num_chunks, + device_temp_ptr, + temp_bytes, + device_uncompressed_ptrs, + device_statuses, + stream, + SnappyDecompressOpts::default(), + ) + } +} + +/// Launches batched Snappy decompression asynchronously with custom options. +/// +/// # Safety +/// +/// Same requirements as [`decompress_async`]. +#[expect(clippy::too_many_arguments)] +pub unsafe fn decompress_async_with_opts( + device_compressed_ptrs: *const *const c_void, + device_compressed_bytes: *const usize, + device_uncompressed_bytes: *const usize, + device_actual_uncompressed_bytes: *mut usize, + num_chunks: usize, + device_temp_ptr: *mut c_void, + temp_bytes: usize, + device_uncompressed_ptrs: *const *mut c_void, + device_statuses: *mut sys::nvcompStatus_t, + stream: sys::cudaStream_t, + opts: SnappyDecompressOpts, +) -> Result<(), NvcompError> { + let library = nvcomp_library()?; + + let status = unsafe { + library.nvcompBatchedSnappyDecompressAsync( + device_compressed_ptrs, + device_compressed_bytes, + device_uncompressed_bytes, + device_actual_uncompressed_bytes, + num_chunks, + device_temp_ptr, + temp_bytes, + device_uncompressed_ptrs, + opts.to_nvcomp(), + device_statuses, + stream, + ) + }; + + check_status(status) +} diff --git a/vortex-cuda/nvcomp/src/zstd.rs b/vortex-cuda/nvcomp/src/zstd.rs index 44111901071..04769d2621a 100644 --- a/vortex-cuda/nvcomp/src/zstd.rs +++ b/vortex-cuda/nvcomp/src/zstd.rs @@ -5,33 +5,13 @@ use std::ffi::c_void; +use crate::backend::AlignmentRequirements; +pub use crate::backend::DecompressBackend; use crate::error::NvcompError; use crate::error::check_status; use crate::nvcomp_library; use crate::sys; -/// Backend selection for nvcomp decompression. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum DecompressBackend { - /// Let nvcomp auto-select the best backend for the hardware. - #[default] - Default, - /// Use hardware decompression - Hardware, - /// Use CUDA - Cuda, -} - -impl DecompressBackend { - fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { - match self { - Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, - Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, - Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, - } - } -} - /// Options for batched ZSTD decompression. #[derive(Debug, Clone, Copy, Default)] pub struct ZstdDecompressOpts { @@ -107,6 +87,29 @@ pub fn get_decompress_temp_size_with_opts( Ok(temp_bytes) } +/// Returns the minimum buffer alignments required by batched ZSTD decompression. +pub fn decompress_alignment_requirements( + opts: ZstdDecompressOpts, +) -> Result { + let library = nvcomp_library()?; + + let mut requirements = sys::nvcompAlignmentRequirements_t { + input: 0, + output: 0, + temp: 0, + }; + + let status = unsafe { + library.nvcompBatchedZstdDecompressGetRequiredAlignments( + opts.to_nvcomp(), + &raw mut requirements, + ) + }; + + check_status(status)?; + Ok(requirements.into()) +} + /// Launches batched ZSTD decompression asynchronously on the GPU. /// /// This function decompresses multiple ZSTD-compressed chunks in parallel on the GPU. From 5eb7a6d622bc3e1b956176e148d49db3ac3a45eb Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 12:50:11 +0000 Subject: [PATCH 02/16] Verify GPU Vortex decode against a separate host scan A CUDA scan hands back arrays whose buffers live in device memory, so decoding those same arrays through the host Arrow path panics rather than producing a CPU reference. Read the file a second time through the ordinary host reader and compare the two scans batch by batch instead. Signed-off-by: Claude --- benchmarks/compress-bench/src/gpu_vortex.rs | 101 ++++++++++++++++---- 1 file changed, 82 insertions(+), 19 deletions(-) diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 0e31d78ec00..d9be0b8730f 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -8,6 +8,7 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; +use anyhow::bail; use anyhow::ensure; use arrow_schema::Field; use async_trait::async_trait; @@ -80,27 +81,19 @@ impl Compressor for GpuVortexCompressor { output.sync_all().await?; drop(output); + if self.verify { + return verify_against_host_scan(gpu_file.path()).await; + } + let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; let start = Instant::now(); - let open_options = SESSION.open_options().with_cuda(); - // Direct IO keeps repeated iterations measuring storage bandwidth rather than - // page-cache hits. It is only available on Linux. - #[cfg(target_os = "linux")] - let open_options = - open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()); - let file = open_options.open_path(gpu_file.path()).await?; + let file = open_gpu(gpu_file.path()).await?; let mut batches = file.scan()?.into_array_stream()?; while let Some(batch) = batches.next().await { let record = batch?.execute::(cuda_ctx.execution_ctx())?; for field in record.iter_unmasked_fields() { - let decoded = field.clone().execute_cuda(&mut cuda_ctx).await?; - if self.verify { - let host = decoded.into_host().await?.into_array(); - verify_field(field, host, cuda_ctx.execution_ctx())?; - } else { - black_box(decoded); - } + black_box(field.clone().execute_cuda(&mut cuda_ctx).await?); } } cuda_ctx.synchronize_stream()?; @@ -109,11 +102,81 @@ impl Compressor for GpuVortexCompressor { } } +/// Opens a Vortex file for CUDA execution. +/// +/// On Linux direct IO keeps repeated iterations measuring storage bandwidth rather than +/// page-cache hits. +async fn open_gpu(path: &Path) -> Result { + let open_options = SESSION.open_options().with_cuda(); + #[cfg(target_os = "linux")] + let open_options = + open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()); + Ok(open_options.open_path(path).await?) +} + +/// Decodes the same file on the GPU and on the CPU and fails on the first difference. +/// +/// The CPU reference comes from a second, host-only scan rather than from re-decoding the +/// GPU scan's arrays: a CUDA scan hands back arrays whose buffers live in device memory, +/// which the host decoders cannot read. +/// +/// Verification runs inline, so the returned duration is not comparable to a plain run. +async fn verify_against_host_scan(path: &Path) -> Result { + let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + let start = Instant::now(); + + let gpu_file = open_gpu(path).await?; + let mut gpu_batches = gpu_file.scan()?.into_array_stream()?; + let host_file = SESSION.open_options().open_path(path).await?; + let mut host_batches = host_file.scan()?.into_array_stream()?; + + let mut fields_checked = 0usize; + loop { + let (gpu_batch, host_batch) = (gpu_batches.next().await, host_batches.next().await); + let (gpu_batch, host_batch) = match (gpu_batch, host_batch) { + (Some(gpu_batch), Some(host_batch)) => (gpu_batch?, host_batch?), + (None, None) => break, + _ => bail!("the GPU and CPU scans of the same file produced different batch counts"), + }; + + let gpu_record = gpu_batch.execute::(cuda_ctx.execution_ctx())?; + let host_record = host_batch.execute::(cuda_ctx.execution_ctx())?; + ensure!( + gpu_record.len() == host_record.len(), + "batch length differs between the GPU and CPU scans: {} vs {}", + gpu_record.len(), + host_record.len() + ); + + let gpu_fields = gpu_record + .iter_unmasked_fields() + .cloned() + .collect::>(); + let host_fields = host_record + .iter_unmasked_fields() + .cloned() + .collect::>(); + ensure!( + gpu_fields.len() == host_fields.len(), + "field count differs between the GPU and CPU scans" + ); + + for (gpu_field, host_field) in gpu_fields.into_iter().zip(host_fields) { + let decoded = gpu_field.execute_cuda(&mut cuda_ctx).await?; + let decoded = decoded.into_host().await?.into_array(); + verify_field(&host_field, decoded, cuda_ctx.execution_ctx())?; + fields_checked += 1; + } + } + cuda_ctx.synchronize_stream()?; + + tracing::info!("verified {fields_checked} GPU-decoded Vortex fields against the CPU decode"); + Ok(start.elapsed()) +} + /// Fails unless a GPU-decoded field matches the same field decoded on the CPU. -fn verify_field(compressed: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { - let expected = SESSION - .arrow() - .execute_arrow(compressed.clone(), None, ctx)?; +fn verify_field(host: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { + let expected = SESSION.arrow().execute_arrow(host.clone(), None, ctx)?; // Pin the Arrow target type so the two sides cannot land on different but equivalent // encodings of the same logical values. let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable()); @@ -122,7 +185,7 @@ fn verify_field(compressed: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> ensure!( expected.to_data() == actual.to_data(), "GPU decode of a {} field does not match the CPU decode", - compressed.encoding_id() + host.encoding_id() ); Ok(()) } From e7a6b4189c3a4d1d4fe8e9f345f5a1d55c684118 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:01:28 +0000 Subject: [PATCH 03/16] Isolate the verification scans and report how a GPU decode differs Two changes to the Vortex GPU verification, after CI reported a `fastlanes.for` mismatch with no detail: - Read the CPU reference from a copy of the file. The session segment cache is keyed by URI and the CUDA reader deliberately bypasses it because its buffers are device-resident, so pointing both scans at one URI risks them sharing entries. - Synchronize the stream before copying a decoded field back, and report the Arrow types, lengths, null counts and the first differing row when the two decodes disagree. Signed-off-by: Claude --- benchmarks/compress-bench/src/gpu_vortex.rs | 92 ++++++++++++++++++--- 1 file changed, 81 insertions(+), 11 deletions(-) diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index d9be0b8730f..9b25b2f6dd4 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -10,6 +10,7 @@ use std::time::Instant; use anyhow::Result; use anyhow::bail; use anyhow::ensure; +use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::Field; use async_trait::async_trait; use futures::StreamExt; @@ -125,12 +126,19 @@ async fn verify_against_host_scan(path: &Path) -> Result { let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; let start = Instant::now(); + // The host scan reads a copy rather than the same path. The session's segment cache is + // keyed by URI, and the CUDA reader deliberately bypasses it because its buffers are + // device-resident; running both scans against one URI risks the two sharing entries. + let host_path = NamedTempFile::new()?; + std::fs::copy(path, host_path.path())?; + let gpu_file = open_gpu(path).await?; let mut gpu_batches = gpu_file.scan()?.into_array_stream()?; - let host_file = SESSION.open_options().open_path(path).await?; + let host_file = SESSION.open_options().open_path(host_path.path()).await?; let mut host_batches = host_file.scan()?.into_array_stream()?; let mut fields_checked = 0usize; + let mut batch_index = 0usize; loop { let (gpu_batch, host_batch) = (gpu_batches.next().await, host_batches.next().await); let (gpu_batch, host_batch) = match (gpu_batch, host_batch) { @@ -143,7 +151,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { let host_record = host_batch.execute::(cuda_ctx.execution_ctx())?; ensure!( gpu_record.len() == host_record.len(), - "batch length differs between the GPU and CPU scans: {} vs {}", + "batch {batch_index} length differs between the GPU and CPU scans: {} vs {}", gpu_record.len(), host_record.len() ); @@ -158,15 +166,28 @@ async fn verify_against_host_scan(path: &Path) -> Result { .collect::>(); ensure!( gpu_fields.len() == host_fields.len(), - "field count differs between the GPU and CPU scans" + "batch {batch_index} field count differs between the GPU and CPU scans" ); - for (gpu_field, host_field) in gpu_fields.into_iter().zip(host_fields) { + for (field_index, (gpu_field, host_field)) in + gpu_fields.into_iter().zip(host_fields).enumerate() + { let decoded = gpu_field.execute_cuda(&mut cuda_ctx).await?; + // The decode is enqueued, not complete: make the writes visible before reading + // the buffers back to the host. + cuda_ctx.synchronize_stream()?; let decoded = decoded.into_host().await?.into_array(); - verify_field(&host_field, decoded, cuda_ctx.execution_ctx())?; + verify_field( + &host_field, + decoded, + cuda_ctx.execution_ctx(), + batch_index, + field_index, + )?; fields_checked += 1; } + + batch_index += 1; } cuda_ctx.synchronize_stream()?; @@ -175,17 +196,66 @@ async fn verify_against_host_scan(path: &Path) -> Result { } /// Fails unless a GPU-decoded field matches the same field decoded on the CPU. -fn verify_field(host: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { +fn verify_field( + host: &ArrayRef, + gpu: ArrayRef, + ctx: &mut ExecutionCtx, + batch_index: usize, + field_index: usize, +) -> Result<()> { let expected = SESSION.arrow().execute_arrow(host.clone(), None, ctx)?; // Pin the Arrow target type so the two sides cannot land on different but equivalent // encodings of the same logical values. let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable()); let actual = SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?; - ensure!( - expected.to_data() == actual.to_data(), - "GPU decode of a {} field does not match the CPU decode", - host.encoding_id() + if expected.to_data() == actual.to_data() { + return Ok(()); + } + + bail!( + "GPU decode of a {} field does not match the CPU decode \ + (batch {batch_index}, field {field_index}){}", + host.encoding_id(), + describe_mismatch(&expected, &actual) + ) +} + +/// Builds a human-readable description of how two Arrow arrays differ. +fn describe_mismatch(expected: &ArrowArrayRef, actual: &ArrowArrayRef) -> String { + let mut description = format!( + "\n cpu: type={:?} len={} nulls={}\n gpu: type={:?} len={} nulls={}", + expected.data_type(), + expected.len(), + expected.null_count(), + actual.data_type(), + actual.len(), + actual.null_count(), ); - Ok(()) + + if expected.data_type() != actual.data_type() || expected.len() != actual.len() { + return description; + } + + // Binary search for the shortest prefix that already differs; its last element is the + // first mismatching row. + let (mut low, mut high) = (0usize, expected.len()); + while low < high { + let mid = low + (high - low) / 2 + 1; + if expected.slice(0, mid).to_data() == actual.slice(0, mid).to_data() { + low = mid; + } else { + high = mid - 1; + } + } + + if low < expected.len() { + description.push_str(&format!( + "\n first difference at row {low}:\n cpu: {:?}\n gpu: {:?}", + expected.slice(low, 1), + actual.slice(low, 1), + )); + } + + description } From 335d14af3e5cf5460af4167a9070bc2d23f519b3 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:19:11 +0000 Subject: [PATCH 04/16] fix(cuda): add the frame of reference to bit-packed patch values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bit-unpack kernel writes patch values straight into the output while the lane decoder adds the frame of reference to every unpacked value. Bit-packing exceptions are stored in the same reference-relative domain as the packed values, so under `FoR(BitPacked)` every patched position came out short by exactly the reference. The existing kernel tests could not catch this: they exercise `BitPacked` directly, where the reference is zero. The new `FoRExecutor` case bit-packs to 8 bits with values that overflow into patches and a non-zero reference. Found by the compression benchmark's new `--gpu-verify` pass, which reported a `fastlanes.for` field decoding row 8038 as 131072 where the CPU produced 393061 — a difference of exactly the 261989 reference. Also thread the dataset name through compress-bench failures, so a benchmark error says which dataset it came from. Signed-off-by: Claude --- benchmarks/compress-bench/src/main.rs | 7 +++- vortex-cuda/kernels/src/bit_unpack_16.cu | 9 ++++- vortex-cuda/kernels/src/bit_unpack_32.cu | 9 ++++- vortex-cuda/kernels/src/bit_unpack_64.cu | 9 ++++- vortex-cuda/kernels/src/bit_unpack_8.cu | 9 ++++- vortex-cuda/src/bit_unpack_gen.rs | 9 ++++- vortex-cuda/src/kernel/encodings/for_.rs | 47 ++++++++++++++++++++++++ 7 files changed, 87 insertions(+), 12 deletions(-) diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index a915112897f..63bcfac7025 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -4,6 +4,7 @@ use std::path::PathBuf; use std::time::Duration; +use anyhow::Context; use clap::Parser; #[cfg(feature = "lance")] use compress_bench::LanceCompressor; @@ -356,7 +357,8 @@ async fn run_benchmark_for_dataset( iterations, bench_name, ) - .await?; + .await + .with_context(|| format!("compressing {bench_name} as {format}"))?; compressed_sizes.insert(*format, result.compressed_size); let all_runs_ns: Vec = result .all_runs @@ -387,7 +389,8 @@ async fn run_benchmark_for_dataset( iterations, bench_name, ) - .await?; + .await + .with_context(|| format!("decompressing {bench_name} as {format}"))?; let all_runs_ns: Vec = result .all_runs .iter() diff --git a/vortex-cuda/kernels/src/bit_unpack_16.cu b/vortex-cuda/kernels/src/bit_unpack_16.cu index a784df201d3..43acff3d08d 100644 --- a/vortex-cuda/kernels/src/bit_unpack_16.cu +++ b/vortex-cuda/kernels/src/bit_unpack_16.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_16_device(const uint16_t *__restrict in, uint16_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_32.cu b/vortex-cuda/kernels/src/bit_unpack_32.cu index 3f8fcb5c227..b9f29c72c3a 100644 --- a/vortex-cuda/kernels/src/bit_unpack_32.cu +++ b/vortex-cuda/kernels/src/bit_unpack_32.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_32_device(const uint32_t *__restrict in, uint32_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_64.cu b/vortex-cuda/kernels/src/bit_unpack_64.cu index ebe0b125369..aebe500c653 100644 --- a/vortex-cuda/kernels/src/bit_unpack_64.cu +++ b/vortex-cuda/kernels/src/bit_unpack_64.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_64_device(const uint64_t *__restrict in, uint64_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 16); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_8.cu b/vortex-cuda/kernels/src/bit_unpack_8.cu index b2fcfd26f04..cabea862e59 100644 --- a/vortex-cuda/kernels/src/bit_unpack_8.cu +++ b/vortex-cuda/kernels/src/bit_unpack_8.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_8_device(const uint8_t *__restrict in, uint8_t *__re } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/src/bit_unpack_gen.rs b/vortex-cuda/src/bit_unpack_gen.rs index 2482c0996b8..7b3ce148753 100644 --- a/vortex-cuda/src/bit_unpack_gen.rs +++ b/vortex-cuda/src/bit_unpack_gen.rs @@ -152,11 +152,16 @@ __device__ void _bit_unpack_{bits}_device(const uint{bits}_t *__restrict in, uin }} __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, {thread_count}); auto patch = cursor.next(); while (patch.index != FL_CHUNK) {{ - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); }} __syncwarp(); diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index fefa4ff5ef6..cc3ceed2869 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -158,6 +158,7 @@ mod tests { use vortex::buffer::Buffer; use vortex::dtype::NativePType; use vortex::encodings::fastlanes::BitPacked; + use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArray; use vortex::error::VortexExpect; @@ -228,4 +229,50 @@ mod tests { assert_arrays_eq!(for_array, gpu_result, &mut ctx); } + + /// Patched positions must pick up the frame of reference, exactly like unpacked ones. + /// + /// The bit-packed exceptions are stored reference-relative, so a decoder that writes them + /// straight into the output leaves every patched value short by the reference. A plain + /// bit-packed array cannot catch that: its reference is zero. + #[rstest] + #[case::u32(100_000u32)] + #[case::u64(1_000_000u64)] + #[crate::test] + async fn test_ffor_patched_values_include_reference(#[case] reference: T) -> VortexResult<()> + where + T: NativePType + Into + From, + { + let mut ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + // Values that fit in 8 bits, with a handful that do not and so become patches. + let mut values = (0..2048u32) + .map(|i| >::from(i % 200)) + .collect::>(); + for index in [7, 1023, 1024, 2047] { + values[index] = >::from(1u32 << 17); + } + + let values = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); + let packed = BitPacked::encode(&values, 8, &mut array_session().create_execution_ctx())?; + assert!( + packed.patches().is_some(), + "test setup expects the exceptions to be stored as patches" + ); + let for_array = FoR::try_new(packed.into_array(), reference.into())?; + + let gpu_result = FoRExecutor + .execute(for_array.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(for_array, gpu_result, &mut ctx); + + Ok(()) + } } From 5f297d4a716f4eff28a84b6178638ad600b57718 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:29:19 +0000 Subject: [PATCH 05/16] Survey every dataset in a GPU verification pass A verification run stopped at the first dataset that failed, so finding the GPU-clean set took one CI cycle per dataset. Run every dataset instead, recording failures and reporting them together at the end, then exit non-zero. Missing CUDA kernel support surfaces as a panic rather than an error, so the survey catches those too. Signed-off-by: Claude --- benchmarks/compress-bench/src/main.rs | 67 +++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 5 deletions(-) diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 63bcfac7025..6fb003f97e3 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; +use std::panic::AssertUnwindSafe; use std::path::PathBuf; use std::time::Duration; @@ -15,6 +17,7 @@ use compress_bench::gpu_vortex::GpuVortexCompressor; use compress_bench::parquet::ParquetCompressor; use compress_bench::parquet_pages::GpuCodec; use compress_bench::vortex::VortexCompressor; +use futures::FutureExt; use indicatif::ProgressBar; use itertools::Itertools; use regex::Regex; @@ -286,18 +289,61 @@ async fn run_compress( let mut measurements = vec![]; let mut v3_records: Vec = Vec::new(); + // A verification pass reports on every dataset rather than stopping at the first failure: + // one run then says exactly which datasets decode correctly on the GPU and which do not. + let survey_all = gpu.is_some_and(|gpu| gpu.verify); + let mut failures: Vec<(String, anyhow::Error)> = Vec::new(); + for dataset_handle in datasets.into_iter() { - let (m, mut records) = - run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle, gpu) - .await?; - measurements.push(m); - v3_records.append(&mut records); + let run = + run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle, gpu); + + // Missing CUDA kernel support surfaces as a panic rather than an error, so the survey + // has to catch those too or the first unsupported dataset ends the run. + let result = if survey_all { + match AssertUnwindSafe(run).catch_unwind().await { + Ok(result) => result, + Err(panic) => Err(anyhow::anyhow!("panicked: {}", panic_message(&panic))), + } + } else { + run.await + }; + + match result { + Ok((m, mut records)) => { + measurements.push(m); + v3_records.append(&mut records); + } + Err(error) if survey_all => { + tracing::error!("{}: {error:#}", dataset_handle.name()); + failures.push((dataset_handle.name().to_string(), error)); + } + Err(error) => return Err(error), + } } let measurements = CompressMeasurements::from_iter(measurements); progress.finish(); + if !failures.is_empty() { + eprintln!( + "\nGPU verification failed for {} dataset(s):", + failures.len() + ); + for (dataset, error) in &failures { + eprintln!(" - {dataset}: {error:#}"); + } + anyhow::bail!( + "GPU verification failed for: {}", + failures + .iter() + .map(|(dataset, _)| dataset.as_str()) + .collect::>() + .join(", ") + ); + } + if let Some(path) = ingest_output { v3::write_jsonl_to_path(&path, &v3_records)?; } @@ -457,3 +503,14 @@ fn push_gpu_ratio( value: vortex_time.as_nanos() as f64 / parquet_time.as_nanos() as f64, }); } + +/// Extracts the message from a caught panic payload. +fn panic_message(panic: &Box) -> String { + if let Some(message) = panic.downcast_ref::<&str>() { + (*message).to_string() + } else if let Some(message) = panic.downcast_ref::() { + message.clone() + } else { + "non-string panic payload".to_string() + } +} From b0eda650245bbde23d5e644ee5c9c7308386f8ea Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:40:24 +0000 Subject: [PATCH 06/16] CI: publish the GPU verification matrix to the PR The per-dataset verification verdicts were only visible by digging through a multi-thousand-line job log. Capture the verification output, publish the per-dataset results to the step summary and a PR comment, and keep failing the job through a separate gate step. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 44 +++++++++++++++++++-- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 45bc59fb317..94fb085eb31 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -45,15 +45,53 @@ jobs: - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh - name: Verify GPU decompression correctness + id: verify shell: bash + continue-on-error: true env: - RUST_BACKTRACE: full + RUST_BACKTRACE: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" # Cross-checks every GPU-decompressed page and field against the CPU decoders before - # any timings are taken. Verification runs inline, so this pass is not timed. + # any timings are taken. Verification runs inline, so this pass is not timed. It runs + # every dataset rather than stopping at the first failure, so one run reports the whole + # matrix; the gate step below still fails the job. run: | + set -o pipefail target/release_debug/compress-bench \ - --gpu-decompress --gpu-verify --iterations 1 -d table + --gpu-decompress --gpu-verify --iterations 1 -d table 2>&1 | tee gpu-verify.txt + - name: Publish verification result + shell: bash + # The per-dataset verdicts are worth surfacing on the PR whether or not they all pass: + # digging them out of a multi-thousand-line job log is otherwise the only way to see + # which encodings decode correctly on the GPU. + run: | + { + echo "# GPU decompression verification" + echo + if [ "${{ steps.verify.outcome }}" = "success" ]; then + echo "All GPU datasets matched the CPU decode." + else + echo "Verification failed. Per-dataset results:" + fi + echo + echo '```text' + grep -E "verified [0-9]+|GPU verification failed|^ - " gpu-verify.txt | tail -40 \ + || tail -40 gpu-verify.txt + echo '```' + } > verify-comment.md + cat verify-comment.md >> "$GITHUB_STEP_SUMMARY" + - name: Comment PR with verification result + if: github.event.pull_request.head.repo.fork == false + uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b # v3 + with: + file-path: verify-comment.md + comment-tag: bench-pr-comment-gpu-verify + - name: Fail if verification failed + if: steps.verify.outcome == 'failure' + shell: bash + run: | + echo "GPU decompression verification failed; see the verification comment." >&2 + exit 1 - name: Run GPU compression benchmark shell: bash env: From 66c75a50aa7148aa9a6628a77524e8dc010db9fd Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:56:54 +0000 Subject: [PATCH 07/16] CI: install DuckDB for the GPU compression benchmark The Public BI datasets build their Parquet fixture through the DuckDB CLI, as in bench-pr.yml. The GPU job never installed it, so all six failed with ENOENT before reaching the GPU at all. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 94fb085eb31..7e60b9e1970 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -32,6 +32,13 @@ 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 bench-pr.yml. Without it those datasets fail to materialise 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" - uses: ./.github/actions/system-info - name: Display NVIDIA GPU details run: | From 64d481c424429539932e0ab1aa557248db7f89b2 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 14:50:15 +0000 Subject: [PATCH 08/16] Time the GPU Parquet number with cuDF instead of nvCOMP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The nvCOMP backend only ran the codec stage on the device: page decoding stayed on the CPU and was excluded from the measurement, so the Parquet figure was an upper bound and the comparison against Vortex was not like-for-like. cuDF's `read_parquet` does the whole read on the device — page header decode, decompression, dictionary/RLE/plain decoding and column assembly — which is the same amount of work the Vortex backend does when it decodes to canonical arrays. It is reached through the prebuilt `cudf-cu12` wheel, so it stays a runtime dependency and never enters the Rust build. Timing is taken inside scripts/cudf-parquet-read.py, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up read runs first. `--gpu-verify` now compares the cuDF frame against a CPU Parquet read. This removes the page scanner, the batched nvCOMP launch path and the nvCOMP Snappy bindings, all of which existed only to serve the codec-stage backend. What remains of the Parquet side is the GPU-friendly writer settings, now in gpu_writer.rs. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 12 + Cargo.lock | 6 +- Cargo.toml | 1 - benchmarks/compress-bench/Cargo.toml | 11 +- benchmarks/compress-bench/README.md | 58 +- benchmarks/compress-bench/src/gpu_parquet.rs | 525 ++--------------- benchmarks/compress-bench/src/gpu_writer.rs | 88 +++ benchmarks/compress-bench/src/lib.rs | 2 +- benchmarks/compress-bench/src/main.rs | 12 +- .../compress-bench/src/parquet_pages.rs | 530 ------------------ scripts/cudf-parquet-read.py | 101 ++++ vortex-cuda/nvcomp/build.rs | 7 - vortex-cuda/nvcomp/src/backend.rs | 52 -- vortex-cuda/nvcomp/src/lib.rs | 4 - vortex-cuda/nvcomp/src/snappy.rs | 222 -------- vortex-cuda/nvcomp/src/zstd.rs | 47 +- 16 files changed, 318 insertions(+), 1360 deletions(-) create mode 100644 benchmarks/compress-bench/src/gpu_writer.rs delete mode 100644 benchmarks/compress-bench/src/parquet_pages.rs create mode 100644 scripts/cudf-parquet-read.py delete mode 100644 vortex-cuda/nvcomp/src/backend.rs delete mode 100644 vortex-cuda/nvcomp/src/snappy.rs diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 7e60b9e1970..36f16223105 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -39,6 +39,18 @@ jobs: 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 full cuDF `read_parquet`. cuDF ships prebuilt manylinux + # wheels on NVIDIA's index, so it stays a runtime dependency and never enters the Rust + # build. pandas and pyarrow back the --gpu-verify cross-check. + 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: | diff --git a/Cargo.lock b/Cargo.lock index e7c6dd2dfdb..d599793b8c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1566,15 +1566,14 @@ dependencies = [ "async-trait", "bytes", "clap", - "cudarc", "futures", "indicatif", "itertools 0.14.0", "lance-bench", "parquet 58.4.0", "regex", - "rstest", - "snap", + "serde", + "serde_json", "tempfile", "tokio", "tracing", @@ -1582,7 +1581,6 @@ dependencies = [ "vortex-arrow", "vortex-bench", "vortex-cuda", - "zstd", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index edf730760b4..997b3281d25 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -245,7 +245,6 @@ similar = "3.0.0" sketches-ddsketch = "0.4.0" smallvec = "1.15.1" smol = "2.0.2" -snap = "1.1.2" spatialbench = "0.2" spatialbench-arrow = "0.2" # spatialbench still pins arrow 56, two majors behind the workspace arrow. Until upstream diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index c20e2861b2d..7b1dadc3209 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -21,14 +21,14 @@ arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } clap = { workspace = true, features = ["derive"] } -cudarc = { workspace = true, optional = true } futures = { workspace = true } indicatif = { workspace = true } itertools = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { workspace = true } -snap = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } tempfile = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } @@ -36,14 +36,9 @@ vortex = { workspace = true } vortex-arrow = { workspace = true } vortex-bench = { workspace = true } vortex-cuda = { workspace = true, optional = true } -zstd = { workspace = true } - -[dev-dependencies] -rstest = { workspace = true } -tempfile = { workspace = true } [features] -cuda = ["dep:cudarc", "dep:tempfile", "dep:vortex-cuda"] +cuda = ["dep:tempfile", "dep:vortex-cuda"] lance = ["dep:lance-bench"] unstable_encodings = ["vortex/unstable_encodings", "vortex-cuda?/unstable_encodings"] diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 132ce153450..e14c117d5a6 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -23,55 +23,58 @@ GPU dataset list in `src/main.rs`. It measures decompression only, for two backe - **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 every - page body is decompressed on the device with nvCOMP's batched Snappy or Zstd entrypoints, - the same decomposition cuDF's Parquet reader uses: column chunks are staged on the device, - then all pages go through a single batched launch. +- **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- 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 (default: snappy) +# 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 ``` -On Linux both backends read through the same pinned, direct-I/O (`O_DIRECT`) reader, so -repeated iterations measure storage bandwidth rather than page-cache hits. +### cuDF -### What the Parquet GPU number does and does not include +cuDF is reached through its prebuilt `cudf-cu12` wheel, so it is a runtime dependency of the +benchmark and never enters the Rust build: -Included: column chunk I/O, the host-to-device transfer, and the batched codec launch. +```bash +uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarrow +``` -Not included: page *decoding* — the dictionary, RLE and plain decoders that turn a -decompressed page into an Arrow array — because there is no Rust GPU Parquet page decoder to -call. The Vortex backend it is compared against decodes all the way to canonical arrays, so -the Parquet figure is an upper bound on what a full GPU Parquet reader could reach, and the -`vortex:parquet- gpu ratio decompress time` metric is biased in Parquet's favour. +`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. -Walking the per-page Thrift headers also happens on the host, once per file, outside the -measurement; a real GPU Parquet reader decodes page headers on the device. +Note that the two backends do not share an I/O path: the Vortex reader uses pinned buffers and +direct I/O (`O_DIRECT`) on Linux, while cuDF does its own host read and host-to-device +transfer. ### GPU-friendly Parquet writer settings -Set in `src/parquet_pages.rs`: +Set in `src/gpu_writer.rs`: | Setting | Value | Why | | --- | --- | --- | -| writer version | `PARQUET_1_0` | v1 pages compress the whole page body, which is the unit nvCOMP decompresses. v2 pages put uncompressed levels ahead of the compressed values in the same body. | -| compression | Snappy (default) or Zstd | The two Parquet codecs nvCOMP implements. Snappy has the higher device throughput and is the Parquet default. | +| 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-chunk setup, small enough to keep every SM fed. Matches the page size cuDF targets. | +| 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 that have to be walked on the host. | +| statistics | chunk-level | Page statistics only inflate the headers a reader has to walk. | ### Correctness `--gpu-verify` cross-checks device output against the CPU decoders on every iteration: -- Parquet: each decompressed page is copied back and compared byte-for-byte against the host - Snappy/Zstd output for the same compressed bytes. +- 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. @@ -83,10 +86,5 @@ cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress --gpu-verify --iterations 1 ``` -Independently of `--gpu-verify`, every Parquet GPU run checks nvCOMP's per-page status and -output-size arrays after the measurement, because a batched launch reports per-page failures -in device memory rather than by failing the launch. - -Page-header scanning is covered by CPU-only unit tests in `src/parquet_pages.rs`, which -assert the located page bodies decompress to exactly the bytes `parquet`'s own page reader -produces. +A verifying 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. diff --git a/benchmarks/compress-bench/src/gpu_parquet.rs b/benchmarks/compress-bench/src/gpu_parquet.rs index b8e6195d784..3ee0297f1a0 100644 --- a/benchmarks/compress-bench/src/gpu_parquet.rs +++ b/benchmarks/compress-bench/src/gpu_parquet.rs @@ -1,89 +1,70 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! GPU Parquet decompression backend. +//! GPU Parquet decompression backend, timed through cuDF. //! -//! Parquet's compressed unit is the page, and every page body is an independent block of -//! Snappy or Zstd. That is precisely the batch shape nvCOMP's device decompressors take, and -//! it is how cuDF's Parquet reader gets pages off the CPU: read the column chunks to the -//! device, then decompress every page in one batched launch. +//! 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. //! -//! What this measures is the decompression stage of a Parquet read — column chunk I/O, -//! host-to-device transfer, and the batched codec launch. Page *decoding* (the dictionary, -//! RLE and plain decoders that turn a decompressed page into an Arrow array) is not -//! included, because there is no Rust GPU Parquet page decoder to call. The Vortex GPU -//! backend it is compared against decodes all the way to canonical arrays, so the Parquet -//! numbers here are an upper bound on what a full GPU Parquet reader could achieve, and the -//! comparison is favourable to Parquet. +//! 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::fs::File; 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::anyhow; +use anyhow::bail; use anyhow::ensure; use arrow_array::RecordBatch; use async_trait::async_trait; -use cudarc::driver::CudaSlice; -use cudarc::driver::DevicePtr; -use cudarc::driver::DevicePtrMut; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::stream; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; -use parquet::file::metadata::ParquetMetaData; -use parquet::file::metadata::ParquetMetaDataReader; +use serde::Deserialize; use tempfile::NamedTempFile; -use vortex::array::buffer::BufferHandle; -use vortex::array::buffer::DeviceBuffer; -use vortex::buffer::Alignment; -use vortex::buffer::Buffer; -use vortex::error::vortex_err; -use vortex::io::VortexReadAt; -use vortex::io::session::RuntimeSessionExt; use vortex_bench::Format; -use vortex_bench::SESSION; use vortex_bench::compress::Compressor; -use vortex_cuda::CudaBufferExt; -use vortex_cuda::CudaDeviceBuffer; -use vortex_cuda::CudaExecutionCtx; -use vortex_cuda::CudaSession; -use vortex_cuda::CudaSessionExt; -use vortex_cuda::PooledFileReadAt; -use vortex_cuda::PooledFileReadAtOptions; -use vortex_cuda::nvcomp::AlignmentRequirements; -use vortex_cuda::nvcomp::snappy; -use vortex_cuda::nvcomp::sys; -use vortex_cuda::nvcomp::sys::nvcompStatus_t; -use vortex_cuda::nvcomp::zstd as nvcomp_zstd; -use crate::parquet_pages::ColumnChunkPages; -use crate::parquet_pages::GpuCodec; -use crate::parquet_pages::gpu_writer_properties; -use crate::parquet_pages::scan_compressed_pages; +use crate::gpu_writer::GpuCodec; +use crate::gpu_writer::gpu_writer_properties; -/// Parquet compressor whose decompression measurement runs the page codec on the GPU. +/// 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 decompresses them with nvCOMP. + /// Create a backend that writes pages with `codec` and times cuDF reading them back. /// - /// When `verify` is set, every decompressed page is copied back and compared against the - /// host codec's output before the measurement is reported. + /// 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(File::open(parquet_path)?)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(std::fs::File::open(parquet_path)?)?; let schema = Arc::clone(builder.schema()); let batches: Vec = builder.build()?.collect::, _>>()?; @@ -101,21 +82,6 @@ impl GpuParquetCompressor { writer.close()?; Ok((output, size)) } - - fn check_codec(&self, metadata: &ParquetMetaData) -> Result<()> { - for row_group in metadata.row_groups() { - for column in row_group.columns() { - ensure!( - self.codec.matches(column.compression()), - "column {} was written with {:?}, expected {}", - column.column_path(), - column.compression(), - self.codec.name() - ); - } - } - Ok(()) - } } #[async_trait] @@ -132,420 +98,43 @@ impl Compressor for GpuParquetCompressor { async fn decompress(&self, parquet_path: &Path) -> Result { let (gpu_file, _) = self.write_gpu_parquet(parquet_path)?; - - let file = File::open(gpu_file.path())?; - let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; - self.check_codec(&metadata)?; - drop(file); - - // The page map is a property of the file, so it is resolved once rather than on every - // iteration. A GPU Parquet reader decodes page headers on the device as part of the - // read; this benchmark cannot, so the host walk is kept out of the measurement. - let file_bytes = std::fs::read(gpu_file.path())?; - let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; - let plan = DecompressPlan::build(&row_groups, self.codec)?; - // Only verification needs the host copy afterwards, to decompress the same bytes. - let file_bytes = self.verify.then_some(file_bytes); - - let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; - let reader = open_reader(gpu_file.path(), &cuda_ctx)?; - - // Row groups are staged and released one at a time so device residency stays bounded - // by the largest row group rather than by the whole file, matching how the Vortex - // backend streams one batch at a time. - let mut elapsed = Duration::ZERO; - let mut checks = Vec::with_capacity(plan.row_groups.len()); - for row_group in &plan.row_groups { - let start = Instant::now(); - let device_chunks = stage_row_group(row_group, &reader).await?; - let output = - decompress_pages(row_group, &device_chunks, self.codec, &mut cuda_ctx).await?; - cuda_ctx.synchronize_stream()?; - elapsed += start.elapsed(); - - let DecompressOutput { - output, - actual_sizes, - statuses, - } = output; - drop(device_chunks); - - if self.verify { - // Verified inline, since holding every row group's decompressed output would - // defeat the point of streaming them. This is why a verifying run's timings - // are not comparable. - let file_bytes = file_bytes - .as_deref() - .ok_or_else(|| anyhow!("verification requires the host file bytes"))?; - verify_against_host(row_group, file_bytes, self.codec, output).await?; - } - checks.push((actual_sizes, statuses)); - } - - // Checked outside the measurement, but on every iteration: a batched launch reports - // per-page failures in device memory rather than by failing the launch, so without - // this a broken run would simply look fast. - for (row_group, (actual_sizes, statuses)) in plan.row_groups.iter().zip(checks) { - check_statuses(row_group, actual_sizes, statuses).await?; - } - - Ok(elapsed) - } -} - -/// Stages a row group's column chunks on the device, bounding concurrent pinned staging. -async fn stage_row_group( - row_group: &RowGroupPlan, - reader: &PooledFileReadAt, -) -> Result> { - let reads = row_group - .chunks - .iter() - .map(|chunk| reader.read_at(chunk.offset, chunk.len, Alignment::of::())) - .collect::>(); - - stream::iter(reads) - .buffered(STAGING_CONCURRENCY) - .try_collect::>() - .await - .map_err(|e| anyhow!("failed to stage Parquet column chunks on the device: {e}")) -} - -/// Column chunks staged on the device at once. Matches `vortex-cuda`'s file read concurrency. -const STAGING_CONCURRENCY: usize = 32; - -/// Opens the GPU Parquet file through the same pinned, direct-I/O reader the Vortex GPU -/// backend uses, so both formats measure storage bandwidth rather than page-cache hits. -fn open_reader(path: &Path, cuda_ctx: &CudaExecutionCtx) -> Result { - let pool = Arc::clone(SESSION.cuda_session().pinned_buffer_pool()); - let options = PooledFileReadAtOptions::default(); - #[cfg(target_os = "linux")] - let options = options.with_direct_io(); - - Ok(PooledFileReadAt::open_with_options( - path, - SESSION.handle(), - pool, - cuda_ctx.stream().clone(), - options, - )?) -} - -/// A page to decompress, addressed relative to its column chunk's device buffer. -struct PlannedPage { - chunk: usize, - offset_in_chunk: usize, - compressed_len: usize, - uncompressed_len: usize, - output_offset: usize, -} - -/// The byte range of a column chunk to stage on the device. -struct PlannedChunk { - offset: u64, - len: usize, -} - -/// Everything needed to issue one batched nvCOMP launch over a row group's pages. -struct RowGroupPlan { - chunks: Vec, - pages: Vec, - output_len: usize, - max_uncompressed: usize, -} - -/// The per-row-group work for one file. -struct DecompressPlan { - row_groups: Vec, -} - -impl DecompressPlan { - fn build(row_groups: &[Vec], codec: GpuCodec) -> Result { - let alignment = decompress_alignments(codec)?; - ensure!( - alignment.output.is_power_of_two(), - "nvcomp reported a non-power-of-two output alignment of {}", - alignment.output - ); - - let row_groups = row_groups - .iter() - .map(|chunks| RowGroupPlan::build(chunks, codec, alignment)) - .collect::>>()?; + let report = run_cudf_read(gpu_file.path(), self.verify)?; ensure!( - row_groups.iter().any(|plan| !plan.pages.is_empty()), - "Parquet file contains no compressed pages" + report.rows > 0 && report.columns > 0, + "cuDF read {} rows and {} columns, expected a non-empty table", + report.rows, + report.columns ); - Ok(Self { row_groups }) + Ok(Duration::from_nanos(report.min_ns)) } } -impl RowGroupPlan { - fn build( - chunks: &[ColumnChunkPages], - codec: GpuCodec, - alignment: AlignmentRequirements, - ) -> Result { - let mut planned_chunks = Vec::with_capacity(chunks.len()); - let mut pages = Vec::new(); - let mut output_len = 0usize; - let mut max_uncompressed = 0usize; - - for (index, chunk) in chunks.iter().enumerate() { - planned_chunks.push(PlannedChunk { - offset: chunk.offset, - len: chunk.len, - }); - - for page in &chunk.pages { - let offset_in_chunk = usize::try_from( - u64::try_from(page.offset)? - .checked_sub(chunk.offset) - .ok_or_else(|| anyhow!("page offset precedes its column chunk"))?, - )?; - // Pages are decompressed in place from their column chunk's device buffer. - // CUDA allocations are at least 256-byte aligned, so a page's device address - // meets nvcomp's requirement exactly when its chunk-relative offset does. - ensure!( - offset_in_chunk.is_multiple_of(alignment.input), - "page at file offset {} sits {offset_in_chunk} bytes into its column chunk, \ - which does not meet nvcomp's {} byte input alignment for {}; \ - use --gpu-parquet-codec snappy", - page.offset, - alignment.input, - codec.name() - ); - - let output_offset = output_len.next_multiple_of(alignment.output); - output_len = output_offset + page.uncompressed_len; - max_uncompressed = max_uncompressed.max(page.uncompressed_len); - - pages.push(PlannedPage { - chunk: index, - offset_in_chunk, - compressed_len: page.compressed_len, - uncompressed_len: page.uncompressed_len, - output_offset, - }); - } - } - - Ok(Self { - chunks: planned_chunks, - pages, - output_len, - max_uncompressed, - }) - } -} - -fn decompress_alignments(codec: GpuCodec) -> Result { - match codec { - GpuCodec::Snappy => { - snappy::decompress_alignment_requirements(snappy::SnappyDecompressOpts::default()) - } - GpuCodec::Zstd => nvcomp_zstd::decompress_alignment_requirements( - nvcomp_zstd::ZstdDecompressOpts::default(), - ), +/// Runs the cuDF read script and returns the timing it measured. +fn run_cudf_read(path: &Path, verify: bool) -> Result { + let mut command = Command::new("python3"); + command.arg(CUDF_SCRIPT).arg(path); + if verify { + command.arg("--verify"); } - .map_err(|e| anyhow!("nvcomp alignment query failed: {e}")) -} - -/// Device-side outputs of a batched decompression launch. -struct DecompressOutput { - output: CudaSlice, - actual_sizes: CudaSlice, - statuses: CudaSlice, -} - -/// Enqueues the batched decompression of every page in `plan` onto the context's stream. -async fn decompress_pages( - plan: &RowGroupPlan, - device_chunks: &[BufferHandle], - codec: GpuCodec, - ctx: &mut CudaExecutionCtx, -) -> Result { - let num_pages = plan.pages.len(); - - let temp_size = match codec { - GpuCodec::Snappy => { - snappy::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) - } - GpuCodec::Zstd => { - nvcomp_zstd::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) - } - } - .map_err(|e| anyhow!("nvcomp temp size query failed: {e}"))?; - - let chunk_bases = device_chunks - .iter() - .map(|handle| handle.cuda_device_ptr()) - .collect::, _>>()?; - - let mut output = ctx.device_alloc::(plan.output_len)?; - // Only the allocation address is needed to build the output pointer table; the device - // write itself is tracked by the guard taken around the launch below. - let output_base = { - let (base, _) = output.device_ptr(ctx.stream()); - base - }; - - let mut compressed_ptrs = Vec::with_capacity(num_pages); - let mut compressed_sizes = Vec::with_capacity(num_pages); - let mut uncompressed_sizes = Vec::with_capacity(num_pages); - let mut output_ptrs = Vec::with_capacity(num_pages); - for page in &plan.pages { - compressed_ptrs.push(chunk_bases[page.chunk] + page.offset_in_chunk as u64); - compressed_sizes.push(page.compressed_len); - uncompressed_sizes.push(page.uncompressed_len); - output_ptrs.push(output_base + page.output_offset as u64); - } - - let (compressed_ptrs, compressed_sizes, uncompressed_sizes, output_ptrs) = futures::try_join!( - ctx.copy_to_device(compressed_ptrs)?, - ctx.copy_to_device(compressed_sizes)?, - ctx.copy_to_device(uncompressed_sizes)?, - ctx.copy_to_device(output_ptrs)? - )?; - - let mut actual_sizes: CudaSlice = ctx.device_alloc(num_pages)?; - let mut statuses: CudaSlice = ctx.device_alloc(num_pages)?; - let mut temp: CudaSlice = ctx.device_alloc(temp_size)?; - - let stream = ctx.stream(); - let compressed_ptrs_view = compressed_ptrs.cuda_view::()?; - let compressed_sizes_view = compressed_sizes.cuda_view::()?; - let uncompressed_sizes_view = uncompressed_sizes.cuda_view::()?; - let output_ptrs_view = output_ptrs.cuda_view::()?; - let (compressed_ptrs_ptr, record_compressed_ptrs) = compressed_ptrs_view.device_ptr(stream); - let (compressed_sizes_ptr, record_compressed_sizes) = compressed_sizes_view.device_ptr(stream); - let (uncompressed_sizes_ptr, record_uncompressed_sizes) = - uncompressed_sizes_view.device_ptr(stream); - let (output_ptrs_ptr, record_output_ptrs) = output_ptrs_view.device_ptr(stream); - let (_output_ptr, record_output) = output.device_ptr_mut(stream); - let (actual_sizes_ptr, record_actual_sizes) = actual_sizes.device_ptr_mut(stream); - let (statuses_ptr, record_statuses) = statuses.device_ptr_mut(stream); - let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); - - ctx.launch_external(plan.output_len, || { - // SAFETY: every pointer is derived from a live device allocation sized by the plan, - // and each batch metadata array holds exactly `num_pages` entries. - unsafe { - match codec { - GpuCodec::Snappy => snappy::decompress_async( - compressed_ptrs_ptr as _, - compressed_sizes_ptr as _, - uncompressed_sizes_ptr as _, - actual_sizes_ptr as _, - num_pages, - temp_ptr as _, - temp_size, - output_ptrs_ptr as _, - statuses_ptr as _, - stream.cu_stream().cast(), - ), - GpuCodec::Zstd => nvcomp_zstd::decompress_async( - compressed_ptrs_ptr as _, - compressed_sizes_ptr as _, - uncompressed_sizes_ptr as _, - actual_sizes_ptr as _, - num_pages, - temp_ptr as _, - temp_size, - output_ptrs_ptr as _, - statuses_ptr as _, - stream.cu_stream().cast(), - ), - } - .map_err(|e| vortex_err!("nvcomp decompress_async failed: {}", e)) - } + let output = command.output().with_context(|| { + format!("failed to run {CUDF_SCRIPT}; is cudf-cu12 installed on this host?") })?; - drop(( - record_compressed_ptrs, - record_compressed_sizes, - record_uncompressed_sizes, - record_output_ptrs, - record_output, - record_actual_sizes, - record_statuses, - record_temp, - )); - // The temporary workspace must outlive the launch, which the stream ordering guarantees - // only while the allocation is alive. - drop(temp); - - Ok(DecompressOutput { - output, - actual_sizes, - statuses, - }) -} - -/// Copies the per-page status and size arrays back and fails on any mismatch. -async fn check_statuses( - plan: &RowGroupPlan, - actual_sizes: CudaSlice, - statuses: CudaSlice, -) -> Result<()> { - let statuses = CudaDeviceBuffer::new(statuses) - .copy_to_host(Alignment::of::())? - .await?; - let actual_sizes = CudaDeviceBuffer::new(actual_sizes) - .copy_to_host(Alignment::of::())? - .await?; - - let statuses = Buffer::::from_byte_buffer(statuses); - let actual_sizes = Buffer::::from_byte_buffer(actual_sizes); - - for (index, page) in plan.pages.iter().enumerate() { - let status = statuses.as_slice()[index]; - ensure!( - status == sys::nvcompStatus_t_nvcompSuccess, - "page {index} failed to decompress with nvcomp status {status}" - ); - let actual = actual_sizes.as_slice()[index]; - ensure!( - actual == page.uncompressed_len, - "page {index} decompressed to {actual} bytes, expected {}", - page.uncompressed_len - ); - } - - Ok(()) -} - -/// Compares every decompressed page against the host codec's output for the same bytes. -async fn verify_against_host( - plan: &RowGroupPlan, - file_bytes: &[u8], - codec: GpuCodec, - output: CudaSlice, -) -> Result<()> { - let device_output = CudaDeviceBuffer::new(output) - .copy_to_host(Alignment::of::())? - .await?; - let device_output = device_output.as_slice(); - - for (index, page) in plan.pages.iter().enumerate() { - let chunk = &plan.chunks[page.chunk]; - let start = usize::try_from(chunk.offset)? + page.offset_in_chunk; - let compressed = &file_bytes[start..start + page.compressed_len]; - let expected = codec.decompress_host(compressed, page.uncompressed_len)?; - let actual = &device_output[page.output_offset..page.output_offset + page.uncompressed_len]; - ensure!( - actual == expected.as_slice(), - "page {index} decompressed on the GPU differs from the host codec output" + if !output.status.success() { + bail!( + "{CUDF_SCRIPT} exited with {}:\n{}", + output.status, + String::from_utf8_lossy(&output.stderr).trim() ); } - tracing::info!( - "verified {} GPU-decompressed {} pages against the host codec", - plan.pages.len(), - codec.name() - ); - Ok(()) + serde_json::from_slice(&output.stdout).with_context(|| { + format!( + "could not parse the report from {CUDF_SCRIPT}: {}", + String::from_utf8_lossy(&output.stdout).trim() + ) + }) } diff --git a/benchmarks/compress-bench/src/gpu_writer.rs b/benchmarks/compress-bench/src/gpu_writer.rs new file mode 100644 index 00000000000..49052a5d6b4 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_writer.rs @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Parquet writer settings for the GPU benchmark. +//! +//! The GPU backend rewrites each dataset before reading it back with cuDF, so the file it +//! reads is written the way a GPU reader wants it rather than the way the CPU suite writes it. + +use clap::ValueEnum; +use parquet::basic::Compression; +use parquet::basic::ZstdLevel; +use parquet::file::properties::EnabledStatistics; +use parquet::file::properties::WriterProperties; +use parquet::file::properties::WriterVersion; + +/// Target size of a data page written for GPU reads. +/// +/// Pages are the unit a GPU reader decompresses in parallel, so they need to be large enough +/// to amortize per-page setup and numerous enough to fill the device. ~1 MiB is the page size +/// cuDF's Parquet reader is tuned around. +pub const GPU_DATA_PAGE_SIZE: usize = 1024 * 1024; + +/// Row cap per data page. +/// +/// `parquet`'s default caps pages at 20k rows, which produces pages far below +/// [`GPU_DATA_PAGE_SIZE`] for narrow columns and leaves the device underfed. +const GPU_DATA_PAGE_ROW_COUNT_LIMIT: usize = 1_000_000; + +/// Parquet page codecs the GPU benchmark can write. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)] +pub enum GpuCodec { + /// The Parquet default, and the codec with the highest device-side throughput. + #[default] + Snappy, + /// Matches the codec used by the CPU Parquet benchmark, at lower device throughput. + Zstd, +} + +impl GpuCodec { + /// The Parquet compression setting for this codec. + pub fn to_parquet(self) -> Compression { + match self { + GpuCodec::Snappy => Compression::SNAPPY, + GpuCodec::Zstd => Compression::ZSTD(ZstdLevel::default()), + } + } + + /// Short lowercase name, used in measurement labels. + pub fn name(self) -> &'static str { + match self { + GpuCodec::Snappy => "snappy", + GpuCodec::Zstd => "zstd", + } + } +} + +/// Writer properties tuned for a GPU read. +pub fn gpu_writer_properties(codec: GpuCodec) -> WriterProperties { + WriterProperties::builder() + // V1 data pages compress the whole page body. V2 pages place uncompressed + // repetition/definition levels ahead of the compressed values in the same body, which + // not every GPU reader path handles. + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_compression(codec.to_parquet()) + // Dictionary encoding keeps the decompressed payload small and is the encoding GPU + // Parquet readers decode fastest. + .set_dictionary_enabled(true) + .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) + .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) + // Per-page statistics only inflate the page headers a reader has to walk. + .set_statistics_enabled(EnabledStatistics::Chunk) + .build() +} + +#[cfg(test)] +mod tests { + use parquet::file::properties::WriterVersion; + + use super::*; + + #[test] + fn gpu_properties_use_v1_pages_and_the_requested_codec() { + let properties = gpu_writer_properties(GpuCodec::Snappy); + assert_eq!(properties.writer_version(), WriterVersion::PARQUET_1_0); + assert_eq!(properties.compression(&"x".into()), Compression::SNAPPY); + assert_eq!(properties.data_page_size_limit(), GPU_DATA_PAGE_SIZE); + } +} diff --git a/benchmarks/compress-bench/src/lib.rs b/benchmarks/compress-bench/src/lib.rs index e82f8574dc6..76e34172f6a 100644 --- a/benchmarks/compress-bench/src/lib.rs +++ b/benchmarks/compress-bench/src/lib.rs @@ -7,6 +7,6 @@ pub use lance_bench::compress::LanceCompressor; pub mod gpu_parquet; #[cfg(feature = "cuda")] pub mod gpu_vortex; +pub mod gpu_writer; pub mod parquet; -pub mod parquet_pages; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 6fb003f97e3..34f67b13e28 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -14,8 +14,8 @@ use compress_bench::LanceCompressor; use compress_bench::gpu_parquet::GpuParquetCompressor; #[cfg(feature = "cuda")] use compress_bench::gpu_vortex::GpuVortexCompressor; +use compress_bench::gpu_writer::GpuCodec; use compress_bench::parquet::ParquetCompressor; -use compress_bench::parquet_pages::GpuCodec; use compress_bench::vortex::VortexCompressor; use futures::FutureExt; use indicatif::ProgressBar; @@ -81,9 +81,9 @@ struct Args { /// decompression only, for both Vortex and Parquet. #[arg(long)] gpu_decompress: bool, - /// Page codec the GPU Parquet backend writes and decompresses with. + /// Page codec the GPU Parquet file is written with. /// - /// Snappy is the Parquet default and the codec nvCOMP decompresses fastest. + /// Snappy is the Parquet default and the codec GPU readers decompress fastest. #[arg(long, value_enum, default_value_t)] gpu_parquet_codec: GpuCodec, /// Cross-check every GPU-decompressed result against the CPU decoder. @@ -147,13 +147,9 @@ async fn main() -> anyhow::Result<()> { /// Settings for the GPU decompression mode. #[derive(Clone, Copy, Debug)] struct GpuOptions { - /// Parquet page codec to write and decompress on the device. + /// Parquet page codec to write the GPU file with. codec: GpuCodec, /// Cross-check decompressed output against the CPU decoders. - #[cfg_attr( - not(feature = "cuda"), - expect(dead_code, reason = "only the CUDA backends read this") - )] verify: bool, } diff --git a/benchmarks/compress-bench/src/parquet_pages.rs b/benchmarks/compress-bench/src/parquet_pages.rs deleted file mode 100644 index d55598cd8dd..00000000000 --- a/benchmarks/compress-bench/src/parquet_pages.rs +++ /dev/null @@ -1,530 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Locating the compressed page bodies inside a Parquet file. -//! -//! Parquet compresses each page body independently with a block codec, which is exactly the -//! shape nvCOMP's batched decompression entrypoints consume: an array of independent -//! compressed chunks with known uncompressed sizes. This module finds those chunks so the -//! GPU backend can hand the whole batch to the device in one launch, the same decomposition -//! cuDF's Parquet reader uses. -//! -//! Column chunk ranges come from the file footer; page boundaries within a chunk are only -//! discoverable by walking the per-page Thrift headers, so a minimal Thrift compact-protocol -//! reader lives here. `parquet::format::PageHeader` is deprecated and scheduled for removal, -//! and `parquet`'s own page-header parser is crate-private, so neither can be used. - -use anyhow::Result; -use anyhow::bail; -use anyhow::ensure; -use clap::ValueEnum; -use parquet::basic::Compression; -use parquet::basic::ZstdLevel; -use parquet::file::metadata::ParquetMetaData; -use parquet::file::properties::EnabledStatistics; -use parquet::file::properties::WriterProperties; -use parquet::file::properties::WriterVersion; - -/// Target size of a data page written for GPU decompression. -/// -/// nvCOMP decompresses one chunk per page, so pages must be large enough to amortize the -/// per-chunk setup yet numerous enough to fill the device. ~1 MiB is the page size cuDF's -/// Parquet reader is tuned around. -pub const GPU_DATA_PAGE_SIZE: usize = 1024 * 1024; - -/// Row cap per data page. -/// -/// `parquet`'s default caps pages at 20k rows, which produces pages far below -/// [`GPU_DATA_PAGE_SIZE`] for narrow columns and leaves the device underfed. -const GPU_DATA_PAGE_ROW_COUNT_LIMIT: usize = 1_000_000; - -/// Parquet page codecs that nvCOMP can decompress on the device. -#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)] -pub enum GpuCodec { - /// The Parquet default, and the codec with the highest device-side throughput. - #[default] - Snappy, - /// Matches the codec used by the CPU Parquet benchmark, at lower device throughput. - Zstd, -} - -impl GpuCodec { - /// The Parquet compression setting for this codec. - pub fn to_parquet(self) -> Compression { - match self { - GpuCodec::Snappy => Compression::SNAPPY, - GpuCodec::Zstd => Compression::ZSTD(ZstdLevel::default()), - } - } - - /// Short lowercase name, used in measurement labels. - pub fn name(self) -> &'static str { - match self { - GpuCodec::Snappy => "snappy", - GpuCodec::Zstd => "zstd", - } - } - - /// Whether a column chunk's codec matches this one. - pub fn matches(self, compression: Compression) -> bool { - matches!( - (self, compression), - (GpuCodec::Snappy, Compression::SNAPPY) | (GpuCodec::Zstd, Compression::ZSTD(_)) - ) - } - - /// Decompress a single page body on the host, for cross-checking device output. - pub fn decompress_host(self, compressed: &[u8], uncompressed_len: usize) -> Result> { - let decompressed = match self { - GpuCodec::Snappy => snap::raw::Decoder::new().decompress_vec(compressed)?, - GpuCodec::Zstd => zstd::bulk::decompress(compressed, uncompressed_len)?, - }; - ensure!( - decompressed.len() == uncompressed_len, - "page decompressed to {} bytes, page header declared {uncompressed_len}", - decompressed.len() - ); - Ok(decompressed) - } -} - -/// Writer properties tuned for GPU decompression. -pub fn gpu_writer_properties(codec: GpuCodec) -> WriterProperties { - WriterProperties::builder() - // V1 data pages compress the entire page body, which is the unit nvCOMP decompresses. - // V2 pages place uncompressed repetition/definition levels ahead of the compressed - // values inside one page body, which the batched entrypoints cannot express. - .set_writer_version(WriterVersion::PARQUET_1_0) - .set_compression(codec.to_parquet()) - // Dictionary encoding keeps the decompressed payload small and is the encoding GPU - // Parquet readers decode fastest. - .set_dictionary_enabled(true) - .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) - .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) - // Per-page statistics only inflate the page headers that have to be walked on the host. - .set_statistics_enabled(EnabledStatistics::Chunk) - .build() -} - -/// A compressed page body located within a Parquet file. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct CompressedPage { - /// Offset of the compressed body, i.e. just past the page header. - pub offset: usize, - /// Length of the compressed body in bytes. - pub compressed_len: usize, - /// Length of the body once decompressed. - pub uncompressed_len: usize, -} - -/// The pages of one column chunk, alongside the byte range the chunk occupies. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ColumnChunkPages { - /// Offset of the column chunk within the file. - pub offset: u64, - /// Length of the column chunk in bytes. - pub len: usize, - /// Compressed pages of this chunk, in file order, with file-absolute offsets. - pub pages: Vec, -} - -/// Walks every column chunk's page headers and returns the compressed page bodies. -/// -/// The outer `Vec` is one entry per row group, which is the unit a reader can stage on the -/// device and release before moving on. Chunks and the pages within them are in file order. -pub fn scan_compressed_pages( - file_bytes: &[u8], - metadata: &ParquetMetaData, -) -> Result>> { - let mut row_groups = Vec::with_capacity(metadata.row_groups().len()); - - for row_group in metadata.row_groups() { - let mut chunks = Vec::with_capacity(row_group.columns().len()); - for column in row_group.columns() { - let (chunk_offset, chunk_len) = column.byte_range(); - let mut pages = Vec::new(); - let (start, len) = (chunk_offset, chunk_len); - let start = usize::try_from(start)?; - let end = start - .checked_add(usize::try_from(len)?) - .filter(|end| *end <= file_bytes.len()) - .ok_or_else(|| { - anyhow::anyhow!( - "column chunk range {start}..+{len} extends past the {} byte file", - file_bytes.len() - ) - })?; - - let mut pos = start; - while pos < end { - let header = read_page_header(&file_bytes[pos..end])?; - let body = pos + header.header_len; - let body_end = body - .checked_add(header.compressed_len) - .filter(|body_end| *body_end <= end) - .ok_or_else(|| { - anyhow::anyhow!( - "page body at {body} of {} bytes overruns its column chunk", - header.compressed_len - ) - })?; - - match header.page_type { - PageType::Data | PageType::Dictionary => pages.push(CompressedPage { - offset: body, - compressed_len: header.compressed_len, - uncompressed_len: header.uncompressed_len, - }), - PageType::DataV2 => bail!( - "v2 data pages are not GPU-decompressible as a single chunk; \ - write the file with WriterVersion::PARQUET_1_0" - ), - // Index pages are not part of the column data and are never written by - // `parquet`; skip over the body rather than decompressing it. - PageType::Index => {} - } - - pos = body_end; - } - - chunks.push(ColumnChunkPages { - offset: chunk_offset, - len: usize::try_from(chunk_len)?, - pages, - }); - } - - row_groups.push(chunks); - } - - Ok(row_groups) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum PageType { - Data, - Index, - Dictionary, - DataV2, -} - -struct PageHeaderInfo { - header_len: usize, - page_type: PageType, - compressed_len: usize, - uncompressed_len: usize, -} - -/// Thrift compact-protocol field types. -mod ttype { - pub(super) const STOP: u8 = 0x00; - pub(super) const BOOL_TRUE: u8 = 0x01; - pub(super) const BOOL_FALSE: u8 = 0x02; - pub(super) const I8: u8 = 0x03; - pub(super) const I16: u8 = 0x04; - pub(super) const I32: u8 = 0x05; - pub(super) const I64: u8 = 0x06; - pub(super) const DOUBLE: u8 = 0x07; - pub(super) const BINARY: u8 = 0x08; - pub(super) const LIST: u8 = 0x09; - pub(super) const SET: u8 = 0x0a; - pub(super) const MAP: u8 = 0x0b; - pub(super) const STRUCT: u8 = 0x0c; - pub(super) const UUID: u8 = 0x0d; -} - -/// Guards against unbounded recursion on malformed headers. -const MAX_STRUCT_DEPTH: u32 = 32; - -/// Reads the `PageHeader` at the start of `buf`, returning its fields and encoded length. -fn read_page_header(buf: &[u8]) -> Result { - let mut reader = CompactReader { buf, pos: 0 }; - let mut page_type = None; - let mut uncompressed_len = None; - let mut compressed_len = None; - let mut last_field_id = 0i16; - - while let Some((field_id, field_type)) = reader.read_field_header(&mut last_field_id)? { - match (field_id, field_type) { - (1, ttype::I32) => page_type = Some(reader.read_i32()?), - (2, ttype::I32) => uncompressed_len = Some(reader.read_i32()?), - (3, ttype::I32) => compressed_len = Some(reader.read_i32()?), - _ => reader.skip_value(field_type, 0)?, - } - } - - let page_type = match page_type { - Some(0) => PageType::Data, - Some(1) => PageType::Index, - Some(2) => PageType::Dictionary, - Some(3) => PageType::DataV2, - Some(other) => bail!("unknown Parquet page type {other}"), - None => bail!("page header is missing its page type"), - }; - - let uncompressed_len = - uncompressed_len.ok_or_else(|| anyhow::anyhow!("page header is missing its size"))?; - let compressed_len = compressed_len - .ok_or_else(|| anyhow::anyhow!("page header is missing its compressed size"))?; - - Ok(PageHeaderInfo { - header_len: reader.pos, - page_type, - compressed_len: usize::try_from(compressed_len)?, - uncompressed_len: usize::try_from(uncompressed_len)?, - }) -} - -/// Minimal reader for the subset of the Thrift compact protocol that page headers use. -struct CompactReader<'a> { - buf: &'a [u8], - pos: usize, -} - -impl CompactReader<'_> { - fn read_u8(&mut self) -> Result { - let byte = *self - .buf - .get(self.pos) - .ok_or_else(|| anyhow::anyhow!("page header ends mid-field"))?; - self.pos += 1; - Ok(byte) - } - - fn advance(&mut self, len: usize) -> Result<()> { - let end = self - .pos - .checked_add(len) - .filter(|end| *end <= self.buf.len()) - .ok_or_else(|| anyhow::anyhow!("page header ends mid-value"))?; - self.pos = end; - Ok(()) - } - - fn read_varint(&mut self) -> Result { - let mut value = 0u64; - for shift in (0..64).step_by(7) { - let byte = self.read_u8()?; - value |= u64::from(byte & 0x7f) << shift; - if byte & 0x80 == 0 { - return Ok(value); - } - } - bail!("varint in page header is not terminated") - } - - fn read_zigzag(&mut self) -> Result { - let encoded = self.read_varint()?; - Ok(((encoded >> 1) as i64) ^ -((encoded & 1) as i64)) - } - - fn read_i32(&mut self) -> Result { - Ok(i32::try_from(self.read_zigzag()?)?) - } - - /// Reads the next field header, or `None` at the struct's STOP byte. - fn read_field_header(&mut self, last_field_id: &mut i16) -> Result> { - let header = self.read_u8()?; - if header == ttype::STOP { - return Ok(None); - } - - let field_type = header & 0x0f; - let delta = header >> 4; - let field_id = if delta == 0 { - i16::try_from(self.read_zigzag()?)? - } else { - last_field_id - .checked_add(i16::from(delta)) - .ok_or_else(|| anyhow::anyhow!("field id overflow in page header"))? - }; - *last_field_id = field_id; - - Ok(Some((field_id, field_type))) - } - - fn skip_struct(&mut self, depth: u32) -> Result<()> { - ensure!( - depth < MAX_STRUCT_DEPTH, - "page header nests structs more than {MAX_STRUCT_DEPTH} deep" - ); - let mut last_field_id = 0i16; - while let Some((_, field_type)) = self.read_field_header(&mut last_field_id)? { - self.skip_value(field_type, depth + 1)?; - } - Ok(()) - } - - /// Skips a field value. Booleans carry their value in the field type, so consume nothing. - fn skip_value(&mut self, field_type: u8, depth: u32) -> Result<()> { - match field_type { - ttype::BOOL_TRUE | ttype::BOOL_FALSE => Ok(()), - ttype::I8 => self.advance(1), - ttype::I16 | ttype::I32 | ttype::I64 => self.read_varint().map(|_| ()), - ttype::DOUBLE => self.advance(8), - ttype::UUID => self.advance(16), - ttype::BINARY => { - let len = usize::try_from(self.read_varint()?)?; - self.advance(len) - } - ttype::LIST | ttype::SET => { - let (len, element_type) = self.read_collection_header()?; - for _ in 0..len { - self.skip_element(element_type, depth + 1)?; - } - Ok(()) - } - ttype::MAP => { - let len = usize::try_from(self.read_varint()?)?; - if len > 0 { - let types = self.read_u8()?; - let (key_type, value_type) = (types >> 4, types & 0x0f); - for _ in 0..len { - self.skip_element(key_type, depth + 1)?; - self.skip_element(value_type, depth + 1)?; - } - } - Ok(()) - } - ttype::STRUCT => self.skip_struct(depth), - other => bail!("unsupported Thrift compact type {other} in page header"), - } - } - - /// Skips one collection element. Unlike fields, booleans here occupy a byte of their own. - fn skip_element(&mut self, element_type: u8, depth: u32) -> Result<()> { - match element_type { - ttype::BOOL_TRUE | ttype::BOOL_FALSE => self.advance(1), - other => self.skip_value(other, depth), - } - } - - fn read_collection_header(&mut self) -> Result<(usize, u8)> { - let header = self.read_u8()?; - let element_type = header & 0x0f; - let len = match header >> 4 { - 0x0f => usize::try_from(self.read_varint()?)?, - short_len => usize::from(short_len), - }; - Ok((len, element_type)) - } -} - -#[cfg(test)] -mod tests { - use std::fs::File; - use std::sync::Arc; - - use arrow_array::Int64Array; - use arrow_array::RecordBatch; - use arrow_array::StringArray; - use arrow_schema::DataType; - use arrow_schema::Field; - use arrow_schema::Schema; - use parquet::arrow::ArrowWriter; - use parquet::file::metadata::ParquetMetaDataReader; - use parquet::file::reader::FileReader; - use parquet::file::reader::SerializedFileReader; - use rstest::rstest; - - use super::*; - - fn sample_batch() -> Result { - let schema = Arc::new(Schema::new(vec![ - Field::new("ints", DataType::Int64, false), - Field::new("strings", DataType::Utf8, false), - ])); - let ints = Int64Array::from_iter_values((0..50_000).map(|i| i % 977)); - let strings = - StringArray::from_iter_values((0..50_000).map(|i| format!("value-{}", i % 1_000))); - Ok(RecordBatch::try_new( - schema, - vec![Arc::new(ints), Arc::new(strings)], - )?) - } - - fn write_sample(path: &std::path::Path, codec: GpuCodec) -> Result<()> { - let batch = sample_batch()?; - let file = File::create(path)?; - let mut writer = - ArrowWriter::try_new(file, batch.schema(), Some(gpu_writer_properties(codec)))?; - writer.write(&batch)?; - writer.close()?; - Ok(()) - } - - /// The page bodies we locate must decompress to exactly the bytes `parquet` itself reads. - #[rstest] - #[case(GpuCodec::Snappy)] - #[case(GpuCodec::Zstd)] - fn scanned_pages_match_parquet_reader(#[case] codec: GpuCodec) -> Result<()> { - let dir = tempfile::tempdir()?; - let path = dir.path().join("sample.parquet"); - write_sample(&path, codec)?; - - let file = File::open(&path)?; - let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; - let file_bytes = std::fs::read(&path)?; - let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; - let pages = row_groups - .iter() - .flatten() - .flat_map(|chunk| chunk.pages.iter()) - .collect::>(); - - let reader = SerializedFileReader::new(File::open(&path)?)?; - let mut expected = Vec::new(); - for row_group in 0..reader.metadata().num_row_groups() { - let row_group_reader = reader.get_row_group(row_group)?; - for column in 0..row_group_reader.num_columns() { - let mut page_reader = row_group_reader.get_column_page_reader(column)?; - while let Some(page) = page_reader.get_next_page()? { - expected.push(page.buffer().to_vec()); - } - } - } - - assert_eq!(pages.len(), expected.len(), "page count mismatch"); - assert!(!pages.is_empty(), "expected the sample file to have pages"); - - for (page, expected) in pages.iter().zip(expected.iter()) { - let compressed = &file_bytes[page.offset..page.offset + page.compressed_len]; - let decompressed = codec.decompress_host(compressed, page.uncompressed_len)?; - assert_eq!(&decompressed, expected); - } - - Ok(()) - } - - /// Pages must tile their column chunks exactly, with no gap left unaccounted for. - #[test] - fn scanned_pages_cover_every_column_chunk() -> Result<()> { - let dir = tempfile::tempdir()?; - let path = dir.path().join("sample.parquet"); - write_sample(&path, GpuCodec::Snappy)?; - - let file = File::open(&path)?; - let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; - let file_bytes = std::fs::read(&path)?; - let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; - let pages = row_groups - .iter() - .flatten() - .flat_map(|chunk| chunk.pages.iter()) - .collect::>(); - - let compressed: usize = pages.iter().map(|page| page.compressed_len).sum(); - let chunk_total: i64 = metadata - .row_groups() - .iter() - .flat_map(|rg| rg.columns()) - .map(|col| col.compressed_size()) - .sum(); - - // The chunk total includes the page headers, so the page bodies must be strictly - // smaller but within a header's worth per page. - assert!(compressed < usize::try_from(chunk_total)?); - assert!(compressed > usize::try_from(chunk_total)? - pages.len() * 256); - Ok(()) - } -} diff --git a/scripts/cudf-parquet-read.py b/scripts/cudf-parquet-read.py new file mode 100644 index 00000000000..d597762e114 --- /dev/null +++ b/scripts/cudf-parquet-read.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Times a full GPU Parquet read with cuDF. + +`cudf.read_parquet` performs the entire read on the device: page header decode, +codec decompression, dictionary/RLE/plain decoding and column assembly. That makes it +the like-for-like opponent for the Vortex GPU backend, which also decodes all the way +to canonical arrays on device. + +Timing excludes interpreter start, `import cudf`, CUDA context creation and any JIT +warm-up, all of which are paid once per process and are not part of a read. A warm-up +read runs first for exactly that reason. + +Emits one JSON object on stdout so the benchmark can parse it. +""" + +import argparse +import json +import sys +import time + + +def synchronize() -> None: + """Block until queued device work finishes. + + `cudf.read_parquet` returns a materialized DataFrame, but synchronizing explicitly + keeps the measurement honest if that ever stops being true. + """ + try: + import cupy + + cupy.cuda.runtime.deviceSynchronize() + except ImportError: + pass + + +def verify(path: str, frame) -> None: + """Fails unless the GPU read matches a CPU Parquet read of the same file.""" + import pandas as pd + from pandas.testing import assert_frame_equal + + expected = pd.read_parquet(path) + actual = frame.to_pandas() + + # cuDF and pyarrow can land on different-but-equivalent dtypes (nullable vs numpy + # backed, for instance), so compare values and leave dtype policy out of it. + assert_frame_equal(actual, expected, check_dtype=False) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", help="Parquet file to read") + parser.add_argument( + "--iterations", type=int, default=1, help="timed reads to perform" + ) + parser.add_argument( + "--verify", + action="store_true", + help="cross-check the GPU read against a CPU Parquet read", + ) + args = parser.parse_args() + + import cudf + + # Warm-up: pays CUDA context creation and any first-call JIT so they stay out of + # the timed reads below. + warmup = cudf.read_parquet(args.path) + synchronize() + + if args.verify: + verify(args.path, warmup) + + rows, columns = warmup.shape + del warmup + + runs_ns = [] + for _ in range(max(args.iterations, 1)): + start = time.perf_counter_ns() + frame = cudf.read_parquet(args.path) + synchronize() + runs_ns.append(time.perf_counter_ns() - start) + del frame + + json.dump( + { + "min_ns": min(runs_ns), + "runs_ns": runs_ns, + "rows": int(rows), + "columns": int(columns), + "verified": bool(args.verify), + }, + sys.stdout, + ) + sys.stdout.write("\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vortex-cuda/nvcomp/build.rs b/vortex-cuda/nvcomp/build.rs index f34507378db..877fc68c58d 100644 --- a/vortex-cuda/nvcomp/build.rs +++ b/vortex-cuda/nvcomp/build.rs @@ -90,20 +90,13 @@ fn main() { let bindings = bindgen::Builder::default() .header(include_dir.join("nvcomp.h").to_string_lossy()) .header(include_dir.join("nvcomp/zstd.h").to_string_lossy()) - .header(include_dir.join("nvcomp/snappy.h").to_string_lossy()) .clang_arg(format!("-I{}", include_dir.display())) .clang_arg(format!("-I{}", cuda_stub_dir.display())) .allowlist_type("nvcompStatus_t") - .allowlist_type("nvcompAlignmentRequirements_t") .allowlist_type("nvcompBatchedZstdDecompressOpts_t") - .allowlist_type("nvcompBatchedSnappyDecompressOpts_t") .allowlist_type("nvcompDecompressBackend_t") .allowlist_function("nvcompBatchedZstdDecompressGetTempSizeAsync") - .allowlist_function("nvcompBatchedZstdDecompressGetRequiredAlignments") .allowlist_function("nvcompBatchedZstdDecompressAsync") - .allowlist_function("nvcompBatchedSnappyDecompressGetTempSizeAsync") - .allowlist_function("nvcompBatchedSnappyDecompressGetRequiredAlignments") - .allowlist_function("nvcompBatchedSnappyDecompressAsync") .dynamic_library_name("NvcompLibrary") .dynamic_link_require_all(true) .wrap_unsafe_ops(true) diff --git a/vortex-cuda/nvcomp/src/backend.rs b/vortex-cuda/nvcomp/src/backend.rs deleted file mode 100644 index 5be1ebbde4c..00000000000 --- a/vortex-cuda/nvcomp/src/backend.rs +++ /dev/null @@ -1,52 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Options and metadata shared by nvcomp's batched decompression APIs. - -use crate::sys; - -/// Backend selection for nvcomp decompression. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum DecompressBackend { - /// Let nvcomp auto-select the best backend for the hardware. - #[default] - Default, - /// Use hardware decompression - Hardware, - /// Use CUDA - Cuda, -} - -impl DecompressBackend { - pub(crate) fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { - match self { - Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, - Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, - Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, - } - } -} - -/// Minimum buffer alignments required by an nvcomp algorithm. -/// -/// Buffers passed to the batched decompression entrypoints must satisfy these alignments. -/// Exceeding them (for example 16- or 32-byte alignment) may improve throughput. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct AlignmentRequirements { - /// Minimum alignment of each compressed input chunk. - pub input: usize, - /// Minimum alignment of each decompressed output chunk. - pub output: usize, - /// Minimum alignment of the temporary workspace buffer. - pub temp: usize, -} - -impl From for AlignmentRequirements { - fn from(value: sys::nvcompAlignmentRequirements_t) -> Self { - Self { - input: value.input, - output: value.output, - temp: value.temp, - } - } -} diff --git a/vortex-cuda/nvcomp/src/lib.rs b/vortex-cuda/nvcomp/src/lib.rs index 859eecbea22..3ab502baa63 100644 --- a/vortex-cuda/nvcomp/src/lib.rs +++ b/vortex-cuda/nvcomp/src/lib.rs @@ -30,13 +30,9 @@ use std::sync::OnceLock; )] pub mod sys; -mod backend; mod error; -pub mod snappy; pub mod zstd; -pub use backend::AlignmentRequirements; -pub use backend::DecompressBackend; pub use error::NvcompError; /// The loaded nvcomp library instance. diff --git a/vortex-cuda/nvcomp/src/snappy.rs b/vortex-cuda/nvcomp/src/snappy.rs deleted file mode 100644 index 175c19d3a13..00000000000 --- a/vortex-cuda/nvcomp/src/snappy.rs +++ /dev/null @@ -1,222 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Wrappers around nvcomp's batched Snappy decompression API. -//! -//! Snappy is the default Parquet page codec and is the fastest of the codecs nvcomp -//! implements on device, which makes it the codec of choice when feeding Parquet pages -//! to the GPU. - -use std::ffi::c_void; - -use crate::backend::AlignmentRequirements; -pub use crate::backend::DecompressBackend; -use crate::error::NvcompError; -use crate::error::check_status; -use crate::nvcomp_library; -use crate::sys; - -/// The largest compressed chunk the Snappy decompressor accepts, in bytes. -pub const MAX_COMPRESSED_CHUNK_SIZE: usize = (1 << 31) - 1; - -/// Options for batched Snappy decompression. -#[derive(Debug, Clone, Copy, Default)] -pub struct SnappyDecompressOpts { - /// Which nvcomp backend performs the decompression. - pub backend: DecompressBackend, - /// Sort chunks by size before submitting them to the hardware decompression engine. - /// - /// Only used when `backend` selects the hardware engine. - pub sort_before_hw_decompress: bool, -} - -impl SnappyDecompressOpts { - fn to_nvcomp(self) -> sys::nvcompBatchedSnappyDecompressOpts_t { - sys::nvcompBatchedSnappyDecompressOpts_t { - backend: self.backend.to_nvcomp(), - sort_before_hw_decompress: i32::from(self.sort_before_hw_decompress), - reserved: [0; 56], - } - } -} - -/// Computes required temporary buffer size for batched Snappy decompression. -/// -/// # Arguments -/// -/// * `num_chunks` - Number of compressed chunks to decompress -/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk -/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks -/// -/// # Returns -/// -/// The required size in bytes for the temporary buffer. -pub fn get_decompress_temp_size( - num_chunks: usize, - max_uncompressed_chunk_bytes: usize, - max_total_uncompressed_bytes: usize, -) -> Result { - get_decompress_temp_size_with_opts( - num_chunks, - max_uncompressed_chunk_bytes, - max_total_uncompressed_bytes, - SnappyDecompressOpts::default(), - ) -} - -/// Computes required temporary buffer size with custom options. -/// -/// # Arguments -/// -/// * `num_chunks` - Number of compressed chunks to decompress -/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk -/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks -/// * `opts` - Decompression options -/// -/// # Returns -/// -/// The required size in bytes for the temporary buffer. -pub fn get_decompress_temp_size_with_opts( - num_chunks: usize, - max_uncompressed_chunk_bytes: usize, - max_total_uncompressed_bytes: usize, - opts: SnappyDecompressOpts, -) -> Result { - let library = nvcomp_library()?; - - let mut temp_bytes: usize = 0; - - let status = unsafe { - library.nvcompBatchedSnappyDecompressGetTempSizeAsync( - num_chunks, - max_uncompressed_chunk_bytes, - opts.to_nvcomp(), - &raw mut temp_bytes, - max_total_uncompressed_bytes, - ) - }; - - check_status(status)?; - Ok(temp_bytes) -} - -/// Returns the minimum buffer alignments required by batched Snappy decompression. -pub fn decompress_alignment_requirements( - opts: SnappyDecompressOpts, -) -> Result { - let library = nvcomp_library()?; - - let mut requirements = sys::nvcompAlignmentRequirements_t { - input: 0, - output: 0, - temp: 0, - }; - - let status = unsafe { - library.nvcompBatchedSnappyDecompressGetRequiredAlignments( - opts.to_nvcomp(), - &raw mut requirements, - ) - }; - - check_status(status)?; - Ok(requirements.into()) -} - -/// Launches batched Snappy decompression asynchronously on the GPU. -/// -/// This function decompresses multiple raw Snappy blocks in parallel on the GPU. All -/// pointer arguments must point to device memory, and the operation is executed -/// asynchronously on the provided CUDA stream. -/// -/// # Arguments -/// -/// * `device_compressed_ptrs` - Device pointer to array of pointers to compressed chunks -/// * `device_compressed_bytes` - Device pointer to array of compressed chunk sizes -/// * `device_uncompressed_bytes` - Device pointer to array of expected uncompressed sizes -/// * `device_actual_uncompressed_bytes` - Device pointer to array for actual uncompressed sizes (output) -/// * `num_chunks` - Number of chunks to decompress -/// * `device_temp_ptr` - Device pointer to temporary workspace buffer -/// * `temp_bytes` - Size of temporary buffer in bytes -/// * `device_uncompressed_ptrs` - Device pointer to array of pointers to output buffers -/// * `device_statuses` - Device pointer to array for per-chunk status codes (output) -/// * `stream` - CUDA stream to execute on -/// -/// # Safety -/// -/// - All device pointers must be valid and point to properly allocated device memory -/// - `device_compressed_ptrs` must point to valid device pointers -/// - `device_uncompressed_ptrs` must point to valid device pointers -/// - Each output buffer must have at least the corresponding `device_uncompressed_bytes` size -/// - `device_temp_ptr` must have at least `temp_bytes` allocated -/// - The stream must be valid -#[expect(clippy::too_many_arguments)] -pub unsafe fn decompress_async( - device_compressed_ptrs: *const *const c_void, - device_compressed_bytes: *const usize, - device_uncompressed_bytes: *const usize, - device_actual_uncompressed_bytes: *mut usize, - num_chunks: usize, - device_temp_ptr: *mut c_void, - temp_bytes: usize, - device_uncompressed_ptrs: *const *mut c_void, - device_statuses: *mut sys::nvcompStatus_t, - stream: sys::cudaStream_t, -) -> Result<(), NvcompError> { - // SAFETY: Caller has to ensure all pointers are valid. - unsafe { - decompress_async_with_opts( - device_compressed_ptrs, - device_compressed_bytes, - device_uncompressed_bytes, - device_actual_uncompressed_bytes, - num_chunks, - device_temp_ptr, - temp_bytes, - device_uncompressed_ptrs, - device_statuses, - stream, - SnappyDecompressOpts::default(), - ) - } -} - -/// Launches batched Snappy decompression asynchronously with custom options. -/// -/// # Safety -/// -/// Same requirements as [`decompress_async`]. -#[expect(clippy::too_many_arguments)] -pub unsafe fn decompress_async_with_opts( - device_compressed_ptrs: *const *const c_void, - device_compressed_bytes: *const usize, - device_uncompressed_bytes: *const usize, - device_actual_uncompressed_bytes: *mut usize, - num_chunks: usize, - device_temp_ptr: *mut c_void, - temp_bytes: usize, - device_uncompressed_ptrs: *const *mut c_void, - device_statuses: *mut sys::nvcompStatus_t, - stream: sys::cudaStream_t, - opts: SnappyDecompressOpts, -) -> Result<(), NvcompError> { - let library = nvcomp_library()?; - - let status = unsafe { - library.nvcompBatchedSnappyDecompressAsync( - device_compressed_ptrs, - device_compressed_bytes, - device_uncompressed_bytes, - device_actual_uncompressed_bytes, - num_chunks, - device_temp_ptr, - temp_bytes, - device_uncompressed_ptrs, - opts.to_nvcomp(), - device_statuses, - stream, - ) - }; - - check_status(status) -} diff --git a/vortex-cuda/nvcomp/src/zstd.rs b/vortex-cuda/nvcomp/src/zstd.rs index 04769d2621a..44111901071 100644 --- a/vortex-cuda/nvcomp/src/zstd.rs +++ b/vortex-cuda/nvcomp/src/zstd.rs @@ -5,13 +5,33 @@ use std::ffi::c_void; -use crate::backend::AlignmentRequirements; -pub use crate::backend::DecompressBackend; use crate::error::NvcompError; use crate::error::check_status; use crate::nvcomp_library; use crate::sys; +/// Backend selection for nvcomp decompression. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum DecompressBackend { + /// Let nvcomp auto-select the best backend for the hardware. + #[default] + Default, + /// Use hardware decompression + Hardware, + /// Use CUDA + Cuda, +} + +impl DecompressBackend { + fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { + match self { + Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, + Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, + Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, + } + } +} + /// Options for batched ZSTD decompression. #[derive(Debug, Clone, Copy, Default)] pub struct ZstdDecompressOpts { @@ -87,29 +107,6 @@ pub fn get_decompress_temp_size_with_opts( Ok(temp_bytes) } -/// Returns the minimum buffer alignments required by batched ZSTD decompression. -pub fn decompress_alignment_requirements( - opts: ZstdDecompressOpts, -) -> Result { - let library = nvcomp_library()?; - - let mut requirements = sys::nvcompAlignmentRequirements_t { - input: 0, - output: 0, - temp: 0, - }; - - let status = unsafe { - library.nvcompBatchedZstdDecompressGetRequiredAlignments( - opts.to_nvcomp(), - &raw mut requirements, - ) - }; - - check_status(status)?; - Ok(requirements.into()) -} - /// Launches batched ZSTD decompression asynchronously on the GPU. /// /// This function decompresses multiple ZSTD-compressed chunks in parallel on the GPU. From c5a74d17e3abb501c2f427eef447f0f39407b82c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 15:07:10 +0000 Subject: [PATCH 09/16] Run the GPU verification reference through a host execution context The reference side of the Vortex verification was executing through the CUDA context: the host scan's batches and both Arrow conversions were handed `cuda_ctx.execution_ctx()`. A CUDA context allocates its outputs in device memory, so the Arrow conversion then read a device buffer from the host and panicked with "unwrap_host called for Device allocation" on the string-heavy Public BI datasets, where canonicalisation goes through the buffer directly. Signed-off-by: Claude --- benchmarks/compress-bench/src/gpu_vortex.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 9b25b2f6dd4..2b9e5399e94 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -18,6 +18,7 @@ use tempfile::NamedTempFile; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::compressor::BtrBlocksCompressorBuilder; @@ -124,6 +125,10 @@ async fn open_gpu(path: &Path) -> Result { /// Verification runs inline, so the returned duration is not comparable to a plain run. async fn verify_against_host_scan(path: &Path) -> Result { let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + // Everything on the reference side — the host scan and both Arrow conversions — has to run + // through a plain host context. A CUDA context allocates its outputs in device memory, and + // the Arrow conversion then reads those buffers on the host. + let mut host_ctx = SESSION.create_execution_ctx(); let start = Instant::now(); // The host scan reads a copy rather than the same path. The session's segment cache is @@ -148,7 +153,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { }; let gpu_record = gpu_batch.execute::(cuda_ctx.execution_ctx())?; - let host_record = host_batch.execute::(cuda_ctx.execution_ctx())?; + let host_record = host_batch.execute::(&mut host_ctx)?; ensure!( gpu_record.len() == host_record.len(), "batch {batch_index} length differs between the GPU and CPU scans: {} vs {}", @@ -180,7 +185,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { verify_field( &host_field, decoded, - cuda_ctx.execution_ctx(), + &mut host_ctx, batch_index, field_index, )?; From c3a0ad928a93d27f3aecb07492c6f9bcb496b62b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 16:15:57 +0000 Subject: [PATCH 10/16] Publish GPU timings even when a dataset fails to decode The verification gate ran before the timed pass, so a single unsupported dataset meant the run produced no numbers at all. Two changes: - the per-dataset survey now applies to any GPU run, not just a verifying one, and the timing tables render before the failure summary, so datasets that do decode still publish their numbers; - the workflow runs the benchmark before the gate and fails the job at the end on either a failed verification or a failed benchmark. Signed-off-by: Claude --- .github/workflows/pr-bench-gpu-compress.yml | 19 +++++--- benchmarks/compress-bench/src/main.rs | 50 +++++++++++---------- 2 files changed, 40 insertions(+), 29 deletions(-) diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index d2975cbbb91..a4e477f3fe3 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -105,14 +105,13 @@ jobs: with: file-path: verify-comment.md comment-tag: bench-pr-comment-gpu-verify - - name: Fail if verification failed - if: steps.verify.outcome == 'failure' - shell: bash - run: | - echo "GPU decompression verification failed; see the verification comment." >&2 - exit 1 - name: Run GPU compression benchmark + id: bench shell: bash + # A dataset the GPU cannot decode is reported per dataset and still fails the run, but + # the timing tables are printed first, so the datasets that do decode publish their + # numbers. The gate at the end of the job turns either failure into a job failure. + continue-on-error: true env: RUST_BACKTRACE: full # Do not enable VORTEX_EXPERIMENTAL_PATCHED_ARRAY here: it rewrites interior @@ -152,3 +151,11 @@ jobs: GPU Compression failed. Check the [workflow run](${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}) for details. comment-tag: bench-pr-comment-gpu-compress + - name: Fail if verification or the benchmark failed + if: steps.verify.outcome == 'failure' || steps.bench.outcome == 'failure' + shell: bash + # Last, so both the verification matrix and the timing tables are already published. + run: | + echo "verify=${{ steps.verify.outcome }} bench=${{ steps.bench.outcome }}" >&2 + echo "GPU decompression failed; see the verification and results comments." >&2 + exit 1 diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 34f67b13e28..fefe61c44f3 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -285,9 +285,9 @@ async fn run_compress( let mut measurements = vec![]; let mut v3_records: Vec = Vec::new(); - // A verification pass reports on every dataset rather than stopping at the first failure: - // one run then says exactly which datasets decode correctly on the GPU and which do not. - let survey_all = gpu.is_some_and(|gpu| gpu.verify); + // A GPU pass reports on every dataset rather than stopping at the first failure, so one run + // says which datasets decode correctly on the GPU and still yields numbers for the rest. + let survey_all = gpu.is_some(); let mut failures: Vec<(String, anyhow::Error)> = Vec::new(); for dataset_handle in datasets.into_iter() { @@ -322,30 +322,14 @@ async fn run_compress( progress.finish(); - if !failures.is_empty() { - eprintln!( - "\nGPU verification failed for {} dataset(s):", - failures.len() - ); - for (dataset, error) in &failures { - eprintln!(" - {dataset}: {error:#}"); - } - anyhow::bail!( - "GPU verification failed for: {}", - failures - .iter() - .map(|(dataset, _)| dataset.as_str()) - .collect::>() - .join(", ") - ); - } - if let Some(path) = ingest_output { v3::write_jsonl_to_path(&path, &v3_records)?; } let mut writer = create_output_writer(&display_format, output_path, BENCHMARK_ID)?; + // The tables render before any failure is reported, so a partially failing GPU matrix still + // publishes the numbers for the datasets that did decode. match display_format { DisplayFormat::Table => { render_table(&mut writer, measurements.timings, &targets)?; @@ -357,13 +341,33 @@ async fn run_compress( } else { vec![] }, - ) + )?; } DisplayFormat::GhJson => { print_measurements_json(&mut writer, measurements.timings, DOC_PATH)?; - print_measurements_json(&mut writer, measurements.ratios, DOC_PATH) + print_measurements_json(&mut writer, measurements.ratios, DOC_PATH)?; } } + + if !failures.is_empty() { + eprintln!( + "\nGPU decompression failed for {} dataset(s):", + failures.len() + ); + for (dataset, error) in &failures { + eprintln!(" - {dataset}: {error:#}"); + } + anyhow::bail!( + "GPU decompression failed for: {}", + failures + .iter() + .map(|(dataset, _)| dataset.as_str()) + .collect::>() + .join(", ") + ); + } + + Ok(()) } async fn run_benchmark_for_dataset( From ca22f1883f46335888801a263186d2437b5002fd Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 16:17:53 +0000 Subject: [PATCH 11/16] Document that a partial GPU matrix still publishes its timings Signed-off-by: Claude --- benchmarks/compress-bench/README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index e14c117d5a6..b1fd9bc6c7c 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -86,5 +86,7 @@ cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress --gpu-verify --iterations 1 ``` -A verifying 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. +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. From 2afc3b79cbbcdf6779908ecd6d4df957526799ce Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 16:58:34 +0000 Subject: [PATCH 12/16] CI: publish the full error detail from a failed GPU verification The per-dataset grep captures only the first line of each error, so a Python traceback from the cuDF script or a Rust backtrace never reached the comment. Attach the tail of the raw output in a collapsed block on failure, and match the summary line's current wording. Signed-off-by: Claude --- .github/workflows/pr-bench-gpu-compress.yml | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index a4e477f3fe3..9131518615b 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -94,9 +94,22 @@ jobs: fi echo echo '```text' - grep -E "verified [0-9]+|GPU verification failed|^ - " gpu-verify.txt | tail -40 \ + grep -E "verified [0-9]+|GPU decompression failed|^ - " gpu-verify.txt | tail -40 \ || tail -40 gpu-verify.txt echo '```' + # The per-dataset lines above carry only the first line of each error. Python + # tracebacks, Rust backtraces and first-differing-row dumps span several lines, so + # the tail of the raw output goes in a collapsed block rather than back in the log. + if [ "${{ steps.verify.outcome }}" != "success" ]; then + echo + echo "
Full error detail" + echo + echo '```text' + tail -200 gpu-verify.txt + echo '```' + echo + echo "
" + fi } > verify-comment.md cat verify-comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR with verification result From 6a0fd9dfc1e13836381100fab251d5debba30fe6 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 17:23:56 +0000 Subject: [PATCH 13/16] Fix three defects the first cuDF comparison run exposed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(cuda): copy validity back to the host in `into_host` `CanonicalCudaExt::into_host` migrated a canonical array's values buffer but passed its validity through untouched, so a nullable array came back to the host half-migrated and the first host read of the mask panicked with "unwrap_host called for Device allocation" — via `Validity::execute_mask` -> `BoolArray::into_bit_buffer`. Non-nullable arrays were unaffected, which is why it only showed up on the Public BI tables. The `Bool` arm already carried a TODO for exactly this. Do not compare Parquet DATE columns across representations pyarrow materialises a DATE column as `datetime.date` objects and cuDF as `datetime64[s]`. The values agree, but `check_dtype=False` does not bridge object-vs-datetime64, so the comparison reported every row as different and failed both TPC-H datasets. Coerce both sides to datetime64 first. Read the Vortex GPU file through the page cache by default cuDF takes an untimed warm-up read, so its timed read is served from the page cache, while the Vortex reader used `O_DIRECT` on every iteration and paid real disk reads each time. That compared a read of the disk against a read of RAM. Direct IO is now off by default and available behind `--gpu-direct-io` for measuring storage bandwidth, which is not a decode comparison. Signed-off-by: Claude --- benchmarks/compress-bench/README.md | 11 +++++-- benchmarks/compress-bench/src/gpu_vortex.rs | 30 +++++++++++------- benchmarks/compress-bench/src/main.rs | 12 +++++++- scripts/cudf-parquet-read.py | 23 ++++++++++++-- vortex-cuda/src/canonical.rs | 34 ++++++++++++++++++--- 5 files changed, 88 insertions(+), 22 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index b1fd9bc6c7c..5033bca34a7 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -53,9 +53,14 @@ uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarro script, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up read runs first for the same reason. -Note that the two backends do not share an I/O path: the Vortex reader uses pinned buffers and -direct I/O (`O_DIRECT`) on Linux, while cuDF does its own host read and host-to-device -transfer. +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 diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 2b9e5399e94..832466978a0 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -42,6 +42,7 @@ use vortex_cuda::layout::register_cuda_layout; /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. pub struct GpuVortexCompressor { verify: bool, + direct_io: bool, } impl GpuVortexCompressor { @@ -50,8 +51,8 @@ impl GpuVortexCompressor { /// When `verify` is set, each GPU-decoded field is copied back to the host and compared /// against the same field decoded on the CPU. Verification runs inline, so timings from a /// verifying run are not comparable to a plain one. - pub fn new(verify: bool) -> Self { - Self { verify } + pub fn new(verify: bool, direct_io: bool) -> Self { + Self { verify, direct_io } } } @@ -84,12 +85,12 @@ impl Compressor for GpuVortexCompressor { drop(output); if self.verify { - return verify_against_host_scan(gpu_file.path()).await; + return verify_against_host_scan(gpu_file.path(), self.direct_io).await; } let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; let start = Instant::now(); - let file = open_gpu(gpu_file.path()).await?; + let file = open_gpu(gpu_file.path(), self.direct_io).await?; let mut batches = file.scan()?.into_array_stream()?; while let Some(batch) = batches.next().await { @@ -106,13 +107,20 @@ impl Compressor for GpuVortexCompressor { /// Opens a Vortex file for CUDA execution. /// -/// On Linux direct IO keeps repeated iterations measuring storage bandwidth rather than -/// page-cache hits. -async fn open_gpu(path: &Path) -> Result { +/// `direct_io` is off by default so this backend is comparable with the cuDF one: cuDF reads +/// through the page cache after an untimed warm-up read, so leaving direct IO on would compare +/// a Vortex read of the disk against a cuDF read of RAM. Turning it on measures storage +/// bandwidth instead, which is a different question and not comparable across the two. +async fn open_gpu(path: &Path, direct_io: bool) -> Result { let open_options = SESSION.open_options().with_cuda(); #[cfg(target_os = "linux")] - let open_options = - open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()); + let open_options = if direct_io { + open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()) + } else { + open_options + }; + #[cfg(not(target_os = "linux"))] + let _ = direct_io; Ok(open_options.open_path(path).await?) } @@ -123,7 +131,7 @@ async fn open_gpu(path: &Path) -> Result { /// which the host decoders cannot read. /// /// Verification runs inline, so the returned duration is not comparable to a plain run. -async fn verify_against_host_scan(path: &Path) -> Result { +async fn verify_against_host_scan(path: &Path, direct_io: bool) -> Result { let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; // Everything on the reference side — the host scan and both Arrow conversions — has to run // through a plain host context. A CUDA context allocates its outputs in device memory, and @@ -137,7 +145,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { let host_path = NamedTempFile::new()?; std::fs::copy(path, host_path.path())?; - let gpu_file = open_gpu(path).await?; + let gpu_file = open_gpu(path, direct_io).await?; let mut gpu_batches = gpu_file.scan()?.into_array_stream()?; let host_file = SESSION.open_options().open_path(host_path.path()).await?; let mut host_batches = host_file.scan()?.into_array_stream()?; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index fefe61c44f3..8232a9a6018 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -92,6 +92,13 @@ struct Args { /// plain one. Intended to be run as its own pass. #[arg(long)] gpu_verify: bool, + /// Read the Vortex GPU file with direct IO, bypassing the page cache. + /// + /// Off by default: cuDF reads through the page cache after an untimed warm-up, so direct + /// IO would compare a Vortex read of the disk against a cuDF read of RAM. Turn it on to + /// measure storage bandwidth instead, and do not read the ratio as a decode comparison. + #[arg(long)] + gpu_direct_io: bool, #[arg(short, long, default_value_t, value_enum)] display_format: DisplayFormat, #[arg(short, long)] @@ -120,6 +127,7 @@ async fn main() -> anyhow::Result<()> { let gpu = args.gpu_decompress.then_some(GpuOptions { codec: args.gpu_parquet_codec, verify: args.gpu_verify, + direct_io: args.gpu_direct_io, }); let (formats, ops) = if gpu.is_some() { @@ -151,6 +159,8 @@ struct GpuOptions { codec: GpuCodec, /// Cross-check decompressed output against the CPU decoders. verify: bool, + /// Read the Vortex file with direct IO instead of through the page cache. + direct_io: bool, } /// Get a compressor for the given format. @@ -159,7 +169,7 @@ fn get_compressor(format: Format, gpu: Option) -> Box { - Box::new(GpuVortexCompressor::new(gpu.verify)) as Box + Box::new(GpuVortexCompressor::new(gpu.verify, gpu.direct_io)) as Box } Format::Parquet => Box::new(GpuParquetCompressor::new(gpu.codec, gpu.verify)), _ => unimplemented!("GPU compress bench not implemented for {format}"), diff --git a/scripts/cudf-parquet-read.py b/scripts/cudf-parquet-read.py index d597762e114..39e435c3a77 100644 --- a/scripts/cudf-parquet-read.py +++ b/scripts/cudf-parquet-read.py @@ -20,6 +20,7 @@ import json import sys import time +from datetime import date def synchronize() -> None: @@ -36,13 +37,31 @@ def synchronize() -> None: pass +def normalize(frame): + """Collapses representation differences that are not value differences. + + A Parquet DATE column comes back from pyarrow as a column of `datetime.date` + objects but from cuDF as `datetime64[s]`. Those hold the same instants, yet + `check_dtype=False` does not bridge them because one side is `object`, so the + comparison reports every row as different. Coercing both sides to datetime64 + compares the dates themselves. + """ + import pandas as pd + + for name in frame.columns: + column = frame[name] + if column.dtype == object and len(column) and isinstance(column.iloc[0], date): + frame[name] = pd.to_datetime(column) + return frame + + def verify(path: str, frame) -> None: """Fails unless the GPU read matches a CPU Parquet read of the same file.""" import pandas as pd from pandas.testing import assert_frame_equal - expected = pd.read_parquet(path) - actual = frame.to_pandas() + expected = normalize(pd.read_parquet(path)) + actual = normalize(frame.to_pandas()) # cuDF and pyarrow can land on different-but-equivalent dtypes (nullable vs numpy # backed, for instance), so compare values and leave dtype policy out of it. diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 9f1fce7e68e..6da059f7cd5 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -23,6 +23,7 @@ use vortex::array::arrays::varbinview::BinaryView; use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::legacy_session; +use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; @@ -36,6 +37,26 @@ pub trait CanonicalCudaExt { Self: Sized; } +/// Copies an array-backed validity mask back to the host. +/// +/// Only [`Validity::Array`] owns a buffer; the other variants are metadata and pass through. +/// Migrating the values of a nullable array without its validity leaves the mask on the +/// device, and the first host read of it — canonicalising to Arrow, say — panics in +/// `BufferHandle::unwrap_host`. +#[allow(clippy::disallowed_methods)] +async fn validity_into_host(validity: Validity) -> VortexResult { + let Validity::Array(array) = validity else { + return Ok(validity); + }; + Ok(Validity::Array( + array + .execute::(&mut legacy_session().create_execution_ctx())? + .into_host() + .await? + .into_array(), + )) +} + #[async_trait] impl CanonicalCudaExt for Canonical { #[allow(clippy::disallowed_methods)] @@ -67,13 +88,11 @@ impl CanonicalCudaExt for Canonical { struct_fields.names().clone(), host_fields, len, - validity, + validity_into_host(validity).await?, ))) } n @ Canonical::Null(_) => Ok(n), Canonical::Bool(bool) => { - // NOTE: update to copy to host when adding buffer handle. - // Also update other method to copy validity to host. let len = bool.len(); let validity = bool.validity()?; let BoolDataParts { bits, meta } = bool.into_data().into_parts(len); @@ -83,7 +102,10 @@ impl CanonicalCudaExt for Canonical { meta.len(), meta.offset(), ); - Ok(Canonical::Bool(BoolArray::new(bits, validity))) + Ok(Canonical::Bool(BoolArray::new( + bits, + validity_into_host(validity).await?, + ))) } Canonical::Primitive(prim) => { let PrimitiveDataParts { @@ -95,7 +117,7 @@ impl CanonicalCudaExt for Canonical { Ok(Canonical::Primitive(PrimitiveArray::from_byte_buffer( buffer.try_into_host()?.await?, ptype, - validity, + validity_into_host(validity).await?, ))) } Canonical::Decimal(decimal) => { @@ -106,6 +128,7 @@ impl CanonicalCudaExt for Canonical { validity, .. } = decimal.into_data_parts(); + let validity = validity_into_host(validity).await?; Ok(Canonical::Decimal(unsafe { DecimalArray::new_unchecked_handle( BufferHandle::new_host(values.try_into_host()?.await?), @@ -136,6 +159,7 @@ impl CanonicalCudaExt for Canonical { let host_buffers = try_join_all(host_buffers).await?; let host_buffers: Arc<[ByteBuffer]> = Arc::from(host_buffers); + let validity = validity_into_host(validity).await?; Ok(Canonical::VarBinView(unsafe { VarBinViewArray::new_unchecked(host_views, host_buffers, dtype, validity) })) From b9fed6ad16f89d1218829355fd38c2155bc16348 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 14 Aug 2026 09:05:58 +0000 Subject: [PATCH 14/16] CI: measure a CPU decompression baseline on the same GPU runner The GPU ratio says which of two GPU readers is faster, not whether either beats the CPU decoders. Run the same binary over the same datasets with the CPU path on the same machine and publish it alongside, so the GPU numbers can be read against something. Also capture the benchmark's exit status rather than letting `shell: bash`'s -e skip the `cat`, which kept the timing tables out of the job log and left them only in the PR comment. Signed-off-by: Claude --- .github/workflows/pr-bench-gpu-compress.yml | 28 ++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index 9131518615b..505a89a2a75 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -135,9 +135,29 @@ jobs: # (cuda.yaml), which also sets only FLAT_LAYOUT_INLINE_ARRAY_NODE. FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" run: | + # `shell: bash` implies -e, so a non-zero benchmark exit would skip the cat and leave + # the timing tables out of the job log entirely. Capture the status instead. + status=0 bash scripts/bench-taskset.sh target/release_debug/compress-bench \ - --gpu-decompress -d table > gpu-compress.txt + --gpu-decompress -d table > gpu-compress.txt || status=$? cat gpu-compress.txt + exit "$status" + - name: Run CPU decompression baseline + id: cpu-bench + continue-on-error: true + shell: bash + env: + RUST_BACKTRACE: full + FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + # Same binary, same datasets, same machine, CPU decoders. Without this the GPU ratio + # says which of two GPU readers is faster but not whether either beats the CPU, which + # is the only number that says whether GPU decoding is worth doing at all. + run: | + bash scripts/bench-taskset.sh target/release_debug/compress-bench \ + --ops decompress \ + --datasets 'TPC-H l_comment|taxi|Arade|Bimbo|CMSprovider|Euro2016|Food|HashTags' \ + -d table > cpu-compress.txt + cat cpu-compress.txt - name: Publish results shell: bash run: | @@ -147,6 +167,12 @@ jobs: echo '```text' cat gpu-compress.txt echo '```' + echo + echo "## CPU baseline (same machine, same datasets)" + echo + echo '```text' + cat cpu-compress.txt || echo "CPU baseline did not produce output" + echo '```' } > comment.md cat comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR From 793f9a0796f72fa20e22700bbbe903ae0fb9d298 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 14 Aug 2026 10:51:52 +0000 Subject: [PATCH 15/16] CI: reduce the GPU workflow change to the two required installs The workflow had grown a verification pass, a CPU baseline pass, two extra PR comments and a separate failure gate. None of that is needed to report a Parquet and a Vortex number. Reset the file to its develop version and add back only what the benchmark cannot run without: - the DuckDB CLI, which builds the Public BI Parquet fixtures; - uv and the cuDF wheel, which the GPU Parquet backend shells out to. Every pre-existing step is now untouched, and `--gpu-verify` remains available as a local flag. Signed-off-by: Claude --- .github/workflows/pr-bench-gpu-compress.yml | 103 +------------------- 1 file changed, 4 insertions(+), 99 deletions(-) diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index 505a89a2a75..f683e2310f0 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -34,7 +34,7 @@ jobs: enable-sccache: "true" - name: Install DuckDB # The Public BI datasets are converted from CSV to Parquet through the DuckDB CLI, as - # in bench-pr.yml. Without it those datasets fail to materialise their fixture. + # 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 @@ -42,9 +42,8 @@ jobs: - name: Install uv uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 - name: Install cuDF - # The GPU Parquet number is a full cuDF `read_parquet`. cuDF ships prebuilt manylinux - # wheels on NVIDIA's index, so it stays a runtime dependency and never enters the Rust - # build. pandas and pyarrow back the --gpu-verify cross-check. + # 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 \ @@ -63,68 +62,8 @@ jobs: cargo build --locked --package compress-bench --profile release_debug --features cuda,unstable_encodings - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh - - name: Verify GPU decompression correctness - id: verify - shell: bash - continue-on-error: true - env: - RUST_BACKTRACE: "1" - FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" - # Cross-checks every GPU-decompressed page and field against the CPU decoders before - # any timings are taken. Verification runs inline, so this pass is not timed. It runs - # every dataset rather than stopping at the first failure, so one run reports the whole - # matrix; the gate step below still fails the job. - run: | - set -o pipefail - target/release_debug/compress-bench \ - --gpu-decompress --gpu-verify --iterations 1 -d table 2>&1 | tee gpu-verify.txt - - name: Publish verification result - shell: bash - # The per-dataset verdicts are worth surfacing on the PR whether or not they all pass: - # digging them out of a multi-thousand-line job log is otherwise the only way to see - # which encodings decode correctly on the GPU. - run: | - { - echo "# GPU decompression verification" - echo - if [ "${{ steps.verify.outcome }}" = "success" ]; then - echo "All GPU datasets matched the CPU decode." - else - echo "Verification failed. Per-dataset results:" - fi - echo - echo '```text' - grep -E "verified [0-9]+|GPU decompression failed|^ - " gpu-verify.txt | tail -40 \ - || tail -40 gpu-verify.txt - echo '```' - # The per-dataset lines above carry only the first line of each error. Python - # tracebacks, Rust backtraces and first-differing-row dumps span several lines, so - # the tail of the raw output goes in a collapsed block rather than back in the log. - if [ "${{ steps.verify.outcome }}" != "success" ]; then - echo - echo "
Full error detail" - echo - echo '```text' - tail -200 gpu-verify.txt - echo '```' - echo - echo "
" - fi - } > verify-comment.md - cat verify-comment.md >> "$GITHUB_STEP_SUMMARY" - - name: Comment PR with verification result - if: github.event.pull_request.head.repo.fork == false - uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b # v3 - with: - file-path: verify-comment.md - comment-tag: bench-pr-comment-gpu-verify - name: Run GPU compression benchmark - id: bench shell: bash - # A dataset the GPU cannot decode is reported per dataset and still fails the run, but - # the timing tables are printed first, so the datasets that do decode publish their - # numbers. The gate at the end of the job turns either failure into a job failure. - continue-on-error: true env: RUST_BACKTRACE: full # Do not enable VORTEX_EXPERIMENTAL_PATCHED_ARRAY here: it rewrites interior @@ -135,29 +74,9 @@ jobs: # (cuda.yaml), which also sets only FLAT_LAYOUT_INLINE_ARRAY_NODE. FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" run: | - # `shell: bash` implies -e, so a non-zero benchmark exit would skip the cat and leave - # the timing tables out of the job log entirely. Capture the status instead. - status=0 bash scripts/bench-taskset.sh target/release_debug/compress-bench \ - --gpu-decompress -d table > gpu-compress.txt || status=$? + --gpu-decompress -d table > gpu-compress.txt cat gpu-compress.txt - exit "$status" - - name: Run CPU decompression baseline - id: cpu-bench - continue-on-error: true - shell: bash - env: - RUST_BACKTRACE: full - FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" - # Same binary, same datasets, same machine, CPU decoders. Without this the GPU ratio - # says which of two GPU readers is faster but not whether either beats the CPU, which - # is the only number that says whether GPU decoding is worth doing at all. - run: | - bash scripts/bench-taskset.sh target/release_debug/compress-bench \ - --ops decompress \ - --datasets 'TPC-H l_comment|taxi|Arade|Bimbo|CMSprovider|Euro2016|Food|HashTags' \ - -d table > cpu-compress.txt - cat cpu-compress.txt - name: Publish results shell: bash run: | @@ -167,12 +86,6 @@ jobs: echo '```text' cat gpu-compress.txt echo '```' - echo - echo "## CPU baseline (same machine, same datasets)" - echo - echo '```text' - cat cpu-compress.txt || echo "CPU baseline did not produce output" - echo '```' } > comment.md cat comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR @@ -190,11 +103,3 @@ jobs: GPU Compression failed. Check the [workflow run](${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}) for details. comment-tag: bench-pr-comment-gpu-compress - - name: Fail if verification or the benchmark failed - if: steps.verify.outcome == 'failure' || steps.bench.outcome == 'failure' - shell: bash - # Last, so both the verification matrix and the timing tables are already published. - run: | - echo "verify=${{ steps.verify.outcome }} bench=${{ steps.bench.outcome }}" >&2 - echo "GPU decompression failed; see the verification and results comments." >&2 - exit 1 From 8e060ca5a2af9406dc507916b381048195379fe1 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 14 Aug 2026 11:01:09 +0000 Subject: [PATCH 16/16] Give both GPU formats the same physical partition size Parquet was read as ~1M-row row groups while Vortex inherited the Arrow reader's ~8K-row batches, so the Vortex file carried hundreds of small chunks where Parquet carried a handful of row groups. A chunk is the unit the reader plans and dispatches over, so that turned single dispatches into hundreds of small kernel launches and made the two numbers measure different shapes of work rather than two decoders. Pin both to `GPU_ROW_GROUP_SIZE` (1,048,576 rows, Parquet's `DEFAULT_MAX_ROW_GROUP_ROW_COUNT`): - state the Parquet row group count explicitly so it moves with the constant; - add `parquet_to_vortex_chunks_with_batch_size`, which concatenates the source batches and re-slices on exact boundaries. Setting the Arrow reader's batch size alone is not enough, because the reader also breaks at the source file's row group boundaries and still emits short batches; - write those batches through as root chunks with `ChunkedLayoutStrategy` and read them back with `SplitBy::RowCount(GPU_ROW_GROUP_SIZE)`. Signed-off-by: Claude --- benchmarks/compress-bench/README.md | 17 +++++ benchmarks/compress-bench/src/gpu_vortex.rs | 46 +++++++++--- benchmarks/compress-bench/src/gpu_writer.rs | 12 ++++ vortex-bench/src/conversions.rs | 79 ++++++++++++++++----- 4 files changed, 128 insertions(+), 26 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 5033bca34a7..ef0515f4216 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -74,6 +74,23 @@ Set in `src/gpu_writer.rs`: | 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 diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 832466978a0..a36f5a249b8 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -24,12 +24,14 @@ use vortex::array::arrays::struct_::StructArrayExt; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; +use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex::layout::layouts::compressed::CompressingStrategy; +use vortex::layout::scan::split_by::SplitBy; use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::compress::Compressor; -use vortex_bench::conversions::parquet_to_vortex_chunks; +use vortex_bench::conversions::parquet_to_vortex_chunks_with_batch_size; use vortex_cuda::CanonicalCudaExt; use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; @@ -39,6 +41,8 @@ use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; use vortex_cuda::layout::register_cuda_layout; +use crate::gpu_writer::GPU_ROW_GROUP_SIZE; + /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. pub struct GpuVortexCompressor { verify: bool, @@ -69,13 +73,24 @@ impl Compressor for GpuVortexCompressor { async fn decompress(&self, parquet_path: &Path) -> Result { register_cuda_layout(&SESSION); - let uncompressed = parquet_to_vortex_chunks(parquet_path.to_path_buf()).await?; + // Rebatch to the same partition size the GPU Parquet file is written with. Left alone, + // the Arrow reader hands back ~8K-row batches, each of which becomes its own Vortex + // chunk and its own set of kernel launches. + let uncompressed = parquet_to_vortex_chunks_with_batch_size( + parquet_path.to_path_buf(), + Some(GPU_ROW_GROUP_SIZE), + ) + .await?; let gpu_file = NamedTempFile::new()?; let mut output = tokio::fs::File::create(gpu_file.path()).await?; - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().only_cuda_compatible()) - .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) - .build(); + // Write those batches straight through as root chunks, so a chunk on disk is one + // partition rather than whatever the default strategy would regroup them into. + let strategy = Arc::new(ChunkedLayoutStrategy::new(CompressingStrategy::new( + CudaFlatLayoutStrategy::default(), + BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .build(), + ))); SESSION .write_options() .with_strategy(strategy) @@ -91,7 +106,12 @@ impl Compressor for GpuVortexCompressor { let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; let start = Instant::now(); let file = open_gpu(gpu_file.path(), self.direct_io).await?; - let mut batches = file.scan()?.into_array_stream()?; + // Split reads on the same boundary the file was written with, so a scan batch is one + // partition instead of a sub-slice of one. + let mut batches = file + .scan()? + .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE)) + .into_array_stream()?; while let Some(batch) = batches.next().await { let record = batch?.execute::(cuda_ctx.execution_ctx())?; @@ -146,9 +166,15 @@ async fn verify_against_host_scan(path: &Path, direct_io: bool) -> Result WriterProperties { .set_dictionary_enabled(true) .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) + // Stated explicitly rather than left to the default, because the Vortex side is + // rebatched to the same constant and the two have to move together. + .set_max_row_group_row_count(Some(GPU_ROW_GROUP_SIZE)) // Per-page statistics only inflate the page headers a reader has to walk. .set_statistics_enabled(EnabledStatistics::Chunk) .build() diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 69cba42c6b0..f9c4ee63734 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -6,6 +6,8 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use arrow_array::RecordBatch; +use arrow_select::concat::concat_batches; use futures::StreamExt; use futures::TryStreamExt; use parquet::arrow::AsyncArrowWriter; @@ -96,18 +98,73 @@ fn calculate_concurrency() -> usize { /// Note: This loads the entire file into memory. For large files, use the streaming conversion like /// in [`parquet_to_vortex_stream`] instead. pub async fn parquet_to_vortex_chunks(parquet_path: PathBuf) -> anyhow::Result { + parquet_to_vortex_chunks_with_batch_size(parquet_path, None).await +} + +/// Read a Parquet file as a Vortex [`ChunkedArray`] with chunks of exactly `batch_size` rows. +/// +/// With `batch_size` set, the source batches are concatenated and re-sliced on exact boundaries, +/// so every chunk but the last has the requested length. Setting the Arrow reader's batch size +/// is not enough on its own: the reader also breaks at the source file's row group boundaries, +/// so a file whose row groups are not a multiple of the batch size still yields short batches. +/// +/// This matters when comparing against a format whose physical partitioning is explicit. Chunk +/// size becomes the Vortex file's partition size, and small chunks mean many small compressed +/// blocks — and, on the GPU, many small kernel launches. +/// +/// `None` keeps whatever batches the Parquet reader produces. +pub async fn parquet_to_vortex_chunks_with_batch_size( + parquet_path: PathBuf, + batch_size: Option, +) -> anyhow::Result { let file = File::open(parquet_path).await?; let builder = ParquetRecordBatchStreamBuilder::new(file).await?; - let reader = builder.build()?; - let chunks: Vec = parquet_to_vortex_stream(reader) - .map(|r| r.map_err(anyhow::Error::from)) + let Some(batch_size) = batch_size.filter(|size| *size > 0) else { + let chunks: Vec = parquet_to_vortex_stream(builder.build()?) + .map(|r| r.map_err(anyhow::Error::from)) + .try_collect() + .await?; + return Ok(ChunkedArray::from_iter(chunks)); + }; + + let batches: Vec = builder + .with_batch_size(batch_size) + .build()? + .map_err(anyhow::Error::from) .try_collect() .await?; + let schema = batches + .first() + .map(RecordBatch::schema) + .ok_or_else(|| anyhow::anyhow!("cannot convert an empty Parquet file"))?; + let combined = concat_batches(&schema, &batches)?; + + let mut chunks = Vec::with_capacity(combined.num_rows().div_ceil(batch_size)); + for start in (0..combined.num_rows()).step_by(batch_size) { + let len = batch_size.min(combined.num_rows() - start); + chunks.push(record_batch_to_vortex(combined.slice(start, len))?); + } + Ok(ChunkedArray::from_iter(chunks)) } +/// Convert one Arrow [`RecordBatch`] into a canonical Vortex array. +fn record_batch_to_vortex(batch: RecordBatch) -> VortexResult { + let schema = batch.schema(); + let chunk = SESSION.arrow().from_arrow_record_batch(batch, &schema)?; + let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); + + // Canonicalize the chunk. + chunk.append_to_builder( + builder.as_mut(), + &mut VortexSession::default().create_execution_ctx(), + )?; + + Ok(builder.finish()) +} + /// Create a streaming Vortex array from a Parquet reader. /// /// Streams record batches and converts them to Vortex arrays on-the-fly, avoiding loading the @@ -116,19 +173,9 @@ pub fn parquet_to_vortex_stream( reader: ParquetRecordBatchStream, ) -> impl futures::Stream> { reader.map(move |result| { - result.map_err(|e| vortex_err!(External: e)).and_then(|rb| { - let schema = rb.schema(); - let chunk = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; - let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); - - // Canonicalize the chunk. - chunk.append_to_builder( - builder.as_mut(), - &mut VortexSession::default().create_execution_ctx(), - )?; - - Ok(builder.finish()) - }) + result + .map_err(|e| vortex_err!(External: e)) + .and_then(record_batch_to_vortex) }) }