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

polars_core/chunked_array/ops/
bits.rs
