From 9a15dd8d58ea1fc924ca1e9a0ee1d071d5df980e Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Thu, 16 Jul 2026 11:18:28 +0100 Subject: [PATCH] Expose CUDA-compatible file sink through FFI Signed-off-by: Onur Satici --- vortex-btrblocks/src/builder.rs | 31 ++++++++++- vortex-cuda/ffi/README.md | 3 ++ vortex-cuda/ffi/cinclude/vortex_cuda.h | 12 +++++ vortex-cuda/ffi/src/lib.rs | 42 ++++++++++++++- vortex-ffi/src/lib.rs | 6 ++- vortex-ffi/src/sink.rs | 75 ++++++++++++++++++-------- 6 files changed, 142 insertions(+), 27 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 26632e70860..e51fac9a9d5 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -165,10 +165,14 @@ impl BtrBlocksCompressorBuilder { // dictionary expansion at decode time, which is incompatible with // pure-GPU decompression paths. Strip whichever string-fragment // scheme is enabled by feature. - #[cfg_attr(not(feature = "unstable_encodings"), allow(unused_mut))] + #[cfg_attr( + not(any(feature = "pco", feature = "unstable_encodings")), + allow(unused_mut) + )] let mut excluded: Vec = vec![ integer::SparseScheme.id(), integer::IntRLEScheme.id(), + float::ALPRDScheme.id(), float::FloatRLEScheme.id(), float::NullDominatedSparseScheme.id(), string::StringDictScheme.id(), @@ -181,6 +185,8 @@ impl BtrBlocksCompressorBuilder { // is incompatible with pure-GPU decompression paths. #[cfg(feature = "unstable_encodings")] excluded.push(integer::DeltaScheme::default().id()); + #[cfg(feature = "pco")] + excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); let builder = self.exclude_schemes(excluded); #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] @@ -223,4 +229,27 @@ mod tests { let builder = BtrBlocksCompressorBuilder::default(); assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } + + #[test] + fn cuda_compatible_excludes_alprd() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + !builder + .schemes + .iter() + .any(|s| s.id() == float::ALPRDScheme.id()) + ); + } + + #[test] + #[cfg(feature = "pco")] + fn cuda_compatible_excludes_pco() { + let builder = BtrBlocksCompressorBuilder::default() + .with_new_scheme(&integer::PcoScheme) + .with_new_scheme(&float::PcoScheme) + .only_cuda_compatible(); + for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { + assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); + } + } } diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 3c030095374..9cc6b30dd94 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -16,3 +16,6 @@ Use this crate as the CUDA-enabled FFI artifact. Include both headers: and link the CUDA FFI library (`vortex_cuda_ffi`). Do not pass Vortex handles between independently linked Rust FFI libraries. Use `vx_cuda_session_new` to initialize CUDA once and reuse it across exports. + +Use `vx_cuda_array_sink_open_file` to open a standard Vortex file sink configured to produce CUDA-readable files. +Push host-resident arrays and close the sink using the standard `vx_array_sink_*` APIs. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 0b06a660afa..d3ca8293f34 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -63,6 +63,18 @@ struct ArrowDeviceArrayStream { */ 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. + */ +vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session, + vx_view path, + const vx_dtype *dtype, + vx_error **error_out); + /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 611f644e535..cfc76ef354a 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -10,10 +10,13 @@ use std::os::raw::c_int; use std::ptr; +use std::sync::Arc; use arrow_schema::ffi::FFI_ArrowSchema; +use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::error::VortexResult; use vortex::error::vortex_ensure; +use vortex::file::WriteStrategyBuilder; use vortex::session::SessionExt; use vortex::session::VortexSession; use vortex_cuda::CudaSession; @@ -21,16 +24,22 @@ 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::register_cuda_layout; use vortex_ffi::ffi_runtime; use vortex_ffi::try_or; use vortex_ffi::vx_array; use vortex_ffi::vx_array_ref; +use vortex_ffi::vx_array_sink; +use vortex_ffi::vx_array_sink_open_file_with_strategy; +use vortex_ffi::vx_dtype; use vortex_ffi::vx_error; use vortex_ffi::vx_partition; use vortex_ffi::vx_partition_into_array_stream; use vortex_ffi::vx_session; use vortex_ffi::vx_session_new_with; use vortex_ffi::vx_session_ref; +use vortex_ffi::vx_view; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; @@ -41,6 +50,7 @@ const VX_CUDA_ERR: c_int = 1; /// returns a new session cloned from `session` with a default [`CudaSession`] attached. fn session_with_cuda(session: &VortexSession) -> VortexResult { session.get::(); + register_cuda_layout(session); Ok(session.clone()) } @@ -59,11 +69,41 @@ pub unsafe extern "C-unwind" fn vx_cuda_session_new( try_or(error_out, ptr::null_mut(), || { let cuda_session = CudaSession::try_default()?; Ok(vx_session_new_with(|session| { - session.with_some(cuda_session) + let session = session.with_some(cuda_session); + register_cuda_layout(&session); + session })) }) } +/// 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. +/// +/// # 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. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( + session: *const vx_session, + path: vx_view, + dtype: *const vx_dtype, + error_out: *mut *mut vx_error, +) -> *mut vx_array_sink { + try_or(error_out, ptr::null_mut(), || { + session_with_cuda(unsafe { vx_session_ref(session) }?)?; + let strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().only_cuda_compatible()) + .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) + .build(); + unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy) } + }) +} + /// 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 diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index ef493a27b03..151359096d4 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -27,6 +27,7 @@ use std::sync::LazyLock; pub use array::vx_array; 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; @@ -37,13 +38,14 @@ pub use session::vx_session; pub use session::vx_session_free; pub use session::vx_session_new_with; pub use session::vx_session_ref; +pub use sink::vx_array_sink; +pub use sink::vx_array_sink_open_file_with_strategy; +pub use string::vx_view; use vortex::dtype::FieldName; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::io::runtime::current::CurrentThreadRuntime; -use crate::string::vx_view; - #[cfg(all(feature = "mimalloc", not(miri)))] #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; diff --git a/vortex-ffi/src/sink.rs b/vortex-ffi/src/sink.rs index f307cfc8bc0..e7a53a0c8dd 100644 --- a/vortex-ffi/src/sink.rs +++ b/vortex-ffi/src/sink.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use futures::SinkExt; use futures::TryStreamExt; use futures::channel::mpsc; @@ -14,10 +16,12 @@ use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; use vortex::file::WriteSummary; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutStrategy; use crate::RUNTIME; use crate::arc_wrapper; @@ -54,6 +58,52 @@ pub struct vx_array_sink { dtype: DType, } +/// Open a file sink with an explicit layout strategy. +/// +/// This is a Rust API for FFI crates layered on top of `vortex-ffi`; the base C API continues to +/// use the default write strategy. File creation and write errors are reported when the returned +/// sink is closed. +/// +/// # Safety +/// +/// `session`, `path`, and `dtype` must satisfy the same requirements as +/// `vx_array_sink_open_file`. +pub unsafe fn vx_array_sink_open_file_with_strategy( + session: *const vx_session, + path: vx_view, + dtype: *const vx_dtype, + strategy: Arc, +) -> VortexResult<*mut vx_array_sink> { + let session = vx_session::as_ref(session).clone(); + + if path.ptr.is_null() { + vortex_bail!("null path"); + } + let path = unsafe { path.as_str() }?.to_string(); + + let file_dtype = vx_dtype::as_ref(dtype); + // The channel size 32 was chosen arbitrarily. + let (sink, rx) = mpsc::channel(32); + let dtype = file_dtype.clone(); + let array_stream = ArrayStreamAdapter::new(dtype.clone(), rx.into_stream()); + + let writer_session = session.clone(); + let writer = session.handle().spawn(async move { + let mut file = async_fs::File::create(path).await?; + writer_session + .write_options() + .with_strategy(strategy) + .write(&mut file, array_stream) + .await + }); + + Ok(Box::into_raw(Box::new(vx_array_sink { + sink, + writer, + dtype, + }))) +} + /// Opens a writable array stream, where sink is used to push values into the stream. /// To close the stream close the sink with `vx_array_sink_close`. /// "path" is copied. @@ -65,29 +115,8 @@ pub unsafe extern "C-unwind" fn vx_array_sink_open_file( error_out: *mut *mut vx_error, ) -> *mut vx_array_sink { try_or_default(error_out, || { - let session = vx_session::as_ref(session).clone(); - - if path.ptr.is_null() { - vortex_bail!("null path"); - } - let path = unsafe { path.as_str() }?.to_string(); - - let file_dtype = vx_dtype::as_ref(dtype); - // The channel size 32 was chosen arbitrarily. - let (sink, rx) = mpsc::channel(32); - let dtype = file_dtype.clone(); - let array_stream = ArrayStreamAdapter::new(dtype.clone(), rx.into_stream()); - - let writer = session.handle().spawn(async move { - let mut file = async_fs::File::create(path).await?; - session.write_options().write(&mut file, array_stream).await - }); - - Ok(Box::into_raw(Box::new(vx_array_sink { - sink, - writer, - dtype, - }))) + let strategy = WriteStrategyBuilder::default().build(); + unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy) } }) }