polars_compute/unique/
boolean.rs

1use arrow::array::{Array, BooleanArray};
2use arrow::bitmap::BitmapBuilder;
3use arrow::datatypes::ArrowDataType;
4
5use super::{GenericUniqueKernel, RangedUniqueKernel};
6
7#[derive(Default)]
8pub struct BooleanUniqueKernelState {
9    seen: u32,
10}
11
12impl BooleanUniqueKernelState {
13    pub fn new() -> Self {
14        Self::default()
15    }
16}
17
18impl RangedUniqueKernel for BooleanUniqueKernelState {
19    type Array = BooleanArray;
20
21    fn has_seen_all(&self) -> bool {
22        self.seen == 0b111
23    }
24
25    fn append(&mut self, array: &Self::Array) {
26        if array.len() == 0 {
27            return;
28        }
29
30        let null_count = array.null_count();
31        self.seen |= u32::from(null_count > 0) << 2;
32        let set_bits = if null_count > 0 {
33            array
34                .values()
35                .num_intersections_with(array.validity().unwrap())
36        } else {
37            array.values().set_bits()
38        };
39
40        self.seen |= u32::from(set_bits != array.len() - null_count);
41        self.seen |= u32::from(set_bits != 0) << 1;
42    }
43
44    fn append_state(&mut self, other: &Self) {
45        self.seen |= other.seen;
46    }
47
48    fn finalize_unique(self) -> Self::Array {
49        let mut values = BitmapBuilder::with_capacity(self.seen.count_ones() as usize);
50
51        if self.seen & 0b001 != 0 {
52            values.push(false);
53        }
54        if self.seen & 0b010 != 0 {
55            values.push(true);
56        }
57        let validity = if self.seen & 0b100 != 0 {
58            let mut validity = BitmapBuilder::with_capacity(values.len() + 1);
59            validity.extend_constant(values.len(), true);
60            validity.push(false);
61            values.push(false);
62            Some(validity.freeze())
63        } else {
64            None
65        };
66
67        let values = values.freeze();
68        BooleanArray::new(ArrowDataType::Boolean, values, validity)
69    }
70
71    fn finalize_n_unique(&self) -> usize {
72        self.seen.count_ones() as usize
73    }
74
75    fn finalize_n_unique_non_null(&self) -> usize {
76        (self.seen & 0b011).count_ones() as usize
77    }
78}
79
80impl GenericUniqueKernel for BooleanArray {
81    fn unique(&self) -> Self {
82        let mut state = BooleanUniqueKernelState::new();
83        state.append(self);
84        state.finalize_unique()
85    }
86
87    fn n_unique(&self) -> usize {
88        let mut state = BooleanUniqueKernelState::new();
89        state.append(self);
90        state.finalize_n_unique()
91    }
92
93    fn n_unique_non_null(&self) -> usize {
94        let mut state = BooleanUniqueKernelState::new();
95        state.append(self);
96        state.finalize_n_unique_non_null()
97    }
98}
99
100#[test]
101fn test_boolean_distinct_count() {
102    use arrow::bitmap::Bitmap;
103    use arrow::datatypes::ArrowDataType;
104
105    macro_rules! assert_bool_dc {
106        ($values:expr, $validity:expr => $dc:expr) => {
107            let validity: Option<Bitmap> =
108                <Option<Vec<bool>>>::map($validity, |v| Bitmap::from_iter(v));
109            let arr =
110                BooleanArray::new(ArrowDataType::Boolean, Bitmap::from_iter($values), validity);
111            assert_eq!(arr.n_unique(), $dc);
112        };
113    }
114
115    assert_bool_dc!(vec![], None => 0);
116    assert_bool_dc!(vec![], Some(vec![]) => 0);
117    assert_bool_dc!(vec![true], None => 1);
118    assert_bool_dc!(vec![true], Some(vec![true]) => 1);
119    assert_bool_dc!(vec![true], Some(vec![false]) => 1);
120    assert_bool_dc!(vec![true, false], None => 2);
121    assert_bool_dc!(vec![true, false, false], None => 2);
122    assert_bool_dc!(vec![true, false, false], Some(vec![true, true, false]) => 3);
123
124    // Copied from https://github.com/pola-rs/polars/pull/16765#discussion_r1629426159
125    assert_bool_dc!(vec![true, true, true, true, true], Some(vec![true, false, true, false, false]) => 2);
126    assert_bool_dc!(vec![false, true, false, true, true], Some(vec![true, false, true, false, false]) => 2);
127    assert_bool_dc!(vec![true, false, true, false, true, true], Some(vec![true, true, false, true, false, false]) => 3);
128}
polars_compute/unique/boolean.rs

polars_compute/unique/
boolean.rs
