feat[buffer]: use BufferHandle in more places (#5567)

Use BufferHandle in segment source

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/serde.rs

@@ -415,6 +415,59 @@ impl ArrayParts {
         this.flatbuffer_loc = root._tab.loc();
         this
     }
+
+    /// Create an [`ArrayParts`] from a pre-existing flatbuffer (ArrayNode) and a segment containing
+    /// only the data buffers (without the flatbuffer suffix).
+    ///
+    /// This is used when the flatbuffer is stored separately in layout metadata (e.g., when
+    /// `FLAT_LAYOUT_INLINE_ARRAY_NODE` is enabled).
+    pub fn from_flatbuffer_and_segment(
+        array_tree: ByteBuffer,
+        segment: BufferHandle,
+    ) -> VortexResult<Self> {
+        // TODO: this can also work with device buffers.
+        let segment = segment.try_to_bytes()?;
+        // We align each buffer individually, so we remove alignment requirements on the buffer.
+        let segment = segment.aligned(Alignment::none());
+
+        let fb_buffer = FlatBuffer::align_from(array_tree);
+
+        // Parse the flatbuffer to extract buffer descriptors and root location.
+        let (flatbuffer_loc, buffers) = {
+            let fb_array = root::<fba::Array>(fb_buffer.as_ref())?;
+            let fb_root = fb_array.root().vortex_expect("Array must have a root node");
+            let flatbuffer_loc = fb_root._tab.loc();
+
+            let mut offset = 0;
+            let buffers: Arc<[_]> = fb_array
+                .buffers()
+                .unwrap_or_default()
+                .iter()
+                .map(|fb_buf| {
+                    // Skip padding
+                    offset += fb_buf.padding() as usize;
+
+                    let buffer_len = fb_buf.length() as usize;
+
+                    // Extract a buffer and ensure it's aligned, copying if necessary
+                    let buffer = segment
+                        .slice(offset..(offset + buffer_len))
+                        .aligned(Alignment::from_exponent(fb_buf.alignment_exponent()));
+
+                    offset += buffer_len;
+                    BufferHandle::Buffer(buffer)
+                })
+                .collect();
+
+            (flatbuffer_loc, buffers)
+        };
+
+        Ok(ArrayParts {
+            flatbuffer: fb_buffer,
+            flatbuffer_loc,
+            buffers,
+        })
+    }
 }
 
 struct ArrayPartsChildren<'a> {
@@ -484,3 +537,11 @@ impl TryFrom<ByteBuffer> for ArrayParts {
         })
     }
 }
+
+impl TryFrom<BufferHandle> for ArrayParts {
+    type Error = VortexError;
+
+    fn try_from(value: BufferHandle) -> Result<Self, Self::Error> {
+        Self::try_from(value.try_to_bytes()?)
+    }
+}

vortex-file/src/segments/source.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 
 use futures::FutureExt;
 use futures::TryFutureExt;
+use vortex_buffer::BufferHandle;
 use vortex_error::VortexError;
 use vortex_error::vortex_err;
 use vortex_io::VortexReadAt;
@@ -40,6 +41,7 @@ impl SegmentSource for FileSegmentSource {
             maybe_fut
                 .ok_or_else(|| vortex_err!("Missing segment: {}", id))?
                 .await
+                .map(BufferHandle::Buffer)
         }
         .boxed()
     }

vortex-layout/src/layouts/flat/mod.rs

@@ -179,6 +179,11 @@ impl FlatLayout {
     pub fn array_ctx(&self) -> &ArrayContext {
         &self.ctx
     }
+
+    #[inline]
+    pub fn array_tree(&self) -> Option<&ByteBuffer> {
+        self.array_tree.as_ref()
+    }
 }
 
 #[derive(prost::Message)]

vortex-layout/src/layouts/flat/reader.rs

@@ -64,11 +64,17 @@ impl FlatReader {
 
         let ctx = self.layout.array_ctx().clone();
         let dtype = self.layout.dtype().clone();
+        let array_tree = self.layout.array_tree().cloned();
         async move {
             let segment = segment_fut.await?;
-            ArrayParts::try_from(segment)?
-                .decode(&ctx, &dtype, row_count)
-                .map_err(Arc::new)
+            let parts = if let Some(array_tree) = array_tree {
+                // Use the pre-stored flatbuffer from layout metadata combined with segment buffers.
+                ArrayParts::from_flatbuffer_and_segment(array_tree, segment)?
+            } else {
+                // Parse the flatbuffer from the segment itself.
+                ArrayParts::try_from(segment)?
+            };
+            parts.decode(&ctx, &dtype, row_count).map_err(Arc::new)
         }
         .boxed()
         .shared()

vortex-layout/src/segments/cache.rs

@@ -9,6 +9,7 @@ use moka::future::Cache;
 use moka::future::CacheBuilder;
 use moka::policy::EvictionPolicy;
 use rustc_hash::FxBuildHasher;
+use vortex_buffer::BufferHandle;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -129,10 +130,13 @@ impl SegmentSource for SegmentCacheSourceAdapter {
         async move {
             if let Ok(Some(segment)) = cache.get(id).await {
                 log::debug!("Resolved segment {} from cache", id);
-                return Ok(segment);
+                return Ok(BufferHandle::Buffer(segment));
             }
             let result = delegate.await?;
-            if let Err(e) = cache.put(id, result.clone()).await {
+            // Cache only CPU buffers; device buffers are not cached.
+            if let BufferHandle::Buffer(ref buffer) = result
+                && let Err(e) = cache.put(id, buffer.clone()).await
+            {
                 log::warn!("Failed to store segment {} in cache: {}", id, e);
             }
             Ok(result)

vortex-layout/src/segments/shared.rs

@@ -7,7 +7,7 @@ use futures::FutureExt;
 use futures::TryFutureExt;
 use futures::future::BoxFuture;
 use futures::future::WeakShared;
-use vortex_buffer::ByteBuffer;
+use vortex_buffer::BufferHandle;
 use vortex_error::SharedVortexResult;
 use vortex_error::VortexError;
 use vortex_error::VortexExpect;
@@ -25,7 +25,7 @@ pub struct SharedSegmentSource<S> {
     in_flight: DashMap<SegmentId, WeakShared<SharedSegmentFuture>>,
 }
 
-type SharedSegmentFuture = BoxFuture<'static, SharedVortexResult<ByteBuffer>>;
+type SharedSegmentFuture = BoxFuture<'static, SharedVortexResult<BufferHandle>>;
 
 impl<S: SegmentSource> SharedSegmentSource<S> {
     /// Create a new `SharedSegmentSource` wrapping the provided inner source.
@@ -111,8 +111,8 @@ mod tests {
 
         // Both futures should resolve to the same data
         let (result1, result2) = futures::join!(future1, future2);
-        assert_eq!(result1.unwrap(), data);
-        assert_eq!(result2.unwrap(), data);
+        assert_eq!(*result1.unwrap().bytes(), data);
+        assert_eq!(*result2.unwrap().bytes(), data);
 
         // The inner source should have been called only once
         assert_eq!(source.request_count.load(Ordering::Relaxed), 1);
@@ -142,7 +142,7 @@ mod tests {
 
         // A new request should still work correctly
         let result = shared_source.request(id).await;
-        assert_eq!(result.unwrap(), data);
+        assert_eq!(*result.unwrap().bytes(), data);
 
         // Should have made 2 requests since the first was dropped before completion
         assert_eq!(source.request_count.load(Ordering::Relaxed), 2);

vortex-layout/src/segments/source.rs

@@ -2,12 +2,12 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use futures::future::BoxFuture;
-use vortex_buffer::ByteBuffer;
+use vortex_buffer::BufferHandle;
 use vortex_error::VortexResult;
 
 use crate::segments::SegmentId;
 /// Static future resolving to a segment byte buffer.
-pub type SegmentFuture = BoxFuture<'static, VortexResult<ByteBuffer>>;
+pub type SegmentFuture = BoxFuture<'static, VortexResult<BufferHandle>>;
 
 /// A trait for providing segment data to a [`crate::LayoutReader`].
 pub trait SegmentSource: 'static + Send + Sync {

vortex-layout/src/segments/test.rs

@@ -6,6 +6,7 @@ use std::sync::Arc;
 use async_trait::async_trait;
 use futures::FutureExt;
 use parking_lot::Mutex;
+use vortex_buffer::BufferHandle;
 use vortex_buffer::ByteBuffer;
 use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexExpect;
@@ -27,7 +28,12 @@ pub struct TestSegments {
 impl SegmentSource for TestSegments {
     fn request(&self, id: SegmentId) -> SegmentFuture {
         let buffer = self.segments.lock().get(*id as usize).cloned();
-        async move { buffer.ok_or_else(|| vortex_err!("Segment not found")) }.boxed()
+        async move {
+            buffer
+                .map(BufferHandle::Buffer)
+                .ok_or_else(|| vortex_err!("Segment not found"))
+        }
+        .boxed()
     }
 }