vortex_array/search_sorted/
primitive.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::cell::RefCell;
5use std::cmp::Ordering;
6use std::marker::PhantomData;
7
8use vortex_error::VortexResult;
9
10use crate::ArrayRef;
11use crate::ExecutionCtx;
12use crate::dtype::NativePType;
13use crate::search_sorted::IndexOrd;
14
15/// A [`SearchSorted`](crate::search_sorted::SearchSorted) adapter over a sorted primitive-typed
16/// array, comparing elements as the native type `T`.
17///
18/// Values can be searched as `T`, `Option<T>`, or `usize`. Searching as `T` or `usize` treats
19/// null elements as `T::zero()`; use `Option<T>` when the array may contain nulls, in which case
20/// nulls sort before all non-null values.
21pub struct SearchSortedPrimitiveArray<'a, T>(
22    &'a ArrayRef,
23    RefCell<&'a mut ExecutionCtx>,
24    PhantomData<T>,
25);
26
27impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> {
28    /// Wraps `array` for searching, panicking if the array's [`PType`](crate::dtype::PType) is
29    /// not `T::PTYPE`.
30    pub fn new(array: &'a ArrayRef, ctx: &'a mut ExecutionCtx) -> Self {
31        assert_eq!(
32            array.dtype().as_ptype(),
33            T::PTYPE,
34            "Array PType must match primitive type"
35        );
36        Self(array, RefCell::new(ctx), PhantomData)
37    }
38
39    /// Returns the value at `idx`, with nulls mapped to `T::zero()`.
40    fn value(&self, idx: usize) -> VortexResult<T> {
41        Ok(self
42            .0
43            .execute_scalar(idx, &mut self.1.borrow_mut())?
44            .as_primitive()
45            .typed_value::<T>()
46            .unwrap_or_else(|| T::zero()))
47    }
48}
49
50impl<T: NativePType> IndexOrd<T> for SearchSortedPrimitiveArray<'_, T> {
51    fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult<Option<Ordering>> {
52        let value = self.value(idx)?;
53        Ok(Some(value.total_compare(*elem)))
54    }
55
56    fn index_len(&self) -> usize {
57        self.0.len()
58    }
59}
60
61impl<T: NativePType> IndexOrd<Option<T>> for SearchSortedPrimitiveArray<'_, T> {
62    fn index_cmp(&self, idx: usize, elem: &Option<T>) -> VortexResult<Option<Ordering>> {
63        // The borrow must end before `self.value` re-borrows the ctx.
64        let valid = self.0.is_valid(idx, &mut self.1.borrow_mut())?;
65        let value = valid.then(|| self.value(idx)).transpose()?;
66
67        Ok(match (value, elem.as_ref()) {
68            (Some(l), Some(r)) => Some(l.total_compare(*r)),
69            (Some(_), None) => Some(Ordering::Greater),
70            (None, Some(_)) => Some(Ordering::Less),
71            (None, None) => Some(Ordering::Equal),
72        })
73    }
74
75    fn index_len(&self) -> usize {
76        self.0.len()
77    }
78}
79
80impl<T: NativePType> IndexOrd<usize> for SearchSortedPrimitiveArray<'_, T> {
81    fn index_cmp(&self, idx: usize, elem: &usize) -> VortexResult<Option<Ordering>> {
82        let value = self.value(idx)?;
83
84        let Some(elem_t) = T::from_usize(*elem) else {
85            return Ok(Some(Ordering::Less));
86        };
87
88        Ok(Some(value.total_compare(elem_t)))
89    }
90
91    fn index_len(&self) -> usize {
92        self.0.len()
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use vortex_buffer::buffer;
99    use vortex_error::VortexResult;
100
101    use crate::IntoArray;
102    use crate::array_session;
103    use crate::arrays::PrimitiveArray;
104    use crate::executor::VortexSessionExecute;
105    use crate::search_sorted::SearchResult;
106    use crate::search_sorted::SearchSorted;
107    use crate::search_sorted::SearchSortedPrimitiveArray;
108    use crate::search_sorted::SearchSortedSide;
109    use crate::validity::Validity;
110
111    #[test]
112    fn search_sorted_optional_value() -> VortexResult<()> {
113        let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::AllValid).into_array();
114        let mut ctx = array_session().create_execution_ctx();
115        let res = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx)
116            .search_sorted(&Some(3i32), SearchSortedSide::Left)?;
117        assert_eq!(res, SearchResult::Found(2));
118        Ok(())
119    }
120
121    #[test]
122    fn search_sorted_optional_value_with_nulls() -> VortexResult<()> {
123        let array =
124            PrimitiveArray::from_option_iter([None, None, Some(2i32), Some(3)]).into_array();
125        let mut ctx = array_session().create_execution_ctx();
126        let searcher = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx);
127        assert_eq!(
128            searcher.search_sorted(&Some(2i32), SearchSortedSide::Left)?,
129            SearchResult::Found(2)
130        );
131        assert_eq!(
132            searcher.search_sorted(&None, SearchSortedSide::Right)?,
133            SearchResult::Found(2)
134        );
135        Ok(())
136    }
137}
vortex_array/search_sorted/primitive.rs

vortex_array/search_sorted/
primitive.rs
