polars_compute/if_then_else/
boolean.rs

1use arrow::array::BooleanArray;
2use arrow::bitmap::{self, Bitmap};
3use arrow::datatypes::ArrowDataType;
4
5use super::{if_then_else_validity, IfThenElseKernel};
6
7impl IfThenElseKernel for BooleanArray {
8    type Scalar<'a> = bool;
9
10    fn if_then_else(mask: &Bitmap, if_true: &Self, if_false: &Self) -> Self {
11        let values = bitmap::ternary(mask, if_true.values(), if_false.values(), |m, t, f| {
12            (m & t) | (!m & f)
13        });
14        let validity = if_then_else_validity(mask, if_true.validity(), if_false.validity());
15        BooleanArray::from(values).with_validity(validity)
16    }
17
18    fn if_then_else_broadcast_true(
19        mask: &Bitmap,
20        if_true: Self::Scalar<'_>,
21        if_false: &Self,
22    ) -> Self {
23        let values = if if_true {
24            bitmap::or(if_false.values(), mask) // (m & true)  | (!m & f)  ->  f | m
25        } else {
26            bitmap::and_not(if_false.values(), mask) // (m & false) | (!m & f)  ->  f & !m
27        };
28        let validity = if_then_else_validity(mask, None, if_false.validity());
29        BooleanArray::from(values).with_validity(validity)
30    }
31
32    fn if_then_else_broadcast_false(
33        mask: &Bitmap,
34        if_true: &Self,
35        if_false: Self::Scalar<'_>,
36    ) -> Self {
37        let values = if if_false {
38            bitmap::or_not(if_true.values(), mask) // (m & t) | (!m & true)   ->  t | !m
39        } else {
40            bitmap::and(if_true.values(), mask) // (m & t) | (!m & false)  ->  t & m
41        };
42        let validity = if_then_else_validity(mask, if_true.validity(), None);
43        BooleanArray::from(values).with_validity(validity)
44    }
45
46    fn if_then_else_broadcast_both(
47        _dtype: ArrowDataType,
48        mask: &Bitmap,
49        if_true: Self::Scalar<'_>,
50        if_false: Self::Scalar<'_>,
51    ) -> Self {
52        let values = match (if_true, if_false) {
53            (false, false) => Bitmap::new_with_value(false, mask.len()),
54            (false, true) => !mask,
55            (true, false) => mask.clone(),
56            (true, true) => Bitmap::new_with_value(true, mask.len()),
57        };
58        BooleanArray::from(values)
59    }
60}
polars_compute/if_then_else/boolean.rs

polars_compute/if_then_else/
boolean.rs
