LazyBuffers: use LazyBufferHandle

vortex-array/src/array/mod.rs

@@ -165,6 +165,9 @@ pub trait DynArray:
 
     /// Replaces the children of the array with the given array references.
     fn with_children(&self, children: Vec<ArrayRef>) -> VortexResult<ArrayRef>;
+
+    /// Replaces the buffers of the array with the given buffer handles.
+    fn with_buffers(&self, buffers: Vec<BufferHandle>) -> VortexResult<ArrayRef>;
 }
 
 impl DynArray for Arc<dyn DynArray> {
@@ -278,6 +281,10 @@ impl DynArray for Arc<dyn DynArray> {
     fn with_children(&self, children: Vec<ArrayRef>) -> VortexResult<ArrayRef> {
         self.as_ref().with_children(children)
     }
+
+    fn with_buffers(&self, buffers: Vec<BufferHandle>) -> VortexResult<ArrayRef> {
+        self.as_ref().with_buffers(buffers)
+    }
 }
 
 /// A reference counted pointer to a dynamic [`DynArray`] trait object.
@@ -663,6 +670,12 @@ impl<V: VTable> DynArray for ArrayAdapter<V> {
         V::with_children(&mut this, children)?;
         Ok(this.into_array())
     }
+
+    fn with_buffers(&self, buffers: Vec<BufferHandle>) -> VortexResult<ArrayRef> {
+        let mut this = self.0.clone();
+        V::with_buffers(&mut this, buffers)?;
+        Ok(this.into_array())
+    }
 }
 
 impl<V: VTable> ArrayHash for ArrayAdapter<V> {

vortex-array/src/vtable/dyn_.rs

@@ -48,6 +48,7 @@ pub trait DynVTable: 'static + private::Sealed + Send + Sync + Debug {
         session: &VortexSession,
     ) -> VortexResult<ArrayRef>;
     fn with_children(&self, array: &ArrayRef, children: Vec<ArrayRef>) -> VortexResult<ArrayRef>;
+    fn with_buffers(&self, array: &ArrayRef, buffers: Vec<BufferHandle>) -> VortexResult<ArrayRef>;
 
     /// See [`VTable::reduce`]
     fn reduce(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
@@ -101,6 +102,12 @@ impl<V: VTable> DynVTable for ArrayVTableAdapter<V> {
         Ok(array.into_array())
     }
 
+    fn with_buffers(&self, array: &ArrayRef, buffers: Vec<BufferHandle>) -> VortexResult<ArrayRef> {
+        let mut array = array.as_::<V>().clone();
+        V::with_buffers(&mut array, buffers)?;
+        Ok(array.into_array())
+    }
+
     fn reduce(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         let Some(reduced) = V::reduce(downcast::<V>(array))? else {
             return Ok(None);

vortex-array/src/vtable/mod.rs

@@ -181,6 +181,23 @@ pub trait VTable: 'static + Sized + Send + Sync + Debug {
     /// of children must be expected.
     fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()>;
 
+    /// Replaces the buffers in `array` with `buffers`.
+    ///
+    /// The default implementation rebuilds the array via [`build()`](VTable::build), which
+    /// re-runs all validation. This is correct for replacing lazy device buffers with
+    /// materialized host buffers.
+    fn with_buffers(array: &mut Self::Array, buffers: Vec<BufferHandle>) -> VortexResult<()> {
+        let metadata = Self::metadata(array)?;
+        let dtype = Self::dtype(array).clone();
+        let len = Self::len(array);
+        let children: Vec<ArrayRef> = (0..Self::nchildren(array))
+            .map(|i| Self::child(array, i))
+            .collect();
+        let children_slice: &[ArrayRef] = &children;
+        *array = Self::build(&dtype, len, &metadata, &buffers, &children_slice)?;
+        Ok(())
+    }
+
     /// Execute this array by returning an [`ExecutionStep`] that tells the scheduler what to
     /// do next.
     ///

vortex-layout/public-api.lock

@@ -12,13 +12,19 @@ pub struct vortex_layout::buffer::LazyBufferHandle
 
 impl vortex_layout::buffer::LazyBufferHandle
 
+pub fn vortex_layout::buffer::LazyBufferHandle::alignment(&self) -> vortex_buffer::alignment::Alignment
+
 pub fn vortex_layout::buffer::LazyBufferHandle::byte_ranges(&self) -> core::option::Option<&[core::ops::range::Range<usize>]>
 
 pub fn vortex_layout::buffer::LazyBufferHandle::filter(&self, ranges: &[core::ops::range::Range<usize>]) -> Self
 
+pub fn vortex_layout::buffer::LazyBufferHandle::is_empty(&self) -> bool
+
+pub fn vortex_layout::buffer::LazyBufferHandle::len(&self) -> usize
+
 pub async fn vortex_layout::buffer::LazyBufferHandle::materialize(&self) -> vortex_error::VortexResult<vortex_array::buffer::BufferHandle>
 
-pub fn vortex_layout::buffer::LazyBufferHandle::new(source: alloc::sync::Arc<dyn vortex_layout::segments::SegmentSource>, segment_id: vortex_layout::segments::SegmentId) -> Self
+pub fn vortex_layout::buffer::LazyBufferHandle::new(source: alloc::sync::Arc<dyn vortex_layout::segments::SegmentSource>, segment_id: vortex_layout::segments::SegmentId, alignment: vortex_buffer::alignment::Alignment) -> Self
 
 pub fn vortex_layout::buffer::LazyBufferHandle::segment_id(&self) -> vortex_layout::segments::SegmentId
 
@@ -28,10 +34,42 @@ impl core::clone::Clone for vortex_layout::buffer::LazyBufferHandle
 
 pub fn vortex_layout::buffer::LazyBufferHandle::clone(&self) -> vortex_layout::buffer::LazyBufferHandle
 
+impl core::cmp::Eq for vortex_layout::buffer::LazyBufferHandle
+
+impl core::cmp::PartialEq for vortex_layout::buffer::LazyBufferHandle
+
+pub fn vortex_layout::buffer::LazyBufferHandle::eq(&self, other: &Self) -> bool
+
 impl core::fmt::Debug for vortex_layout::buffer::LazyBufferHandle
 
 pub fn vortex_layout::buffer::LazyBufferHandle::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
+impl core::hash::Hash for vortex_layout::buffer::LazyBufferHandle
+
+pub fn vortex_layout::buffer::LazyBufferHandle::hash<H: core::hash::Hasher>(&self, state: &mut H)
+
+impl vortex_array::buffer::DeviceBuffer for vortex_layout::buffer::LazyBufferHandle
+
+pub fn vortex_layout::buffer::LazyBufferHandle::aligned(self: alloc::sync::Arc<Self>, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<alloc::sync::Arc<dyn vortex_array::buffer::DeviceBuffer>>
+
+pub fn vortex_layout::buffer::LazyBufferHandle::alignment(&self) -> vortex_buffer::alignment::Alignment
+
+pub fn vortex_layout::buffer::LazyBufferHandle::as_any(&self) -> &dyn core::any::Any
+
+pub fn vortex_layout::buffer::LazyBufferHandle::copy_to_host(&self, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<futures_core::future::BoxFuture<'static, vortex_error::VortexResult<vortex_buffer::ByteBuffer>>>
+
+pub fn vortex_layout::buffer::LazyBufferHandle::copy_to_host_sync(&self, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<vortex_buffer::ByteBuffer>
+
+pub fn vortex_layout::buffer::LazyBufferHandle::filter(&self, ranges: &[core::ops::range::Range<usize>]) -> vortex_error::VortexResult<alloc::sync::Arc<dyn vortex_array::buffer::DeviceBuffer>>
+
+pub fn vortex_layout::buffer::LazyBufferHandle::len(&self) -> usize
+
+pub fn vortex_layout::buffer::LazyBufferHandle::slice(&self, range: core::ops::range::Range<usize>) -> alloc::sync::Arc<dyn vortex_array::buffer::DeviceBuffer>
+
+pub fn vortex_layout::buffer::create_lazy_array_parts(array_tree: vortex_buffer::ByteBuffer, source: alloc::sync::Arc<dyn vortex_layout::segments::SegmentSource>, segment_id: vortex_layout::segments::SegmentId) -> vortex_error::VortexResult<vortex_array::serde::ArrayParts>
+
+pub async fn vortex_layout::buffer::materialize_recursive(array: &vortex_array::array::ArrayRef) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
+
 pub mod vortex_layout::display
 
 pub struct vortex_layout::display::DisplayLayoutTree