From f23ee3732c666085a72e67e54d6a0f69fcb8af39 Mon Sep 17 00:00:00 2001
From: Raunak Bhagat <r@rabh.io>
Date: Tue, 10 Dec 2024 18:00:36 -0800
Subject: [PATCH] feat: Add iterators to more types (#3539)

# Overview
This PR adds more iterators to daft-core datatypes.

Co-authored-by: Raunak Bhagat <rabhagat31@gmail.com>
---
 src/daft-core/src/array/iterator.rs | 70 +++++++++++++++++++++++++----
 1 file changed, 61 insertions(+), 9 deletions(-)

diff --git a/src/daft-core/src/array/iterator.rs b/src/daft-core/src/array/iterator.rs
index 0474f2b560..8ddc0bffc1 100644
--- a/src/daft-core/src/array/iterator.rs
+++ b/src/daft-core/src/array/iterator.rs
@@ -1,26 +1,78 @@
-use arrow2::bitmap::utils::{BitmapIter, ZipValidity};
+use std::{
+    iter::{repeat, Repeat, Take},
+    slice::{ChunksExact, Iter},
+};
 
-use super::{ops::as_arrow::AsArrow, DataArray};
-use crate::datatypes::{BooleanArray, DaftPrimitiveType};
+use arrow2::{
+    array::ArrayValuesIter,
+    bitmap::utils::{BitmapIter, ZipValidity},
+};
+
+use crate::{
+    array::{
+        ops::as_arrow::AsArrow,
+        prelude::{NullArray, Utf8Array},
+        DataArray,
+    },
+    datatypes::{BinaryArray, BooleanArray, DaftPrimitiveType, FixedSizeBinaryArray},
+};
+
+macro_rules! impl_into_iter {
+    (
+        $array:ty
+        , $into_iter:ty
+        $(,)?
+    ) => {
+        impl<'a> IntoIterator for &'a $array {
+            type IntoIter = $into_iter;
+            type Item = <Self::IntoIter as IntoIterator>::Item;
+
+            #[inline]
+            fn into_iter(self) -> Self::IntoIter {
+                self.as_arrow().into_iter()
+            }
+        }
+    };
+}
+
+// yields `bool`s
+impl_into_iter!(BooleanArray, ZipValidity<bool, BitmapIter<'a>, BitmapIter<'a>>);
+
+// both yield `&[u8]`s
+impl_into_iter!(
+    BinaryArray,
+    ZipValidity<&'a [u8], ArrayValuesIter<'a, arrow2::array::BinaryArray<i64>>, BitmapIter<'a>>,
+);
+impl_into_iter!(
+    FixedSizeBinaryArray,
+    ZipValidity<&'a [u8], ChunksExact<'a, u8>, BitmapIter<'a>>,
+);
+
+// yields `&str`s
+impl_into_iter!(
+    Utf8Array,
+    ZipValidity<&'a str, ArrayValuesIter<'a, arrow2::array::Utf8Array<i64>>, BitmapIter<'a>>,
+);
 
 impl<'a, T> IntoIterator for &'a DataArray<T>
 where
     T: DaftPrimitiveType,
 {
-    type Item = Option<&'a T::Native>;
-    type IntoIter = ZipValidity<&'a T::Native, std::slice::Iter<'a, T::Native>, BitmapIter<'a>>;
+    type IntoIter = ZipValidity<&'a T::Native, Iter<'a, T::Native>, BitmapIter<'a>>;
+    type Item = <Self::IntoIter as IntoIterator>::Item;
+
     #[inline]
     fn into_iter(self) -> Self::IntoIter {
         self.as_arrow().into_iter()
     }
 }
 
-impl<'a> IntoIterator for &'a BooleanArray {
-    type Item = Option<bool>;
-    type IntoIter = ZipValidity<bool, BitmapIter<'a>, BitmapIter<'a>>;
+impl<'a> IntoIterator for &'a NullArray {
+    type IntoIter = Take<Repeat<Option<()>>>;
+    type Item = <Self::IntoIter as IntoIterator>::Item;
 
     #[inline]
     fn into_iter(self) -> Self::IntoIter {
-        self.as_arrow().into_iter()
+        repeat(None).take(self.len())
     }
 }