From ad262a36f130c5c617c6be82b170501aac89b5f7 Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Thu, 16 Jul 2026 11:23:11 +0100 Subject: [PATCH 1/2] Expose pooled CUDA file scans through FFI Signed-off-by: Onur Satici --- vortex-cuda/ffi/README.md | 4 ++ vortex-cuda/ffi/cinclude/vortex_cuda.h | 15 +++++++ vortex-cuda/ffi/src/lib.rs | 58 ++++++++++++++++++++++++++ vortex-cuda/gpu-scan-cli/src/main.rs | 11 ++--- vortex-cuda/src/pinned.rs | 9 ++++ vortex-cuda/src/session.rs | 9 ++++ vortex-ffi/src/string.rs | 4 +- 7 files changed, 101 insertions(+), 9 deletions(-) diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 9cc6b30dd94..1b07e12f0df 100644 --- a/vortex-cuda/ffi/README.md +++ b/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. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index d3ca8293f34..971ad9872f4 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/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. * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index cfc76ef354a..12bc584dad6 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/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::(); + let stream = cuda_session.stream()?; + let pool = Arc::clone(cuda_session.pinned_buffer_pool()); + drop(cuda_session); + + let reader: Arc = 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 diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 9078e30c691..29fc55c9155 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/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::(); + 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 { diff --git a/vortex-cuda/src/pinned.rs b/vortex-cuda/src/pinned.rs index cd97db63164..52b1aa46e2a 100644 --- a/vortex-cuda/src/pinned.rs +++ b/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, buffer: PinnedByteBuffer, diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index 5bc4e55de69..df3e3c13d55 100644 --- a/vortex-cuda/src/session.rs +++ b/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, kernel_loader: Arc, stream_pool: Arc, + pinned_buffer_pool: Arc, } 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 { + &self.pinned_buffer_pool + } + /// Registers CUDA support for an array encoding. /// /// # Arguments diff --git a/vortex-ffi/src/string.rs b/vortex-ffi/src/string.rs index dd77508f048..8e0dfd71900 100644 --- a/vortex-ffi/src/string.rs +++ b/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}")) } } From 61345d8f614c7bc96d899a49fd69770ac3437c05 Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Thu, 16 Jul 2026 14:22:19 +0100 Subject: [PATCH 2/2] import Signed-off-by: Onur Satici --- vortex-cuda/gpu-scan-cli/src/main.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 29fc55c9155..00f6cd1ac1c 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -29,6 +29,7 @@ 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; use vortex_cuda::CudaSession; use vortex_cuda::PooledByteBufferReadAt;