1#![allow(unsafe_op_in_unsafe_fn)]
2use std::ops::Deref;
3
4use either::Either;
5use polars_buffer::{Buffer, SharedStorage};
6use polars_error::{PolarsResult, polars_bail};
7use polars_utils::relaxed_cell::RelaxedCell;
8
9use super::utils::{self, BitChunk, BitChunks, BitmapIter, count_zeros, fmt, get_bit_unchecked};
10use super::{IntoIter, MutableBitmap, chunk_iter_to_vec, num_intersections_with};
11use crate::array::Splitable;
12use crate::bitmap::BitmapBuilder;
13use crate::bitmap::aligned::AlignedBitmapSlice;
14use crate::bitmap::iterator::{
15 FastU32BitmapIter, FastU56BitmapIter, FastU64BitmapIter, TrueIdxIter,
16};
17use crate::bitmap::utils::bytes_for;
18use crate::legacy::utils::FromTrustedLenIterator;
19use crate::trusted_len::TrustedLen;
20
21const UNKNOWN_BIT_COUNT: u64 = u64::MAX;
22
23#[derive(Default, Clone)]
56pub struct Bitmap {
57 storage: SharedStorage<u8>,
58 offset: usize,
61 length: usize,
62
63 unset_bit_count_cache: RelaxedCell<u64>,
68}
69
70#[inline(always)]
71fn has_cached_unset_bit_count(ubcc: u64) -> bool {
72 ubcc >> 63 == 0
73}
74
75impl std::fmt::Debug for Bitmap {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 let (bytes, offset, len) = self.as_slice();
78 fmt(bytes, offset, len, f)
79 }
80}
81
82pub(super) fn check(bytes: &[u8], offset: usize, length: usize) -> PolarsResult<()> {
83 if offset + length > bytes.len().saturating_mul(8) {
84 polars_bail!(InvalidOperation:
85 "The offset + length of the bitmap ({}) must be `<=` to the number of bytes times 8 ({})",
86 offset + length,
87 bytes.len().saturating_mul(8)
88 );
89 }
90 Ok(())
91}
92
93impl Bitmap {
94 #[inline]
96 pub fn new() -> Self {
97 Self::default()
98 }
99
100 #[inline]
104 pub fn try_new(bytes: Vec<u8>, length: usize) -> PolarsResult<Self> {
105 check(&bytes, 0, length)?;
106 Ok(Self {
107 storage: SharedStorage::from_vec(bytes),
108 length,
109 offset: 0,
110 unset_bit_count_cache: RelaxedCell::from(if length == 0 {
111 0
112 } else {
113 UNKNOWN_BIT_COUNT
114 }),
115 })
116 }
117
118 #[inline]
120 pub fn len(&self) -> usize {
121 self.length
122 }
123
124 #[inline]
126 pub fn is_empty(&self) -> bool {
127 self.len() == 0
128 }
129
130 pub fn iter(&self) -> BitmapIter<'_> {
132 BitmapIter::new(&self.storage, self.offset, self.length)
133 }
134
135 pub fn chunks<T: BitChunk>(&self) -> BitChunks<'_, T> {
139 BitChunks::new(&self.storage, self.offset, self.length)
140 }
141
142 pub fn fast_iter_u32(&self) -> FastU32BitmapIter<'_> {
145 FastU32BitmapIter::new(&self.storage, self.offset, self.length)
146 }
147
148 pub fn fast_iter_u56(&self) -> FastU56BitmapIter<'_> {
151 FastU56BitmapIter::new(&self.storage, self.offset, self.length)
152 }
153
154 pub fn fast_iter_u64(&self) -> FastU64BitmapIter<'_> {
157 FastU64BitmapIter::new(&self.storage, self.offset, self.length)
158 }
159
160 pub fn true_idx_iter(&self) -> TrueIdxIter<'_> {
162 TrueIdxIter::new(self.len(), Some(self))
163 }
164
165 pub fn aligned<T: BitChunk>(&self) -> AlignedBitmapSlice<'_, T> {
167 AlignedBitmapSlice::new(&self.storage, self.offset, self.length)
168 }
169
170 pub fn to_aligned_bitmap(&self) -> Bitmap {
172 if self.offset.is_multiple_of(8) {
173 self.clone()
174 } else {
175 Bitmap::from_trusted_len_iter(self.iter())
176 }
177 }
178
179 #[inline]
187 pub fn as_slice(&self) -> (&[u8], usize, usize) {
188 let start = self.offset / 8;
189 let len = (self.offset % 8 + self.length).saturating_add(7) / 8;
190 (
191 &self.storage[start..start + len],
192 self.offset % 8,
193 self.length,
194 )
195 }
196
197 pub fn as_buffer(&self) -> (Buffer<u8>, usize, usize) {
205 let start = self.offset / 8;
206 let len = (self.offset % 8 + self.length).saturating_add(7) / 8;
207 (
208 Buffer::from_storage(self.storage.clone()).sliced(start..start + len),
209 self.offset % 8,
210 self.length,
211 )
212 }
213
214 #[inline]
218 pub fn set_bits(&self) -> usize {
219 self.length - self.unset_bits()
220 }
221
222 #[inline]
226 pub fn lazy_set_bits(&self) -> Option<usize> {
227 Some(self.length - self.lazy_unset_bits()?)
228 }
229
230 pub fn unset_bits(&self) -> usize {
239 self.lazy_unset_bits().unwrap_or_else(|| {
240 let zeros = count_zeros(&self.storage, self.offset, self.length);
241 self.unset_bit_count_cache.store(zeros as u64);
242 zeros
243 })
244 }
245
246 pub fn lazy_unset_bits(&self) -> Option<usize> {
250 let cache = self.unset_bit_count_cache.load();
251 has_cached_unset_bit_count(cache).then_some(cache as usize)
252 }
253
254 pub unsafe fn update_bit_count(&mut self, bits_set: usize) {
260 assert!(bits_set <= self.length);
261 let zeros = self.length - bits_set;
262 self.unset_bit_count_cache.store(zeros as u64);
263 }
264
265 #[inline]
270 pub fn slice(&mut self, offset: usize, length: usize) {
271 assert!(offset + length <= self.length);
272 unsafe { self.slice_unchecked(offset, length) }
273 }
274
275 #[inline]
280 pub unsafe fn slice_unchecked(&mut self, offset: usize, length: usize) {
281 if offset == 0 && length == self.length {
283 return;
284 }
285
286 let unset_bit_count_cache = self.unset_bit_count_cache.get_mut();
288 if *unset_bit_count_cache == 0 || *unset_bit_count_cache == self.length as u64 {
289 let new_count = if *unset_bit_count_cache > 0 {
290 length as u64
291 } else {
292 0
293 };
294 *unset_bit_count_cache = new_count;
295 self.offset += offset;
296 self.length = length;
297 return;
298 }
299
300 if has_cached_unset_bit_count(*unset_bit_count_cache) {
301 let small_portion = (self.length / 5).max(32);
305 if length + small_portion >= self.length {
306 let slice_end = self.offset + offset + length;
308 let head_count = count_zeros(&self.storage, self.offset, offset);
309 let tail_count =
310 count_zeros(&self.storage, slice_end, self.length - length - offset);
311 let new_count = *unset_bit_count_cache - head_count as u64 - tail_count as u64;
312 *unset_bit_count_cache = new_count;
313 } else {
314 *unset_bit_count_cache = UNKNOWN_BIT_COUNT;
315 }
316 }
317
318 self.offset += offset;
319 self.length = length;
320 }
321
322 #[inline]
327 #[must_use]
328 pub fn sliced(self, offset: usize, length: usize) -> Self {
329 assert!(offset + length <= self.length);
330 unsafe { self.sliced_unchecked(offset, length) }
331 }
332
333 #[inline]
338 #[must_use]
339 pub unsafe fn sliced_unchecked(mut self, offset: usize, length: usize) -> Self {
340 self.slice_unchecked(offset, length);
341 self
342 }
343
344 #[inline]
348 pub fn get_bit(&self, i: usize) -> bool {
349 assert!(i < self.len());
350 unsafe { self.get_bit_unchecked(i) }
351 }
352
353 #[inline]
358 pub unsafe fn get_bit_unchecked(&self, i: usize) -> bool {
359 debug_assert!(i < self.len());
360 get_bit_unchecked(&self.storage, self.offset + i)
361 }
362
363 pub(crate) fn as_ptr(&self) -> *const u8 {
366 self.storage.deref().as_ptr()
367 }
368
369 pub fn as_aligned_ptr(&self) -> Option<*const u8> {
372 self.offset
373 .is_multiple_of(8)
374 .then(|| unsafe { self.as_ptr().add(self.offset / 8) })
375 }
376
377 pub(crate) fn offset(&self) -> usize {
380 self.offset
381 }
382
383 pub fn into_mut(mut self) -> Either<Self, MutableBitmap> {
391 match self.storage.try_into_vec() {
392 Ok(v) => Either::Right(MutableBitmap::from_vec(v, self.length)),
393 Err(storage) => {
394 self.storage = storage;
395 Either::Left(self)
396 },
397 }
398 }
399
400 pub fn make_mut(self) -> MutableBitmap {
403 match self.into_mut() {
404 Either::Left(data) => {
405 if data.offset > 0 {
406 let chunks = data.chunks::<u64>();
408 let remainder = chunks.remainder();
409 let vec = chunk_iter_to_vec(chunks.chain(std::iter::once(remainder)));
410 MutableBitmap::from_vec(vec, data.length)
411 } else {
412 let len = bytes_for(data.length);
413 MutableBitmap::from_vec(data.storage[0..len].to_vec(), data.length)
414 }
415 },
416 Either::Right(data) => data,
417 }
418 }
419
420 #[inline]
422 pub fn new_zeroed(length: usize) -> Self {
423 let bytes_needed = length.div_ceil(8);
424 let storage = Buffer::zeroed(bytes_needed).into_storage();
425 Self {
426 storage,
427 offset: 0,
428 length,
429 unset_bit_count_cache: RelaxedCell::from(length as u64),
430 }
431 }
432
433 #[inline]
435 pub fn new_with_value(value: bool, length: usize) -> Self {
436 if !value {
437 return Self::new_zeroed(length);
438 }
439
440 unsafe {
441 Bitmap::from_inner_unchecked(
442 SharedStorage::from_vec(vec![u8::MAX; length.saturating_add(7) / 8]),
443 0,
444 length,
445 Some(0),
446 )
447 }
448 }
449
450 #[inline]
452 pub fn null_count_range(&self, offset: usize, length: usize) -> usize {
453 count_zeros(&self.storage, self.offset + offset, length)
454 }
455
456 #[inline]
460 pub fn from_u8_slice<T: AsRef<[u8]>>(slice: T, length: usize) -> Self {
461 Bitmap::try_new(slice.as_ref().to_vec(), length).unwrap()
462 }
463
464 #[inline]
469 pub fn from_u8_vec(vec: Vec<u8>, length: usize) -> Self {
470 Bitmap::try_new(vec, length).unwrap()
471 }
472
473 #[inline]
475 pub fn get(&self, i: usize) -> Option<bool> {
476 if i < self.len() {
477 Some(unsafe { self.get_bit_unchecked(i) })
478 } else {
479 None
480 }
481 }
482
483 pub unsafe fn from_inner_unchecked(
489 storage: SharedStorage<u8>,
490 offset: usize,
491 length: usize,
492 unset_bits: Option<usize>,
493 ) -> Self {
494 debug_assert!(check(&storage[..], offset, length).is_ok());
495
496 let unset_bit_count_cache = if let Some(n) = unset_bits {
497 RelaxedCell::from(n as u64)
498 } else {
499 RelaxedCell::from(UNKNOWN_BIT_COUNT)
500 };
501 Self {
502 storage,
503 offset,
504 length,
505 unset_bit_count_cache,
506 }
507 }
508
509 pub fn intersects_with(&self, other: &Self) -> bool {
513 self.num_intersections_with(other) != 0
514 }
515
516 pub fn num_intersections_with(&self, other: &Self) -> usize {
518 num_intersections_with(
519 super::bitmask::BitMask::from_bitmap(self),
520 super::bitmask::BitMask::from_bitmap(other),
521 )
522 }
523
524 pub fn select(&self, truthy: &Self, falsy: &Self) -> Self {
530 super::bitmap_ops::select(self, truthy, falsy)
531 }
532
533 pub fn select_constant(&self, truthy: &Self, falsy: bool) -> Self {
539 super::bitmap_ops::select_constant(self, truthy, falsy)
540 }
541
542 pub fn num_edges(&self) -> usize {
544 super::bitmap_ops::num_edges(self)
545 }
546
547 pub fn leading_zeros(&self) -> usize {
549 utils::leading_zeros(&self.storage, self.offset, self.length)
550 }
551 pub fn leading_ones(&self) -> usize {
553 utils::leading_ones(&self.storage, self.offset, self.length)
554 }
555 pub fn trailing_zeros(&self) -> usize {
557 utils::trailing_zeros(&self.storage, self.offset, self.length)
558 }
559 pub fn trailing_ones(&self) -> usize {
561 utils::trailing_ones(&self.storage, self.offset, self.length)
562 }
563
564 pub fn take_leading_zeros(&mut self) -> usize {
567 if self
568 .lazy_unset_bits()
569 .is_some_and(|unset_bits| unset_bits == self.length)
570 {
571 let leading_zeros = self.length;
572 self.offset += self.length;
573 self.length = 0;
574 *self.unset_bit_count_cache.get_mut() = 0;
575 return leading_zeros;
576 }
577
578 let leading_zeros = self.leading_zeros();
579 self.offset += leading_zeros;
580 self.length -= leading_zeros;
581 if has_cached_unset_bit_count(*self.unset_bit_count_cache.get_mut()) {
582 *self.unset_bit_count_cache.get_mut() -= leading_zeros as u64;
583 }
584 leading_zeros
585 }
586 pub fn take_leading_ones(&mut self) -> usize {
589 if self
590 .lazy_unset_bits()
591 .is_some_and(|unset_bits| unset_bits == 0)
592 {
593 let leading_ones = self.length;
594 self.offset += self.length;
595 self.length = 0;
596 *self.unset_bit_count_cache.get_mut() = 0;
597 return leading_ones;
598 }
599
600 let leading_ones = self.leading_ones();
601 self.offset += leading_ones;
602 self.length -= leading_ones;
603 leading_ones
605 }
606 pub fn take_trailing_zeros(&mut self) -> usize {
609 if self
610 .lazy_unset_bits()
611 .is_some_and(|unset_bits| unset_bits == self.length)
612 {
613 let trailing_zeros = self.length;
614 self.length = 0;
615 *self.unset_bit_count_cache.get_mut() = 0;
616 return trailing_zeros;
617 }
618
619 let trailing_zeros = self.trailing_zeros();
620 self.length -= trailing_zeros;
621 if has_cached_unset_bit_count(*self.unset_bit_count_cache.get_mut()) {
622 *self.unset_bit_count_cache.get_mut() -= trailing_zeros as u64;
623 }
624 trailing_zeros
625 }
626 pub fn take_trailing_ones(&mut self) -> usize {
629 if self
630 .lazy_unset_bits()
631 .is_some_and(|unset_bits| unset_bits == 0)
632 {
633 let trailing_ones = self.length;
634 self.length = 0;
635 *self.unset_bit_count_cache.get_mut() = 0;
636 return trailing_ones;
637 }
638
639 let trailing_ones = self.trailing_ones();
640 self.length -= trailing_ones;
641 trailing_ones
643 }
644}
645
646impl<P: AsRef<[bool]>> From<P> for Bitmap {
647 fn from(slice: P) -> Self {
648 Self::from_trusted_len_iter(slice.as_ref().iter().copied())
649 }
650}
651
652impl FromIterator<bool> for Bitmap {
653 fn from_iter<I>(iter: I) -> Self
654 where
655 I: IntoIterator<Item = bool>,
656 {
657 MutableBitmap::from_iter(iter).into()
658 }
659}
660
661impl FromTrustedLenIterator<bool> for Bitmap {
662 fn from_iter_trusted_length<T: IntoIterator<Item = bool>>(iter: T) -> Self
663 where
664 T::IntoIter: TrustedLen,
665 {
666 MutableBitmap::from_trusted_len_iter(iter.into_iter()).into()
667 }
668}
669
670impl Bitmap {
671 pub fn opt_from_iter<I: Iterator<Item = bool>>(mut iterator: I) -> Option<Self> {
673 let mut num_true = 0;
674 loop {
675 match iterator.next() {
676 Some(true) => num_true += 1,
677 Some(false) => break,
678 None => return None, }
680 }
681
682 let mut bm = BitmapBuilder::with_capacity(num_true + 1 + iterator.size_hint().0);
683 bm.extend_constant(num_true, true);
684 bm.push(false);
685 for x in iterator {
686 bm.push(x);
687 }
688 bm.into_opt_validity()
689 }
690
691 #[inline]
696 pub unsafe fn from_trusted_len_iter_unchecked<I: Iterator<Item = bool>>(iterator: I) -> Self {
697 MutableBitmap::from_trusted_len_iter_unchecked(iterator).into()
698 }
699
700 #[inline]
702 pub fn from_trusted_len_iter<I: TrustedLen<Item = bool>>(iterator: I) -> Self {
703 MutableBitmap::from_trusted_len_iter(iterator).into()
704 }
705
706 #[inline]
708 pub fn try_from_trusted_len_iter<E, I: TrustedLen<Item = std::result::Result<bool, E>>>(
709 iterator: I,
710 ) -> std::result::Result<Self, E> {
711 Ok(MutableBitmap::try_from_trusted_len_iter(iterator)?.into())
712 }
713
714 #[inline]
719 pub unsafe fn try_from_trusted_len_iter_unchecked<
720 E,
721 I: Iterator<Item = std::result::Result<bool, E>>,
722 >(
723 iterator: I,
724 ) -> std::result::Result<Self, E> {
725 Ok(MutableBitmap::try_from_trusted_len_iter_unchecked(iterator)?.into())
726 }
727}
728
729impl<'a> IntoIterator for &'a Bitmap {
730 type Item = bool;
731 type IntoIter = BitmapIter<'a>;
732
733 fn into_iter(self) -> Self::IntoIter {
734 BitmapIter::<'a>::new(&self.storage, self.offset, self.length)
735 }
736}
737
738impl IntoIterator for Bitmap {
739 type Item = bool;
740 type IntoIter = IntoIter;
741
742 fn into_iter(self) -> Self::IntoIter {
743 IntoIter::new(self)
744 }
745}
746
747impl Splitable for Bitmap {
748 #[inline(always)]
749 fn check_bound(&self, offset: usize) -> bool {
750 offset <= self.len()
751 }
752
753 unsafe fn _split_at_unchecked(&self, offset: usize) -> (Self, Self) {
754 if offset == 0 {
755 return (Self::new(), self.clone());
756 }
757 if offset == self.len() {
758 return (self.clone(), Self::new());
759 }
760
761 let ubcc = self.unset_bit_count_cache.load();
762
763 let lhs_length = offset;
764 let rhs_length = self.length - offset;
765
766 let mut lhs_ubcc = UNKNOWN_BIT_COUNT;
767 let mut rhs_ubcc = UNKNOWN_BIT_COUNT;
768
769 if has_cached_unset_bit_count(ubcc) {
770 if ubcc == 0 {
771 lhs_ubcc = 0;
772 rhs_ubcc = 0;
773 } else if ubcc == self.length as u64 {
774 lhs_ubcc = offset as u64;
775 rhs_ubcc = (self.length - offset) as u64;
776 } else {
777 let small_portion = (self.length / 4).max(32);
781
782 if lhs_length <= rhs_length {
783 if rhs_length + small_portion >= self.length {
784 let count = count_zeros(&self.storage, self.offset, lhs_length) as u64;
785 lhs_ubcc = count;
786 rhs_ubcc = ubcc - count;
787 }
788 } else if lhs_length + small_portion >= self.length {
789 let count = count_zeros(&self.storage, self.offset + offset, rhs_length) as u64;
790 lhs_ubcc = ubcc - count;
791 rhs_ubcc = count;
792 }
793 }
794 }
795
796 debug_assert!(lhs_ubcc == UNKNOWN_BIT_COUNT || lhs_ubcc <= ubcc);
797 debug_assert!(rhs_ubcc == UNKNOWN_BIT_COUNT || rhs_ubcc <= ubcc);
798
799 (
800 Self {
801 storage: self.storage.clone(),
802 offset: self.offset,
803 length: lhs_length,
804 unset_bit_count_cache: RelaxedCell::from(lhs_ubcc),
805 },
806 Self {
807 storage: self.storage.clone(),
808 offset: self.offset + offset,
809 length: rhs_length,
810 unset_bit_count_cache: RelaxedCell::from(rhs_ubcc),
811 },
812 )
813 }
814}