vortex_array/compute/
is_sorted.rs

1use vortex_error::{VortexExpect, VortexResult, vortex_bail};
2
3use crate::arrays::{ConstantArray, NullArray};
4use crate::stats::{Precision, Stat};
5use crate::{Array, ArrayExt, Encoding};
6
7pub trait IsSortedFn<A> {
8    /// # Preconditions
9    /// - The array's length is > 1.
10    /// - The array is not encoded as `NullArray` or `ConstantArray`.
11    /// - If doing a `strict` check, if the array is nullable, it'll have at most 1 null element
12    ///   as the first item in the array.
13    fn is_sorted(&self, array: A) -> VortexResult<bool>;
14
15    fn is_strict_sorted(&self, array: A) -> VortexResult<bool>;
16}
17
18impl<E: Encoding> IsSortedFn<&dyn Array> for E
19where
20    E: for<'a> IsSortedFn<&'a E::Array>,
21{
22    fn is_sorted(&self, array: &dyn Array) -> VortexResult<bool> {
23        let array_ref = array
24            .as_any()
25            .downcast_ref::<E::Array>()
26            .vortex_expect("Failed to downcast array");
27        IsSortedFn::is_sorted(self, array_ref)
28    }
29
30    fn is_strict_sorted(&self, array: &dyn Array) -> VortexResult<bool> {
31        let array_ref = array
32            .as_any()
33            .downcast_ref::<E::Array>()
34            .vortex_expect("Failed to downcast array");
35        IsSortedFn::is_strict_sorted(self, array_ref)
36    }
37}
38
39#[allow(clippy::wrong_self_convention)]
40/// Helper methods to check sortedness with strictness
41pub trait IsSortedIteratorExt: Iterator
42where
43    <Self as Iterator>::Item: PartialOrd,
44{
45    fn is_strict_sorted(self) -> bool
46    where
47        Self: Sized,
48        Self::Item: PartialOrd,
49    {
50        self.is_sorted_by(|a, b| a < b)
51    }
52}
53
54impl<T> IsSortedIteratorExt for T
55where
56    T: Iterator + ?Sized,
57    T::Item: PartialOrd,
58{
59}
60
61pub fn is_sorted(array: &dyn Array) -> VortexResult<bool> {
62    // We currently don't support sorting struct arrays.
63    if array.dtype().is_struct() {
64        return Ok(false);
65    }
66
67    if let Some(Precision::Exact(value)) = array.statistics().get_as::<bool>(Stat::IsSorted) {
68        return Ok(value);
69    }
70
71    let is_sorted = is_sorted_impl(array, false)?;
72    let array_stats = array.statistics();
73
74    if is_sorted {
75        array_stats.set(Stat::IsSorted, Precision::Exact(true.into()));
76    } else {
77        array_stats.set(Stat::IsSorted, Precision::Exact(false.into()));
78        array_stats.set(Stat::IsStrictSorted, Precision::Exact(false.into()));
79    }
80
81    Ok(is_sorted)
82}
83pub fn is_strict_sorted(array: &dyn Array) -> VortexResult<bool> {
84    // We currently don't support sorting struct arrays.
85    if array.dtype().is_struct() {
86        return Ok(false);
87    }
88
89    if let Some(Precision::Exact(value)) = array.statistics().get_as::<bool>(Stat::IsStrictSorted) {
90        return Ok(value);
91    }
92
93    let is_sorted = is_sorted_impl(array, true)?;
94    let array_stats = array.statistics();
95
96    if is_sorted {
97        array_stats.set(Stat::IsSorted, Precision::Exact(true.into()));
98        array_stats.set(Stat::IsStrictSorted, Precision::Exact(true.into()));
99    } else {
100        array_stats.set(Stat::IsStrictSorted, Precision::Exact(false.into()));
101    }
102
103    Ok(is_sorted)
104}
105
106fn is_sorted_impl(array: &dyn Array, strict: bool) -> VortexResult<bool> {
107    // Arrays with 0 or 1 elements are strict sorted.
108    if array.len() <= 1 {
109        return Ok(true);
110    }
111
112    // Constant and null arrays are always sorted, but not strict sorted.
113    if array.is::<ConstantArray>() || array.is::<NullArray>() {
114        return Ok(!strict);
115    }
116
117    let invalid_count = array.invalid_count()?;
118
119    // Enforce strictness before we even try to check if the array is sorted.
120    if strict {
121        match invalid_count {
122            // We can keep going
123            0 => {}
124            // If we have a potential null value - it has to be the first one.
125            1 => {
126                if !array.is_invalid(0)? {
127                    return Ok(false);
128                }
129            }
130            _ => return Ok(false),
131        }
132    }
133
134    let is_sorted = if let Some(vtable_fn) = array.vtable().is_sorted_fn() {
135        if strict {
136            vtable_fn.is_strict_sorted(array)?
137        } else {
138            vtable_fn.is_sorted(array)?
139        }
140    } else {
141        log::debug!("No is_sorted implementation found for {}", array.encoding());
142        let array = array.to_canonical()?;
143
144        if let Some(vtable_fn) = array.as_ref().vtable().is_sorted_fn() {
145            let array = array.as_ref();
146            if strict {
147                vtable_fn.is_strict_sorted(array)?
148            } else {
149                vtable_fn.is_sorted(array)?
150            }
151        } else {
152            vortex_bail!(
153                "No is_sorted function for canonical array: {}",
154                array.as_ref().encoding(),
155            )
156        }
157    };
158
159    Ok(is_sorted)
160}
161
162#[cfg(test)]
163mod tests {
164    use vortex_buffer::buffer;
165
166    use crate::Array;
167    use crate::arrays::{BoolArray, PrimitiveArray};
168    use crate::compute::{is_sorted, is_strict_sorted};
169    use crate::validity::Validity;
170
171    #[test]
172    fn test_is_sorted() {
173        assert!(
174            is_sorted(&PrimitiveArray::new(
175                buffer!(0, 1, 2, 3),
176                Validity::AllValid
177            ))
178            .unwrap()
179        );
180        assert!(
181            is_sorted(&PrimitiveArray::new(
182                buffer!(0, 1, 2, 3),
183                Validity::Array(BoolArray::from_iter([false, true, true, true]).into_array())
184            ))
185            .unwrap()
186        );
187        assert!(
188            is_sorted(&PrimitiveArray::new(
189                buffer!(0, 1, 2, 3),
190                Validity::Array(BoolArray::from_iter([true, false, true, true]).into_array())
191            ))
192            .unwrap()
193        );
194
195        assert!(
196            !is_sorted(&PrimitiveArray::new(
197                buffer!(0, 1, 3, 2),
198                Validity::AllValid
199            ))
200            .unwrap()
201        );
202        assert!(
203            !is_sorted(&PrimitiveArray::new(
204                buffer!(0, 1, 3, 2),
205                Validity::Array(BoolArray::from_iter([false, true, true, true]).into_array()),
206            ))
207            .unwrap(),
208        );
209    }
210
211    #[test]
212    fn test_is_strict_sorted() {
213        assert!(
214            is_strict_sorted(&PrimitiveArray::new(
215                buffer!(0, 1, 2, 3),
216                Validity::AllValid
217            ))
218            .unwrap()
219        );
220        assert!(
221            is_strict_sorted(&PrimitiveArray::new(
222                buffer!(0, 1, 2, 3),
223                Validity::Array(BoolArray::from_iter([false, true, true, true]).into_array())
224            ))
225            .unwrap()
226        );
227        assert!(
228            !is_strict_sorted(&PrimitiveArray::new(
229                buffer!(0, 1, 2, 3),
230                Validity::Array(BoolArray::from_iter([true, false, true, true]).into_array()),
231            ))
232            .unwrap(),
233        );
234
235        assert!(
236            !is_strict_sorted(&PrimitiveArray::new(
237                buffer!(0, 1, 3, 2),
238                Validity::Array(BoolArray::from_iter([false, true, true, true]).into_array()),
239            ))
240            .unwrap(),
241        );
242    }
243}
vortex_array/compute/is_sorted.rs

vortex_array/compute/
is_sorted.rs
