polars_arrow/bitmap/utils/chunk_iterator/
chunks_exact.rs

1use std::slice::ChunksExact;
2
3use super::{BitChunk, BitChunkIterExact};
4use crate::trusted_len::TrustedLen;
5
6/// An iterator over a slice of bytes in [`BitChunk`]s.
7#[derive(Debug)]
8pub struct BitChunksExact<'a, T: BitChunk> {
9    iter: ChunksExact<'a, u8>,
10    remainder: &'a [u8],
11    remainder_len: usize,
12    phantom: std::marker::PhantomData<T>,
13}
14
15impl<'a, T: BitChunk> BitChunksExact<'a, T> {
16    /// Creates a new [`BitChunksExact`].
17    #[inline]
18    pub fn new(bitmap: &'a [u8], length: usize) -> Self {
19        assert!(length <= bitmap.len() * 8);
20        let size_of = size_of::<T>();
21
22        let bitmap = &bitmap[..length.saturating_add(7) / 8];
23
24        let split = (length / 8 / size_of) * size_of;
25        let (chunks, remainder) = bitmap.split_at(split);
26        let remainder_len = length - chunks.len() * 8;
27        let iter = chunks.chunks_exact(size_of);
28
29        Self {
30            iter,
31            remainder,
32            remainder_len,
33            phantom: std::marker::PhantomData,
34        }
35    }
36
37    /// Returns the number of chunks of this iterator
38    #[inline]
39    pub fn len(&self) -> usize {
40        self.iter.len()
41    }
42
43    /// Returns whether there are still elements in this iterator
44    #[inline]
45    pub fn is_empty(&self) -> bool {
46        self.len() == 0
47    }
48
49    /// Returns the remaining [`BitChunk`]. It is zero iff `len / 8 == 0`.
50    #[inline]
51    pub fn remainder(&self) -> T {
52        let remainder_bytes = self.remainder;
53        if remainder_bytes.is_empty() {
54            return T::zero();
55        }
56        let remainder = match remainder_bytes.try_into() {
57            Ok(a) => a,
58            Err(_) => {
59                let mut remainder = T::zero().to_ne_bytes();
60                remainder_bytes
61                    .iter()
62                    .enumerate()
63                    .for_each(|(index, b)| remainder[index] = *b);
64                remainder
65            },
66        };
67        T::from_ne_bytes(remainder)
68    }
69}
70
71impl<T: BitChunk> Iterator for BitChunksExact<'_, T> {
72    type Item = T;
73
74    #[inline]
75    fn next(&mut self) -> Option<Self::Item> {
76        self.iter.next().map(|x| match x.try_into() {
77            Ok(a) => T::from_ne_bytes(a),
78            Err(_) => unreachable!(),
79        })
80    }
81
82    #[inline]
83    fn size_hint(&self) -> (usize, Option<usize>) {
84        self.iter.size_hint()
85    }
86}
87
88unsafe impl<T: BitChunk> TrustedLen for BitChunksExact<'_, T> {}
89
90impl<T: BitChunk> BitChunkIterExact<T> for BitChunksExact<'_, T> {
91    #[inline]
92    fn remainder(&self) -> T {
93        self.remainder()
94    }
95
96    #[inline]
97    fn remainder_len(&self) -> usize {
98        self.remainder_len
99    }
100}
polars_arrow/bitmap/utils/chunk_iterator/chunks_exact.rs

polars_arrow/bitmap/utils/chunk_iterator/
chunks_exact.rs
