polars-compute 0.53.0

Private compute kernels for the Polars DataFrame library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

use arrow::array::{Array, DictionaryArray, DictionaryKey};
use arrow::bitmap::{Bitmap, BitmapBuilder};

use super::TotalEqKernel;
use crate::comparisons::dyn_array::{array_tot_eq_missing_kernel, array_tot_ne_missing_kernel};

impl<K: DictionaryKey> TotalEqKernel for DictionaryArray<K> {
    type Scalar = Box<dyn Array>;

    fn tot_eq_kernel(&self, other: &Self) -> Bitmap {
        assert_eq!(self.len(), other.len());

        let mut bitmap = BitmapBuilder::with_capacity(self.len());

        for i in 0..self.len() {
            let lval = self.validity().is_none_or(|v| v.get(i).unwrap());
            let rval = other.validity().is_none_or(|v| v.get(i).unwrap());

            if !lval || !rval {
                bitmap.push(true);
                continue;
            }

            let lkey = self.key_value(i);
            let rkey = other.key_value(i);

            let mut lhs_value = self.values().clone();
            lhs_value.slice(lkey, 1);
            let mut rhs_value = other.values().clone();
            rhs_value.slice(rkey, 1);

            let result = array_tot_eq_missing_kernel(lhs_value.as_ref(), rhs_value.as_ref());
            bitmap.push(result.unset_bits() == 0);
        }

        bitmap.freeze()
    }

    fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
        assert_eq!(self.len(), other.len());

        let mut bitmap = BitmapBuilder::with_capacity(self.len());

        for i in 0..self.len() {
            let lval = self.validity().is_none_or(|v| v.get(i).unwrap());
            let rval = other.validity().is_none_or(|v| v.get(i).unwrap());

            if !lval || !rval {
                bitmap.push(false);
                continue;
            }

            let lkey = self.key_value(i);
            let rkey = other.key_value(i);

            let mut lhs_value = self.values().clone();
            lhs_value.slice(lkey, 1);
            let mut rhs_value = other.values().clone();
            rhs_value.slice(rkey, 1);

            let result = array_tot_ne_missing_kernel(lhs_value.as_ref(), rhs_value.as_ref());
            bitmap.push(result.set_bits() > 0);
        }

        bitmap.freeze()
    }

    fn tot_eq_kernel_broadcast(&self, _other: &Self::Scalar) -> arrow::bitmap::Bitmap {
        todo!()
    }

    fn tot_ne_kernel_broadcast(&self, _other: &Self::Scalar) -> arrow::bitmap::Bitmap {
        todo!()
    }
}