arrow2/types/
bit_chunk.rs

1use std::{
2    fmt::Binary,
3    ops::{BitAndAssign, Not, Shl, ShlAssign, ShrAssign},
4};
5
6use num_traits::PrimInt;
7
8use super::NativeType;
9
10/// A chunk of bits. This is used to create masks of a given length
11/// whose width is `1` bit. In `portable_simd` notation, this corresponds to `m1xY`.
12///
13/// This (sealed) trait is implemented for [`u8`], [`u16`], [`u32`] and [`u64`].
14pub trait BitChunk:
15    super::private::Sealed
16    + PrimInt
17    + NativeType
18    + Binary
19    + ShlAssign
20    + Not<Output = Self>
21    + ShrAssign<usize>
22    + ShlAssign<usize>
23    + Shl<usize, Output = Self>
24    + BitAndAssign
25{
26    /// convert itself into bytes.
27    fn to_ne_bytes(self) -> Self::Bytes;
28    /// convert itself from bytes.
29    fn from_ne_bytes(v: Self::Bytes) -> Self;
30}
31
32macro_rules! bit_chunk {
33    ($ty:ty) => {
34        impl BitChunk for $ty {
35            #[inline(always)]
36            fn to_ne_bytes(self) -> Self::Bytes {
37                self.to_ne_bytes()
38            }
39
40            #[inline(always)]
41            fn from_ne_bytes(v: Self::Bytes) -> Self {
42                Self::from_ne_bytes(v)
43            }
44        }
45    };
46}
47
48bit_chunk!(u8);
49bit_chunk!(u16);
50bit_chunk!(u32);
51bit_chunk!(u64);
52
53/// An [`Iterator<Item=bool>`] over a [`BitChunk`]. This iterator is often
54/// compiled to SIMD.
55/// The [LSB](https://en.wikipedia.org/wiki/Bit_numbering#Least_significant_bit) corresponds
56/// to the first slot, as defined by the arrow specification.
57/// # Example
58/// ```
59/// use arrow2::types::BitChunkIter;
60/// let a = 0b00010000u8;
61/// let iter = BitChunkIter::new(a, 7);
62/// let r = iter.collect::<Vec<_>>();
63/// assert_eq!(r, vec![false, false, false, false, true, false, false]);
64/// ```
65pub struct BitChunkIter<T: BitChunk> {
66    value: T,
67    mask: T,
68    remaining: usize,
69}
70
71impl<T: BitChunk> BitChunkIter<T> {
72    /// Creates a new [`BitChunkIter`] with `len` bits.
73    #[inline]
74    pub fn new(value: T, len: usize) -> Self {
75        assert!(len <= std::mem::size_of::<T>() * 8);
76        Self {
77            value,
78            remaining: len,
79            mask: T::one(),
80        }
81    }
82}
83
84impl<T: BitChunk> Iterator for BitChunkIter<T> {
85    type Item = bool;
86
87    #[inline]
88    fn next(&mut self) -> Option<Self::Item> {
89        if self.remaining == 0 {
90            return None;
91        };
92        let result = Some(self.value & self.mask != T::zero());
93        self.remaining -= 1;
94        self.mask <<= 1;
95        result
96    }
97
98    #[inline]
99    fn size_hint(&self) -> (usize, Option<usize>) {
100        (self.remaining, Some(self.remaining))
101    }
102}
103
104// # Safety
105// a mathematical invariant of this iterator
106unsafe impl<T: BitChunk> crate::trusted_len::TrustedLen for BitChunkIter<T> {}
107
108/// An [`Iterator<Item=usize>`] over a [`BitChunk`] returning the index of each bit set in the chunk
109/// See <https://lemire.me/blog/2018/03/08/iterating-over-set-bits-quickly-simd-edition/> for details
110/// # Example
111/// ```
112/// use arrow2::types::BitChunkOnes;
113/// let a = 0b00010000u8;
114/// let iter = BitChunkOnes::new(a);
115/// let r = iter.collect::<Vec<_>>();
116/// assert_eq!(r, vec![4]);
117/// ```
118pub struct BitChunkOnes<T: BitChunk> {
119    value: T,
120    remaining: usize,
121}
122
123impl<T: BitChunk> BitChunkOnes<T> {
124    /// Creates a new [`BitChunkOnes`] with `len` bits.
125    #[inline]
126    pub fn new(value: T) -> Self {
127        Self {
128            value,
129            remaining: value.count_ones() as usize,
130        }
131    }
132
133    #[inline]
134    #[cfg(feature = "compute_filter")]
135    pub(crate) fn from_known_count(value: T, remaining: usize) -> Self {
136        Self { value, remaining }
137    }
138}
139
140impl<T: BitChunk> Iterator for BitChunkOnes<T> {
141    type Item = usize;
142
143    #[inline]
144    fn next(&mut self) -> Option<Self::Item> {
145        if self.remaining == 0 {
146            return None;
147        }
148        let v = self.value.trailing_zeros() as usize;
149        self.value &= self.value - T::one();
150
151        self.remaining -= 1;
152        Some(v)
153    }
154
155    #[inline]
156    fn size_hint(&self) -> (usize, Option<usize>) {
157        (self.remaining, Some(self.remaining))
158    }
159}
160
161// # Safety
162// a mathematical invariant of this iterator
163unsafe impl<T: BitChunk> crate::trusted_len::TrustedLen for BitChunkOnes<T> {}
arrow2/types/bit_chunk.rs

arrow2/types/
bit_chunk.rs
