• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

vortex-data / vortex / 16935267080

13 Aug 2025 11:00AM UTC coverage: 24.312% (-63.3%) from 87.658%
16935267080

Pull #4226

github

web-flow
Merge 81b48c7fb into baa6ea202
Pull Request #4226: Support converting TimestampTZ to and from duckdb

0 of 2 new or added lines in 1 file covered. (0.0%)

20666 existing lines in 469 files now uncovered.

8726 of 35892 relevant lines covered (24.31%)

147.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0
/vortex-array/src/search_sorted.rs
1
// SPDX-License-Identifier: Apache-2.0
2
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3

4
use std::cmp::Ordering;
5
use std::cmp::Ordering::{Equal, Greater, Less};
6
use std::fmt::{Debug, Display, Formatter};
7
use std::hint;
8

9
use vortex_scalar::Scalar;
10

11
use crate::Array;
12

13
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
14
pub enum SearchSortedSide {
15
    Left,
16
    Right,
17
}
18

19
impl Display for SearchSortedSide {
20
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
×
21
        match self {
×
22
            SearchSortedSide::Left => write!(f, "left"),
×
23
            SearchSortedSide::Right => write!(f, "right"),
×
24
        }
25
    }
×
26
}
27

28
/// Result of performing search_sorted on an Array
29
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
30
pub enum SearchResult {
31
    /// Result for a found element was found in the array and another one could be inserted at the given position
32
    /// in the sorted order
33
    Found(usize),
34

35
    /// Result for an element not found, but that could be inserted at the given position
36
    /// in the sorted order.
37
    NotFound(usize),
38
}
39

40
impl SearchResult {
41
    /// Convert search result to an index only if the value have been found
UNCOV
42
    pub fn to_found(self) -> Option<usize> {
×
UNCOV
43
        match self {
×
UNCOV
44
            Self::Found(i) => Some(i),
×
UNCOV
45
            Self::NotFound(_) => None,
×
46
        }
UNCOV
47
    }
×
48

49
    /// Extract index out of search result regardless of whether the value have been found or not
UNCOV
50
    pub fn to_index(self) -> usize {
×
UNCOV
51
        match self {
×
UNCOV
52
            Self::Found(i) => i,
×
UNCOV
53
            Self::NotFound(i) => i,
×
54
        }
UNCOV
55
    }
×
56

57
    /// Convert search result into an index suitable for searching array of offset indices, i.e. first element starts at 0.
58
    ///
59
    /// For example for a ChunkedArray with chunk offsets array [0, 3, 8, 10] you can use this method to
60
    /// obtain index suitable for indexing into it after performing a search
UNCOV
61
    pub fn to_offsets_index(self, len: usize, side: SearchSortedSide) -> usize {
×
UNCOV
62
        match self {
×
UNCOV
63
            SearchResult::Found(i) => {
×
UNCOV
64
                if side == SearchSortedSide::Right || i == len {
×
UNCOV
65
                    i.saturating_sub(1)
×
66
                } else {
UNCOV
67
                    i
×
68
                }
69
            }
UNCOV
70
            SearchResult::NotFound(i) => i.saturating_sub(1),
×
71
        }
UNCOV
72
    }
×
73

74
    /// Convert search result into an index suitable for searching array of end indices without 0 offset,
75
    /// i.e. first element implicitly covers 0..0th-element range.
76
    ///
77
    /// For example for a RunEndArray with ends array [3, 8, 10], you can use this method to obtain index suitable for
78
    /// indexing into it after performing a search
UNCOV
79
    pub fn to_ends_index(self, len: usize) -> usize {
×
UNCOV
80
        let idx = self.to_index();
×
UNCOV
81
        if idx == len { idx - 1 } else { idx }
×
UNCOV
82
    }
×
83
}
84

85
impl Display for SearchResult {
86
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
×
87
        match self {
×
88
            SearchResult::Found(i) => write!(f, "Found({i})"),
×
89
            SearchResult::NotFound(i) => write!(f, "NotFound({i})"),
×
90
        }
91
    }
×
92
}
93

94
pub trait IndexOrd<V> {
95
    /// PartialOrd of the value at index `idx` with `elem`.
96
    /// For example, if self\[idx\] > elem, return Some(Greater).
97
    fn index_cmp(&self, idx: usize, elem: &V) -> Option<Ordering>;
98

UNCOV
99
    fn index_lt(&self, idx: usize, elem: &V) -> bool {
×
UNCOV
100
        matches!(self.index_cmp(idx, elem), Some(Less))
×
UNCOV
101
    }
×
102

UNCOV
103
    fn index_le(&self, idx: usize, elem: &V) -> bool {
×
UNCOV
104
        matches!(self.index_cmp(idx, elem), Some(Less | Equal))
×
UNCOV
105
    }
×
106

107
    fn index_gt(&self, idx: usize, elem: &V) -> bool {
×
108
        matches!(self.index_cmp(idx, elem), Some(Greater))
×
109
    }
×
110

111
    fn index_ge(&self, idx: usize, elem: &V) -> bool {
×
112
        matches!(self.index_cmp(idx, elem), Some(Greater | Equal))
×
113
    }
×
114

115
    /// Get the length of the underlying ordered collection
116
    fn index_len(&self) -> usize;
117
}
118

119
/// Searches for value assuming the array is sorted.
120
///
121
/// Returned indices satisfy following condition if the search for value was to be inserted into the array at found positions
122
///
123
/// |side |result satisfies|
124
/// |-----|----------------|
125
/// |left |array\[i-1\] < value <= array\[i\]|
126
/// |right|array\[i-1\] <= value < array\[i\]|
127
pub trait SearchSorted<T> {
128
    fn search_sorted_many<I: IntoIterator<Item = T>>(
×
129
        &self,
×
130
        values: I,
×
131
        side: SearchSortedSide,
×
132
    ) -> impl Iterator<Item = SearchResult>
×
133
    where
×
134
        Self: IndexOrd<T>,
×
135
    {
136
        values
×
137
            .into_iter()
×
138
            .map(move |value| self.search_sorted(&value, side))
×
139
    }
×
140

UNCOV
141
    fn search_sorted(&self, value: &T, side: SearchSortedSide) -> SearchResult
×
UNCOV
142
    where
×
UNCOV
143
        Self: IndexOrd<T>,
×
144
    {
UNCOV
145
        match side {
×
UNCOV
146
            SearchSortedSide::Left => self.search_sorted_by(
×
UNCOV
147
                |idx| self.index_cmp(idx, value).unwrap_or(Less),
×
UNCOV
148
                |idx| {
×
UNCOV
149
                    if self.index_lt(idx, value) {
×
UNCOV
150
                        Less
×
151
                    } else {
UNCOV
152
                        Greater
×
153
                    }
UNCOV
154
                },
×
UNCOV
155
                side,
×
156
            ),
UNCOV
157
            SearchSortedSide::Right => self.search_sorted_by(
×
UNCOV
158
                |idx| self.index_cmp(idx, value).unwrap_or(Less),
×
UNCOV
159
                |idx| {
×
UNCOV
160
                    if self.index_le(idx, value) {
×
UNCOV
161
                        Less
×
162
                    } else {
UNCOV
163
                        Greater
×
164
                    }
UNCOV
165
                },
×
UNCOV
166
                side,
×
167
            ),
168
        }
UNCOV
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<F: FnMut(usize) -> Ordering, N: FnMut(usize) -> Ordering>(
174
        &self,
175
        find: F,
176
        side_find: N,
177
        side: SearchSortedSide,
178
    ) -> SearchResult;
179
}
180

181
// Default implementation for types that implement IndexOrd.
182
impl<S, T> SearchSorted<T> for S
183
where
184
    S: IndexOrd<T> + ?Sized,
185
{
UNCOV
186
    fn search_sorted_by<F: FnMut(usize) -> Ordering, N: FnMut(usize) -> Ordering>(
×
UNCOV
187
        &self,
×
UNCOV
188
        find: F,
×
UNCOV
189
        side_find: N,
×
UNCOV
190
        side: SearchSortedSide,
×
UNCOV
191
    ) -> SearchResult {
×
UNCOV
192
        match search_sorted_side_idx(find, 0, self.index_len()) {
×
UNCOV
193
            SearchResult::Found(found) => {
×
UNCOV
194
                let idx_search = match side {
×
UNCOV
195
                    SearchSortedSide::Left => search_sorted_side_idx(side_find, 0, found),
×
196
                    SearchSortedSide::Right => {
UNCOV
197
                        search_sorted_side_idx(side_find, found, self.index_len())
×
198
                    }
199
                };
UNCOV
200
                match idx_search {
×
UNCOV
201
                    SearchResult::NotFound(i) => SearchResult::Found(i),
×
202
                    _ => unreachable!(
×
203
                        "searching amongst equal values should never return Found result"
204
                    ),
205
                }
206
            }
UNCOV
207
            s => s,
×
208
        }
UNCOV
209
    }
×
210
}
211

212
// Code adapted from Rust standard library slice::binary_search_by
UNCOV
213
fn search_sorted_side_idx<F: FnMut(usize) -> Ordering>(
×
UNCOV
214
    mut find: F,
×
UNCOV
215
    from: usize,
×
UNCOV
216
    to: usize,
×
UNCOV
217
) -> SearchResult {
×
UNCOV
218
    let mut size = to - from;
×
UNCOV
219
    if size == 0 {
×
UNCOV
220
        return SearchResult::NotFound(0);
×
UNCOV
221
    }
×
UNCOV
222
    let mut base = from;
×
223

224
    // This loop intentionally doesn't have an early exit if the comparison
225
    // returns Equal. We want the number of loop iterations to depend *only*
226
    // on the size of the input slice so that the CPU can reliably predict
227
    // the loop count.
UNCOV
228
    while size > 1 {
×
UNCOV
229
        let half = size / 2;
×
UNCOV
230
        let mid = base + half;
×
231

232
        // SAFETY: the call is made safe by the following inconstants:
233
        // - `mid >= 0`: by definition
234
        // - `mid < size`: `mid = size / 2 + size / 4 + size / 8 ...`
UNCOV
235
        let cmp = find(mid);
×
236

237
        // Binary search interacts poorly with branch prediction, so force
238
        // the compiler to use conditional moves if supported by the target
239
        // architecture.
UNCOV
240
        base = if cmp == Greater { base } else { mid };
×
241

242
        // This is imprecise in the case where `size` is odd and the
243
        // comparison returns Greater: the mid element still gets included
244
        // by `size` even though it's known to be larger than the element
245
        // being searched for.
246
        // This is fine though: we gain more performance by keeping the
247
        // loop iteration count invariant (and thus predictable) than we
248
        // lose from considering one additional element.
UNCOV
249
        size -= half;
×
250
    }
251

252
    // SAFETY: base is always in [0, size) because base <= mid.
UNCOV
253
    let cmp = find(base);
×
UNCOV
254
    if cmp == Equal {
×
255
        // SAFETY: same as the call to `find` above.
UNCOV
256
        unsafe { hint::assert_unchecked(base < to) };
×
UNCOV
257
        SearchResult::Found(base)
×
258
    } else {
UNCOV
259
        let result = base + (cmp == Less) as usize;
×
260
        // SAFETY: same as the call to `find` above.
261
        // Note that this is `<=`, unlike the assert in the `Found` path.
UNCOV
262
        unsafe { hint::assert_unchecked(result <= to) };
×
UNCOV
263
        SearchResult::NotFound(result)
×
264
    }
UNCOV
265
}
×
266

267
impl IndexOrd<Scalar> for dyn Array + '_ {
268
    fn index_cmp(&self, idx: usize, elem: &Scalar) -> Option<Ordering> {
×
269
        let scalar_a = self.scalar_at(idx).ok()?;
×
270
        scalar_a.partial_cmp(elem)
×
271
    }
×
272

273
    fn index_len(&self) -> usize {
×
274
        Self::len(self)
×
275
    }
×
276
}
277

278
impl<T: PartialOrd> IndexOrd<T> for [T] {
UNCOV
279
    fn index_cmp(&self, idx: usize, elem: &T) -> Option<Ordering> {
×
280
        // SAFETY: Used in search_sorted_by same as the standard library. The search_sorted ensures idx is in bounds
UNCOV
281
        unsafe { self.get_unchecked(idx) }.partial_cmp(elem)
×
UNCOV
282
    }
×
283

UNCOV
284
    fn index_len(&self) -> usize {
×
UNCOV
285
        self.len()
×
UNCOV
286
    }
×
287
}
288

289
#[cfg(test)]
290
mod test {
291
    use crate::search_sorted::{SearchResult, SearchSorted, SearchSortedSide};
292

293
    #[test]
294
    fn left_side_equal() {
295
        let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9];
296
        let res = arr.search_sorted(&2, SearchSortedSide::Left);
297
        assert_eq!(arr[res.to_index()], 2);
298
        assert_eq!(res, SearchResult::Found(2));
299
    }
300

301
    #[test]
302
    fn right_side_equal() {
303
        let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9];
304
        let res = arr.search_sorted(&2, SearchSortedSide::Right);
305
        assert_eq!(arr[res.to_index() - 1], 2);
306
        assert_eq!(res, SearchResult::Found(6));
307
    }
308

309
    #[test]
310
    fn left_side_equal_beginning() {
311
        let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
312
        let res = arr.search_sorted(&0, SearchSortedSide::Left);
313
        assert_eq!(arr[res.to_index()], 0);
314
        assert_eq!(res, SearchResult::Found(0));
315
    }
316

317
    #[test]
318
    fn right_side_equal_beginning() {
319
        let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
320
        let res = arr.search_sorted(&0, SearchSortedSide::Right);
321
        assert_eq!(arr[res.to_index() - 1], 0);
322
        assert_eq!(res, SearchResult::Found(4));
323
    }
324

325
    #[test]
326
    fn left_side_equal_end() {
327
        let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9];
328
        let res = arr.search_sorted(&9, SearchSortedSide::Left);
329
        assert_eq!(arr[res.to_index()], 9);
330
        assert_eq!(res, SearchResult::Found(9));
331
    }
332

333
    #[test]
334
    fn right_side_equal_end() {
335
        let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9];
336
        let res = arr.search_sorted(&9, SearchSortedSide::Right);
337
        assert_eq!(arr[res.to_index() - 1], 9);
338
        assert_eq!(res, SearchResult::Found(13));
339
    }
340
}
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc