polars_core/chunked_array/ops/
compare_inner.rs

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
92impl<T> TotalEqInner for T
93where
94    T: GetInner + Send + Sync,
95    T::Item: TotalEq,
96{
97    #[inline]
98    unsafe fn eq_element_unchecked(&self, idx_a: usize, idx_b: usize) -> bool {
99        self.get_unchecked(idx_a).tot_eq(&self.get_unchecked(idx_b))
100    }
101}
102
103impl TotalEqInner for &NullChunked {
104    unsafe fn eq_element_unchecked(&self, _idx_a: usize, _idx_b: usize) -> bool {
105        true
106    }
107}
108
109/// Create a type that implements TotalEqInner.
110pub(crate) trait IntoTotalEqInner<'a> {
111    /// Create a type that implements `TakeRandom`.
112    fn into_total_eq_inner(self) -> Box<dyn TotalEqInner + 'a>;
113}
114
115impl<'a> IntoTotalEqInner<'a> for &'a NullChunked {
116    fn into_total_eq_inner(self) -> Box<dyn TotalEqInner + 'a> {
117        Box::new(self)
118    }
119}
120
121/// We use a trait object because we want to call this from Series and cannot use a typed enum.
122impl<'a, T> IntoTotalEqInner<'a> for &'a ChunkedArray<T>
123where
124    T: PolarsDataType,
125    T::Physical<'a>: TotalEq,
126{
127    fn into_total_eq_inner(self) -> Box<dyn TotalEqInner + 'a> {
128        match self.layout() {
129            ChunkedArrayLayout::SingleNoNull(arr) => Box::new(NonNull(arr)),
130            ChunkedArrayLayout::Single(arr) => Box::new(arr),
131            ChunkedArrayLayout::MultiNoNull(ca) => Box::new(NonNull(ca)),
132            ChunkedArrayLayout::Multi(ca) => Box::new(ca),
133        }
134    }
135}
136
137pub trait TotalOrdInner: Send + Sync {
138    /// # Safety
139    /// Does not do any bound checks.
140    unsafe fn cmp_element_unchecked(
141        &self,
142        idx_a: usize,
143        idx_b: usize,
144        nulls_last: bool,
145    ) -> Ordering;
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
166impl TotalOrdInner for &NullChunked {
167    #[inline]
168    unsafe fn cmp_element_unchecked(
169        &self,
170        _idx_a: usize,
171        _idx_b: usize,
172        _nulls_last: bool,
173    ) -> Ordering {
174        Ordering::Equal
175    }
176}
177
178/// Create a type that implements TotalOrdInner.
179pub(crate) trait IntoTotalOrdInner<'a> {
180    /// Create a type that implements `TakeRandom`.
181    fn into_total_ord_inner(self) -> Box<dyn TotalOrdInner + 'a>;
182}
183
184impl<'a, T> IntoTotalOrdInner<'a> for &'a ChunkedArray<T>
185where
186    T: PolarsDataType,
187    T::Physical<'a>: TotalOrd,
188{
189    fn into_total_ord_inner(self) -> Box<dyn TotalOrdInner + 'a> {
190        match self.layout() {
191            ChunkedArrayLayout::SingleNoNull(arr) => Box::new(NonNull(arr)),
192            ChunkedArrayLayout::Single(arr) => Box::new(arr),
193            ChunkedArrayLayout::MultiNoNull(ca) => Box::new(NonNull(ca)),
194            ChunkedArrayLayout::Multi(ca) => Box::new(ca),
195        }
196    }
197}
198
199impl<'a> IntoTotalOrdInner<'a> for &'a NullChunked {
200    fn into_total_ord_inner(self) -> Box<dyn TotalOrdInner + 'a> {
201        Box::new(self)
202    }
203}
204
205#[cfg(feature = "dtype-categorical")]
206struct LexicalCategorical<'a, T: PolarsCategoricalType> {
207    mapping: &'a CategoricalMapping,
208    cats: &'a ChunkedArray<T::PolarsPhysical>,
209}
210
211#[cfg(feature = "dtype-categorical")]
212impl<'a, T: PolarsCategoricalType> GetInner for LexicalCategorical<'a, T> {
213    type Item = Option<&'a str>;
214    unsafe fn get_unchecked(&self, idx: usize) -> Self::Item {
215        let cat = self.cats.get_unchecked(idx)?;
216        Some(self.mapping.cat_to_str_unchecked(cat.as_cat()))
217    }
218}
219
220#[cfg(feature = "dtype-categorical")]
221impl<'a, T: PolarsCategoricalType> IntoTotalOrdInner<'a> for &'a CategoricalChunked<T> {
222    fn into_total_ord_inner(self) -> Box<dyn TotalOrdInner + 'a> {
223        if self.uses_lexical_ordering() {
224            Box::new(LexicalCategorical::<T> {
225                mapping: self.get_mapping(),
226                cats: &self.phys,
227            })
228        } else {
229            self.phys.into_total_ord_inner()
230        }
231    }
232}
polars_core/chunked_array/ops/compare_inner.rs

polars_core/chunked_array/ops/
compare_inner.rs
