diff --git a/.github/workflows/cuda.yaml b/.github/workflows/cuda.yaml index 4e14ea370a5..5aa30f436a9 100644 --- a/.github/workflows/cuda.yaml +++ b/.github/workflows/cuda.yaml @@ -45,6 +45,8 @@ jobs: - "vortex-test/**" - "pyproject.toml" - "uv.lock" + - "Cargo.toml" + - "Cargo.lock" - ".github/workflows/**" cuda-build-lint: @@ -71,6 +73,8 @@ jobs: 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 \ diff --git a/Cargo.lock b/Cargo.lock index 275db7ef8ee..92f52bd0198 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10911,7 +10911,10 @@ name = "vortex-cuda-ffi" version = "0.1.0" dependencies = [ "arrow-schema 59.3.0", + "cbindgen", + "cudarc", "futures", + "tempfile", "vortex", "vortex-cuda", "vortex-cuda-macros", diff --git a/vortex-cuda/ffi/Cargo.toml b/vortex-cuda/ffi/Cargo.toml index c1e1efdb23c..6ec839c100f 100644 --- a/vortex-cuda/ffi/Cargo.toml +++ b/vortex-cuda/ffi/Cargo.toml @@ -22,8 +22,13 @@ vortex-cuda = { path = ".." } vortex-ffi = { path = "../../vortex-ffi" } [dev-dependencies] +cudarc = { workspace = true } +tempfile = { workspace = true } 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..80e442852f7 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -27,3 +27,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. + +## Generated header + +`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/build.rs b/vortex-cuda/ffi/build.rs new file mode 100644 index 00000000000..7a915f65296 --- /dev/null +++ b/vortex-cuda/ffi/build.rs @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: Apache-2.0 +// 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; + +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 header = "cinclude/vortex_cuda.h"; + let mut generated = Vec::new(); + // 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); + match fs::read(header) { + 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, &generated) + .map_err(|error| format!("failed to publish {header}: {error}"))?; + Ok(()) +} + +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()); + 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), + } + }; + + // 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() + && let Err(error) = fs::remove_file(&temporary) + { + eprintln!("failed to remove temporary header {temporary}: {error}"); + } + result +} diff --git a/vortex-cuda/ffi/cbindgen.toml b/vortex-cuda/ffi/cbindgen.toml new file mode 100644 index 00000000000..d529d4db99c --- /dev/null +++ b/vortex-cuda/ffi/cbindgen.toml @@ -0,0 +1,77 @@ +# 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 + +// clang-format off + +#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 +""" + +trailer = "// clang-format on" + +# 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..2fb8ae8054e 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -1,7 +1,12 @@ // 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 + #include #include @@ -10,10 +15,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 +57,56 @@ 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 rejected. + */ + uint32_t flags; + /** + * 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; + +#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. + * Returns an owned handle with reusable CUDA state, 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, @@ -79,14 +116,13 @@ 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. + * 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 * - * 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. + * 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 +131,14 @@ 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. - * - * 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. + * 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`. * - * 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 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, @@ -126,14 +146,11 @@ 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. + * Like `vx_cuda_scan_path_arrow_device_stream`, with `batch_rows` as in `vx_cuda_scan_options`. * - * `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`. + * # Safety * - * 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. + * 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 +159,60 @@ 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. * - * 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. + * 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`. */ 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. + * + * 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 + * + * 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, + 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, @@ -173,18 +221,16 @@ 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. * - * This function takes ownership of `partition`. Callers must not free or reuse - * it after calling this function, regardless of success or failure. + * Consumes `partition`, even on error. Return codes and output ownership follow + * `vx_cuda_scan_path_arrow_device_stream`. * - * 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 +238,7 @@ int vx_cuda_partition_scan_arrow_device_stream(const vx_session *session, vx_error **error_out); #ifdef __cplusplus -} -#endif +} // extern "C" +#endif // __cplusplus + +// clang-format on diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 44fc7d87938..0647b2936a7 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -10,29 +10,35 @@ 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::editions::ComponentKind; -use vortex::editions::EditionSessionExt; +use vortex::dtype::FieldName; +use vortex::dtype::FieldNames; 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::WriteStrategyBuilder; +use vortex::file::VortexFile; use vortex::io::runtime::BlockingRuntime; +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; +use vortex_cuda::DictionaryExport; use vortex_cuda::PooledFileReadAtOptions; 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; @@ -52,8 +58,10 @@ 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; +/// 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; + const VX_CUDA_SCAN_KNOWN_FLAGS: u32 = VX_CUDA_SCAN_FLAG_DIRECT_IO; /// Options for scanning a CUDA-compatible Vortex file. @@ -62,26 +70,23 @@ 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 rejected. pub flags: u32, - /// Number of rows in each output batch. Zero uses layout-derived splitting. + /// 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, } -/// 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 { +fn session_with_cuda(session: &VortexSession) -> &VortexSession { session.get::(); register_cuda_layout(session); - Ok(session.clone()) + session } /// 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 /// @@ -102,15 +107,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, @@ -118,23 +121,19 @@ 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) } } /// 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. -/// -/// 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`]. +/// 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 /// -/// `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,53 +143,28 @@ 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), + // 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, + 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. -/// -/// 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. -/// -/// On error returns `1` and, when `error_out` is non-null, writes a `vx_error` (free 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, @@ -198,6 +172,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, @@ -209,19 +184,11 @@ 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 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`]. +/// Like `vx_cuda_scan_path_arrow_device_stream`, with `batch_rows` as in `vx_cuda_scan_options`. /// /// # 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, @@ -234,6 +201,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, @@ -245,18 +213,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 +228,159 @@ 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. +/// +/// 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 +/// +/// 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, + 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) }?)?; - let options = unsafe { scan_options(options) }?; - let array_stream = ffi_runtime().block_on(async { - let file = 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() .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()); + // SAFETY: The output is non-null and the caller guarantees writable storage. 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); + let mut seen = HashSet::<&str>::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!(seen.insert(name), "duplicate CUDA scan column: {name:?}"); + names.push(FieldName::from(name)); + } + Ok(names.into()) +} + +fn projected_scan( + file: &VortexFile, + columns: FieldNames, + batch_rows: usize, +) -> VortexResult> { + let mut scan = file.scan()?; + if !columns.is_empty() { + let projection = select(columns, root()).optimize(file.dtype())?; + scan = scan.with_projection(projection.bind(file.dtype())?); + } + let split_by = if batch_rows == 0 { + SplitBy::Layout + } else { + SplitBy::RowCount(batch_rows) + }; + Ok(scan.with_split_by(split_by)) +} + struct CudaScanOptions { read_at_options: PooledFileReadAtOptions, batch_rows: usize, } +/// 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 rejected. +/// +/// # 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 (flags, batch_rows) = if options.is_null() { - (0, 0) - } else { - let options = unsafe { &*options }; - (options.flags, options.batch_rows) - }; + 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!( - flags & !VX_CUDA_SCAN_KNOWN_FLAGS == 0, + options.flags & !VX_CUDA_SCAN_KNOWN_FLAGS == 0, "unsupported CUDA scan option flags: {:#x}", - flags & !VX_CUDA_SCAN_KNOWN_FLAGS + options.flags & !VX_CUDA_SCAN_KNOWN_FLAGS ); - 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 +397,15 @@ unsafe fn scan_options(options: *const vx_cuda_scan_options) -> VortexResult *mut vx_session { + Box::into_raw(Box::new(session)).cast::() } - #[test] - fn rejects_unknown_scan_option_flags() { - let options = vx_cuda_scan_options { - flags: 1 << 31, - ..Default::default() + 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!(unsafe { scan_options(&raw const options) }.is_err()); + assert_eq!(status, VX_CUDA_OK); + assert!(error.is_null()); + (schema, device_array) } - #[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() + 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!( - unsafe { scan_options(&raw const options) }?.read_at_options, - PooledFileReadAtOptions::default().with_direct_io() + 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 maps_batch_rows_scan_option() -> VortexResult<()> { - let options = vx_cuda_scan_options { - batch_rows: 8192, - ..Default::default() + fn test_projection_rejects_invalid_names_and_counts() { + let invalid_utf8 = vx_view { + ptr: [0xffu8].as_ptr().cast(), + len: 1, }; - assert_eq!( - unsafe { scan_options(&raw const options) }?.batch_rows, - 8192 - ); + 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(()) } - fn test_session(session: VortexSession) -> *mut vx_session { - Box::into_raw(Box::new(session)).cast::() + #[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(()) } - unsafe fn free_test_session(session: *mut vx_session) { - unsafe { drop(Box::from_raw(session.cast::())) }; + #[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(()) } - fn test_array(array: impl IntoArray) -> *const vx_array { - Arc::into_raw(Arc::new(array.into_array())).cast::() + #[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}" + ); } - unsafe fn free_test_array(array: *const vx_array) { - unsafe { Arc::decrement_strong_count(array.cast::()) }; + #[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(()) } - unsafe fn release_schema(schema: &mut FFI_ArrowSchema) { - unsafe { - if let Some(release) = schema.release { - release(schema); - } + #[test] + fn test_scan_options_default_to_buffered_io() -> VortexResult<()> { + let options = vx_cuda_scan_options::default(); + + 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(()) } - unsafe fn release_device_array(array: &mut ArrowDeviceArray) { - unsafe { - if let Some(release) = array.array.release { - release(&raw mut array.array); - } + #[test] + fn test_maps_scan_options() -> VortexResult<()> { + let buffered = PooledFileReadAtOptions::default(); + for (flags, batch_rows, read_at_options) in [ + (0, 8192, buffered), + #[cfg(target_os = "linux")] + (VX_CUDA_SCAN_FLAG_DIRECT_IO, 0, 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(()) } - 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], + #[test] + fn test_scan_options_reject_unknown_flags() { + for flags in [1 << 1, VX_CUDA_SCAN_FLAG_DIRECT_IO | (1 << 1)] { + 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_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 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(), + )); + let array = DictArray::try_new( + PrimitiveArray::from_iter([1u8, 0, 1]).into_array(), + PrimitiveArray::from_iter([10i32, 20]).into_array(), + )? + .into_array(); + let mut exported = + ffi_runtime().block_on(array.export_device_array_with_schema(&mut ctx))?; + release_device_array(&mut exported.array); + assert_eq!( + Field::try_from(&exported.schema)?.data_type(), + &DataType::Int32 + ); + Ok(()) + } + #[cuda_test] fn test_export_primitive_arrow_device() { - 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(); - - 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 field = Field::try_from(&schema).expect("schema should be a field"); assert_eq!(field.name(), ""); @@ -559,7 +1182,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(), @@ -570,21 +1192,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); @@ -624,20 +1233,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); @@ -655,7 +1252,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 { @@ -670,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); } 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/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 4a96187a9db..e872cfdab51 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,13 @@ 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())? + } + // Decode schemas use only dtype; keep encodings for recursion and 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 +212,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 +1505,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 +1516,11 @@ 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::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; use crate::session::CudaSession; use crate::session::VarBinExportLayout; @@ -1752,39 +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()) - } - - 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()) + Ok(Buffer::::from_byte_buffer(private_data_buffer_bytes(array, buffer_idx)?).to_vec()) } // Assert Arrow Binary export uses the standard null bitmap, i32 offsets, and values layout. @@ -2590,9 +2580,27 @@ 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)?; + // 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); 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)?), + 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..a985601debb --- /dev/null +++ b/vortex-cuda/src/arrow/dictionary_tests.rs @@ -0,0 +1,421 @@ +// 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::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::error::vortex_bail; + +use super::tests::last_error; +use super::tests::private_data_buffer_bytes as buffer; +use super::*; +use crate::CudaSession; + +/// 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() { + 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) } +} + +/// 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 = Buffer::::from_byte_buffer(buffer(array, 1)?).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(Field::try_from(&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. +fn upload_chunks( + chunks: Vec, + ctx: &mut CudaExecutionCtx, +) -> VortexResult>> { + let mut device_chunks = Vec::new(); + for chunk in chunks { + let chunk = upload(chunk, ctx)?; + 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, + #[values(false, true)] schema_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 = upload_chunks(chunks, &mut ctx)?; + let mut stream = ArrayStreamAdapter::new(expected.dtype().clone(), stream::iter(chunks)) + .boxed() + .export_device_array_stream(&session, &runtime)?; + let plain_schema = Field::try_from(&arrow_schema_for_array(&expected, &mut ctx)?)?; + if schema_first { + assert_eq!(get_schema(&mut stream)?, plain_schema); + } + 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); + } + 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()); + // SAFETY: This is the live stream's final use. + unsafe { stream.release.expect("missing release")(&raw mut stream) }; + Ok(()) +} + +#[rstest] +#[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, + #[case] 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 { + assert_eq!( + get_schema(&mut stream)?, + 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(); + 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 + .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"); + state.get_or_init_schema()?; + + let error = state + .export_stream_array(PrimitiveArray::from_iter([10u32, 20, 30]).into_array()) + .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) + .expect_err("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() + .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)?; + 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 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); + 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()); + release_device_array(&mut preserved); + + 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); + 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)?; + + 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 = 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)?; + assert_eq!( + get_schema(&mut stream)?.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..b7e584c32f7 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 { @@ -111,10 +114,8 @@ impl ArrowArray { } impl ArrowDeviceArray { - /// A zeroed device array: an empty Arrow array with no device. Used as a - /// callback output placeholder and as the basis for the end-of-stream - /// marker. - fn empty() -> Self { + /// Create a released array with zeroed device metadata for callback output storage. + pub fn empty() -> Self { Self { array: ArrowArray::empty(), device_id: 0, @@ -362,12 +363,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( @@ -383,15 +387,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() + }), } } @@ -404,36 +412,55 @@ impl DeviceArrayStreamPrivateData { array.dtype() ); + 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); + } + return Ok(device_array); + } + 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 { - Ok(exported_schema) => exported_schema, + // 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); + } + 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 self.schema.is_none() { + 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) } - /// 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 {}", @@ -445,19 +472,7 @@ impl DeviceArrayStreamPrivateData { self.device_id, device_array.device_id ); - - 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) + Ok(()) } } @@ -474,14 +489,12 @@ 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. /// - /// 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. + /// 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 /// underlying scan tasks and per-array exports. @@ -499,38 +512,39 @@ 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 { + /// Like [`DeviceArrayStreamExt::export_device_array_stream`], using an owned execution context. + 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(), + } } } @@ -547,15 +561,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 { @@ -570,23 +575,7 @@ 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, 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, @@ -600,17 +589,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. @@ -627,24 +616,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. @@ -690,6 +672,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 +758,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) @@ -943,6 +919,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; @@ -952,17 +929,34 @@ 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; - fn last_error(stream: &mut ArrowDeviceArrayStream) -> VortexResult { + /// 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, + ) -> 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 .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..7cc7d0322da 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -7,6 +7,7 @@ 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; @@ -28,12 +29,21 @@ 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; +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; use vortex::error::vortex_panic; +use vortex::file::WriteStrategyBuilder; use vortex::layout::Layout; use vortex::layout::LayoutChildType; use vortex::layout::LayoutDeserializeArgs; @@ -55,11 +65,14 @@ 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; 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; @@ -534,12 +547,309 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { result } -/// Register the [`CudaFlatLayoutEncoding`] in the session's layout registry. +/// Build a CUDA-flat writer using only CUDA-compatible, session-enabled array encodings. /// -/// Call this alongside [`crate::initialize_cuda`] when setting up a CUDA-enabled session. +/// 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 + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + 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 { + strategy.with_btrblocks_builder(builder).build() + } else { + // An opaque compressor keeps IntDict; disabling the probe avoids u16-sized outer blocks. + strategy + .with_compressor(builder.build()) + .with_probe_compressor(BtrBlocksCompressorBuilder::empty().build()) + .with_row_block_size(block_rows) + .with_data_block_target_bytes(None) + .build() + } +} + +#[derive(Clone, Debug)] +struct CudaLayoutRegistration(Arc); + +impl Default for CudaLayoutRegistration { + fn default() -> Self { + Self(Arc::new(Once::new())) + } +} + +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 { + id: CUDA_EDITION, + min_library_version: None, + }, + added: &[EditionMember::layout(&"vortex.cuda_flat")], +}; + +/// Register [`CudaFlat`] and its draft `cuda` edition once per session. +/// +/// Enables a newly registered edition only if no `cuda` edition is selected; otherwise preserves +/// writer policy, including on repeated calls. The draft has no cross-version compatibility +/// guarantee. Readers must also register the layout. +/// +/// Call alongside [`crate::initialize_cuda`]; registration itself needs no GPU. pub fn register_cuda_layout(session: &VortexSession) { - use vortex::layout::session::LayoutSessionExt; - session - .layouts() - .register(LayoutEncodingRef::new_ref(&CudaFlat)); + // Editions are published before their members; concurrent callers must wait for both. + session.get::().0.call_once(|| { + session + .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"); + 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)] +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::editions::CORE_2025_05_0; + 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::scan::split_by::SplitBy; + + use super::*; + + fn repeated_ids(unique: i64, rows: usize) -> VortexResult { + // Wide, shuffled 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, + block_rows: usize, + ) -> VortexResult { + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .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_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()?) { + // Exclude zone maps and dictionary values from data row counts. + if !matches!(kind, LayoutChildType::Auxiliary(_)) { + rows.extend(data_block_rows(&child)?); + } + } + Ok(rows) + } + + #[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(), block_rows).await?; + + let batches: Vec<_> = file + .scan()? + .with_split_by(SplitBy::Layout) + .into_array_stream()? + .try_collect() + .await?; + 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 { + // 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; + } + + 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 u16 cardinality while remaining eligible for outer dictionaries. + let block_rows = 70_000 * 8; + let input = repeated_ids(70_000, block_rows)?; + let file = write_file(&session, input, block_rows).await?; + assert_eq!( + data_block_rows(file.footer().layout())?, + [block_rows as u64] + ); + Ok(()) + }) + } + + #[test] + fn test_concurrent_cuda_registration_preserves_edition_policy() -> VortexResult<()> { + let session = VortexSession::default(); + session.enable_edition(CORE_2025_05_0)?; + let mut expected_editions = session.enabled_editions().editions(); + expected_editions.push(CUDA_EDITION); + expected_editions.sort_unstable(); + 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 + .enabled_component_ids(ComponentKind::Layout) + .contains(&CudaFlat.id()) + ); + }); + } + }); + + 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::Array), + expected_arrays + ); + 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); + + 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 + ); + Ok(()) + } + + #[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)?; + 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); + + 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 + ); + 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..2d4ff54e65f 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -40,6 +40,26 @@ 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. +#[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. + #[default] + Preserve, + /// Expand dictionaries into plain values for a stable 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 +70,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 +98,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 +115,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 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); 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) diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 1ef423bb383..6646d0740fc 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -27,11 +27,13 @@ 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; 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;