diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index cd97755e636..f683e2310f0 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -32,6 +32,24 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} enable-sccache: "true" + - name: Install DuckDB + # The Public BI datasets are converted from CSV to Parquet through the DuckDB CLI, as + # in pr-bench-compress.yml. Without it those datasets cannot build their fixture. + run: | + wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.5/duckdb_cli-linux-amd64.zip | funzip > duckdb + chmod +x duckdb + echo "$PWD" >> "$GITHUB_PATH" + - name: Install uv + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + - name: Install cuDF + # The GPU Parquet number is a cuDF `read_parquet`. cuDF ships prebuilt manylinux wheels + # on NVIDIA's index, so it stays a runtime dependency and never enters the Rust build. + run: | + uv venv --python 3.12 .venv-cudf + uv pip install --python .venv-cudf \ + --extra-index-url https://pypi.nvidia.com \ + cudf-cu12 pandas pyarrow + echo "$PWD/.venv-cudf/bin" >> "$GITHUB_PATH" - uses: ./.github/actions/system-info - name: Display NVIDIA GPU details run: | diff --git a/Cargo.lock b/Cargo.lock index fee1e3ebd34..a80d6f85f02 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1578,6 +1578,8 @@ dependencies = [ "lance-bench", "parquet 58.4.0", "regex", + "serde", + "serde_json", "tempfile", "tokio", "tracing", diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 4046a12d42e..7b1dadc3209 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -27,6 +27,8 @@ itertools = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } tempfile = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } @@ -45,7 +47,6 @@ name = "compress-bench" test = false [lib] -test = false [lints] workspace = true diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index bf2d3efc1db..ef0515f4216 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -15,13 +15,100 @@ See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--format cargo run -p compress-bench --profile release_debug ``` -GPU decompression is opt-in and runs only the existing benchmark names allow-listed in -`src/main.rs`: +## GPU decompression + +`--gpu-decompress` is opt-in, requires the `cuda` feature, and restricts the suite to the +GPU dataset list in `src/main.rs`. It measures decompression only, for two backends: + +- **Vortex** — the file is written with CUDA-compatible BtrBlocks encodings only + (`only_cuda_compatible`) and a CUDA flat layout, then decoded on the device all the way to + canonical arrays. +- **Parquet** — the file is rewritten with GPU-friendly writer settings (see below) and read + back with [cuDF](https://github.com/rapidsai/cudf)'s `read_parquet`, which performs the + whole read on the device: page header decode, codec decompression, dictionary/RLE/plain + decoding and column assembly. + +Both sides therefore decode all the way to device-resident arrays, which is what makes the +`vortex:parquet- gpu ratio decompress time` metric a like-for-like comparison. ```bash cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress + +# pick the Parquet page codec the GPU file is written with (default: snappy) +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress --gpu-parquet-codec zstd +``` + +### cuDF + +cuDF is reached through its prebuilt `cudf-cu12` wheel, so it is a runtime dependency of the +benchmark and never enters the Rust build: + +```bash +uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarrow +``` + +`scripts/cudf-parquet-read.py` performs and times the read. Timing is taken inside that +script, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up +read runs first for the same reason. + +Both backends read a warm file by default. cuDF runs an untimed warm-up read before the timed +one, so its timed read hits the page cache; the Vortex reader therefore does **not** use direct +I/O by default, because `O_DIRECT` would bypass the page cache and compare a Vortex read of the +disk against a cuDF read of RAM. `--gpu-direct-io` turns it back on to measure storage bandwidth +instead — a different question, and the resulting ratio is not a decode comparison. + +The remaining asymmetry is the transfer path: the Vortex reader uses pinned buffers, while cuDF +does its own host read and host-to-device copy. + +### GPU-friendly Parquet writer settings + +Set in `src/gpu_writer.rs`: + +| Setting | Value | Why | +| --- | --- | --- | +| writer version | `PARQUET_1_0` | v1 pages compress the whole page body; v2 pages put uncompressed levels ahead of the compressed values in the same body. | +| compression | Snappy (default) or Zstd | Snappy is the Parquet default and has the higher device-side throughput. | +| dictionary | enabled | Keeps the decompressed payload small; the encoding GPU Parquet readers decode fastest. | +| data page size | 1 MiB | Large enough to amortize per-page setup, small enough to keep every SM fed. Matches the page size cuDF targets. | +| data page row limit | 1,000,000 | The 20k-row default caps narrow columns' pages far below 1 MiB. | +| statistics | chunk-level | Page statistics only inflate the headers a reader has to walk. | +| row group size | 1,048,576 rows | Shared with the Vortex side as `GPU_ROW_GROUP_SIZE` — see below. | + +### Matching physical partitions + +A Parquet row group and a Vortex chunk are the same thing for this comparison: the unit the +reader plans and dispatches over. Both formats are pinned to `GPU_ROW_GROUP_SIZE` +(1,048,576 rows, Parquet's `DEFAULT_MAX_ROW_GROUP_ROW_COUNT`). + +Without this the two are not comparable. Parquet reads ~1M-row row groups, while the Vortex +side inherits the Arrow reader's ~8K-row batches — each of which becomes its own chunk, its own +compressed blocks and its own kernel launches, so a single dispatch turns into hundreds. + +Setting the Arrow reader's batch size alone is not enough: the reader also breaks at the source +file's row group boundaries, so short batches survive. `parquet_to_vortex_chunks_with_batch_size` +therefore concatenates the source batches and re-slices them on exact boundaries. Those batches +are written straight through as root chunks via `ChunkedLayoutStrategy`, and read back with +`SplitBy::RowCount(GPU_ROW_GROUP_SIZE)` so a scan batch is one whole partition. + +### Correctness + +`--gpu-verify` cross-checks device output against the CPU decoders on every iteration: + +- Parquet: the cuDF-read frame is compared against a CPU Parquet read of the same file. +- Vortex: each GPU-decoded field is copied back and compared against the same field decoded + on the CPU, through Arrow with a pinned target type. + +Verification runs inline, so timings from a verifying run are not comparable to a plain one — +run it as its own pass: + +```bash +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress --gpu-verify --iterations 1 ``` -On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure -storage bandwidth rather than page-cache hits. +Any `--gpu-decompress` run reports on every dataset rather than stopping at the first failure, so +one run shows which datasets decode correctly on the GPU and which do not. The timing tables are +rendered before the failure summary, so a dataset the GPU cannot decode still leaves the rest of +the matrix with numbers — the process exits non-zero either way. diff --git a/benchmarks/compress-bench/src/gpu_parquet.rs b/benchmarks/compress-bench/src/gpu_parquet.rs new file mode 100644 index 00000000000..3ee0297f1a0 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_parquet.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! GPU Parquet decompression backend, timed through cuDF. +//! +//! cuDF's `read_parquet` performs the whole read on the device — page header decode, +//! codec decompression, dictionary/RLE/plain decoding and column assembly — which makes it +//! the like-for-like opponent for the Vortex GPU backend, which likewise decodes all the way +//! to canonical arrays on device. +//! +//! cuDF is reached through its prebuilt `cudf-cu12` wheel rather than by linking libcudf, so +//! it stays a runtime dependency of this benchmark and never enters the Rust build. The +//! measurement is taken inside [`CUDF_SCRIPT`], so interpreter start, `import cudf` and CUDA +//! context creation are excluded; only the reads themselves are timed. + +use std::path::Path; +use std::process::Command; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Context; +use anyhow::Result; +use anyhow::bail; +use anyhow::ensure; +use arrow_array::RecordBatch; +use async_trait::async_trait; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use serde::Deserialize; +use tempfile::NamedTempFile; +use vortex_bench::Format; +use vortex_bench::compress::Compressor; + +use crate::gpu_writer::GpuCodec; +use crate::gpu_writer::gpu_writer_properties; + +/// Repo-relative path of the script that performs and times the cuDF read. +const CUDF_SCRIPT: &str = "scripts/cudf-parquet-read.py"; + +/// Parquet compressor whose decompression measurement is a full cuDF GPU read. +pub struct GpuParquetCompressor { + codec: GpuCodec, + verify: bool, +} + +/// What the cuDF script reports back. +#[derive(Debug, Deserialize)] +struct CudfReadReport { + /// Fastest timed read, in nanoseconds. + min_ns: u64, + rows: u64, + columns: u64, +} + +impl GpuParquetCompressor { + /// Create a backend that writes pages with `codec` and times cuDF reading them back. + /// + /// When `verify` is set, the GPU read is cross-checked against a CPU Parquet read of the + /// same file before the measurement is reported. + pub fn new(codec: GpuCodec, verify: bool) -> Self { + Self { codec, verify } + } + + /// Rewrite the source Parquet file with GPU-friendly writer settings. + fn write_gpu_parquet(&self, parquet_path: &Path) -> Result<(NamedTempFile, u64)> { + let builder = ParquetRecordBatchReaderBuilder::try_new(std::fs::File::open(parquet_path)?)?; + let schema = Arc::clone(builder.schema()); + let batches: Vec = 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)) + } +} + +#[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 report = run_cudf_read(gpu_file.path(), self.verify)?; + + ensure!( + report.rows > 0 && report.columns > 0, + "cuDF read {} rows and {} columns, expected a non-empty table", + report.rows, + report.columns + ); + + Ok(Duration::from_nanos(report.min_ns)) + } +} + +/// Runs the cuDF read script and returns the timing it measured. +fn run_cudf_read(path: &Path, verify: bool) -> Result { + let mut command = Command::new("python3"); + command.arg(CUDF_SCRIPT).arg(path); + if verify { + command.arg("--verify"); + } + + let output = command.output().with_context(|| { + format!("failed to run {CUDF_SCRIPT}; is cudf-cu12 installed on this host?") + })?; + + if !output.status.success() { + bail!( + "{CUDF_SCRIPT} exited with {}:\n{}", + output.status, + String::from_utf8_lossy(&output.stderr).trim() + ); + } + + serde_json::from_slice(&output.stdout).with_context(|| { + format!( + "could not parse the report from {CUDF_SCRIPT}: {}", + String::from_utf8_lossy(&output.stdout).trim() + ) + }) +} diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 3dbb68bc7a8..a36f5a249b8 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -8,20 +8,31 @@ use std::time::Duration; 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; 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; 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; #[cfg(target_os = "linux")] @@ -30,8 +41,24 @@ 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; +pub struct GpuVortexCompressor { + verify: bool, + direct_io: 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, direct_io: bool) -> Self { + Self { verify, direct_io } + } +} #[async_trait] impl Compressor for GpuVortexCompressor { @@ -46,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) @@ -61,16 +99,19 @@ impl Compressor for GpuVortexCompressor { output.sync_all().await?; drop(output); + if self.verify { + 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 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 mut batches = file.scan()?.into_array_stream()?; + let file = open_gpu(gpu_file.path(), self.direct_io).await?; + // 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())?; @@ -83,3 +124,177 @@ impl Compressor for GpuVortexCompressor { Ok(start.elapsed()) } } + +/// Opens a Vortex file for CUDA execution. +/// +/// `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 = 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?) +} + +/// 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, 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 + // 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 + // 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, direct_io).await?; + let mut gpu_batches = gpu_file + .scan()? + .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE)) + .into_array_stream()?; + let host_file = SESSION.open_options().open_path(host_path.path()).await?; + let mut host_batches = host_file + .scan()? + .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE)) + .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) { + (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::(&mut host_ctx)?; + ensure!( + gpu_record.len() == host_record.len(), + "batch {batch_index} 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(), + "batch {batch_index} field count differs between the GPU and CPU scans" + ); + + 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, + &mut host_ctx, + batch_index, + field_index, + )?; + fields_checked += 1; + } + + batch_index += 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( + 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)?; + + 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(), + ); + + 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 +} diff --git a/benchmarks/compress-bench/src/gpu_writer.rs b/benchmarks/compress-bench/src/gpu_writer.rs new file mode 100644 index 00000000000..f7a495250bc --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_writer.rs @@ -0,0 +1,100 @@ +// 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::DEFAULT_MAX_ROW_GROUP_ROW_COUNT; +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; + +/// Rows per physical partition in both GPU benchmark formats. +/// +/// A Parquet row group and a Vortex chunk are the same thing for this comparison: the unit the +/// reader plans and dispatches over. Pinning both to one value is what makes the two numbers +/// comparable — otherwise Parquet reads ~1M-row row groups while Vortex inherits the Arrow +/// reader's ~8K-row batches, turning one launch into hundreds. +pub const GPU_ROW_GROUP_SIZE: usize = DEFAULT_MAX_ROW_GROUP_ROW_COUNT; + +/// 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) + // 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() +} + +#[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 68039996605..76e34172f6a 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 gpu_writer; pub mod parquet; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 1b1603e52c8..8232a9a6018 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -1,16 +1,23 @@ // 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; +use anyhow::Context; 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::gpu_writer::GpuCodec; use compress_bench::parquet::ParquetCompressor; use compress_bench::vortex::VortexCompressor; +use futures::FutureExt; use indicatif::ProgressBar; use itertools::Itertools; use regex::Regex; @@ -35,6 +42,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 +75,30 @@ 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 file is written with. + /// + /// 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. + /// + /// 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, + /// 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)] @@ -97,8 +124,17 @@ 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, + direct_io: args.gpu_direct_io, + }); + + let (formats, ops) = if gpu.is_some() { + ( + vec![Format::Parquet, Format::OnDiskVortex], + vec![CompressOp::Decompress], + ) } else { (args.formats, args.ops) }; @@ -108,7 +144,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 +152,33 @@ async fn main() -> anyhow::Result<()> { .await } +/// Settings for the GPU decompression mode. +#[derive(Clone, Copy, Debug)] +struct GpuOptions { + /// Parquet page codec to write the GPU file with. + 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. -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, gpu.direct_io)) 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 +205,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 +232,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 +271,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 { @@ -225,18 +295,37 @@ async fn run_compress( let mut measurements = vec![]; let mut v3_records: Vec = Vec::new(); + // 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() { - let (m, mut records) = run_benchmark_for_dataset( - &progress, - &formats, - &ops, - iterations, - dataset_handle, - gpu_decompress, - ) - .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); @@ -249,6 +338,8 @@ async fn run_compress( 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)?; @@ -260,13 +351,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( @@ -275,7 +386,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 +402,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 { @@ -302,7 +413,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 @@ -333,7 +445,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() @@ -342,7 +455,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 +474,53 @@ 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, + }); +} + +/// 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() + } +} diff --git a/scripts/cudf-parquet-read.py b/scripts/cudf-parquet-read.py new file mode 100644 index 00000000000..39e435c3a77 --- /dev/null +++ b/scripts/cudf-parquet-read.py @@ -0,0 +1,120 @@ +#!/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 +from datetime import date + + +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 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 = 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. + 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-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) }) } 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/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) })) 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(()) + } }