polars_compute/comparisons/
mod.rs

1use arrow::array::Array;
2use arrow::bitmap::{self, Bitmap};
3
4pub trait TotalEqKernel: Sized + Array {
5    type Scalar: ?Sized;
6
7    // These kernels ignore validity entirely (results for nulls are unspecified
8    // but initialized).
9    fn tot_eq_kernel(&self, other: &Self) -> Bitmap;
10    fn tot_ne_kernel(&self, other: &Self) -> Bitmap;
11    fn tot_eq_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
12    fn tot_ne_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
13
14    // These kernels treat null as any other value equal to itself but unequal
15    // to anything else.
16    fn tot_eq_missing_kernel(&self, other: &Self) -> Bitmap {
17        let q = self.tot_eq_kernel(other);
18        match (self.validity(), other.validity()) {
19            (None, None) => q,
20            (None, Some(r)) => &q & r,
21            (Some(l), None) => &q & l,
22            (Some(l), Some(r)) => bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | !(l | r)),
23        }
24    }
25
26    fn tot_ne_missing_kernel(&self, other: &Self) -> Bitmap {
27        let q = self.tot_ne_kernel(other);
28        match (self.validity(), other.validity()) {
29            (None, None) => q,
30            (None, Some(r)) => &q | &!r,
31            (Some(l), None) => &q | &!l,
32            (Some(l), Some(r)) => bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | (l ^ r)),
33        }
34    }
35    fn tot_eq_missing_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
36        let q = self.tot_eq_kernel_broadcast(other);
37        if let Some(valid) = self.validity() {
38            bitmap::binary(&q, valid, |q, v| q & v)
39        } else {
40            q
41        }
42    }
43
44    fn tot_ne_missing_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
45        let q = self.tot_ne_kernel_broadcast(other);
46        if let Some(valid) = self.validity() {
47            bitmap::binary(&q, valid, |q, v| q | !v)
48        } else {
49            q
50        }
51    }
52}
53
54// Low-level comparison kernel.
55pub trait TotalOrdKernel: Sized + Array {
56    type Scalar: ?Sized;
57
58    // These kernels ignore validity entirely (results for nulls are unspecified
59    // but initialized).
60    fn tot_lt_kernel(&self, other: &Self) -> Bitmap;
61    fn tot_le_kernel(&self, other: &Self) -> Bitmap;
62    fn tot_gt_kernel(&self, other: &Self) -> Bitmap {
63        other.tot_lt_kernel(self)
64    }
65    fn tot_ge_kernel(&self, other: &Self) -> Bitmap {
66        other.tot_le_kernel(self)
67    }
68
69    // These kernels ignore validity entirely (results for nulls are unspecified
70    // but initialized).
71    fn tot_lt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
72    fn tot_le_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
73    fn tot_gt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
74    fn tot_ge_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
75}
76
77mod binary;
78mod boolean;
79mod dictionary;
80mod dyn_array;
81mod list;
82mod null;
83mod scalar;
84mod struct_;
85mod utf8;
86mod view;
87
88#[cfg(feature = "simd")]
89mod _simd_dtypes {
90    use arrow::types::{days_ms, f16, i256, months_days_ns};
91
92    use crate::NotSimdPrimitive;
93
94    impl NotSimdPrimitive for f16 {}
95    impl NotSimdPrimitive for i256 {}
96    impl NotSimdPrimitive for days_ms {}
97    impl NotSimdPrimitive for months_days_ns {}
98}
99
100#[cfg(feature = "simd")]
101mod simd;
102
103#[cfg(feature = "dtype-array")]
104mod array;