polars_arrow/bitmap/
assign_ops.rs

1use super::utils::{BitChunk, BitChunkIterExact, BitChunksExact};
2use crate::bitmap::{Bitmap, MutableBitmap};
3
4/// Applies a function to every bit of this [`MutableBitmap`] in chunks
5///
6/// This function can be for operations like `!` to a [`MutableBitmap`].
7pub fn unary_assign<T: BitChunk, F: Fn(T) -> T>(bitmap: &mut MutableBitmap, op: F) {
8    let mut chunks = bitmap.bitchunks_exact_mut::<T>();
9
10    chunks.by_ref().for_each(|chunk| {
11        let new_chunk: T = match (chunk as &[u8]).try_into() {
12            Ok(a) => T::from_ne_bytes(a),
13            Err(_) => unreachable!(),
14        };
15        let new_chunk = op(new_chunk);
16        chunk.copy_from_slice(new_chunk.to_ne_bytes().as_ref());
17    });
18
19    if chunks.remainder().is_empty() {
20        return;
21    }
22    let mut new_remainder = T::zero().to_ne_bytes();
23    chunks
24        .remainder()
25        .iter()
26        .enumerate()
27        .for_each(|(index, b)| new_remainder[index] = *b);
28    new_remainder = op(T::from_ne_bytes(new_remainder)).to_ne_bytes();
29
30    let len = chunks.remainder().len();
31    chunks
32        .remainder()
33        .copy_from_slice(&new_remainder.as_ref()[..len]);
34}
35
36impl std::ops::Not for MutableBitmap {
37    type Output = Self;
38
39    #[inline]
40    fn not(mut self) -> Self {
41        unary_assign(&mut self, |a: u64| !a);
42        self
43    }
44}
45
46fn binary_assign_impl<I, T, F>(lhs: &mut MutableBitmap, mut rhs: I, op: F)
47where
48    I: BitChunkIterExact<T>,
49    T: BitChunk,
50    F: Fn(T, T) -> T,
51{
52    let mut lhs_chunks = lhs.bitchunks_exact_mut::<T>();
53
54    lhs_chunks
55        .by_ref()
56        .zip(rhs.by_ref())
57        .for_each(|(lhs, rhs)| {
58            let new_chunk: T = match (lhs as &[u8]).try_into() {
59                Ok(a) => T::from_ne_bytes(a),
60                Err(_) => unreachable!(),
61            };
62            let new_chunk = op(new_chunk, rhs);
63            lhs.copy_from_slice(new_chunk.to_ne_bytes().as_ref());
64        });
65
66    let rem_lhs = lhs_chunks.remainder();
67    let rem_rhs = rhs.remainder();
68    if rem_lhs.is_empty() {
69        return;
70    }
71    let mut new_remainder = T::zero().to_ne_bytes();
72    lhs_chunks
73        .remainder()
74        .iter()
75        .enumerate()
76        .for_each(|(index, b)| new_remainder[index] = *b);
77    new_remainder = op(T::from_ne_bytes(new_remainder), rem_rhs).to_ne_bytes();
78
79    let len = lhs_chunks.remainder().len();
80    lhs_chunks
81        .remainder()
82        .copy_from_slice(&new_remainder.as_ref()[..len]);
83}
84
85/// Apply a bitwise binary operation to a [`MutableBitmap`].
86///
87/// This function can be used for operations like `&=` to a [`MutableBitmap`].
88/// # Panics
89/// This function panics iff `lhs.len() != `rhs.len()`
90pub fn binary_assign<T: BitChunk, F>(lhs: &mut MutableBitmap, rhs: &Bitmap, op: F)
91where
92    F: Fn(T, T) -> T,
93{
94    assert_eq!(lhs.len(), rhs.len());
95
96    let (slice, offset, length) = rhs.as_slice();
97    if offset == 0 {
98        let iter = BitChunksExact::<T>::new(slice, length);
99        binary_assign_impl(lhs, iter, op)
100    } else {
101        let rhs_chunks = rhs.chunks::<T>();
102        binary_assign_impl(lhs, rhs_chunks, op)
103    }
104}
105
106/// Apply a bitwise binary operation to a [`MutableBitmap`].
107///
108/// This function can be used for operations like `&=` to a [`MutableBitmap`].
109/// # Panics
110/// This function panics iff `lhs.len() != `rhs.len()`
111pub fn binary_assign_mut<T: BitChunk, F>(lhs: &mut MutableBitmap, rhs: &MutableBitmap, op: F)
112where
113    F: Fn(T, T) -> T,
114{
115    assert_eq!(lhs.len(), rhs.len());
116
117    let slice = rhs.as_slice();
118    let iter = BitChunksExact::<T>::new(slice, rhs.len());
119    binary_assign_impl(lhs, iter, op)
120}
121
122#[inline]
123/// Compute bitwise OR operation in-place
124fn or_assign<T: BitChunk>(lhs: &mut MutableBitmap, rhs: &Bitmap) {
125    if rhs.unset_bits() == 0 {
126        assert_eq!(lhs.len(), rhs.len());
127        lhs.clear();
128        lhs.extend_constant(rhs.len(), true);
129    } else if rhs.unset_bits() == rhs.len() {
130        // bitmap remains
131    } else {
132        binary_assign(lhs, rhs, |x: T, y| x | y)
133    }
134}
135
136#[inline]
137/// Compute bitwise OR operation in-place
138fn or_assign_mut<T: BitChunk>(lhs: &mut MutableBitmap, rhs: &MutableBitmap) {
139    if rhs.unset_bits() == 0 {
140        assert_eq!(lhs.len(), rhs.len());
141        lhs.clear();
142        lhs.extend_constant(rhs.len(), true);
143    } else if rhs.unset_bits() == rhs.len() {
144        // bitmap remains
145    } else {
146        binary_assign_mut(lhs, rhs, |x: T, y| x | y)
147    }
148}
149
150impl<'a> std::ops::BitOrAssign<&'a MutableBitmap> for &mut MutableBitmap {
151    #[inline]
152    fn bitor_assign(&mut self, rhs: &'a MutableBitmap) {
153        or_assign_mut::<u64>(self, rhs)
154    }
155}
156
157impl<'a> std::ops::BitOrAssign<&'a Bitmap> for &mut MutableBitmap {
158    #[inline]
159    fn bitor_assign(&mut self, rhs: &'a Bitmap) {
160        or_assign::<u64>(self, rhs)
161    }
162}
163
164impl<'a> std::ops::BitOr<&'a Bitmap> for MutableBitmap {
165    type Output = Self;
166
167    #[inline]
168    fn bitor(mut self, rhs: &'a Bitmap) -> Self {
169        or_assign::<u64>(&mut self, rhs);
170        self
171    }
172}
173
174#[inline]
175/// Compute bitwise `&` between `lhs` and `rhs`, assigning it to `lhs`
176fn and_assign<T: BitChunk>(lhs: &mut MutableBitmap, rhs: &Bitmap) {
177    if rhs.unset_bits() == 0 {
178        // bitmap remains
179    }
180    if rhs.unset_bits() == rhs.len() {
181        assert_eq!(lhs.len(), rhs.len());
182        lhs.clear();
183        lhs.extend_constant(rhs.len(), false);
184    } else {
185        binary_assign(lhs, rhs, |x: T, y| x & y)
186    }
187}
188
189impl<'a> std::ops::BitAndAssign<&'a Bitmap> for &mut MutableBitmap {
190    #[inline]
191    fn bitand_assign(&mut self, rhs: &'a Bitmap) {
192        and_assign::<u64>(self, rhs)
193    }
194}
195
196impl<'a> std::ops::BitAnd<&'a Bitmap> for MutableBitmap {
197    type Output = Self;
198
199    #[inline]
200    fn bitand(mut self, rhs: &'a Bitmap) -> Self {
201        and_assign::<u64>(&mut self, rhs);
202        self
203    }
204}
205
206#[inline]
207/// Compute bitwise XOR operation
208fn xor_assign<T: BitChunk>(lhs: &mut MutableBitmap, rhs: &Bitmap) {
209    binary_assign(lhs, rhs, |x: T, y| x ^ y)
210}
211
212impl<'a> std::ops::BitXorAssign<&'a Bitmap> for &mut MutableBitmap {
213    #[inline]
214    fn bitxor_assign(&mut self, rhs: &'a Bitmap) {
215        xor_assign::<u64>(self, rhs)
216    }
217}
218
219impl<'a> std::ops::BitXor<&'a Bitmap> for MutableBitmap {
220    type Output = Self;
221
222    #[inline]
223    fn bitxor(mut self, rhs: &'a Bitmap) -> Self {
224        xor_assign::<u64>(&mut self, rhs);
225        self
226    }
227}
polars_arrow/bitmap/assign_ops.rs

polars_arrow/bitmap/
assign_ops.rs
