polars_core/datatypes/
static_array_collect.rs

1use arrow::array::ArrayFromIter;
2use arrow::bitmap::BitmapBuilder;
3
4use crate::chunked_array::object::{ObjectArray, PolarsObject};
5
6// TODO: more efficient implementations, I really took the short path here.
7impl<'a, T: PolarsObject> ArrayFromIter<&'a T> for ObjectArray<T> {
8    fn arr_from_iter<I: IntoIterator<Item = &'a T>>(iter: I) -> Self {
9        Self::try_arr_from_iter(iter.into_iter().map(|o| -> Result<_, ()> { Ok(Some(o)) })).unwrap()
10    }
11
12    fn try_arr_from_iter<E, I: IntoIterator<Item = Result<&'a T, E>>>(iter: I) -> Result<Self, E> {
13        Self::try_arr_from_iter(iter.into_iter().map(|o| Ok(Some(o?))))
14    }
15}
16
17impl<'a, T: PolarsObject> ArrayFromIter<Option<&'a T>> for ObjectArray<T> {
18    fn arr_from_iter<I: IntoIterator<Item = Option<&'a T>>>(iter: I) -> Self {
19        Self::try_arr_from_iter(iter.into_iter().map(|o| -> Result<_, ()> { Ok(o) })).unwrap()
20    }
21
22    fn try_arr_from_iter<E, I: IntoIterator<Item = Result<Option<&'a T>, E>>>(
23        iter: I,
24    ) -> Result<Self, E> {
25        let iter = iter.into_iter();
26        let size = iter.size_hint().0;
27
28        let mut null_mask_builder = BitmapBuilder::with_capacity(size);
29        let values: Vec<T> = iter
30            .map(|value| match value? {
31                Some(value) => {
32                    null_mask_builder.push(true);
33                    Ok(value.clone())
34                },
35                None => {
36                    null_mask_builder.push(false);
37                    Ok(T::default())
38                },
39            })
40            .collect::<Result<Vec<T>, E>>()?;
41
42        Ok(ObjectArray::from(values).with_validity(null_mask_builder.into_opt_validity()))
43    }
44}