Speed up fixed-sized array iteration

crates/re_log_types/src/field_types/mod.rs

@@ -4,7 +4,7 @@
 //! schemas are additionally documented in doctests.
 
 use arrow2::{
-    array::{FixedSizeListArray, MutableFixedSizeListArray},
+    array::{FixedSizeListArray, MutableFixedSizeListArray, PrimitiveArray},
     datatypes::{DataType, Field},
 };
 use arrow2_convert::{
@@ -123,8 +123,8 @@ pub type Result<T> = std::result::Result<T, FieldError>;
 ///
 /// #[derive(ArrowField, ArrowSerialize, ArrowDeserialize)]
 /// pub struct ConvertibleType {
-///     #[arrow_field(type = "FixedSizeArrayField<bool,2>")]
-///     data: [bool; 2],
+///     #[arrow_field(type = "FixedSizeArrayField<u32,2>")]
+///     data: [u32; 2],
 /// }
 /// ```
 pub struct FixedSizeArrayField<T, const SIZE: usize>(std::marker::PhantomData<T>);
@@ -169,27 +169,92 @@ where
     }
 }
 
+pub struct FastFixedSizeArrayIter<'a, T, const SIZE: usize>
+where
+    T: arrow2::types::NativeType,
+{
+    offset: usize,
+    array: &'a FixedSizeListArray,
+    values: &'a PrimitiveArray<T>,
+}
+
+impl<'a, T, const SIZE: usize> Iterator for FastFixedSizeArrayIter<'a, T, SIZE>
+where
+    T: arrow2::types::NativeType,
+{
+    type Item = Option<[T; SIZE]>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.offset < self.array.len() {
+            if let Some(validity) = self.array.validity() {
+                if !validity.get_bit(self.offset) {
+                    self.offset += 1;
+                    return Some(None);
+                }
+            }
+
+            let out: [T; SIZE] =
+                array_init::array_init(|i: usize| self.values.value(self.offset * SIZE + i));
+            self.offset += 1;
+            Some(Some(out))
+        } else {
+            None
+        }
+    }
+}
+
+pub struct FastFixedSizeListArray<T, const SIZE: usize>(std::marker::PhantomData<T>);
+
+impl<'a, T, const SIZE: usize> IntoIterator for &'a FastFixedSizeListArray<T, SIZE>
+where
+    T: arrow2::types::NativeType,
+{
+    type Item = Option<[T; SIZE]>;
+
+    type IntoIter = FastFixedSizeArrayIter<'a, T, SIZE>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        panic!("Use iter_from_array_ref. This is a quirk of the way the traits work in arrow2_convert.");
+    }
+}
+
+impl<T, const SIZE: usize> ArrowArray for FastFixedSizeListArray<T, SIZE>
+where
+    T: arrow2::types::NativeType,
+{
+    type BaseArrayType = FixedSizeListArray;
+
+    fn iter_from_array_ref(b: &dyn arrow2::array::Array) -> <&Self as IntoIterator>::IntoIter {
+        let array = b.as_any().downcast_ref::<Self::BaseArrayType>().unwrap();
+        let values = array
+            .values()
+            .as_any()
+            .downcast_ref::<PrimitiveArray<T>>()
+            .unwrap();
+        FastFixedSizeArrayIter::<T, SIZE> {
+            offset: 0,
+            array,
+            values,
+        }
+    }
+}
+
 impl<T, const SIZE: usize> ArrowDeserialize for FixedSizeArrayField<T, SIZE>
 where
-    T: ArrowDeserialize + ArrowEnableVecForType + ArrowField<Type = T> + 'static,
+    T: arrow2::types::NativeType
+        + ArrowDeserialize
+        + ArrowEnableVecForType
+        + ArrowField<Type = T>
+        + 'static,
     <T as ArrowDeserialize>::ArrayType: 'static,
     for<'b> &'b <T as ArrowDeserialize>::ArrayType: IntoIterator,
 {
-    type ArrayType = FixedSizeListArray;
+    type ArrayType = FastFixedSizeListArray<T, SIZE>;
 
+    #[inline]
     fn arrow_deserialize(
         v: <&Self::ArrayType as IntoIterator>::Item,
     ) -> Option<<Self as ArrowField>::Type> {
-        if let Some(array) = v {
-            let mut iter = <<T as ArrowDeserialize>::ArrayType as ArrowArray>::iter_from_array_ref(
-                array.as_ref(),
-            )
-            .map(<T as ArrowDeserialize>::arrow_deserialize_internal);
-            let out: Result<[T; SIZE]> =
-                array_init::try_array_init(|_i: usize| iter.next().ok_or(FieldError::BadValue));
-            out.ok()
-        } else {
-            None
-        }
+        v
     }
 }