polars_compute/comparisons/
scalar.rs

1use arrow::array::PrimitiveArray;
2use arrow::bitmap::Bitmap;
3use polars_utils::total_ord::TotalOrd;
4
5use super::{TotalEqKernel, TotalOrdKernel};
6use crate::NotSimdPrimitive;
7
8impl<T: NotSimdPrimitive + TotalOrd> TotalEqKernel for PrimitiveArray<T> {
9    type Scalar = T;
10
11    fn tot_eq_kernel(&self, other: &Self) -> Bitmap {
12        assert!(self.len() == other.len());
13        self.values()
14            .iter()
15            .zip(other.values().iter())
16            .map(|(l, r)| l.tot_eq(r))
17            .collect()
18    }
19
20    fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
21        assert!(self.len() == other.len());
22        self.values()
23            .iter()
24            .zip(other.values().iter())
25            .map(|(l, r)| l.tot_ne(r))
26            .collect()
27    }
28
29    fn tot_eq_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
30        self.values().iter().map(|l| l.tot_eq(other)).collect()
31    }
32
33    fn tot_ne_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
34        self.values().iter().map(|l| l.tot_ne(other)).collect()
35    }
36}
37
38impl<T: NotSimdPrimitive + TotalOrd> TotalOrdKernel for PrimitiveArray<T> {
39    type Scalar = T;
40
41    fn tot_lt_kernel(&self, other: &Self) -> Bitmap {
42        assert!(self.len() == other.len());
43        self.values()
44            .iter()
45            .zip(other.values().iter())
46            .map(|(l, r)| l.tot_lt(r))
47            .collect()
48    }
49
50    fn tot_le_kernel(&self, other: &Self) -> Bitmap {
51        assert!(self.len() == other.len());
52        self.values()
53            .iter()
54            .zip(other.values().iter())
55            .map(|(l, r)| l.tot_le(r))
56            .collect()
57    }
58
59    fn tot_lt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
60        self.values().iter().map(|l| l.tot_lt(other)).collect()
61    }
62
63    fn tot_le_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
64        self.values().iter().map(|l| l.tot_le(other)).collect()
65    }
66
67    fn tot_gt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
68        self.values().iter().map(|l| l.tot_gt(other)).collect()
69    }
70
71    fn tot_ge_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
72        self.values().iter().map(|l| l.tot_ge(other)).collect()
73    }
74}