Expose CUDA-compatible file sink through FFI

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<SchemeId> = 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));
+        }
+    }
 }

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.

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.
  *

vortex-cuda/ffi/src/lib.rs

@@ -10,27 +10,36 @@
 
 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;
 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<VortexSession> {
     session.get::<CudaSession>();
+    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

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;

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<dyn LayoutStrategy>,
+) -> 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) }
     })
 }