Route CUDA pruning io through host memory

vortex-cuda/ffi/src/lib.rs

@@ -19,13 +19,11 @@ 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::CudaOpenOptionsExt;
 use vortex_cuda::CudaSession;
-use vortex_cuda::PooledFileReadAt;
 use vortex_cuda::arrow::ArrowDeviceArray;
 use vortex_cuda::arrow::ArrowDeviceArrayStream;
 use vortex_cuda::arrow::DeviceArrayExt;
@@ -110,7 +108,10 @@ 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.
+/// Scan a local Vortex file 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
@@ -140,19 +141,8 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream(
 
         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?;
+            let file = session.open_options().with_cuda().open_path(path).await?;
             Ok::<_, vortex::error::VortexError>(file.scan()?.into_array_stream()?.boxed())
         })?;
         let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?;

vortex-cuda/gpu-scan-cli/src/main.rs

@@ -31,9 +31,9 @@ use vortex::file::WriteStrategyBuilder;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::session::SessionExt;
 use vortex::session::VortexSession;
+use vortex_cuda::CudaOpenOptionsExt;
 use vortex_cuda::CudaSession;
 use vortex_cuda::PooledByteBufferReadAt;
-use vortex_cuda::PooledFileReadAt;
 use vortex_cuda::TracingLaunchStrategy;
 use vortex_cuda::executor::CudaArrayExt;
 use vortex_cuda::layout::CudaFlatLayoutStrategy;
@@ -158,11 +158,9 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes
     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 {
-        let reader = PooledFileReadAt::open(&path, handle, Arc::clone(&pool), cuda_stream)?;
-        session.open_options().open(Arc::new(reader)).await?
+        session.open_options().with_cuda().open_path(&path).await?
     } else {
         let (recompressed, footer) = recompress_for_gpu(&path, &session).await?;
         let reader = PooledByteBufferReadAt::new(recompressed, Arc::clone(&pool), cuda_stream);

vortex-cuda/src/file.rs

@@ -0,0 +1,235 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! CUDA-aware Vortex file opening.
+
+use std::path::Path;
+use std::sync::Arc;
+
+use vortex::error::VortexResult;
+use vortex::error::vortex_bail;
+use vortex::file::VortexFile;
+use vortex::file::VortexOpenOptions;
+use vortex::io::session::RuntimeSessionExt;
+use vortex::layout::Layout;
+use vortex::layout::layouts::zoned::LegacyStats;
+use vortex::layout::layouts::zoned::Zoned;
+use vortex::layout::segments::SegmentFuture;
+use vortex::layout::segments::SegmentId;
+use vortex::layout::segments::SegmentSource;
+
+use crate::CudaSessionExt;
+use crate::PooledFileReadAt;
+use crate::layout::register_cuda_layout;
+
+/// Extension trait for opening CUDA-readable files from [`VortexOpenOptions`].
+pub trait CudaOpenOptionsExt {
+    /// Configure the opener to keep control-plane segments on the host and read data-plane
+    /// segments into CUDA device memory.
+    fn with_cuda(self) -> CudaOpenOptions;
+}
+
+impl CudaOpenOptionsExt for VortexOpenOptions {
+    fn with_cuda(self) -> CudaOpenOptions {
+        CudaOpenOptions { inner: self }
+    }
+}
+
+/// CUDA-aware Vortex file open options.
+///
+/// Construct this adapter with [`CudaOpenOptionsExt::with_cuda`]. Standard file options should be
+/// configured before calling `with_cuda`.
+pub struct CudaOpenOptions {
+    inner: VortexOpenOptions,
+}
+
+impl CudaOpenOptions {
+    /// Open a local Vortex file for CUDA execution.
+    ///
+    /// The footer and zone-map segments are read through the ordinary host path. All other file
+    /// segments are staged through the session-scoped pinned buffer pool and transferred to the
+    /// device.
+    pub async fn open_path(self, path: impl AsRef<Path>) -> VortexResult<VortexFile> {
+        let path = path.as_ref().to_owned();
+        register_cuda_layout(self.inner.session());
+
+        // A host cache cannot front the data source: cached host buffers would violate the CUDA
+        // source's device-placement contract. It remains enabled on the control-plane source.
+        let data_options = self.inner.clone().without_segment_cache();
+        let control_file = self.inner.open_path(&path).await?;
+        let footer = control_file.footer().clone();
+
+        let session = control_file.session().clone();
+        let cuda_session = session.cuda_session();
+        let stream = cuda_session.stream()?;
+        let pool = Arc::clone(cuda_session.pinned_buffer_pool());
+        drop(cuda_session);
+
+        let reader = Arc::new(PooledFileReadAt::open(
+            &path,
+            session.handle(),
+            pool,
+            stream,
+        )?);
+        let data_file = data_options
+            .with_footer(footer.clone())
+            .open(reader)
+            .await?;
+
+        let control_segments =
+            control_plane_segments(footer.layout().as_ref(), footer.segment_map().len())?;
+        let routed = Arc::new(RoutingSegmentSource {
+            control_segments: control_segments.into(),
+            control: control_file.segment_source(),
+            data: data_file.segment_source(),
+        });
+
+        Ok(data_file.with_segment_source(routed))
+    }
+}
+
+struct RoutingSegmentSource {
+    control_segments: Arc<[bool]>,
+    control: Arc<dyn SegmentSource>,
+    data: Arc<dyn SegmentSource>,
+}
+
+impl SegmentSource for RoutingSegmentSource {
+    fn request(&self, id: SegmentId) -> SegmentFuture {
+        let Ok(idx) = usize::try_from(*id) else {
+            return self.data.request(id);
+        };
+        if self.control_segments.get(idx).copied().unwrap_or(false) {
+            self.control.request(id)
+        } else {
+            self.data.request(id)
+        }
+    }
+}
+
+fn control_plane_segments(layout: &dyn Layout, segment_count: usize) -> VortexResult<Vec<bool>> {
+    let mut control_segments = vec![false; segment_count];
+    mark_zone_map_segments(layout, &mut control_segments)?;
+    Ok(control_segments)
+}
+
+fn mark_zone_map_segments(layout: &dyn Layout, control_segments: &mut [bool]) -> VortexResult<()> {
+    if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
+        mark_layout_segments(layout.child(1)?.as_ref(), control_segments)?;
+        mark_zone_map_segments(layout.child(0)?.as_ref(), control_segments)?;
+        return Ok(());
+    }
+
+    for child in layout.children()? {
+        mark_zone_map_segments(child.as_ref(), control_segments)?;
+    }
+    Ok(())
+}
+
+fn mark_layout_segments(layout: &dyn Layout, control_segments: &mut [bool]) -> VortexResult<()> {
+    for id in layout.segment_ids() {
+        let idx = usize::try_from(*id)?;
+        let Some(is_control) = control_segments.get_mut(idx) else {
+            vortex_bail!("layout references out-of-range segment {id}");
+        };
+        *is_control = true;
+    }
+
+    for child in layout.children()? {
+        mark_layout_segments(child.as_ref(), control_segments)?;
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::num::NonZeroUsize;
+
+    use futures::FutureExt;
+    use parking_lot::Mutex;
+    use vortex::array::buffer::BufferHandle;
+    use vortex::array::dtype::DType;
+    use vortex::array::dtype::Nullability::NonNullable;
+    use vortex::array::dtype::PType;
+    use vortex::array::dtype::StructFields;
+    use vortex::buffer::ByteBuffer;
+    use vortex::layout::IntoLayout;
+    use vortex::layout::layouts::flat::FlatLayout;
+    use vortex::layout::layouts::zoned::ZonedLayout;
+    use vortex::session::registry::ReadContext;
+
+    use super::*;
+
+    #[derive(Clone)]
+    struct RecordingSource {
+        requests: Arc<Mutex<Vec<SegmentId>>>,
+        value: u8,
+    }
+
+    impl RecordingSource {
+        fn new(value: u8) -> Self {
+            Self {
+                requests: Arc::new(Mutex::new(Vec::new())),
+                value,
+            }
+        }
+    }
+
+    impl SegmentSource for RecordingSource {
+        fn request(&self, id: SegmentId) -> SegmentFuture {
+            self.requests.lock().push(id);
+            let value = self.value;
+            async move { Ok(BufferHandle::new_host(ByteBuffer::from(vec![value]))) }.boxed()
+        }
+    }
+
+    #[tokio::test]
+    async fn routes_control_and_data_segments() {
+        let control = RecordingSource::new(1);
+        let data = RecordingSource::new(2);
+        let source = RoutingSegmentSource {
+            control_segments: Arc::from([false, true, false]),
+            control: Arc::new(control.clone()),
+            data: Arc::new(data.clone()),
+        };
+
+        let control_buffer = source
+            .request(SegmentId::from(1))
+            .await
+            .unwrap()
+            .unwrap_host();
+        let data_buffer = source
+            .request(SegmentId::from(2))
+            .await
+            .unwrap()
+            .unwrap_host();
+
+        assert_eq!(control_buffer.as_ref(), &[1]);
+        assert_eq!(data_buffer.as_ref(), &[2]);
+        assert_eq!(*control.requests.lock(), [SegmentId::from(1)]);
+        assert_eq!(*data.requests.lock(), [SegmentId::from(2)]);
+    }
+
+    #[test]
+    fn marks_only_zone_map_subtree_segments_as_control_plane() -> VortexResult<()> {
+        let read_ctx = ReadContext::new([]);
+        let data_dtype = DType::Primitive(PType::I32, NonNullable);
+        let data =
+            FlatLayout::new(16, data_dtype, SegmentId::from(0), read_ctx.clone()).into_layout();
+        let zones = FlatLayout::new(
+            1,
+            DType::Struct(StructFields::empty(), NonNullable),
+            SegmentId::from(1),
+            read_ctx,
+        )
+        .into_layout();
+        let layout =
+            ZonedLayout::try_new(data, zones, NonZeroUsize::new(16).unwrap(), Arc::new([]))?
+                .into_layout();
+
+        let control_segments = control_plane_segments(layout.as_ref(), 2)?;
+
+        assert_eq!(control_segments, [false, true]);
+        Ok(())
+    }
+}

vortex-cuda/src/lib.rs

@@ -14,6 +14,7 @@ mod device_buffer;
 mod device_read_at;
 pub mod dynamic_dispatch;
 pub mod executor;
+mod file;
 pub mod hybrid_dispatch;
 mod kernel;
 pub mod layout;
@@ -34,6 +35,8 @@ pub use device_read_at::CopyDeviceReadAt;
 pub use executor::CudaDispatchMode;
 pub use executor::CudaExecutionCtx;
 pub use executor::CudaKernelEvents;
+pub use file::CudaOpenOptions;
+pub use file::CudaOpenOptionsExt;
 use kernel::ALPExecutor;
 use kernel::BitPackedExecutor;
 use kernel::ConstantNumericExecutor;

vortex-cuda/src/pooled_read_at.rs

@@ -40,6 +40,9 @@ pub const DEFAULT_OBJECT_STORE_CONCURRENCY: usize = 192;
 ///
 /// Reads into a pooled pinned (page-locked) buffer, then submits a non-blocking
 /// H2D DMA transfer and returns a device `BufferHandle`.
+///
+/// This is a data-plane reader. To open a complete local Vortex file, prefer
+/// [`crate::CudaOpenOptionsExt::with_cuda`], which keeps the footer and zone maps on the host.
 #[derive(Clone)]
 pub struct PooledFileReadAt {
     uri: Arc<str>,

vortex-file/src/file.rs

@@ -125,6 +125,18 @@ impl VortexFile {
         Arc::clone(&self.segment_source)
     }
 
+    /// Replace the segment source used by this file.
+    ///
+    /// Any cached layout reader is cleared so that subsequent scans construct readers over the
+    /// replacement source.
+    pub fn with_segment_source(mut self, segment_source: Arc<dyn SegmentSource>) -> Self {
+        self.segment_source = segment_source;
+        if self.layout_reader_cache.is_some() {
+            self.layout_reader_cache = Some(OnceLock::new());
+        }
+        self
+    }
+
     /// Returns a reference to the Vortex session used to open this file.
     pub fn session(&self) -> &VortexSession {
         &self.session