1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

use std::cmp::Ordering;

use vortex_array::accessor::ArrayAccessor;
use vortex_array::array::{BoolArray, PrimitiveArray, VarBinViewArray};
use vortex_array::compute::scalar_at;
use vortex_array::validity::ArrayValidity;
use vortex_array::{ArrayDType, ArrayData, IntoArrayData, IntoArrayVariant};
use vortex_dtype::{match_each_native_ptype, DType, NativePType};

use crate::take::take_canonical_array;

pub fn sort_canonical_array(array: &ArrayData) -> ArrayData {
    match array.dtype() {
        DType::Bool(_) => {
            let bool_array = array.clone().into_bool().unwrap();
            let mut opt_values = bool_array
                .boolean_buffer()
                .iter()
                .zip(
                    bool_array
                        .logical_validity()
                        .into_array()
                        .into_bool()
                        .unwrap()
                        .boolean_buffer()
                        .iter(),
                )
                .map(|(b, v)| v.then_some(b))
                .collect::<Vec<_>>();
            sort_opt_slice(&mut opt_values);
            BoolArray::from_iter(opt_values).into_array()
        }
        DType::Primitive(p, _) => {
            let primitive_array = array.clone().into_primitive().unwrap();
            match_each_native_ptype!(p, |$P| {
                let mut opt_values = primitive_array
                    .maybe_null_slice::<$P>()
                    .iter()
                    .copied()
                    .zip(
                        primitive_array
                            .logical_validity()
                            .into_array()
                            .into_bool()
                            .unwrap()
                            .boolean_buffer()
                            .iter(),
                    )
                    .map(|(p, v)| v.then_some(p))
                    .collect::<Vec<_>>();
                sort_primitive_slice(&mut opt_values);
                PrimitiveArray::from_nullable_vec(opt_values).into_array()
            })
        }
        DType::Utf8(_) | DType::Binary(_) => {
            let utf8 = array.clone().into_varbinview().unwrap();
            let mut opt_values = utf8
                .with_iterator(|iter| iter.map(|v| v.map(|u| u.to_vec())).collect::<Vec<_>>())
                .unwrap();
            sort_opt_slice(&mut opt_values);
            VarBinViewArray::from_iter(opt_values, array.dtype().clone()).into_array()
        }
        DType::Struct(..) => {
            let mut sort_indices = (0..array.len()).collect::<Vec<_>>();
            sort_indices.sort_by(|a, b| {
                scalar_at(array, *a)
                    .unwrap()
                    .partial_cmp(&scalar_at(array, *b).unwrap())
                    .unwrap()
            });
            take_canonical_array(array, &sort_indices)
        }
        _ => unreachable!("Not a canonical array"),
    }
}

fn sort_primitive_slice<T: NativePType>(s: &mut [Option<T>]) {
    s.sort_by(|a, b| match (a, b) {
        (Some(v), Some(w)) => v.total_compare(*w),
        (None, None) => Ordering::Equal,
        (None, Some(_)) => Ordering::Greater,
        (Some(_), None) => Ordering::Less,
    });
}

/// Reverse sorting of Option<T> such that None is last (Greatest)
fn sort_opt_slice<T: Ord>(s: &mut [Option<T>]) {
    s.sort_by(|a, b| match (a, b) {
        (Some(v), Some(w)) => v.cmp(w),
        (None, None) => Ordering::Equal,
        (None, Some(_)) => Ordering::Greater,
        (Some(_), None) => Ordering::Less,
    });
}