polars_arrow/legacy/
utils.rs

1use crate::array::PrimitiveArray;
2use crate::bitmap::utils::set_bit_unchecked;
3use crate::bitmap::MutableBitmap;
4use crate::datatypes::ArrowDataType;
5use crate::legacy::trusted_len::{FromIteratorReversed, TrustedLenPush};
6use crate::trusted_len::{TrustMyLength, TrustedLen};
7use crate::types::NativeType;
8
9pub trait CustomIterTools: Iterator {
10    /// Turn any iterator in a trusted length iterator
11    ///
12    /// # Safety
13    /// The given length must be correct.
14    #[inline]
15    unsafe fn trust_my_length(self, length: usize) -> TrustMyLength<Self, Self::Item>
16    where
17        Self: Sized,
18    {
19        unsafe { TrustMyLength::new(self, length) }
20    }
21
22    fn collect_trusted<T: FromTrustedLenIterator<Self::Item>>(self) -> T
23    where
24        Self: Sized + TrustedLen,
25    {
26        FromTrustedLenIterator::from_iter_trusted_length(self)
27    }
28
29    fn collect_reversed<T: FromIteratorReversed<Self::Item>>(self) -> T
30    where
31        Self: Sized + TrustedLen,
32    {
33        FromIteratorReversed::from_trusted_len_iter_rev(self)
34    }
35}
36
37pub trait CustomIterToolsSized: Iterator + Sized {}
38
39impl<T: ?Sized> CustomIterTools for T where T: Iterator {}
40
41pub trait FromTrustedLenIterator<A>: Sized {
42    fn from_iter_trusted_length<T: IntoIterator<Item = A>>(iter: T) -> Self
43    where
44        T::IntoIter: TrustedLen;
45}
46
47impl<T> FromTrustedLenIterator<T> for Vec<T> {
48    fn from_iter_trusted_length<I: IntoIterator<Item = T>>(iter: I) -> Self
49    where
50        I::IntoIter: TrustedLen,
51    {
52        let iter = iter.into_iter();
53        let len = iter.size_hint().0;
54        let mut v = Vec::with_capacity(len);
55        v.extend_trusted_len(iter);
56        v
57    }
58}
59
60impl<T: NativeType> FromTrustedLenIterator<Option<T>> for PrimitiveArray<T> {
61    fn from_iter_trusted_length<I: IntoIterator<Item = Option<T>>>(iter: I) -> Self
62    where
63        I::IntoIter: TrustedLen,
64    {
65        let iter = iter.into_iter();
66        unsafe { PrimitiveArray::from_trusted_len_iter_unchecked(iter) }
67    }
68}
69
70impl<T: NativeType> FromTrustedLenIterator<T> for PrimitiveArray<T> {
71    fn from_iter_trusted_length<I: IntoIterator<Item = T>>(iter: I) -> Self
72    where
73        I::IntoIter: TrustedLen,
74    {
75        let iter = iter.into_iter();
76        unsafe { PrimitiveArray::from_trusted_len_values_iter_unchecked(iter) }
77    }
78}
79
80impl<T> FromIteratorReversed<T> for Vec<T> {
81    fn from_trusted_len_iter_rev<I: TrustedLen<Item = T>>(iter: I) -> Self {
82        unsafe {
83            let len = iter.size_hint().1.unwrap();
84            let mut out: Vec<T> = Vec::with_capacity(len);
85            let mut idx = len;
86            for x in iter {
87                debug_assert!(idx > 0);
88                idx -= 1;
89                out.as_mut_ptr().add(idx).write(x);
90            }
91            debug_assert!(idx == 0);
92            out.set_len(len);
93            out
94        }
95    }
96}
97
98impl<T: NativeType> FromIteratorReversed<T> for PrimitiveArray<T> {
99    fn from_trusted_len_iter_rev<I: TrustedLen<Item = T>>(iter: I) -> Self {
100        let vals: Vec<T> = iter.collect_reversed();
101        PrimitiveArray::new(ArrowDataType::from(T::PRIMITIVE), vals.into(), None)
102    }
103}
104
105impl<T: NativeType> FromIteratorReversed<Option<T>> for PrimitiveArray<T> {
106    fn from_trusted_len_iter_rev<I: TrustedLen<Item = Option<T>>>(iter: I) -> Self {
107        let size = iter.size_hint().1.unwrap();
108
109        let mut vals: Vec<T> = Vec::with_capacity(size);
110        let mut validity = MutableBitmap::with_capacity(size);
111        validity.extend_constant(size, true);
112        let validity_slice = validity.as_mut_slice();
113        unsafe {
114            // Set to end of buffer.
115            let mut ptr = vals.as_mut_ptr().add(size);
116            let mut offset = size;
117
118            iter.for_each(|opt_item| {
119                offset -= 1;
120                ptr = ptr.sub(1);
121                match opt_item {
122                    Some(item) => {
123                        std::ptr::write(ptr, item);
124                    },
125                    None => {
126                        std::ptr::write(ptr, T::default());
127                        set_bit_unchecked(validity_slice, offset, false);
128                    },
129                }
130            });
131            vals.set_len(size)
132        }
133        PrimitiveArray::new(
134            ArrowDataType::from(T::PRIMITIVE),
135            vals.into(),
136            Some(validity.into()),
137        )
138    }
139}
polars_arrow/legacy/utils.rs

polars_arrow/legacy/
utils.rs
