chore[visitor]: add a visit buffer handle with name (#6191)

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

encodings/bytebool/src/array.rs

@@ -33,7 +33,7 @@ use vortex_array::vtable::VisitorVTable;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
-use vortex_error::VortexExpect;
+use vortex_error::VortexExpect as _;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -226,7 +226,7 @@ impl OperationsVTable<ByteBoolVTable> for ByteBoolVTable {
 
 impl VisitorVTable<ByteBoolVTable> for ByteBoolVTable {
     fn visit_buffers(array: &ByteBoolArray, visitor: &mut dyn ArrayBufferVisitor) {
-        visitor.visit_buffer(array.buffer().as_host());
+        visitor.visit_buffer_handle("values", array.buffer());
     }
 
     fn visit_children(array: &ByteBoolArray, visitor: &mut dyn ArrayChildVisitor) {

encodings/fastlanes/src/bitpacking/vtable/visitor.rs

@@ -11,7 +11,7 @@ use crate::BitPackedVTable;
 
 impl VisitorVTable<BitPackedVTable> for BitPackedVTable {
     fn visit_buffers(array: &BitPackedArray, visitor: &mut dyn ArrayBufferVisitor) {
-        visitor.visit_buffer(array.packed().as_host());
+        visitor.visit_buffer_handle("packed", array.packed());
     }
 
     fn visit_children(array: &BitPackedArray, visitor: &mut dyn ArrayChildVisitor) {

encodings/fsst/src/array.rs

@@ -427,8 +427,14 @@ impl ValidityChild<FSSTVTable> for FSSTVTable {
 
 impl VisitorVTable<FSSTVTable> for FSSTVTable {
     fn visit_buffers(array: &FSSTArray, visitor: &mut dyn ArrayBufferVisitor) {
-        visitor.visit_buffer(&array.symbols().clone().into_byte_buffer());
-        visitor.visit_buffer(&array.symbol_lengths().clone().into_byte_buffer());
+        visitor.visit_buffer_handle(
+            "symbols",
+            &BufferHandle::new_host(array.symbols().clone().into_byte_buffer()),
+        );
+        visitor.visit_buffer_handle(
+            "symbol_lengths",
+            &BufferHandle::new_host(array.symbol_lengths().clone().into_byte_buffer()),
+        );
     }
 
     fn visit_children(array: &FSSTArray, visitor: &mut dyn ArrayChildVisitor) {

encodings/pco/src/array.rs

@@ -52,7 +52,7 @@ use vortex_dtype::DType;
 use vortex_dtype::PType;
 use vortex_dtype::half;
 use vortex_error::VortexError;
-use vortex_error::VortexExpect;
+use vortex_error::VortexExpect as _;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -533,11 +533,17 @@ impl OperationsVTable<PcoVTable> for PcoVTable {
 
 impl VisitorVTable<PcoVTable> for PcoVTable {
     fn visit_buffers(array: &PcoArray, visitor: &mut dyn ArrayBufferVisitor) {
-        for buffer in &array.chunk_metas {
-            visitor.visit_buffer(buffer);
+        for (i, buffer) in array.chunk_metas.iter().enumerate() {
+            visitor.visit_buffer_handle(
+                &format!("chunk_meta_{i}"),
+                &BufferHandle::new_host(buffer.clone()),
+            );
         }
-        for buffer in &array.pages {
-            visitor.visit_buffer(buffer);
+        for (i, buffer) in array.pages.iter().enumerate() {
+            visitor.visit_buffer_handle(
+                &format!("page_{i}"),
+                &BufferHandle::new_host(buffer.clone()),
+            );
         }
     }
 

encodings/sparse/src/lib.rs

@@ -489,7 +489,7 @@ impl VisitorVTable<SparseVTable> for SparseVTable {
             .value()
             .to_protobytes::<ByteBufferMut>()
             .freeze();
-        visitor.visit_buffer(&fill_value_buffer);
+        visitor.visit_buffer_handle("fill_value", &BufferHandle::new_host(fill_value_buffer));
     }
 
     fn visit_children(array: &SparseArray, visitor: &mut dyn ArrayChildVisitor) {

encodings/zstd/src/array.rs

@@ -897,10 +897,13 @@ impl OperationsVTable<ZstdVTable> for ZstdVTable {
 impl VisitorVTable<ZstdVTable> for ZstdVTable {
     fn visit_buffers(array: &ZstdArray, visitor: &mut dyn ArrayBufferVisitor) {
         if let Some(buffer) = &array.dictionary {
-            visitor.visit_buffer(buffer);
+            visitor.visit_buffer_handle("dictionary", &BufferHandle::new_host(buffer.clone()));
         }
-        for buffer in &array.frames {
-            visitor.visit_buffer(buffer);
+        for (i, buffer) in array.frames.iter().enumerate() {
+            visitor.visit_buffer_handle(
+                &format!("frame_{i}"),
+                &BufferHandle::new_host(buffer.clone()),
+            );
         }
     }
 

vortex-array/src/array/mod.rs

@@ -46,6 +46,7 @@ use crate::arrays::SliceArray;
 use crate::arrays::StructVTable;
 use crate::arrays::VarBinVTable;
 use crate::arrays::VarBinViewVTable;
+use crate::buffer::BufferHandle;
 use crate::builders::ArrayBuilder;
 use crate::compute;
 use crate::compute::ComputeFn;
@@ -778,8 +779,8 @@ impl<V: VTable> ArrayVisitor for ArrayAdapter<V> {
         }
 
         impl ArrayBufferVisitor for BufferCollector {
-            fn visit_buffer(&mut self, buffer: &ByteBuffer) {
-                self.buffers.push(buffer.clone());
+            fn visit_buffer_handle(&mut self, _name: &str, handle: &BufferHandle) {
+                self.buffers.push(handle.to_host_sync());
             }
         }
 
@@ -790,6 +791,46 @@ impl<V: VTable> ArrayVisitor for ArrayAdapter<V> {
         collector.buffers
     }
 
+    fn buffer_handles(&self) -> Vec<BufferHandle> {
+        struct BufferHandleCollector {
+            handles: Vec<BufferHandle>,
+        }
+
+        impl ArrayBufferVisitor for BufferHandleCollector {
+            fn visit_buffer_handle(&mut self, _name: &str, handle: &BufferHandle) {
+                self.handles.push(handle.clone());
+            }
+        }
+
+        let mut collector = BufferHandleCollector {
+            handles: Vec::new(),
+        };
+        <V::VisitorVTable as VisitorVTable<V>>::visit_buffers(&self.0, &mut collector);
+        collector.handles
+    }
+
+    fn buffer_names(&self) -> Vec<String> {
+        <V::VisitorVTable as VisitorVTable<V>>::buffer_names(&self.0)
+    }
+
+    fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
+        struct NamedBufferCollector {
+            buffers: Vec<(String, BufferHandle)>,
+        }
+
+        impl ArrayBufferVisitor for NamedBufferCollector {
+            fn visit_buffer_handle(&mut self, name: &str, handle: &BufferHandle) {
+                self.buffers.push((name.to_string(), handle.clone()));
+            }
+        }
+
+        let mut collector = NamedBufferCollector {
+            buffers: Vec::new(),
+        };
+        <V::VisitorVTable as VisitorVTable<V>>::visit_buffers(&self.0, &mut collector);
+        collector.buffers
+    }
+
     fn nbuffers(&self) -> usize {
         <V::VisitorVTable as VisitorVTable<V>>::nbuffers(&self.0)
     }

vortex-array/src/array/visitor.rs

@@ -31,6 +31,15 @@ pub trait ArrayVisitor {
     /// Returns the buffers of the array.
     fn buffers(&self) -> Vec<ByteBuffer>;
 
+    /// Returns the buffer handles of the array.
+    fn buffer_handles(&self) -> Vec<BufferHandle>;
+
+    /// Returns the names of the buffers of the array.
+    fn buffer_names(&self) -> Vec<String>;
+
+    /// Returns the array's buffers with their names.
+    fn named_buffers(&self) -> Vec<(String, BufferHandle)>;
+
     /// Returns the number of buffers of the array.
     fn nbuffers(&self) -> usize;
 
@@ -63,6 +72,18 @@ impl ArrayVisitor for Arc<dyn Array> {
         self.as_ref().buffers()
     }
 
+    fn buffer_handles(&self) -> Vec<BufferHandle> {
+        self.as_ref().buffer_handles()
+    }
+
+    fn buffer_names(&self) -> Vec<String> {
+        self.as_ref().buffer_names()
+    }
+
+    fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
+        self.as_ref().named_buffers()
+    }
+
     fn nbuffers(&self) -> usize {
         self.as_ref().nbuffers()
     }
@@ -114,11 +135,7 @@ pub trait ArrayVisitorExt: Array {
 impl<A: Array + ?Sized> ArrayVisitorExt for A {}
 
 pub trait ArrayBufferVisitor {
-    fn visit_buffer_handle(&mut self, handle: &BufferHandle) -> VortexResult<()> {
-        self.visit_buffer(&handle.clone().try_to_host_sync()?);
-        Ok(())
-    }
-    fn visit_buffer(&mut self, buffer: &ByteBuffer);
+    fn visit_buffer_handle(&mut self, _name: &str, handle: &BufferHandle);
 }
 
 pub trait ArrayChildVisitor {

vortex-array/src/arrays/bool/vtable/visitor.rs

@@ -9,7 +9,7 @@ use crate::vtable::VisitorVTable;
 
 impl VisitorVTable<BoolVTable> for BoolVTable {
     fn visit_buffers(array: &BoolArray, visitor: &mut dyn ArrayBufferVisitor) {
-        visitor.visit_buffer(array.bits.as_host())
+        visitor.visit_buffer_handle("bits", &array.bits);
     }
 
     fn visit_children(array: &BoolArray, visitor: &mut dyn ArrayChildVisitor) {

vortex-array/src/arrays/constant/vtable/visitor.rs

@@ -7,6 +7,7 @@ use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
 use crate::arrays::ConstantArray;
 use crate::arrays::ConstantVTable;
+use crate::buffer::BufferHandle;
 use crate::vtable::VisitorVTable;
 
 impl VisitorVTable<ConstantVTable> for ConstantVTable {
@@ -16,7 +17,7 @@ impl VisitorVTable<ConstantVTable> for ConstantVTable {
             .value()
             .to_protobytes::<ByteBufferMut>()
             .freeze();
-        visitor.visit_buffer(&buffer);
+        visitor.visit_buffer_handle("scalar", &BufferHandle::new_host(buffer));
     }
 
     fn visit_children(_array: &ConstantArray, _visitor: &mut dyn ArrayChildVisitor) {}