polars_compute/comparisons/
dictionary.rs

1use arrow::array::{Array, DictionaryArray, DictionaryKey};
2use arrow::bitmap::{Bitmap, BitmapBuilder};
3
4use super::TotalEqKernel;
5use crate::comparisons::dyn_array::{array_tot_eq_missing_kernel, array_tot_ne_missing_kernel};
6
7impl<K: DictionaryKey> TotalEqKernel for DictionaryArray<K> {
8    type Scalar = Box<dyn Array>;
9
10    fn tot_eq_kernel(&self, other: &Self) -> Bitmap {
11        assert_eq!(self.len(), other.len());
12
13        let mut bitmap = BitmapBuilder::with_capacity(self.len());
14
15        for i in 0..self.len() {
16            let lval = self.validity().is_none_or(|v| v.get(i).unwrap());
17            let rval = other.validity().is_none_or(|v| v.get(i).unwrap());
18
19            if !lval || !rval {
20                bitmap.push(true);
21                continue;
22            }
23
24            let lkey = self.key_value(i);
25            let rkey = other.key_value(i);
26
27            let mut lhs_value = self.values().clone();
28            lhs_value.slice(lkey, 1);
29            let mut rhs_value = other.values().clone();
30            rhs_value.slice(rkey, 1);
31
32            let result = array_tot_eq_missing_kernel(lhs_value.as_ref(), rhs_value.as_ref());
33            bitmap.push(result.unset_bits() == 0);
34        }
35
36        bitmap.freeze()
37    }
38
39    fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
40        assert_eq!(self.len(), other.len());
41
42        let mut bitmap = BitmapBuilder::with_capacity(self.len());
43
44        for i in 0..self.len() {
45            let lval = self.validity().is_none_or(|v| v.get(i).unwrap());
46            let rval = other.validity().is_none_or(|v| v.get(i).unwrap());
47
48            if !lval || !rval {
49                bitmap.push(false);
50                continue;
51            }
52
53            let lkey = self.key_value(i);
54            let rkey = other.key_value(i);
55
56            let mut lhs_value = self.values().clone();
57            lhs_value.slice(lkey, 1);
58            let mut rhs_value = other.values().clone();
59            rhs_value.slice(rkey, 1);
60
61            let result = array_tot_ne_missing_kernel(lhs_value.as_ref(), rhs_value.as_ref());
62            bitmap.push(result.set_bits() > 0);
63        }
64
65        bitmap.freeze()
66    }
67
68    fn tot_eq_kernel_broadcast(&self, _other: &Self::Scalar) -> arrow::bitmap::Bitmap {
69        todo!()
70    }
71
72    fn tot_ne_kernel_broadcast(&self, _other: &Self::Scalar) -> arrow::bitmap::Bitmap {
73        todo!()
74    }
75}
polars_compute/comparisons/dictionary.rs

polars_compute/comparisons/
dictionary.rs
