vortex_array/search_sorted/
scalar.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::cell::RefCell;
5use std::cmp::Ordering;
6
7use vortex_error::VortexResult;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::RepeatedArrayProbe;
12use crate::scalar::Scalar;
13use crate::search_sorted::IndexOrd;
14
15/// A [`SearchSorted`](crate::search_sorted::SearchSorted) adapter over a sorted array of any
16/// encoding, comparing elements as [`Scalar`]s.
17///
18/// Reads go through a [`RepeatedArrayProbe`], so the whole search shares one execution context
19/// and whatever the first comparison resolved.
20///
21/// Prefer [`SearchSortedPrimitiveArray`](crate::search_sorted::SearchSortedPrimitiveArray) when
22/// the element type is known: comparing through `Scalar` cannot read the values directly and
23/// builds a scalar per probe.
24pub struct SearchSortedArray<'a> {
25    probe: RefCell<RepeatedArrayProbe>,
26    len: usize,
27    ctx: RefCell<&'a mut ExecutionCtx>,
28}
29
30impl<'a> SearchSortedArray<'a> {
31    /// Wraps `array` for searching.
32    pub fn new(array: &ArrayRef, ctx: &'a mut ExecutionCtx) -> Self {
33        Self {
34            probe: RefCell::new(array.repeated_probe()),
35            len: array.len(),
36            ctx: RefCell::new(ctx),
37        }
38    }
39}
40
41impl IndexOrd<Scalar> for SearchSortedArray<'_> {
42    /// The probe and the context are separate cells, so the two borrows never overlap.
43    fn index_cmp(&self, idx: usize, elem: &Scalar) -> VortexResult<Option<Ordering>> {
44        let scalar = self
45            .probe
46            .borrow_mut()
47            .execute_scalar(idx, &mut self.ctx.borrow_mut())?;
48        Ok(scalar.partial_cmp(elem))
49    }
50
51    fn index_len(&self) -> usize {
52        self.len
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use vortex_buffer::buffer;
59    use vortex_error::VortexResult;
60
61    use crate::IntoArray;
62    use crate::array_session;
63    use crate::arrays::PrimitiveArray;
64    use crate::executor::VortexSessionExecute;
65    use crate::scalar::Scalar;
66    use crate::search_sorted::SearchResult;
67    use crate::search_sorted::SearchSorted;
68    use crate::search_sorted::SearchSortedArray;
69    use crate::search_sorted::SearchSortedSide;
70    use crate::validity::Validity;
71
72    #[test]
73    fn search_sorted_scalar() -> VortexResult<()> {
74        let array =
75            PrimitiveArray::new(buffer![1i32, 2, 3, 3, 4], Validity::NonNullable).into_array();
76        let mut ctx = array_session().create_execution_ctx();
77        let searcher = SearchSortedArray::new(&array, &mut ctx);
78        assert_eq!(
79            searcher.search_sorted(&Scalar::from(3i32), SearchSortedSide::Left)?,
80            SearchResult::Found(2)
81        );
82        assert_eq!(
83            searcher.search_sorted(&Scalar::from(3i32), SearchSortedSide::Right)?,
84            SearchResult::Found(4)
85        );
86        assert_eq!(
87            searcher.search_sorted(&Scalar::from(5i32), SearchSortedSide::Left)?,
88            SearchResult::NotFound(5)
89        );
90        Ok(())
91    }
92
93    #[test]
94    fn search_sorted_scalar_with_nulls() -> VortexResult<()> {
95        let array =
96            PrimitiveArray::from_option_iter([None, None, Some(2i32), Some(3)]).into_array();
97        let mut ctx = array_session().create_execution_ctx();
98        let searcher = SearchSortedArray::new(&array, &mut ctx);
99        assert_eq!(
100            searcher.search_sorted(&Scalar::from(Some(2i32)), SearchSortedSide::Left)?,
101            SearchResult::Found(2)
102        );
103        Ok(())
104    }
105}
vortex_array/search_sorted/scalar.rs

vortex_array/search_sorted/
scalar.rs
