polars_arrow/legacy/trusted_len/
boolean.rs

1use crate::array::BooleanArray;
2use crate::bitmap::utils::set_bit_unchecked;
3use crate::bitmap::{BitmapBuilder, MutableBitmap};
4use crate::datatypes::ArrowDataType;
5use crate::legacy::array::default_arrays::FromData;
6use crate::legacy::trusted_len::FromIteratorReversed;
7use crate::legacy::utils::FromTrustedLenIterator;
8use crate::trusted_len::TrustedLen;
9
10impl FromTrustedLenIterator<Option<bool>> for BooleanArray {
11    fn from_iter_trusted_length<I: IntoIterator<Item = Option<bool>>>(iter: I) -> Self
12    where
13        I::IntoIter: TrustedLen,
14    {
15        // Soundness
16        // Trait system bounded to TrustedLen
17        unsafe { BooleanArray::from_trusted_len_iter_unchecked(iter.into_iter()) }
18    }
19}
20impl FromTrustedLenIterator<bool> for BooleanArray {
21    fn from_iter_trusted_length<I: IntoIterator<Item = bool>>(iter: I) -> Self
22    where
23        I::IntoIter: TrustedLen,
24    {
25        // Soundness
26        // Trait system bounded to TrustedLen
27        unsafe {
28            BooleanArray::from_data_default(
29                BitmapBuilder::from_trusted_len_iter(iter.into_iter()).freeze(),
30                None,
31            )
32        }
33    }
34}
35
36impl FromIteratorReversed<bool> for BooleanArray {
37    fn from_trusted_len_iter_rev<I: TrustedLen<Item = bool>>(iter: I) -> Self {
38        let size = iter.size_hint().1.unwrap();
39
40        let mut vals = MutableBitmap::from_len_zeroed(size);
41        let vals_slice = vals.as_mut_slice();
42        unsafe {
43            let mut offset = size;
44            iter.for_each(|item| {
45                offset -= 1;
46                if item {
47                    set_bit_unchecked(vals_slice, offset, true);
48                }
49            });
50        }
51        BooleanArray::new(ArrowDataType::Boolean, vals.into(), None)
52    }
53}
54
55impl FromIteratorReversed<Option<bool>> for BooleanArray {
56    fn from_trusted_len_iter_rev<I: TrustedLen<Item = Option<bool>>>(iter: I) -> Self {
57        let size = iter.size_hint().1.unwrap();
58
59        let mut vals = MutableBitmap::from_len_zeroed(size);
60        let mut validity = MutableBitmap::with_capacity(size);
61        validity.extend_constant(size, true);
62        let validity_slice = validity.as_mut_slice();
63        let vals_slice = vals.as_mut_slice();
64        unsafe {
65            let mut offset = size;
66
67            iter.for_each(|opt_item| {
68                offset -= 1;
69                match opt_item {
70                    Some(item) => {
71                        if item {
72                            // Set value (validity bit is already true).
73                            set_bit_unchecked(vals_slice, offset, true);
74                        }
75                    },
76                    None => {
77                        // Unset validity bit.
78                        set_bit_unchecked(validity_slice, offset, false)
79                    },
80                }
81            });
82        }
83        BooleanArray::new(ArrowDataType::Boolean, vals.into(), Some(validity.into()))
84    }
85}