vortex_array/search_sorted/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod primitive;
5
6use std::cmp::Ordering;
7use std::cmp::Ordering::Equal;
8use std::cmp::Ordering::Greater;
9use std::cmp::Ordering::Less;
10use std::fmt::Debug;
11use std::fmt::Display;
12use std::fmt::Formatter;
13use std::hint;
14
15pub use primitive::*;
16use vortex_error::VortexResult;
17
18use crate::ArrayRef;
19use crate::VortexSessionExecute;
20use crate::legacy_session;
21use crate::scalar::Scalar;
22
23#[derive(Debug, Copy, Clone, Eq, PartialEq)]
24pub enum SearchSortedSide {
25    Left,
26    Right,
27}
28
29impl Display for SearchSortedSide {
30    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
31        match self {
32            SearchSortedSide::Left => write!(f, "left"),
33            SearchSortedSide::Right => write!(f, "right"),
34        }
35    }
36}
37
38/// Result of performing search_sorted on an Array
39#[derive(Debug, Copy, Clone, PartialEq, Eq)]
40pub enum SearchResult {
41    /// Result for a found element was found in the array and another one could be inserted at the given position
42    /// in the sorted order
43    Found(usize),
44
45    /// Result for an element not found, but that could be inserted at the given position
46    /// in the sorted order.
47    NotFound(usize),
48}
49
50impl SearchResult {
51    /// Convert search result to an index only if the value have been found
52    pub fn to_found(self) -> Option<usize> {
53        match self {
54            Self::Found(i) => Some(i),
55            Self::NotFound(_) => None,
56        }
57    }
58
59    /// Extract index out of search result regardless of whether the value have been found or not
60    pub fn to_index(self) -> usize {
61        match self {
62            Self::Found(i) => i,
63            Self::NotFound(i) => i,
64        }
65    }
66
67    /// Convert search result into an index suitable for searching array of offset indices, i.e. first element starts at 0.
68    ///
69    /// For example for a ChunkedArray with chunk offsets array [0, 3, 8, 10] you can use this method to
70    /// obtain index suitable for indexing into it after performing a search
71    pub fn to_offsets_index(self, len: usize, side: SearchSortedSide) -> usize {
72        match self {
73            SearchResult::Found(i) => {
74                if side == SearchSortedSide::Right || i == len {
75                    i.saturating_sub(1)
76                } else {
77                    i
78                }
79            }
80            SearchResult::NotFound(i) => i.saturating_sub(1),
81        }
82    }
83
84    /// Convert search result into an index suitable for searching array of end indices without 0 offset,
85    /// i.e. first element implicitly covers 0..0th-element range.
86    ///
87    /// For example for a RunEndArray with ends array [3, 8, 10], you can use this method to obtain index suitable for
88    /// indexing into it after performing a search
89    pub fn to_ends_index(self, len: usize) -> usize {
90        let idx = self.to_index();
91        if idx == len { idx - 1 } else { idx }
92    }
93}
94
95impl Display for SearchResult {
96    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
97        match self {
98            SearchResult::Found(i) => write!(f, "Found({i})"),
99            SearchResult::NotFound(i) => write!(f, "NotFound({i})"),
100        }
101    }
102}
103
104pub trait IndexOrd<V> {
105    /// PartialOrd of the value at index `idx` with `elem`.
106    /// For example, if self\[idx\] > elem, return Some(Greater).
107    fn index_cmp(&self, idx: usize, elem: &V) -> VortexResult<Option<Ordering>>;
108
109    #[inline]
110    fn index_lt(&self, idx: usize, elem: &V) -> VortexResult<bool> {
111        Ok(matches!(self.index_cmp(idx, elem)?, Some(Less)))
112    }
113
114    #[inline]
115    fn index_le(&self, idx: usize, elem: &V) -> VortexResult<bool> {
116        Ok(matches!(self.index_cmp(idx, elem)?, Some(Less | Equal)))
117    }
118
119    fn index_gt(&self, idx: usize, elem: &V) -> VortexResult<bool> {
120        Ok(matches!(self.index_cmp(idx, elem)?, Some(Greater)))
121    }
122
123    fn index_ge(&self, idx: usize, elem: &V) -> VortexResult<bool> {
124        Ok(matches!(self.index_cmp(idx, elem)?, Some(Greater | Equal)))
125    }
126
127    /// Get the length of the underlying ordered collection
128    fn index_len(&self) -> usize;
129}
130
131/// Searches for value assuming the array is sorted.
132///
133/// Returned indices satisfy following condition if the search for value was to be inserted into the array at found positions
134///
135/// |side |result satisfies|
136/// |-----|----------------|
137/// |left |array\[i-1\] < value <= array\[i\]|
138/// |right|array\[i-1\] <= value < array\[i\]|
139pub trait SearchSorted<T> {
140    #[inline]
141    fn search_sorted(&self, value: &T, side: SearchSortedSide) -> VortexResult<SearchResult>
142    where
143        Self: IndexOrd<T>,
144    {
145        match side {
146            SearchSortedSide::Left => self.search_sorted_by(
147                |idx| Ok(self.index_cmp(idx, value)?.unwrap_or(Less)),
148                |idx| {
149                    Ok(if self.index_lt(idx, value)? {
150                        Less
151                    } else {
152                        Greater
153                    })
154                },
155                side,
156            ),
157            SearchSortedSide::Right => self.search_sorted_by(
158                |idx| Ok(self.index_cmp(idx, value)?.unwrap_or(Less)),
159                |idx| {
160                    Ok(if self.index_le(idx, value)? {
161                        Less
162                    } else {
163                        Greater
164                    })
165                },
166                side,
167            ),
168        }
169    }
170
171    /// find function is used to find the element if it exists, if element exists side_find will be
172    /// used to find desired index amongst equal values
173    fn search_sorted_by<
174        F: FnMut(usize) -> VortexResult<Ordering>,
175        N: FnMut(usize) -> VortexResult<Ordering>,
176    >(
177        &self,
178        find: F,
179        side_find: N,
180        side: SearchSortedSide,
181    ) -> VortexResult<SearchResult>;
182}
183
184// Default implementation for types that implement IndexOrd.
185impl<S, T> SearchSorted<T> for S
186where
187    S: IndexOrd<T> + ?Sized,
188{
189    #[inline]
190    fn search_sorted_by<
191        F: FnMut(usize) -> VortexResult<Ordering>,
192        N: FnMut(usize) -> VortexResult<Ordering>,
193    >(
194        &self,
195        find: F,
196        side_find: N,
197        side: SearchSortedSide,
198    ) -> VortexResult<SearchResult> {
199        match search_sorted_side_idx(find, 0, self.index_len())? {
200            SearchResult::Found(found) => {
201                let idx_search = match side {
202                    SearchSortedSide::Left => search_sorted_side_idx(side_find, 0, found)?,
203                    SearchSortedSide::Right => {
204                        search_sorted_side_idx(side_find, found, self.index_len())?
205                    }
206                };
207                match idx_search {
208                    SearchResult::NotFound(i) => Ok(SearchResult::Found(i)),
209                    _ => unreachable!(
210                        "searching amongst equal values should never return Found result"
211                    ),
212                }
213            }
214            s => Ok(s),
215        }
216    }
217}
218
219// Code adapted from Rust standard library slice::binary_search_by
220#[inline]
221fn search_sorted_side_idx<F: FnMut(usize) -> VortexResult<Ordering>>(
222    mut find: F,
223    from: usize,
224    to: usize,
225) -> VortexResult<SearchResult> {
226    let mut size = to - from;
227    if size == 0 {
228        return Ok(SearchResult::NotFound(0));
229    }
230    let mut base = from;
231
232    // This loop intentionally doesn't have an early exit if the comparison
233    // returns Equal. We want the number of loop iterations to depend *only*
234    // on the size of the input slice so that the CPU can reliably predict
235    // the loop count.
236    while size > 1 {
237        let half = size / 2;
238        let mid = base + half;
239
240        // SAFETY: the call is made safe by the following inconstants:
241        // - `mid >= 0`: by definition
242        // - `mid < size`: `mid = size / 2 + size / 4 + size / 8 ...`
243        let cmp = find(mid)?;
244
245        // Binary search interacts poorly with branch prediction, so force
246        // the compiler to use conditional moves if supported by the target
247        // architecture.
248        base = hint::select_unpredictable(cmp == Greater, base, mid);
249
250        // This is imprecise in the case where `size` is odd and the
251        // comparison returns Greater: the mid element still gets included
252        // by `size` even though it's known to be larger than the element
253        // being searched for.
254        // This is fine though: we gain more performance by keeping the
255        // loop iteration count invariant (and thus predictable) than we
256        // lose from considering one additional element.
257        size -= half;
258    }
259
260    // SAFETY: base is always in [0, size) because base <= mid.
261    let cmp = find(base)?;
262    if cmp == Equal {
263        // SAFETY: same as the call to `find` above.
264        unsafe { hint::assert_unchecked(base < to) };
265        Ok(SearchResult::Found(base))
266    } else {
267        let result = base + (cmp == Less) as usize;
268        // SAFETY: same as the call to `find` above.
269        // Note that this is `<=`, unlike the assert in the `Found` path.
270        unsafe { hint::assert_unchecked(result <= to) };
271        Ok(SearchResult::NotFound(result))
272    }
273}
274
275impl IndexOrd<Scalar> for ArrayRef {
276    #[allow(clippy::disallowed_methods)]
277    fn index_cmp(&self, idx: usize, elem: &Scalar) -> VortexResult<Option<Ordering>> {
278        let scalar_a = self.execute_scalar(idx, &mut legacy_session().create_execution_ctx())?;
279        Ok(scalar_a.partial_cmp(elem))
280    }
281
282    fn index_len(&self) -> usize {
283        Self::len(self)
284    }
285}
286
287impl<T: PartialOrd> IndexOrd<T> for [T] {
288    #[inline]
289    fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult<Option<Ordering>> {
290        // SAFETY: Used in search_sorted_by same as the standard library. The search_sorted ensures idx is in bounds
291        Ok(unsafe { self.get_unchecked(idx) }.partial_cmp(elem))
292    }
293
294    #[inline]
295    fn index_len(&self) -> usize {
296        self.len()
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use vortex_error::VortexResult;
303
304    use crate::search_sorted::SearchResult;
305    use crate::search_sorted::SearchSorted;
306    use crate::search_sorted::SearchSortedSide;
307
308    #[test]
309    fn left_side_equal() -> VortexResult<()> {
310        let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9];
311        let res = arr.search_sorted(&2, SearchSortedSide::Left)?;
312        assert_eq!(arr[res.to_index()], 2);
313        assert_eq!(res, SearchResult::Found(2));
314        Ok(())
315    }
316
317    #[test]
318    fn right_side_equal() -> VortexResult<()> {
319        let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9];
320        let res = arr.search_sorted(&2, SearchSortedSide::Right)?;
321        assert_eq!(arr[res.to_index() - 1], 2);
322        assert_eq!(res, SearchResult::Found(6));
323        Ok(())
324    }
325
326    #[test]
327    fn left_side_equal_beginning() -> VortexResult<()> {
328        let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
329        let res = arr.search_sorted(&0, SearchSortedSide::Left)?;
330        assert_eq!(arr[res.to_index()], 0);
331        assert_eq!(res, SearchResult::Found(0));
332        Ok(())
333    }
334
335    #[test]
336    fn right_side_equal_beginning() -> VortexResult<()> {
337        let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
338        let res = arr.search_sorted(&0, SearchSortedSide::Right)?;
339        assert_eq!(arr[res.to_index() - 1], 0);
340        assert_eq!(res, SearchResult::Found(4));
341        Ok(())
342    }
343
344    #[test]
345    fn left_side_equal_end() -> VortexResult<()> {
346        let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9];
347        let res = arr.search_sorted(&9, SearchSortedSide::Left)?;
348        assert_eq!(arr[res.to_index()], 9);
349        assert_eq!(res, SearchResult::Found(9));
350        Ok(())
351    }
352
353    #[test]
354    fn right_side_equal_end() -> VortexResult<()> {
355        let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9];
356        let res = arr.search_sorted(&9, SearchSortedSide::Right)?;
357        assert_eq!(arr[res.to_index() - 1], 9);
358        assert_eq!(res, SearchResult::Found(13));
359        Ok(())
360    }
361}
vortex_array/search_sorted/mod.rs

vortex_array/search_sorted/
mod.rs
