Cleanup ArrayData::buffers

apache · Jul 29, 2023 · bd8b8e3 · bd8b8e3
1 parent 8c85d34
commit bd8b8e3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 102 deletions.
diff --git a/arrow-data/src/data/mod.rs → arrow-data/src/data.rs b/arrow-data/src/data/mod.rs → arrow-data/src/data.rs
@@ -29,8 +29,8 @@ use std::sync::Arc;
 
 use crate::equal;
 
-mod buffers;
-pub use buffers::*;
+/// A collection of [`Buffer`]
+pub type Buffers<'a> = &'a [Buffer];
 
 #[inline]
 pub(crate) fn contains_nulls(
@@ -109,6 +109,10 @@ pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuff
             buffer.push(0i64);
             [buffer, MutableBuffer::new(capacity * mem::size_of::<u8>())]
         }
+        DataType::Utf8View | DataType::BinaryView => [
+            MutableBuffer::new(capacity * mem::size_of::<u128>()),
+            empty_buffer,
+        ],
         DataType::List(_) | DataType::Map(_, _) => {
             // offset buffer always starts with a zero
             let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i32>());
@@ -345,10 +349,9 @@ impl ArrayData {
         &self.data_type
     }
 
-    /// Returns the [`Buffers`] storing data for this [`ArrayData`]
-    pub fn buffers(&self) -> Buffers<'_> {
-        // In future ArrayData won't store data contiguously as `Vec<Buffer>` (#1799)
-        Buffers::from_slice(&self.buffers)
+    /// Returns the [`Buffer`] storing data for this [`ArrayData`]
+    pub fn buffers(&self) -> &[Buffer] {
+        &self.buffers
     }
 
     /// Returns a slice of children [`ArrayData`]. This will be non
@@ -1486,6 +1489,9 @@ impl ArrayData {
 
 /// Return the expected [`DataTypeLayout`] Arrays of this data
 /// type are expected to have
+///
+/// For types with a variadic number of buffers, such as [`DataType::Utf8View`],
+/// this only returns the required buffers
 pub fn layout(data_type: &DataType) -> DataTypeLayout {
     // based on C/C++ implementation in
     // https://github.com/apache/arrow/blob/661c7d749150905a63dd3b52e0a04dac39030d95/cpp/src/arrow/type.h (and .cc)
@@ -1527,8 +1533,10 @@ pub fn layout(data_type: &DataType) -> DataTypeLayout {
             DataTypeLayout::new_fixed_width(bytes_per_value)
         }
         DataType::LargeBinary => DataTypeLayout::new_binary(size_of::<i64>()),
+        DataType::BinaryView => DataTypeLayout::new_fixed_width(size_of::<u128>()),
         DataType::Utf8 => DataTypeLayout::new_binary(size_of::<i32>()),
         DataType::LargeUtf8 => DataTypeLayout::new_binary(size_of::<i64>()),
+        DataType::Utf8View => DataTypeLayout::new_fixed_width(size_of::<u128>()),
         DataType::List(_) => DataTypeLayout::new_fixed_width(size_of::<i32>()),
         DataType::FixedSizeList(_, _) => DataTypeLayout::new_empty(), // all in child data
         DataType::LargeList(_) => DataTypeLayout::new_fixed_width(size_of::<i64>()),

diff --git a/arrow-data/src/data/buffers.rs b/arrow-data/src/data/buffers.rs