Skip to content

Commit

Permalink
feat: Add iterators to more types (#3539)
Browse files Browse the repository at this point in the history
# Overview
This PR adds more iterators to daft-core datatypes.

Co-authored-by: Raunak Bhagat <[email protected]>
  • Loading branch information
raunakab and Raunak Bhagat authored Dec 11, 2024
1 parent d103573 commit f23ee37
Showing 1 changed file with 61 additions and 9 deletions.
70 changes: 61 additions & 9 deletions src/daft-core/src/array/iterator.rs
Original file line number Diff line number Diff line change
@@ -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())
}
}

0 comments on commit f23ee37

Please sign in to comment.