Expose pooled CUDA file scans through FFI

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

vortex-cuda/ffi/README.md

@@ -19,3 +19,7 @@ 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.
+
+Use `vx_cuda_scan_path_arrow_device_stream` to read such a local file through pinned host buffers
+and receive an Arrow C Device stream. Reuse the same CUDA session across scans so the pinned buffer
+pool and CUDA state are reused as well.

vortex-cuda/ffi/cinclude/vortex_cuda.h

@@ -75,6 +75,21 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session,
                                             const vx_dtype *dtype,
                                             vx_error **error_out);
 
+/**
+ * Scan a local CUDA-compatible Vortex file as an Arrow C Device stream.
+ *
+ * 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.
+ *
+ * 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.
+ */
+int vx_cuda_scan_path_arrow_device_stream(const vx_session *session,
+                                          vx_view path,
+                                          struct ArrowDeviceArrayStream *out_stream,
+                                          vx_error **error_out);
+
 /**
  * Export a borrowed Vortex array for cuDF's Arrow Device import path.
  *

vortex-cuda/ffi/src/lib.rs

@@ -13,13 +13,19 @@ use std::ptr;
 use std::sync::Arc;
 
 use arrow_schema::ffi::FFI_ArrowSchema;
+use vortex::array::stream::ArrayStreamExt;
 use vortex::compressor::BtrBlocksCompressorBuilder;
 use vortex::error::VortexResult;
 use vortex::error::vortex_ensure;
+use vortex::file::OpenOptionsSessionExt;
 use vortex::file::WriteStrategyBuilder;
+use vortex::io::VortexReadAt;
+use vortex::io::runtime::BlockingRuntime;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::session::SessionExt;
 use vortex::session::VortexSession;
 use vortex_cuda::CudaSession;
+use vortex_cuda::PooledFileReadAt;
 use vortex_cuda::arrow::ArrowDeviceArray;
 use vortex_cuda::arrow::ArrowDeviceArrayStream;
 use vortex_cuda::arrow::DeviceArrayExt;
@@ -104,6 +110,58 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file(
     })
 }
 
+/// Scan a local Vortex file through pinned host buffers and export an Arrow C Device stream.
+///
+/// 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`).
+///
+/// # 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.
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream(
+    session: *const vx_session,
+    path: vx_view,
+    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 cuda_session = session.get::<CudaSession>();
+        let stream = cuda_session.stream()?;
+        let pool = Arc::clone(cuda_session.pinned_buffer_pool());
+        drop(cuda_session);
+
+        let reader: Arc<dyn VortexReadAt> = Arc::new(PooledFileReadAt::open(
+            path,
+            session.handle(),
+            pool,
+            stream,
+        )?);
+        let array_stream = ffi_runtime().block_on(async {
+            let file = session.open_options().open(reader).await?;
+            Ok::<_, vortex::error::VortexError>(file.scan()?.into_array_stream()?.boxed())
+        })?;
+        let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?;
+
+        unsafe { ptr::write(out_stream, device_stream) };
+        Ok(VX_CUDA_OK)
+    })
+}
+
 /// 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-cuda/gpu-scan-cli/src/main.rs

@@ -31,11 +31,9 @@ use vortex::file::WriteStrategyBuilder;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::session::VortexSession;
 use vortex_cuda::CudaSession;
-use vortex_cuda::PinnedByteBufferPool;
 use vortex_cuda::PooledByteBufferReadAt;
 use vortex_cuda::PooledFileReadAt;
 use vortex_cuda::TracingLaunchStrategy;
-use vortex_cuda::VortexCudaStreamPool;
 use vortex_cuda::executor::CudaArrayExt;
 use vortex_cuda::layout::CudaFlatLayoutStrategy;
 use vortex_cuda::layout::register_cuda_layout;
@@ -155,11 +153,10 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes
     let mut cuda_ctx = CudaSession::create_execution_ctx(&session)?
         .with_launch_strategy(Arc::new(TracingLaunchStrategy));
 
-    let pool = Arc::new(PinnedByteBufferPool::new(Arc::clone(
-        cuda_ctx.stream().context(),
-    )));
-    let cuda_stream =
-        VortexCudaStreamPool::new(Arc::clone(cuda_ctx.stream().context()), 1).stream()?;
+    let cuda_session = session.get::<CudaSession>();
+    let pool = Arc::clone(cuda_session.pinned_buffer_pool());
+    let cuda_stream = cuda_session.stream()?;
+    drop(cuda_session);
     let handle = session.handle();
 
     let gpu_file_handle = if gpu_file {

vortex-cuda/src/pinned.rs

@@ -113,6 +113,15 @@ pub struct PinnedByteBufferPool {
     puts: AtomicU64,
 }
 
+impl std::fmt::Debug for PinnedByteBufferPool {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("PinnedByteBufferPool")
+            .field("max_keep_per_size", &self.max_keep_per_size)
+            .field("stats", &self.stats())
+            .finish_non_exhaustive()
+    }
+}
+
 struct InflightPinnedBuffer {
     event: Arc<CudaEvent>,
     buffer: PinnedByteBuffer,

vortex-cuda/src/session.rs

@@ -23,6 +23,7 @@ use crate::executor::CudaExecute;
 pub use crate::executor::CudaExecutionCtx;
 use crate::initialize_cuda;
 use crate::kernel::KernelLoader;
+use crate::pinned::PinnedByteBufferPool;
 use crate::stream::VortexCudaStream;
 use crate::stream_pool::VortexCudaStreamPool;
 
@@ -40,6 +41,7 @@ pub struct CudaSession {
     export_device_array: Arc<dyn ExportDeviceArray>,
     kernel_loader: Arc<KernelLoader>,
     stream_pool: Arc<VortexCudaStreamPool>,
+    pinned_buffer_pool: Arc<PinnedByteBufferPool>,
 }
 
 impl CudaSession {
@@ -57,12 +59,14 @@ impl CudaSession {
             Arc::clone(&context),
             stream_pool_capacity,
         ));
+        let pinned_buffer_pool = Arc::new(PinnedByteBufferPool::new(Arc::clone(&context)));
         Self {
             context,
             kernels: Arc::new(DashMap::default()),
             kernel_loader: Arc::new(KernelLoader::new()),
             export_device_array: Arc::new(CanonicalDeviceArrayExport),
             stream_pool,
+            pinned_buffer_pool,
         }
     }
 
@@ -105,6 +109,11 @@ impl CudaSession {
         self.stream_pool.stream()
     }
 
+    /// Returns the session-scoped pool used for staging file reads in pinned host memory.
+    pub fn pinned_buffer_pool(&self) -> &Arc<PinnedByteBufferPool> {
+        &self.pinned_buffer_pool
+    }
+
     /// Registers CUDA support for an array encoding.
     ///
     /// # Arguments

vortex-ffi/src/string.rs

@@ -61,8 +61,8 @@ impl vx_view {
     ///
     /// # Safety
     ///
-    /// Same requirements as in as_bytes
-    pub(crate) unsafe fn as_str<'a>(&self) -> VortexResult<&'a str> {
+    /// `self.ptr` must be valid for `self.len` reads, or null when `self.len` is zero.
+    pub unsafe fn as_str<'a>(&self) -> VortexResult<&'a str> {
         str::from_utf8(unsafe { self.as_bytes() }?).map_err(|e| vortex_err!("invalid utf-8: {e}"))
     }
 }