feat!: add IntoIterator for VariableSizeBinaryArray (#247)

mbrobbel · Sep 12, 2024 · 3e72841 · 3e72841
1 parent 81e92bd
commit 3e72841
Show file tree

Hide file tree

Showing 5 changed files with 113 additions and 67 deletions.
diff --git a/examples/parquet.rs b/examples/parquet.rs
@@ -53,7 +53,7 @@ fn main() {
                 vec![1, 2, 3, 4, 42],
             )])),
             m: NaiveDate::MAX,
-            n: TimeDelta::max_value(),
+            n: TimeDelta::seconds(12345),
         },
         Foo {
             a: 42,
@@ -69,7 +69,7 @@ fn main() {
             k: Utc::now().time(),
             l: None,
             m: NaiveDate::MIN,
-            n: TimeDelta::min_value(),
+            n: TimeDelta::minutes(1234),
         },
     ];
 

diff --git a/src/array/string.rs b/src/array/string.rs
@@ -1,6 +1,6 @@
 //! Array with string values.
 
-use std::str;
+use std::{iter::Map, str};
 
 use super::{Array, VariableSizeBinaryArray};
 use crate::{
@@ -239,51 +239,6 @@ where
     }
 }
 
-/// An iterator over items in a [`StringArray`].
-pub struct StringIntoIter<const NULLABLE: bool, OffsetItem: OffsetElement, Buffer: BufferType>
-where
-    <Buffer as BufferType>::Buffer<OffsetItem>: Validity<NULLABLE>,
-{
-    /// Reference to the array.
-    array: StringArray<NULLABLE, OffsetItem, Buffer>,
-    /// Current index.
-    index: usize,
-}
-
-impl<OffsetItem: OffsetElement, Buffer: BufferType> Iterator
-    for StringIntoIter<false, OffsetItem, Buffer>
-{
-    type Item = String;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.array
-            .index(self.index)
-            .into_iter()
-            .inspect(|_| {
-                self.index += 1;
-            })
-            .next()
-            .map(ToOwned::to_owned)
-    }
-}
-
-impl<OffsetItem: OffsetElement, Buffer: BufferType> Iterator
-    for StringIntoIter<true, OffsetItem, Buffer>
-{
-    type Item = Option<String>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.array
-            .index(self.index)
-            .into_iter()
-            .inspect(|_| {
-                self.index += 1;
-            })
-            .next()
-            .map(|opt| opt.map(ToOwned::to_owned))
-    }
-}
-
 impl<'a, const NULLABLE: bool, OffsetItem: OffsetElement, Buffer: BufferType> IntoIterator
     for &'a StringArray<NULLABLE, OffsetItem, Buffer>
 where
@@ -301,20 +256,47 @@ where
     }
 }
 
-impl<const NULLABLE: bool, OffsetItem: OffsetElement, Buffer: BufferType> IntoIterator
-    for StringArray<NULLABLE, OffsetItem, Buffer>
+impl<OffsetItem: OffsetElement, Buffer: BufferType> IntoIterator
+    for StringArray<false, OffsetItem, Buffer>
 where
-    <Buffer as BufferType>::Buffer<OffsetItem>: Validity<NULLABLE>,
-    StringIntoIter<NULLABLE, OffsetItem, Buffer>: Iterator,
+    VariableSizeBinaryArray<false, OffsetItem, Buffer>: IntoIterator<Item = Vec<u8>>,
 {
-    type Item = <StringIntoIter<NULLABLE, OffsetItem, Buffer> as Iterator>::Item;
-    type IntoIter = StringIntoIter<NULLABLE, OffsetItem, Buffer>;
+    type Item = String;
+    type IntoIter = Map<
+        <VariableSizeBinaryArray<false, OffsetItem, Buffer> as IntoIterator>::IntoIter,
+        fn(<VariableSizeBinaryArray<false, OffsetItem, Buffer> as IntoIterator>::Item) -> String,
+    >;
 
     fn into_iter(self) -> Self::IntoIter {
-        StringIntoIter {
-            array: self,
-            index: 0,
-        }
+        self.0.into_iter().map(|bytes| {
+            // SAFETY:
+            // - String arrays contain valid UTF8.
+            unsafe { String::from_utf8_unchecked(bytes) }
+        })
+    }
+}
+
+impl<OffsetItem: OffsetElement, Buffer: BufferType> IntoIterator
+    for StringArray<true, OffsetItem, Buffer>
+where
+    VariableSizeBinaryArray<true, OffsetItem, Buffer>: IntoIterator<Item = Option<Vec<u8>>>,
+{
+    type Item = Option<String>;
+    type IntoIter = Map<
+        <VariableSizeBinaryArray<true, OffsetItem, Buffer> as IntoIterator>::IntoIter,
+        fn(
+            <VariableSizeBinaryArray<true, OffsetItem, Buffer> as IntoIterator>::Item,
+        ) -> Option<String>,
+    >;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter().map(|opt| {
+            opt.map(|bytes| {
+                // SAFETY:
+                // - String arrays contain valid UTF8.
+                unsafe { String::from_utf8_unchecked(bytes) }
+            })
+        })
     }
 }
 

diff --git a/src/array/variable_size_binary.rs b/src/array/variable_size_binary.rs
@@ -22,7 +22,7 @@ where
 pub type BinaryArray<const NULLABLE: bool = false, Buffer = VecBuffer> =
     VariableSizeBinaryArray<NULLABLE, i32, Buffer>;
 
-/// Variable-size binary elements, using `i64` offset value.
+/// Variable-size binary elements, using `i64` offset values.
 pub type LargeBinaryArray<const NULLABLE: bool = false, Buffer = VecBuffer> =
     VariableSizeBinaryArray<NULLABLE, i64, Buffer>;
 
@@ -192,6 +192,26 @@ where
     }
 }
 
+impl<const NULLABLE: bool, OffsetItem: OffsetElement, Buffer: BufferType> IntoIterator
+    for VariableSizeBinaryArray<NULLABLE, OffsetItem, Buffer>
+where
+    <Buffer as BufferType>::Buffer<u8>: Validity<NULLABLE>,
+    <Buffer as BufferType>::Buffer<OffsetItem>: Validity<NULLABLE>,
+    Offset<FixedSizePrimitiveArray<u8, false, Buffer>, NULLABLE, OffsetItem, Buffer>: IntoIterator,
+{
+    type Item = <Offset<FixedSizePrimitiveArray<u8, false, Buffer>, NULLABLE, OffsetItem, Buffer> as IntoIterator>::Item;
+    type IntoIter = <Offset<
+        FixedSizePrimitiveArray<u8, false, Buffer>,
+        NULLABLE,
+        OffsetItem,
+        Buffer,
+    > as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
 impl<const NULLABLE: bool, OffsetItem: OffsetElement, Buffer: BufferType> Length
     for VariableSizeBinaryArray<NULLABLE, OffsetItem, Buffer>
 where
@@ -259,6 +279,22 @@ mod tests {
         assert_eq!(array_vec.0.offsets, &[0, 1, 1, 3, 4]);
     }
 
+    #[test]
+    fn into_iter() {
+        let input: [&[u8]; 4] = [&[1], &[2, 3], &[4, 5, 6], &[7, 8, 9, 0]];
+        let array = input.into_iter().collect::<VariableSizeBinaryArray>();
+        let output = array.into_iter().collect::<Vec<_>>();
+        assert_eq!(output, input);
+
+        let input_vec = vec![vec![1], vec![], vec![2, 3], vec![4]];
+        let array_vec = input_vec
+            .clone()
+            .into_iter()
+            .collect::<VariableSizeBinaryArray>();
+        let output_vec = array_vec.into_iter().collect::<Vec<_>>();
+        assert_eq!(output_vec, input_vec);
+    }
+
     #[test]
     fn from_iter_nullable() {
         let input: [Option<&[u8]>; 4] = [Some(&[1]), None, Some(&[4, 5, 6]), Some(&[7, 8, 9, 0])];
@@ -286,6 +322,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn into_iter_nullable() {
+        let input: [Option<&[u8]>; 4] = [Some(&[1]), None, Some(&[4, 5, 6]), Some(&[7, 8, 9, 0])];
+        let array = input.into_iter().collect::<VariableSizeBinaryArray<true>>();
+        let output = array.into_iter().collect::<Vec<_>>();
+        assert_eq!(output, input.map(|opt| opt.map(<[u8]>::to_vec)));
+
+        let input_vec = vec![Some(vec![1]), None, Some(vec![2, 3]), Some(vec![4])];
+        let array_vec = input_vec
+            .clone()
+            .into_iter()
+            .collect::<VariableSizeBinaryArray<true>>();
+        let output_vec = array_vec.into_iter().collect::<Vec<_>>();
+        assert_eq!(output_vec, input_vec);
+    }
+
     #[test]
     fn index() {
         let input: [&[u8]; 4] = [&[1], &[2, 3], &[4, 5, 6], &[7, 8, 9, 0]];

diff --git a/src/arrow/array/fixed_size_list.rs b/src/arrow/array/fixed_size_list.rs
@@ -204,9 +204,10 @@ mod tests {
                 .iter()
                 .flatten()
                 .flat_map(|dyn_array| {
-                    let array: StringArray<false, i32, crate::arrow::buffer::ScalarBuffer> =
-                        dyn_array.into();
-                    array.into_iter().collect::<Vec<_>>()
+                    StringArray::<false, i32, crate::arrow::buffer::ScalarBuffer>::from(dyn_array)
+                        .into_iter()
+                        .map(ToOwned::to_owned)
+                        .collect::<Vec<_>>()
                 })
                 .collect::<Vec<_>>(),
             INPUT_NULLABLE
@@ -282,10 +283,6 @@ mod tests {
         let fixed_size_list_array_nullable =
             arrow_array::FixedSizeListArray::from(fixed_size_list_array_nullable_input);
 
-        let owned_input_nullable = INPUT_NULLABLE
-            .into_iter()
-            .map(|item| item.map(|[first, second]| [first.to_owned(), second.to_owned()]))
-            .collect::<Vec<_>>();
         assert_eq!(
             FixedSizeListArray::<
                 2,
@@ -294,7 +291,7 @@ mod tests {
             >::from(fixed_size_list_array_nullable)
             .into_iter()
             .collect::<Vec<_>>(),
-            owned_input_nullable
+            INPUT_NULLABLE
         );
     }
 }
diff --git a/src/offset.rs b/src/offset.rs
@@ -204,6 +204,21 @@ where
     }
 }
 
+impl<const NULLABLE: bool, T, OffsetItem: OffsetElement, Buffer: BufferType> Clone
+    for Offset<T, NULLABLE, OffsetItem, Buffer>
+where
+    T: Clone,
+    <Buffer as BufferType>::Buffer<OffsetItem>: Validity<NULLABLE>,
+    <<Buffer as BufferType>::Buffer<OffsetItem> as Validity<NULLABLE>>::Storage<Buffer>: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            data: self.data.clone(),
+            offsets: self.offsets.clone(),
+        }
+    }
+}
+
 impl<T: Default, OffsetItem: OffsetElement, Buffer: BufferType> Default
     for Offset<T, false, OffsetItem, Buffer>
 where