Add some flexibility to Nullable

mbrobbel · Feb 14, 2023 · 4949c03 · 4949c03
1 parent 411b13d
commit 4949c03
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 15 deletions.
diff --git a/src/array/fixed_size_primitive.rs b/src/array/fixed_size_primitive.rs
@@ -246,11 +246,11 @@ mod tests {
     fn buffers() {
         let input = [1u64, 2, 3, 4, 5, 6, 7];
         let array = input.iter().copied().collect::<Uint64Array>();
-        assert_eq!(array.buffer_ref().as_slice(), input);
+        assert_eq!(array.buffer_ref(), input);
 
         let input = [Some(1u8), None, Some(3), Some(4)];
         let array = input.iter().copied().collect::<Uint8Array<true>>();
-        assert_eq!(array.buffer_ref().as_slice(), [1u8, u8::default(), 3, 4]);
+        assert_eq!(array.buffer_ref(), [1u8, u8::default(), 3, 4]);
         assert_eq!(
             array.validity_bitmap(),
             [true, false, true, true].as_slice()

diff --git a/src/buffer.rs b/src/buffer.rs
@@ -38,7 +38,7 @@ where
 
 // Any type that can be borrowed as a slice of some `Primitive` can be used as a
 // Buffer.
-impl<T, U> Buffer<T> for U
+impl<T, U: ?Sized> Buffer<T> for U
 where
     T: Primitive,
     U: Borrow<[T]>,
@@ -114,7 +114,7 @@ where
 /// A reference to a buffer.
 pub trait BufferRef {
     type Element: Primitive;
-    type Buffer: Buffer<Self::Element>;
+    type Buffer: ?Sized + Buffer<Self::Element>;
 
     /// Returns a reference to the buffer.
     fn buffer_ref(&self) -> &Self::Buffer;
@@ -124,11 +124,25 @@ impl<T> BufferRef for Vec<T>
 where
     T: Primitive,
 {
-    type Buffer = Vec<T>;
+    type Buffer = [T];
     type Element = T;
 
     fn buffer_ref(&self) -> &Self::Buffer {
-        self
+        self.as_slice()
+    }
+}
+
+impl<T, const N: usize> BufferRef for Vec<[T; N]>
+where
+    T: Primitive,
+{
+    type Buffer = [T];
+    type Element = T;
+
+    fn buffer_ref(&self) -> &Self::Buffer {
+        // self.flatten() is nightly
+        // SAFETY: `[T]` is layout-identical to `[T; N]`
+        unsafe { std::slice::from_raw_parts(self.as_ptr().cast(), self.len() * N) }
     }
 }
 

diff --git a/src/nullable.rs b/src/nullable.rs
@@ -53,19 +53,20 @@ where
     }
 }
 
-impl<T, Data, BitmapBuffer> FromIterator<(bool, T)> for Nullable<Data, BitmapBuffer>
+impl<T, U, Data, BitmapBuffer> FromIterator<(bool, U)> for Nullable<Data, BitmapBuffer>
 where
     T: Default,
+    U: IntoIterator<Item = T>,
     Data: Default + Extend<T>,
     BitmapBuffer: BufferAlloc<u8>,
 {
-    fn from_iter<I: IntoIterator<Item = (bool, T)>>(iter: I) -> Self {
+    fn from_iter<I: IntoIterator<Item = (bool, U)>>(iter: I) -> Self {
         let mut data = Data::default();
         data.extend(Some(T::default()));
         let validity = iter
             .into_iter()
             .map(|(valid, item)| {
-                data.extend(Some(item));
+                data.extend(item);
                 valid
             })
             .collect();
@@ -163,6 +164,34 @@ mod tests {
         let nullable = input.into_iter().collect::<Nullable<Vec<_>>>();
         assert_eq!(nullable.buffer_ref(), &[1, 2, 3, 4, u32::default(), 42]);
         assert_eq!(nullable.validity_bitmap().buffer_ref(), &[0b00101111u8]);
+
+        let input = [
+            Some([1, 1]),
+            Some([2, 2]),
+            Some([3, 3]),
+            Some([4, 4]),
+            None,
+            Some([42, 42]),
+        ];
+        let nullable = input.into_iter().collect::<Nullable<Vec<_>>>();
+        assert_eq!(
+            nullable.buffer_ref(),
+            &[
+                1,
+                1,
+                2,
+                2,
+                3,
+                3,
+                4,
+                4,
+                u32::default(),
+                u32::default(),
+                42,
+                42
+            ]
+        );
+        assert_eq!(nullable.validity_bitmap().buffer_ref(), &[0b00101111u8]);
     }
 
     #[test]

diff --git a/src/offset/buffer.rs b/src/offset/buffer.rs
@@ -6,7 +6,7 @@ pub trait OffsetBuffer<T>
 where
     T: OffsetElement,
 {
-    type Buffer: Buffer<T>;
+    type Buffer: ?Sized + Buffer<T>;
 
     fn offset_buffer(&self) -> &Self::Buffer;
 }
diff --git a/src/offset/mod.rs b/src/offset/mod.rs
@@ -138,10 +138,10 @@ where
             .map(|opt| match opt {
                 Some(item) => {
                     state += OffsetElement::try_from(item.len()).unwrap();
-                    data.extend(item.into_iter());
-                    (true, state)
+                    data.extend(item);
+                    (true, std::iter::once(state))
                 }
-                None => (false, state),
+                None => (false, std::iter::once(state)),
             })
             .collect();
         Self {
@@ -159,9 +159,12 @@ where
     OffsetElement: self::OffsetElement,
     OffsetBuffer: Buffer<OffsetElement> + Validity<NULLABLE>,
     BitmapBuffer: Buffer<u8>,
-    <OffsetBuffer as Validity<NULLABLE>>::Storage<BitmapBuffer>: BufferRef<Buffer = OffsetBuffer>,
+    <OffsetBuffer as Validity<NULLABLE>>::Storage<BitmapBuffer>: BufferRef,
+    <<OffsetBuffer as Validity<NULLABLE>>::Storage<BitmapBuffer> as BufferRef>::Buffer:
+        Buffer<OffsetElement>,
 {
-    type Buffer = OffsetBuffer;
+    type Buffer =
+        <<OffsetBuffer as Validity<NULLABLE>>::Storage<BitmapBuffer> as BufferRef>::Buffer;
 
     fn offset_buffer(&self) -> &Self::Buffer {
         self.offsets.buffer_ref()