polars_compute/gather/
boolean.rs

1use arrow::array::{Array, BooleanArray, PrimitiveArray};
2use arrow::bitmap::{Bitmap, BitmapBuilder};
3use polars_utils::IdxSize;
4
5use super::bitmap::{take_bitmap_nulls_unchecked, take_bitmap_unchecked};
6
7// Take implementation when neither values nor indices contain nulls.
8unsafe fn take_no_validity(values: &Bitmap, indices: &[IdxSize]) -> (Bitmap, Option<Bitmap>) {
9    (take_bitmap_unchecked(values, indices), None)
10}
11
12// Take implementation when only values contain nulls.
13unsafe fn take_values_validity(
14    values: &BooleanArray,
15    indices: &[IdxSize],
16) -> (Bitmap, Option<Bitmap>) {
17    let validity_values = values.validity().unwrap();
18    let validity = take_bitmap_unchecked(validity_values, indices);
19
20    let values_values = values.values();
21    let buffer = take_bitmap_unchecked(values_values, indices);
22
23    (buffer, validity.into())
24}
25
26// Take implementation when only indices contain nulls.
27unsafe fn take_indices_validity(
28    values: &Bitmap,
29    indices: &PrimitiveArray<IdxSize>,
30) -> (Bitmap, Option<Bitmap>) {
31    let buffer = take_bitmap_nulls_unchecked(values, indices);
32    (buffer, indices.validity().cloned())
33}
34
35// Take implementation when both values and indices contain nulls.
36unsafe fn take_values_indices_validity(
37    values: &BooleanArray,
38    indices: &PrimitiveArray<IdxSize>,
39) -> (Bitmap, Option<Bitmap>) {
40    let mut validity = BitmapBuilder::with_capacity(indices.len());
41
42    let values_validity = values.validity().unwrap();
43
44    let values_values = values.values();
45    let values = indices.iter().map(|index| match index {
46        Some(&index) => {
47            let index = index as usize;
48            debug_assert!(index < values.len());
49            validity.push(values_validity.get_bit_unchecked(index));
50            values_values.get_bit_unchecked(index)
51        },
52        None => {
53            validity.push(false);
54            false
55        },
56    });
57    let values = Bitmap::from_trusted_len_iter(values);
58    (values, validity.into_opt_validity())
59}
60
61/// `take` implementation for boolean arrays
62/// # Safety
63/// The indices must be in-bounds.
64pub unsafe fn take_unchecked(
65    values: &BooleanArray,
66    indices: &PrimitiveArray<IdxSize>,
67) -> BooleanArray {
68    let dtype = values.dtype().clone();
69    let indices_has_validity = indices.null_count() > 0;
70    let values_has_validity = values.null_count() > 0;
71
72    let (values, validity) = match (values_has_validity, indices_has_validity) {
73        (false, false) => take_no_validity(values.values(), indices.values()),
74        (true, false) => take_values_validity(values, indices.values()),
75        (false, true) => take_indices_validity(values.values(), indices),
76        (true, true) => take_values_indices_validity(values, indices),
77    };
78
79    BooleanArray::new(dtype, values, validity)
80}
polars_compute/gather/boolean.rs

polars_compute/gather/
boolean.rs
