polars_core/chunked_array/ops/
bits.rs

1use super::BooleanChunked;
2
3fn first_true_idx_impl(ca: &BooleanChunked, invert: bool) -> Option<usize> {
4    let null_count = ca.null_count();
5    if null_count == ca.len() {
6        return None;
7    }
8
9    if (ca.is_sorted_ascending_flag() && invert) || (ca.is_sorted_descending_flag() && !invert) {
10        return ca.first_non_null();
11    }
12
13    let invert_mask = if invert { u64::MAX } else { 0 };
14    let mut offset = 0;
15    for arr in ca.downcast_iter() {
16        let values = arr.values();
17        if let Some(validity) = arr.validity() {
18            let mut x_it = values.fast_iter_u56();
19            let mut v_it = validity.fast_iter_u56();
20            for (x, v) in x_it.by_ref().zip(v_it.by_ref()) {
21                let n = ((x ^ invert_mask) & v).trailing_zeros() as usize;
22                if n < 56 {
23                    return Some(offset + n);
24                }
25                offset += 56;
26            }
27
28            let (x, rest_len) = x_it.remainder();
29            let (v, _rest_len) = v_it.remainder();
30            let n = ((x ^ invert_mask) & v).trailing_zeros() as usize;
31            if n < rest_len {
32                return Some(offset + n);
33            }
34            offset += rest_len;
35        } else {
36            let n = if invert {
37                values.leading_ones()
38            } else {
39                values.leading_zeros()
40            };
41            if n < values.len() {
42                return Some(offset + n);
43            }
44            offset += values.len();
45        }
46    }
47
48    None
49}
50
51impl BooleanChunked {
52    pub fn num_trues(&self) -> usize {
53        self.downcast_iter()
54            .map(|arr| match arr.validity() {
55                None => arr.values().set_bits(),
56                Some(validity) => arr.values().num_intersections_with(validity),
57            })
58            .sum()
59    }
60
61    pub fn num_falses(&self) -> usize {
62        self.downcast_iter()
63            .map(|arr| match arr.validity() {
64                None => arr.values().unset_bits(),
65                Some(validity) => (!arr.values()).num_intersections_with(validity),
66            })
67            .sum()
68    }
69
70    pub fn first_true_idx(&self) -> Option<usize> {
71        first_true_idx_impl(self, false)
72    }
73
74    pub fn first_false_idx(&self) -> Option<usize> {
75        first_true_idx_impl(self, true)
76    }
77}
polars_core/chunked_array/ops/bits.rs

polars_core/chunked_array/ops/
bits.rs
