polars_utils/
sort.rs

1use std::cmp::Ordering;
2use std::mem::MaybeUninit;
3use std::ops::Deref;
4
5use num_traits::FromPrimitive;
6
7use crate::itertools::Itertools;
8use crate::nulls::IsNull;
9use crate::total_ord::TotalOrd;
10
11unsafe fn assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T] {
12    unsafe { &mut *(slice as *mut [MaybeUninit<T>] as *mut [T]) }
13}
14
15pub fn arg_sort_ascending<'a, T: TotalOrd + Copy + 'a, Idx, I: IntoIterator<Item = T>>(
16    v: I,
17    scratch: &'a mut Vec<u8>,
18    n: usize,
19) -> &'a mut [Idx]
20where
21    Idx: FromPrimitive + Copy,
22{
23    let upper_bound = size_of::<(T, Idx)>() * n + align_of::<(T, Idx)>();
24    scratch.reserve(upper_bound);
25    let scratch_slice = unsafe {
26        let cap_slice = scratch.spare_capacity_mut();
27        let (_, scratch_slice, _) = cap_slice.align_to_mut::<MaybeUninit<(T, Idx)>>();
28        &mut scratch_slice[..n]
29    };
30
31    for ((i, v), dst) in v.into_iter().enumerate().zip_eq(scratch_slice.iter_mut()) {
32        *dst = MaybeUninit::new((v, Idx::from_usize(i).unwrap()));
33    }
34
35    let scratch_slice = unsafe { assume_init_mut(scratch_slice) };
36    scratch_slice.sort_by(|key1, key2| key1.0.tot_cmp(&key2.0));
37
38    // Now we write the indexes in the same array. So from (T, Idxsize) to just IdxSize.
39    unsafe {
40        let src = scratch_slice.as_ptr();
41        let (_, scratch_slice_aligned_to_idx, _) = scratch_slice.align_to_mut::<Idx>();
42        let dst = scratch_slice_aligned_to_idx.as_mut_ptr();
43
44        for i in 0..n {
45            dst.add(i).write((*src.add(i)).1);
46        }
47
48        &mut scratch_slice_aligned_to_idx[..n]
49    }
50}
51
52#[derive(PartialEq, Eq, Clone, Hash)]
53#[repr(transparent)]
54pub struct ReorderWithNulls<T, const DESCENDING: bool, const NULLS_LAST: bool>(pub Option<T>);
55
56impl<T, const DESCENDING: bool, const NULLS_LAST: bool>
57    ReorderWithNulls<T, DESCENDING, NULLS_LAST>
58{
59    pub fn as_deref(&self) -> ReorderWithNulls<&<T as Deref>::Target, DESCENDING, NULLS_LAST>
60    where
61        T: Deref,
62    {
63        let x = self.0.as_deref();
64        ReorderWithNulls(x)
65    }
66}
67
68impl<T: PartialOrd, const DESCENDING: bool, const NULLS_LAST: bool> PartialOrd
69    for ReorderWithNulls<T, DESCENDING, NULLS_LAST>
70{
71    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
72        match (&self.0, &other.0) {
73            (None, None) => Some(Ordering::Equal),
74            (None, Some(_)) => {
75                if NULLS_LAST {
76                    Some(Ordering::Greater)
77                } else {
78                    Some(Ordering::Less)
79                }
80            },
81            (Some(_), None) => {
82                if NULLS_LAST {
83                    Some(Ordering::Less)
84                } else {
85                    Some(Ordering::Greater)
86                }
87            },
88            (Some(l), Some(r)) => {
89                if DESCENDING {
90                    r.partial_cmp(l)
91                } else {
92                    l.partial_cmp(r)
93                }
94            },
95        }
96    }
97}
98
99impl<T: Ord, const DESCENDING: bool, const NULLS_LAST: bool> Ord
100    for ReorderWithNulls<T, DESCENDING, NULLS_LAST>
101{
102    fn cmp(&self, other: &Self) -> Ordering {
103        reorder_cmp(&self.0, &other.0, DESCENDING, NULLS_LAST)
104    }
105}
106
107/// Compare two values with support for sort direction and nulls position.
108///
109/// # Panics
110///
111/// Panics if `T::partial_cmp(lhs, rhs)` returns `None`.
112#[inline]
113pub fn reorder_cmp<T: PartialOrd + IsNull>(
114    lhs: &T,
115    rhs: &T,
116    descending: bool,
117    nulls_last: bool,
118) -> Ordering {
119    match PartialOrd::partial_cmp(lhs, rhs).expect("expected total ordering") {
120        Ordering::Equal => Ordering::Equal,
121        _ if lhs.is_null() && nulls_last => Ordering::Greater,
122        _ if rhs.is_null() && nulls_last => Ordering::Less,
123        _ if lhs.is_null() => Ordering::Less,
124        _ if rhs.is_null() => Ordering::Greater,
125        ord if descending => ord.reverse(),
126        ord => ord,
127    }
128}
polars_utils/sort.rs

polars_utils/
sort.rs
