polars_compute/gather/
bitmap.rs

1use arrow::array::Array;
2use arrow::bitmap::Bitmap;
3use arrow::datatypes::IdxArr;
4use polars_utils::IdxSize;
5
6/// # Safety
7/// Doesn't do any bound checks.
8pub unsafe fn take_bitmap_unchecked(values: &Bitmap, indices: &[IdxSize]) -> Bitmap {
9    let values = indices.iter().map(|&index| {
10        debug_assert!((index as usize) < values.len());
11        values.get_bit_unchecked(index as usize)
12    });
13    Bitmap::from_trusted_len_iter(values)
14}
15
16/// # Safety
17/// Doesn't check bounds for non-null elements.
18pub unsafe fn take_bitmap_nulls_unchecked(values: &Bitmap, indices: &IdxArr) -> Bitmap {
19    // Fast-path: no need to bother with null indices.
20    if indices.null_count() == 0 {
21        return take_bitmap_unchecked(values, indices.values());
22    }
23
24    if values.is_empty() {
25        // Nothing can be in-bounds, assume indices is full-null.
26        debug_assert!(indices.null_count() == indices.len());
27        return Bitmap::new_zeroed(indices.len());
28    }
29
30    let values = indices.iter().map(|opt_index| {
31        // We checked that values.len() > 0 so we can use index 0 for nulls.
32        let index = opt_index.copied().unwrap_or(0) as usize;
33        debug_assert!(index < values.len());
34        values.get_bit_unchecked(index)
35    });
36    Bitmap::from_trusted_len_iter(values)
37}
polars_compute/gather/bitmap.rs

polars_compute/gather/
bitmap.rs
