Return Buffers from ArrayData::buffers instead of slice (#1799)

arrow-array/src/array/primitive_array.rs

@@ -1236,7 +1236,7 @@ mod tests {
     fn test_primitive_array_from_vec() {
         let buf = Buffer::from_slice_ref([0, 1, 2, 3, 4]);
         let arr = Int32Array::from(vec![0, 1, 2, 3, 4]);
-        assert_eq!(buf, arr.data.buffers()[0]);
+        assert_eq!(buf, *arr.data.buffers()[0]);
         assert_eq!(5, arr.len());
         assert_eq!(0, arr.offset());
         assert_eq!(0, arr.null_count());
@@ -1742,7 +1742,7 @@ mod tests {
             .build()
             .unwrap();
         let arr = Int32Array::from(data);
-        assert_eq!(buf2, arr.data.buffers()[0]);
+        assert_eq!(buf2, *arr.data.buffers()[0]);
         assert_eq!(5, arr.len());
         assert_eq!(0, arr.null_count());
         for i in 0..3 {

arrow-array/src/array/union_array.rs

@@ -409,7 +409,7 @@ mod tests {
 
         // Check type ids
         assert_eq!(
-            union.data().buffers()[0],
+            *union.data().buffers()[0],
             Buffer::from_slice_ref(&expected_type_ids)
         );
         for (i, id) in expected_type_ids.iter().enumerate() {
@@ -418,7 +418,7 @@ mod tests {
 
         // Check offsets
         assert_eq!(
-            union.data().buffers()[1],
+            *union.data().buffers()[1],
             Buffer::from_slice_ref(&expected_value_offsets)
         );
         for (i, id) in expected_value_offsets.iter().enumerate() {
@@ -427,15 +427,15 @@ mod tests {
 
         // Check data
         assert_eq!(
-            union.data().child_data()[0].buffers()[0],
+            *union.data().child_data()[0].buffers()[0],
             Buffer::from_slice_ref([1_i32, 4, 6])
         );
         assert_eq!(
-            union.data().child_data()[1].buffers()[0],
+            *union.data().child_data()[1].buffers()[0],
             Buffer::from_slice_ref([2_i32, 7])
         );
         assert_eq!(
-            union.data().child_data()[2].buffers()[0],
+            *union.data().child_data()[2].buffers()[0],
             Buffer::from_slice_ref([3_i32, 5]),
         );
 
@@ -467,7 +467,7 @@ mod tests {
 
         // Check type ids
         assert_eq!(
-            union.data().buffers()[0],
+            *union.data().buffers()[0],
             Buffer::from_slice_ref(&expected_type_ids)
         );
         for (i, id) in expected_type_ids.iter().enumerate() {
@@ -476,7 +476,7 @@ mod tests {
 
         // Check offsets
         assert_eq!(
-            union.data().buffers()[1],
+            *union.data().buffers()[1],
             Buffer::from_slice_ref(&expected_value_offsets)
         );
         for (i, id) in expected_value_offsets.iter().enumerate() {
@@ -660,15 +660,15 @@ mod tests {
         .unwrap();
 
         // Check type ids
-        assert_eq!(Buffer::from_slice_ref(type_ids), array.data().buffers()[0]);
+        assert_eq!(Buffer::from_slice_ref(type_ids), *array.data().buffers()[0]);
         for (i, id) in type_ids.iter().enumerate() {
             assert_eq!(id, &array.type_id(i));
         }
 
         // Check offsets
         assert_eq!(
             Buffer::from_slice_ref(value_offsets),
-            array.data().buffers()[1]
+            *array.data().buffers()[1]
         );
         for (i, id) in value_offsets.iter().enumerate() {
             assert_eq!(id, &array.value_offset(i));
@@ -736,7 +736,7 @@ mod tests {
         // Check type ids
         assert_eq!(
             Buffer::from_slice_ref(&expected_type_ids),
-            union.data().buffers()[0]
+            *union.data().buffers()[0]
         );
         for (i, id) in expected_type_ids.iter().enumerate() {
             assert_eq!(id, &union.type_id(i));
@@ -747,16 +747,16 @@ mod tests {
 
         // Check data
         assert_eq!(
-            union.data().child_data()[0].buffers()[0],
+            *union.data().child_data()[0].buffers()[0],
             Buffer::from_slice_ref([1_i32, 0, 0, 4, 0, 6, 0]),
         );
         assert_eq!(
             Buffer::from_slice_ref([0_i32, 2_i32, 0, 0, 0, 0, 7]),
-            union.data().child_data()[1].buffers()[0]
+            *union.data().child_data()[1].buffers()[0]
         );
         assert_eq!(
             Buffer::from_slice_ref([0_i32, 0, 3_i32, 0, 5, 0, 0]),
-            union.data().child_data()[2].buffers()[0]
+            *union.data().child_data()[2].buffers()[0]
         );
 
         assert_eq!(expected_array_values.len(), union.len());
@@ -785,7 +785,7 @@ mod tests {
         // Check type ids
         assert_eq!(
             Buffer::from_slice_ref(&expected_type_ids),
-            union.data().buffers()[0]
+            *union.data().buffers()[0]
         );
         for (i, id) in expected_type_ids.iter().enumerate() {
             assert_eq!(id, &union.type_id(i));
@@ -847,7 +847,7 @@ mod tests {
         // Check type ids
         assert_eq!(
             Buffer::from_slice_ref(&expected_type_ids),
-            union.data().buffers()[0]
+            *union.data().buffers()[0]
         );
         for (i, id) in expected_type_ids.iter().enumerate() {
             assert_eq!(id, &union.type_id(i));

arrow-data/src/data/buffers.rs

@@ -0,0 +1,97 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow_buffer::Buffer;
+use std::iter::Chain;
+use std::ops::Index;
+
+/// A collection of [`Buffer`]
+#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
+pub struct Buffers<'a>([Option<&'a Buffer>; 2]);
+
+impl<'a> Buffers<'a> {
+    /// Temporary will be removed once ArrayData does not store `Vec<Buffer>` directly (#3769)
+    #[inline]
+    pub(crate) fn from_slice(a: &'a [Buffer]) -> Self {
+        match a.len() {
+            0 => Self([None, None]),
+            1 => Self([Some(&a[0]), None]),
+            _ => Self([Some(&a[0]), Some(&a[1])]),
+        }
+    }
+
+    /// Returns the number of [`Buffer`] in this collection
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.0[0].is_some() as usize + self.0[1].is_some() as usize
+    }
+
+    /// Returns `true` if this collection is empty
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.0[0].is_none() && self.0[1].is_none()
+    }
+
+    #[inline]
+    pub fn iter(&self) -> IntoIter<'a> {
+        self.into_iter()
+    }
+
+    /// Converts this [`Buffers`] to a `Vec<Buffer>`
+    #[inline]
+    pub fn to_vec(&self) -> Vec<Buffer> {
+        self.iter().cloned().collect()
+    }
+}
+
+impl<'a> Index<usize> for Buffers<'a> {
+    type Output = &'a Buffer;
+
+    #[inline]
+    fn index(&self, index: usize) -> &Self::Output {
+        self.0[index].as_ref().unwrap()
+    }
+}
+
+impl<'a> IntoIterator for Buffers<'a> {
+    type Item = &'a Buffer;
+    type IntoIter = IntoIter<'a>;
+
+    #[inline]
+    fn into_iter(self) -> Self::IntoIter {
+        IntoIter(self.0[0].into_iter().chain(self.0[1].into_iter()))
+    }
+}
+
+type OptionIter<'a> = std::option::IntoIter<&'a Buffer>;
+
+/// [`Iterator`] for [`Buffers`]
+pub struct IntoIter<'a>(Chain<OptionIter<'a>, OptionIter<'a>>);
+
+impl<'a> Iterator for IntoIter<'a> {
+    type Item = &'a Buffer;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
+}

arrow-data/src/data/mod.rs

@@ -29,6 +29,9 @@ use std::sync::Arc;
 
 use crate::equal;
 
+mod buffers;
+pub use buffers::*;
+
 #[allow(unused)] // Private until ready (#1176)
 mod bytes;
 #[allow(unused)] // Private until ready (#1176)
@@ -366,9 +369,10 @@ impl ArrayData {
         &self.data_type
     }
 
-    /// Returns a slice of the [`Buffer`]s that hold the data.
-    pub fn buffers(&self) -> &[Buffer] {
-        &self.buffers[..]
+    /// 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 a slice of children [`ArrayData`]. This will be non

arrow-json/src/reader.rs

@@ -2210,7 +2210,7 @@ mod tests {
                 .unwrap();
             // test that the list offsets are correct
             assert_eq!(
-                cc.data().buffers()[0],
+                *cc.data().buffers()[0],
                 Buffer::from_slice_ref([0i32, 2, 2, 4, 5])
             );
             let cc = as_boolean_array(cc.values());
@@ -2230,7 +2230,7 @@ mod tests {
                 .unwrap();
             // test that the list offsets are correct
             assert_eq!(
-                dd.data().buffers()[0],
+                *dd.data().buffers()[0],
                 Buffer::from_slice_ref([0i32, 1, 1, 2, 6])
             );
 
@@ -2374,7 +2374,7 @@ mod tests {
         let read: &ListArray = read.as_any().downcast_ref::<ListArray>().unwrap();
         let expected = expected.as_any().downcast_ref::<ListArray>().unwrap();
         assert_eq!(
-            read.data().buffers()[0],
+            *read.data().buffers()[0],
             Buffer::from_slice_ref([0i32, 2, 3, 6, 6, 6, 7])
         );
         // compare list null buffers

arrow/src/ffi.rs

@@ -1158,7 +1158,7 @@ mod tests {
         // Check type ids
         assert_eq!(
             Buffer::from_slice_ref(&expected_type_ids),
-            array.data().buffers()[0]
+            *array.data().buffers()[0]
         );
         for (i, id) in expected_type_ids.iter().enumerate() {
             assert_eq!(id, &array.type_id(i));
@@ -1222,7 +1222,7 @@ mod tests {
         // Check type ids
         assert_eq!(
             Buffer::from_slice_ref(&expected_type_ids),
-            array.data().buffers()[0]
+            *array.data().buffers()[0]
         );
         for (i, id) in expected_type_ids.iter().enumerate() {
             assert_eq!(id, &array.type_id(i));