Arrays have multiple buffers

encodings/bytebool/src/array.rs

@@ -46,7 +46,7 @@ impl ByteBoolArray {
             Arc::new(ByteBoolMetadata {
                 validity: validity.to_metadata(length)?,
             }),
-            Some(buffer.into_byte_buffer()),
+            [buffer.into_byte_buffer()].into(),
             validity.into_array().into_iter().collect::<Vec<_>>().into(),
             StatsSet::default(),
         )?
@@ -63,7 +63,7 @@ impl ByteBoolArray {
 
     pub fn buffer(&self) -> &ByteBuffer {
         self.as_ref()
-            .byte_buffer()
+            .byte_buffer(0)
             .vortex_expect("ByteBoolArray is missing the underlying buffer")
     }
 

encodings/fastlanes/src/bitpacking/mod.rs

@@ -150,7 +150,7 @@ impl BitPackedArray {
             dtype,
             length,
             Arc::new(metadata),
-            Some(packed),
+            [packed].into(),
             children.into(),
             StatsSet::default(),
         )?
@@ -160,7 +160,7 @@ impl BitPackedArray {
     #[inline]
     pub fn packed(&self) -> &ByteBuffer {
         self.as_ref()
-            .byte_buffer()
+            .byte_buffer(0)
             .vortex_expect("BitPackedArray must contain packed buffer")
     }
 

encodings/roaring/src/boolean/mod.rs

@@ -55,7 +55,7 @@ impl RoaringBoolArray {
             DType::Bool(Nullability::NonNullable),
             length,
             Arc::new(RoaringBoolMetadata),
-            Some(ByteBuffer::from(bitmap.serialize::<Native>())),
+            [ByteBuffer::from(bitmap.serialize::<Native>())].into(),
             vec![].into(),
             stats,
         )?
@@ -77,7 +77,7 @@ impl RoaringBoolArray {
 
     pub fn buffer(&self) -> &ByteBuffer {
         self.as_ref()
-            .byte_buffer()
+            .byte_buffer(0)
             .vortex_expect("Missing buffer in PrimitiveArray")
     }
 }

encodings/roaring/src/integer/mod.rs

@@ -68,7 +68,7 @@ impl RoaringIntArray {
             DType::Primitive(ptype, NonNullable),
             length,
             Arc::new(RoaringIntMetadata { ptype }),
-            Some(ByteBuffer::from(bitmap.serialize::<Portable>())),
+            [ByteBuffer::from(bitmap.serialize::<Portable>())].into(),
             vec![].into(),
             stats,
         )?
@@ -78,7 +78,7 @@ impl RoaringIntArray {
     pub fn owned_bitmap(&self) -> Bitmap {
         Bitmap::deserialize::<Portable>(
             self.as_ref()
-                .byte_buffer()
+                .byte_buffer(0)
                 .vortex_expect("RoaringBoolArray buffer is missing")
                 .as_ref(),
         )
@@ -139,7 +139,7 @@ impl VisitorVTable<RoaringIntArray> for RoaringIntEncoding {
         visitor.visit_buffer(
             array
                 .as_ref()
-                .byte_buffer()
+                .byte_buffer(0)
                 .vortex_expect("Missing buffer in RoaringIntArray"),
         )
     }

vortex-array/src/array/bool/mod.rs

@@ -42,14 +42,14 @@ impl BoolArray {
     /// Access internal array buffer
     pub fn buffer(&self) -> &ByteBuffer {
         self.as_ref()
-            .byte_buffer()
+            .byte_buffer(0)
             .vortex_expect("Missing buffer in BoolArray")
     }
 
     /// Convert array into its internal buffer
     pub fn into_buffer(self) -> ByteBuffer {
         self.into_array()
-            .into_byte_buffer()
+            .into_byte_buffer(0)
             .vortex_expect("BoolArray must have a buffer")
     }
 
@@ -127,7 +127,7 @@ impl BoolArray {
                 validity: validity.to_metadata(buffer_len)?,
                 first_byte_bit_offset,
             }),
-            Some(ByteBuffer::from_arrow_buffer(inner, Alignment::of::<u8>())),
+            vec![ByteBuffer::from_arrow_buffer(inner, Alignment::of::<u8>())].into(),
             validity.into_array().into_iter().collect(),
             StatsSet::default(),
         )?

vortex-array/src/array/primitive/mod.rs

@@ -47,7 +47,7 @@ impl PrimitiveArray {
             Arc::new(PrimitiveMetadata {
                 validity: validity.to_metadata(len).vortex_expect("Invalid validity"),
             }),
-            Some(buffer.into_byte_buffer()),
+            [buffer.into_byte_buffer()].into(),
             validity.into_array().into_iter().collect(),
             StatsSet::default(),
         )
@@ -102,16 +102,14 @@ impl PrimitiveArray {
     }
 
     pub fn byte_buffer(&self) -> &ByteBuffer {
-        let buffer = self
-            .as_ref()
-            .byte_buffer()
-            .vortex_expect("Missing buffer in PrimitiveArray");
-        buffer
+        self.as_ref()
+            .byte_buffer(0)
+            .vortex_expect("Missing buffer in PrimitiveArray")
     }
 
     pub fn into_byte_buffer(self) -> ByteBuffer {
         self.into_array()
-            .into_byte_buffer()
+            .into_byte_buffer(0)
             .vortex_expect("PrimitiveArray must have a buffer")
     }
 
@@ -136,7 +134,7 @@ impl PrimitiveArray {
         }
         Buffer::from_byte_buffer(
             self.into_array()
-                .into_byte_buffer()
+                .into_byte_buffer(0)
                 .vortex_expect("PrimitiveArray must have a buffer"),
         )
     }

vortex-array/src/data/mod.rs

@@ -61,7 +61,7 @@ impl ArrayData {
         dtype: DType,
         len: usize,
         metadata: Arc<dyn ArrayMetadata>,
-        buffer: Option<ByteBuffer>,
+        buffers: Arc<[ByteBuffer]>,
         children: Arc<[ArrayData]>,
         statistics: StatsSet,
     ) -> VortexResult<Self> {
@@ -70,7 +70,7 @@ impl ArrayData {
             dtype,
             len,
             metadata,
-            buffer,
+            buffers,
             children,
             stats_set: Arc::new(RwLock::new(statistics)),
             #[cfg(feature = "canonical_counter")]
@@ -247,7 +247,7 @@ impl ArrayData {
             .iter()
             .map(|child| child.cumulative_nbuffers())
             .sum::<usize>()
-            + if self.byte_buffer().is_some() { 1 } else { 0 }
+            + self.nbuffers()
     }
 
     /// Return the buffer offsets and the total length of all buffers, assuming the given alignment.
@@ -257,7 +257,7 @@ impl ArrayData {
         let mut offset = 0;
 
         for col_data in self.depth_first_traversal() {
-            if let Some(buffer) = col_data.byte_buffer() {
+            for buffer in col_data.byte_buffers() {
                 offsets.push(offset as u64);
 
                 let buffer_size = buffer.len();
@@ -320,17 +320,30 @@ impl ArrayData {
         }
     }
 
-    pub fn byte_buffer(&self) -> Option<&ByteBuffer> {
+    pub fn nbuffers(&self) -> usize {
         match &self.0 {
-            InnerArrayData::Owned(d) => d.byte_buffer(),
-            InnerArrayData::Viewed(v) => v.byte_buffer(),
+            InnerArrayData::Owned(o) => o.buffers.len(),
+            InnerArrayData::Viewed(v) => v.nbuffers(),
         }
     }
 
-    pub fn into_byte_buffer(self) -> Option<ByteBuffer> {
+    pub fn byte_buffer(&self, index: usize) -> Option<&ByteBuffer> {
+        match &self.0 {
+            InnerArrayData::Owned(d) => d.byte_buffer(index),
+            InnerArrayData::Viewed(v) => v.buffer(index),
+        }
+    }
+
+    pub fn byte_buffers(&self) -> impl Iterator<Item = ByteBuffer> + '_ {
+        (0..self.nbuffers())
+            .map(|i| self.byte_buffer(i).vortex_expect("missing declared buffer"))
+            .cloned()
+    }
+
+    pub fn into_byte_buffer(self, index: usize) -> Option<ByteBuffer> {
         match self.0 {
-            InnerArrayData::Owned(d) => d.into_byte_buffer(),
-            InnerArrayData::Viewed(v) => v.byte_buffer().cloned(),
+            InnerArrayData::Owned(d) => d.into_byte_buffer(index),
+            InnerArrayData::Viewed(v) => v.buffer(index).cloned(),
         }
     }
 

vortex-array/src/data/owned.rs

@@ -16,7 +16,7 @@ pub(super) struct OwnedArrayData {
     pub(super) dtype: DType, // FIXME(ngates): Arc?
     pub(super) len: usize,
     pub(super) metadata: Arc<dyn ArrayMetadata>,
-    pub(super) buffer: Option<ByteBuffer>,
+    pub(super) buffers: Arc<[ByteBuffer]>,
     pub(super) children: Arc<[ArrayData]>,
     pub(super) stats_set: Arc<RwLock<StatsSet>>,
     #[cfg(feature = "canonical_counter")]
@@ -28,12 +28,14 @@ impl OwnedArrayData {
         &self.metadata
     }
 
-    pub fn byte_buffer(&self) -> Option<&ByteBuffer> {
-        self.buffer.as_ref()
+    pub fn byte_buffer(&self, index: usize) -> Option<&ByteBuffer> {
+        self.buffers.get(index)
     }
 
-    pub fn into_byte_buffer(self) -> Option<ByteBuffer> {
-        self.buffer
+    pub fn into_byte_buffer(self, index: usize) -> Option<ByteBuffer> {
+        // While this does require a clone, we're still "into" because it makes sure the self
+        // reference is dropped.
+        self.buffers.get(index).cloned()
     }
 
     // We want to allow these panics because they are indicative of implementation error.

vortex-array/src/data/viewed.rs

@@ -103,12 +103,22 @@ impl ViewedArrayData {
         collector.children()
     }
 
-    pub fn byte_buffer(&self) -> Option<&ByteBuffer> {
+    pub fn nbuffers(&self) -> usize {
         self.flatbuffer()
             .buffers()
-            .and_then(|buffers| {
-                assert!(buffers.len() <= 1, "Array: expected at most one buffer");
-                (!buffers.is_empty()).then(|| buffers.get(0) as usize)
+            .map(|b| b.len())
+            .unwrap_or_default()
+    }
+
+    pub fn buffer(&self, index: usize) -> Option<&ByteBuffer> {
+        self.flatbuffer()
+            .buffers()
+            .map(|buffers| {
+                assert!(
+                    index < buffers.len(),
+                    "ArrayView buffer index out of bounds"
+                );
+                buffers.get(index) as usize
             })
             .map(|idx| &self.buffers[idx])
     }

vortex-array/src/macros.rs

@@ -54,7 +54,7 @@ macro_rules! impl_encoding {
                             dtype,
                             len,
                             std::sync::Arc::new(metadata),
-                            None,
+                            vec![].into(),
                             children,
                             stats
                     )?)

vortex-ipc/src/messages/decoder.rs

@@ -306,7 +306,7 @@ mod test {
         // Constant arrays have no buffers
         let array = ConstantArray::new(10i32, 20).into_array();
         assert!(
-            array.byte_buffer().is_none(),
+            array.byte_buffer(0).is_none(),
             "Array should have no buffers"
         );
         write_and_read(array);

vortex-ipc/src/messages/encoder.rs

@@ -86,7 +86,7 @@ impl MessageEncoder {
 
                 let mut fb_buffers = vec![];
                 for child in array.depth_first_traversal() {
-                    if let Some(buffer) = child.byte_buffer() {
+                    for buffer in child.byte_buffers() {
                         let end_excl_padding = self.pos + buffer.len();
                         let end_incl_padding = end_excl_padding.next_multiple_of(self.alignment);
                         let padding = u16::try_from(end_incl_padding - end_excl_padding)
@@ -198,12 +198,9 @@ impl WriteFlatBuffer for ArrayWriter<'_> {
         // The second tuple element holds the buffer_index for this Array subtree. If this array
         // has a buffer, that is its buffer index. If it does not, that buffer index belongs
         // to one of the children.
-        let child_buffer_idx = self.buffer_idx
-            + if self.array.byte_buffer().is_some() {
-                1
-            } else {
-                0
-            };
+        let nbuffers = u16::try_from(self.array.nbuffers())
+            .vortex_expect("Array can have at most u16::MAX buffers");
+        let child_buffer_idx = self.buffer_idx + nbuffers;
 
         let children = self
             .array
@@ -225,12 +222,7 @@ impl WriteFlatBuffer for ArrayWriter<'_> {
             .collect_vec();
         let children = Some(fbb.create_vector(&children));
 
-        let buffers = self
-            .array
-            .byte_buffer()
-            .is_some()
-            .then_some(self.buffer_idx)
-            .map(|buffer_idx| fbb.create_vector_from_iter(std::iter::once(buffer_idx)));
+        let buffers = Some(fbb.create_vector_from_iter((0..nbuffers).map(|i| i + self.buffer_idx)));
 
         let stats = Some(self.array.statistics().write_flatbuffer(fbb));