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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

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

use super::TotalEqKernel;
use crate::comparisons::dyn_array::array_tot_eq_missing_kernel;

impl TotalEqKernel for StructArray {
    type Scalar = Box<dyn Array>;

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

        if lhs.fields() != rhs.fields() {
            return Bitmap::new_zeroed(lhs.len());
        }

        let ln = lhs.validity();
        let rn = rhs.validity();

        let lv = lhs.values();
        let rv = rhs.values();

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

        for i in 0..lhs.len() {
            let mut is_equal = true;

            if !ln.is_none_or(|v| v.get(i).unwrap()) || !rn.is_none_or(|v| v.get(i).unwrap()) {
                bitmap.push(true);
                continue;
            }

            for j in 0..lhs.values().len() {
                if lv[j].len() != rv[j].len() {
                    is_equal = false;
                    break;
                }

                let result = array_tot_eq_missing_kernel(lv[j].as_ref(), rv[j].as_ref());
                if result.unset_bits() != 0 {
                    is_equal = false;
                    break;
                }
            }

            bitmap.push(is_equal);
        }

        bitmap.freeze()
    }

    fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
        let lhs = self;
        let rhs = other;

        if lhs.fields() != rhs.fields() {
            return Bitmap::new_with_value(true, lhs.len());
        }

        if lhs.values().len() != rhs.values().len() {
            return Bitmap::new_with_value(true, lhs.len());
        }

        let ln = lhs.validity();
        let rn = rhs.validity();

        let lv = lhs.values();
        let rv = rhs.values();

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

        for i in 0..lhs.len() {
            let mut is_equal = true;

            if !ln.is_none_or(|v| v.get(i).unwrap()) || !rn.is_none_or(|v| v.get(i).unwrap()) {
                bitmap.push(false);
                continue;
            }

            for j in 0..lhs.values().len() {
                if lv[j].len() != rv[j].len() {
                    is_equal = false;
                    break;
                }

                let result = array_tot_eq_missing_kernel(lv[j].as_ref(), rv[j].as_ref());
                if result.unset_bits() != 0 {
                    is_equal = false;
                    break;
                }
            }

            bitmap.push(!is_equal);
        }

        bitmap.freeze()
    }

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

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