From 9f250999d40a2213f91bda93e67797f5f2ad34b8 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 14:19:06 +0000 Subject: [PATCH 01/35] feat(cuda): add projected dictionary-free file scans Signed-off-by: Alexander Droste --- Cargo.lock | 1 + vortex-cuda/ffi/Cargo.toml | 3 + vortex-cuda/ffi/README.md | 14 +- vortex-cuda/ffi/build.rs | 34 ++ vortex-cuda/ffi/cbindgen.toml | 72 ++++ vortex-cuda/ffi/cinclude/vortex_cuda.h | 182 ++++++--- vortex-cuda/ffi/src/lib.rs | 461 +++++++++++++++------- vortex-cuda/ffi/src/tests/projection.rs | 419 ++++++++++++++++++++ vortex-cuda/src/arrow/canonical.rs | 41 +- vortex-cuda/src/arrow/dictionary_tests.rs | 350 ++++++++++++++++ vortex-cuda/src/arrow/mod.rs | 97 ++--- vortex-cuda/src/executor.rs | 7 + vortex-cuda/src/layout.rs | 135 ++++++- vortex-cuda/src/lib.rs | 1 + vortex-cuda/src/session.rs | 25 ++ 15 files changed, 1576 insertions(+), 266 deletions(-) create mode 100644 vortex-cuda/ffi/build.rs create mode 100644 vortex-cuda/ffi/cbindgen.toml create mode 100644 vortex-cuda/ffi/src/tests/projection.rs create mode 100644 vortex-cuda/src/arrow/dictionary_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 275db7ef8ee..1f52ea42973 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10911,6 +10911,7 @@ name = "vortex-cuda-ffi" version = "0.1.0" dependencies = [ "arrow-schema 59.3.0", + "cbindgen", "futures", "vortex", "vortex-cuda", diff --git a/vortex-cuda/ffi/Cargo.toml b/vortex-cuda/ffi/Cargo.toml index c1e1efdb23c..32ec00124ac 100644 --- a/vortex-cuda/ffi/Cargo.toml +++ b/vortex-cuda/ffi/Cargo.toml @@ -24,6 +24,9 @@ vortex-ffi = { path = "../../vortex-ffi" } [dev-dependencies] vortex-cuda-macros = { workspace = true } +[build-dependencies] +cbindgen = { workspace = true } + [lib] name = "vortex_cuda_ffi" crate-type = ["rlib", "staticlib", "cdylib"] diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index d83974921ec..aedd25046e9 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -1,10 +1,9 @@ # vortex-cuda-ffi -CUDA-specific C FFI helpers for cuDF interop. - -This crate keeps CUDA out of the base `vortex-ffi` crate. Its public C API exports a borrowed `vx_array` as an `ArrowSchema + ArrowDeviceArray` pair. - -It does not create cuDF objects itself. The caller passes the exported Arrow Device structs to cuDF and releases them after cuDF is done importing. +CUDA-specific C FFI helpers for cuDF interop, keeping CUDA out of the base `vortex-ffi` +crate. The public C API exports a borrowed `vx_array` as an `ArrowSchema + ArrowDeviceArray` +pair, not cuDF objects. The caller passes these structs to cuDF and releases them after +cuDF finishes importing. Use this crate as the CUDA-enabled FFI artifact. Include both headers: @@ -27,3 +26,8 @@ pool and CUDA state are reused as well. On Linux, use `vx_cuda_scan_path_arrow_device_stream_with_options` with `vx_cuda_scan_options.flags = VX_CUDA_SCAN_FLAG_DIRECT_IO` to bypass the operating system page cache for pooled data-plane reads. Footer and zone-map reads remain buffered on the host. + +`build.rs` generates `cinclude/vortex_cuda.h` with cbindgen on stable Rust, without macro +expansion. Edit the API and docs in `src/lib.rs`, not the generated header; commit +regenerated headers with API changes. `cbindgen.toml` supplies the standard Arrow Device +interface compatibility preamble. diff --git a/vortex-cuda/ffi/build.rs b/vortex-cuda/ffi/build.rs new file mode 100644 index 00000000000..10218d19e94 --- /dev/null +++ b/vortex-cuda/ffi/build.rs @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::env; +use std::error::Error; +use std::path::PathBuf; +use std::process::Command; + +fn main() -> Result<(), Box> { + println!("cargo:rerun-if-changed=src"); + println!("cargo:rerun-if-changed=cbindgen.toml"); + println!("cargo:rerun-if-changed=build.rs"); + + let crate_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR")?); + let header = crate_dir.join("cinclude/vortex_cuda.h"); + // The CUDA API needs no macro expansion or dependency parsing, so generate on stable Rust + // without recursively building the CUDA implementation. + cbindgen::Builder::new() + .with_src(crate_dir.join("src/lib.rs")) + .with_config(cbindgen::Config::from_file( + crate_dir.join("cbindgen.toml"), + )?) + .generate()? + .write_to_file(&header); + if !Command::new("clang-format") + .args(["--style=file", "-i"]) + .arg(&header) + .status() + .is_ok_and(|status| status.success()) + { + println!("cargo:warning=clang-format unavailable or failed; CUDA header left unformatted"); + } + Ok(()) +} diff --git a/vortex-cuda/ffi/cbindgen.toml b/vortex-cuda/ffi/cbindgen.toml new file mode 100644 index 00000000000..1343336dd3e --- /dev/null +++ b/vortex-cuda/ffi/cbindgen.toml @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +language = "C" +braces = "SameLine" +cpp_compat = true +usize_is_size_t = true +style = "both" +no_includes = true + +header = """ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#pragma once + +// THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY + +#include +#include + +#include "vortex.h" + +/* Link against the CUDA-enabled FFI library that provides both the base Vortex FFI and these CUDA + * entry points. Do not pass Vortex handles between independently linked Rust FFI libraries. */ + +/* Definitions from the Arrow C Device data interface. Define USE_OWN_ARROW_DEVICE to skip them. + * https://arrow.apache.org/docs/format/CDeviceDataInterface.html */ +#if !defined(ARROW_C_DEVICE_DATA_INTERFACE) && !defined(USE_OWN_ARROW_DEVICE) +#define ARROW_C_DEVICE_DATA_INTERFACE + +typedef int32_t ArrowDeviceType; +#define ARROW_DEVICE_CPU 1 +#define ARROW_DEVICE_CUDA 2 +#define ARROW_DEVICE_CUDA_HOST 3 +#define ARROW_DEVICE_OPENCL 4 +#define ARROW_DEVICE_VULKAN 7 +#define ARROW_DEVICE_METAL 8 +#define ARROW_DEVICE_VPI 9 +#define ARROW_DEVICE_ROCM 10 +#define ARROW_DEVICE_ROCM_HOST 11 +#define ARROW_DEVICE_EXT_DEV 12 +#define ARROW_DEVICE_CUDA_MANAGED 13 +#define ARROW_DEVICE_ONEAPI 14 +#define ARROW_DEVICE_WEBGPU 15 +#define ARROW_DEVICE_HEXAGON 16 + +struct ArrowDeviceArray { + struct ArrowArray array; + int64_t device_id; + ArrowDeviceType device_type; + void *sync_event; + int64_t reserved[3]; +}; +#endif + +#if !defined(ARROW_C_DEVICE_STREAM_INTERFACE) && !defined(USE_OWN_ARROW_DEVICE) +#define ARROW_C_DEVICE_STREAM_INTERFACE +struct ArrowDeviceArrayStream { + ArrowDeviceType device_type; + int (*get_schema)(struct ArrowDeviceArrayStream *, struct ArrowSchema *out); + int (*get_next)(struct ArrowDeviceArrayStream *, struct ArrowDeviceArray *out); + const char *(*get_last_error)(struct ArrowDeviceArrayStream *); + void (*release)(struct ArrowDeviceArrayStream *); + void *private_data; +}; +#endif +""" + +# These externally defined Arrow ABI types use struct tags, not typedef names. +[export.rename] +"ArrowDeviceArray" = "struct ArrowDeviceArray" +"ArrowDeviceArrayStream" = "struct ArrowDeviceArrayStream" diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 3c83b9b9ffd..f0e13d29f8b 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #pragma once +// THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY + #include #include @@ -10,10 +12,6 @@ /* Link against the CUDA-enabled FFI library that provides both the base Vortex FFI and these CUDA * entry points. Do not pass Vortex handles between independently linked Rust FFI libraries. */ -#ifdef __cplusplus -extern "C" { -#endif - /* Definitions from the Arrow C Device data interface. Define USE_OWN_ARROW_DEVICE to skip them. * https://arrow.apache.org/docs/format/CDeviceDataInterface.html */ #if !defined(ARROW_C_DEVICE_DATA_INTERFACE) && !defined(USE_OWN_ARROW_DEVICE) @@ -56,20 +54,55 @@ struct ArrowDeviceArrayStream { }; #endif +/** + * Bypass the operating system page cache for pooled data-plane reads. + * Footer and zone-map reads remain buffered. Supported only on Linux. + */ +#define VX_CUDA_SCAN_FLAG_DIRECT_IO (1u << 0) + +/** + * Options for scanning a CUDA-compatible Vortex file. + * + * Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. + */ +typedef struct vx_cuda_scan_options { + /** + * A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are ignored. + */ + uint32_t flags; + /** + * Maximum rows in each output batch. Zero uses layout-derived splitting. + * Physical layout boundaries may produce shorter batches. + */ + size_t batch_rows; +} vx_cuda_scan_options; + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + /** * Create a CUDA Vortex session. * - * Repeated `vx_cuda_array_export_arrow_device` calls reuse this CUDA state. Returns an owned - * session handle, or NULL and an optional `vx_error` on failure. + * Repeated [`vx_cuda_array_export_arrow_device`] calls reuse this CUDA state. Returns an owned + * session handle, or null and an optional `vx_error` on failure. + * + * # Safety + * + * If `error_out` is non-null, it must be valid for writing one error pointer. */ vx_session *vx_cuda_session_new(vx_error **error_out); /** * Open a Vortex file sink configured to produce CUDA-readable files. * - * Push host-resident arrays and close or abort the returned sink with the standard - * `vx_array_sink_*` functions. This API configures the on-disk encodings and layout; it does not - * move arrays to the GPU during the write. + * Push host arrays and close/abort with `vx_array_sink_*`. Only on-disk encodings and layouts + * change; writing does not move arrays to the GPU. + * + * # Safety + * + * `session`, `path`, and `dtype` follow `vx_array_sink_open_file`'s requirements. + * Non-null `error_out` must be writable for one error pointer. */ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, vx_view path, @@ -80,13 +113,15 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, * Open a CUDA-readable Vortex file sink with a fixed row block size. * * `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default - * writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Passing any - * nonzero value disables this byte-size coalescing, so passing 8,192 is not equivalent to passing - * zero. + * writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Any nonzero + * value disables byte-size coalescing and outer layout dictionaries, so passing 8,192 is not + * equivalent to passing zero. + * + * Write and scan sizing are independent; scan batches preserve on-disk layout boundaries. * - * Write and scan sizing are independent. To align on-disk row blocks with scan batches, pass the - * same nonzero value to this function and `vx_cuda_scan_path_arrow_device_stream_batch_rows`; the - * API does not enforce a match. + * # Safety + * + * Same requirements as [`vx_cuda_array_sink_open_file`]. */ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session, vx_view path, @@ -95,30 +130,24 @@ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session vx_error **error_out); /** - * Options for scanning a CUDA-compatible Vortex file. + * Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. * - * Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. - */ -/** Bypass the operating system page cache for pooled data-plane reads. - * Footer and zone-map reads remain buffered. Supported only on Linux. */ -#define VX_CUDA_SCAN_FLAG_DIRECT_IO (UINT32_C(1) << 0) - -typedef struct vx_cuda_scan_options { - /** Bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. */ - uint32_t flags; - /** Number of rows in each output ArrowDeviceArray. Zero uses layout-derived splitting. */ - size_t batch_rows; -} vx_cuda_scan_options; - -/** - * Scan a local CUDA-compatible Vortex file as an Arrow C Device stream. + * Requires CUDA-supported encodings/layouts, such as files from [`vx_cuda_array_sink_open_file`]. + * Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers + * reused across scans with the same CUDA session. + * + * Dictionaries, including nested children, decode on CUDA for a stable plain Arrow schema, + * without changing session policy. Decoding can increase device memory use and requires CUDA + * support for device-resident dictionaries. + * + * Returns `0` with an owned `out_stream`; release it and each batch via their Arrow callbacks. + * Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. * - * Files written by `vx_cuda_array_sink_open_file` are compatible with this path. Reusing the same - * CUDA session across calls also reuses the pinned host buffers used to stage file reads. + * # Safety * - * On success returns 0 and writes an owned `ArrowDeviceArrayStream` to `out_stream`. The caller - * must release the stream and each produced `ArrowDeviceArray` through their embedded Arrow - * release callbacks. On error returns 1 and writes a `vx_error` to `error_out` when non-NULL. + * `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the + * duration of this call and contain UTF-8. `out_stream` must be a valid writable pointer. If + * `error_out` is non-null, it must be valid for writing one error pointer. */ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, vx_view path, @@ -126,14 +155,15 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, vx_error **error_out); /** - * Scan a local CUDA-compatible Vortex file with fixed-size row batches. + * Scan a local Vortex file with bounded row batches. * - * `batch_rows` controls the number of rows in each output `ArrowDeviceArray`. Pass zero to use the - * layout-derived splitting of `vx_cuda_scan_path_arrow_device_stream`. + * Uses [`vx_cuda_scan_path_arrow_device_stream`]'s export and ownership rules. + * `batch_rows` caps output rows; zero uses layout splitting. Physical boundaries may shorten + * batches. Scan and write sizing are independent; scans preserve on-disk layout boundaries. * - * Scan and write sizing are independent. To align scan batches with on-disk row blocks, pass the - * same nonzero value to this function and `vx_cuda_array_sink_open_file_block_rows`; the API does - * not enforce a match. + * # Safety + * + * Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]. */ int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, vx_view path, @@ -142,29 +172,61 @@ int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, vx_error **error_out); /** - * Scan a local CUDA-compatible Vortex file with explicit options. + * Like [`vx_cuda_scan_path_arrow_device_stream`], with explicit scan options. + * + * Null or zero-initialized `options` selects buffered I/O and layout-derived batch splitting. + * + * # Safety * - * This has the same ownership and file compatibility requirements as - * `vx_cuda_scan_path_arrow_device_stream`. Pass NULL or a zero-initialized options struct to use - * buffered file I/O and layout-derived batch splitting. + * Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]; non-null `options` must + * point to a valid [`vx_cuda_scan_options`]. */ int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session, vx_view path, - const vx_cuda_scan_options *options, + const struct vx_cuda_scan_options *options, struct ArrowDeviceArrayStream *out_stream, vx_error **error_out); +/** + * Scan a local Vortex file with ordered top-level column projection. + * + * Same options, ownership, and file requirements as + * [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes column I/O/decoding. + * Names are literal and case-sensitive; unknown/duplicate names and non-struct files are rejected. + * `ncolumns == 0` ignores `columns` and selects all. Names are copied; empty files retain the + * projected schema. Errors leave `out_stream` unchanged. + * + * # Safety + * + * In addition to [`vx_cuda_scan_path_arrow_device_stream_with_options`]'s requirements, + * nonzero `ncolumns` requires that many initialized, aligned [`vx_view`] values at `columns`. + * Each name borrows `len` readable UTF-8 bytes for this call; null is allowed only for zero length. + */ +int vx_cuda_scan_path_arrow_device_stream_projected(const vx_session *session, + vx_view path, + const struct vx_cuda_scan_options *options, + const vx_view *columns, + size_t ncolumns, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); + /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. * - * On success returns 0 and writes independently releasable `out_schema` and `out_array`; the caller - * passes them to cuDF and releases both via their embedded Arrow callbacks after import. On error - * returns 1 and, when `error_out` is non-NULL, writes a `vx_error` (free with `vx_error_free`). + * Returns `0` with independently owned `out_schema` and `out_array`. Pass them to cuDF, then + * release both via their Arrow callbacks after import. Returns `1` on error, writing a `vx_error` + * if `error_out` is non-null; free it with `vx_error_free`. * * `out_array` is exported on `ARROW_DEVICE_CUDA`; struct arrays become table-shaped schemas, * non-struct arrays a single column field. * * Export is stream-ordered; `out_array->sync_event` is valid until `out_array` is released. + * + * # Safety + * + * `session` and `array` must be valid borrowed handles created by `vortex-ffi`. `out_schema` + * and `out_array` must be valid writable pointers. If `error_out` is non-null, it must be valid + * for writing one error pointer. */ int vx_cuda_array_export_arrow_device(const vx_session *session, const vx_array *array, @@ -175,16 +237,16 @@ int vx_cuda_array_export_arrow_device(const vx_session *session, /** * Consume a Vortex partition and scan it as an Arrow C Device stream. * - * This function takes ownership of `partition`. Callers must not free or reuse - * it after calling this function, regardless of success or failure. + * Consumes `partition` on success or failure; never free or reuse it afterward. + * Returns `0` with an owned `out_stream` retaining the scan iterator. Release the stream and + * each produced batch via their Arrow release callbacks. + * Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. * - * On success returns 0 and writes an owned `ArrowDeviceArrayStream` to - * `out_stream`. The stream owns the resulting scan iterator. The caller must - * release the stream through its embedded Arrow `release` callback, and must - * release each produced `ArrowDeviceArray` through its embedded - * `ArrowArray.release` callback. + * # Safety * - * On error returns 1 and writes a `vx_error` to `error_out` when non-NULL. + * `session` must be a valid borrowed handle created by `vortex-ffi`. `partition` must be an owned + * partition handle created by `vortex-ffi`. `out_stream` must be a valid writable pointer. If + * `error_out` is non-null, it must be valid for writing one error pointer. */ int vx_cuda_partition_scan_arrow_device_stream(const vx_session *session, vx_partition *partition, @@ -192,5 +254,5 @@ int vx_cuda_partition_scan_arrow_device_stream(const vx_session *session, vx_error **error_out); #ifdef __cplusplus -} -#endif +} // extern "C" +#endif // __cplusplus diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 44fc7d87938..050eadfa0c6 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -13,20 +13,31 @@ use std::ptr; use std::sync::Arc; use arrow_schema::ffi::FFI_ArrowSchema; +use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamExt; use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::dtype::FieldName; +use vortex::dtype::FieldNames; use vortex::editions::ComponentKind; use vortex::editions::EditionSessionExt; use vortex::error::VortexResult; use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex::expr::root; +use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; use vortex::file::WriteStrategyBuilder; use vortex::io::runtime::BlockingRuntime; +use vortex::layout::LayoutStrategy; +use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::layout::scan::split_by::SplitBy; use vortex::session::SessionExt; use vortex::session::VortexSession; +use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; +use vortex_cuda::DictionaryExport; use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::arrow::ArrowDeviceArray; use vortex_cuda::arrow::ArrowDeviceArrayStream; @@ -52,9 +63,9 @@ use vortex_ffi::vx_view; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; -/// Enable direct I/O for pooled CUDA file reads. -pub const VX_CUDA_SCAN_FLAG_DIRECT_IO: u32 = 1 << 0; -const VX_CUDA_SCAN_KNOWN_FLAGS: u32 = VX_CUDA_SCAN_FLAG_DIRECT_IO; +/// Bypass the operating system page cache for pooled data-plane reads. +/// Footer and zone-map reads remain buffered. Supported only on Linux. +pub const VX_CUDA_SCAN_FLAG_DIRECT_IO: u32 = 1u32 << 0; /// Options for scanning a CUDA-compatible Vortex file. /// @@ -62,20 +73,42 @@ const VX_CUDA_SCAN_KNOWN_FLAGS: u32 = VX_CUDA_SCAN_FLAG_DIRECT_IO; #[repr(C)] #[derive(Default)] pub struct vx_cuda_scan_options { - /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. + /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are ignored. pub flags: u32, - /// Number of rows in each output batch. Zero uses layout-derived splitting. + /// Maximum rows in each output batch. Zero uses layout-derived splitting. + /// Physical layout boundaries may produce shorter batches. pub batch_rows: usize, } -/// Return a Vortex session with a [`CudaSession`] session variable. -/// -/// If `session` already has CUDA support, this returns a clone of it. Otherwise it -/// returns a new session cloned from `session` with a default [`CudaSession`] attached. -fn session_with_cuda(session: &VortexSession) -> VortexResult { +/// Initialize CUDA support on `session` and return the same borrow. +fn session_with_cuda(session: &VortexSession) -> &VortexSession { session.get::(); register_cuda_layout(session); - Ok(session.clone()) + session +} + +/// Build a CUDA-flat writer using only session-enabled encodings. +fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { + let allowed_encodings = session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + let mut strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder( + BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .retain_allowed_encodings(&allowed_encodings), + ) + .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); + if block_rows > 0 { + // Preserve explicit row blocks: outer layout dictionaries can split a high-cardinality + // block into u16-sized dictionary runs, while a byte target can coalesce adjacent blocks. + strategy = strategy + .with_probe_compressor(BtrBlocksCompressorBuilder::empty().build()) + .with_row_block_size(block_rows) + .with_data_block_target_bytes(None); + } + strategy.build() } /// Create a CUDA Vortex session. @@ -102,15 +135,13 @@ pub unsafe extern "C-unwind" fn vx_cuda_session_new( /// Open a Vortex file sink configured to produce CUDA-readable files. /// -/// Push host-resident arrays and close or abort the returned sink with the standard -/// `vx_array_sink_*` functions. This function configures the on-disk encodings and layout; it does -/// not move arrays to the GPU during the write. +/// Push host arrays and close/abort with `vx_array_sink_*`. Only on-disk encodings and layouts +/// change; writing does not move arrays to the GPU. /// /// # Safety /// -/// `session`, `path`, and `dtype` must satisfy the same requirements as -/// `vx_array_sink_open_file`. If `error_out` is non-null, it must be valid for writing one error -/// pointer. +/// `session`, `path`, and `dtype` follow `vx_array_sink_open_file`'s requirements. +/// Non-null `error_out` must be writable for one error pointer. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( session: *const vx_session, @@ -123,18 +154,16 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( /// Open a CUDA-readable Vortex file sink with a fixed row block size. /// -/// `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero preserves the -/// default writer strategy used by [`vx_cuda_array_sink_open_file`]. Any nonzero value disables -/// byte-size coalescing so data blocks retain the requested row granularity. +/// `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default +/// writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Any nonzero +/// value disables byte-size coalescing and outer layout dictionaries, so passing 8,192 is not +/// equivalent to passing zero. /// -/// Write and scan sizing are independent. To align on-disk row blocks with scan batches, pass the -/// same nonzero value to this function and [`vx_cuda_scan_path_arrow_device_stream_batch_rows`]. +/// Write and scan sizing are independent; scan batches preserve on-disk layout boundaries. /// /// # Safety /// -/// `session`, `path`, and `dtype` must satisfy the same requirements as -/// `vx_array_sink_open_file`. If `error_out` is non-null, it must be valid for writing one error -/// pointer. +/// Same requirements as [`vx_cuda_array_sink_open_file`]. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( session: *const vx_session, @@ -144,47 +173,30 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( error_out: *mut *mut vx_error, ) -> *mut vx_array_sink { try_or(error_out, ptr::null_mut(), || { - let vortex_session = unsafe { vx_session_ref(session) }?; - session_with_cuda(vortex_session)?; - let allowed_encodings = vortex_session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - let mut strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default() - .only_cuda_compatible() - .retain_allowed_encodings(&allowed_encodings), + let vortex_session = session_with_cuda(unsafe { vx_session_ref(session) }?); + unsafe { + vx_array_sink_open_file_with_strategy( + session, + path, + dtype, + cuda_write_strategy(vortex_session, block_rows), ) - .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); - if block_rows > 0 { - // The default byte-size target can coalesce several row blocks into one data block. - // A scan using the same row count would then split inside that data block, defeating - // the requested alignment. The explicit block row count already defines the desired - // granularity for this opt-in path, so a separate byte-size target is unnecessary. - strategy = strategy - .with_row_block_size(block_rows) - .with_data_block_target_bytes(None); } - unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy.build()) } }) } /// Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. /// -/// Footer and zone-map reads remain on the host. Data segments are staged through pinned host -/// buffers and transferred directly to the GPU. -/// -/// The file must use encodings and layouts supported by the CUDA execution path, such as files -/// written by [`vx_cuda_array_sink_open_file`]. Pinned staging buffers are reused across scans made -/// with the same CUDA session. +/// Requires CUDA-supported encodings/layouts, such as files from [`vx_cuda_array_sink_open_file`]. +/// Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers +/// reused across scans with the same CUDA session. /// -/// On success returns `0` and writes an owned [`ArrowDeviceArrayStream`] to `out_stream`. The -/// caller must release the stream and each array produced by it through their embedded Arrow -/// release callbacks. +/// Dictionaries, including nested children, decode on CUDA for a stable plain Arrow schema, +/// without changing session policy. Decoding can increase device memory use and requires CUDA +/// support for device-resident dictionaries. /// -/// On error returns `1` and, when `error_out` is non-null, writes a `vx_error` (free with -/// `vx_error_free`). +/// Returns `0` with an owned `out_stream`; release it and each batch via their Arrow callbacks. +/// Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. /// /// # Safety /// @@ -209,19 +221,15 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( } } -/// Scan a local Vortex file and export an Arrow C Device stream with fixed-size row batches. -/// -/// `batch_rows` controls the number of rows in each output batch. Passing zero preserves the -/// layout-derived splitting used by [`vx_cuda_scan_path_arrow_device_stream`]. +/// Scan a local Vortex file with bounded row batches. /// -/// Scan and write sizing are independent. To align scan batches with on-disk row blocks, pass the -/// same nonzero value to this function and [`vx_cuda_array_sink_open_file_block_rows`]. +/// Uses [`vx_cuda_scan_path_arrow_device_stream`]'s export and ownership rules. +/// `batch_rows` caps output rows; zero uses layout splitting. Physical boundaries may shorten +/// batches. Scan and write sizing are independent; scans preserve on-disk layout boundaries. /// /// # Safety /// -/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the -/// duration of this call and contain UTF-8. `out_stream` must be a valid writable pointer. If -/// `error_out` is non-null, it must be valid for writing one error pointer. +/// Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows( session: *const vx_session, @@ -245,18 +253,14 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows } } -/// Scan a local Vortex file with explicit options and export an Arrow C Device stream. +/// Like [`vx_cuda_scan_path_arrow_device_stream`], with explicit scan options. /// -/// This has the same ownership and file compatibility requirements as -/// [`vx_cuda_scan_path_arrow_device_stream`]. Pass a null `options` pointer or a zero-initialized -/// [`vx_cuda_scan_options`] to use buffered file I/O and layout-derived batch splitting. +/// Null or zero-initialized `options` selects buffered I/O and layout-derived batch splitting. /// /// # Safety /// -/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the -/// duration of this call and contain UTF-8. `options`, when non-null, must point to a valid -/// [`vx_cuda_scan_options`]. `out_stream` must be a valid writable pointer. If `error_out` is -/// non-null, it must be valid for writing one error pointer. +/// Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]; non-null `options` must +/// point to a valid [`vx_cuda_scan_options`]. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_options( session: *const vx_session, @@ -264,55 +268,160 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_optio options: *const vx_cuda_scan_options, out_stream: *mut ArrowDeviceArrayStream, error_out: *mut *mut vx_error, +) -> c_int { + // SAFETY: The caller supplies valid borrowed inputs and writable outputs. + unsafe { + vx_cuda_scan_path_arrow_device_stream_projected( + session, + path, + options, + ptr::null(), + 0, + out_stream, + error_out, + ) + } +} + +/// Scan a local Vortex file with ordered top-level column projection. +/// +/// Same options, ownership, and file requirements as +/// [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes column I/O/decoding. +/// Names are literal and case-sensitive; unknown/duplicate names and non-struct files are rejected. +/// `ncolumns == 0` ignores `columns` and selects all. Names are copied; empty files retain the +/// projected schema. Errors leave `out_stream` unchanged. +/// +/// # Safety +/// +/// In addition to [`vx_cuda_scan_path_arrow_device_stream_with_options`]'s requirements, +/// nonzero `ncolumns` requires that many initialized, aligned [`vx_view`] values at `columns`. +/// Each name borrows `len` readable UTF-8 bytes for this call; null is allowed only for zero length. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_projected( + session: *const vx_session, + path: vx_view, + options: *const vx_cuda_scan_options, + columns: *const vx_view, + ncolumns: usize, + out_stream: *mut ArrowDeviceArrayStream, + error_out: *mut *mut vx_error, ) -> c_int { try_or(error_out, VX_CUDA_ERR, || { vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); - let path = unsafe { path.as_str() }?.to_owned(); - let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; + // SAFETY: The caller keeps the borrowed options and column views alive for this call. let options = unsafe { scan_options(options) }?; - let array_stream = ffi_runtime().block_on(async { - let file = session + let columns = unsafe { scan_columns(columns, ncolumns) }?; + let path = unsafe { path.as_str() }?; + let session = session_with_cuda(unsafe { vx_session_ref(session) }?); + let file = ffi_runtime().block_on( + session .open_options() .with_cuda() .with_read_at_options(options.read_at_options) - .open_path(path) - .await?; - let scan = file.scan()?; - let scan = if options.batch_rows == 0 { - scan - } else { - scan.with_split_by(SplitBy::RowCount(options.batch_rows)) - }; - Ok::<_, vortex::error::VortexError>(scan.into_array_stream()?.boxed()) - })?; - let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?; + .open_path(path), + )?; + let scan = projected_scan(&file, columns, options.batch_rows)?; + let array_stream = scan.into_array_stream()?.boxed(); + let ctx = scan_export_ctx(session)?; + let device_stream = ArrowDeviceArrayStream::new(array_stream, ctx, ffi_runtime()); unsafe { ptr::write(out_stream, device_stream) }; Ok(VX_CUDA_OK) }) } +/// Copy ordered, literal field names, rejecting invalid views and duplicates. Zero selects all. +/// +/// # Safety +/// +/// Inputs passing null/alignment/size checks must reference `ncolumns` live views at `columns`, +/// with `len` readable bytes at each non-null name pointer. +unsafe fn scan_columns(columns: *const vx_view, ncolumns: usize) -> VortexResult { + if ncolumns == 0 { + return Ok(FieldNames::default()); + } + vortex_ensure!( + !columns.is_null(), + "null CUDA scan columns with nonzero count" + ); + vortex_ensure!(columns.is_aligned(), "unaligned CUDA scan columns pointer"); + vortex_ensure!( + ncolumns <= isize::MAX as usize / size_of::(), + "CUDA scan column count is too large" + ); + // SAFETY: Null, alignment, and size were checked; the caller guarantees readable views. + let columns = unsafe { std::slice::from_raw_parts(columns, ncolumns) }; + let mut names = Vec::::with_capacity(ncolumns); + for (index, column) in columns.iter().enumerate() { + vortex_ensure!( + column.len <= isize::MAX as usize, + "CUDA scan column {index} name is too long" + ); + // SAFETY: The caller guarantees readable name bytes. as_str checks null and UTF-8. + let name = unsafe { column.as_str() } + .map_err(|error| vortex_err!("invalid CUDA scan column {index}: {error}"))?; + vortex_ensure!( + !names.iter().any(|existing| existing.as_ref() == name), + "duplicate CUDA scan column: {name:?}" + ); + names.push(FieldName::from(name)); + } + Ok(names.into()) +} + +/// Apply projection before column reads; row limits subdivide, never merge, layout splits. +fn projected_scan( + file: &VortexFile, + columns: FieldNames, + batch_rows: usize, +) -> VortexResult> { + let mut scan = file.scan()?; + if !columns.is_empty() { + let fields = file.dtype().as_struct_fields_opt().ok_or_else(|| { + vortex_err!("CUDA scan column projection requires a struct file dtype") + })?; + for name in columns.iter() { + vortex_ensure!( + fields.find(name).is_some(), + "unknown CUDA scan column: {name:?}" + ); + } + let projection = select(columns, root()).optimize_recursive(file.dtype())?; + scan = scan.with_projection(projection.bind(file.dtype())?); + } + if batch_rows != 0 { + let max_rows = u64::try_from(batch_rows) + .map_err(|_| vortex_err!("CUDA scan batch row count is too large"))?; + scan = scan.with_split_by(SplitBy::LayoutSubSplitting { max_rows }); + } + Ok(scan) +} + struct CudaScanOptions { read_at_options: PooledFileReadAtOptions, batch_rows: usize, } -unsafe fn scan_options(options: *const vx_cuda_scan_options) -> VortexResult { - let (flags, batch_rows) = if options.is_null() { - (0, 0) - } else { - let options = unsafe { &*options }; - (options.flags, options.batch_rows) - }; - vortex_ensure!( - flags & !VX_CUDA_SCAN_KNOWN_FLAGS == 0, - "unsupported CUDA scan option flags: {:#x}", - flags & !VX_CUDA_SCAN_KNOWN_FLAGS - ); +/// Select plain Arrow export for one scan without mutating the shared session. +fn scan_export_ctx(session: &VortexSession) -> VortexResult { + Ok( + CudaSession::create_execution_ctx(session)? + .with_dictionary_export(DictionaryExport::Decode), + ) +} +/// Parse scan settings; null selects defaults and unknown flags are ignored. +/// +/// # Safety +/// +/// Non-null `options` must point to an initialized, aligned [`vx_cuda_scan_options`]. +unsafe fn scan_options(options: *const vx_cuda_scan_options) -> VortexResult { + let defaults = vx_cuda_scan_options::default(); + // SAFETY: The caller guarantees that a non-null options pointer is valid for this call. + let options = unsafe { options.as_ref() }.unwrap_or(&defaults); let read_at_options = PooledFileReadAtOptions::default(); - let read_at_options = if flags & VX_CUDA_SCAN_FLAG_DIRECT_IO == 0 { + let read_at_options = if options.flags & VX_CUDA_SCAN_FLAG_DIRECT_IO == 0 { read_at_options } else { #[cfg(target_os = "linux")] @@ -329,16 +438,15 @@ unsafe fn scan_options(options: *const vx_cuda_scan_options) -> VortexResult VortexResult<()> { let options = vx_cuda_scan_options::default(); assert_eq!(options.flags, 0); assert_eq!(options.batch_rows, 0); + for pointer in [ptr::null(), &raw const options] { + // SAFETY: Each pointer is either null or points to the live options above. + let parsed = unsafe { scan_options(pointer) }?; + assert_eq!(parsed.read_at_options, PooledFileReadAtOptions::default()); + assert_eq!(parsed.batch_rows, 0); + } + Ok(()) } #[test] - fn rejects_unknown_scan_option_flags() { - let options = vx_cuda_scan_options { - flags: 1 << 31, - ..Default::default() - }; - assert!(unsafe { scan_options(&raw const options) }.is_err()); - } - - #[cfg(target_os = "linux")] - #[test] - fn maps_direct_io_scan_option_to_pooled_reader() -> VortexResult<()> { - let options = vx_cuda_scan_options { - flags: VX_CUDA_SCAN_FLAG_DIRECT_IO, - ..Default::default() - }; - assert_eq!( - unsafe { scan_options(&raw const options) }?.read_at_options, - PooledFileReadAtOptions::default().with_direct_io() - ); + fn maps_scan_options_and_ignores_unknown_flags() -> VortexResult<()> { + let buffered = PooledFileReadAtOptions::default(); + for (flags, batch_rows, read_at_options) in [ + (0, 8192, buffered), + (1 << 1, 0, buffered), + #[cfg(target_os = "linux")] + (VX_CUDA_SCAN_FLAG_DIRECT_IO, 0, buffered.with_direct_io()), + #[cfg(target_os = "linux")] + ( + VX_CUDA_SCAN_FLAG_DIRECT_IO | (1 << 1), + 8192, + buffered.with_direct_io(), + ), + ] { + let options = vx_cuda_scan_options { flags, batch_rows }; + // SAFETY: options lives for the duration of parsing. + let parsed = unsafe { scan_options(&raw const options) }?; + assert_eq!(parsed.read_at_options, read_at_options, "flags={flags}"); + assert_eq!(parsed.batch_rows, batch_rows, "flags={flags}"); + } Ok(()) } - #[test] - fn maps_batch_rows_scan_option() -> VortexResult<()> { - let options = vx_cuda_scan_options { - batch_rows: 8192, - ..Default::default() - }; + #[cuda_test] + fn scan_decodes_dictionaries_and_reuses_session_resources() -> VortexResult<()> { + // A distinct allocator identity detects accidental reconstruction of a default session. + let allocator = BufferAllocatorRef::new(StaticBufferAllocator); + let session = VortexSession::default() + .with_some(CudaSession::try_default()?) + .with_allocator(allocator.clone()); + let mut export_ctx = scan_export_ctx(session_with_cuda(&session))?; + assert!(export_ctx.execution_ctx().allocator().ptr_eq(&allocator)); + let export_session = export_ctx.execution_ctx().session(); + assert!(Arc::ptr_eq( + session.get::().pinned_buffer_pool(), + export_session.get::().pinned_buffer_pool(), + )); + let array = DictArray::try_new( + PrimitiveArray::from_iter([1u8, 0, 1]).into_array(), + PrimitiveArray::from_iter([10i32, 20]).into_array(), + )? + .into_array(); + for (ctx, expected_type, preserved) in [ + (export_ctx, DataType::Int32, false), + ( + CudaSession::create_execution_ctx(&session)?, + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Int32)), + true, + ), + ] { + let mut stream = ArrowDeviceArrayStream::new( + array.clone().to_array_stream().boxed(), + ctx, + ffi_runtime(), + ); + let get_schema = stream.get_schema.expect("missing get_schema"); + let get_next = stream.get_next.expect("missing get_next"); + let release = stream.release.expect("missing release"); + let mut schema = FFI_ArrowSchema::empty(); + let mut exported = empty_device_array(); + // SAFETY: The live stream owns these callbacks, and both outputs are writable. + unsafe { + assert_eq!(get_schema(&raw mut stream, (&raw mut schema).cast()), 0); + assert_eq!(get_next(&raw mut stream, &raw mut exported), 0); + } + assert_eq!(Field::try_from(&schema)?.data_type(), &expected_type); + assert_eq!(!exported.array.dictionary.is_null(), preserved); + // SAFETY: The batch and stream are live and released exactly once. + unsafe { + release_device_array(&mut exported); + release(&raw mut stream); + } + } assert_eq!( - unsafe { scan_options(&raw const options) }?.batch_rows, - 8192 + session.get::().dictionary_export(), + DictionaryExport::Preserve ); + assert!(session.allocator().ptr_eq(&allocator)); + Ok(()) } diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs new file mode 100644 index 00000000000..bbdd8213a03 --- /dev/null +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -0,0 +1,419 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ffi::CStr; +use std::io::Write; +use std::mem::MaybeUninit; +use std::path::PathBuf; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use futures::TryStreamExt; +use vortex::array::VortexSessionExecute; +use vortex::arrow::ArrowSessionExt; +use vortex::buffer::ByteBuffer; +use vortex::buffer::ByteBufferMut; +use vortex::file::WriteOptionsSessionExt; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutStrategy; +use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex::layout::layouts::table::TableStrategy; +use vortex::layout::segments::SegmentFuture; +use vortex::layout::segments::SegmentId; +use vortex::layout::segments::SegmentSource; + +use super::*; + +fn view(value: &str) -> vx_view { + vx_view { + ptr: value.as_ptr().cast(), + len: value.len(), + } +} + +fn names(values: &[&str]) -> VortexResult { + let views: Vec<_> = values.iter().map(|name| view(name)).collect(); + // SAFETY: All views and their string bytes remain live throughout parsing. + unsafe { scan_columns(views.as_ptr(), views.len()) } +} + +fn assert_error(result: VortexResult, message: &str) { + let error = result.err().expect("expected an error"); + assert!(error.to_string().contains(message), "{error}"); +} + +#[test] +fn test_projection_names_are_owned_and_zero_count_means_all() -> VortexResult<()> { + let parsed = { + let name = String::from("値.x"); + names(&[&name, ""])? + }; + assert_eq!(parsed, ["値.x", ""]); + // SAFETY: Zero count ignores the pointer, including a null pointer. + assert!(unsafe { scan_columns(ptr::null(), 0) }?.is_empty()); + let empty = vx_view { + ptr: ptr::null(), + len: 0, + }; + // SAFETY: A null, zero-length view is a valid empty name. + assert_eq!(unsafe { scan_columns(&raw const empty, 1) }?, [""]); + Ok(()) +} + +#[test] +fn test_projection_rejects_invalid_names_and_counts() { + let invalid_utf8 = vx_view { + ptr: [0xffu8].as_ptr().cast(), + len: 1, + }; + let null_name = vx_view { + ptr: ptr::null(), + len: 1, + }; + let long_name = vx_view { + ptr: "x".as_ptr().cast(), + len: usize::MAX, + }; + let aligned = [view("x"), view("x")]; + let misaligned = aligned.as_ptr().cast::().wrapping_add(1).cast(); + for (columns, count, message) in [ + (ptr::null(), 1, "null CUDA scan columns"), + (aligned.as_ptr(), usize::MAX, "column count is too large"), + (misaligned, 1, "unaligned CUDA scan columns"), + (&raw const invalid_utf8, 1, "invalid utf-8"), + (&raw const null_name, 1, "null vx_view pointer"), + (&raw const long_name, 1, "name is too long"), + (aligned.as_ptr(), 2, "duplicate CUDA scan column: \"x\""), + ] { + // SAFETY: Invalid pointer/length combinations must be rejected before dereferencing; + // all remaining views and bytes are live. + assert_error(unsafe { scan_columns(columns, count) }, message); + } +} + +fn session() -> VortexSession { + VortexSession::default().with_handle(ffi_runtime().handle()) +} + +fn table() -> VortexResult { + Ok(StructArray::try_new( + ["ids", "unused", "値.x"].into(), + vec![ + PrimitiveArray::from_iter(0u32..5).into_array(), + PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]).into_array(), + PrimitiveArray::from_option_iter([Some(10i64), None, Some(30), None, Some(50)]) + .into_array(), + ], + 5, + Validity::NonNullable, + )? + .into_array()) +} + +fn file_bytes( + session: &VortexSession, + array: ArrayRef, + cuda_block_rows: Option, +) -> VortexResult { + let strategy: Arc = if let Some(block_rows) = cuda_block_rows { + register_cuda_layout(session); + cuda_write_strategy(session, block_rows) + } else { + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + Arc::new(TableStrategy::new(Arc::clone(&flat), flat)) + }; + let mut bytes = ByteBufferMut::empty(); + ffi_runtime().block_on( + session + .write_options() + .with_strategy(strategy) + .write(&mut bytes, array.to_array_stream()), + )?; + Ok(bytes.freeze()) +} + +fn open_file( + session: &VortexSession, + array: ArrayRef, + cuda_block_rows: Option, +) -> VortexResult { + session + .open_options() + .open_buffer(file_bytes(session, array, cuda_block_rows)?) +} + +fn assert_arrow_eq( + session: &VortexSession, + actual: ArrayRef, + expected: ArrayRef, +) -> VortexResult<()> { + let mut ctx = session.create_execution_ctx(); + let mut to_data = |array| { + session + .arrow() + .execute_arrow(array, None, &mut ctx) + .map(|array| array.to_data()) + }; + assert_eq!(to_data(actual)?, to_data(expected)?); + Ok(()) +} + +#[test] +fn test_cuda_write_strategy_preserves_high_cardinality_row_blocks() -> VortexResult<()> { + let session = session(); + let unique = 70_000u32; + let ids = PrimitiveArray::from_iter((0..unique).chain(0..unique)).into_array(); + let rows = ids.len(); + let input = + StructArray::try_new(["ids"].into(), vec![ids], rows, Validity::NonNullable)?.into_array(); + let file = open_file(&session, input, Some(rows))?; + let lengths: Vec<_> = ffi_runtime().block_on( + projected_scan(&file, names(&["ids"])?, rows)? + .into_array_stream()? + .map_ok(|batch| batch.len()) + .try_collect(), + )?; + assert_eq!(lengths, [rows]); + Ok(()) +} + +struct RejectSegments { + inner: Arc, + forbidden: Vec, + rejected: AtomicUsize, +} + +impl SegmentSource for RejectSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + if self.forbidden.contains(&id) { + self.rejected.fetch_add(1, Ordering::Relaxed); + return Box::pin(async move { + Err(vortex_err!("unselected column segment requested: {id}")) + }); + } + self.inner.request(id) + } +} + +#[test] +fn test_projection_cpu_never_requests_unselected_column_segments() -> VortexResult<()> { + let session = session(); + let input = table()?; + let columns = names(&["値.x", "ids"])?; + let expected = input + .clone() + .execute::(&mut session.create_execution_ctx())? + .project(columns.as_ref())? + .into_array(); + let file = open_file(&session, input, None)?; + // TableStrategy writes one flat child per column, so child 1 is exactly the unused column. + let children = file.footer().layout().children()?; + let forbidden = children[1].segment_ids(); + assert!(!forbidden.is_empty()); + let source = Arc::new(RejectSegments { + inner: file.segment_source(), + forbidden, + rejected: AtomicUsize::new(0), + }); + let file = file.with_segment_source(Arc::::clone(&source)); + let actual = ffi_runtime().block_on( + projected_scan(&file, columns, 2)? + .into_array_stream()? + .read_all(), + )?; + assert_arrow_eq(&session, actual, expected)?; + assert_eq!(source.rejected.load(Ordering::Relaxed), 0); + // The same reader must fail without projection, proving the guard actually observes reads. + assert_error( + ffi_runtime().block_on( + projected_scan(&file, names(&[])?, 0)? + .into_array_stream()? + .read_all(), + ), + "unselected column segment requested", + ); + assert!(source.rejected.load(Ordering::Relaxed) > 0); + Ok(()) +} + +#[test] +fn test_projection_ffi_validation_without_cuda() { + let mut stream = ArrowDeviceArrayStream { + device_type: -1, + get_schema: None, + get_next: None, + get_last_error: None, + release: None, + private_data: ptr::null_mut(), + }; + let mut error = ptr::null_mut(); + // SAFETY: Output pointers are writable; invalid columns must fail before session/path use. + let status = unsafe { + vx_cuda_scan_path_arrow_device_stream_projected( + ptr::null(), + view(""), + ptr::null(), + ptr::null(), + 1, + &raw mut stream, + &raw mut error, + ) + }; + assert_eq!(status, VX_CUDA_ERR); + assert_eq!(stream.device_type, -1); + assert!(stream.release.is_none()); + assert!(!error.is_null()); + // SAFETY: This call owns the returned error and frees it exactly once. + unsafe { vortex_ffi::vx_error_free(error) }; + // SAFETY: Null output is rejected before any other input is used; error output is optional. + assert_eq!( + unsafe { + vx_cuda_scan_path_arrow_device_stream_projected( + ptr::null(), + view(""), + ptr::null(), + ptr::null(), + 0, + ptr::null_mut(), + ptr::null_mut(), + ) + }, + VX_CUDA_ERR + ); +} + +struct LocalFile(PathBuf); + +impl LocalFile { + fn new(bytes: &[u8]) -> VortexResult { + static NEXT_ID: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "vortex-cuda-ffi-projection-{}-{}.vortex", + std::process::id(), + NEXT_ID.fetch_add(1, Ordering::Relaxed) + )); + let mut file = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&path)?; + let result = Self(path); + file.write_all(bytes)?; + Ok(result) + } +} + +impl Drop for LocalFile { + fn drop(&mut self) { + drop(std::fs::remove_file(&self.0)); + } +} + +fn stream_error(stream: &mut ArrowDeviceArrayStream) -> String { + // SAFETY: The callback and returned C string belong to this live stream. + unsafe { + stream + .get_last_error + .and_then(|callback| callback(stream).as_ref()) + .map(|message| CStr::from_ptr(message).to_string_lossy().into_owned()) + .unwrap_or_default() + } +} + +fn open_stream( + session: &VortexSession, + path: &str, + options: &vx_cuda_scan_options, +) -> ArrowDeviceArrayStream { + let mut output = MaybeUninit::::uninit(); + let mut error = ptr::null_mut(); + let handle = test_session(session.clone()); + let columns = [view("値.x"), view("ids")]; + // SAFETY: All borrowed inputs and writable outputs are live for this call. + let status = unsafe { + vx_cuda_scan_path_arrow_device_stream_projected( + handle, + view(path), + options, + columns.as_ptr(), + columns.len(), + output.as_mut_ptr(), + &raw mut error, + ) + }; + // SAFETY: This is the sole release of the borrowed session handle. + unsafe { free_test_session(handle) }; + assert_eq!(status, VX_CUDA_OK); + assert!(error.is_null()); + // SAFETY: A successful call initialized the stream, which owns its session state. + unsafe { output.assume_init() } +} + +fn stream_schema(stream: &mut ArrowDeviceArrayStream) -> FFI_ArrowSchema { + let mut schema = FFI_ArrowSchema::empty(); + let get_schema = stream.get_schema.expect("missing get_schema"); + // SAFETY: This live stream owns the callback; schema is writable. + assert_eq!( + unsafe { get_schema(stream, (&raw mut schema).cast()) }, + 0, + "{}", + stream_error(stream) + ); + schema +} + +fn batch_lengths(stream: &mut ArrowDeviceArrayStream) -> Vec { + let get_next = stream.get_next.expect("missing get_next"); + let mut lengths = Vec::new(); + loop { + let mut array = empty_device_array(); + // SAFETY: This live stream owns the callback; array is writable. + assert_eq!( + unsafe { get_next(stream, &raw mut array) }, + 0, + "{}", + stream_error(stream) + ); + if array.array.release.is_none() { + break; + } + assert_eq!(array.device_type, ARROW_DEVICE_CUDA); + assert_eq!(array.array.n_children, 2); + lengths.push(array.array.length); + // SAFETY: Each live batch is released exactly once, before requesting the next one. + unsafe { release_device_array(&mut array) }; + } + lengths +} + +#[cuda_test] +fn test_projection_gpu_local_file_schema_and_batch_boundaries() -> VortexResult<()> { + for (block_rows, batch_rows) in [(0, 2), (2, 3)] { + let session = session().with_some(CudaSession::try_default()?); + let file = LocalFile::new(&file_bytes(&session, table()?, Some(block_rows))?)?; + let path = file + .0 + .to_str() + .ok_or_else(|| vortex_err!("non-UTF-8 test path"))?; + let options = vx_cuda_scan_options { + batch_rows, + ..Default::default() + }; + let mut stream = open_stream(&session, path, &options); + // The stream must retain its session state after the caller releases its session. + drop(session); + let mut schema = stream_schema(&mut stream); + let expected_fields = vec![ + Field::new("値.x", DataType::Int64, true), + Field::new("ids", DataType::UInt32, false), + ]; + assert_eq!(Schema::try_from(&schema)?, Schema::new(expected_fields)); + assert_eq!(batch_lengths(&mut stream), [2, 2, 1]); + let release = stream.release.expect("missing release"); + // SAFETY: Both objects are live and released exactly once. + unsafe { + release_schema(&mut schema); + release(&raw mut stream); + } + } + Ok(()) +} diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 4a96187a9db..bed7d072b27 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -65,6 +65,7 @@ use vortex_onpair::OnPair; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::CudaExecutionCtx; +use crate::DictionaryExport; use crate::VarBinExportLayout; use crate::arrow::ARROW_DEVICE_CUDA; use crate::arrow::ArrowArray; @@ -114,7 +115,14 @@ impl ExportDeviceArray for CanonicalDeviceArrayExport { array: ArrayRef, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let array = rebuild_array_for_export_schema(array, ctx.execution_ctx())?; + let array = match ctx.cuda_session().dictionary_export() { + DictionaryExport::Preserve => { + rebuild_array_for_export_schema(array, ctx.execution_ctx())? + } + // Dictionary layouts no longer affect the schema. Preserve other encodings for + // structural recursion and direct FSST/OnPair export. + DictionaryExport::Decode => array, + }; let schema = arrow_schema_for_array(&array, ctx)?; let array = self.export_device_array(array, ctx).await?; Ok(ArrowDeviceArrayWithSchema { schema, array }) @@ -205,7 +213,10 @@ fn export_array( ) -> BoxFuture<'_, VortexResult<(ArrowArray, SyncEvent)>> { Box::pin(async { let array = match array.try_downcast::() { - Ok(dict) => return export_dict(dict, ctx).await, + Ok(dict) if ctx.cuda_session().dictionary_export() == DictionaryExport::Preserve => { + return export_dict(dict, ctx).await; + } + Ok(dict) => dict.into_array(), Err(array) => array, }; let array = match array.try_downcast::() { @@ -1495,7 +1506,9 @@ mod tests { use vortex::extension::datetime::TimeUnit; use crate::CudaBufferExt; + use crate::CudaDispatchMode; use crate::CudaExecutionCtx; + use crate::DictionaryExport; use crate::arrow::ARROW_DEVICE_CUDA; use crate::arrow::ArrowArray; use crate::arrow::ArrowDeviceArray; @@ -1504,8 +1517,10 @@ mod tests { use crate::arrow::arrow_schema_for_array; use crate::arrow::canonical::export_arrow_validity_buffer; use crate::arrow::canonical::repack_arrow_validity_buffer; + use crate::arrow::dictionary_tests::upload; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::device_buffer::cuda_backing_allocation; + use crate::executor::CudaArrayExt; use crate::session::CudaSession; use crate::session::VarBinExportLayout; @@ -2590,9 +2605,29 @@ mod tests { async fn test_export_fsst_varbin_contents( #[case] values: Vec>, #[case] dtype: DType, + #[values(DictionaryExport::Preserve, DictionaryExport::Decode)] policy: DictionaryExport, ) -> VortexResult<()> { - let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let session = + array_session().with_some(CudaSession::try_default()?.with_dictionary_export(policy)); + // Direct FSST varbin export must work when execute_cuda rejects standalone FSST, + // ruling out eager canonicalization. + let mut ctx = CudaSession::create_execution_ctx(&session)? + .with_dispatch_mode(CudaDispatchMode::DynDispatchOnly); let fsst = fsst_array_from(&values, dtype.clone(), &mut ctx)?; + // CUDA FSST needs a host symbol table. Upload codes and lengths to prevent CPU fallback. + let mut slots = Vec::new(); + for slot in fsst.slots().iter() { + slots.push(match slot { + Some(child) => Some(upload(child.clone(), &mut ctx).await?), + None => None, + }); + } + // SAFETY: Child values are unchanged; upload only changes buffer placement. + let fsst = unsafe { fsst.with_slots(slots.into()) }?; + if !fsst.is_empty() { + assert!(!fsst.is_host()); + assert!(fsst.clone().execute_cuda(&mut ctx).await.is_err()); + } let mut exported = fsst.export_device_array_with_schema(&mut ctx).await?; let expected_data_type = if matches!(dtype, DType::Utf8(_)) { diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs new file mode 100644 index 00000000000..a078f694448 --- /dev/null +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -0,0 +1,350 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use futures::future::BoxFuture; +use futures::stream; +use rstest::rstest; +use vortex::array::IntoArray; +use vortex::array::arrays::Constant; +use vortex::array::arrays::DictArray; +use vortex::array::arrays::ListViewArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::arrays::VarBinArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::assert_arrays_eq; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::array::validity::Validity; +use vortex::buffer::BitBuffer; +use vortex::buffer::Buffer; +use vortex::buffer::ByteBuffer; +use vortex::error::vortex_bail; + +use super::tests::last_error; +use super::*; +use crate::CudaSession; + +/// Preserve encodings while moving all buffers, including validity, to CUDA so unsupported +/// decoding errors instead of falling back to the CPU. +pub(super) fn upload( + array: ArrayRef, + ctx: &mut CudaExecutionCtx, +) -> BoxFuture<'_, VortexResult> { + Box::pin(async move { + // Constants store scalar metadata, not replaceable data buffers. + if array.as_opt::().is_some() { + return Ok(array); + } + let mut slots = Vec::new(); + for slot in array.slots().iter() { + slots.push(match slot { + Some(child) => Some(upload(child.clone(), ctx).await?), + None => None, + }); + } + let mut buffers = Vec::new(); + for buffer in array.buffer_handles() { + buffers.push(ctx.ensure_on_device(buffer).await?); + } + // SAFETY: Slots and buffers are byte-for-byte copies; only their placement changes. + unsafe { array.with_slots(slots.into())?.with_buffers(buffers) } + }) +} + +/// Copy a device buffer from a live, unreleased array produced by this exporter to the host. +fn buffer(array: &ArrowArray, index: usize) -> VortexResult { + // SAFETY: Only called on live arrays produced by our exporter, before their release. + let private = unsafe { &*array.private_data.cast::() }; + let buffer = private.buffers[index] + .as_ref() + .ok_or_else(|| vortex_err!("missing exported buffer {index}"))?; + buffer.cuda_device_ptr()?; + buffer.try_to_host_sync() +} + +/// Rebuild supported zero-offset, dictionary-free exports as host arrays for comparison. +/// Requires live arrays from this exporter and their matching logical dtype. +fn read_plain(array: &ArrowArray, dtype: &DType) -> VortexResult { + assert!(array.dictionary.is_null()); + assert_eq!(array.offset, 0); + let len = usize::try_from(array.length)?; + let validity = if !dtype.is_nullable() { + assert_eq!(array.null_count, 0); + Validity::NonNullable + } else if array.null_count == 0 { + Validity::AllValid + } else { + Validity::from(BitBuffer::new(buffer(array, 0)?, len)) + }; + match dtype { + DType::Primitive(ptype, _) => { + assert_eq!(array.n_buffers, 2); + assert_eq!(array.n_children, 0); + Ok(PrimitiveArray::from_byte_buffer(buffer(array, 1)?, *ptype, validity).into_array()) + } + DType::Utf8(_) => { + assert_eq!(array.n_buffers, 3); + assert_eq!(array.n_children, 0); + let offsets = PrimitiveArray::from_byte_buffer( + buffer(array, 1)?, + PType::I32, + Validity::NonNullable, + ) + .into_array(); + Ok(VarBinArray::try_new( + offsets, + buffer(array, 2)?.slice_unaligned(..), + dtype.clone(), + validity, + )? + .into_array()) + } + DType::Struct(fields, _) => { + assert_eq!(array.n_buffers, 1); + assert_eq!(usize::try_from(array.n_children)?, fields.nfields()); + let mut children = Vec::new(); + for (index, dtype) in fields.fields().enumerate() { + // SAFETY: The live struct owns exactly n_children child pointers. + let child = unsafe { &**array.children.add(index) }; + children.push(read_plain(child, &dtype)?); + } + Ok(StructArray::try_new(fields.names().clone(), children, len, validity)?.into_array()) + } + _ => vortex_bail!("unsupported test dtype {dtype}"), + } +} + +fn wrap_struct(array: ArrayRef) -> ArrayRef { + let len = array.len(); + let inner = StructArray::new(["value"].into(), vec![array], len, Validity::NonNullable); + StructArray::new( + ["nested"].into(), + vec![inner.into_array()], + len, + Validity::NonNullable, + ) + .into_array() +} + +fn values_and_expected(strings: bool) -> (ArrayRef, ArrayRef) { + if strings { + let long = "an out-of-line dictionary value"; + ( + VarBinViewArray::from_iter_nullable_str([Some("short"), None, Some(long)]).into_array(), + VarBinViewArray::from_iter_nullable_str([Some(long), None, Some("short"), Some(long)]) + .into_array(), + ) + } else { + ( + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + PrimitiveArray::from_iter([30i32, 20, 10, 30]).into_array(), + ) + } +} + +fn dictionary(values: ArrayRef, width: PType) -> VortexResult { + let codes = match width { + PType::U8 => PrimitiveArray::from_iter([2u8, 1, 0, 2]).into_array(), + PType::U16 => PrimitiveArray::from_iter([2u16, 1, 0, 2]).into_array(), + PType::U32 => PrimitiveArray::from_iter([2u32, 1, 0, 2]).into_array(), + _ => vortex_bail!("unsupported test index width {width}"), + }; + Ok(DictArray::try_new(codes, values)?.into_array()) +} + +fn get_schema(stream: &mut ArrowDeviceArrayStream) -> VortexResult { + let callback = stream.get_schema.expect("missing get_schema"); + let mut schema = FFI_ArrowSchema::empty(); + // SAFETY: The stream and output schema are live and writable. + let status = unsafe { callback(stream, (&raw mut schema).cast()) }; + assert_eq!(status, 0, "{}", last_error(stream)?); + Ok(schema) +} + +fn get_next(stream: &mut ArrowDeviceArrayStream) -> (i32, ArrowDeviceArray) { + let callback = stream.get_next.expect("missing get_next"); + let mut array = ArrowDeviceArray::empty(); + // SAFETY: The stream and output array are live and writable. + let status = unsafe { callback(stream, &raw mut array) }; + (status, array) +} + +/// Upload chunks and synchronize before handing them to a separate export context. +async fn upload_chunks( + chunks: Vec, + ctx: &mut CudaExecutionCtx, +) -> VortexResult>> { + let mut device_chunks = Vec::new(); + for chunk in chunks { + let chunk = upload(chunk, ctx).await?; + assert!(!chunk.is_host()); + device_chunks.push(Ok(chunk)); + } + ctx.synchronize_stream()?; + Ok(device_chunks) +} + +#[rstest] +#[case::plain_first_primitives(false, false, true)] +#[case::dictionary_first_nested_strings(true, true, false)] +#[crate::test] +fn test_decode_mixed_dictionary_device_stream( + #[case] strings: bool, + #[case] nested: bool, + #[case] plain_first: bool, +) -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_dictionary_export(DictionaryExport::Decode)); + let mut ctx = CudaSession::create_execution_ctx(&session)?; + let (values, expected) = values_and_expected(strings); + let nested_values = DictArray::try_new( + PrimitiveArray::from_iter([0u8, 1, 2]).into_array(), + values.clone(), + )? + .into_array(); + let mut chunks = vec![ + dictionary(values.clone(), PType::U8)?, + dictionary(values, PType::U16)?, + dictionary(nested_values, PType::U32)?, + expected.clone(), + ]; + if plain_first { + chunks.rotate_right(1); + } + let wrap = |array| if nested { wrap_struct(array) } else { array }; + let expected = wrap(expected); + let chunks = chunks.into_iter().map(wrap).collect(); + let chunks = runtime.block_on(upload_chunks(chunks, &mut ctx))?; + let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) + .boxed() + .export_device_array_stream(&session, &runtime)?; + let schema = get_schema(&mut stream)?; + let mut plain = runtime.block_on(expected.clone().export_device_array_with_schema(&mut ctx))?; + assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain.schema)?); + release_device_array(&mut plain.array); + for _ in 0..4 { + let (status, mut array) = get_next(&mut stream); + assert_eq!(status, 0, "{}", last_error(&mut stream)?); + assert_eq!(array.device_type, ARROW_DEVICE_CUDA); + let actual = read_plain(&array.array, expected.dtype())?; + assert_arrays_eq!(actual, expected, ctx.execution_ctx()); + release_device_array(&mut array); + } + let (status, eos) = get_next(&mut stream); + assert_eq!(status, 0); + assert!(eos.array.release.is_none()); + // SAFETY: This is the live stream's final use. + unsafe { stream.release.expect("missing release")(&raw mut stream) }; + Ok(()) +} + +#[crate::test] +async fn test_decode_non_contiguous_dictionary_list_view() -> VortexResult<()> { + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_dictionary_export(DictionaryExport::Decode)); + let mut ctx = CudaSession::create_execution_ctx(&session)?; + let (values, expected) = values_and_expected(true); + let array = ListViewArray::new( + dictionary(values, PType::U8)?, + PrimitiveArray::from_iter([2i32, 0]).into_array(), + PrimitiveArray::from_iter([2i32, 2]).into_array(), + Validity::NonNullable, + ) + .into_array(); + let expected = expected.take(PrimitiveArray::from_iter([2u32, 3, 0, 1]).into_array())?; + let array = upload(array, &mut ctx).await?; + let mut exported = array.export_device_array_with_schema(&mut ctx).await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new_list( + "", + Field::new(Field::LIST_FIELD_DEFAULT_NAME, DataType::Utf8, true), + false, + ) + ); + assert_eq!(exported.array.array.length, 2); + assert_eq!(exported.array.array.n_children, 1); + assert_eq!( + Buffer::::from_byte_buffer(buffer(&exported.array.array, 1)?).as_ref(), + &[0, 2, 4] + ); + // SAFETY: This live list array owns the single child checked above. + let values = unsafe { &**exported.array.array.children }; + let actual = read_plain(values, expected.dtype())?; + assert_arrays_eq!(actual, expected, ctx.execution_ctx()); + release_device_array(&mut exported.array); + Ok(()) +} + +#[crate::test] +async fn test_decode_unsupported_device_dictionary_does_not_fall_back_to_cpu() -> VortexResult<()> { + let cuda = CudaSession::try_default()?; + let session = vortex::array::array_session().with_some(cuda.clone()); + let mut ctx = CudaSession::create_execution_ctx(&session)?; + // A dictionary of structs can be preserved, but has no CUDA gather kernel today. + let (values, _) = values_and_expected(false); + let array = upload(dictionary(wrap_struct(values), PType::U8)?, &mut ctx).await?; + assert!(!array.is_host()); + let mut preserved = array.clone().export_device_array(&mut ctx).await?; + assert!(!preserved.array.dictionary.is_null()); + release_device_array(&mut preserved); + + let session = vortex::array::array_session() + .with_some(cuda.with_dictionary_export(DictionaryExport::Decode)); + let mut ctx = CudaSession::create_execution_ctx(&session)?; + let error = match array.export_device_array_with_schema(&mut ctx).await { + Ok(mut exported) => { + release_device_array(&mut exported.array); + vortex_bail!("unsupported device dictionary unexpectedly decoded"); + } + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("CPU fallback with device-resident buffers is not supported") + ); + Ok(()) +} + +#[rstest] +#[case::different_dictionary_width(Some(PType::U16))] +#[case::plain_chunk(None)] +#[crate::test] +fn test_default_dictionary_device_stream(#[case] second_width: Option) -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = crate::cuda_session(); + let mut ctx = CudaSession::create_execution_ctx(&session)?; + assert_eq!( + ctx.cuda_session().dictionary_export(), + DictionaryExport::Preserve + ); + let (values, expected) = values_and_expected(false); + let first = dictionary(values.clone(), PType::U8)?; + let second = match second_width { + Some(width) => dictionary(values, width)?, + None => expected.clone(), + }; + let chunks = runtime.block_on(upload_chunks(vec![first, second], &mut ctx))?; + let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) + .boxed() + .export_device_array_stream(&session, &runtime)?; + let schema = get_schema(&mut stream)?; + assert_eq!( + Field::try_from(&schema)?.data_type(), + &DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Int32),) + ); + let (status, mut array) = get_next(&mut stream); + assert_eq!(status, 0); + assert!(!array.array.dictionary.is_null()); + release_device_array(&mut array); + let (status, rejected) = get_next(&mut stream); + assert_eq!(status, LIBC_EIO); + assert!(last_error(&mut stream)?.contains("Arrow schema changed")); + assert!(rejected.array.release.is_none()); + // SAFETY: This is the live stream's final use. + unsafe { stream.release.expect("missing release")(&raw mut stream) }; + Ok(()) +} diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index f606e623d4b..8d9e053aebd 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -9,6 +9,8 @@ //! More documentation at mod canonical; +#[cfg(test)] +mod dictionary_tests; mod list_view; mod offsets; @@ -62,6 +64,7 @@ use vortex_arrow::ArrowSessionExt; use crate::CudaBufferExt; use crate::CudaExecutionCtx; +use crate::DictionaryExport; use crate::VarBinExportLayout; mod arrow_c_abi { @@ -482,6 +485,8 @@ pub trait DeviceArrayStreamExt { /// The Arrow Device stream contract requires all arrays to share the schema reported by /// `get_schema`. The schema is derived from the first array, or from the logical dtype /// for an empty stream. Chunks that export to different Arrow types are rejected mid-stream. + /// Set [`DictionaryExport::Decode`] on the [`crate::CudaSession`] to export logical plain + /// types even when chunks vary between dictionary/plain encodings or dictionary index widths. /// /// Drive the returned stream from one thread. `runtime` must be the runtime that owns the /// underlying scan tasks and per-array exports. @@ -499,38 +504,42 @@ impl DeviceArrayStreamExt for SendableArrayStream { session: &VortexSession, runtime: &CurrentThreadRuntime, ) -> VortexResult { - let dtype = self.dtype().clone(); let ctx = crate::CudaSession::create_execution_ctx(session)?; - let array_iter = Box::new(runtime.block_on_stream(self)); - Ok(device_array_stream(array_iter, dtype, ctx, runtime.clone())) + Ok(ArrowDeviceArrayStream::new(self, ctx, runtime)) } } -/// Build the Arrow Device stream that owns `array_iter` and exports its arrays through `ctx`. -fn device_array_stream( - array_iter: ArrayStreamIterator, - dtype: DType, - ctx: CudaExecutionCtx, - runtime: CurrentThreadRuntime, -) -> ArrowDeviceArrayStream { - let private_data = Box::new(DeviceArrayStreamPrivateData { - device_id: ctx.stream().context().ordinal() as i64, - array_iter, - ctx, - runtime, - dtype, - schema: None, - pending_array: None, - last_error: None, - }); - - ArrowDeviceArrayStream { - device_type: ARROW_DEVICE_CUDA, - get_schema: Some(device_stream_get_schema), - get_next: Some(device_stream_get_next), - get_last_error: Some(device_stream_get_last_error), - release: Some(device_stream_release), - private_data: Box::into_raw(private_data).cast(), +impl ArrowDeviceArrayStream { + /// Export a stream using an owned context, retaining its session and per-context configuration. + /// + /// The schema, runtime, and release requirements of + /// [`DeviceArrayStreamExt::export_device_array_stream`] also apply here. + pub fn new( + array_stream: SendableArrayStream, + ctx: CudaExecutionCtx, + runtime: &CurrentThreadRuntime, + ) -> Self { + let dtype = array_stream.dtype().clone(); + let array_iter = Box::new(runtime.block_on_stream(array_stream)); + let private_data = Box::new(DeviceArrayStreamPrivateData { + device_id: ctx.stream().context().ordinal() as i64, + array_iter, + ctx, + runtime: runtime.clone(), + dtype, + schema: None, + pending_array: None, + last_error: None, + }); + + Self { + device_type: ARROW_DEVICE_CUDA, + get_schema: Some(device_stream_get_schema), + get_next: Some(device_stream_get_next), + get_last_error: Some(device_stream_get_last_error), + release: Some(device_stream_release), + private_data: Box::into_raw(private_data).cast(), + } } } @@ -690,6 +699,10 @@ pub(crate) fn arrow_schema_for_array( array: &ArrayRef, ctx: &mut CudaExecutionCtx, ) -> VortexResult { + if ctx.cuda_session().dictionary_export() == DictionaryExport::Decode { + return ArrowDeviceStreamSchema::from_dtype(array.dtype(), ctx)?.to_ffi(); + } + if let Some(struct_array) = array.as_opt::() { return Ok(FFI_ArrowSchema::try_from(Schema::new( arrow_device_export_struct_fields_for_array( @@ -772,31 +785,21 @@ fn arrow_device_export_field_for_array( )); } + let mut list_field = |elements: &ArrayRef| { + let element = + arrow_device_export_field_for_array(Field::LIST_FIELD_DEFAULT_NAME, elements, ctx)?; + Ok(Field::new_list(name, element, array.dtype().is_nullable())) + }; if let Some(list) = array.as_opt::() { - let element = arrow_device_export_field_for_array( - Field::LIST_FIELD_DEFAULT_NAME, - list.elements(), - ctx, - )?; - return Ok(Field::new_list(name, element, array.dtype().is_nullable())); + return list_field(list.elements()); } if let Some(list) = array.as_opt::() { - let element = arrow_device_export_field_for_array( - Field::LIST_FIELD_DEFAULT_NAME, - list.elements(), - ctx, - )?; - return Ok(Field::new_list(name, element, array.dtype().is_nullable())); + return list_field(list.elements()); } if let Some(list) = array.as_opt::() { - let element = arrow_device_export_field_for_array( - Field::LIST_FIELD_DEFAULT_NAME, - list.elements(), - ctx, - )?; - return Ok(Field::new_list(name, element, array.dtype().is_nullable())); + return list_field(list.elements()); } arrow_device_export_field(name, &arrow_device_export_dtype(array.dtype()), ctx) @@ -962,7 +965,7 @@ mod tests { use crate::arrow::release_device_array; use crate::arrow::release_schema; - fn last_error(stream: &mut ArrowDeviceArrayStream) -> VortexResult { + pub(super) fn last_error(stream: &mut ArrowDeviceArrayStream) -> VortexResult { let get_last_error = stream .get_last_error .ok_or_else(|| vortex_err!("stream missing get_last_error callback"))?; diff --git a/vortex-cuda/src/executor.rs b/vortex-cuda/src/executor.rs index 1c56bdab090..0d28b79c292 100644 --- a/vortex-cuda/src/executor.rs +++ b/vortex-cuda/src/executor.rs @@ -40,6 +40,7 @@ use vortex::error::vortex_ensure; use vortex::error::vortex_err; use crate::CudaSession; +use crate::DictionaryExport; use crate::ExportDeviceArray; use crate::hybrid_dispatch; use crate::kernel::DefaultLaunchStrategy; @@ -135,6 +136,12 @@ impl CudaExecutionCtx { self } + /// Override the dictionary export policy for this context without changing its backing session. + pub fn with_dictionary_export(mut self, policy: DictionaryExport) -> Self { + self.cuda_session = self.cuda_session.with_dictionary_export(policy); + self + } + /// Perform an external kernel launch, with events created and logged via the configured /// [`LaunchStrategy`]. /// diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 216b3369bd1..ee96f690524 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -13,6 +13,7 @@ use async_trait::async_trait; use futures::FutureExt; use futures::StreamExt; use futures::future::BoxFuture; +use parking_lot::Mutex; use vortex::array::ArrayRef; use vortex::array::ArrayVTable; use vortex::array::MaskFuture; @@ -30,6 +31,12 @@ use vortex::buffer::BufferString; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::FieldMask; +use vortex::editions::Edition; +use vortex::editions::EditionDeclaration; +use vortex::editions::EditionFamily; +use vortex::editions::EditionId; +use vortex::editions::EditionMember; +use vortex::editions::EditionSessionExt; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -55,6 +62,7 @@ use vortex::layout::segments::SegmentSinkRef; use vortex::layout::segments::SegmentSource; use vortex::layout::sequence::SendableSequentialStream; use vortex::layout::sequence::SequencePointer; +use vortex::layout::session::LayoutSessionExt; use vortex::mask::Mask; use vortex::scalar::Scalar; use vortex::scalar::ScalarTruncation; @@ -534,12 +542,135 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { result } -/// Register the [`CudaFlatLayoutEncoding`] in the session's layout registry. +const CUDA_EDITION: EditionId = EditionId::new("cuda", 2026, 9, 0); +static CUDA_EDITION_DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: CUDA_EDITION, + min_library_version: None, + }, + added: &[EditionMember::layout(&"vortex.cuda_flat")], +}; + +/// Register the [`CudaFlatLayoutEncoding`] and enable its draft `cuda` edition for writing. +/// +/// Other edition selections and checks are unchanged. The draft has no cross-version +/// compatibility guarantee; readers must register the CUDA layout. /// /// Call this alongside [`crate::initialize_cuda`] when setting up a CUDA-enabled session. +/// Registration itself does not require a GPU. pub fn register_cuda_layout(session: &VortexSession) { - use vortex::layout::session::LayoutSessionExt; session .layouts() .register(LayoutEncodingRef::new_ref(&CudaFlat)); + + // Concurrent CUDA FFI calls may register session clones; serialize the check and registration + // because edition declarations reject duplicates. + static REGISTRATION_LOCK: Mutex<()> = Mutex::new(()); + let _guard = REGISTRATION_LOCK.lock(); + if session.editions().find(&CUDA_EDITION).is_none() { + session + .editions() + .declare_family(&EditionFamily { + name: "cuda", + origin: "vortex-cuda", + doc: "CUDA-readable layouts, enabled only when CUDA layout support is registered.", + }) + .vortex_expect("CUDA edition family is valid"); + session + .register_edition(&CUDA_EDITION_DECLARATION) + .vortex_expect("CUDA edition declaration is valid"); + } + session + .enable_edition(CUDA_EDITION) + .vortex_expect("CUDA edition is registered"); +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex::VortexSessionDefault; + use vortex::array::IntoArray; + use vortex::buffer::ByteBufferMut; + use vortex::buffer::buffer; + use vortex::editions::CORE_2025_05_0; + use vortex::editions::ComponentKind; + use vortex::editions::DEFAULT_CORE_EDITION; + use vortex::file::WriteOptionsSessionExt; + use vortex::io::runtime::BlockingRuntime; + use vortex::io::runtime::current::CurrentThreadRuntime; + use vortex::io::session::RuntimeSessionExt; + + use super::*; + + #[rstest] + fn test_cuda_registration_preserves_edition_policy( + #[values(DEFAULT_CORE_EDITION, CORE_2025_05_0)] core: EditionId, + ) -> VortexResult<()> { + let session = VortexSession::default(); + session.enable_edition(core)?; + let kinds = [ + ComponentKind::Array, + ComponentKind::Layout, + ComponentKind::DType, + ComponentKind::Aggregate, + ]; + let expected_ids = kinds.map(|kind| { + let mut ids = session.enabled_component_ids(kind); + if kind == ComponentKind::Layout { + assert!(!ids.contains(&CudaFlat.id())); + ids.push(CudaFlat.id()); + ids.sort_unstable(); + } + ids + }); + + std::thread::scope(|scope| { + for _ in 0..4 { + let session = session.clone(); + scope.spawn(move || register_cuda_layout(&session)); + } + }); + register_cuda_layout(&session); + + for (kind, expected) in kinds.into_iter().zip(expected_ids) { + assert_eq!(session.enabled_component_ids(kind), expected); + } + let mut enabled_editions = session.enabled_editions().editions(); + enabled_editions.sort_unstable(); + assert_eq!(enabled_editions, [core, CUDA_EDITION]); + session.editions().validate()?; + assert!( + !VortexSession::default() + .enabled_component_ids(ComponentKind::Layout) + .contains(&CudaFlat.id()) + ); + Ok(()) + } + + #[test] + fn test_registry_alone_does_not_permit_cuda_flat() -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with_handle(runtime.handle()); + session + .layouts() + .register(LayoutEncodingRef::new_ref(&CudaFlat)); + runtime.block_on(async { + let array = buffer![1i32, 4, 9, 16].into_array(); + let mut buffer = ByteBufferMut::empty(); + let error = session + .write_options() + .with_strategy(Arc::new(CudaFlatLayoutStrategy::default())) + .write(&mut buffer, array.to_array_stream()) + .await + .err() + .expect("write permitted an uneditioned CUDA layout"); + assert!( + error + .to_string() + .contains("Layout encoding vortex.cuda_flat not permitted by ctx"), + "unexpected error: {error}" + ); + Ok(()) + }) + } } diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 2447dfdbad5..48aa39bef82 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -69,6 +69,7 @@ pub use pooled_read_at::PooledFileReadAtOptions; pub use pooled_read_at::PooledObjectStoreReadAt; pub use session::CudaSession; pub use session::CudaSessionExt; +pub use session::DictionaryExport; pub use session::VarBinExportLayout; pub use stream::VortexCudaStream; pub use stream_pool::VortexCudaStreamPool; diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index 0924854405f..9b377384952 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -40,6 +40,18 @@ pub enum VarBinExportLayout { VarBinView, } +/// Controls whether Arrow Device exports preserve dictionaries or expand them to plain values, +/// including nested children. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum DictionaryExport { + /// Preserve dictionary values and indices in the Arrow schema and device array. + #[default] + Preserve, + /// Decode dictionaries on CUDA to keep one plain Arrow schema across batches. + /// May increase device memory use. Device-resident inputs require CUDA decoding support. + Decode, +} + /// CUDA session for GPU accelerated execution. /// /// Maintains a registry of CUDA kernel implementations for array encodings. @@ -50,6 +62,7 @@ pub struct CudaSession { kernels: Arc>, export_device_array: Arc, varbin_export_layout: VarBinExportLayout, + dictionary_export: DictionaryExport, kernel_loader: Arc, stream_pool: Arc, pinned_buffer_pool: Arc, @@ -77,6 +90,7 @@ impl CudaSession { kernel_loader: Arc::new(KernelLoader::new()), export_device_array: Arc::new(CanonicalDeviceArrayExport), varbin_export_layout: VarBinExportLayout::default(), + dictionary_export: DictionaryExport::default(), stream_pool, pinned_buffer_pool, } @@ -93,6 +107,17 @@ impl CudaSession { self.varbin_export_layout } + /// Selects whether Arrow Device exports preserve or decode dictionaries. + pub fn with_dictionary_export(mut self, policy: DictionaryExport) -> Self { + self.dictionary_export = policy; + self + } + + /// Returns the dictionary policy used for Arrow Device exports. + pub fn dictionary_export(&self) -> DictionaryExport { + self.dictionary_export + } + /// Creates a default CUDA session using device 0, with all GPU array kernels preloaded. /// /// Unlike [`Default::default`], this returns an error instead of panicking when CUDA cannot be From 11c9af9813ecf1a96e4f1fca2b33c21d594ca249 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 14:55:50 +0000 Subject: [PATCH 02/35] refactor(cuda): simplify scan setup and export tests Signed-off-by: Alexander Droste --- vortex-cuda/ffi/build.rs | 15 +++------ vortex-cuda/ffi/src/lib.rs | 2 +- vortex-cuda/ffi/src/tests/projection.rs | 41 +++++++---------------- vortex-cuda/src/arrow/canonical.rs | 5 ++- vortex-cuda/src/arrow/dictionary_tests.rs | 12 +++---- 5 files changed, 24 insertions(+), 51 deletions(-) diff --git a/vortex-cuda/ffi/build.rs b/vortex-cuda/ffi/build.rs index 10218d19e94..ab74d21cfca 100644 --- a/vortex-cuda/ffi/build.rs +++ b/vortex-cuda/ffi/build.rs @@ -1,9 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::env; use std::error::Error; -use std::path::PathBuf; use std::process::Command; fn main() -> Result<(), Box> { @@ -11,20 +9,17 @@ fn main() -> Result<(), Box> { println!("cargo:rerun-if-changed=cbindgen.toml"); println!("cargo:rerun-if-changed=build.rs"); - let crate_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR")?); - let header = crate_dir.join("cinclude/vortex_cuda.h"); + let header = "cinclude/vortex_cuda.h"; // The CUDA API needs no macro expansion or dependency parsing, so generate on stable Rust // without recursively building the CUDA implementation. cbindgen::Builder::new() - .with_src(crate_dir.join("src/lib.rs")) - .with_config(cbindgen::Config::from_file( - crate_dir.join("cbindgen.toml"), - )?) + .with_src("src/lib.rs") + .with_config(cbindgen::Config::from_file("cbindgen.toml")?) .generate()? - .write_to_file(&header); + .write_to_file(header); if !Command::new("clang-format") .args(["--style=file", "-i"]) - .arg(&header) + .arg(header) .status() .is_ok_and(|status| status.success()) { diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 050eadfa0c6..f99885b0911 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -387,7 +387,7 @@ fn projected_scan( "unknown CUDA scan column: {name:?}" ); } - let projection = select(columns, root()).optimize_recursive(file.dtype())?; + let projection = select(columns, root()).optimize(file.dtype())?; scan = scan.with_projection(projection.bind(file.dtype())?); } if batch_rows != 0 { diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index bbdd8213a03..52d09454517 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -10,7 +10,7 @@ use std::sync::atomic::Ordering; use futures::TryStreamExt; use vortex::array::VortexSessionExecute; -use vortex::arrow::ArrowSessionExt; +use vortex::array::assert_arrays_eq; use vortex::buffer::ByteBuffer; use vortex::buffer::ByteBufferMut; use vortex::file::WriteOptionsSessionExt; @@ -95,8 +95,8 @@ fn session() -> VortexSession { VortexSession::default().with_handle(ffi_runtime().handle()) } -fn table() -> VortexResult { - Ok(StructArray::try_new( +fn table() -> VortexResult { + StructArray::try_new( ["ids", "unused", "値.x"].into(), vec![ PrimitiveArray::from_iter(0u32..5).into_array(), @@ -106,8 +106,7 @@ fn table() -> VortexResult { ], 5, Validity::NonNullable, - )? - .into_array()) + ) } fn file_bytes( @@ -142,22 +141,6 @@ fn open_file( .open_buffer(file_bytes(session, array, cuda_block_rows)?) } -fn assert_arrow_eq( - session: &VortexSession, - actual: ArrayRef, - expected: ArrayRef, -) -> VortexResult<()> { - let mut ctx = session.create_execution_ctx(); - let mut to_data = |array| { - session - .arrow() - .execute_arrow(array, None, &mut ctx) - .map(|array| array.to_data()) - }; - assert_eq!(to_data(actual)?, to_data(expected)?); - Ok(()) -} - #[test] fn test_cuda_write_strategy_preserves_high_cardinality_row_blocks() -> VortexResult<()> { let session = session(); @@ -200,12 +183,8 @@ fn test_projection_cpu_never_requests_unselected_column_segments() -> VortexResu let session = session(); let input = table()?; let columns = names(&["値.x", "ids"])?; - let expected = input - .clone() - .execute::(&mut session.create_execution_ctx())? - .project(columns.as_ref())? - .into_array(); - let file = open_file(&session, input, None)?; + let expected = input.project(columns.as_ref())?.into_array(); + let file = open_file(&session, input.into_array(), None)?; // TableStrategy writes one flat child per column, so child 1 is exactly the unused column. let children = file.footer().layout().children()?; let forbidden = children[1].segment_ids(); @@ -221,7 +200,7 @@ fn test_projection_cpu_never_requests_unselected_column_segments() -> VortexResu .into_array_stream()? .read_all(), )?; - assert_arrow_eq(&session, actual, expected)?; + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); assert_eq!(source.rejected.load(Ordering::Relaxed), 0); // The same reader must fail without projection, proving the guard actually observes reads. assert_error( @@ -389,7 +368,11 @@ fn batch_lengths(stream: &mut ArrowDeviceArrayStream) -> Vec { fn test_projection_gpu_local_file_schema_and_batch_boundaries() -> VortexResult<()> { for (block_rows, batch_rows) in [(0, 2), (2, 3)] { let session = session().with_some(CudaSession::try_default()?); - let file = LocalFile::new(&file_bytes(&session, table()?, Some(block_rows))?)?; + let file = LocalFile::new(&file_bytes( + &session, + table()?.into_array(), + Some(block_rows), + )?)?; let path = file .0 .to_str() diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index bed7d072b27..f2948fb223f 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -2607,11 +2607,10 @@ mod tests { #[case] dtype: DType, #[values(DictionaryExport::Preserve, DictionaryExport::Decode)] policy: DictionaryExport, ) -> VortexResult<()> { - let session = - array_session().with_some(CudaSession::try_default()?.with_dictionary_export(policy)); // Direct FSST varbin export must work when execute_cuda rejects standalone FSST, // ruling out eager canonicalization. - let mut ctx = CudaSession::create_execution_ctx(&session)? + let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)? + .with_dictionary_export(policy) .with_dispatch_mode(CudaDispatchMode::DynDispatchOnly); let fsst = fsst_array_from(&values, dtype.clone(), &mut ctx)?; // CUDA FSST needs a host symbol table. Upload codes and lengths to prevent CPU fallback. diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index a078f694448..e2faf24dc83 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -221,9 +221,8 @@ fn test_decode_mixed_dictionary_device_stream( .boxed() .export_device_array_stream(&session, &runtime)?; let schema = get_schema(&mut stream)?; - let mut plain = runtime.block_on(expected.clone().export_device_array_with_schema(&mut ctx))?; - assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain.schema)?); - release_device_array(&mut plain.array); + let plain_schema = arrow_schema_for_array(&expected, &mut ctx)?; + assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain_schema)?); for _ in 0..4 { let (status, mut array) = get_next(&mut stream); assert_eq!(status, 0, "{}", last_error(&mut stream)?); @@ -280,8 +279,7 @@ async fn test_decode_non_contiguous_dictionary_list_view() -> VortexResult<()> { #[crate::test] async fn test_decode_unsupported_device_dictionary_does_not_fall_back_to_cpu() -> VortexResult<()> { - let cuda = CudaSession::try_default()?; - let session = vortex::array::array_session().with_some(cuda.clone()); + let session = vortex::array::array_session().with_some(CudaSession::try_default()?); let mut ctx = CudaSession::create_execution_ctx(&session)?; // A dictionary of structs can be preserved, but has no CUDA gather kernel today. let (values, _) = values_and_expected(false); @@ -291,9 +289,7 @@ async fn test_decode_unsupported_device_dictionary_does_not_fall_back_to_cpu() - assert!(!preserved.array.dictionary.is_null()); release_device_array(&mut preserved); - let session = vortex::array::array_session() - .with_some(cuda.with_dictionary_export(DictionaryExport::Decode)); - let mut ctx = CudaSession::create_execution_ctx(&session)?; + let mut ctx = ctx.with_dictionary_export(DictionaryExport::Decode); let error = match array.export_device_array_with_schema(&mut ctx).await { Ok(mut exported) => { release_device_array(&mut exported.array); From 2b3d8df8c12e7075fc84309d159f28bc4cf79a75 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 15:00:36 +0000 Subject: [PATCH 03/35] perf(cuda): reuse decoded stream schemas across batches Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/dictionary_tests.rs | 48 ++++++++++++++++++++++- vortex-cuda/src/arrow/mod.rs | 40 ++++++++++++++++--- 2 files changed, 80 insertions(+), 8 deletions(-) diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index e2faf24dc83..cf44202b75a 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -193,6 +193,7 @@ fn test_decode_mixed_dictionary_device_stream( #[case] strings: bool, #[case] nested: bool, #[case] plain_first: bool, + #[values(false, true)] schema_first: bool, ) -> VortexResult<()> { let runtime = CurrentThreadRuntime::new(); let session = vortex::array::array_session() @@ -220,9 +221,11 @@ fn test_decode_mixed_dictionary_device_stream( let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) .boxed() .export_device_array_stream(&session, &runtime)?; - let schema = get_schema(&mut stream)?; let plain_schema = arrow_schema_for_array(&expected, &mut ctx)?; - assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain_schema)?); + if schema_first { + let schema = get_schema(&mut stream)?; + assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain_schema)?); + } for _ in 0..4 { let (status, mut array) = get_next(&mut stream); assert_eq!(status, 0, "{}", last_error(&mut stream)?); @@ -231,6 +234,8 @@ fn test_decode_mixed_dictionary_device_stream( assert_arrays_eq!(actual, expected, ctx.execution_ctx()); release_device_array(&mut array); } + let schema = get_schema(&mut stream)?; + assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain_schema)?); let (status, eos) = get_next(&mut stream); assert_eq!(status, 0); assert!(eos.array.release.is_none()); @@ -239,6 +244,45 @@ fn test_decode_mixed_dictionary_device_stream( Ok(()) } +#[crate::test] +fn test_decode_stream_validates_dtype_and_device() -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_dictionary_export(DictionaryExport::Decode)); + let array = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); + let mut stream = array + .clone() + .to_array_stream() + .boxed() + .export_device_array_stream(&session, &runtime)?; + // SAFETY: The stream is live and exclusively borrowed until the state is no longer used. + let state = unsafe { device_stream_private_data(&raw mut stream) }.expect("missing state"); + let mut first = state.export_stream_array(array.clone())?; + release_device_array(&mut first); + assert!(state.schema.is_some()); + + let error = state + .export_stream_array(PrimitiveArray::from_iter([10u32, 20, 30]).into_array()) + .err() + .expect("accepted a different dtype"); + assert!(error.to_string().contains("stream array dtype changed")); + let error = state.check_device(&ArrowDeviceArray::empty()).unwrap_err(); + assert!(error.to_string().contains("non-CUDA device type")); + state.device_id = -1; + let error = state + .export_stream_array(array) + .err() + .expect("accepted a different device"); + assert!( + error + .to_string() + .contains("stream array moved from CUDA device") + ); + // SAFETY: This is the live stream's final use; no state borrow remains. + unsafe { stream.release.expect("missing release")(&raw mut stream) }; + Ok(()) +} + #[crate::test] async fn test_decode_non_contiguous_dictionary_list_view() -> VortexResult<()> { let session = vortex::array::array_session() diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 8d9e053aebd..7e22c734083 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -407,6 +407,30 @@ impl DeviceArrayStreamPrivateData { array.dtype() ); + if self.ctx.cuda_session().dictionary_export() == DictionaryExport::Decode { + // The canonical exporter uses only the dtype and fixed context settings in this mode. + // Avoid constructing and parsing a temporary C schema for every batch. + let schema = if self.schema.is_none() { + Some(ArrowDeviceStreamSchema::from_dtype( + &self.dtype, + &mut self.ctx, + )?) + } else { + None + }; + let mut device_array = self + .runtime + .block_on(array.export_device_array(&mut self.ctx))?; + if let Err(error) = self.check_device(&device_array) { + release_device_array(&mut device_array); + return Err(error); + } + if let Some(schema) = schema { + self.schema = Some(schema); + } + return Ok(device_array); + } + let ArrowDeviceArrayWithSchema { schema: mut ffi_schema, array: mut device_array, @@ -431,12 +455,7 @@ impl DeviceArrayStreamPrivateData { Ok(device_array) } - /// Check that a freshly exported device array matches the stream schema and CUDA device. - fn check_stream_array( - &self, - ffi_schema: &FFI_ArrowSchema, - device_array: &ArrowDeviceArray, - ) -> VortexResult { + fn check_device(&self, device_array: &ArrowDeviceArray) -> VortexResult<()> { vortex_ensure!( device_array.device_type == ARROW_DEVICE_CUDA, "stream array exported on non-CUDA device type {}", @@ -448,7 +467,16 @@ impl DeviceArrayStreamPrivateData { self.device_id, device_array.device_id ); + Ok(()) + } + /// Check that a freshly exported device array matches the stream schema and CUDA device. + fn check_stream_array( + &self, + ffi_schema: &FFI_ArrowSchema, + device_array: &ArrowDeviceArray, + ) -> VortexResult { + self.check_device(device_array)?; let exported_schema = ArrowDeviceStreamSchema::from_ffi(ffi_schema, &self.dtype)?; if let Some(stream_schema) = &self.schema { vortex_ensure!( From 1f9a40be8fc9dba2ba74969be6e412bee9b9eaf6 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 15:02:52 +0000 Subject: [PATCH 04/35] perf(cuda): derive decoded stream schemas without polling Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/dictionary_tests.rs | 65 +++++++++++++++++++++++ vortex-cuda/src/arrow/mod.rs | 22 ++++---- 2 files changed, 78 insertions(+), 9 deletions(-) diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index cf44202b75a..99d41dea160 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; + use futures::future::BoxFuture; use futures::stream; use rstest::rstest; @@ -244,6 +248,67 @@ fn test_decode_mixed_dictionary_device_stream( Ok(()) } +#[rstest] +#[case::values(true, false)] +#[case::error(true, true)] +#[case::empty(false, false)] +#[crate::test] +fn test_decode_stream_schema_does_not_poll( + #[case] has_batch: bool, + #[case] fails: bool, + #[values(false, true)] consume: bool, +) -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_dictionary_export(DictionaryExport::Decode)); + let array = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); + let dtype = array.dtype().clone(); + let mut batch = has_batch.then(|| { + if fails { + Err(vortex_err!("deferred scan error")) + } else { + Ok(array) + } + }); + let polls = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&polls); + let input = stream::poll_fn(move |_| { + counter.fetch_add(1, Ordering::Relaxed); + Poll::Ready(batch.take()) + }); + let mut stream = ArrayStreamAdapter::new(dtype, input) + .boxed() + .export_device_array_stream(&session, &runtime)?; + for _ in 0..2 { + let schema = get_schema(&mut stream)?; + assert_eq!( + Field::try_from(&schema)?, + Field::new("", DataType::Int32, false) + ); + } + assert_eq!(polls.load(Ordering::Relaxed), 0); + if consume { + let (status, mut array) = get_next(&mut stream); + assert_eq!(polls.load(Ordering::Relaxed), 1); + if fails { + assert_eq!(status, LIBC_EIO); + assert!(last_error(&mut stream)?.contains("deferred scan error")); + assert!(array.array.release.is_none()); + } else { + assert_eq!(status, 0, "{}", last_error(&mut stream)?); + assert_eq!(array.array.release.is_some(), has_batch); + if has_batch { + assert_eq!(array.array.length, 3); + } + release_device_array(&mut array); + } + } + // SAFETY: This is the live stream's final use, including schema-only consumers. + unsafe { stream.release.expect("missing release")(&raw mut stream) }; + assert_eq!(polls.load(Ordering::Relaxed), usize::from(consume)); + Ok(()) +} + #[crate::test] fn test_decode_stream_validates_dtype_and_device() -> VortexResult<()> { let runtime = CurrentThreadRuntime::new(); diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 7e22c734083..ed031d52f1e 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -365,12 +365,15 @@ impl DeviceArrayStreamPrivateData { code } - /// Return the stream schema, exporting the first stream array to derive it if needed. - /// - /// A first array is held in `pending_array` so the following `get_next` returns it. + /// Derive decoded schemas from the dtype without pulling a batch. Preserved dictionaries + /// require the first array, held in `pending_array` for the following `get_next`. fn get_or_init_schema(&mut self) -> VortexResult<&ArrowDeviceStreamSchema> { if self.schema.is_none() { - match self.array_iter.next() { + let first = match self.ctx.cuda_session().dictionary_export() { + DictionaryExport::Preserve => self.array_iter.next(), + DictionaryExport::Decode => None, + }; + match first { Some(array) => self.pending_array = Some(self.export_stream_array(array?)?), None => { self.schema = Some(ArrowDeviceStreamSchema::from_dtype( @@ -511,10 +514,11 @@ pub trait DeviceArrayStreamExt { /// embedded `release` callback. /// /// The Arrow Device stream contract requires all arrays to share the schema reported by - /// `get_schema`. The schema is derived from the first array, or from the logical dtype - /// for an empty stream. Chunks that export to different Arrow types are rejected mid-stream. - /// Set [`DictionaryExport::Decode`] on the [`crate::CudaSession`] to export logical plain - /// types even when chunks vary between dictionary/plain encodings or dictionary index widths. + /// `get_schema`. By default, the schema is derived from the first array, or from the logical + /// dtype for an empty stream. Chunks exporting different Arrow types are rejected mid-stream. + /// With [`DictionaryExport::Decode`], the logical dtype determines a stable plain schema even + /// when chunks vary between dictionary/plain encodings or dictionary index widths. In this + /// mode, `get_schema` does not pull or export a batch; read/decode errors surface in `get_next`. /// /// Drive the returned stream from one thread. `runtime` must be the runtime that owns the /// underlying scan tasks and per-array exports. @@ -623,7 +627,7 @@ fn device_stream_callback( } } -/// Write the stream's Arrow schema, initializing it from the first stream array if unset. +/// Write the stream's Arrow schema, deriving it from the dtype or first array as needed. unsafe extern "C" fn device_stream_get_schema( stream: *mut ArrowDeviceArrayStream, out: *mut ArrowSchema, From ee8eaa1774cdb64c1f8d002e3b405642493fa4c8 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 15:03:24 +0000 Subject: [PATCH 05/35] perf(cuda): use a set for duplicate projection names Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/lib.rs | 7 +++---- vortex-cuda/ffi/src/tests/projection.rs | 22 ++++++++++++++++++++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index f99885b0911..d2448df551e 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -34,6 +34,7 @@ use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::layout::scan::split_by::SplitBy; use vortex::session::SessionExt; use vortex::session::VortexSession; +use vortex::utils::aliases::hash_set::HashSet; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; @@ -353,6 +354,7 @@ unsafe fn scan_columns(columns: *const vx_view, ncolumns: usize) -> VortexResult // SAFETY: Null, alignment, and size were checked; the caller guarantees readable views. let columns = unsafe { std::slice::from_raw_parts(columns, ncolumns) }; let mut names = Vec::::with_capacity(ncolumns); + let mut seen = HashSet::<&str>::with_capacity(ncolumns); for (index, column) in columns.iter().enumerate() { vortex_ensure!( column.len <= isize::MAX as usize, @@ -361,10 +363,7 @@ unsafe fn scan_columns(columns: *const vx_view, ncolumns: usize) -> VortexResult // SAFETY: The caller guarantees readable name bytes. as_str checks null and UTF-8. let name = unsafe { column.as_str() } .map_err(|error| vortex_err!("invalid CUDA scan column {index}: {error}"))?; - vortex_ensure!( - !names.iter().any(|existing| existing.as_ref() == name), - "duplicate CUDA scan column: {name:?}" - ); + vortex_ensure!(seen.insert(name), "duplicate CUDA scan column: {name:?}"); names.push(FieldName::from(name)); } Ok(names.into()) diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index 52d09454517..f4cc7424ae2 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -60,6 +60,28 @@ fn test_projection_names_are_owned_and_zero_count_means_all() -> VortexResult<() Ok(()) } +#[test] +fn test_projection_wide_order_and_first_late_duplicate() -> VortexResult<()> { + let columns: Vec<_> = (0..1024).rev().map(|i| format!("column_{i}")).collect(); + let mut projection: Vec<_> = columns.iter().map(String::as_str).collect(); + let parsed = names(&projection)?; + assert_eq!( + parsed + .iter() + .map(|name| name.as_ref()) + .collect::>(), + projection + ); + + let duplicate = String::from("column_512"); + projection.extend([duplicate.as_str(), "column_1023"]); + assert_error( + names(&projection), + "duplicate CUDA scan column: \"column_512\"", + ); + Ok(()) +} + #[test] fn test_projection_rejects_invalid_names_and_counts() { let invalid_utf8 = vx_view { From 55fbac685bb127a1a319d00979948f70fba7411f Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 15:54:00 +0000 Subject: [PATCH 06/35] refactor(cuda): simplify schema cleanup and test helpers Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/tests/projection.rs | 8 +------- vortex-cuda/src/arrow/dictionary_tests.rs | 18 +++++++----------- vortex-cuda/src/arrow/mod.rs | 9 +++------ 3 files changed, 11 insertions(+), 24 deletions(-) diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index f4cc7424ae2..3bf1ef9e455 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -65,13 +65,7 @@ fn test_projection_wide_order_and_first_late_duplicate() -> VortexResult<()> { let columns: Vec<_> = (0..1024).rev().map(|i| format!("column_{i}")).collect(); let mut projection: Vec<_> = columns.iter().map(String::as_str).collect(); let parsed = names(&projection)?; - assert_eq!( - parsed - .iter() - .map(|name| name.as_ref()) - .collect::>(), - projection - ); + assert_eq!(parsed, projection.as_slice()); let duplicate = String::from("column_512"); projection.extend([duplicate.as_str(), "column_1023"]); diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index 99d41dea160..ecdf2556dfd 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -157,13 +157,13 @@ fn dictionary(values: ArrayRef, width: PType) -> VortexResult { Ok(DictArray::try_new(codes, values)?.into_array()) } -fn get_schema(stream: &mut ArrowDeviceArrayStream) -> VortexResult { +fn get_schema(stream: &mut ArrowDeviceArrayStream) -> VortexResult { let callback = stream.get_schema.expect("missing get_schema"); let mut schema = FFI_ArrowSchema::empty(); // SAFETY: The stream and output schema are live and writable. let status = unsafe { callback(stream, (&raw mut schema).cast()) }; assert_eq!(status, 0, "{}", last_error(stream)?); - Ok(schema) + Ok(Field::try_from(&schema)?) } fn get_next(stream: &mut ArrowDeviceArrayStream) -> (i32, ArrowDeviceArray) { @@ -225,10 +225,9 @@ fn test_decode_mixed_dictionary_device_stream( let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) .boxed() .export_device_array_stream(&session, &runtime)?; - let plain_schema = arrow_schema_for_array(&expected, &mut ctx)?; + let plain_schema = Field::try_from(&arrow_schema_for_array(&expected, &mut ctx)?)?; if schema_first { - let schema = get_schema(&mut stream)?; - assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain_schema)?); + assert_eq!(get_schema(&mut stream)?, plain_schema); } for _ in 0..4 { let (status, mut array) = get_next(&mut stream); @@ -238,8 +237,7 @@ fn test_decode_mixed_dictionary_device_stream( assert_arrays_eq!(actual, expected, ctx.execution_ctx()); release_device_array(&mut array); } - let schema = get_schema(&mut stream)?; - assert_eq!(Field::try_from(&schema)?, Field::try_from(&plain_schema)?); + assert_eq!(get_schema(&mut stream)?, plain_schema); let (status, eos) = get_next(&mut stream); assert_eq!(status, 0); assert!(eos.array.release.is_none()); @@ -280,9 +278,8 @@ fn test_decode_stream_schema_does_not_poll( .boxed() .export_device_array_stream(&session, &runtime)?; for _ in 0..2 { - let schema = get_schema(&mut stream)?; assert_eq!( - Field::try_from(&schema)?, + get_schema(&mut stream)?, Field::new("", DataType::Int32, false) ); } @@ -436,9 +433,8 @@ fn test_default_dictionary_device_stream(#[case] second_width: Option) -> let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) .boxed() .export_device_array_stream(&session, &runtime)?; - let schema = get_schema(&mut stream)?; assert_eq!( - Field::try_from(&schema)?.data_type(), + get_schema(&mut stream)?.data_type(), &DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Int32),) ); let (status, mut array) = get_next(&mut stream); diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index ed031d52f1e..8f35f6af947 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -435,17 +435,14 @@ impl DeviceArrayStreamPrivateData { } let ArrowDeviceArrayWithSchema { - schema: mut ffi_schema, + schema: ffi_schema, array: mut device_array, } = self .runtime .block_on(array.export_device_array_with_schema(&mut self.ctx))?; - // Release the schema we no longer need, and on failure release the array we will not - // return. - let checked = self.check_stream_array(&ffi_schema, &device_array); - release_schema(&mut ffi_schema); - let exported_schema = match checked { + // Schemas release themselves on drop; rejected device arrays need explicit release. + let exported_schema = match self.check_stream_array(&ffi_schema, &device_array) { Ok(exported_schema) => exported_schema, Err(error) => { release_device_array(&mut device_array); From ea9d2beaf0f3dab10d231ad604c5b325868b0c91 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 16:22:13 +0000 Subject: [PATCH 07/35] refactor(cuda): reduce stream and test scaffolding Signed-off-by: Alexander Droste --- Cargo.lock | 1 + vortex-cuda/ffi/Cargo.toml | 1 + vortex-cuda/ffi/README.md | 9 +- vortex-cuda/ffi/src/lib.rs | 91 ++++-------- vortex-cuda/ffi/src/tests/projection.rs | 38 +---- vortex-cuda/src/arrow/canonical.rs | 20 +-- vortex-cuda/src/arrow/dictionary_tests.rs | 61 ++++---- vortex-cuda/src/arrow/mod.rs | 173 ++++++++-------------- vortex-ffi/src/lib.rs | 1 + 9 files changed, 134 insertions(+), 261 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1f52ea42973..c5f0d08fdff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10913,6 +10913,7 @@ dependencies = [ "arrow-schema 59.3.0", "cbindgen", "futures", + "tempfile", "vortex", "vortex-cuda", "vortex-cuda-macros", diff --git a/vortex-cuda/ffi/Cargo.toml b/vortex-cuda/ffi/Cargo.toml index 32ec00124ac..939a38b1b63 100644 --- a/vortex-cuda/ffi/Cargo.toml +++ b/vortex-cuda/ffi/Cargo.toml @@ -22,6 +22,7 @@ vortex-cuda = { path = ".." } vortex-ffi = { path = "../../vortex-ffi" } [dev-dependencies] +tempfile = { workspace = true } vortex-cuda-macros = { workspace = true } [build-dependencies] diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index aedd25046e9..5064d0a9035 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -1,9 +1,10 @@ # vortex-cuda-ffi -CUDA-specific C FFI helpers for cuDF interop, keeping CUDA out of the base `vortex-ffi` -crate. The public C API exports a borrowed `vx_array` as an `ArrowSchema + ArrowDeviceArray` -pair, not cuDF objects. The caller passes these structs to cuDF and releases them after -cuDF finishes importing. +CUDA-specific C FFI helpers for cuDF interop. + +This crate keeps CUDA out of the base `vortex-ffi` crate. Its public C API exports a borrowed `vx_array` as an `ArrowSchema + ArrowDeviceArray` pair. + +It does not create cuDF objects itself. The caller passes the exported Arrow Device structs to cuDF and releases them after cuDF is done importing. Use this crate as the CUDA-enabled FFI artifact. Include both headers: diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index d2448df551e..d45d38ddb35 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -539,8 +539,12 @@ mod tests { use vortex::array::validity::Validity; use vortex::error::VortexResult; use vortex_cuda::arrow::ARROW_DEVICE_CUDA; + use vortex_cuda::arrow::release_device_array; + use vortex_cuda::arrow::release_schema; use vortex_cuda_macros::cuda_not_available; use vortex_cuda_macros::test as cuda_test; + use vortex_ffi::vx_array_free as free_test_array; + use vortex_ffi::vx_session_free as free_test_session; use super::*; @@ -614,14 +618,12 @@ mod tests { ctx, ffi_runtime(), ); - let get_schema = stream.get_schema.expect("missing get_schema"); let get_next = stream.get_next.expect("missing get_next"); let release = stream.release.expect("missing release"); - let mut schema = FFI_ArrowSchema::empty(); + let schema = projection::stream_schema(&mut stream); let mut exported = empty_device_array(); - // SAFETY: The live stream owns these callbacks, and both outputs are writable. + // SAFETY: The live stream owns the callback, and the output is writable. unsafe { - assert_eq!(get_schema(&raw mut stream, (&raw mut schema).cast()), 0); assert_eq!(get_next(&raw mut stream, &raw mut exported), 0); } assert_eq!(Field::try_from(&schema)?.data_type(), &expected_type); @@ -645,32 +647,8 @@ mod tests { Box::into_raw(Box::new(session)).cast::() } - unsafe fn free_test_session(session: *mut vx_session) { - unsafe { drop(Box::from_raw(session.cast::())) }; - } - fn test_array(array: impl IntoArray) -> *const vx_array { - Arc::into_raw(Arc::new(array.into_array())).cast::() - } - - unsafe fn free_test_array(array: *const vx_array) { - unsafe { Arc::decrement_strong_count(array.cast::()) }; - } - - unsafe fn release_schema(schema: &mut FFI_ArrowSchema) { - unsafe { - if let Some(release) = schema.release { - release(schema); - } - } - } - - unsafe fn release_device_array(array: &mut ArrowDeviceArray) { - unsafe { - if let Some(release) = array.array.release { - release(&raw mut array.array); - } - } + Box::into_raw(Box::new(array.into_array())).cast::() } fn empty_device_array() -> ArrowDeviceArray { @@ -683,14 +661,16 @@ mod tests { } } - #[cuda_test] - fn test_export_primitive_arrow_device() { + /// # Safety + /// `session` and `array` must be valid borrowed FFI handles for the duration of the call. + unsafe fn export_array( + session: *const vx_session, + array: *const vx_array, + ) -> (FFI_ArrowSchema, ArrowDeviceArray) { let mut error = ptr::null_mut(); - let session = test_session(VortexSession::default()); - let array = test_array(PrimitiveArray::from_iter(0u32..5)); let mut schema = FFI_ArrowSchema::empty(); let mut device_array = empty_device_array(); - + // SAFETY: The caller guarantees valid handles; all outputs are live and writable. let status = unsafe { vx_cuda_array_export_arrow_device( session, @@ -702,6 +682,15 @@ mod tests { }; assert_eq!(status, VX_CUDA_OK); assert!(error.is_null()); + (schema, device_array) + } + + #[cuda_test] + fn test_export_primitive_arrow_device() { + let session = test_session(VortexSession::default()); + let array = test_array(PrimitiveArray::from_iter(0u32..5)); + // SAFETY: Both handles remain live until cleanup below. + let (mut schema, mut device_array) = unsafe { export_array(session, array) }; let field = Field::try_from(&schema).expect("schema should be a field"); assert_eq!(field.name(), ""); @@ -721,7 +710,6 @@ mod tests { #[cuda_test] fn test_export_struct_arrow_device_table() -> VortexResult<()> { - let mut error = ptr::null_mut(); let session = test_session(VortexSession::default()); let array = test_array(StructArray::try_new( ["ids", "values"].into(), @@ -732,21 +720,8 @@ mod tests { 3, Validity::NonNullable, )?); - - let mut schema = FFI_ArrowSchema::empty(); - let mut device_array = empty_device_array(); - - let status = unsafe { - vx_cuda_array_export_arrow_device( - session, - array, - &raw mut schema, - &raw mut device_array, - &raw mut error, - ) - }; - assert_eq!(status, VX_CUDA_OK); - assert!(error.is_null()); + // SAFETY: Both handles remain live until cleanup below. + let (mut schema, mut device_array) = unsafe { export_array(session, array) }; let arrow_schema = Schema::try_from(&schema)?; assert_eq!(arrow_schema.fields().len(), 2); @@ -786,20 +761,8 @@ mod tests { assert!(!session.is_null()); let array = test_array(PrimitiveArray::from_iter(0u32..5)); - let mut schema = FFI_ArrowSchema::empty(); - let mut device_array = empty_device_array(); - - let status = unsafe { - vx_cuda_array_export_arrow_device( - session, - array, - &raw mut schema, - &raw mut device_array, - &raw mut error, - ) - }; - assert_eq!(status, VX_CUDA_OK); - assert!(error.is_null()); + // SAFETY: Both handles remain live until cleanup below. + let (mut schema, mut device_array) = unsafe { export_array(session, array) }; assert_eq!(device_array.array.length, 5); assert_eq!(device_array.device_type, ARROW_DEVICE_CUDA); diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index 3bf1ef9e455..52ee9e13c8a 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -4,11 +4,11 @@ use std::ffi::CStr; use std::io::Write; use std::mem::MaybeUninit; -use std::path::PathBuf; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use futures::TryStreamExt; +use tempfile::NamedTempFile; use vortex::array::VortexSessionExecute; use vortex::array::assert_arrays_eq; use vortex::buffer::ByteBuffer; @@ -277,32 +277,6 @@ fn test_projection_ffi_validation_without_cuda() { ); } -struct LocalFile(PathBuf); - -impl LocalFile { - fn new(bytes: &[u8]) -> VortexResult { - static NEXT_ID: AtomicUsize = AtomicUsize::new(0); - let path = std::env::temp_dir().join(format!( - "vortex-cuda-ffi-projection-{}-{}.vortex", - std::process::id(), - NEXT_ID.fetch_add(1, Ordering::Relaxed) - )); - let mut file = std::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&path)?; - let result = Self(path); - file.write_all(bytes)?; - Ok(result) - } -} - -impl Drop for LocalFile { - fn drop(&mut self) { - drop(std::fs::remove_file(&self.0)); - } -} - fn stream_error(stream: &mut ArrowDeviceArrayStream) -> String { // SAFETY: The callback and returned C string belong to this live stream. unsafe { @@ -343,7 +317,7 @@ fn open_stream( unsafe { output.assume_init() } } -fn stream_schema(stream: &mut ArrowDeviceArrayStream) -> FFI_ArrowSchema { +pub(super) fn stream_schema(stream: &mut ArrowDeviceArrayStream) -> FFI_ArrowSchema { let mut schema = FFI_ArrowSchema::empty(); let get_schema = stream.get_schema.expect("missing get_schema"); // SAFETY: This live stream owns the callback; schema is writable. @@ -374,8 +348,7 @@ fn batch_lengths(stream: &mut ArrowDeviceArrayStream) -> Vec { assert_eq!(array.device_type, ARROW_DEVICE_CUDA); assert_eq!(array.array.n_children, 2); lengths.push(array.array.length); - // SAFETY: Each live batch is released exactly once, before requesting the next one. - unsafe { release_device_array(&mut array) }; + release_device_array(&mut array); } lengths } @@ -384,13 +357,14 @@ fn batch_lengths(stream: &mut ArrowDeviceArrayStream) -> Vec { fn test_projection_gpu_local_file_schema_and_batch_boundaries() -> VortexResult<()> { for (block_rows, batch_rows) in [(0, 2), (2, 3)] { let session = session().with_some(CudaSession::try_default()?); - let file = LocalFile::new(&file_bytes( + let mut file = NamedTempFile::new()?; + file.write_all(&file_bytes( &session, table()?.into_array(), Some(block_rows), )?)?; let path = file - .0 + .path() .to_str() .ok_or_else(|| vortex_err!("non-UTF-8 test path"))?; let options = vx_cuda_scan_options { diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index f2948fb223f..b441bdfcfb5 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -1767,28 +1767,14 @@ mod tests { array: &ArrowArray, buffer_idx: usize, ) -> VortexResult> { - let private_data = unsafe { &*array.private_data.cast::() }; - let buffer = private_data.buffers[buffer_idx] - .as_ref() - .vortex_expect("buffer should be present"); - Ok(Buffer::::from_byte_buffer(buffer.to_host_sync()) - .iter() - .copied() - .collect()) + Ok(Buffer::::from_byte_buffer(private_data_buffer_bytes(array, buffer_idx)?).to_vec()) } fn private_data_buffer_i16_values( array: &ArrowArray, buffer_idx: usize, ) -> VortexResult> { - let private_data = unsafe { &*array.private_data.cast::() }; - let buffer = private_data.buffers[buffer_idx] - .as_ref() - .vortex_expect("buffer should be present"); - Ok(Buffer::::from_byte_buffer(buffer.to_host_sync()) - .iter() - .copied() - .collect()) + Ok(Buffer::::from_byte_buffer(private_data_buffer_bytes(array, buffer_idx)?).to_vec()) } fn private_data_buffer_bytes( @@ -2617,7 +2603,7 @@ mod tests { let mut slots = Vec::new(); for slot in fsst.slots().iter() { slots.push(match slot { - Some(child) => Some(upload(child.clone(), &mut ctx).await?), + Some(child) => Some(upload(child.clone(), &mut ctx)?), None => None, }); } diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index ecdf2556dfd..e33fe183157 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -5,7 +5,6 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Poll; -use futures::future::BoxFuture; use futures::stream; use rstest::rstest; use vortex::array::IntoArray; @@ -31,29 +30,24 @@ use crate::CudaSession; /// Preserve encodings while moving all buffers, including validity, to CUDA so unsupported /// decoding errors instead of falling back to the CPU. -pub(super) fn upload( - array: ArrayRef, - ctx: &mut CudaExecutionCtx, -) -> BoxFuture<'_, VortexResult> { - Box::pin(async move { - // Constants store scalar metadata, not replaceable data buffers. - if array.as_opt::().is_some() { - return Ok(array); - } - let mut slots = Vec::new(); - for slot in array.slots().iter() { - slots.push(match slot { - Some(child) => Some(upload(child.clone(), ctx).await?), - None => None, - }); - } - let mut buffers = Vec::new(); - for buffer in array.buffer_handles() { - buffers.push(ctx.ensure_on_device(buffer).await?); - } - // SAFETY: Slots and buffers are byte-for-byte copies; only their placement changes. - unsafe { array.with_slots(slots.into())?.with_buffers(buffers) } - }) +pub(super) fn upload(array: ArrayRef, ctx: &mut CudaExecutionCtx) -> VortexResult { + // Constants store scalar metadata, not replaceable data buffers. + if array.as_opt::().is_some() { + return Ok(array); + } + let mut slots = Vec::new(); + for slot in array.slots().iter() { + slots.push(match slot { + Some(child) => Some(upload(child.clone(), ctx)?), + None => None, + }); + } + let mut buffers = Vec::new(); + for buffer in array.buffer_handles() { + buffers.push(ctx.ensure_on_device_sync(buffer)?); + } + // SAFETY: Slots and buffers are byte-for-byte copies; only their placement changes. + unsafe { array.with_slots(slots.into())?.with_buffers(buffers) } } /// Copy a device buffer from a live, unreleased array produced by this exporter to the host. @@ -90,12 +84,7 @@ fn read_plain(array: &ArrowArray, dtype: &DType) -> VortexResult { DType::Utf8(_) => { assert_eq!(array.n_buffers, 3); assert_eq!(array.n_children, 0); - let offsets = PrimitiveArray::from_byte_buffer( - buffer(array, 1)?, - PType::I32, - Validity::NonNullable, - ) - .into_array(); + let offsets = Buffer::::from_byte_buffer(buffer(array, 1)?).into_array(); Ok(VarBinArray::try_new( offsets, buffer(array, 2)?.slice_unaligned(..), @@ -175,13 +164,13 @@ fn get_next(stream: &mut ArrowDeviceArrayStream) -> (i32, ArrowDeviceArray) { } /// Upload chunks and synchronize before handing them to a separate export context. -async fn upload_chunks( +fn upload_chunks( chunks: Vec, ctx: &mut CudaExecutionCtx, ) -> VortexResult>> { let mut device_chunks = Vec::new(); for chunk in chunks { - let chunk = upload(chunk, ctx).await?; + let chunk = upload(chunk, ctx)?; assert!(!chunk.is_host()); device_chunks.push(Ok(chunk)); } @@ -221,7 +210,7 @@ fn test_decode_mixed_dictionary_device_stream( let wrap = |array| if nested { wrap_struct(array) } else { array }; let expected = wrap(expected); let chunks = chunks.into_iter().map(wrap).collect(); - let chunks = runtime.block_on(upload_chunks(chunks, &mut ctx))?; + let chunks = upload_chunks(chunks, &mut ctx)?; let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) .boxed() .export_device_array_stream(&session, &runtime)?; @@ -359,7 +348,7 @@ async fn test_decode_non_contiguous_dictionary_list_view() -> VortexResult<()> { ) .into_array(); let expected = expected.take(PrimitiveArray::from_iter([2u32, 3, 0, 1]).into_array())?; - let array = upload(array, &mut ctx).await?; + let array = upload(array, &mut ctx)?; let mut exported = array.export_device_array_with_schema(&mut ctx).await?; assert_eq!( Field::try_from(&exported.schema)?, @@ -389,7 +378,7 @@ async fn test_decode_unsupported_device_dictionary_does_not_fall_back_to_cpu() - let mut ctx = CudaSession::create_execution_ctx(&session)?; // A dictionary of structs can be preserved, but has no CUDA gather kernel today. let (values, _) = values_and_expected(false); - let array = upload(dictionary(wrap_struct(values), PType::U8)?, &mut ctx).await?; + let array = upload(dictionary(wrap_struct(values), PType::U8)?, &mut ctx)?; assert!(!array.is_host()); let mut preserved = array.clone().export_device_array(&mut ctx).await?; assert!(!preserved.array.dictionary.is_null()); @@ -429,7 +418,7 @@ fn test_default_dictionary_device_stream(#[case] second_width: Option) -> Some(width) => dictionary(values, width)?, None => expected.clone(), }; - let chunks = runtime.block_on(upload_chunks(vec![first, second], &mut ctx))?; + let chunks = upload_chunks(vec![first, second], &mut ctx)?; let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) .boxed() .export_device_array_stream(&session, &runtime)?; diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 8f35f6af947..fe7243b4f6f 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -389,15 +389,19 @@ impl DeviceArrayStreamPrivateData { .ok_or_else(|| vortex_err!("ArrowDeviceArrayStream schema was not initialized")) } - /// Export and return the next Arrow device array, or `None` at end of stream. - fn next_array(&mut self) -> VortexResult> { + /// Export the next array, or return a released array at end of stream. + fn next_array(&mut self) -> VortexResult { if let Some(array) = self.pending_array.take() { - return Ok(Some(array)); + return Ok(array); } match self.array_iter.next() { - Some(array) => self.export_stream_array(array?).map(Some), - None => Ok(None), + Some(array) => self.export_stream_array(array?), + None => Ok(ArrowDeviceArray { + device_id: self.device_id, + device_type: ARROW_DEVICE_CUDA, + ..ArrowDeviceArray::empty() + }), } } @@ -410,47 +414,53 @@ impl DeviceArrayStreamPrivateData { array.dtype() ); - if self.ctx.cuda_session().dictionary_export() == DictionaryExport::Decode { - // The canonical exporter uses only the dtype and fixed context settings in this mode. - // Avoid constructing and parsing a temporary C schema for every batch. - let schema = if self.schema.is_none() { - Some(ArrowDeviceStreamSchema::from_dtype( - &self.dtype, - &mut self.ctx, - )?) + let mut staged_schema = None; + let (mut device_array, ffi_schema) = + if self.ctx.cuda_session().dictionary_export() == DictionaryExport::Decode { + // Decode schemas depend only on dtype and fixed context settings. Stage before + // export to preserve error ordering, but cache only after device validation. + if self.schema.is_none() { + staged_schema = Some(ArrowDeviceStreamSchema::from_dtype( + &self.dtype, + &mut self.ctx, + )?); + } + let device_array = self + .runtime + .block_on(array.export_device_array(&mut self.ctx))?; + (device_array, None) } else { - None + let exported = self + .runtime + .block_on(array.export_device_array_with_schema(&mut self.ctx))?; + (exported.array, Some(exported.schema)) }; - let mut device_array = self - .runtime - .block_on(array.export_device_array(&mut self.ctx))?; - if let Err(error) = self.check_device(&device_array) { - release_device_array(&mut device_array); - return Err(error); - } - if let Some(schema) = schema { - self.schema = Some(schema); - } - return Ok(device_array); - } - - let ArrowDeviceArrayWithSchema { - schema: ffi_schema, - array: mut device_array, - } = self - .runtime - .block_on(array.export_device_array_with_schema(&mut self.ctx))?; // Schemas release themselves on drop; rejected device arrays need explicit release. - let exported_schema = match self.check_stream_array(&ffi_schema, &device_array) { - Ok(exported_schema) => exported_schema, - Err(error) => { - release_device_array(&mut device_array); - return Err(error); + let validation = (|| { + self.check_device(&device_array)?; + if let Some(ffi_schema) = &ffi_schema { + let exported_schema = ArrowDeviceStreamSchema::from_ffi(ffi_schema, &self.dtype)?; + if let Some(stream_schema) = &self.schema { + vortex_ensure!( + stream_schema == &exported_schema, + "stream array Arrow schema changed from {:?} to {:?}; an Arrow C device stream \ + requires every array to share one schema, so chunks must not vary their \ + encoding (for example a dictionary-encoded chunk among plain chunks)", + stream_schema, + exported_schema + ); + } + staged_schema = Some(exported_schema); } - }; + Ok(()) + })(); + if let Err(error) = validation { + release_device_array(&mut device_array); + return Err(error); + } if self.schema.is_none() { - self.schema = Some(exported_schema); + self.schema = staged_schema; } Ok(device_array) } @@ -469,27 +479,6 @@ impl DeviceArrayStreamPrivateData { ); Ok(()) } - - /// Check that a freshly exported device array matches the stream schema and CUDA device. - fn check_stream_array( - &self, - ffi_schema: &FFI_ArrowSchema, - device_array: &ArrowDeviceArray, - ) -> VortexResult { - self.check_device(device_array)?; - let exported_schema = ArrowDeviceStreamSchema::from_ffi(ffi_schema, &self.dtype)?; - if let Some(stream_schema) = &self.schema { - vortex_ensure!( - stream_schema == &exported_schema, - "stream array Arrow schema changed from {:?} to {:?}; an Arrow C device stream \ - requires every array to share one schema, so chunks must not vary their \ - encoding (for example a dictionary-encoded chunk among plain chunks)", - stream_schema, - exported_schema - ); - } - Ok(exported_schema) - } } impl Drop for DeviceArrayStreamPrivateData { @@ -585,15 +574,6 @@ unsafe fn device_stream_private_data<'a>( } } -/// Create the Arrow end-of-stream marker for the stream's CUDA device. -fn released_device_array(device_id: i64) -> ArrowDeviceArray { - ArrowDeviceArray { - device_id, - device_type: ARROW_DEVICE_CUDA, - ..ArrowDeviceArray::empty() - } -} - /// Release an Arrow C schema if it is live. pub fn release_schema(schema: &mut FFI_ArrowSchema) { if let Some(release) = schema.release { @@ -608,22 +588,6 @@ pub fn release_device_array(array: &mut ArrowDeviceArray) { } } -/// Runs an Arrow stream callback body. -/// -/// Returns an Arrow callback status code and stores failures in `last_error`. -fn device_stream_callback( - state: &mut DeviceArrayStreamPrivateData, - panic_message: &'static str, - callback: impl FnOnce(&mut DeviceArrayStreamPrivateData) -> VortexResult<()>, -) -> c_int { - let result = catch_unwind(AssertUnwindSafe(|| callback(state))); - match result { - Ok(Ok(())) => 0, - Ok(Err(err)) => state.set_error(err, LIBC_EIO), - Err(_) => state.set_error(panic_message, LIBC_EIO), - } -} - /// Write the stream's Arrow schema, deriving it from the dtype or first array as needed. unsafe extern "C" fn device_stream_get_schema( stream: *mut ArrowDeviceArrayStream, @@ -638,17 +602,17 @@ unsafe extern "C" fn device_stream_get_schema( return state.set_error("null ArrowSchema output", LIBC_EINVAL); } - fn body(state: &mut DeviceArrayStreamPrivateData, out: *mut ArrowSchema) -> VortexResult<()> { + let result = catch_unwind(AssertUnwindSafe(|| -> VortexResult<()> { let schema = state.get_or_init_schema()?.to_ffi()?; + // SAFETY: out is non-null; the caller provides writable ArrowSchema storage. unsafe { ptr::write(out.cast::(), schema) }; Ok(()) + })); + match result { + Ok(Ok(())) => 0, + Ok(Err(err)) => state.set_error(err, LIBC_EIO), + Err(_) => state.set_error("panic in ArrowDeviceArrayStream::get_schema", LIBC_EIO), } - - device_stream_callback( - state, - "panic in ArrowDeviceArrayStream::get_schema", - |state| body(state, out), - ) } /// Write the next exported Arrow device array, or a released array at end of stream. @@ -665,24 +629,17 @@ unsafe extern "C" fn device_stream_get_next( return state.set_error("null ArrowDeviceArray output", LIBC_EINVAL); } - // Keep the fallible part in a local function so `device_stream_callback` handles callback - // status and error reporting consistently. - fn body( - state: &mut DeviceArrayStreamPrivateData, - out: *mut ArrowDeviceArray, - ) -> VortexResult<()> { - let array = state - .next_array()? - .unwrap_or_else(|| released_device_array(state.device_id)); + let result = catch_unwind(AssertUnwindSafe(|| -> VortexResult<()> { + let array = state.next_array()?; + // SAFETY: out is non-null; the caller provides writable ArrowDeviceArray storage. unsafe { ptr::write(out, array) }; Ok(()) + })); + match result { + Ok(Ok(())) => 0, + Ok(Err(err)) => state.set_error(err, LIBC_EIO), + Err(_) => state.set_error("panic in ArrowDeviceArrayStream::get_next", LIBC_EIO), } - - device_stream_callback( - state, - "panic in ArrowDeviceArrayStream::get_next", - |state| body(state, out), - ) } /// Return the most recent callback error message, or null if no error is stored. diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 1ef423bb383..1b1a3f1bcf6 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -27,6 +27,7 @@ use std::sync::Arc; use std::sync::LazyLock; pub use array::vx_array; +pub use array::vx_array_free; pub use array::vx_array_ref; pub use dtype::vx_dtype; pub use error::try_or; From 78f6eb63313072580ae84ecf5a1f21661d324931 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Thu, 17 Sep 2026 16:28:11 +0000 Subject: [PATCH 08/35] docs(cuda): clarify dictionary export modes with an example Signed-off-by: Alexander Droste --- vortex-cuda/src/session.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index 9b377384952..9d46af4683e 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -40,15 +40,24 @@ pub enum VarBinExportLayout { VarBinView, } -/// Controls whether Arrow Device exports preserve dictionaries or expand them to plain values, -/// including nested children. +/// Controls whether Arrow Device exports keep dictionary encoding or fully decode it, +/// including dictionaries nested inside structs and lists. +/// +/// For example, indices `[0, 1, 0]` and dictionary values `["apple", "pear"]` export as: +/// +/// - [`Preserve`](Self::Preserve): separate index and dictionary-value arrays, with an Arrow +/// dictionary type. +/// - [`Decode`](Self::Decode): the plain string array `["apple", "pear", "apple"]`, with no +/// dictionary indices or dictionary child. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum DictionaryExport { - /// Preserve dictionary values and indices in the Arrow schema and device array. + /// Keep separate index and dictionary-value arrays rather than expanding repeated values. + /// The Arrow schema retains the dictionary type, including its index type. #[default] Preserve, - /// Decode dictionaries on CUDA to keep one plain Arrow schema across batches. - /// May increase device memory use. Device-resident inputs require CUDA decoding support. + /// Fully decode dictionary encoding into plain values on CUDA, including repeated values. + /// This keeps one plain Arrow schema across batches with different dictionary encodings, + /// but may increase device memory use. Device-resident inputs require CUDA decoding support. Decode, } From 744abf1d9264f3dcc88a80a9e604b6e83ef8e3f0 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Fri, 18 Sep 2026 13:44:53 +0000 Subject: [PATCH 09/35] Fix CUDA test lints and edition policy expectations Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/dictionary_tests.rs | 7 ++----- vortex-cuda/src/layout.rs | 7 ++++++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index e33fe183157..ca7eba4d3ee 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -302,7 +302,6 @@ fn test_decode_stream_validates_dtype_and_device() -> VortexResult<()> { .with_some(CudaSession::try_default()?.with_dictionary_export(DictionaryExport::Decode)); let array = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let mut stream = array - .clone() .to_array_stream() .boxed() .export_device_array_stream(&session, &runtime)?; @@ -314,16 +313,14 @@ fn test_decode_stream_validates_dtype_and_device() -> VortexResult<()> { let error = state .export_stream_array(PrimitiveArray::from_iter([10u32, 20, 30]).into_array()) - .err() - .expect("accepted a different dtype"); + .expect_err("accepted a different dtype"); assert!(error.to_string().contains("stream array dtype changed")); let error = state.check_device(&ArrowDeviceArray::empty()).unwrap_err(); assert!(error.to_string().contains("non-CUDA device type")); state.device_id = -1; let error = state .export_stream_array(array) - .err() - .expect("accepted a different device"); + .expect_err("accepted a different device"); assert!( error .to_string() diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index ee96f690524..86da0a921ab 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -608,6 +608,11 @@ mod tests { ) -> VortexResult<()> { let session = VortexSession::default(); session.enable_edition(core)?; + let mut expected_editions = session.enabled_editions().editions(); + assert!(expected_editions.contains(&core)); + assert!(!expected_editions.contains(&CUDA_EDITION)); + expected_editions.push(CUDA_EDITION); + expected_editions.sort_unstable(); let kinds = [ ComponentKind::Array, ComponentKind::Layout, @@ -637,7 +642,7 @@ mod tests { } let mut enabled_editions = session.enabled_editions().editions(); enabled_editions.sort_unstable(); - assert_eq!(enabled_editions, [core, CUDA_EDITION]); + assert_eq!(enabled_editions, expected_editions); session.editions().validate()?; assert!( !VortexSession::default() From b5da45bbacecec87213aafcdc7f4e5fe68dfdaf6 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 09:44:25 +0000 Subject: [PATCH 10/35] test(cuda): clarify projected scan fixtures and verify device values Signed-off-by: Alexander Droste --- Cargo.lock | 1 + vortex-cuda/ffi/Cargo.toml | 1 + vortex-cuda/ffi/src/lib.rs | 27 ++- vortex-cuda/ffi/src/tests/projection.rs | 297 +++++++++++++++++------- vortex-ffi/src/lib.rs | 1 + 5 files changed, 236 insertions(+), 91 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c5f0d08fdff..92f52bd0198 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10912,6 +10912,7 @@ version = "0.1.0" dependencies = [ "arrow-schema 59.3.0", "cbindgen", + "cudarc", "futures", "tempfile", "vortex", diff --git a/vortex-cuda/ffi/Cargo.toml b/vortex-cuda/ffi/Cargo.toml index 939a38b1b63..6ec839c100f 100644 --- a/vortex-cuda/ffi/Cargo.toml +++ b/vortex-cuda/ffi/Cargo.toml @@ -22,6 +22,7 @@ vortex-cuda = { path = ".." } vortex-ffi = { path = "../../vortex-ffi" } [dev-dependencies] +cudarc = { workspace = true } tempfile = { workspace = true } vortex-cuda-macros = { workspace = true } diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index d45d38ddb35..9a206e79ba3 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -521,6 +521,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_partition_scan_arrow_device_stream( mod tests { mod projection; + use std::ffi::CStr; use std::ptr; use std::sync::Arc; @@ -620,7 +621,7 @@ mod tests { ); let get_next = stream.get_next.expect("missing get_next"); let release = stream.release.expect("missing release"); - let schema = projection::stream_schema(&mut stream); + let schema = stream_schema(&mut stream); let mut exported = empty_device_array(); // SAFETY: The live stream owns the callback, and the output is writable. unsafe { @@ -651,6 +652,30 @@ mod tests { Box::into_raw(Box::new(array.into_array())).cast::() } + fn stream_error(stream: &mut ArrowDeviceArrayStream) -> String { + // SAFETY: The callback and returned C string belong to this live stream. + unsafe { + stream + .get_last_error + .and_then(|callback| callback(stream).as_ref()) + .map(|message| CStr::from_ptr(message).to_string_lossy().into_owned()) + .unwrap_or_default() + } + } + + fn stream_schema(stream: &mut ArrowDeviceArrayStream) -> FFI_ArrowSchema { + let mut schema = FFI_ArrowSchema::empty(); + let get_schema = stream.get_schema.expect("missing get_schema"); + // SAFETY: This live stream owns the callback; schema is writable. + assert_eq!( + unsafe { get_schema(stream, (&raw mut schema).cast()) }, + 0, + "{}", + stream_error(stream) + ); + schema + } + fn empty_device_array() -> ArrowDeviceArray { ArrowDeviceArray { array: vortex_cuda::arrow::ArrowArray::empty(), diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index 52ee9e13c8a..3e747d5761d 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -1,18 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ffi::CStr; use std::io::Write; use std::mem::MaybeUninit; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use cudarc::driver::CudaContext; +use cudarc::driver::result; +use cudarc::driver::sys::CUevent; use futures::TryStreamExt; use tempfile::NamedTempFile; use vortex::array::VortexSessionExecute; +use vortex::array::arrays::ChunkedArray; use vortex::array::assert_arrays_eq; +use vortex::buffer::BitBuffer; +use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; use vortex::buffer::ByteBufferMut; +use vortex::dtype::NativePType; +use vortex::dtype::Nullability; use vortex::file::WriteOptionsSessionExt; use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutStrategy; @@ -21,6 +28,9 @@ use vortex::layout::layouts::table::TableStrategy; use vortex::layout::segments::SegmentFuture; use vortex::layout::segments::SegmentId; use vortex::layout::segments::SegmentSource; +use vortex_cuda::arrow::ArrowArray; +use vortex_ffi::vx_error_free; +use vortex_ffi::vx_error_message; use super::*; @@ -128,15 +138,8 @@ fn table() -> VortexResult { fn file_bytes( session: &VortexSession, array: ArrayRef, - cuda_block_rows: Option, + strategy: Arc, ) -> VortexResult { - let strategy: Arc = if let Some(block_rows) = cuda_block_rows { - register_cuda_layout(session); - cuda_write_strategy(session, block_rows) - } else { - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - Arc::new(TableStrategy::new(Arc::clone(&flat), flat)) - }; let mut bytes = ByteBufferMut::empty(); ffi_runtime().block_on( session @@ -150,11 +153,11 @@ fn file_bytes( fn open_file( session: &VortexSession, array: ArrayRef, - cuda_block_rows: Option, + strategy: Arc, ) -> VortexResult { session .open_options() - .open_buffer(file_bytes(session, array, cuda_block_rows)?) + .open_buffer(file_bytes(session, array, strategy)?) } #[test] @@ -165,7 +168,8 @@ fn test_cuda_write_strategy_preserves_high_cardinality_row_blocks() -> VortexRes let rows = ids.len(); let input = StructArray::try_new(["ids"].into(), vec![ids], rows, Validity::NonNullable)?.into_array(); - let file = open_file(&session, input, Some(rows))?; + register_cuda_layout(&session); + let file = open_file(&session, input, cuda_write_strategy(&session, rows))?; let lengths: Vec<_> = ffi_runtime().block_on( projected_scan(&file, names(&["ids"])?, rows)? .into_array_stream()? @@ -200,7 +204,9 @@ fn test_projection_cpu_never_requests_unselected_column_segments() -> VortexResu let input = table()?; let columns = names(&["値.x", "ids"])?; let expected = input.project(columns.as_ref())?.into_array(); - let file = open_file(&session, input.into_array(), None)?; + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let strategy = Arc::new(TableStrategy::new(Arc::clone(&flat), flat)); + let file = open_file(&session, input.into_array(), strategy)?; // TableStrategy writes one flat child per column, so child 1 is exactly the unused column. let children = file.footer().layout().children()?; let forbidden = children[1].segment_ids(); @@ -254,13 +260,17 @@ fn test_projection_ffi_validation_without_cuda() { &raw mut error, ) }; - assert_eq!(status, VX_CUDA_ERR); + // SAFETY: This call owns the returned error and frees it exactly once. + let message = unsafe { take_error_message(error) }.expect("missing FFI error"); + assert_eq!(status, VX_CUDA_ERR, "{message}"); + assert!( + message.contains("null CUDA scan columns with nonzero count"), + "{message}" + ); assert_eq!(stream.device_type, -1); assert!(stream.release.is_none()); - assert!(!error.is_null()); - // SAFETY: This call owns the returned error and frees it exactly once. - unsafe { vortex_ffi::vx_error_free(error) }; - // SAFETY: Null output is rejected before any other input is used; error output is optional. + error = ptr::null_mut(); + // SAFETY: Null output is rejected before any other input is used; error is writable. assert_eq!( unsafe { vx_cuda_scan_path_arrow_device_stream_projected( @@ -270,33 +280,43 @@ fn test_projection_ffi_validation_without_cuda() { ptr::null(), 0, ptr::null_mut(), - ptr::null_mut(), + &raw mut error, ) }, VX_CUDA_ERR ); + // SAFETY: This call owns the returned error and frees it exactly once. + let message = unsafe { take_error_message(error) }.expect("missing FFI error"); + assert!( + message.contains("null ArrowDeviceArrayStream output"), + "{message}" + ); } -fn stream_error(stream: &mut ArrowDeviceArrayStream) -> String { - // SAFETY: The callback and returned C string belong to this live stream. - unsafe { - stream - .get_last_error - .and_then(|callback| callback(stream).as_ref()) - .map(|message| CStr::from_ptr(message).to_string_lossy().into_owned()) - .unwrap_or_default() +/// # Safety +/// `error` must be null or an owned FFI error, consumed by this call. +unsafe fn take_error_message(error: *mut vx_error) -> Option { + if error.is_null() { + return None; } + // SAFETY: The error remains live while its message is copied, then is freed exactly once. + let message = unsafe { vx_error_message(error).as_str() } + .map(str::to_owned) + .unwrap_or_else(|error| error.to_string()); + unsafe { vx_error_free(error) }; + Some(message) } fn open_stream( session: &VortexSession, path: &str, options: &vx_cuda_scan_options, + columns: &[&str], ) -> ArrowDeviceArrayStream { let mut output = MaybeUninit::::uninit(); let mut error = ptr::null_mut(); let handle = test_session(session.clone()); - let columns = [view("値.x"), view("ids")]; + let columns: Vec<_> = columns.iter().map(|name| view(name)).collect(); // SAFETY: All borrowed inputs and writable outputs are live for this call. let status = unsafe { vx_cuda_scan_path_arrow_device_stream_projected( @@ -309,84 +329,181 @@ fn open_stream( &raw mut error, ) }; - // SAFETY: This is the sole release of the borrowed session handle. - unsafe { free_test_session(handle) }; - assert_eq!(status, VX_CUDA_OK); - assert!(error.is_null()); + // SAFETY: This call owns both the session handle and any returned error. + let message = unsafe { + free_test_session(handle); + take_error_message(error) + }; + assert_eq!( + status, + VX_CUDA_OK, + "{}", + message.as_deref().unwrap_or("no FFI error") + ); + assert!(message.is_none(), "unexpected FFI error: {message:?}"); // SAFETY: A successful call initialized the stream, which owns its session state. unsafe { output.assume_init() } } -pub(super) fn stream_schema(stream: &mut ArrowDeviceArrayStream) -> FFI_ArrowSchema { - let mut schema = FFI_ArrowSchema::empty(); - let get_schema = stream.get_schema.expect("missing get_schema"); - // SAFETY: This live stream owns the callback; schema is writable. - assert_eq!( - unsafe { get_schema(stream, (&raw mut schema).cast()) }, - 0, - "{}", - stream_error(stream) - ); - schema +/// # Safety +/// `array` must be a live Arrow primitive array of `T` on the current CUDA context. +/// Its producer's sync event must have completed before this call. +unsafe fn read_primitive( + array: &ArrowArray, + nullability: Nullability, +) -> VortexResult { + assert!(array.release.is_some()); + assert!(array.dictionary.is_null()); + assert_eq!(array.n_buffers, 2); + assert_eq!(array.n_children, 0); + assert!(!array.buffers.is_null()); + let len = usize::try_from(array.length)?; + let offset = usize::try_from(array.offset)?; + // SAFETY: The live primitive array owns two buffer pointers. + let buffers = unsafe { std::slice::from_raw_parts(array.buffers, 2) }; + let mut values = vec![T::default(); offset + len]; + // SAFETY: The synchronized data buffer contains offset + len values of T and remains live. + unsafe { result::memcpy_dtoh_sync(&mut values, buffers[1] as u64) } + .map_err(|error| vortex_err!("copying Arrow values: {error}"))?; + + let validity = if buffers[0].is_null() { + assert_eq!(array.null_count, 0); + match nullability { + Nullability::NonNullable => Validity::NonNullable, + Nullability::Nullable => Validity::AllValid, + } + } else { + let mut bytes = vec![0u8; (offset + len).div_ceil(8)]; + // SAFETY: The synchronized bitmap covers offset + len bits and remains live. + unsafe { result::memcpy_dtoh_sync(&mut bytes, buffers[0] as u64) } + .map_err(|error| vortex_err!("copying Arrow validity: {error}"))?; + let bits = BitBuffer::new_with_offset(ByteBuffer::from(bytes), len, offset); + assert_eq!(array.null_count, i64::try_from(len - bits.true_count())?); + match nullability { + Nullability::NonNullable => { + assert_eq!(array.null_count, 0); + Validity::NonNullable + } + Nullability::Nullable => Validity::from(bits), + } + }; + Ok(PrimitiveArray::new(Buffer::from(values).slice(offset..), validity).into_array()) } -fn batch_lengths(stream: &mut ArrowDeviceArrayStream) -> Vec { +/// Read the fixture's projected Int64/UInt32 columns through the public Arrow Device ABI. +/// +/// # Safety +/// `array` must be a live batch from the fixture's projected stream, with its schema checked. +unsafe fn read_projected_batch(array: &ArrowDeviceArray) -> VortexResult { + assert_eq!(array.device_type, ARROW_DEVICE_CUDA); + let context = CudaContext::new(usize::try_from(array.device_id)?) + .map_err(|error| vortex_err!("opening Arrow device context: {error}"))?; + context + .bind_to_thread() + .map_err(|error| vortex_err!("binding Arrow device context: {error}"))?; + if !array.sync_event.is_null() { + // SAFETY: Arrow's CUDA sync_event points to a live event handle owned by this batch. + unsafe { result::event::synchronize(*array.sync_event.cast::()) } + .map_err(|error| vortex_err!("waiting for Arrow device batch: {error}"))?; + } + let array = &array.array; + assert!(array.dictionary.is_null()); + assert_eq!(array.offset, 0); + assert_eq!(array.null_count, 0); + assert_eq!(array.n_children, 2); + assert!(!array.children.is_null()); + // SAFETY: The live struct owns two children matching the previously checked schema. + let fields = unsafe { + let values = (*array.children).as_ref().expect("missing values child"); + let ids = (*array.children.add(1)) + .as_ref() + .expect("missing ids child"); + assert_eq!(values.length, array.length); + assert_eq!(ids.length, array.length); + vec![ + read_primitive::(values, Nullability::Nullable)?, + read_primitive::(ids, Nullability::NonNullable)?, + ] + }; + Ok(StructArray::try_new( + ["値.x", "ids"].into(), + fields, + usize::try_from(array.length)?, + Validity::NonNullable, + )? + .into_array()) +} + +fn read_projected_batches(stream: &mut ArrowDeviceArrayStream) -> VortexResult> { let get_next = stream.get_next.expect("missing get_next"); - let mut lengths = Vec::new(); + let mut batches = Vec::new(); loop { let mut array = empty_device_array(); // SAFETY: This live stream owns the callback; array is writable. - assert_eq!( - unsafe { get_next(stream, &raw mut array) }, - 0, - "{}", - stream_error(stream) - ); + let status = unsafe { get_next(stream, &raw mut array) }; + vortex_ensure!(status == 0, "get_next failed: {}", stream_error(stream)); if array.array.release.is_none() { break; } - assert_eq!(array.device_type, ARROW_DEVICE_CUDA); - assert_eq!(array.array.n_children, 2); - lengths.push(array.array.length); + // SAFETY: The fixture's schema was checked, and this batch remains live during readback. + let batch = unsafe { read_projected_batch(&array) }; release_device_array(&mut array); + batches.push(batch?); } - lengths + Ok(batches) } -#[cuda_test] -fn test_projection_gpu_local_file_schema_and_batch_boundaries() -> VortexResult<()> { - for (block_rows, batch_rows) in [(0, 2), (2, 3)] { - let session = session().with_some(CudaSession::try_default()?); - let mut file = NamedTempFile::new()?; - file.write_all(&file_bytes( - &session, - table()?.into_array(), - Some(block_rows), - )?)?; - let path = file - .path() - .to_str() - .ok_or_else(|| vortex_err!("non-UTF-8 test path"))?; - let options = vx_cuda_scan_options { - batch_rows, - ..Default::default() - }; - let mut stream = open_stream(&session, path, &options); - // The stream must retain its session state after the caller releases its session. - drop(session); - let mut schema = stream_schema(&mut stream); - let expected_fields = vec![ - Field::new("値.x", DataType::Int64, true), - Field::new("ids", DataType::UInt32, false), - ]; - assert_eq!(Schema::try_from(&schema)?, Schema::new(expected_fields)); - assert_eq!(batch_lengths(&mut stream), [2, 2, 1]); - let release = stream.release.expect("missing release"); - // SAFETY: Both objects are live and released exactly once. - unsafe { - release_schema(&mut schema); - release(&raw mut stream); - } - } +fn check_projected_file(block_rows: usize, batch_rows: usize) -> VortexResult<()> { + let session = session().with_some(CudaSession::try_default()?); + register_cuda_layout(&session); + let input = table()?; + let columns = ["値.x", "ids"]; + let expected = input.project(names(&columns)?.as_ref())?.into_array(); + let mut file = NamedTempFile::new()?; + file.write_all(&file_bytes( + &session, + input.into_array(), + cuda_write_strategy(&session, block_rows), + )?)?; + let path = file + .path() + .to_str() + .ok_or_else(|| vortex_err!("non-UTF-8 test path"))?; + let options = vx_cuda_scan_options { + batch_rows, + ..Default::default() + }; + let mut stream = open_stream(&session, path, &options, &columns); + // The stream must retain its session state after the caller releases its session. + drop(session); + let schema = stream_schema(&mut stream); + let expected_fields = vec![ + Field::new("値.x", DataType::Int64, true), + Field::new("ids", DataType::UInt32, false), + ]; + assert_eq!(Schema::try_from(&schema)?, Schema::new(expected_fields)); + let batches = read_projected_batches(&mut stream); + let release = stream.release.expect("missing release"); + // SAFETY: The stream is live and released exactly once, including on readback errors. + unsafe { release(&raw mut stream) }; + let batches = batches?; + let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); + assert_eq!(lengths, [2, 2, 1]); + let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); + assert_arrays_eq!( + actual, + expected, + &mut VortexSession::default().create_execution_ctx() + ); Ok(()) } + +#[cuda_test] +fn test_projection_gpu_subdivides_large_blocks() -> VortexResult<()> { + check_projected_file(0, 2) +} + +#[cuda_test] +fn test_projection_gpu_preserves_small_block_boundaries() -> VortexResult<()> { + check_projected_file(2, 3) +} diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 1b1a3f1bcf6..6646d0740fc 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -33,6 +33,7 @@ pub use dtype::vx_dtype; pub use error::try_or; pub use error::vx_error; pub use error::vx_error_free; +pub use error::vx_error_message; pub use log::vx_log_level; pub use scan::vx_partition; pub use scan::vx_partition_into_array_stream; From 8a36a1c1098981a94ddb666d4b3c94e402a677d7 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:10 +0000 Subject: [PATCH 11/35] fix(cuda): publish formatted headers only on change and enforce drift in CI Address review item 4. Signed-off-by: Alexander Droste --- .github/workflows/cuda.yaml | 16 +++++++ vortex-cuda/ffi/README.md | 21 +++++++-- vortex-cuda/ffi/build.rs | 92 +++++++++++++++++++++++++++++++++---- 3 files changed, 118 insertions(+), 11 deletions(-) diff --git a/.github/workflows/cuda.yaml b/.github/workflows/cuda.yaml index 4e14ea370a5..d27bc7b1452 100644 --- a/.github/workflows/cuda.yaml +++ b/.github/workflows/cuda.yaml @@ -45,6 +45,9 @@ jobs: - "vortex-test/**" - "pyproject.toml" - "uv.lock" + - "Cargo.toml" + - "Cargo.lock" + - ".clang-format" - ".github/workflows/**" cuda-build-lint: @@ -65,12 +68,20 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} enable-sccache: "true" + - name: Install clang-format + run: | + # Match the Ubuntu 24.04 formatter used by the C/C++ lint job. + sudo apt-get update + sudo apt-get install -y clang-format-18 + echo "/usr/lib/llvm-18/bin" >> "$GITHUB_PATH" - uses: ./.github/actions/check-rebuild with: command: >- cargo build --profile ci --locked --all-features --all-targets -p vortex-cuda -p vortex-cuda-ffi -p vortex-cub -p vortex-nvcomp -p gpu-scan-cli -p vortex-test-e2e-cuda -p vortex-python-cuda + - name: Verify generated CUDA FFI header is up to date + run: git --no-pager diff --exit-code -- vortex-cuda/ffi/cinclude/vortex_cuda.h - name: Clippy CUDA crates run: | cargo clippy --profile ci --locked --all-features --all-targets \ @@ -111,6 +122,11 @@ jobs: uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 with: sync: false + - name: Install clang-format + run: | + sudo apt-get update + sudo apt-get install -y clang-format-18 + echo "/usr/lib/llvm-18/bin" >> "$GITHUB_PATH" - name: Install nextest uses: taiki-e/install-action@3f74d7c16a4242f1c95561e98edc25d36adb4375 # v2 with: diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 5064d0a9035..96bd030f4ae 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -28,7 +28,22 @@ On Linux, use `vx_cuda_scan_path_arrow_device_stream_with_options` with `vx_cuda_scan_options.flags = VX_CUDA_SCAN_FLAG_DIRECT_IO` to bypass the operating system page cache for pooled data-plane reads. Footer and zone-map reads remain buffered on the host. +## Generated header + `build.rs` generates `cinclude/vortex_cuda.h` with cbindgen on stable Rust, without macro -expansion. Edit the API and docs in `src/lib.rs`, not the generated header; commit -regenerated headers with API changes. `cbindgen.toml` supplies the standard Arrow Device -interface compatibility preamble. +expansion or compiling the CUDA implementation. Edit the API and docs in `src/lib.rs`, not +the generated header; commit regenerated headers with API changes. `cbindgen.toml` supplies +the standard Arrow Device interface compatibility preamble. + +Header generation requires `clang-format` on `PATH`, with support for the repository's +`.clang-format` configuration. The build script generates bytes in memory, then formats +through clang-format's stdin/stdout using `--style=file` and +`--assume-filename=cinclude/vortex_cuda.h` from this crate's directory. A missing or failing +formatter fails the build before changing the existing header; there is no unformatted +fallback. + +Only the fully formatted bytes are compared with the committed header. Identical output +leaves the file and its timestamp untouched. Changed output is written to a uniquely created +temporary file in the header's directory and published by atomic rename, so concurrent +builds do not expose a truncated header. The build script tracks changes to `src/`, +`cbindgen.toml`, `build.rs`, and the repository configuration at `../../.clang-format`. diff --git a/vortex-cuda/ffi/build.rs b/vortex-cuda/ffi/build.rs index ab74d21cfca..4fe2033d382 100644 --- a/vortex-cuda/ffi/build.rs +++ b/vortex-cuda/ffi/build.rs @@ -2,28 +2,104 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::error::Error; +use std::fs; +use std::fs::OpenOptions; +use std::io; +use std::io::Write; +use std::process; use std::process::Command; +use std::process::Stdio; +use std::thread; fn main() -> Result<(), Box> { println!("cargo:rerun-if-changed=src"); println!("cargo:rerun-if-changed=cbindgen.toml"); println!("cargo:rerun-if-changed=build.rs"); + println!("cargo:rerun-if-changed=../../.clang-format"); let header = "cinclude/vortex_cuda.h"; + let mut generated = Vec::new(); // The CUDA API needs no macro expansion or dependency parsing, so generate on stable Rust // without recursively building the CUDA implementation. cbindgen::Builder::new() .with_src("src/lib.rs") .with_config(cbindgen::Config::from_file("cbindgen.toml")?) .generate()? - .write_to_file(header); - if !Command::new("clang-format") - .args(["--style=file", "-i"]) - .arg(header) - .status() - .is_ok_and(|status| status.success()) - { - println!("cargo:warning=clang-format unavailable or failed; CUDA header left unformatted"); + .write(&mut generated); + let formatted = format_header(header, generated)?; + match fs::read(header) { + Ok(existing) if existing == formatted => return Ok(()), + Ok(_) => {} + Err(error) if error.kind() == io::ErrorKind::NotFound => {} + Err(error) => return Err(format!("failed to read {header}: {error}").into()), } + publish_header(header, &formatted) + .map_err(|error| format!("failed to publish {header}: {error}"))?; Ok(()) } + +fn format_header(header: &str, generated: Vec) -> Result, Box> { + // Cargo runs this script in the crate directory; the assumed header path lets clang-format + // find the repository's .clang-format while reading from stdin. + let mut child = Command::new("clang-format") + .arg("--style=file") + .arg(format!("--assume-filename={header}")) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .map_err(|error| { + format!("failed to start clang-format (required on PATH to generate {header}): {error}") + })?; + let mut stdin = child + .stdin + .take() + .ok_or("clang-format stdin is unavailable")?; + // Feed stdin while wait_with_output drains stdout and stderr, even when a header exceeds + // pipe capacity. Dropping stdin in the writer signals EOF to clang-format. + let writer = thread::spawn(move || stdin.write_all(&generated)); + let output = child.wait_with_output(); + let written = writer + .join() + .map_err(|_| "clang-format stdin writer panicked")?; + let output = + output.map_err(|error| format!("failed to collect clang-format output: {error}"))?; + if !output.status.success() { + return Err(format!( + "clang-format failed for {header} ({}): {}", + output.status, + String::from_utf8_lossy(&output.stderr) + ) + .into()); + } + written.map_err(|error| format!("failed to write header to clang-format stdin: {error}"))?; + Ok(output.stdout) +} + +fn publish_header(header: &str, formatted: &[u8]) -> io::Result<()> { + let mut attempt = 0_u64; + let (temporary, mut file) = loop { + let temporary = format!("{header}.{}.{attempt}.tmp", process::id()); + match OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary) + { + Ok(file) => break (temporary, file), + Err(error) if error.kind() == io::ErrorKind::AlreadyExists => attempt += 1, + Err(error) => return Err(error), + } + }; + + // A unique sibling file and atomic rename prevent concurrent builds from exposing a + // truncated shared header. Close the file before renaming or cleaning it up. + let written = file.write_all(formatted); + drop(file); + let result = written.and_then(|()| fs::rename(&temporary, header)); + if result.is_err() + && let Err(error) = fs::remove_file(&temporary) + { + eprintln!("failed to remove temporary header {temporary}: {error}"); + } + result +} From 8ad7a2033d198572d43e45076ea21b66216b8ad4 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:10 +0000 Subject: [PATCH 12/35] fix(cuda): register layouts once without replacing edition policy Address review item 5. Signed-off-by: Alexander Droste --- vortex-cuda/src/layout.rs | 165 ++++++++++++++++++++++++++++++++------ 1 file changed, 140 insertions(+), 25 deletions(-) diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 86da0a921ab..bd9908b3da6 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -7,13 +7,13 @@ use std::any::Any; use std::ops::BitAnd; use std::ops::Range; use std::sync::Arc; +use std::sync::Once; use std::sync::OnceLock; use async_trait::async_trait; use futures::FutureExt; use futures::StreamExt; use futures::future::BoxFuture; -use parking_lot::Mutex; use vortex::array::ArrayRef; use vortex::array::ArrayVTable; use vortex::array::MaskFuture; @@ -68,6 +68,8 @@ use vortex::scalar::Scalar; use vortex::scalar::ScalarTruncation; use vortex::scalar::lower_bound; use vortex::scalar::upper_bound; +use vortex::session::SessionExt; +use vortex::session::SessionVar; use vortex::session::VortexSession; use vortex::session::registry::CachedId; use vortex::session::registry::ReadContext; @@ -542,6 +544,24 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { result } +#[derive(Clone, Debug, Default)] +struct CudaLayoutRegistration(Arc); + +impl SessionVar for CudaLayoutRegistration { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +const CUDA_EDITION_FAMILY: EditionFamily = EditionFamily { + name: "cuda", + origin: "vortex-cuda", + doc: "CUDA-readable layouts, enabled only when CUDA layout support is registered.", +}; const CUDA_EDITION: EditionId = EditionId::new("cuda", 2026, 9, 0); static CUDA_EDITION_DECLARATION: EditionDeclaration = EditionDeclaration { edition: Edition { @@ -551,38 +571,45 @@ static CUDA_EDITION_DECLARATION: EditionDeclaration = EditionDeclaration { added: &[EditionMember::layout(&"vortex.cuda_flat")], }; -/// Register the [`CudaFlatLayoutEncoding`] and enable its draft `cuda` edition for writing. +/// Register [`CudaFlat`] and its draft `cuda` edition once per session. /// -/// Other edition selections and checks are unchanged. The draft has no cross-version -/// compatibility guarantee; readers must register the CUDA layout. +/// A newly registered edition is enabled for writing only if no `cuda` edition is selected. +/// A pre-registered edition and subsequent calls leave writer policy unchanged. Other edition +/// selections and checks are unchanged. The draft has no cross-version compatibility guarantee; +/// readers must register the CUDA layout. /// /// Call this alongside [`crate::initialize_cuda`] when setting up a CUDA-enabled session. /// Registration itself does not require a GPU. pub fn register_cuda_layout(session: &VortexSession) { - session - .layouts() - .register(LayoutEncodingRef::new_ref(&CudaFlat)); - - // Concurrent CUDA FFI calls may register session clones; serialize the check and registration - // because edition declarations reject duplicates. - static REGISTRATION_LOCK: Mutex<()> = Mutex::new(()); - let _guard = REGISTRATION_LOCK.lock(); - if session.editions().find(&CUDA_EDITION).is_none() { + // Edition declarations publish the edition before its members. All callers must wait for + // initialization to finish rather than treating an unlocked edition lookup as completion. + session.get::().0.call_once(|| { session - .editions() - .declare_family(&EditionFamily { - name: "cuda", - origin: "vortex-cuda", - doc: "CUDA-readable layouts, enabled only when CUDA layout support is registered.", - }) - .vortex_expect("CUDA edition family is valid"); + .layouts() + .register(LayoutEncodingRef::new_ref(&CudaFlat)); + if session.editions().find(&CUDA_EDITION).is_some() { + return; + } + if session.editions().find_family("cuda").is_none() { + session + .editions() + .declare_family(&CUDA_EDITION_FAMILY) + .vortex_expect("CUDA edition family is valid"); + } session .register_edition(&CUDA_EDITION_DECLARATION) .vortex_expect("CUDA edition declaration is valid"); - } - session - .enable_edition(CUDA_EDITION) - .vortex_expect("CUDA edition is registered"); + if !session + .enabled_editions() + .editions() + .iter() + .any(|edition| edition.family == CUDA_EDITION.family) + { + session + .enable_edition(CUDA_EDITION) + .vortex_expect("CUDA edition is registered"); + } + }); } #[cfg(test)] @@ -632,7 +659,14 @@ mod tests { std::thread::scope(|scope| { for _ in 0..4 { let session = session.clone(); - scope.spawn(move || register_cuda_layout(&session)); + scope.spawn(move || { + register_cuda_layout(&session); + assert!( + session + .enabled_component_ids(ComponentKind::Layout) + .contains(&CudaFlat.id()) + ); + }); } }); register_cuda_layout(&session); @@ -652,6 +686,87 @@ mod tests { Ok(()) } + #[rstest] + fn test_cuda_registration_preserves_selected_cuda_edition( + #[values(false, true)] register_first: bool, + ) -> VortexResult<()> { + const OTHER_CUDA_EDITION: EditionId = EditionId::new("cuda", 2026, 8, 0); + let session = VortexSession::default(); + if register_first { + register_cuda_layout(&session); + } else { + session.editions().declare_family(&CUDA_EDITION_FAMILY)?; + } + session.register_edition(&EditionDeclaration { + edition: Edition { + id: OTHER_CUDA_EDITION, + min_library_version: None, + }, + added: &[], + })?; + session.enable_edition(OTHER_CUDA_EDITION)?; + let mut expected_editions = session.enabled_editions().editions(); + expected_editions.sort_unstable(); + let expected_layouts = session.enabled_component_ids(ComponentKind::Layout); + assert!(!expected_layouts.contains(&CudaFlat.id())); + + std::thread::scope(|scope| { + for _ in 0..4 { + let session = session.clone(); + let expected_editions = &expected_editions; + let expected_layouts = &expected_layouts; + scope.spawn(move || { + register_cuda_layout(&session); + let mut enabled_editions = session.enabled_editions().editions(); + enabled_editions.sort_unstable(); + assert_eq!(&enabled_editions, expected_editions); + assert_eq!( + &session.enabled_component_ids(ComponentKind::Layout), + expected_layouts + ); + }); + } + }); + register_cuda_layout(&session); + assert!(session.editions().find(&CUDA_EDITION).is_some()); + assert!( + session + .enabled_editions() + .editions() + .contains(&OTHER_CUDA_EDITION) + ); + session.editions().validate()?; + Ok(()) + } + + #[rstest] + fn test_cuda_registration_preserves_pre_registered_edition_policy( + #[values(false, true)] enabled: bool, + ) -> VortexResult<()> { + let session = VortexSession::default(); + session.editions().declare_family(&CUDA_EDITION_FAMILY)?; + session.register_edition(&CUDA_EDITION_DECLARATION)?; + if enabled { + session.enable_edition(CUDA_EDITION)?; + } + let mut expected_editions = session.enabled_editions().editions(); + expected_editions.sort_unstable(); + let expected_layouts = session.enabled_component_ids(ComponentKind::Layout); + + register_cuda_layout(&session); + register_cuda_layout(&session); + + let mut enabled_editions = session.enabled_editions().editions(); + enabled_editions.sort_unstable(); + assert_eq!(enabled_editions, expected_editions); + assert_eq!( + session.enabled_component_ids(ComponentKind::Layout), + expected_layouts + ); + session.editions().validate()?; + Ok(()) + } + #[test] fn test_registry_alone_does_not_permit_cuda_flat() -> VortexResult<()> { let runtime = CurrentThreadRuntime::new(); From 984759bbbd64ac2bc579b1554d536b803eaa951e Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:10 +0000 Subject: [PATCH 13/35] refactor(cuda): reuse decoded stream schema initialization Address review item 7. Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/mod.rs | 83 +++++++++++++++++------------------- 1 file changed, 40 insertions(+), 43 deletions(-) diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index fe7243b4f6f..6e999f079c0 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -414,53 +414,50 @@ impl DeviceArrayStreamPrivateData { array.dtype() ); - let mut staged_schema = None; - let (mut device_array, ffi_schema) = - if self.ctx.cuda_session().dictionary_export() == DictionaryExport::Decode { - // Decode schemas depend only on dtype and fixed context settings. Stage before - // export to preserve error ordering, but cache only after device validation. - if self.schema.is_none() { - staged_schema = Some(ArrowDeviceStreamSchema::from_dtype( - &self.dtype, - &mut self.ctx, - )?); - } - let device_array = self - .runtime - .block_on(array.export_device_array(&mut self.ctx))?; - (device_array, None) - } else { - let exported = self - .runtime - .block_on(array.export_device_array_with_schema(&mut self.ctx))?; - (exported.array, Some(exported.schema)) - }; - - // Schemas release themselves on drop; rejected device arrays need explicit release. - let validation = (|| { - self.check_device(&device_array)?; - if let Some(ffi_schema) = &ffi_schema { - let exported_schema = ArrowDeviceStreamSchema::from_ffi(ffi_schema, &self.dtype)?; - if let Some(stream_schema) = &self.schema { - vortex_ensure!( - stream_schema == &exported_schema, - "stream array Arrow schema changed from {:?} to {:?}; an Arrow C device stream \ - requires every array to share one schema, so chunks must not vary their \ - encoding (for example a dictionary-encoded chunk among plain chunks)", - stream_schema, - exported_schema - ); - } - staged_schema = Some(exported_schema); + if self.ctx.cuda_session().dictionary_export() == DictionaryExport::Decode { + self.get_or_init_schema()?; + let mut device_array = self + .runtime + .block_on(array.export_device_array(&mut self.ctx))?; + if let Err(error) = self.check_device(&device_array) { + release_device_array(&mut device_array); + return Err(error); } - Ok(()) - })(); - if let Err(error) = validation { + return Ok(device_array); + } + + let ArrowDeviceArrayWithSchema { + schema: ffi_schema, + array: mut device_array, + } = self + .runtime + .block_on(array.export_device_array_with_schema(&mut self.ctx))?; + + // The FFI schema releases itself on drop; rejected arrays need explicit release. + if let Err(error) = self.check_device(&device_array) { release_device_array(&mut device_array); return Err(error); } - if self.schema.is_none() { - self.schema = staged_schema; + let exported_schema = match ArrowDeviceStreamSchema::from_ffi(&ffi_schema, &self.dtype) { + Ok(schema) => schema, + Err(error) => { + release_device_array(&mut device_array); + return Err(error); + } + }; + if let Some(stream_schema) = &self.schema { + if stream_schema != &exported_schema { + release_device_array(&mut device_array); + return Err(vortex_err!( + "stream array Arrow schema changed from {:?} to {:?}; an Arrow C device stream \ + requires every array to share one schema, so chunks must not vary their \ + encoding (for example a dictionary-encoded chunk among plain chunks)", + stream_schema, + exported_schema + )); + } + } else { + self.schema = Some(exported_schema); } Ok(device_array) } From 8ca1a99d964172fa01cbbc6ad08ffea2a383fd9f Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:21 +0000 Subject: [PATCH 14/35] test(cuda): share exported device buffer readback helper Address review item 11. Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/canonical.rs | 12 +----------- vortex-cuda/src/arrow/dictionary_tests.rs | 13 +------------ vortex-cuda/src/arrow/mod.rs | 18 ++++++++++++++++++ 3 files changed, 20 insertions(+), 23 deletions(-) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index b441bdfcfb5..a57528af765 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -1518,6 +1518,7 @@ mod tests { use crate::arrow::canonical::export_arrow_validity_buffer; use crate::arrow::canonical::repack_arrow_validity_buffer; use crate::arrow::dictionary_tests::upload; + use crate::arrow::tests::private_data_buffer_bytes; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::device_buffer::cuda_backing_allocation; use crate::executor::CudaArrayExt; @@ -1777,17 +1778,6 @@ mod tests { Ok(Buffer::::from_byte_buffer(private_data_buffer_bytes(array, buffer_idx)?).to_vec()) } - fn private_data_buffer_bytes( - array: &ArrowArray, - buffer_idx: usize, - ) -> VortexResult { - let private_data = unsafe { &*array.private_data.cast::() }; - let buffer = private_data.buffers[buffer_idx] - .as_ref() - .vortex_expect("buffer should be present"); - Ok(buffer.to_host_sync()) - } - // Assert Arrow Binary export uses the standard null bitmap, i32 offsets, and values layout. fn assert_binary_layout( array: &ArrowArray, diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index ca7eba4d3ee..c39fb7e286a 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -21,10 +21,10 @@ use vortex::array::stream::ArrayStreamExt; use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; -use vortex::buffer::ByteBuffer; use vortex::error::vortex_bail; use super::tests::last_error; +use super::tests::private_data_buffer_bytes as buffer; use super::*; use crate::CudaSession; @@ -50,17 +50,6 @@ pub(super) fn upload(array: ArrayRef, ctx: &mut CudaExecutionCtx) -> VortexResul unsafe { array.with_slots(slots.into())?.with_buffers(buffers) } } -/// Copy a device buffer from a live, unreleased array produced by this exporter to the host. -fn buffer(array: &ArrowArray, index: usize) -> VortexResult { - // SAFETY: Only called on live arrays produced by our exporter, before their release. - let private = unsafe { &*array.private_data.cast::() }; - let buffer = private.buffers[index] - .as_ref() - .ok_or_else(|| vortex_err!("missing exported buffer {index}"))?; - buffer.cuda_device_ptr()?; - buffer.try_to_host_sync() -} - /// Rebuild supported zero-offset, dictionary-free exports as host arrays for comparison. /// Requires live arrays from this exporter and their matching logical dtype. fn read_plain(array: &ArrowArray, dtype: &DType) -> VortexResult { diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index 6e999f079c0..d72e2221be2 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -929,6 +929,7 @@ mod tests { use vortex::array::arrays::PrimitiveArray; use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; + use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; @@ -938,16 +939,33 @@ mod tests { use vortex::session::VortexSession; use vortex_cuda_macros::test as cuda_test; + use crate::CudaBufferExt; use crate::CudaSession; use crate::arrow::ARROW_DEVICE_CUDA; + use crate::arrow::ArrowArray; use crate::arrow::ArrowDeviceArray; use crate::arrow::ArrowDeviceArrayStream; use crate::arrow::ArrowSchema; use crate::arrow::DeviceArrayStreamExt; use crate::arrow::LIBC_EINVAL; + use crate::arrow::PrivateData; use crate::arrow::release_device_array; use crate::arrow::release_schema; + /// Copy a CUDA buffer from a live, unreleased array produced by this exporter to the host. + pub(super) fn private_data_buffer_bytes( + array: &ArrowArray, + index: usize, + ) -> VortexResult { + // SAFETY: Only called on live arrays produced by our exporter, before their release. + let private = unsafe { &*array.private_data.cast::() }; + let buffer = private.buffers[index] + .as_ref() + .ok_or_else(|| vortex_err!("missing exported buffer {index}"))?; + buffer.cuda_device_ptr()?; + buffer.try_to_host_sync() + } + pub(super) fn last_error(stream: &mut ArrowDeviceArrayStream) -> VortexResult { let get_last_error = stream .get_last_error From 7e7db60c9170391e7808973e9b293220f58acf90 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:21 +0000 Subject: [PATCH 15/35] refactor(cuda): expose and reuse the empty device array constructor Address review item 10. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/lib.rs | 16 +++------------- vortex-cuda/ffi/src/tests/projection.rs | 2 +- vortex-cuda/src/arrow/mod.rs | 9 +++++---- 3 files changed, 9 insertions(+), 18 deletions(-) diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 9a206e79ba3..93b3f28a663 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -622,7 +622,7 @@ mod tests { let get_next = stream.get_next.expect("missing get_next"); let release = stream.release.expect("missing release"); let schema = stream_schema(&mut stream); - let mut exported = empty_device_array(); + let mut exported = ArrowDeviceArray::empty(); // SAFETY: The live stream owns the callback, and the output is writable. unsafe { assert_eq!(get_next(&raw mut stream, &raw mut exported), 0); @@ -676,16 +676,6 @@ mod tests { schema } - fn empty_device_array() -> ArrowDeviceArray { - ArrowDeviceArray { - array: vortex_cuda::arrow::ArrowArray::empty(), - device_id: 0, - device_type: 0, - sync_event: ptr::null_mut(), - reserved: [0; 3], - } - } - /// # Safety /// `session` and `array` must be valid borrowed FFI handles for the duration of the call. unsafe fn export_array( @@ -694,7 +684,7 @@ mod tests { ) -> (FFI_ArrowSchema, ArrowDeviceArray) { let mut error = ptr::null_mut(); let mut schema = FFI_ArrowSchema::empty(); - let mut device_array = empty_device_array(); + let mut device_array = ArrowDeviceArray::empty(); // SAFETY: The caller guarantees valid handles; all outputs are live and writable. let status = unsafe { vx_cuda_array_export_arrow_device( @@ -805,7 +795,7 @@ mod tests { let session = test_session(VortexSession::default()); let array = test_array(PrimitiveArray::from_iter(0u32..5)); let mut schema = FFI_ArrowSchema::empty(); - let mut device_array = empty_device_array(); + let mut device_array = ArrowDeviceArray::empty(); let mut error = ptr::null_mut(); let status = unsafe { diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index 3e747d5761d..91b403f0ffa 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -438,7 +438,7 @@ fn read_projected_batches(stream: &mut ArrowDeviceArrayStream) -> VortexResult Self { + /// Create an empty, released array with zeroed device metadata and a null sync event. + /// + /// Use this as storage for an Arrow C device callback output or as the basis for an + /// end-of-stream marker. No CUDA device is selected; `device_id` and `device_type` are zero. + pub fn empty() -> Self { Self { array: ArrowArray::empty(), device_id: 0, From 03fad7d5eb144d113ae770d519c16ed30738c809 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:22 +0000 Subject: [PATCH 16/35] docs(cuda): explain safety of FFI pointer operations Address review item 12. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/lib.rs | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 93b3f28a663..245f0c8bffd 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -150,6 +150,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( dtype: *const vx_dtype, error_out: *mut *mut vx_error, ) -> *mut vx_array_sink { + // SAFETY: The forwarded pointers satisfy the same requirements as this wrapper. unsafe { vx_cuda_array_sink_open_file_block_rows(session, path, dtype, 0, error_out) } } @@ -174,7 +175,9 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( error_out: *mut *mut vx_error, ) -> *mut vx_array_sink { try_or(error_out, ptr::null_mut(), || { + // SAFETY: The caller supplies a live borrowed session handle. let vortex_session = session_with_cuda(unsafe { vx_session_ref(session) }?); + // SAFETY: All borrowed inputs satisfy the underlying sink's requirements. unsafe { vx_array_sink_open_file_with_strategy( session, @@ -211,6 +214,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( out_stream: *mut ArrowDeviceArrayStream, error_out: *mut *mut vx_error, ) -> c_int { + // SAFETY: The forwarded pointers satisfy this wrapper's requirements; null options is valid. unsafe { vx_cuda_scan_path_arrow_device_stream_with_options( session, @@ -243,6 +247,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows batch_rows, ..Default::default() }; + // SAFETY: The caller supplies valid pointers; the local options remain live during the call. unsafe { vx_cuda_scan_path_arrow_device_stream_with_options( session, @@ -310,11 +315,17 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_projected( try_or(error_out, VX_CUDA_ERR, || { vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); - // SAFETY: The caller keeps the borrowed options and column views alive for this call. - let options = unsafe { scan_options(options) }?; - let columns = unsafe { scan_columns(columns, ncolumns) }?; - let path = unsafe { path.as_str() }?; - let session = session_with_cuda(unsafe { vx_session_ref(session) }?); + // SAFETY: The caller keeps options, column views and their bytes, path bytes, and the + // borrowed session handle valid for this call. + let (options, columns, path, session) = unsafe { + ( + scan_options(options)?, + scan_columns(columns, ncolumns)?, + path.as_str()?, + vx_session_ref(session)?, + ) + }; + let session = session_with_cuda(session); let file = ffi_runtime().block_on( session .open_options() @@ -327,6 +338,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_projected( let ctx = scan_export_ctx(session)?; let device_stream = ArrowDeviceArrayStream::new(array_stream, ctx, ffi_runtime()); + // SAFETY: The output is non-null and the caller guarantees writable storage. unsafe { ptr::write(out_stream, device_stream) }; Ok(VX_CUDA_OK) }) @@ -469,12 +481,15 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_export_arrow_device( vortex_ensure!(!out_schema.is_null(), "null ArrowSchema output"); vortex_ensure!(!out_array.is_null(), "null ArrowDeviceArray output"); + // SAFETY: The caller supplies a live borrowed session handle. let session = session_with_cuda(unsafe { vx_session_ref(session) }?); + // SAFETY: The caller supplies a live borrowed array handle. let array = unsafe { vx_array_ref(array) }?.clone(); let mut ctx = CudaSession::create_execution_ctx(session)?; let exported = futures::executor::block_on(array.export_device_array_with_schema(&mut ctx))?; + // SAFETY: Both outputs are non-null and the caller guarantees writable storage. unsafe { ptr::write(out_schema, exported.schema); ptr::write(out_array, exported.array); @@ -505,13 +520,16 @@ pub unsafe extern "C-unwind" fn vx_cuda_partition_scan_arrow_device_stream( try_or(error_out, VX_CUDA_ERR, || { vortex_ensure!(!partition.is_null(), "null vx_partition"); + // SAFETY: The caller transfers ownership of this non-null partition handle. let array_stream = unsafe { vx_partition_into_array_stream(partition) }?; vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); + // SAFETY: The caller supplies a live borrowed session handle. let session = session_with_cuda(unsafe { vx_session_ref(session) }?); // Drive the stream on the same runtime the partition's scan spawned its work onto. let device_stream = array_stream.export_device_array_stream(session, ffi_runtime())?; + // SAFETY: The output is non-null and the caller guarantees writable storage. unsafe { ptr::write(out_stream, device_stream) }; Ok(VX_CUDA_OK) }) From 4bcf47039598737ccbef7640ba558b76d3d31b6d Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:22 +0000 Subject: [PATCH 17/35] test(cuda): use the required test function prefix Address review item 13. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 245f0c8bffd..efa32680973 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -568,7 +568,7 @@ mod tests { use super::*; #[test] - fn scan_options_default_to_buffered_io() -> VortexResult<()> { + fn test_scan_options_default_to_buffered_io() -> VortexResult<()> { let options = vx_cuda_scan_options::default(); assert_eq!(options.flags, 0); assert_eq!(options.batch_rows, 0); @@ -582,7 +582,7 @@ mod tests { } #[test] - fn maps_scan_options_and_ignores_unknown_flags() -> VortexResult<()> { + fn test_maps_scan_options() -> VortexResult<()> { let buffered = PooledFileReadAtOptions::default(); for (flags, batch_rows, read_at_options) in [ (0, 8192, buffered), @@ -606,7 +606,7 @@ mod tests { } #[cuda_test] - fn scan_decodes_dictionaries_and_reuses_session_resources() -> VortexResult<()> { + fn test_scan_decodes_dictionaries_and_reuses_session_resources() -> VortexResult<()> { // A distinct allocator identity detects accidental reconstruction of a default session. let allocator = BufferAllocatorRef::new(StaticBufferAllocator); let session = VortexSession::default() From f88635701ece271bac90ddba06c9be65e4e03b14 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:48 +0000 Subject: [PATCH 18/35] test(cuda): release projected streams on assertion and error paths Address review item 15. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/tests/projection.rs | 27 ++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index 91b403f0ffa..fec5578b725 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -307,12 +307,23 @@ unsafe fn take_error_message(error: *mut vx_error) -> Option { Some(message) } +struct OwnedDeviceStream(ArrowDeviceArrayStream); + +impl Drop for OwnedDeviceStream { + fn drop(&mut self) { + if let Some(release) = self.0.release { + // SAFETY: This owner holds the live stream and releases it exactly once. + unsafe { release(&raw mut self.0) }; + } + } +} + fn open_stream( session: &VortexSession, path: &str, options: &vx_cuda_scan_options, columns: &[&str], -) -> ArrowDeviceArrayStream { +) -> OwnedDeviceStream { let mut output = MaybeUninit::::uninit(); let mut error = ptr::null_mut(); let handle = test_session(session.clone()); @@ -340,9 +351,11 @@ fn open_stream( "{}", message.as_deref().unwrap_or("no FFI error") ); - assert!(message.is_none(), "unexpected FFI error: {message:?}"); // SAFETY: A successful call initialized the stream, which owns its session state. - unsafe { output.assume_init() } + let stream = OwnedDeviceStream(unsafe { output.assume_init() }); + assert!(message.is_none(), "unexpected FFI error: {message:?}"); + assert!(stream.0.release.is_some(), "missing release"); + stream } /// # Safety @@ -476,17 +489,13 @@ fn check_projected_file(block_rows: usize, batch_rows: usize) -> VortexResult<() let mut stream = open_stream(&session, path, &options, &columns); // The stream must retain its session state after the caller releases its session. drop(session); - let schema = stream_schema(&mut stream); + let schema = stream_schema(&mut stream.0); let expected_fields = vec![ Field::new("値.x", DataType::Int64, true), Field::new("ids", DataType::UInt32, false), ]; assert_eq!(Schema::try_from(&schema)?, Schema::new(expected_fields)); - let batches = read_projected_batches(&mut stream); - let release = stream.release.expect("missing release"); - // SAFETY: The stream is live and released exactly once, including on readback errors. - unsafe { release(&raw mut stream) }; - let batches = batches?; + let batches = read_projected_batches(&mut stream.0)?; let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); assert_eq!(lengths, [2, 2, 1]); let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); From 4d0ccf2c6161f0c233ced69a4f8d5f8d99cd5266 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:48 +0000 Subject: [PATCH 19/35] refactor(cuda): let projection planning validate struct fields Address review item 8. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 9 ++++--- vortex-cuda/ffi/src/lib.rs | 18 ++++--------- vortex-cuda/ffi/src/tests/projection.rs | 34 +++++++++++++++++++++++++ 3 files changed, 44 insertions(+), 17 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index f0e13d29f8b..eefafcbf70f 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -191,10 +191,11 @@ int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session * Scan a local Vortex file with ordered top-level column projection. * * Same options, ownership, and file requirements as - * [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes column I/O/decoding. - * Names are literal and case-sensitive; unknown/duplicate names and non-struct files are rejected. - * `ncolumns == 0` ignores `columns` and selects all. Names are copied; empty files retain the - * projected schema. Errors leave `out_stream` unchanged. + * [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes decoding and skips + * unselected column I/O when the file layout stores columns separately. + * Names are literal and case-sensitive; a nonempty projection rejects unknown/duplicate names + * and non-struct files. `ncolumns == 0` ignores `columns` and selects all. Names are copied; + * empty files retain the projected schema. Errors leave `out_stream` unchanged. * * # Safety * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index efa32680973..29c0ace36a8 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -292,10 +292,11 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_optio /// Scan a local Vortex file with ordered top-level column projection. /// /// Same options, ownership, and file requirements as -/// [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes column I/O/decoding. -/// Names are literal and case-sensitive; unknown/duplicate names and non-struct files are rejected. -/// `ncolumns == 0` ignores `columns` and selects all. Names are copied; empty files retain the -/// projected schema. Errors leave `out_stream` unchanged. +/// [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes decoding and skips +/// unselected column I/O when the file layout stores columns separately. +/// Names are literal and case-sensitive; a nonempty projection rejects unknown/duplicate names +/// and non-struct files. `ncolumns == 0` ignores `columns` and selects all. Names are copied; +/// empty files retain the projected schema. Errors leave `out_stream` unchanged. /// /// # Safety /// @@ -389,15 +390,6 @@ fn projected_scan( ) -> VortexResult> { let mut scan = file.scan()?; if !columns.is_empty() { - let fields = file.dtype().as_struct_fields_opt().ok_or_else(|| { - vortex_err!("CUDA scan column projection requires a struct file dtype") - })?; - for name in columns.iter() { - vortex_ensure!( - fields.find(name).is_some(), - "unknown CUDA scan column: {name:?}" - ); - } let projection = select(columns, root()).optimize(file.dtype())?; scan = scan.with_projection(projection.bind(file.dtype())?); } diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index fec5578b725..a62648de6d9 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -160,6 +160,14 @@ fn open_file( .open_buffer(file_bytes(session, array, strategy)?) } +fn flat_ids_file(session: &VortexSession, rows: u32) -> VortexResult { + let ids = PrimitiveArray::from_iter(0..rows).into_array(); + let rows = ids.len(); + let input = + StructArray::try_new(["ids"].into(), vec![ids], rows, Validity::NonNullable)?.into_array(); + open_file(session, input, Arc::new(FlatLayoutStrategy::default())) +} + #[test] fn test_cuda_write_strategy_preserves_high_cardinality_row_blocks() -> VortexResult<()> { let session = session(); @@ -180,6 +188,32 @@ fn test_cuda_write_strategy_preserves_high_cardinality_row_blocks() -> VortexRes Ok(()) } +#[test] +fn test_projected_scan_rejects_unknown_field() -> VortexResult<()> { + let session = session(); + let file = flat_ids_file(&session, 5)?; + assert_error( + projected_scan(&file, names(&["missing"])?, 0), + "must be a subset of child fields", + ); + Ok(()) +} + +#[test] +fn test_projected_scan_rejects_nonstruct_projection() -> VortexResult<()> { + let session = session(); + let file = open_file( + &session, + PrimitiveArray::from_iter(0u32..5).into_array(), + Arc::new(FlatLayoutStrategy::default()), + )?; + assert_error( + projected_scan(&file, names(&["ids"])?, 0), + "Select child must return a struct dtype", + ); + Ok(()) +} + struct RejectSegments { inner: Arc, forbidden: Vec, From 41a96040ad5ecceb2096f1c4247661cc17f5e08d Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:48 +0000 Subject: [PATCH 20/35] fix(cuda): preserve layout boundaries when batch rows is zero Address review item 3. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 2 +- vortex-cuda/ffi/src/lib.rs | 12 +++++++----- vortex-cuda/ffi/src/tests/projection.rs | 13 ++++++++++++- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index eefafcbf70f..78775772f19 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -71,7 +71,7 @@ typedef struct vx_cuda_scan_options { */ uint32_t flags; /** - * Maximum rows in each output batch. Zero uses layout-derived splitting. + * Maximum rows in each output batch. Zero preserves layout boundaries without a row cap. * Physical layout boundaries may produce shorter batches. */ size_t batch_rows; diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 29c0ace36a8..7cd9ea6a0ae 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -76,7 +76,7 @@ pub const VX_CUDA_SCAN_FLAG_DIRECT_IO: u32 = 1u32 << 0; pub struct vx_cuda_scan_options { /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are ignored. pub flags: u32, - /// Maximum rows in each output batch. Zero uses layout-derived splitting. + /// Maximum rows in each output batch. Zero preserves layout boundaries without a row cap. /// Physical layout boundaries may produce shorter batches. pub batch_rows: usize, } @@ -393,12 +393,14 @@ fn projected_scan( let projection = select(columns, root()).optimize(file.dtype())?; scan = scan.with_projection(projection.bind(file.dtype())?); } - if batch_rows != 0 { + let split_by = if batch_rows == 0 { + SplitBy::Layout + } else { let max_rows = u64::try_from(batch_rows) .map_err(|_| vortex_err!("CUDA scan batch row count is too large"))?; - scan = scan.with_split_by(SplitBy::LayoutSubSplitting { max_rows }); - } - Ok(scan) + SplitBy::LayoutSubSplitting { max_rows } + }; + Ok(scan.with_split_by(split_by)) } struct CudaScanOptions { diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index a62648de6d9..f3233db4739 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -214,6 +214,17 @@ fn test_projected_scan_rejects_nonstruct_projection() -> VortexResult<()> { Ok(()) } +#[test] +fn test_projected_scan_zero_batch_rows_preserves_large_layout_span() -> VortexResult<()> { + let session = session(); + let file = flat_ids_file(&session, 1_000_000)?; + for columns in [names(&[])?, names(&["ids"])?] { + let splits = projected_scan(&file, columns, 0)?.full_file_splits()?; + assert_eq!(splits, [0, 1_000_000]); + } + Ok(()) +} + struct RejectSegments { inner: Arc, forbidden: Vec, @@ -548,5 +559,5 @@ fn test_projection_gpu_subdivides_large_blocks() -> VortexResult<()> { #[cuda_test] fn test_projection_gpu_preserves_small_block_boundaries() -> VortexResult<()> { - check_projected_file(2, 3) + check_projected_file(2, 0) } From a12ab0a81336aac3ebb5a24238dd38e52de3d163 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:01:48 +0000 Subject: [PATCH 21/35] fix(cuda): restore exact nonzero scan batch row counts Address review item 2. Restore the pre-PR RowCount policy rather than silently subdividing layout spans evenly. Add exact-size and cross-layout host regressions plus GPU coverage for splitting a single block. General CUDA Chunked concatenation remains a documented pre-existing limitation. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 21 ++++-- vortex-cuda/ffi/src/lib.rs | 27 +++++--- vortex-cuda/ffi/src/tests/projection.rs | 90 +++++++++++++++++++++++-- 3 files changed, 117 insertions(+), 21 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 78775772f19..1b58ec44a24 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -71,8 +71,11 @@ typedef struct vx_cuda_scan_options { */ uint32_t flags; /** - * Maximum rows in each output batch. Zero preserves layout boundaries without a row cap. - * Physical layout boundaries may produce shorter batches. + * Rows in each output batch, except for a possibly smaller final batch. + * Zero preserves layout boundaries without a row cap. Nonzero values split at exact row + * counts independently of layout boundaries: 1,000 rows with 300 yields 300/300/300/100. + * Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of + * `Chunked` arrays is currently unsupported. */ size_t batch_rows; } vx_cuda_scan_options; @@ -117,7 +120,10 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, * value disables byte-size coalescing and outer layout dictionaries, so passing 8,192 is not * equivalent to passing zero. * - * Write and scan sizing are independent; scan batches preserve on-disk layout boundaries. + * Write and scan sizing are independent. Zero scan `batch_rows` preserves on-disk layout + * boundaries; nonzero values request exact row counts with a possibly smaller final batch. + * Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of `Chunked` + * arrays is currently unsupported. * * # Safety * @@ -155,11 +161,14 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, vx_error **error_out); /** - * Scan a local Vortex file with bounded row batches. + * Scan a local Vortex file with exact row batches and a possibly smaller final batch. * * Uses [`vx_cuda_scan_path_arrow_device_stream`]'s export and ownership rules. - * `batch_rows` caps output rows; zero uses layout splitting. Physical boundaries may shorten - * batches. Scan and write sizing are independent; scans preserve on-disk layout boundaries. + * Zero preserves layout boundaries without a row cap, so batches may be large. Nonzero + * `batch_rows` splits at exact row counts independently of layout boundaries. For example, + * 1,000 rows with `batch_rows = 300` yields batches of 300/300/300/100 rows. + * Scan and write sizing are independent. Cross-layout batches still require CUDA-supported + * encodings; CUDA concatenation of `Chunked` arrays is currently unsupported. * * # Safety * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 7cd9ea6a0ae..7ec0c2c6d46 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -76,8 +76,11 @@ pub const VX_CUDA_SCAN_FLAG_DIRECT_IO: u32 = 1u32 << 0; pub struct vx_cuda_scan_options { /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are ignored. pub flags: u32, - /// Maximum rows in each output batch. Zero preserves layout boundaries without a row cap. - /// Physical layout boundaries may produce shorter batches. + /// Rows in each output batch, except for a possibly smaller final batch. + /// Zero preserves layout boundaries without a row cap. Nonzero values split at exact row + /// counts independently of layout boundaries: 1,000 rows with 300 yields 300/300/300/100. + /// Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of + /// `Chunked` arrays is currently unsupported. pub batch_rows: usize, } @@ -161,7 +164,10 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( /// value disables byte-size coalescing and outer layout dictionaries, so passing 8,192 is not /// equivalent to passing zero. /// -/// Write and scan sizing are independent; scan batches preserve on-disk layout boundaries. +/// Write and scan sizing are independent. Zero scan `batch_rows` preserves on-disk layout +/// boundaries; nonzero values request exact row counts with a possibly smaller final batch. +/// Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of `Chunked` +/// arrays is currently unsupported. /// /// # Safety /// @@ -226,11 +232,14 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( } } -/// Scan a local Vortex file with bounded row batches. +/// Scan a local Vortex file with exact row batches and a possibly smaller final batch. /// /// Uses [`vx_cuda_scan_path_arrow_device_stream`]'s export and ownership rules. -/// `batch_rows` caps output rows; zero uses layout splitting. Physical boundaries may shorten -/// batches. Scan and write sizing are independent; scans preserve on-disk layout boundaries. +/// Zero preserves layout boundaries without a row cap, so batches may be large. Nonzero +/// `batch_rows` splits at exact row counts independently of layout boundaries. For example, +/// 1,000 rows with `batch_rows = 300` yields batches of 300/300/300/100 rows. +/// Scan and write sizing are independent. Cross-layout batches still require CUDA-supported +/// encodings; CUDA concatenation of `Chunked` arrays is currently unsupported. /// /// # Safety /// @@ -382,7 +391,7 @@ unsafe fn scan_columns(columns: *const vx_view, ncolumns: usize) -> VortexResult Ok(names.into()) } -/// Apply projection before column reads; row limits subdivide, never merge, layout splits. +/// Apply projection before column reads; zero preserves layouts, nonzero splits by row count. fn projected_scan( file: &VortexFile, columns: FieldNames, @@ -396,9 +405,7 @@ fn projected_scan( let split_by = if batch_rows == 0 { SplitBy::Layout } else { - let max_rows = u64::try_from(batch_rows) - .map_err(|_| vortex_err!("CUDA scan batch row count is too large"))?; - SplitBy::LayoutSubSplitting { max_rows } + SplitBy::RowCount(batch_rows) }; Ok(scan.with_split_by(split_by)) } diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index f3233db4739..604de3aa415 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -23,6 +23,7 @@ use vortex::dtype::Nullability; use vortex::file::WriteOptionsSessionExt; use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutStrategy; +use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; use vortex::layout::layouts::table::TableStrategy; use vortex::layout::segments::SegmentFuture; @@ -225,6 +226,64 @@ fn test_projected_scan_zero_batch_rows_preserves_large_layout_span() -> VortexRe Ok(()) } +#[test] +fn test_projected_scan_exact_batch_rows_with_final_tail() -> VortexResult<()> { + let session = session(); + let file = flat_ids_file(&session, 1_000)?; + let expected = StructArray::try_new( + ["ids"].into(), + vec![PrimitiveArray::from_iter(0u32..1_000).into_array()], + 1_000, + Validity::NonNullable, + )? + .into_array(); + for columns in [names(&[])?, names(&["ids"])?] { + let batches: Vec = ffi_runtime().block_on( + projected_scan(&file, columns, 300)? + .into_array_stream()? + .try_collect(), + )?; + let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); + assert_eq!(lengths, [300, 300, 300, 100]); + let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + } + Ok(()) +} + +#[test] +fn test_projected_scan_exact_batch_rows_crosses_layout_blocks() -> VortexResult<()> { + let session = session(); + let input = table()?; + let columns = names(&["値.x", "ids"])?; + let expected = input.project(columns.as_ref())?.into_array(); + let input = input.into_array(); + let chunks = ChunkedArray::try_new( + vec![input.slice(0..2)?, input.slice(2..4)?, input.slice(4..5)?], + input.dtype().clone(), + )? + .into_array(); + let file = open_file( + &session, + chunks, + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())), + )?; + assert_eq!( + projected_scan(&file, columns.clone(), 0)?.full_file_splits()?, + [0, 2, 4, 5] + ); + let batches: Vec = ffi_runtime().block_on( + projected_scan(&file, columns, 3)? + .into_array_stream()? + .try_collect(), + )?; + let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); + assert_eq!(lengths, [3, 2]); + let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + Ok(()) +} + struct RejectSegments { inner: Arc, forbidden: Vec, @@ -511,10 +570,14 @@ fn read_projected_batches(stream: &mut ArrowDeviceArrayStream) -> VortexResult VortexResult<()> { +fn check_projected_file( + input: StructArray, + block_rows: usize, + batch_rows: usize, + expected_lengths: &[usize], +) -> VortexResult<()> { let session = session().with_some(CudaSession::try_default()?); register_cuda_layout(&session); - let input = table()?; let columns = ["値.x", "ids"]; let expected = input.project(names(&columns)?.as_ref())?.into_array(); let mut file = NamedTempFile::new()?; @@ -542,7 +605,7 @@ fn check_projected_file(block_rows: usize, batch_rows: usize) -> VortexResult<() assert_eq!(Schema::try_from(&schema)?, Schema::new(expected_fields)); let batches = read_projected_batches(&mut stream.0)?; let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); - assert_eq!(lengths, [2, 2, 1]); + assert_eq!(lengths, expected_lengths); let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); assert_arrays_eq!( actual, @@ -554,10 +617,27 @@ fn check_projected_file(block_rows: usize, batch_rows: usize) -> VortexResult<() #[cuda_test] fn test_projection_gpu_subdivides_large_blocks() -> VortexResult<()> { - check_projected_file(0, 2) + check_projected_file(table()?, 5, 2, &[2, 2, 1]) } #[cuda_test] fn test_projection_gpu_preserves_small_block_boundaries() -> VortexResult<()> { - check_projected_file(2, 0) + check_projected_file(table()?, 2, 0, &[2, 2, 1]) +} + +#[cuda_test] +fn test_projection_gpu_exact_batch_rows_with_final_tail() -> VortexResult<()> { + let input = StructArray::try_new( + ["ids", "値.x"].into(), + vec![ + PrimitiveArray::from_iter(0u32..1_000).into_array(), + PrimitiveArray::from_option_iter( + (0i64..1_000).map(|value| (value % 2 == 0).then_some(value)), + ) + .into_array(), + ], + 1_000, + Validity::NonNullable, + )?; + check_projected_file(input, 1_000, 300, &[300, 300, 300, 100]) } From b65236e453c37cf8825fadc67c318cb0a7d88e34 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:02:51 +0000 Subject: [PATCH 22/35] refactor(cuda): share the CUDA file writer strategy with the CLI Address review item 9. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/lib.rs | 32 +--------------------------- vortex-cuda/gpu-scan-cli/src/main.rs | 27 +++-------------------- vortex-cuda/src/layout.rs | 28 +++++++++++++++++++++++- 3 files changed, 31 insertions(+), 56 deletions(-) diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 7ec0c2c6d46..7d3683802df 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -10,16 +10,12 @@ use std::os::raw::c_int; use std::ptr; -use std::sync::Arc; use arrow_schema::ffi::FFI_ArrowSchema; use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamExt; -use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::FieldName; use vortex::dtype::FieldNames; -use vortex::editions::ComponentKind; -use vortex::editions::EditionSessionExt; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::error::vortex_err; @@ -27,9 +23,7 @@ use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; -use vortex::file::WriteStrategyBuilder; use vortex::io::runtime::BlockingRuntime; -use vortex::layout::LayoutStrategy; use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::layout::scan::split_by::SplitBy; use vortex::session::SessionExt; @@ -44,7 +38,7 @@ use vortex_cuda::arrow::ArrowDeviceArray; use vortex_cuda::arrow::ArrowDeviceArrayStream; use vortex_cuda::arrow::DeviceArrayExt; use vortex_cuda::arrow::DeviceArrayStreamExt; -use vortex_cuda::layout::CudaFlatLayoutStrategy; +use vortex_cuda::layout::cuda_write_strategy; use vortex_cuda::layout::register_cuda_layout; use vortex_ffi::ffi_runtime; use vortex_ffi::try_or; @@ -91,30 +85,6 @@ fn session_with_cuda(session: &VortexSession) -> &VortexSession { session } -/// Build a CUDA-flat writer using only session-enabled encodings. -fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { - let allowed_encodings = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - let mut strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default() - .only_cuda_compatible() - .retain_allowed_encodings(&allowed_encodings), - ) - .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); - if block_rows > 0 { - // Preserve explicit row blocks: outer layout dictionaries can split a high-cardinality - // block into u16-sized dictionary runs, while a byte target can coalesce adjacent blocks. - strategy = strategy - .with_probe_compressor(BtrBlocksCompressorBuilder::empty().build()) - .with_row_block_size(block_rows) - .with_data_block_target_bytes(None); - } - strategy.build() -} - /// Create a CUDA Vortex session. /// /// Repeated [`vx_cuda_array_export_arrow_device`] calls reuse this CUDA state. Returns an owned diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 2aa32729ef3..74ff881a0ea 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -23,13 +23,9 @@ use vortex::array::arrays::Dict; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::buffer::ByteBufferMut; -use vortex::compressor::BtrBlocksCompressorBuilder; -use vortex::editions::ComponentKind; -use vortex::editions::EditionSessionExt; use vortex::error::VortexResult; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; use vortex::io::session::RuntimeSessionExt; use vortex::session::SessionExt; use vortex::session::VortexSession; @@ -38,7 +34,7 @@ use vortex_cuda::CudaSession; use vortex_cuda::PooledByteBufferReadAt; use vortex_cuda::TracingLaunchStrategy; use vortex_cuda::executor::CudaArrayExt; -use vortex_cuda::layout::CudaFlatLayoutStrategy; +use vortex_cuda::layout::cuda_write_strategy; use vortex_cuda::layout::register_cuda_layout; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; @@ -92,23 +88,6 @@ async fn main() -> VortexResult<()> { } } -/// Build the write strategy used for CUDA-compatible file output. -#[cuda_available] -fn cuda_write_strategy(session: &VortexSession) -> Arc { - let allowed_encodings = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - WriteStrategyBuilder::default() - .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default() - .only_cuda_compatible() - .retain_allowed_encodings(&allowed_encodings), - ) - .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) - .build() -} - /// Convert an input Vortex file to CUDA-compatible encodings and write to disk. #[cuda_available] async fn cmd_convert(input: PathBuf, output: PathBuf) -> VortexResult<()> { @@ -121,7 +100,7 @@ async fn cmd_convert(input: PathBuf, output: PathBuf) -> VortexResult<()> { let mut out = tokio::fs::File::create(&output).await?; session .write_options() - .with_strategy(cuda_write_strategy(&session)) + .with_strategy(cuda_write_strategy(&session, 0)) .write(&mut out, scan) .await?; @@ -233,7 +212,7 @@ async fn recompress_for_gpu( let mut out = ByteBufferMut::empty(); let result = session .write_options() - .with_strategy(cuda_write_strategy(session)) + .with_strategy(cuda_write_strategy(session, 0)) .write(&mut out, scan) .await?; diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index bd9908b3da6..8810044ebed 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -29,8 +29,10 @@ use vortex::array::serde::SerializedArray; use vortex::array::stats::StatsSetRef; use vortex::buffer::BufferString; use vortex::buffer::ByteBuffer; +use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; use vortex::dtype::FieldMask; +use vortex::editions::ComponentKind; use vortex::editions::Edition; use vortex::editions::EditionDeclaration; use vortex::editions::EditionFamily; @@ -41,6 +43,7 @@ use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_panic; +use vortex::file::WriteStrategyBuilder; use vortex::layout::Layout; use vortex::layout::LayoutChildType; use vortex::layout::LayoutDeserializeArgs; @@ -544,6 +547,30 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { result } +/// Build a CUDA-flat writer using only session-enabled encodings. +pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { + let allowed_encodings = session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + let mut strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder( + BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .retain_allowed_encodings(&allowed_encodings), + ) + .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); + if block_rows > 0 { + // Preserve explicit row blocks: outer layout dictionaries can split a high-cardinality + // block into u16-sized dictionary runs, while a byte target can coalesce adjacent blocks. + strategy = strategy + .with_probe_compressor(BtrBlocksCompressorBuilder::empty().build()) + .with_row_block_size(block_rows) + .with_data_block_target_bytes(None); + } + strategy.build() +} + #[derive(Clone, Debug, Default)] struct CudaLayoutRegistration(Arc); @@ -620,7 +647,6 @@ mod tests { use vortex::buffer::ByteBufferMut; use vortex::buffer::buffer; use vortex::editions::CORE_2025_05_0; - use vortex::editions::ComponentKind; use vortex::editions::DEFAULT_CORE_EDITION; use vortex::file::WriteOptionsSessionExt; use vortex::io::runtime::BlockingRuntime; From 8b227f901fac1224fd45a0f7fa08203668a3cb36 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:03:08 +0000 Subject: [PATCH 23/35] fix(cuda): retain integer dictionaries with explicit writer blocks Address review item 1. Keep per-block IntDict compression when disabling outer layout dictionaries; move and strengthen writer regressions alongside layout tests. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 4 +- vortex-cuda/ffi/src/lib.rs | 4 +- vortex-cuda/ffi/src/tests/projection.rs | 72 ++++----- vortex-cuda/src/layout.rs | 189 ++++++++++++++++++++++-- 4 files changed, 206 insertions(+), 63 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 1b58ec44a24..5a94db9c50a 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -117,8 +117,8 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, * * `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default * writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Any nonzero - * value disables byte-size coalescing and outer layout dictionaries, so passing 8,192 is not - * equivalent to passing zero. + * value disables byte-size coalescing and outer layout dictionaries, but retains per-block + * dictionary compression. Passing 8,192 is therefore not equivalent to passing zero. * * Write and scan sizing are independent. Zero scan `batch_rows` preserves on-disk layout * boundaries; nonzero values request exact row counts with a possibly smaller final batch. diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 7d3683802df..2179d810c71 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -131,8 +131,8 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( /// /// `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default /// writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Any nonzero -/// value disables byte-size coalescing and outer layout dictionaries, so passing 8,192 is not -/// equivalent to passing zero. +/// value disables byte-size coalescing and outer layout dictionaries, but retains per-block +/// dictionary compression. Passing 8,192 is therefore not equivalent to passing zero. /// /// Write and scan sizing are independent. Zero scan `batch_rows` preserves on-disk layout /// boundaries; nonzero values request exact row counts with a possibly smaller final batch. diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs index 604de3aa415..dd4164192d0 100644 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ b/vortex-cuda/ffi/src/tests/projection.rs @@ -169,52 +169,6 @@ fn flat_ids_file(session: &VortexSession, rows: u32) -> VortexResult open_file(session, input, Arc::new(FlatLayoutStrategy::default())) } -#[test] -fn test_cuda_write_strategy_preserves_high_cardinality_row_blocks() -> VortexResult<()> { - let session = session(); - let unique = 70_000u32; - let ids = PrimitiveArray::from_iter((0..unique).chain(0..unique)).into_array(); - let rows = ids.len(); - let input = - StructArray::try_new(["ids"].into(), vec![ids], rows, Validity::NonNullable)?.into_array(); - register_cuda_layout(&session); - let file = open_file(&session, input, cuda_write_strategy(&session, rows))?; - let lengths: Vec<_> = ffi_runtime().block_on( - projected_scan(&file, names(&["ids"])?, rows)? - .into_array_stream()? - .map_ok(|batch| batch.len()) - .try_collect(), - )?; - assert_eq!(lengths, [rows]); - Ok(()) -} - -#[test] -fn test_projected_scan_rejects_unknown_field() -> VortexResult<()> { - let session = session(); - let file = flat_ids_file(&session, 5)?; - assert_error( - projected_scan(&file, names(&["missing"])?, 0), - "must be a subset of child fields", - ); - Ok(()) -} - -#[test] -fn test_projected_scan_rejects_nonstruct_projection() -> VortexResult<()> { - let session = session(); - let file = open_file( - &session, - PrimitiveArray::from_iter(0u32..5).into_array(), - Arc::new(FlatLayoutStrategy::default()), - )?; - assert_error( - projected_scan(&file, names(&["ids"])?, 0), - "Select child must return a struct dtype", - ); - Ok(()) -} - #[test] fn test_projected_scan_zero_batch_rows_preserves_large_layout_span() -> VortexResult<()> { let session = session(); @@ -284,6 +238,32 @@ fn test_projected_scan_exact_batch_rows_crosses_layout_blocks() -> VortexResult< Ok(()) } +#[test] +fn test_projected_scan_rejects_unknown_field() -> VortexResult<()> { + let session = session(); + let file = flat_ids_file(&session, 5)?; + assert_error( + projected_scan(&file, names(&["missing"])?, 0), + "must be a subset of child fields", + ); + Ok(()) +} + +#[test] +fn test_projected_scan_rejects_nonstruct_projection() -> VortexResult<()> { + let session = session(); + let file = open_file( + &session, + PrimitiveArray::from_iter(0u32..5).into_array(), + Arc::new(FlatLayoutStrategy::default()), + )?; + assert_error( + projected_scan(&file, names(&["ids"])?, 0), + "Select child must return a struct dtype", + ); + Ok(()) +} + struct RejectSegments { inner: Arc, forbidden: Vec, diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 8810044ebed..a2ae4a9fd8a 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -547,28 +547,34 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { result } -/// Build a CUDA-flat writer using only session-enabled encodings. +/// Build a CUDA-flat writer using only CUDA-compatible, session-enabled array encodings. +/// +/// Register CUDA layout support with [`register_cuda_layout`] before writing. A zero `block_rows` +/// uses the default writer's row sizing and dictionary policy. A nonzero value sets explicit row +/// blocks, disables outer layout dictionaries and byte-size coalescing, and retains per-block +/// dictionary compression. pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { let allowed_encodings = session .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - let mut strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default() - .only_cuda_compatible() - .retain_allowed_encodings(&allowed_encodings), - ) + let builder = BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .retain_allowed_encodings(&allowed_encodings); + let strategy = WriteStrategyBuilder::default() .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); - if block_rows > 0 { - // Preserve explicit row blocks: outer layout dictionaries can split a high-cardinality - // block into u16-sized dictionary runs, while a byte target can coalesce adjacent blocks. - strategy = strategy + if block_rows == 0 { + strategy.with_btrblocks_builder(builder).build() + } else { + // Outer dictionaries can split blocks into u16-sized runs. Disable their probe, but pass + // an opaque compressor so the writer does not also exclude per-block IntDict compression. + strategy + .with_compressor(builder.build()) .with_probe_compressor(BtrBlocksCompressorBuilder::empty().build()) .with_row_block_size(block_rows) - .with_data_block_target_bytes(None); + .with_data_block_target_bytes(None) + .build() } - strategy.build() } #[derive(Clone, Debug, Default)] @@ -641,20 +647,177 @@ pub fn register_cuda_layout(session: &VortexSession) { #[cfg(test)] mod tests { + use futures::TryStreamExt; use rstest::rstest; use vortex::VortexSessionDefault; use vortex::array::IntoArray; + use vortex::array::arrays::Dict; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::arrays::StructArray; + use vortex::array::arrays::struct_::StructArrayExt; + use vortex::array::assert_arrays_eq; use vortex::buffer::ByteBufferMut; use vortex::buffer::buffer; use vortex::editions::CORE_2025_05_0; use vortex::editions::DEFAULT_CORE_EDITION; + use vortex::file::OpenOptionsSessionExt; + use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; use vortex::io::session::RuntimeSessionExt; + use vortex::layout::layouts::dict::Dict as DictLayout; + use vortex::layout::scan::split_by::SplitBy; use super::*; + fn repeated_ids(unique: i64, rows: usize) -> VortexResult { + // Wide, irregularly ordered values favor dictionaries over bitpacking and FoR. + let ids = PrimitiveArray::from_iter( + (0..unique) + .cycle() + .take(rows) + .map(|id| (id * 7_919 % unique).wrapping_mul(0x5851_f42d_4c95_7f2d)), + ); + Ok(StructArray::from_fields(&[("ids", ids.into_array())])?.into_array()) + } + + async fn write_file( + session: &VortexSession, + array: ArrayRef, + strategy: Arc, + ) -> VortexResult { + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(strategy) + .write(&mut buffer, array.to_array_stream()) + .await?; + session.open_options().open_buffer(buffer.freeze()) + } + + fn data_layouts(layout: &LayoutRef) -> VortexResult> { + let mut layouts = vec![layout.clone()]; + for (kind, child) in layout.child_types().zip(layout.children()?) { + // Zone maps and dictionary values do not describe physical data row blocks. + if !matches!(kind, LayoutChildType::Auxiliary(_)) { + layouts.extend(data_layouts(&child)?); + } + } + Ok(layouts) + } + + #[rstest] + fn test_cuda_write_strategy_preserves_integer_dictionary_compression( + #[values(1024, 4096)] block_rows: usize, + ) -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with_handle(runtime.handle()); + register_cuda_layout(&session); + runtime.block_on(async { + let input = repeated_ids(8, 2 * block_rows + 137)?; + let file = write_file( + &session, + input.clone(), + cuda_write_strategy(&session, block_rows), + ) + .await?; + let layouts = data_layouts(file.footer().layout())?; + assert!(!layouts.iter().any(|layout| layout.is::())); + let physical_rows: Vec<_> = layouts + .iter() + .filter(|layout| layout.is::()) + .map(|layout| layout.row_count()) + .collect(); + assert_eq!(physical_rows, [block_rows as u64, block_rows as u64, 137]); + + let batches: Vec<_> = file + .scan()? + .with_split_by(SplitBy::Layout) + .into_array_stream()? + .try_collect() + .await?; + assert_eq!(batches.len(), 3); + let mut ctx = session.create_execution_ctx(); + let mut offset = 0; + for batch in batches { + // Execute only the struct wrapper, leaving its encoded child intact. + let batch = batch.execute::(&mut ctx)?; + assert!(batch.unmasked_field(0).is::()); + let end = offset + batch.len(); + assert_arrays_eq!(batch.into_array(), input.slice(offset..end)?, &mut ctx); + offset = end; + } + assert_eq!(offset, input.len()); + Ok(()) + }) + } + + #[test] + fn test_cuda_write_strategy_preserves_high_cardinality_row_blocks() -> VortexResult<()> { + let runtime = CurrentThreadRuntime::new(); + let session = VortexSession::default().with_handle(runtime.handle()); + register_cuda_layout(&session); + runtime.block_on(async { + // Exceed the outer dictionary's u16 limit, with enough repetitions to be eligible. + let block_rows = 70_000 * 8; + let input = repeated_ids(70_000, block_rows)?; + let allowed_encodings = session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + let compressor = BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .retain_allowed_encodings(&allowed_encodings) + .build(); + // Positive control: identical row sizing and compression, but with the probe enabled. + let control = WriteStrategyBuilder::default() + .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) + .with_compressor(compressor) + .with_row_block_size(block_rows) + .with_data_block_target_bytes(None) + .build(); + let control = write_file(&session, input.clone(), control).await?; + let layouts = data_layouts(control.footer().layout())?; + assert!(layouts.iter().any(|layout| layout.is::())); + let control_rows: Vec<_> = layouts + .iter() + .filter(|layout| layout.is::()) + .map(|layout| layout.row_count()) + .collect(); + assert!(control_rows.len() > 1); + assert_eq!(control_rows.iter().sum::(), block_rows as u64); + + let file = write_file( + &session, + input.clone(), + cuda_write_strategy(&session, block_rows), + ) + .await?; + let layouts = data_layouts(file.footer().layout())?; + assert!(!layouts.iter().any(|layout| layout.is::())); + let physical_rows: Vec<_> = layouts + .iter() + .filter(|layout| layout.is::()) + .map(|layout| layout.row_count()) + .collect(); + assert_eq!(physical_rows, [block_rows as u64]); + + let mut batches: Vec<_> = file + .scan()? + .with_split_by(SplitBy::Layout) + .into_array_stream()? + .try_collect() + .await?; + assert_eq!(batches.len(), 1); + let mut ctx = session.create_execution_ctx(); + let batch = batches.remove(0).execute::(&mut ctx)?; + assert!(batch.unmasked_field(0).is::()); + assert_arrays_eq!(batch.into_array(), input, &mut ctx); + Ok(()) + }) + } + #[rstest] fn test_cuda_registration_preserves_edition_policy( #[values(DEFAULT_CORE_EDITION, CORE_2025_05_0)] core: EditionId, From d018878c11bd062da7fb6e619a7c0b85e2ea8b02 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:03:08 +0000 Subject: [PATCH 24/35] fix(cuda): reject unknown scan option flag bits Address review item 6. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 2 +- vortex-cuda/ffi/src/lib.rs | 40 +++++++++++++++++++------- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 5a94db9c50a..329ae075c36 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -67,7 +67,7 @@ struct ArrowDeviceArrayStream { */ typedef struct vx_cuda_scan_options { /** - * A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are ignored. + * A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are rejected. */ uint32_t flags; /** diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 2179d810c71..1e02cf933b3 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -62,13 +62,15 @@ const VX_CUDA_ERR: c_int = 1; /// Footer and zone-map reads remain buffered. Supported only on Linux. pub const VX_CUDA_SCAN_FLAG_DIRECT_IO: u32 = 1u32 << 0; +const VX_CUDA_SCAN_KNOWN_FLAGS: u32 = VX_CUDA_SCAN_FLAG_DIRECT_IO; + /// Options for scanning a CUDA-compatible Vortex file. /// /// Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. #[repr(C)] #[derive(Default)] pub struct vx_cuda_scan_options { - /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are ignored. + /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are rejected. pub flags: u32, /// Rows in each output batch, except for a possibly smaller final batch. /// Zero preserves layout boundaries without a row cap. Nonzero values split at exact row @@ -393,15 +395,20 @@ fn scan_export_ctx(session: &VortexSession) -> VortexResult { ) } -/// Parse scan settings; null selects defaults and unknown flags are ignored. +/// Parse scan settings; null selects defaults and unknown flags are rejected. /// /// # Safety /// -/// Non-null `options` must point to an initialized, aligned [`vx_cuda_scan_options`]. +/// Non-null `options` must point to an initialized, aligned `vx_cuda_scan_options`. unsafe fn scan_options(options: *const vx_cuda_scan_options) -> VortexResult { let defaults = vx_cuda_scan_options::default(); // SAFETY: The caller guarantees that a non-null options pointer is valid for this call. let options = unsafe { options.as_ref() }.unwrap_or(&defaults); + vortex_ensure!( + options.flags & !VX_CUDA_SCAN_KNOWN_FLAGS == 0, + "unsupported CUDA scan option flags: {:#x}", + options.flags & !VX_CUDA_SCAN_KNOWN_FLAGS + ); let read_at_options = PooledFileReadAtOptions::default(); let read_at_options = if options.flags & VX_CUDA_SCAN_FLAG_DIRECT_IO == 0 { read_at_options @@ -557,15 +564,8 @@ mod tests { let buffered = PooledFileReadAtOptions::default(); for (flags, batch_rows, read_at_options) in [ (0, 8192, buffered), - (1 << 1, 0, buffered), #[cfg(target_os = "linux")] (VX_CUDA_SCAN_FLAG_DIRECT_IO, 0, buffered.with_direct_io()), - #[cfg(target_os = "linux")] - ( - VX_CUDA_SCAN_FLAG_DIRECT_IO | (1 << 1), - 8192, - buffered.with_direct_io(), - ), ] { let options = vx_cuda_scan_options { flags, batch_rows }; // SAFETY: options lives for the duration of parsing. @@ -576,6 +576,26 @@ mod tests { Ok(()) } + #[test] + fn test_scan_options_reject_unknown_flags() { + for flags in [1 << 1, VX_CUDA_SCAN_FLAG_DIRECT_IO | (1 << 1), u32::MAX] { + let options = vx_cuda_scan_options { + flags, + batch_rows: 0, + }; + // SAFETY: options remains live throughout parsing. + let error = unsafe { scan_options(&raw const options) } + .err() + .expect("unknown flags must be rejected"); + assert!( + error + .to_string() + .contains("unsupported CUDA scan option flags"), + "{error}" + ); + } + } + #[cuda_test] fn test_scan_decodes_dictionaries_and_reuses_session_resources() -> VortexResult<()> { // A distinct allocator identity detects accidental reconstruction of a default session. From d67340c19f740f324863e10ec6faa8660b1e59c3 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:03:08 +0000 Subject: [PATCH 25/35] docs(cuda): avoid Rustdoc link brackets in the generated C header Address review item 14. Signed-off-by: Alexander Droste --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 22 +++++++++++----------- vortex-cuda/ffi/src/lib.rs | 22 +++++++++++----------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 329ae075c36..bf96c5b682d 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -87,7 +87,7 @@ extern "C" { /** * Create a CUDA Vortex session. * - * Repeated [`vx_cuda_array_export_arrow_device`] calls reuse this CUDA state. Returns an owned + * Repeated `vx_cuda_array_export_arrow_device` calls reuse this CUDA state. Returns an owned * session handle, or null and an optional `vx_error` on failure. * * # Safety @@ -127,7 +127,7 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, * * # Safety * - * Same requirements as [`vx_cuda_array_sink_open_file`]. + * Same requirements as `vx_cuda_array_sink_open_file`. */ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session, vx_view path, @@ -138,7 +138,7 @@ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session /** * Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. * - * Requires CUDA-supported encodings/layouts, such as files from [`vx_cuda_array_sink_open_file`]. + * Requires CUDA-supported encodings/layouts, such as files from `vx_cuda_array_sink_open_file`. * Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers * reused across scans with the same CUDA session. * @@ -163,7 +163,7 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, /** * Scan a local Vortex file with exact row batches and a possibly smaller final batch. * - * Uses [`vx_cuda_scan_path_arrow_device_stream`]'s export and ownership rules. + * Uses `vx_cuda_scan_path_arrow_device_stream`'s export and ownership rules. * Zero preserves layout boundaries without a row cap, so batches may be large. Nonzero * `batch_rows` splits at exact row counts independently of layout boundaries. For example, * 1,000 rows with `batch_rows = 300` yields batches of 300/300/300/100 rows. @@ -172,7 +172,7 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, * * # Safety * - * Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]. + * Same requirements as `vx_cuda_scan_path_arrow_device_stream`. */ int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, vx_view path, @@ -181,14 +181,14 @@ int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, vx_error **error_out); /** - * Like [`vx_cuda_scan_path_arrow_device_stream`], with explicit scan options. + * Like `vx_cuda_scan_path_arrow_device_stream`, with explicit scan options. * * Null or zero-initialized `options` selects buffered I/O and layout-derived batch splitting. * * # Safety * - * Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]; non-null `options` must - * point to a valid [`vx_cuda_scan_options`]. + * Same requirements as `vx_cuda_scan_path_arrow_device_stream`; non-null `options` must + * point to a valid `vx_cuda_scan_options`. */ int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session, vx_view path, @@ -200,7 +200,7 @@ int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session * Scan a local Vortex file with ordered top-level column projection. * * Same options, ownership, and file requirements as - * [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes decoding and skips + * `vx_cuda_scan_path_arrow_device_stream_with_options`. Projection precedes decoding and skips * unselected column I/O when the file layout stores columns separately. * Names are literal and case-sensitive; a nonempty projection rejects unknown/duplicate names * and non-struct files. `ncolumns == 0` ignores `columns` and selects all. Names are copied; @@ -208,8 +208,8 @@ int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session * * # Safety * - * In addition to [`vx_cuda_scan_path_arrow_device_stream_with_options`]'s requirements, - * nonzero `ncolumns` requires that many initialized, aligned [`vx_view`] values at `columns`. + * In addition to `vx_cuda_scan_path_arrow_device_stream_with_options`'s requirements, + * nonzero `ncolumns` requires that many initialized, aligned `vx_view` values at `columns`. * Each name borrows `len` readable UTF-8 bytes for this call; null is allowed only for zero length. */ int vx_cuda_scan_path_arrow_device_stream_projected(const vx_session *session, diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 1e02cf933b3..d0cb94d9fed 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -89,7 +89,7 @@ fn session_with_cuda(session: &VortexSession) -> &VortexSession { /// Create a CUDA Vortex session. /// -/// Repeated [`vx_cuda_array_export_arrow_device`] calls reuse this CUDA state. Returns an owned +/// Repeated `vx_cuda_array_export_arrow_device` calls reuse this CUDA state. Returns an owned /// session handle, or null and an optional `vx_error` on failure. /// /// # Safety @@ -143,7 +143,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( /// /// # Safety /// -/// Same requirements as [`vx_cuda_array_sink_open_file`]. +/// Same requirements as `vx_cuda_array_sink_open_file`. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( session: *const vx_session, @@ -169,7 +169,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( /// Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. /// -/// Requires CUDA-supported encodings/layouts, such as files from [`vx_cuda_array_sink_open_file`]. +/// Requires CUDA-supported encodings/layouts, such as files from `vx_cuda_array_sink_open_file`. /// Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers /// reused across scans with the same CUDA session. /// @@ -206,7 +206,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( /// Scan a local Vortex file with exact row batches and a possibly smaller final batch. /// -/// Uses [`vx_cuda_scan_path_arrow_device_stream`]'s export and ownership rules. +/// Uses `vx_cuda_scan_path_arrow_device_stream`'s export and ownership rules. /// Zero preserves layout boundaries without a row cap, so batches may be large. Nonzero /// `batch_rows` splits at exact row counts independently of layout boundaries. For example, /// 1,000 rows with `batch_rows = 300` yields batches of 300/300/300/100 rows. @@ -215,7 +215,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( /// /// # Safety /// -/// Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]. +/// Same requirements as `vx_cuda_scan_path_arrow_device_stream`. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows( session: *const vx_session, @@ -240,14 +240,14 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows } } -/// Like [`vx_cuda_scan_path_arrow_device_stream`], with explicit scan options. +/// Like `vx_cuda_scan_path_arrow_device_stream`, with explicit scan options. /// /// Null or zero-initialized `options` selects buffered I/O and layout-derived batch splitting. /// /// # Safety /// -/// Same requirements as [`vx_cuda_scan_path_arrow_device_stream`]; non-null `options` must -/// point to a valid [`vx_cuda_scan_options`]. +/// Same requirements as `vx_cuda_scan_path_arrow_device_stream`; non-null `options` must +/// point to a valid `vx_cuda_scan_options`. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_options( session: *const vx_session, @@ -273,7 +273,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_optio /// Scan a local Vortex file with ordered top-level column projection. /// /// Same options, ownership, and file requirements as -/// [`vx_cuda_scan_path_arrow_device_stream_with_options`]. Projection precedes decoding and skips +/// `vx_cuda_scan_path_arrow_device_stream_with_options`. Projection precedes decoding and skips /// unselected column I/O when the file layout stores columns separately. /// Names are literal and case-sensitive; a nonempty projection rejects unknown/duplicate names /// and non-struct files. `ncolumns == 0` ignores `columns` and selects all. Names are copied; @@ -281,8 +281,8 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_optio /// /// # Safety /// -/// In addition to [`vx_cuda_scan_path_arrow_device_stream_with_options`]'s requirements, -/// nonzero `ncolumns` requires that many initialized, aligned [`vx_view`] values at `columns`. +/// In addition to `vx_cuda_scan_path_arrow_device_stream_with_options`'s requirements, +/// nonzero `ncolumns` requires that many initialized, aligned `vx_view` values at `columns`. /// Each name borrows `len` readable UTF-8 bytes for this call; null is allowed only for zero length. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_projected( From a0861d19503f2854e0acff1124fb8dfb77bf4794 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:34:54 +0000 Subject: [PATCH 26/35] build(cuda): leave generated FFI headers in cbindgen format Signed-off-by: Alexander Droste --- .github/workflows/cuda.yaml | 12 ------ vortex-cuda/ffi/README.md | 23 +++-------- vortex-cuda/ffi/build.rs | 57 +++----------------------- vortex-cuda/ffi/cbindgen.toml | 5 +++ vortex-cuda/ffi/cinclude/vortex_cuda.h | 34 ++++++++------- 5 files changed, 37 insertions(+), 94 deletions(-) diff --git a/.github/workflows/cuda.yaml b/.github/workflows/cuda.yaml index d27bc7b1452..5aa30f436a9 100644 --- a/.github/workflows/cuda.yaml +++ b/.github/workflows/cuda.yaml @@ -47,7 +47,6 @@ jobs: - "uv.lock" - "Cargo.toml" - "Cargo.lock" - - ".clang-format" - ".github/workflows/**" cuda-build-lint: @@ -68,12 +67,6 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} enable-sccache: "true" - - name: Install clang-format - run: | - # Match the Ubuntu 24.04 formatter used by the C/C++ lint job. - sudo apt-get update - sudo apt-get install -y clang-format-18 - echo "/usr/lib/llvm-18/bin" >> "$GITHUB_PATH" - uses: ./.github/actions/check-rebuild with: command: >- @@ -122,11 +115,6 @@ jobs: uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 with: sync: false - - name: Install clang-format - run: | - sudo apt-get update - sudo apt-get install -y clang-format-18 - echo "/usr/lib/llvm-18/bin" >> "$GITHUB_PATH" - name: Install nextest uses: taiki-e/install-action@3f74d7c16a4242f1c95561e98edc25d36adb4375 # v2 with: diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 96bd030f4ae..1b55b9b740c 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -30,20 +30,9 @@ cache for pooled data-plane reads. Footer and zone-map reads remain buffered on ## Generated header -`build.rs` generates `cinclude/vortex_cuda.h` with cbindgen on stable Rust, without macro -expansion or compiling the CUDA implementation. Edit the API and docs in `src/lib.rs`, not -the generated header; commit regenerated headers with API changes. `cbindgen.toml` supplies -the standard Arrow Device interface compatibility preamble. - -Header generation requires `clang-format` on `PATH`, with support for the repository's -`.clang-format` configuration. The build script generates bytes in memory, then formats -through clang-format's stdin/stdout using `--style=file` and -`--assume-filename=cinclude/vortex_cuda.h` from this crate's directory. A missing or failing -formatter fails the build before changing the existing header; there is no unformatted -fallback. - -Only the fully formatted bytes are compared with the committed header. Identical output -leaves the file and its timestamp untouched. Changed output is written to a uniquely created -temporary file in the header's directory and published by atomic rename, so concurrent -builds do not expose a truncated header. The build script tracks changes to `src/`, -`cbindgen.toml`, `build.rs`, and the repository configuration at `../../.clang-format`. +`build.rs` generates `cinclude/vortex_cuda.h` using cbindgen on stable Rust. Edit `src/lib.rs` +or `cbindgen.toml`, not the header, and commit regenerated output. The header keeps cbindgen's +formatting and is excluded from clang-format through generated markers. + +Unchanged output is not rewritten; changed output is published atomically. CUDA CI checks +for header drift after building the FFI crate. diff --git a/vortex-cuda/ffi/build.rs b/vortex-cuda/ffi/build.rs index 4fe2033d382..7a915f65296 100644 --- a/vortex-cuda/ffi/build.rs +++ b/vortex-cuda/ffi/build.rs @@ -7,76 +7,32 @@ use std::fs::OpenOptions; use std::io; use std::io::Write; use std::process; -use std::process::Command; -use std::process::Stdio; -use std::thread; fn main() -> Result<(), Box> { println!("cargo:rerun-if-changed=src"); println!("cargo:rerun-if-changed=cbindgen.toml"); println!("cargo:rerun-if-changed=build.rs"); - println!("cargo:rerun-if-changed=../../.clang-format"); let header = "cinclude/vortex_cuda.h"; let mut generated = Vec::new(); - // The CUDA API needs no macro expansion or dependency parsing, so generate on stable Rust - // without recursively building the CUDA implementation. + // Parse only the FFI source to avoid macro expansion and recursive CUDA builds. cbindgen::Builder::new() .with_src("src/lib.rs") .with_config(cbindgen::Config::from_file("cbindgen.toml")?) .generate()? .write(&mut generated); - let formatted = format_header(header, generated)?; match fs::read(header) { - Ok(existing) if existing == formatted => return Ok(()), + Ok(existing) if existing == generated => return Ok(()), Ok(_) => {} Err(error) if error.kind() == io::ErrorKind::NotFound => {} Err(error) => return Err(format!("failed to read {header}: {error}").into()), } - publish_header(header, &formatted) + publish_header(header, &generated) .map_err(|error| format!("failed to publish {header}: {error}"))?; Ok(()) } -fn format_header(header: &str, generated: Vec) -> Result, Box> { - // Cargo runs this script in the crate directory; the assumed header path lets clang-format - // find the repository's .clang-format while reading from stdin. - let mut child = Command::new("clang-format") - .arg("--style=file") - .arg(format!("--assume-filename={header}")) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .map_err(|error| { - format!("failed to start clang-format (required on PATH to generate {header}): {error}") - })?; - let mut stdin = child - .stdin - .take() - .ok_or("clang-format stdin is unavailable")?; - // Feed stdin while wait_with_output drains stdout and stderr, even when a header exceeds - // pipe capacity. Dropping stdin in the writer signals EOF to clang-format. - let writer = thread::spawn(move || stdin.write_all(&generated)); - let output = child.wait_with_output(); - let written = writer - .join() - .map_err(|_| "clang-format stdin writer panicked")?; - let output = - output.map_err(|error| format!("failed to collect clang-format output: {error}"))?; - if !output.status.success() { - return Err(format!( - "clang-format failed for {header} ({}): {}", - output.status, - String::from_utf8_lossy(&output.stderr) - ) - .into()); - } - written.map_err(|error| format!("failed to write header to clang-format stdin: {error}"))?; - Ok(output.stdout) -} - -fn publish_header(header: &str, formatted: &[u8]) -> io::Result<()> { +fn publish_header(header: &str, generated: &[u8]) -> io::Result<()> { let mut attempt = 0_u64; let (temporary, mut file) = loop { let temporary = format!("{header}.{}.{attempt}.tmp", process::id()); @@ -91,9 +47,8 @@ fn publish_header(header: &str, formatted: &[u8]) -> io::Result<()> { } }; - // A unique sibling file and atomic rename prevent concurrent builds from exposing a - // truncated shared header. Close the file before renaming or cleaning it up. - let written = file.write_all(formatted); + // Publish atomically so concurrent builds never see a truncated header. + let written = file.write_all(generated); drop(file); let result = written.and_then(|()| fs::rename(&temporary, header)); if result.is_err() diff --git a/vortex-cuda/ffi/cbindgen.toml b/vortex-cuda/ffi/cbindgen.toml index 1343336dd3e..d529d4db99c 100644 --- a/vortex-cuda/ffi/cbindgen.toml +++ b/vortex-cuda/ffi/cbindgen.toml @@ -11,6 +11,9 @@ no_includes = true header = """ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors + +// clang-format off + #pragma once // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY @@ -66,6 +69,8 @@ struct ArrowDeviceArrayStream { #endif """ +trailer = "// clang-format on" + # These externally defined Arrow ABI types use struct tags, not typedef names. [export.rename] "ArrowDeviceArray" = "struct ArrowDeviceArray" diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index bf96c5b682d..0143e34e1d2 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -1,5 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors + +// clang-format off + #pragma once // THIS FILE IS AUTO-GENERATED, DO NOT MAKE EDITS DIRECTLY @@ -54,6 +57,7 @@ struct ArrowDeviceArrayStream { }; #endif + /** * Bypass the operating system page cache for pooled data-plane reads. * Footer and zone-map reads remain buffered. Supported only on Linux. @@ -66,18 +70,18 @@ struct ArrowDeviceArrayStream { * Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. */ typedef struct vx_cuda_scan_options { - /** - * A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are rejected. - */ - uint32_t flags; - /** - * Rows in each output batch, except for a possibly smaller final batch. - * Zero preserves layout boundaries without a row cap. Nonzero values split at exact row - * counts independently of layout boundaries: 1,000 rows with 300 yields 300/300/300/100. - * Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of - * `Chunked` arrays is currently unsupported. - */ - size_t batch_rows; + /** + * A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are rejected. + */ + uint32_t flags; + /** + * Rows in each output batch, except for a possibly smaller final batch. + * Zero preserves layout boundaries without a row cap. Nonzero values split at exact row + * counts independently of layout boundaries: 1,000 rows with 300 yields 300/300/300/100. + * Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of + * `Chunked` arrays is currently unsupported. + */ + size_t batch_rows; } vx_cuda_scan_options; #ifdef __cplusplus @@ -264,5 +268,7 @@ int vx_cuda_partition_scan_arrow_device_stream(const vx_session *session, vx_error **error_out); #ifdef __cplusplus -} // extern "C" -#endif // __cplusplus +} // extern "C" +#endif // __cplusplus + +// clang-format on From 362d3ef696295388ef217f850a03e47d31f2590e Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:34:54 +0000 Subject: [PATCH 27/35] docs(cuda): trim repeated export and batching explanations Signed-off-by: Alexander Droste --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 35 ++++++++---------------- vortex-cuda/ffi/src/lib.rs | 38 ++++++++------------------ vortex-cuda/src/arrow/canonical.rs | 6 ++-- vortex-cuda/src/arrow/mod.rs | 19 ++++++------- vortex-cuda/src/layout.rs | 21 ++++++-------- vortex-cuda/src/session.rs | 18 +++--------- 6 files changed, 45 insertions(+), 92 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 0143e34e1d2..b4118ab1910 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -75,11 +75,9 @@ typedef struct vx_cuda_scan_options { */ uint32_t flags; /** - * Rows in each output batch, except for a possibly smaller final batch. - * Zero preserves layout boundaries without a row cap. Nonzero values split at exact row - * counts independently of layout boundaries: 1,000 rows with 300 yields 300/300/300/100. - * Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of - * `Chunked` arrays is currently unsupported. + * Rows per batch, except for a possibly smaller final batch. Zero preserves layout boundaries. + * Nonzero counts ignore layout boundaries and may require unsupported CUDA `Chunked` + * concatenation. */ size_t batch_rows; } vx_cuda_scan_options; @@ -91,8 +89,7 @@ extern "C" { /** * Create a CUDA Vortex session. * - * Repeated `vx_cuda_array_export_arrow_device` calls reuse this CUDA state. Returns an owned - * session handle, or null and an optional `vx_error` on failure. + * Returns an owned handle with reusable CUDA state, or null and an optional `vx_error` on failure. * * # Safety * @@ -119,15 +116,9 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, /** * Open a CUDA-readable Vortex file sink with a fixed row block size. * - * `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default - * writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Any nonzero - * value disables byte-size coalescing and outer layout dictionaries, but retains per-block - * dictionary compression. Passing 8,192 is therefore not equivalent to passing zero. - * - * Write and scan sizing are independent. Zero scan `batch_rows` preserves on-disk layout - * boundaries; nonzero values request exact row counts with a possibly smaller final batch. - * Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of `Chunked` - * arrays is currently unsupported. + * Zero `block_rows` uses default writer sizing. Nonzero values disable byte-size coalescing and + * outer layout dictionaries, but retain per-block dictionary compression. + * Write sizing is independent of scan `batch_rows`; see `vx_cuda_scan_options`. * * # Safety * @@ -146,9 +137,8 @@ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session * Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers * reused across scans with the same CUDA session. * - * Dictionaries, including nested children, decode on CUDA for a stable plain Arrow schema, - * without changing session policy. Decoding can increase device memory use and requires CUDA - * support for device-resident dictionaries. + * Dictionaries, including nested children, decode on CUDA to a stable plain Arrow schema without + * changing session policy. Decoding requires CUDA support and may increase device memory use. * * Returns `0` with an owned `out_stream`; release it and each batch via their Arrow callbacks. * Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. @@ -168,11 +158,8 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, * Scan a local Vortex file with exact row batches and a possibly smaller final batch. * * Uses `vx_cuda_scan_path_arrow_device_stream`'s export and ownership rules. - * Zero preserves layout boundaries without a row cap, so batches may be large. Nonzero - * `batch_rows` splits at exact row counts independently of layout boundaries. For example, - * 1,000 rows with `batch_rows = 300` yields batches of 300/300/300/100 rows. - * Scan and write sizing are independent. Cross-layout batches still require CUDA-supported - * encodings; CUDA concatenation of `Chunked` arrays is currently unsupported. + * `batch_rows` follows `vx_cuda_scan_options`: zero preserves layout boundaries; nonzero counts + * ignore them and may require unsupported CUDA `Chunked` concatenation. * * # Safety * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index d0cb94d9fed..e38ed91cdce 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -72,15 +72,12 @@ const VX_CUDA_SCAN_KNOWN_FLAGS: u32 = VX_CUDA_SCAN_FLAG_DIRECT_IO; pub struct vx_cuda_scan_options { /// A bitwise combination of `VX_CUDA_SCAN_FLAG_*` values. Unknown bits are rejected. pub flags: u32, - /// Rows in each output batch, except for a possibly smaller final batch. - /// Zero preserves layout boundaries without a row cap. Nonzero values split at exact row - /// counts independently of layout boundaries: 1,000 rows with 300 yields 300/300/300/100. - /// Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of - /// `Chunked` arrays is currently unsupported. + /// Rows per batch, except for a possibly smaller final batch. Zero preserves layout boundaries. + /// Nonzero counts ignore layout boundaries and may require unsupported CUDA `Chunked` + /// concatenation. pub batch_rows: usize, } -/// Initialize CUDA support on `session` and return the same borrow. fn session_with_cuda(session: &VortexSession) -> &VortexSession { session.get::(); register_cuda_layout(session); @@ -89,8 +86,7 @@ fn session_with_cuda(session: &VortexSession) -> &VortexSession { /// Create a CUDA Vortex session. /// -/// Repeated `vx_cuda_array_export_arrow_device` calls reuse this CUDA state. Returns an owned -/// session handle, or null and an optional `vx_error` on failure. +/// Returns an owned handle with reusable CUDA state, or null and an optional `vx_error` on failure. /// /// # Safety /// @@ -131,15 +127,9 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( /// Open a CUDA-readable Vortex file sink with a fixed row block size. /// -/// `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default -/// writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Any nonzero -/// value disables byte-size coalescing and outer layout dictionaries, but retains per-block -/// dictionary compression. Passing 8,192 is therefore not equivalent to passing zero. -/// -/// Write and scan sizing are independent. Zero scan `batch_rows` preserves on-disk layout -/// boundaries; nonzero values request exact row counts with a possibly smaller final batch. -/// Cross-layout batches still require CUDA-supported encodings; CUDA concatenation of `Chunked` -/// arrays is currently unsupported. +/// Zero `block_rows` uses default writer sizing. Nonzero values disable byte-size coalescing and +/// outer layout dictionaries, but retain per-block dictionary compression. +/// Write sizing is independent of scan `batch_rows`; see `vx_cuda_scan_options`. /// /// # Safety /// @@ -173,9 +163,8 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( /// Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers /// reused across scans with the same CUDA session. /// -/// Dictionaries, including nested children, decode on CUDA for a stable plain Arrow schema, -/// without changing session policy. Decoding can increase device memory use and requires CUDA -/// support for device-resident dictionaries. +/// Dictionaries, including nested children, decode on CUDA to a stable plain Arrow schema without +/// changing session policy. Decoding requires CUDA support and may increase device memory use. /// /// Returns `0` with an owned `out_stream`; release it and each batch via their Arrow callbacks. /// Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. @@ -207,11 +196,8 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( /// Scan a local Vortex file with exact row batches and a possibly smaller final batch. /// /// Uses `vx_cuda_scan_path_arrow_device_stream`'s export and ownership rules. -/// Zero preserves layout boundaries without a row cap, so batches may be large. Nonzero -/// `batch_rows` splits at exact row counts independently of layout boundaries. For example, -/// 1,000 rows with `batch_rows = 300` yields batches of 300/300/300/100 rows. -/// Scan and write sizing are independent. Cross-layout batches still require CUDA-supported -/// encodings; CUDA concatenation of `Chunked` arrays is currently unsupported. +/// `batch_rows` follows `vx_cuda_scan_options`: zero preserves layout boundaries; nonzero counts +/// ignore them and may require unsupported CUDA `Chunked` concatenation. /// /// # Safety /// @@ -363,7 +349,7 @@ unsafe fn scan_columns(columns: *const vx_view, ncolumns: usize) -> VortexResult Ok(names.into()) } -/// Apply projection before column reads; zero preserves layouts, nonzero splits by row count. +/// Apply projection before column reads. fn projected_scan( file: &VortexFile, columns: FieldNames, diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index a57528af765..e872cfdab51 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -119,8 +119,7 @@ impl ExportDeviceArray for CanonicalDeviceArrayExport { DictionaryExport::Preserve => { rebuild_array_for_export_schema(array, ctx.execution_ctx())? } - // Dictionary layouts no longer affect the schema. Preserve other encodings for - // structural recursion and direct FSST/OnPair export. + // Decode schemas use only dtype; keep encodings for recursion and FSST/OnPair export. DictionaryExport::Decode => array, }; let schema = arrow_schema_for_array(&array, ctx)?; @@ -2583,8 +2582,7 @@ mod tests { #[case] dtype: DType, #[values(DictionaryExport::Preserve, DictionaryExport::Decode)] policy: DictionaryExport, ) -> VortexResult<()> { - // Direct FSST varbin export must work when execute_cuda rejects standalone FSST, - // ruling out eager canonicalization. + // Reject standalone FSST execution to catch eager canonicalization before direct export. let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)? .with_dictionary_export(policy) .with_dispatch_mode(CudaDispatchMode::DynDispatchOnly); diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index b61110be461..d5e035ada1d 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -114,10 +114,8 @@ impl ArrowArray { } impl ArrowDeviceArray { - /// Create an empty, released array with zeroed device metadata and a null sync event. - /// - /// Use this as storage for an Arrow C device callback output or as the basis for an - /// end-of-stream marker. No CUDA device is selected; `device_id` and `device_type` are zero. + /// Create a released array with zeroed device metadata and a null sync event. + /// Use as callback output storage or an end-of-stream marker base; no device is selected. pub fn empty() -> Self { Self { array: ArrowArray::empty(), @@ -497,12 +495,11 @@ pub trait DeviceArrayStreamExt { /// returned [`ArrowDeviceArrayStream`] owns the Vortex stream and must be released through its /// embedded `release` callback. /// - /// The Arrow Device stream contract requires all arrays to share the schema reported by - /// `get_schema`. By default, the schema is derived from the first array, or from the logical - /// dtype for an empty stream. Chunks exporting different Arrow types are rejected mid-stream. - /// With [`DictionaryExport::Decode`], the logical dtype determines a stable plain schema even - /// when chunks vary between dictionary/plain encodings or dictionary index widths. In this - /// mode, `get_schema` does not pull or export a batch; read/decode errors surface in `get_next`. + /// All arrays must share the `get_schema` schema. By default, it comes from the first array + /// (the logical dtype for empty streams); chunks exporting different Arrow types are rejected. + /// With [`DictionaryExport::Decode`], the dtype determines a plain schema independent of chunk + /// encoding or dictionary index width. `get_schema` does not pull or export a batch; + /// read/decode errors surface in `get_next`. /// /// Drive the returned stream from one thread. `runtime` must be the runtime that owns the /// underlying scan tasks and per-array exports. @@ -953,7 +950,7 @@ mod tests { use crate::arrow::release_device_array; use crate::arrow::release_schema; - /// Copy a CUDA buffer from a live, unreleased array produced by this exporter to the host. + /// Copy a CUDA buffer to the host; requires a live array from this exporter. pub(super) fn private_data_buffer_bytes( array: &ArrowArray, index: usize, diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index a2ae4a9fd8a..ef6e0f86682 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -549,9 +549,8 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { /// Build a CUDA-flat writer using only CUDA-compatible, session-enabled array encodings. /// -/// Register CUDA layout support with [`register_cuda_layout`] before writing. A zero `block_rows` -/// uses the default writer's row sizing and dictionary policy. A nonzero value sets explicit row -/// blocks, disables outer layout dictionaries and byte-size coalescing, and retains per-block +/// Requires [`register_cuda_layout`]. Zero `block_rows` uses default sizing and dictionary policy; +/// nonzero sets row blocks without outer dictionaries or byte coalescing, retaining per-block /// dictionary compression. pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { let allowed_encodings = session @@ -566,8 +565,7 @@ pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc().0.call_once(|| { session .layouts() diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index 9d46af4683e..498b0516626 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -40,24 +40,14 @@ pub enum VarBinExportLayout { VarBinView, } -/// Controls whether Arrow Device exports keep dictionary encoding or fully decode it, -/// including dictionaries nested inside structs and lists. -/// -/// For example, indices `[0, 1, 0]` and dictionary values `["apple", "pear"]` export as: -/// -/// - [`Preserve`](Self::Preserve): separate index and dictionary-value arrays, with an Arrow -/// dictionary type. -/// - [`Decode`](Self::Decode): the plain string array `["apple", "pear", "apple"]`, with no -/// dictionary indices or dictionary child. +/// Dictionary policy for Arrow Device exports, including dictionaries nested in structs and lists. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum DictionaryExport { - /// Keep separate index and dictionary-value arrays rather than expanding repeated values. - /// The Arrow schema retains the dictionary type, including its index type. + /// Keep separate indices and dictionary values; the Arrow schema includes the index type. #[default] Preserve, - /// Fully decode dictionary encoding into plain values on CUDA, including repeated values. - /// This keeps one plain Arrow schema across batches with different dictionary encodings, - /// but may increase device memory use. Device-resident inputs require CUDA decoding support. + /// Expand dictionaries into plain values for a stable schema across batches. + /// May increase device memory use; device-resident inputs require CUDA decoding support. Decode, } From de89ec498bc76005d3f9c7004f7b300633d7ded1 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:34:54 +0000 Subject: [PATCH 28/35] test(cuda): focus layout and dictionary regression coverage Signed-off-by: Alexander Droste --- vortex-cuda/src/arrow/dictionary_tests.rs | 21 +- vortex-cuda/src/layout.rs | 223 +++++----------------- 2 files changed, 54 insertions(+), 190 deletions(-) diff --git a/vortex-cuda/src/arrow/dictionary_tests.rs b/vortex-cuda/src/arrow/dictionary_tests.rs index c39fb7e286a..a985601debb 100644 --- a/vortex-cuda/src/arrow/dictionary_tests.rs +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -28,8 +28,7 @@ use super::tests::private_data_buffer_bytes as buffer; use super::*; use crate::CudaSession; -/// Preserve encodings while moving all buffers, including validity, to CUDA so unsupported -/// decoding errors instead of falling back to the CPU. +/// Move all buffers, including validity, to CUDA without decoding, preventing CPU fallback. pub(super) fn upload(array: ArrayRef, ctx: &mut CudaExecutionCtx) -> VortexResult { // Constants store scalar metadata, not replaceable data buffers. if array.as_opt::().is_some() { @@ -225,14 +224,15 @@ fn test_decode_mixed_dictionary_device_stream( } #[rstest] -#[case::values(true, false)] -#[case::error(true, true)] -#[case::empty(false, false)] +#[case::values(true, false, true)] +#[case::error(true, true, true)] +#[case::empty(false, false, true)] +#[case::schema_only(true, true, false)] #[crate::test] fn test_decode_stream_schema_does_not_poll( #[case] has_batch: bool, #[case] fails: bool, - #[values(false, true)] consume: bool, + #[case] consume: bool, ) -> VortexResult<()> { let runtime = CurrentThreadRuntime::new(); let session = vortex::array::array_session() @@ -296,9 +296,7 @@ fn test_decode_stream_validates_dtype_and_device() -> VortexResult<()> { .export_device_array_stream(&session, &runtime)?; // SAFETY: The stream is live and exclusively borrowed until the state is no longer used. let state = unsafe { device_stream_private_data(&raw mut stream) }.expect("missing state"); - let mut first = state.export_stream_array(array.clone())?; - release_device_array(&mut first); - assert!(state.schema.is_some()); + state.get_or_init_schema()?; let error = state .export_stream_array(PrimitiveArray::from_iter([10u32, 20, 30]).into_array()) @@ -394,10 +392,7 @@ fn test_default_dictionary_device_stream(#[case] second_width: Option) -> let runtime = CurrentThreadRuntime::new(); let session = crate::cuda_session(); let mut ctx = CudaSession::create_execution_ctx(&session)?; - assert_eq!( - ctx.cuda_session().dictionary_export(), - DictionaryExport::Preserve - ); + let (values, expected) = values_and_expected(false); let first = dictionary(values.clone(), PType::U8)?; let second = match second_width { diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index ef6e0f86682..a8ddf10eb23 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -652,22 +652,19 @@ mod tests { use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::assert_arrays_eq; use vortex::buffer::ByteBufferMut; - use vortex::buffer::buffer; use vortex::editions::CORE_2025_05_0; - use vortex::editions::DEFAULT_CORE_EDITION; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; use vortex::io::session::RuntimeSessionExt; - use vortex::layout::layouts::dict::Dict as DictLayout; use vortex::layout::scan::split_by::SplitBy; use super::*; fn repeated_ids(unique: i64, rows: usize) -> VortexResult { - // Wide, irregularly ordered values favor dictionaries over bitpacking and FoR. + // Wide, shuffled values favor dictionaries over bitpacking and FoR. let ids = PrimitiveArray::from_iter( (0..unique) .cycle() @@ -680,51 +677,40 @@ mod tests { async fn write_file( session: &VortexSession, array: ArrayRef, - strategy: Arc, + block_rows: usize, ) -> VortexResult { let mut buffer = ByteBufferMut::empty(); session .write_options() - .with_strategy(strategy) + .with_strategy(cuda_write_strategy(session, block_rows)) .write(&mut buffer, array.to_array_stream()) .await?; session.open_options().open_buffer(buffer.freeze()) } - fn data_layouts(layout: &LayoutRef) -> VortexResult> { - let mut layouts = vec![layout.clone()]; + fn data_block_rows(layout: &LayoutRef) -> VortexResult> { + let mut rows = Vec::new(); + if layout.is::() { + rows.push(layout.row_count()); + } for (kind, child) in layout.child_types().zip(layout.children()?) { - // Zone maps and dictionary values do not describe physical data row blocks. + // Exclude zone maps and dictionary values from data row counts. if !matches!(kind, LayoutChildType::Auxiliary(_)) { - layouts.extend(data_layouts(&child)?); + rows.extend(data_block_rows(&child)?); } } - Ok(layouts) + Ok(rows) } - #[rstest] - fn test_cuda_write_strategy_preserves_integer_dictionary_compression( - #[values(1024, 4096)] block_rows: usize, - ) -> VortexResult<()> { + #[test] + fn test_cuda_write_strategy_preserves_integer_dictionary_compression() -> VortexResult<()> { + let block_rows = 1024; let runtime = CurrentThreadRuntime::new(); let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); runtime.block_on(async { let input = repeated_ids(8, 2 * block_rows + 137)?; - let file = write_file( - &session, - input.clone(), - cuda_write_strategy(&session, block_rows), - ) - .await?; - let layouts = data_layouts(file.footer().layout())?; - assert!(!layouts.iter().any(|layout| layout.is::())); - let physical_rows: Vec<_> = layouts - .iter() - .filter(|layout| layout.is::()) - .map(|layout| layout.row_count()) - .collect(); - assert_eq!(physical_rows, [block_rows as u64, block_rows as u64, 137]); + let file = write_file(&session, input.clone(), block_rows).await?; let batches: Vec<_> = file .scan()? @@ -732,18 +718,21 @@ mod tests { .into_array_stream()? .try_collect() .await?; - assert_eq!(batches.len(), 3); + assert_eq!( + batches.iter().map(|batch| batch.len()).collect::>(), + [block_rows, block_rows, 137] + ); let mut ctx = session.create_execution_ctx(); let mut offset = 0; for batch in batches { - // Execute only the struct wrapper, leaving its encoded child intact. + // Keep the child encoded to detect loss of IntDict compression. let batch = batch.execute::(&mut ctx)?; assert!(batch.unmasked_field(0).is::()); let end = offset + batch.len(); assert_arrays_eq!(batch.into_array(), input.slice(offset..end)?, &mut ctx); offset = end; } - assert_eq!(offset, input.len()); + Ok(()) }) } @@ -754,96 +743,34 @@ mod tests { let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); runtime.block_on(async { - // Exceed the outer dictionary's u16 limit, with enough repetitions to be eligible. + // Exceed u16 cardinality while remaining eligible for outer dictionaries. let block_rows = 70_000 * 8; let input = repeated_ids(70_000, block_rows)?; - let allowed_encodings = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - let compressor = BtrBlocksCompressorBuilder::default() - .only_cuda_compatible() - .retain_allowed_encodings(&allowed_encodings) - .build(); - // Positive control: identical row sizing and compression, but with the probe enabled. - let control = WriteStrategyBuilder::default() - .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) - .with_compressor(compressor) - .with_row_block_size(block_rows) - .with_data_block_target_bytes(None) - .build(); - let control = write_file(&session, input.clone(), control).await?; - let layouts = data_layouts(control.footer().layout())?; - assert!(layouts.iter().any(|layout| layout.is::())); - let control_rows: Vec<_> = layouts - .iter() - .filter(|layout| layout.is::()) - .map(|layout| layout.row_count()) - .collect(); - assert!(control_rows.len() > 1); - assert_eq!(control_rows.iter().sum::(), block_rows as u64); - - let file = write_file( - &session, - input.clone(), - cuda_write_strategy(&session, block_rows), - ) - .await?; - let layouts = data_layouts(file.footer().layout())?; - assert!(!layouts.iter().any(|layout| layout.is::())); - let physical_rows: Vec<_> = layouts - .iter() - .filter(|layout| layout.is::()) - .map(|layout| layout.row_count()) - .collect(); - assert_eq!(physical_rows, [block_rows as u64]); - - let mut batches: Vec<_> = file - .scan()? - .with_split_by(SplitBy::Layout) - .into_array_stream()? - .try_collect() - .await?; - assert_eq!(batches.len(), 1); - let mut ctx = session.create_execution_ctx(); - let batch = batches.remove(0).execute::(&mut ctx)?; - assert!(batch.unmasked_field(0).is::()); - assert_arrays_eq!(batch.into_array(), input, &mut ctx); + let file = write_file(&session, input, block_rows).await?; + assert_eq!( + data_block_rows(file.footer().layout())?, + [block_rows as u64] + ); Ok(()) }) } - #[rstest] - fn test_cuda_registration_preserves_edition_policy( - #[values(DEFAULT_CORE_EDITION, CORE_2025_05_0)] core: EditionId, - ) -> VortexResult<()> { + #[test] + fn test_concurrent_cuda_registration_preserves_edition_policy() -> VortexResult<()> { let session = VortexSession::default(); - session.enable_edition(core)?; + session.enable_edition(CORE_2025_05_0)?; let mut expected_editions = session.enabled_editions().editions(); - assert!(expected_editions.contains(&core)); - assert!(!expected_editions.contains(&CUDA_EDITION)); expected_editions.push(CUDA_EDITION); expected_editions.sort_unstable(); - let kinds = [ - ComponentKind::Array, - ComponentKind::Layout, - ComponentKind::DType, - ComponentKind::Aggregate, - ]; - let expected_ids = kinds.map(|kind| { - let mut ids = session.enabled_component_ids(kind); - if kind == ComponentKind::Layout { - assert!(!ids.contains(&CudaFlat.id())); - ids.push(CudaFlat.id()); - ids.sort_unstable(); - } - ids - }); + let expected_arrays = session.enabled_component_ids(ComponentKind::Array); + let barrier = std::sync::Barrier::new(4); std::thread::scope(|scope| { for _ in 0..4 { let session = session.clone(); + let barrier = &barrier; scope.spawn(move || { + barrier.wait(); register_cuda_layout(&session); assert!( session @@ -853,19 +780,13 @@ mod tests { }); } }); - register_cuda_layout(&session); - for (kind, expected) in kinds.into_iter().zip(expected_ids) { - assert_eq!(session.enabled_component_ids(kind), expected); - } let mut enabled_editions = session.enabled_editions().editions(); enabled_editions.sort_unstable(); assert_eq!(enabled_editions, expected_editions); - session.editions().validate()?; - assert!( - !VortexSession::default() - .enabled_component_ids(ComponentKind::Layout) - .contains(&CudaFlat.id()) + assert_eq!( + session.enabled_component_ids(ComponentKind::Array), + expected_arrays ); Ok(()) } @@ -892,52 +813,28 @@ mod tests { let mut expected_editions = session.enabled_editions().editions(); expected_editions.sort_unstable(); let expected_layouts = session.enabled_component_ids(ComponentKind::Layout); - assert!(!expected_layouts.contains(&CudaFlat.id())); - std::thread::scope(|scope| { - for _ in 0..4 { - let session = session.clone(); - let expected_editions = &expected_editions; - let expected_layouts = &expected_layouts; - scope.spawn(move || { - register_cuda_layout(&session); - let mut enabled_editions = session.enabled_editions().editions(); - enabled_editions.sort_unstable(); - assert_eq!(&enabled_editions, expected_editions); - assert_eq!( - &session.enabled_component_ids(ComponentKind::Layout), - expected_layouts - ); - }); - } - }); register_cuda_layout(&session); - assert!(session.editions().find(&CUDA_EDITION).is_some()); - assert!( - session - .enabled_editions() - .editions() - .contains(&OTHER_CUDA_EDITION) + + let mut enabled_editions = session.enabled_editions().editions(); + enabled_editions.sort_unstable(); + assert_eq!(enabled_editions, expected_editions); + assert_eq!( + session.enabled_component_ids(ComponentKind::Layout), + expected_layouts ); - session.editions().validate()?; Ok(()) } - #[rstest] - fn test_cuda_registration_preserves_pre_registered_edition_policy( - #[values(false, true)] enabled: bool, - ) -> VortexResult<()> { + #[test] + fn test_cuda_registration_preserves_disabled_pre_registered_edition() -> VortexResult<()> { let session = VortexSession::default(); session.editions().declare_family(&CUDA_EDITION_FAMILY)?; session.register_edition(&CUDA_EDITION_DECLARATION)?; - if enabled { - session.enable_edition(CUDA_EDITION)?; - } let mut expected_editions = session.enabled_editions().editions(); expected_editions.sort_unstable(); let expected_layouts = session.enabled_component_ids(ComponentKind::Layout); - register_cuda_layout(&session); register_cuda_layout(&session); let mut enabled_editions = session.enabled_editions().editions(); @@ -947,34 +844,6 @@ mod tests { session.enabled_component_ids(ComponentKind::Layout), expected_layouts ); - session.editions().validate()?; Ok(()) } - - #[test] - fn test_registry_alone_does_not_permit_cuda_flat() -> VortexResult<()> { - let runtime = CurrentThreadRuntime::new(); - let session = VortexSession::default().with_handle(runtime.handle()); - session - .layouts() - .register(LayoutEncodingRef::new_ref(&CudaFlat)); - runtime.block_on(async { - let array = buffer![1i32, 4, 9, 16].into_array(); - let mut buffer = ByteBufferMut::empty(); - let error = session - .write_options() - .with_strategy(Arc::new(CudaFlatLayoutStrategy::default())) - .write(&mut buffer, array.to_array_stream()) - .await - .err() - .expect("write permitted an uneditioned CUDA layout"); - assert!( - error - .to_string() - .contains("Layout encoding vortex.cuda_flat not permitted by ctx"), - "unexpected error: {error}" - ); - Ok(()) - }) - } } From 556b223822bd16c8137ce4d103207b510a901f86 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:34:54 +0000 Subject: [PATCH 29/35] test(cuda): consolidate focused FFI tests into the inline module Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/lib.rs | 678 ++++++++++++++++++++---- vortex-cuda/ffi/src/tests/projection.rs | 623 ---------------------- 2 files changed, 579 insertions(+), 722 deletions(-) delete mode 100644 vortex-cuda/ffi/src/tests/projection.rs diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index e38ed91cdce..2f38cb30b57 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -501,41 +501,600 @@ pub unsafe extern "C-unwind" fn vx_cuda_partition_scan_arrow_device_stream( #[cfg(test)] mod tests { - mod projection; - use std::ffi::CStr; + use std::io::Write; + use std::mem::MaybeUninit; use std::ptr; use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Schema; + use cudarc::driver::CudaContext; + use cudarc::driver::result; + use cudarc::driver::sys::CUevent; + use futures::TryStreamExt; + use tempfile::NamedTempFile; use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; + use vortex::array::assert_arrays_eq; use vortex::array::memory::BufferAllocatorRef; use vortex::array::memory::MemorySessionExt; use vortex::array::memory::StaticBufferAllocator; use vortex::array::validity::Validity; + use vortex::buffer::BitBuffer; + use vortex::buffer::Buffer; + use vortex::buffer::ByteBuffer; + use vortex::buffer::ByteBufferMut; + use vortex::dtype::NativePType; + use vortex::dtype::Nullability; use vortex::error::VortexResult; + use vortex::file::WriteOptionsSessionExt; + use vortex::io::session::RuntimeSessionExt; + use vortex::layout::LayoutStrategy; + use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; + use vortex::layout::layouts::table::TableStrategy; + use vortex::layout::segments::SegmentFuture; + use vortex::layout::segments::SegmentId; + use vortex::layout::segments::SegmentSource; use vortex_cuda::arrow::ARROW_DEVICE_CUDA; + use vortex_cuda::arrow::ArrowArray; use vortex_cuda::arrow::release_device_array; use vortex_cuda::arrow::release_schema; use vortex_cuda_macros::cuda_not_available; use vortex_cuda_macros::test as cuda_test; use vortex_ffi::vx_array_free as free_test_array; + use vortex_ffi::vx_error_free; + use vortex_ffi::vx_error_message; use vortex_ffi::vx_session_free as free_test_session; use super::*; + fn test_session(session: VortexSession) -> *mut vx_session { + Box::into_raw(Box::new(session)).cast::() + } + + fn test_array(array: impl IntoArray) -> *const vx_array { + Box::into_raw(Box::new(array.into_array())).cast::() + } + + fn stream_error(stream: &mut ArrowDeviceArrayStream) -> String { + // SAFETY: The callback and returned C string belong to this live stream. + unsafe { + stream + .get_last_error + .and_then(|callback| callback(stream).as_ref()) + .map(|message| CStr::from_ptr(message).to_string_lossy().into_owned()) + .unwrap_or_default() + } + } + + fn stream_schema(stream: &mut ArrowDeviceArrayStream) -> FFI_ArrowSchema { + let mut schema = FFI_ArrowSchema::empty(); + let get_schema = stream.get_schema.expect("missing get_schema"); + // SAFETY: This live stream owns the callback; schema is writable. + assert_eq!( + unsafe { get_schema(stream, (&raw mut schema).cast()) }, + 0, + "{}", + stream_error(stream) + ); + schema + } + + /// # Safety + /// `session` and `array` must be valid borrowed FFI handles for the duration of the call. + unsafe fn export_array( + session: *const vx_session, + array: *const vx_array, + ) -> (FFI_ArrowSchema, ArrowDeviceArray) { + let mut error = ptr::null_mut(); + let mut schema = FFI_ArrowSchema::empty(); + let mut device_array = ArrowDeviceArray::empty(); + // SAFETY: The caller guarantees valid handles; all outputs are live and writable. + let status = unsafe { + vx_cuda_array_export_arrow_device( + session, + array, + &raw mut schema, + &raw mut device_array, + &raw mut error, + ) + }; + assert_eq!(status, VX_CUDA_OK); + assert!(error.is_null()); + (schema, device_array) + } + + fn view(value: &str) -> vx_view { + vx_view { + ptr: value.as_ptr().cast(), + len: value.len(), + } + } + + fn names(values: &[&str]) -> VortexResult { + let views: Vec<_> = values.iter().map(|name| view(name)).collect(); + // SAFETY: All views and their string bytes remain live throughout parsing. + unsafe { scan_columns(views.as_ptr(), views.len()) } + } + + fn assert_error(result: VortexResult, message: &str) { + let error = result.err().expect("expected an error"); + assert!(error.to_string().contains(message), "{error}"); + } + + fn session() -> VortexSession { + VortexSession::default().with_handle(ffi_runtime().handle()) + } + + fn table() -> VortexResult { + StructArray::try_new( + ["ids", "unused", "値.x"].into(), + vec![ + PrimitiveArray::from_iter(0u32..5).into_array(), + PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]).into_array(), + PrimitiveArray::from_option_iter([Some(10i64), None, Some(30), None, Some(50)]) + .into_array(), + ], + 5, + Validity::NonNullable, + ) + } + + fn file_bytes( + session: &VortexSession, + array: ArrayRef, + strategy: Arc, + ) -> VortexResult { + let mut bytes = ByteBufferMut::empty(); + ffi_runtime().block_on( + session + .write_options() + .with_strategy(strategy) + .write(&mut bytes, array.to_array_stream()), + )?; + Ok(bytes.freeze()) + } + + fn open_file( + session: &VortexSession, + array: ArrayRef, + strategy: Arc, + ) -> VortexResult { + session + .open_options() + .open_buffer(file_bytes(session, array, strategy)?) + } + + fn flat_ids_file(session: &VortexSession, rows: u32) -> VortexResult { + let ids = PrimitiveArray::from_iter(0..rows).into_array(); + let rows = ids.len(); + let input = StructArray::try_new(["ids"].into(), vec![ids], rows, Validity::NonNullable)? + .into_array(); + open_file(session, input, Arc::new(FlatLayoutStrategy::default())) + } + + struct RejectSegments { + inner: Arc, + forbidden: Vec, + rejected: AtomicUsize, + } + + impl SegmentSource for RejectSegments { + fn request(&self, id: SegmentId) -> SegmentFuture { + if self.forbidden.contains(&id) { + self.rejected.fetch_add(1, Ordering::Relaxed); + return Box::pin(async move { + Err(vortex_err!("unselected column segment requested: {id}")) + }); + } + self.inner.request(id) + } + } + + /// # Safety + /// `error` must be null or an owned FFI error, consumed by this call. + unsafe fn take_error_message(error: *mut vx_error) -> Option { + if error.is_null() { + return None; + } + // SAFETY: The error remains live while its message is copied, then is freed exactly once. + let message = unsafe { vx_error_message(error).as_str() } + .map(str::to_owned) + .unwrap_or_else(|error| error.to_string()); + unsafe { vx_error_free(error) }; + Some(message) + } + + struct OwnedDeviceStream(ArrowDeviceArrayStream); + + impl Drop for OwnedDeviceStream { + fn drop(&mut self) { + if let Some(release) = self.0.release { + // SAFETY: This owner holds the live stream and releases it exactly once. + unsafe { release(&raw mut self.0) }; + } + } + } + + fn open_stream( + session: &VortexSession, + path: &str, + options: &vx_cuda_scan_options, + columns: &[&str], + ) -> OwnedDeviceStream { + let mut output = MaybeUninit::::uninit(); + let mut error = ptr::null_mut(); + let handle = test_session(session.clone()); + let columns: Vec<_> = columns.iter().map(|name| view(name)).collect(); + // SAFETY: All borrowed inputs and writable outputs are live for this call. + let status = unsafe { + vx_cuda_scan_path_arrow_device_stream_projected( + handle, + view(path), + options, + columns.as_ptr(), + columns.len(), + output.as_mut_ptr(), + &raw mut error, + ) + }; + // SAFETY: This call owns both the session handle and any returned error. + let message = unsafe { + free_test_session(handle); + take_error_message(error) + }; + assert_eq!( + status, + VX_CUDA_OK, + "{}", + message.as_deref().unwrap_or("no FFI error") + ); + // SAFETY: A successful call initialized the stream, which owns its session state. + let stream = OwnedDeviceStream(unsafe { output.assume_init() }); + assert!(message.is_none(), "unexpected FFI error: {message:?}"); + assert!(stream.0.release.is_some(), "missing release"); + stream + } + + /// # Safety + /// `array` must be a live Arrow primitive array of `T` on the current CUDA context. + /// Its producer's sync event must have completed before this call. + unsafe fn read_primitive( + array: &ArrowArray, + nullability: Nullability, + ) -> VortexResult { + assert!(array.release.is_some()); + assert!(array.dictionary.is_null()); + assert_eq!(array.n_buffers, 2); + assert_eq!(array.n_children, 0); + assert!(!array.buffers.is_null()); + let len = usize::try_from(array.length)?; + let offset = usize::try_from(array.offset)?; + // SAFETY: The live primitive array owns two buffer pointers. + let buffers = unsafe { std::slice::from_raw_parts(array.buffers, 2) }; + let mut values = vec![T::default(); offset + len]; + // SAFETY: The synchronized data buffer contains offset + len values of T and remains live. + unsafe { result::memcpy_dtoh_sync(&mut values, buffers[1] as u64) } + .map_err(|error| vortex_err!("copying Arrow values: {error}"))?; + + let validity = if buffers[0].is_null() { + assert_eq!(array.null_count, 0); + match nullability { + Nullability::NonNullable => Validity::NonNullable, + Nullability::Nullable => Validity::AllValid, + } + } else { + let mut bytes = vec![0u8; (offset + len).div_ceil(8)]; + // SAFETY: The synchronized bitmap covers offset + len bits and remains live. + unsafe { result::memcpy_dtoh_sync(&mut bytes, buffers[0] as u64) } + .map_err(|error| vortex_err!("copying Arrow validity: {error}"))?; + let bits = BitBuffer::new_with_offset(ByteBuffer::from(bytes), len, offset); + assert_eq!(array.null_count, i64::try_from(len - bits.true_count())?); + match nullability { + Nullability::NonNullable => { + assert_eq!(array.null_count, 0); + Validity::NonNullable + } + Nullability::Nullable => Validity::from(bits), + } + }; + Ok(PrimitiveArray::new(Buffer::from(values).slice(offset..), validity).into_array()) + } + + /// Read the fixture's projected Int64/UInt32 columns through the public Arrow Device ABI. + /// + /// # Safety + /// `array` must be a live batch from the fixture's projected stream, with its schema checked. + unsafe fn read_projected_batch(array: &ArrowDeviceArray) -> VortexResult { + assert_eq!(array.device_type, ARROW_DEVICE_CUDA); + let context = CudaContext::new(usize::try_from(array.device_id)?) + .map_err(|error| vortex_err!("opening Arrow device context: {error}"))?; + context + .bind_to_thread() + .map_err(|error| vortex_err!("binding Arrow device context: {error}"))?; + if !array.sync_event.is_null() { + // SAFETY: Arrow's CUDA sync_event points to a live event handle owned by this batch. + unsafe { result::event::synchronize(*array.sync_event.cast::()) } + .map_err(|error| vortex_err!("waiting for Arrow device batch: {error}"))?; + } + let array = &array.array; + assert!(array.dictionary.is_null()); + assert_eq!(array.offset, 0); + assert_eq!(array.null_count, 0); + assert_eq!(array.n_children, 2); + assert!(!array.children.is_null()); + // SAFETY: The live struct owns two children matching the previously checked schema. + let fields = unsafe { + let values = (*array.children).as_ref().expect("missing values child"); + let ids = (*array.children.add(1)) + .as_ref() + .expect("missing ids child"); + assert_eq!(values.length, array.length); + assert_eq!(ids.length, array.length); + vec![ + read_primitive::(values, Nullability::Nullable)?, + read_primitive::(ids, Nullability::NonNullable)?, + ] + }; + Ok(StructArray::try_new( + ["値.x", "ids"].into(), + fields, + usize::try_from(array.length)?, + Validity::NonNullable, + )? + .into_array()) + } + + fn read_projected_batches(stream: &mut ArrowDeviceArrayStream) -> VortexResult> { + let get_next = stream.get_next.expect("missing get_next"); + let mut batches = Vec::new(); + loop { + let mut array = ArrowDeviceArray::empty(); + // SAFETY: This live stream owns the callback; array is writable. + let status = unsafe { get_next(stream, &raw mut array) }; + vortex_ensure!(status == 0, "get_next failed: {}", stream_error(stream)); + if array.array.release.is_none() { + break; + } + // SAFETY: The fixture's schema was checked, and this batch remains live during readback. + let batch = unsafe { read_projected_batch(&array) }; + release_device_array(&mut array); + batches.push(batch?); + } + Ok(batches) + } + + #[test] + fn test_projection_names_are_owned_and_zero_count_means_all() -> VortexResult<()> { + let parsed = { + let name = String::from("値.x"); + names(&[&name, "ids", ""])? + }; + assert_eq!(parsed, ["値.x", "ids", ""]); + // SAFETY: Zero count ignores the pointer, including a null pointer. + assert!(unsafe { scan_columns(ptr::null(), 0) }?.is_empty()); + let empty = vx_view { + ptr: ptr::null(), + len: 0, + }; + // SAFETY: A null, zero-length view is a valid empty name. + assert_eq!(unsafe { scan_columns(&raw const empty, 1) }?, [""]); + Ok(()) + } + + #[test] + fn test_projection_rejects_invalid_names_and_counts() { + let invalid_utf8 = vx_view { + ptr: [0xffu8].as_ptr().cast(), + len: 1, + }; + let null_name = vx_view { + ptr: ptr::null(), + len: 1, + }; + let long_name = vx_view { + ptr: "x".as_ptr().cast(), + len: usize::MAX, + }; + let duplicate = String::from("x"); + let aligned = [view("x"), view(&duplicate)]; + let misaligned = aligned.as_ptr().cast::().wrapping_add(1).cast(); + for (columns, count, message) in [ + (ptr::null(), 1, "null CUDA scan columns"), + (aligned.as_ptr(), usize::MAX, "column count is too large"), + (misaligned, 1, "unaligned CUDA scan columns"), + (&raw const invalid_utf8, 1, "invalid utf-8"), + (&raw const null_name, 1, "null vx_view pointer"), + (&raw const long_name, 1, "name is too long"), + (aligned.as_ptr(), 2, "duplicate CUDA scan column: \"x\""), + ] { + // SAFETY: Invalid pointer/length combinations must be rejected before dereferencing; + // all remaining views and bytes are live. + assert_error(unsafe { scan_columns(columns, count) }, message); + } + } + + #[test] + fn test_projected_scan_zero_batch_rows_preserves_large_layout_span() -> VortexResult<()> { + let session = session(); + // Exceed the default scan split cap to catch accidentally leaving it enabled. + let file = flat_ids_file(&session, 1_000_000)?; + for columns in [names(&[])?, names(&["ids"])?] { + let splits = projected_scan(&file, columns, 0)?.full_file_splits()?; + assert_eq!(splits, [0, 1_000_000]); + } + Ok(()) + } + + #[test] + fn test_projected_scan_exact_batch_rows_with_final_tail() -> VortexResult<()> { + let session = session(); + let file = flat_ids_file(&session, 10)?; + for columns in [names(&[])?, names(&["ids"])?] { + let batches: Vec = ffi_runtime().block_on( + projected_scan(&file, columns, 3)? + .into_array_stream()? + .try_collect(), + )?; + let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); + assert_eq!(lengths, [3, 3, 3, 1]); + } + Ok(()) + } + + #[test] + fn test_projection_cpu_never_requests_unselected_column_segments() -> VortexResult<()> { + let session = session(); + let input = table()?; + let columns = names(&["値.x", "ids"])?; + let expected = input.project(columns.as_ref())?.into_array(); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let strategy = Arc::new(TableStrategy::new(Arc::clone(&flat), flat)); + let file = open_file(&session, input.into_array(), strategy)?; + // TableStrategy writes one flat child per column, so child 1 is exactly the unused column. + let children = file.footer().layout().children()?; + let forbidden = children[1].segment_ids(); + assert!(!forbidden.is_empty()); + let source = Arc::new(RejectSegments { + inner: file.segment_source(), + forbidden, + rejected: AtomicUsize::new(0), + }); + let file = file.with_segment_source(Arc::::clone(&source)); + let actual = ffi_runtime().block_on( + projected_scan(&file, columns, 2)? + .into_array_stream()? + .read_all(), + )?; + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + assert_eq!(source.rejected.load(Ordering::Relaxed), 0); + // The same reader must fail without projection, proving the guard actually observes reads. + assert_error( + ffi_runtime().block_on( + projected_scan(&file, names(&[])?, 0)? + .into_array_stream()? + .read_all(), + ), + "unselected column segment requested", + ); + assert!(source.rejected.load(Ordering::Relaxed) > 0); + Ok(()) + } + + #[test] + fn test_projection_ffi_validation_without_cuda() { + let mut stream = ArrowDeviceArrayStream { + device_type: -1, + get_schema: None, + get_next: None, + get_last_error: None, + release: None, + private_data: ptr::null_mut(), + }; + let mut error = ptr::null_mut(); + // SAFETY: Output pointers are writable; invalid columns must fail before session/path use. + let status = unsafe { + vx_cuda_scan_path_arrow_device_stream_projected( + ptr::null(), + view(""), + ptr::null(), + ptr::null(), + 1, + &raw mut stream, + &raw mut error, + ) + }; + // SAFETY: This call owns the returned error and frees it exactly once. + let message = unsafe { take_error_message(error) }.expect("missing FFI error"); + assert_eq!(status, VX_CUDA_ERR, "{message}"); + assert!( + message.contains("null CUDA scan columns with nonzero count"), + "{message}" + ); + assert_eq!(stream.device_type, -1); + assert!(stream.release.is_none()); + error = ptr::null_mut(); + // SAFETY: Null output is rejected before any other input is used; error is writable. + assert_eq!( + unsafe { + vx_cuda_scan_path_arrow_device_stream_projected( + ptr::null(), + view(""), + ptr::null(), + ptr::null(), + 0, + ptr::null_mut(), + &raw mut error, + ) + }, + VX_CUDA_ERR + ); + // SAFETY: This call owns the returned error and frees it exactly once. + let message = unsafe { take_error_message(error) }.expect("missing FFI error"); + assert!( + message.contains("null ArrowDeviceArrayStream output"), + "{message}" + ); + } + + #[cuda_test] + fn test_projection_gpu_values_and_validity() -> VortexResult<()> { + let session = session().with_some(CudaSession::try_default()?); + register_cuda_layout(&session); + let input = table()?; + let columns = ["値.x", "ids"]; + let expected = input.project(names(&columns)?.as_ref())?.into_array(); + let mut file = NamedTempFile::new()?; + file.write_all(&file_bytes( + &session, + input.into_array(), + cuda_write_strategy(&session, 5), + )?)?; + let path = file + .path() + .to_str() + .ok_or_else(|| vortex_err!("non-UTF-8 test path"))?; + let options = vx_cuda_scan_options { + batch_rows: 2, + ..Default::default() + }; + let mut stream = open_stream(&session, path, &options, &columns); + // The stream must retain its session state after the caller releases its session. + drop(session); + let schema = stream_schema(&mut stream.0); + let expected_fields = vec![ + Field::new("値.x", DataType::Int64, true), + Field::new("ids", DataType::UInt32, false), + ]; + assert_eq!(Schema::try_from(&schema)?, Schema::new(expected_fields)); + let batches = read_projected_batches(&mut stream.0)?; + assert_eq!( + batches.iter().map(|batch| batch.len()).collect::>(), + [2, 2, 1] + ); + let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); + assert_arrays_eq!( + actual, + expected, + &mut VortexSession::default().create_execution_ctx() + ); + Ok(()) + } + #[test] fn test_scan_options_default_to_buffered_io() -> VortexResult<()> { let options = vx_cuda_scan_options::default(); - assert_eq!(options.flags, 0); - assert_eq!(options.batch_rows, 0); + for pointer in [ptr::null(), &raw const options] { // SAFETY: Each pointer is either null or points to the live options above. let parsed = unsafe { scan_options(pointer) }?; @@ -564,7 +1123,7 @@ mod tests { #[test] fn test_scan_options_reject_unknown_flags() { - for flags in [1 << 1, VX_CUDA_SCAN_FLAG_DIRECT_IO | (1 << 1), u32::MAX] { + for flags in [1 << 1, VX_CUDA_SCAN_FLAG_DIRECT_IO | (1 << 1)] { let options = vx_cuda_scan_options { flags, batch_rows: 0, @@ -583,15 +1142,20 @@ mod tests { } #[cuda_test] - fn test_scan_decodes_dictionaries_and_reuses_session_resources() -> VortexResult<()> { - // A distinct allocator identity detects accidental reconstruction of a default session. + fn test_scan_context_preserves_session_resources_and_policy() -> VortexResult<()> { + // A distinct allocator detects accidental reconstruction of a default session. let allocator = BufferAllocatorRef::new(StaticBufferAllocator); let session = VortexSession::default() .with_some(CudaSession::try_default()?) .with_allocator(allocator.clone()); - let mut export_ctx = scan_export_ctx(session_with_cuda(&session))?; - assert!(export_ctx.execution_ctx().allocator().ptr_eq(&allocator)); - let export_session = export_ctx.execution_ctx().session(); + let mut ctx = scan_export_ctx(session_with_cuda(&session))?; + + assert_eq!( + session.get::().dictionary_export(), + DictionaryExport::Preserve + ); + assert!(ctx.execution_ctx().allocator().ptr_eq(&allocator)); + let export_session = ctx.execution_ctx().session(); assert!(Arc::ptr_eq( session.get::().pinned_buffer_pool(), export_session.get::().pinned_buffer_pool(), @@ -601,100 +1165,16 @@ mod tests { PrimitiveArray::from_iter([10i32, 20]).into_array(), )? .into_array(); - for (ctx, expected_type, preserved) in [ - (export_ctx, DataType::Int32, false), - ( - CudaSession::create_execution_ctx(&session)?, - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Int32)), - true, - ), - ] { - let mut stream = ArrowDeviceArrayStream::new( - array.clone().to_array_stream().boxed(), - ctx, - ffi_runtime(), - ); - let get_next = stream.get_next.expect("missing get_next"); - let release = stream.release.expect("missing release"); - let schema = stream_schema(&mut stream); - let mut exported = ArrowDeviceArray::empty(); - // SAFETY: The live stream owns the callback, and the output is writable. - unsafe { - assert_eq!(get_next(&raw mut stream, &raw mut exported), 0); - } - assert_eq!(Field::try_from(&schema)?.data_type(), &expected_type); - assert_eq!(!exported.array.dictionary.is_null(), preserved); - // SAFETY: The batch and stream are live and released exactly once. - unsafe { - release_device_array(&mut exported); - release(&raw mut stream); - } - } + let mut exported = + ffi_runtime().block_on(array.export_device_array_with_schema(&mut ctx))?; + release_device_array(&mut exported.array); assert_eq!( - session.get::().dictionary_export(), - DictionaryExport::Preserve + Field::try_from(&exported.schema)?.data_type(), + &DataType::Int32 ); - assert!(session.allocator().ptr_eq(&allocator)); - Ok(()) } - fn test_session(session: VortexSession) -> *mut vx_session { - Box::into_raw(Box::new(session)).cast::() - } - - fn test_array(array: impl IntoArray) -> *const vx_array { - Box::into_raw(Box::new(array.into_array())).cast::() - } - - fn stream_error(stream: &mut ArrowDeviceArrayStream) -> String { - // SAFETY: The callback and returned C string belong to this live stream. - unsafe { - stream - .get_last_error - .and_then(|callback| callback(stream).as_ref()) - .map(|message| CStr::from_ptr(message).to_string_lossy().into_owned()) - .unwrap_or_default() - } - } - - fn stream_schema(stream: &mut ArrowDeviceArrayStream) -> FFI_ArrowSchema { - let mut schema = FFI_ArrowSchema::empty(); - let get_schema = stream.get_schema.expect("missing get_schema"); - // SAFETY: This live stream owns the callback; schema is writable. - assert_eq!( - unsafe { get_schema(stream, (&raw mut schema).cast()) }, - 0, - "{}", - stream_error(stream) - ); - schema - } - - /// # Safety - /// `session` and `array` must be valid borrowed FFI handles for the duration of the call. - unsafe fn export_array( - session: *const vx_session, - array: *const vx_array, - ) -> (FFI_ArrowSchema, ArrowDeviceArray) { - let mut error = ptr::null_mut(); - let mut schema = FFI_ArrowSchema::empty(); - let mut device_array = ArrowDeviceArray::empty(); - // SAFETY: The caller guarantees valid handles; all outputs are live and writable. - let status = unsafe { - vx_cuda_array_export_arrow_device( - session, - array, - &raw mut schema, - &raw mut device_array, - &raw mut error, - ) - }; - assert_eq!(status, VX_CUDA_OK); - assert!(error.is_null()); - (schema, device_array) - } - #[cuda_test] fn test_export_primitive_arrow_device() { let session = test_session(VortexSession::default()); diff --git a/vortex-cuda/ffi/src/tests/projection.rs b/vortex-cuda/ffi/src/tests/projection.rs deleted file mode 100644 index dd4164192d0..00000000000 --- a/vortex-cuda/ffi/src/tests/projection.rs +++ /dev/null @@ -1,623 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::io::Write; -use std::mem::MaybeUninit; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; - -use cudarc::driver::CudaContext; -use cudarc::driver::result; -use cudarc::driver::sys::CUevent; -use futures::TryStreamExt; -use tempfile::NamedTempFile; -use vortex::array::VortexSessionExecute; -use vortex::array::arrays::ChunkedArray; -use vortex::array::assert_arrays_eq; -use vortex::buffer::BitBuffer; -use vortex::buffer::Buffer; -use vortex::buffer::ByteBuffer; -use vortex::buffer::ByteBufferMut; -use vortex::dtype::NativePType; -use vortex::dtype::Nullability; -use vortex::file::WriteOptionsSessionExt; -use vortex::io::session::RuntimeSessionExt; -use vortex::layout::LayoutStrategy; -use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; -use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; -use vortex::layout::layouts::table::TableStrategy; -use vortex::layout::segments::SegmentFuture; -use vortex::layout::segments::SegmentId; -use vortex::layout::segments::SegmentSource; -use vortex_cuda::arrow::ArrowArray; -use vortex_ffi::vx_error_free; -use vortex_ffi::vx_error_message; - -use super::*; - -fn view(value: &str) -> vx_view { - vx_view { - ptr: value.as_ptr().cast(), - len: value.len(), - } -} - -fn names(values: &[&str]) -> VortexResult { - let views: Vec<_> = values.iter().map(|name| view(name)).collect(); - // SAFETY: All views and their string bytes remain live throughout parsing. - unsafe { scan_columns(views.as_ptr(), views.len()) } -} - -fn assert_error(result: VortexResult, message: &str) { - let error = result.err().expect("expected an error"); - assert!(error.to_string().contains(message), "{error}"); -} - -#[test] -fn test_projection_names_are_owned_and_zero_count_means_all() -> VortexResult<()> { - let parsed = { - let name = String::from("値.x"); - names(&[&name, ""])? - }; - assert_eq!(parsed, ["値.x", ""]); - // SAFETY: Zero count ignores the pointer, including a null pointer. - assert!(unsafe { scan_columns(ptr::null(), 0) }?.is_empty()); - let empty = vx_view { - ptr: ptr::null(), - len: 0, - }; - // SAFETY: A null, zero-length view is a valid empty name. - assert_eq!(unsafe { scan_columns(&raw const empty, 1) }?, [""]); - Ok(()) -} - -#[test] -fn test_projection_wide_order_and_first_late_duplicate() -> VortexResult<()> { - let columns: Vec<_> = (0..1024).rev().map(|i| format!("column_{i}")).collect(); - let mut projection: Vec<_> = columns.iter().map(String::as_str).collect(); - let parsed = names(&projection)?; - assert_eq!(parsed, projection.as_slice()); - - let duplicate = String::from("column_512"); - projection.extend([duplicate.as_str(), "column_1023"]); - assert_error( - names(&projection), - "duplicate CUDA scan column: \"column_512\"", - ); - Ok(()) -} - -#[test] -fn test_projection_rejects_invalid_names_and_counts() { - let invalid_utf8 = vx_view { - ptr: [0xffu8].as_ptr().cast(), - len: 1, - }; - let null_name = vx_view { - ptr: ptr::null(), - len: 1, - }; - let long_name = vx_view { - ptr: "x".as_ptr().cast(), - len: usize::MAX, - }; - let aligned = [view("x"), view("x")]; - let misaligned = aligned.as_ptr().cast::().wrapping_add(1).cast(); - for (columns, count, message) in [ - (ptr::null(), 1, "null CUDA scan columns"), - (aligned.as_ptr(), usize::MAX, "column count is too large"), - (misaligned, 1, "unaligned CUDA scan columns"), - (&raw const invalid_utf8, 1, "invalid utf-8"), - (&raw const null_name, 1, "null vx_view pointer"), - (&raw const long_name, 1, "name is too long"), - (aligned.as_ptr(), 2, "duplicate CUDA scan column: \"x\""), - ] { - // SAFETY: Invalid pointer/length combinations must be rejected before dereferencing; - // all remaining views and bytes are live. - assert_error(unsafe { scan_columns(columns, count) }, message); - } -} - -fn session() -> VortexSession { - VortexSession::default().with_handle(ffi_runtime().handle()) -} - -fn table() -> VortexResult { - StructArray::try_new( - ["ids", "unused", "値.x"].into(), - vec![ - PrimitiveArray::from_iter(0u32..5).into_array(), - PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]).into_array(), - PrimitiveArray::from_option_iter([Some(10i64), None, Some(30), None, Some(50)]) - .into_array(), - ], - 5, - Validity::NonNullable, - ) -} - -fn file_bytes( - session: &VortexSession, - array: ArrayRef, - strategy: Arc, -) -> VortexResult { - let mut bytes = ByteBufferMut::empty(); - ffi_runtime().block_on( - session - .write_options() - .with_strategy(strategy) - .write(&mut bytes, array.to_array_stream()), - )?; - Ok(bytes.freeze()) -} - -fn open_file( - session: &VortexSession, - array: ArrayRef, - strategy: Arc, -) -> VortexResult { - session - .open_options() - .open_buffer(file_bytes(session, array, strategy)?) -} - -fn flat_ids_file(session: &VortexSession, rows: u32) -> VortexResult { - let ids = PrimitiveArray::from_iter(0..rows).into_array(); - let rows = ids.len(); - let input = - StructArray::try_new(["ids"].into(), vec![ids], rows, Validity::NonNullable)?.into_array(); - open_file(session, input, Arc::new(FlatLayoutStrategy::default())) -} - -#[test] -fn test_projected_scan_zero_batch_rows_preserves_large_layout_span() -> VortexResult<()> { - let session = session(); - let file = flat_ids_file(&session, 1_000_000)?; - for columns in [names(&[])?, names(&["ids"])?] { - let splits = projected_scan(&file, columns, 0)?.full_file_splits()?; - assert_eq!(splits, [0, 1_000_000]); - } - Ok(()) -} - -#[test] -fn test_projected_scan_exact_batch_rows_with_final_tail() -> VortexResult<()> { - let session = session(); - let file = flat_ids_file(&session, 1_000)?; - let expected = StructArray::try_new( - ["ids"].into(), - vec![PrimitiveArray::from_iter(0u32..1_000).into_array()], - 1_000, - Validity::NonNullable, - )? - .into_array(); - for columns in [names(&[])?, names(&["ids"])?] { - let batches: Vec = ffi_runtime().block_on( - projected_scan(&file, columns, 300)? - .into_array_stream()? - .try_collect(), - )?; - let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); - assert_eq!(lengths, [300, 300, 300, 100]); - let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); - assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); - } - Ok(()) -} - -#[test] -fn test_projected_scan_exact_batch_rows_crosses_layout_blocks() -> VortexResult<()> { - let session = session(); - let input = table()?; - let columns = names(&["値.x", "ids"])?; - let expected = input.project(columns.as_ref())?.into_array(); - let input = input.into_array(); - let chunks = ChunkedArray::try_new( - vec![input.slice(0..2)?, input.slice(2..4)?, input.slice(4..5)?], - input.dtype().clone(), - )? - .into_array(); - let file = open_file( - &session, - chunks, - Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())), - )?; - assert_eq!( - projected_scan(&file, columns.clone(), 0)?.full_file_splits()?, - [0, 2, 4, 5] - ); - let batches: Vec = ffi_runtime().block_on( - projected_scan(&file, columns, 3)? - .into_array_stream()? - .try_collect(), - )?; - let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); - assert_eq!(lengths, [3, 2]); - let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); - assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); - Ok(()) -} - -#[test] -fn test_projected_scan_rejects_unknown_field() -> VortexResult<()> { - let session = session(); - let file = flat_ids_file(&session, 5)?; - assert_error( - projected_scan(&file, names(&["missing"])?, 0), - "must be a subset of child fields", - ); - Ok(()) -} - -#[test] -fn test_projected_scan_rejects_nonstruct_projection() -> VortexResult<()> { - let session = session(); - let file = open_file( - &session, - PrimitiveArray::from_iter(0u32..5).into_array(), - Arc::new(FlatLayoutStrategy::default()), - )?; - assert_error( - projected_scan(&file, names(&["ids"])?, 0), - "Select child must return a struct dtype", - ); - Ok(()) -} - -struct RejectSegments { - inner: Arc, - forbidden: Vec, - rejected: AtomicUsize, -} - -impl SegmentSource for RejectSegments { - fn request(&self, id: SegmentId) -> SegmentFuture { - if self.forbidden.contains(&id) { - self.rejected.fetch_add(1, Ordering::Relaxed); - return Box::pin(async move { - Err(vortex_err!("unselected column segment requested: {id}")) - }); - } - self.inner.request(id) - } -} - -#[test] -fn test_projection_cpu_never_requests_unselected_column_segments() -> VortexResult<()> { - let session = session(); - let input = table()?; - let columns = names(&["値.x", "ids"])?; - let expected = input.project(columns.as_ref())?.into_array(); - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let strategy = Arc::new(TableStrategy::new(Arc::clone(&flat), flat)); - let file = open_file(&session, input.into_array(), strategy)?; - // TableStrategy writes one flat child per column, so child 1 is exactly the unused column. - let children = file.footer().layout().children()?; - let forbidden = children[1].segment_ids(); - assert!(!forbidden.is_empty()); - let source = Arc::new(RejectSegments { - inner: file.segment_source(), - forbidden, - rejected: AtomicUsize::new(0), - }); - let file = file.with_segment_source(Arc::::clone(&source)); - let actual = ffi_runtime().block_on( - projected_scan(&file, columns, 2)? - .into_array_stream()? - .read_all(), - )?; - assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); - assert_eq!(source.rejected.load(Ordering::Relaxed), 0); - // The same reader must fail without projection, proving the guard actually observes reads. - assert_error( - ffi_runtime().block_on( - projected_scan(&file, names(&[])?, 0)? - .into_array_stream()? - .read_all(), - ), - "unselected column segment requested", - ); - assert!(source.rejected.load(Ordering::Relaxed) > 0); - Ok(()) -} - -#[test] -fn test_projection_ffi_validation_without_cuda() { - let mut stream = ArrowDeviceArrayStream { - device_type: -1, - get_schema: None, - get_next: None, - get_last_error: None, - release: None, - private_data: ptr::null_mut(), - }; - let mut error = ptr::null_mut(); - // SAFETY: Output pointers are writable; invalid columns must fail before session/path use. - let status = unsafe { - vx_cuda_scan_path_arrow_device_stream_projected( - ptr::null(), - view(""), - ptr::null(), - ptr::null(), - 1, - &raw mut stream, - &raw mut error, - ) - }; - // SAFETY: This call owns the returned error and frees it exactly once. - let message = unsafe { take_error_message(error) }.expect("missing FFI error"); - assert_eq!(status, VX_CUDA_ERR, "{message}"); - assert!( - message.contains("null CUDA scan columns with nonzero count"), - "{message}" - ); - assert_eq!(stream.device_type, -1); - assert!(stream.release.is_none()); - error = ptr::null_mut(); - // SAFETY: Null output is rejected before any other input is used; error is writable. - assert_eq!( - unsafe { - vx_cuda_scan_path_arrow_device_stream_projected( - ptr::null(), - view(""), - ptr::null(), - ptr::null(), - 0, - ptr::null_mut(), - &raw mut error, - ) - }, - VX_CUDA_ERR - ); - // SAFETY: This call owns the returned error and frees it exactly once. - let message = unsafe { take_error_message(error) }.expect("missing FFI error"); - assert!( - message.contains("null ArrowDeviceArrayStream output"), - "{message}" - ); -} - -/// # Safety -/// `error` must be null or an owned FFI error, consumed by this call. -unsafe fn take_error_message(error: *mut vx_error) -> Option { - if error.is_null() { - return None; - } - // SAFETY: The error remains live while its message is copied, then is freed exactly once. - let message = unsafe { vx_error_message(error).as_str() } - .map(str::to_owned) - .unwrap_or_else(|error| error.to_string()); - unsafe { vx_error_free(error) }; - Some(message) -} - -struct OwnedDeviceStream(ArrowDeviceArrayStream); - -impl Drop for OwnedDeviceStream { - fn drop(&mut self) { - if let Some(release) = self.0.release { - // SAFETY: This owner holds the live stream and releases it exactly once. - unsafe { release(&raw mut self.0) }; - } - } -} - -fn open_stream( - session: &VortexSession, - path: &str, - options: &vx_cuda_scan_options, - columns: &[&str], -) -> OwnedDeviceStream { - let mut output = MaybeUninit::::uninit(); - let mut error = ptr::null_mut(); - let handle = test_session(session.clone()); - let columns: Vec<_> = columns.iter().map(|name| view(name)).collect(); - // SAFETY: All borrowed inputs and writable outputs are live for this call. - let status = unsafe { - vx_cuda_scan_path_arrow_device_stream_projected( - handle, - view(path), - options, - columns.as_ptr(), - columns.len(), - output.as_mut_ptr(), - &raw mut error, - ) - }; - // SAFETY: This call owns both the session handle and any returned error. - let message = unsafe { - free_test_session(handle); - take_error_message(error) - }; - assert_eq!( - status, - VX_CUDA_OK, - "{}", - message.as_deref().unwrap_or("no FFI error") - ); - // SAFETY: A successful call initialized the stream, which owns its session state. - let stream = OwnedDeviceStream(unsafe { output.assume_init() }); - assert!(message.is_none(), "unexpected FFI error: {message:?}"); - assert!(stream.0.release.is_some(), "missing release"); - stream -} - -/// # Safety -/// `array` must be a live Arrow primitive array of `T` on the current CUDA context. -/// Its producer's sync event must have completed before this call. -unsafe fn read_primitive( - array: &ArrowArray, - nullability: Nullability, -) -> VortexResult { - assert!(array.release.is_some()); - assert!(array.dictionary.is_null()); - assert_eq!(array.n_buffers, 2); - assert_eq!(array.n_children, 0); - assert!(!array.buffers.is_null()); - let len = usize::try_from(array.length)?; - let offset = usize::try_from(array.offset)?; - // SAFETY: The live primitive array owns two buffer pointers. - let buffers = unsafe { std::slice::from_raw_parts(array.buffers, 2) }; - let mut values = vec![T::default(); offset + len]; - // SAFETY: The synchronized data buffer contains offset + len values of T and remains live. - unsafe { result::memcpy_dtoh_sync(&mut values, buffers[1] as u64) } - .map_err(|error| vortex_err!("copying Arrow values: {error}"))?; - - let validity = if buffers[0].is_null() { - assert_eq!(array.null_count, 0); - match nullability { - Nullability::NonNullable => Validity::NonNullable, - Nullability::Nullable => Validity::AllValid, - } - } else { - let mut bytes = vec![0u8; (offset + len).div_ceil(8)]; - // SAFETY: The synchronized bitmap covers offset + len bits and remains live. - unsafe { result::memcpy_dtoh_sync(&mut bytes, buffers[0] as u64) } - .map_err(|error| vortex_err!("copying Arrow validity: {error}"))?; - let bits = BitBuffer::new_with_offset(ByteBuffer::from(bytes), len, offset); - assert_eq!(array.null_count, i64::try_from(len - bits.true_count())?); - match nullability { - Nullability::NonNullable => { - assert_eq!(array.null_count, 0); - Validity::NonNullable - } - Nullability::Nullable => Validity::from(bits), - } - }; - Ok(PrimitiveArray::new(Buffer::from(values).slice(offset..), validity).into_array()) -} - -/// Read the fixture's projected Int64/UInt32 columns through the public Arrow Device ABI. -/// -/// # Safety -/// `array` must be a live batch from the fixture's projected stream, with its schema checked. -unsafe fn read_projected_batch(array: &ArrowDeviceArray) -> VortexResult { - assert_eq!(array.device_type, ARROW_DEVICE_CUDA); - let context = CudaContext::new(usize::try_from(array.device_id)?) - .map_err(|error| vortex_err!("opening Arrow device context: {error}"))?; - context - .bind_to_thread() - .map_err(|error| vortex_err!("binding Arrow device context: {error}"))?; - if !array.sync_event.is_null() { - // SAFETY: Arrow's CUDA sync_event points to a live event handle owned by this batch. - unsafe { result::event::synchronize(*array.sync_event.cast::()) } - .map_err(|error| vortex_err!("waiting for Arrow device batch: {error}"))?; - } - let array = &array.array; - assert!(array.dictionary.is_null()); - assert_eq!(array.offset, 0); - assert_eq!(array.null_count, 0); - assert_eq!(array.n_children, 2); - assert!(!array.children.is_null()); - // SAFETY: The live struct owns two children matching the previously checked schema. - let fields = unsafe { - let values = (*array.children).as_ref().expect("missing values child"); - let ids = (*array.children.add(1)) - .as_ref() - .expect("missing ids child"); - assert_eq!(values.length, array.length); - assert_eq!(ids.length, array.length); - vec![ - read_primitive::(values, Nullability::Nullable)?, - read_primitive::(ids, Nullability::NonNullable)?, - ] - }; - Ok(StructArray::try_new( - ["値.x", "ids"].into(), - fields, - usize::try_from(array.length)?, - Validity::NonNullable, - )? - .into_array()) -} - -fn read_projected_batches(stream: &mut ArrowDeviceArrayStream) -> VortexResult> { - let get_next = stream.get_next.expect("missing get_next"); - let mut batches = Vec::new(); - loop { - let mut array = ArrowDeviceArray::empty(); - // SAFETY: This live stream owns the callback; array is writable. - let status = unsafe { get_next(stream, &raw mut array) }; - vortex_ensure!(status == 0, "get_next failed: {}", stream_error(stream)); - if array.array.release.is_none() { - break; - } - // SAFETY: The fixture's schema was checked, and this batch remains live during readback. - let batch = unsafe { read_projected_batch(&array) }; - release_device_array(&mut array); - batches.push(batch?); - } - Ok(batches) -} - -fn check_projected_file( - input: StructArray, - block_rows: usize, - batch_rows: usize, - expected_lengths: &[usize], -) -> VortexResult<()> { - let session = session().with_some(CudaSession::try_default()?); - register_cuda_layout(&session); - let columns = ["値.x", "ids"]; - let expected = input.project(names(&columns)?.as_ref())?.into_array(); - let mut file = NamedTempFile::new()?; - file.write_all(&file_bytes( - &session, - input.into_array(), - cuda_write_strategy(&session, block_rows), - )?)?; - let path = file - .path() - .to_str() - .ok_or_else(|| vortex_err!("non-UTF-8 test path"))?; - let options = vx_cuda_scan_options { - batch_rows, - ..Default::default() - }; - let mut stream = open_stream(&session, path, &options, &columns); - // The stream must retain its session state after the caller releases its session. - drop(session); - let schema = stream_schema(&mut stream.0); - let expected_fields = vec![ - Field::new("値.x", DataType::Int64, true), - Field::new("ids", DataType::UInt32, false), - ]; - assert_eq!(Schema::try_from(&schema)?, Schema::new(expected_fields)); - let batches = read_projected_batches(&mut stream.0)?; - let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect(); - assert_eq!(lengths, expected_lengths); - let actual = ChunkedArray::try_new(batches, expected.dtype().clone())?.into_array(); - assert_arrays_eq!( - actual, - expected, - &mut VortexSession::default().create_execution_ctx() - ); - Ok(()) -} - -#[cuda_test] -fn test_projection_gpu_subdivides_large_blocks() -> VortexResult<()> { - check_projected_file(table()?, 5, 2, &[2, 2, 1]) -} - -#[cuda_test] -fn test_projection_gpu_preserves_small_block_boundaries() -> VortexResult<()> { - check_projected_file(table()?, 2, 0, &[2, 2, 1]) -} - -#[cuda_test] -fn test_projection_gpu_exact_batch_rows_with_final_tail() -> VortexResult<()> { - let input = StructArray::try_new( - ["ids", "値.x"].into(), - vec![ - PrimitiveArray::from_iter(0u32..1_000).into_array(), - PrimitiveArray::from_option_iter( - (0i64..1_000).map(|value| (value % 2 == 0).then_some(value)), - ) - .into_array(), - ], - 1_000, - Validity::NonNullable, - )?; - check_projected_file(input, 1_000, 300, &[300, 300, 300, 100]) -} From 9085f45ac2ea28cd3fdc947ecb5cac172a9d8748 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:36:31 +0000 Subject: [PATCH 30/35] docs Signed-off-by: Alexander Droste --- vortex-cuda/src/session.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index 498b0516626..2d4ff54e65f 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -40,10 +40,19 @@ pub enum VarBinExportLayout { VarBinView, } -/// Dictionary policy for Arrow Device exports, including dictionaries nested in structs and lists. +/// Controls whether Arrow Device exports keep dictionary encoding or fully decode it, +/// including dictionaries nested inside structs and lists. +/// +/// For example, indices `[0, 1, 0]` and dictionary values `["apple", "pear"]` export as: +/// +/// - [`Preserve`](Self::Preserve): separate index and dictionary-value arrays, with an Arrow +/// dictionary type. +/// - [`Decode`](Self::Decode): the plain string array `["apple", "pear", "apple"]`, with no +/// dictionary indices or dictionary child. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum DictionaryExport { - /// Keep separate indices and dictionary values; the Arrow schema includes the index type. + /// Keep separate index and dictionary-value arrays rather than expanding repeated values. + /// The Arrow schema retains the dictionary type, including its index type. #[default] Preserve, /// Expand dictionaries into plain values for a stable schema across batches. From c70e408689dab3bd296da8fbb6f102e58d42e8e8 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:50:06 +0000 Subject: [PATCH 31/35] docs(cuda): shorten scan and export API documentation Signed-off-by: Alexander Droste --- vortex-cuda/ffi/README.md | 8 ++--- vortex-cuda/ffi/cinclude/vortex_cuda.h | 49 +++++++++---------------- vortex-cuda/ffi/src/lib.rs | 50 +++++++++----------------- vortex-cuda/src/arrow/mod.rs | 20 ++++------- 4 files changed, 40 insertions(+), 87 deletions(-) diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 1b55b9b740c..80e442852f7 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -30,9 +30,5 @@ cache for pooled data-plane reads. Footer and zone-map reads remain buffered on ## Generated header -`build.rs` generates `cinclude/vortex_cuda.h` using cbindgen on stable Rust. Edit `src/lib.rs` -or `cbindgen.toml`, not the header, and commit regenerated output. The header keeps cbindgen's -formatting and is excluded from clang-format through generated markers. - -Unchanged output is not rewritten; changed output is published atomically. CUDA CI checks -for header drift after building the FFI crate. +`build.rs` generates `cinclude/vortex_cuda.h` with cbindgen. Edit `src/lib.rs` or `cbindgen.toml` +and commit the regenerated header. It skips clang-format; CUDA CI checks for header drift. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index b4118ab1910..2fb8ae8054e 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -131,23 +131,14 @@ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session vx_error **error_out); /** - * Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. - * - * Requires CUDA-supported encodings/layouts, such as files from `vx_cuda_array_sink_open_file`. - * Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers - * reused across scans with the same CUDA session. - * - * Dictionaries, including nested children, decode on CUDA to a stable plain Arrow schema without - * changing session policy. Decoding requires CUDA support and may increase device memory use. - * - * Returns `0` with an owned `out_stream`; release it and each batch via their Arrow callbacks. - * Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. + * Scan a local CUDA-readable file with buffered I/O into an Arrow C Device stream. + * Dictionaries export as plain values. Returns `0` on success, `1` on error. + * Release the stream and batches via Arrow callbacks; free errors with `vx_error_free`. * * # Safety * - * `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the - * duration of this call and contain UTF-8. `out_stream` must be a valid writable pointer. If - * `error_out` is non-null, it must be valid for writing one error pointer. + * `session` must be a live borrowed `vortex-ffi` handle; `path` must contain readable UTF-8 + * for this call. `out_stream` and non-null `error_out` must point to writable output storage. */ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, vx_view path, @@ -155,11 +146,7 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session, vx_error **error_out); /** - * Scan a local Vortex file with exact row batches and a possibly smaller final batch. - * - * Uses `vx_cuda_scan_path_arrow_device_stream`'s export and ownership rules. - * `batch_rows` follows `vx_cuda_scan_options`: zero preserves layout boundaries; nonzero counts - * ignore them and may require unsupported CUDA `Chunked` concatenation. + * Like `vx_cuda_scan_path_arrow_device_stream`, with `batch_rows` as in `vx_cuda_scan_options`. * * # Safety * @@ -190,18 +177,16 @@ int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session /** * Scan a local Vortex file with ordered top-level column projection. * - * Same options, ownership, and file requirements as - * `vx_cuda_scan_path_arrow_device_stream_with_options`. Projection precedes decoding and skips - * unselected column I/O when the file layout stores columns separately. - * Names are literal and case-sensitive; a nonempty projection rejects unknown/duplicate names - * and non-struct files. `ncolumns == 0` ignores `columns` and selects all. Names are copied; - * empty files retain the projected schema. Errors leave `out_stream` unchanged. + * Otherwise follows `vx_cuda_scan_path_arrow_device_stream_with_options`. + * Names are literal, case-sensitive, and unique; nonempty projections require a struct file. + * Zero `ncolumns` ignores `columns` and selects all. Unknown names fail. + * Skips unselected column I/O where the layout allows. Errors leave `out_stream` unchanged. * * # Safety * - * In addition to `vx_cuda_scan_path_arrow_device_stream_with_options`'s requirements, - * nonzero `ncolumns` requires that many initialized, aligned `vx_view` values at `columns`. - * Each name borrows `len` readable UTF-8 bytes for this call; null is allowed only for zero length. + * Same requirements as `vx_cuda_scan_path_arrow_device_stream_with_options`. + * For nonzero `ncolumns`, `columns` must reference that many valid, aligned `vx_view` values. + * Name bytes must be readable UTF-8 for this call; null is allowed only for zero length. */ int vx_cuda_scan_path_arrow_device_stream_projected(const vx_session *session, vx_view path, @@ -236,12 +221,10 @@ int vx_cuda_array_export_arrow_device(const vx_session *session, vx_error **error_out); /** - * Consume a Vortex partition and scan it as an Arrow C Device stream. + * Scan a Vortex partition as an Arrow C Device stream. * - * Consumes `partition` on success or failure; never free or reuse it afterward. - * Returns `0` with an owned `out_stream` retaining the scan iterator. Release the stream and - * each produced batch via their Arrow release callbacks. - * Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. + * Consumes `partition`, even on error. Return codes and output ownership follow + * `vx_cuda_scan_path_arrow_device_stream`. * * # Safety * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 2f38cb30b57..3d218412cb1 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -157,23 +157,14 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( }) } -/// Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. -/// -/// Requires CUDA-supported encodings/layouts, such as files from `vx_cuda_array_sink_open_file`. -/// Footer/zone-map reads stay on the host; data reaches the GPU through pinned staging buffers -/// reused across scans with the same CUDA session. -/// -/// Dictionaries, including nested children, decode on CUDA to a stable plain Arrow schema without -/// changing session policy. Decoding requires CUDA support and may increase device memory use. -/// -/// Returns `0` with an owned `out_stream`; release it and each batch via their Arrow callbacks. -/// Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. +/// Scan a local CUDA-readable file with buffered I/O into an Arrow C Device stream. +/// Dictionaries export as plain values. Returns `0` on success, `1` on error. +/// Release the stream and batches via Arrow callbacks; free errors with `vx_error_free`. /// /// # Safety /// -/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the -/// duration of this call and contain UTF-8. `out_stream` must be a valid writable pointer. If -/// `error_out` is non-null, it must be valid for writing one error pointer. +/// `session` must be a live borrowed `vortex-ffi` handle; `path` must contain readable UTF-8 +/// for this call. `out_stream` and non-null `error_out` must point to writable output storage. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( session: *const vx_session, @@ -193,11 +184,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( } } -/// Scan a local Vortex file with exact row batches and a possibly smaller final batch. -/// -/// Uses `vx_cuda_scan_path_arrow_device_stream`'s export and ownership rules. -/// `batch_rows` follows `vx_cuda_scan_options`: zero preserves layout boundaries; nonzero counts -/// ignore them and may require unsupported CUDA `Chunked` concatenation. +/// Like `vx_cuda_scan_path_arrow_device_stream`, with `batch_rows` as in `vx_cuda_scan_options`. /// /// # Safety /// @@ -258,18 +245,16 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_optio /// Scan a local Vortex file with ordered top-level column projection. /// -/// Same options, ownership, and file requirements as -/// `vx_cuda_scan_path_arrow_device_stream_with_options`. Projection precedes decoding and skips -/// unselected column I/O when the file layout stores columns separately. -/// Names are literal and case-sensitive; a nonempty projection rejects unknown/duplicate names -/// and non-struct files. `ncolumns == 0` ignores `columns` and selects all. Names are copied; -/// empty files retain the projected schema. Errors leave `out_stream` unchanged. +/// Otherwise follows `vx_cuda_scan_path_arrow_device_stream_with_options`. +/// Names are literal, case-sensitive, and unique; nonempty projections require a struct file. +/// Zero `ncolumns` ignores `columns` and selects all. Unknown names fail. +/// Skips unselected column I/O where the layout allows. Errors leave `out_stream` unchanged. /// /// # Safety /// -/// In addition to `vx_cuda_scan_path_arrow_device_stream_with_options`'s requirements, -/// nonzero `ncolumns` requires that many initialized, aligned `vx_view` values at `columns`. -/// Each name borrows `len` readable UTF-8 bytes for this call; null is allowed only for zero length. +/// Same requirements as `vx_cuda_scan_path_arrow_device_stream_with_options`. +/// For nonzero `ncolumns`, `columns` must reference that many valid, aligned `vx_view` values. +/// Name bytes must be readable UTF-8 for this call; null is allowed only for zero length. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_projected( session: *const vx_session, @@ -349,7 +334,6 @@ unsafe fn scan_columns(columns: *const vx_view, ncolumns: usize) -> VortexResult Ok(names.into()) } -/// Apply projection before column reads. fn projected_scan( file: &VortexFile, columns: FieldNames, @@ -462,12 +446,10 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_export_arrow_device( }) } -/// Consume a Vortex partition and scan it as an Arrow C Device stream. +/// Scan a Vortex partition as an Arrow C Device stream. /// -/// Consumes `partition` on success or failure; never free or reuse it afterward. -/// Returns `0` with an owned `out_stream` retaining the scan iterator. Release the stream and -/// each produced batch via their Arrow release callbacks. -/// Returns `1` on error, writing a `vx_error` if `error_out` is non-null; free it with `vx_error_free`. +/// Consumes `partition`, even on error. Return codes and output ownership follow +/// `vx_cuda_scan_path_arrow_device_stream`. /// /// # Safety /// diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index d5e035ada1d..b7e584c32f7 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -114,8 +114,7 @@ impl ArrowArray { } impl ArrowDeviceArray { - /// Create a released array with zeroed device metadata and a null sync event. - /// Use as callback output storage or an end-of-stream marker base; no device is selected. + /// Create a released array with zeroed device metadata for callback output storage. pub fn empty() -> Self { Self { array: ArrowArray::empty(), @@ -490,15 +489,11 @@ impl Drop for DeviceArrayStreamPrivateData { pub trait DeviceArrayStreamExt { /// Export this stream as an [`ArrowDeviceArrayStream`]. /// - /// Arrays are exported by reusing one [`CudaExecutionCtx`], and every produced - /// [`ArrowDeviceArray`] must remain on the CUDA device captured at stream construction. The - /// returned [`ArrowDeviceArrayStream`] owns the Vortex stream and must be released through its - /// embedded `release` callback. + /// Reuses one [`CudaExecutionCtx`] and rejects changes of device or Arrow schema. + /// The returned stream owns the input; release it through its `release` callback. /// - /// All arrays must share the `get_schema` schema. By default, it comes from the first array - /// (the logical dtype for empty streams); chunks exporting different Arrow types are rejected. - /// With [`DictionaryExport::Decode`], the dtype determines a plain schema independent of chunk - /// encoding or dictionary index width. `get_schema` does not pull or export a batch; + /// By default, the schema comes from the first batch, or the dtype for empty streams. + /// With [`DictionaryExport::Decode`], it comes from the dtype without polling a batch; /// read/decode errors surface in `get_next`. /// /// Drive the returned stream from one thread. `runtime` must be the runtime that owns the @@ -523,10 +518,7 @@ impl DeviceArrayStreamExt for SendableArrayStream { } impl ArrowDeviceArrayStream { - /// Export a stream using an owned context, retaining its session and per-context configuration. - /// - /// The schema, runtime, and release requirements of - /// [`DeviceArrayStreamExt::export_device_array_stream`] also apply here. + /// Like [`DeviceArrayStreamExt::export_device_array_stream`], using an owned execution context. pub fn new( array_stream: SendableArrayStream, ctx: CudaExecutionCtx, From 9c800c0660addb278f52e7f134ca3fd0943e1a3e Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:54:31 +0000 Subject: [PATCH 32/35] fix(cuda): initialize layout registration guard explicitly Signed-off-by: Alexander Droste --- vortex-cuda/src/layout.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index a8ddf10eb23..7cc7d0322da 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -575,9 +575,15 @@ pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc); +impl Default for CudaLayoutRegistration { + fn default() -> Self { + Self(Arc::new(Once::new())) + } +} + impl SessionVar for CudaLayoutRegistration { fn as_any(&self) -> &dyn Any { self From 9eb44758d93644e2223d7f6de8c8020ce7fad523 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 10:59:16 +0000 Subject: [PATCH 33/35] fix(cuda): remove redundant FFI error free qualification Signed-off-by: Alexander Droste --- vortex-cuda/ffi/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 3d218412cb1..0647b2936a7 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -1267,7 +1267,7 @@ mod tests { assert_eq!(status, VX_CUDA_ERR); assert!(!error.is_null()); unsafe { - vortex_ffi::vx_error_free(error); + vx_error_free(error); free_test_array(array); free_test_session(session); } From aaa140b3a30be59c61ebd43a271127d73b5d4cb7 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 11:00:21 +0000 Subject: [PATCH 34/35] docs(ffi): document error message safety requirements Signed-off-by: Alexander Droste --- vortex-ffi/src/error.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vortex-ffi/src/error.rs b/vortex-ffi/src/error.rs index 9fac9470727..8f13a699e59 100644 --- a/vortex-ffi/src/error.rs +++ b/vortex-ffi/src/error.rs @@ -159,8 +159,11 @@ pub fn try_or( } } -/// Return error message for this error. -/// Returned view is valid while "error" is valid. +/// Return a message view borrowed from `error`. +/// +/// # Safety +/// +/// `error` must be a non-null, live `vx_error` handle and remain live while the view is used. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_error_message(error: *const vx_error) -> vx_view { vx_view::from_str(&vx_error::as_ref(error).message) From 39ddfa5b1a8a94b3b7637c00c04392438b6730ad Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Tue, 22 Sep 2026 11:08:43 +0000 Subject: [PATCH 35/35] fix(ffi): regenerate header after error safety docs Signed-off-by: Alexander Droste --- vortex-ffi/cinclude/vortex.h | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index fa214e8efbe..384dc29305c 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -1044,8 +1044,11 @@ vx_dtype_from_arrow_schema(const vx_session *session, FFI_ArrowSchema *schema, v void vx_error_free(const vx_error *ptr); /** - * Return error message for this error. - * Returned view is valid while "error" is valid. + * Return a message view borrowed from `error`. + * + * # Safety + * + * `error` must be a non-null, live `vx_error` handle and remain live while the view is used. */ vx_view vx_error_message(const vx_error *error);