polars_core/chunked_array/ops/
compare_inner.rs

Help
1#![allow(unsafe_op_in_unsafe_fn)]
2//! Used to speed up TotalEq and TotalOrd of elements within an array.
3
4use std::cmp::Ordering;
5
6use crate::chunked_array::ChunkedArrayLayout;
7use crate::prelude::*;
8use crate::series::implementations::null::NullChunked;
9
10#[repr(transparent)]
11#[derive(Copy, Clone)]
12pub struct NonNull<T>(pub T);
13
14impl<T: TotalEq> TotalEq for NonNull<T> {
15    fn tot_eq(&self, other: &Self) -> bool {
16        self.0.tot_eq(&other.0)
17    }
18}
19
20pub trait NullOrderCmp {
21    fn null_order_cmp(&self, other: &Self, nulls_last: bool) -> Ordering;
22}
23
24impl<T: TotalOrd> NullOrderCmp for Option<T> {
25    fn null_order_cmp(&self, other: &Self, nulls_last: bool) -> Ordering {
26        match (self, other) {
27            (None, None) => Ordering::Equal,
28            (None, Some(_)) => {
29                if nulls_last {
30                    Ordering::Greater
31                } else {
32                    Ordering::Less
33                }
34            },
35            (Some(_), None) => {
36                if nulls_last {
37                    Ordering::Less
38                } else {
39                    Ordering::Greater
40                }
41            },
42            (Some(l), Some(r)) => l.tot_cmp(r),
43        }
44    }
45}
46
47impl<T: TotalOrd> NullOrderCmp for NonNull<T> {
48    fn null_order_cmp(&self, other: &Self, _nulls_last: bool) -> Ordering {
49        self.0.tot_cmp(&other.0)
50    }
51}
52
53trait GetInner {
54    type Item;
55    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item;
56}
57
58impl<'a, T: PolarsDataType> GetInner for &'a ChunkedArray<T> {
59    type Item = Option<T::Physical<'a>>;
60    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item {
61        ChunkedArray::get_unchecked(self, idx)
62    }
63}
64
65impl<'a, T: StaticArray> GetInner for &'a T {
66    type Item = Option<T::ValueT<'a>>;
67    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item {
68        <T as StaticArray>::get_unchecked(self, idx)
69    }
70}
71
72impl<'a, T: PolarsDataType> GetInner for NonNull<&'a ChunkedArray<T>> {
73    type Item = NonNull<T::Physical<'a>>;
74    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item {
75        NonNull(self.0.value_unchecked(idx))
76    }
77}
78
79impl<'a, T: StaticArray> GetInner for NonNull<&'a T> {
80    type Item = NonNull<T::ValueT<'a>>;
81    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item {
82        NonNull(self.0.value_unchecked(idx))
83    }
84}
85
86pub trait TotalEqInner: Send + Sync {
87    /// # Safety
88    /// Does not do any bound checks.
89    unsafe fn eq_element_unchecked(&self, idx_a: usize, idx_b: usize) -> bool;
90}
91
92pub trait TotalOrdInner: Send + Sync {
93    /// # Safety
94    /// Does not do any bound checks.
95    unsafe fn cmp_element_unchecked(
96        &self,
97        idx_a: usize,
98        idx_b: usize,
99        nulls_last: bool,
100    ) -> Ordering;
101}
102
103impl<T> TotalEqInner for T
104where
105    T: GetInner + Send + Sync,
106    T::Item: TotalEq,
107{
108    #[inline]
109    unsafe fn eq_element_unchecked(&self, idx_a: usize, idx_b: usize) -> bool {
110        self.get_unchecked(idx_a).tot_eq(&self.get_unchecked(idx_b))
111    }
112}
113
114impl TotalEqInner for &NullChunked {
115    unsafe fn eq_element_unchecked(&self, _idx_a: usize, _idx_b: usize) -> bool {
116        true
117    }
118}
119
120/// Create a type that implements TotalEqInner.
121pub(crate) trait IntoTotalEqInner<'a> {
122    /// Create a type that implements `TakeRandom`.
123    fn into_total_eq_inner(self) -> Box<dyn TotalEqInner + 'a>;
124}
125
126impl<'a> IntoTotalEqInner<'a> for &'a NullChunked {
127    fn into_total_eq_inner(self) -> Box<dyn TotalEqInner + 'a> {
128        Box::new(self)
129    }
130}
131
132/// We use a trait object because we want to call this from Series and cannot use a typed enum.
133impl<'a, T> IntoTotalEqInner<'a> for &'a ChunkedArray<T>
134where
135    T: PolarsDataType,
136    T::Physical<'a>: TotalEq,
137{
138    fn into_total_eq_inner(self) -> Box<dyn TotalEqInner + 'a> {
139        match self.layout() {
140            ChunkedArrayLayout::SingleNoNull(arr) => Box::new(NonNull(arr)),
141            ChunkedArrayLayout::Single(arr) => Box::new(arr),
142            ChunkedArrayLayout::MultiNoNull(ca) => Box::new(NonNull(ca)),
143            ChunkedArrayLayout::Multi(ca) => Box::new(ca),
144        }
145    }
146}
147
148impl<T> TotalOrdInner for T
149where
150    T: GetInner + Send + Sync,
151    T::Item: NullOrderCmp,
152{
153    #[inline]
154    unsafe fn cmp_element_unchecked(
155        &self,
156        idx_a: usize,
157        idx_b: usize,
158        nulls_last: bool,
159    ) -> Ordering {
160        let a = self.get_unchecked(idx_a);
161        let b = self.get_unchecked(idx_b);
162        a.null_order_cmp(&b, nulls_last)
163    }
164}
165
166/// Create a type that implements TotalOrdInner.
167pub(crate) trait IntoTotalOrdInner<'a> {
168    /// Create a type that implements `TakeRandom`.
169    fn into_total_ord_inner(self) -> Box<dyn TotalOrdInner + 'a>;
170}
171
172impl<'a, T> IntoTotalOrdInner<'a> for &'a ChunkedArray<T>
173where
174    T: PolarsDataType,
175    T::Physical<'a>: TotalOrd,
176{
177    fn into_total_ord_inner(self) -> Box<dyn TotalOrdInner + 'a> {
178        match self.layout() {
179            ChunkedArrayLayout::SingleNoNull(arr) => Box::new(NonNull(arr)),
180            ChunkedArrayLayout::Single(arr) => Box::new(arr),
181            ChunkedArrayLayout::MultiNoNull(ca) => Box::new(NonNull(ca)),
182            ChunkedArrayLayout::Multi(ca) => Box::new(ca),
183        }
184    }
185}
186
187#[cfg(feature = "dtype-categorical")]
188struct LocalCategorical<'a> {
189    rev_map: &'a Utf8ViewArray,
190    cats: &'a UInt32Chunked,
191}
192
193#[cfg(feature = "dtype-categorical")]
194impl<'a> GetInner for LocalCategorical<'a> {
195    type Item = Option<&'a str>;
196    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item {
197        let cat = self.cats.get_unchecked(idx)?;
198        Some(self.rev_map.value_unchecked(cat as usize))
199    }
200}
201
202#[cfg(feature = "dtype-categorical")]
203struct GlobalCategorical<'a> {
204    p1: &'a PlHashMap<u32, u32>,
205    p2: &'a Utf8ViewArray,
206    cats: &'a UInt32Chunked,
207}
208
209#[cfg(feature = "dtype-categorical")]
210impl<'a> GetInner for GlobalCategorical<'a> {
211    type Item = Option<&'a str>;
212    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item {
213        let cat = self.cats.get_unchecked(idx)?;
214        let idx = self.p1.get(&cat).unwrap();
215        Some(self.p2.value_unchecked(*idx as usize))
216    }
217}
218
219#[cfg(feature = "dtype-categorical")]
220impl<'a> IntoTotalOrdInner<'a> for &'a CategoricalChunked {
221    fn into_total_ord_inner(self) -> Box<dyn TotalOrdInner + 'a> {
222        let cats = self.physical();
223        match &**self.get_rev_map() {
224            RevMapping::Global(p1, p2, _) => Box::new(GlobalCategorical { p1, p2, cats }),
225            RevMapping::Local(rev_map, _) => Box::new(LocalCategorical { rev_map, cats }),
226        }
227    }
228}
polars_core/chunked_array/ops/compare_inner.rs

polars_core/chunked_array/ops/
compare_inner.rs
