Address review findings on SearchSortedPrimitiveArray

- Fix RefCell double-borrow panic in IndexOrd<Option<T>>::index_cmp: the
  RefMut guard from is_valid was still alive when value() re-borrowed the
  ctx, so searching any valid element with an optional needle panicked
  with BorrowMutError. Add regression tests covering Option<T> searches
  with and without nulls.
- Stop swallowing the underlying error in Patches
  search_index_binary_search_scalar; the map_err replaced real execution
  errors with a misleading "indices must be a primitive array" message.
- Pass the in-scope ExecutionCtx to find_physical_index in RunEnd
  scalar_at and slice instead of spinning up a fresh LEGACY_SESSION ctx
  per call.
- Add doc comments to the new public SearchSortedPrimitiveArray API and
  document the null-as-zero semantics of plain T searches.
- Restore VortexResult<()> + ? style in test_search_sorted_primitive.

Signed-off-by: "Robert Kruszewski" <robert@spiraldb.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: robert3005

encodings/runend/src/kernel.rs

@@ -25,6 +25,7 @@ use vortex_session::VortexSession;
 use crate::RunEnd;
 use crate::array::RunEndArrayExt;
 use crate::compute::take_from::RunEndTakeFrom;
+use crate::ops::find_physical_index;
 use crate::ops::find_slice_end_index;
 
 pub(super) fn initialize(session: &VortexSession) {
@@ -64,7 +65,7 @@ fn slice(
 ) -> VortexResult<ArrayRef> {
     let new_length = range.len();
 
-    let slice_begin = array.find_physical_index(range.start)?;
+    let slice_begin = find_physical_index(array.ends(), range.start + array.offset(), ctx)?;
     let slice_end = find_slice_end_index(array.ends(), range.end + array.offset(), ctx)?;
 
     // If the sliced range contains only a single run, opt to return a ConstantArray.

encodings/runend/src/ops.rs

@@ -22,9 +22,8 @@ impl OperationsVTable<RunEnd> for RunEnd {
         index: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar> {
-        array
-            .values()
-            .execute_scalar(array.find_physical_index(index)?, ctx)
+        let physical_index = find_physical_index(array.ends(), index + array.offset(), ctx)?;
+        array.values().execute_scalar(physical_index, ctx)
     }
 }
 

vortex-array/src/arrays/primitive/tests.rs

@@ -4,7 +4,6 @@
 use std::sync::LazyLock;
 
 use vortex_buffer::buffer;
-use vortex_error::VortexExpect;
 use vortex_session::VortexSession;
 
 use crate::ArrayRef;
@@ -32,11 +31,11 @@ fn test_search_sorted_primitive(
     #[case] value: i32,
     #[case] side: SearchSortedSide,
     #[case] expected: SearchResult,
-) {
+) -> vortex_error::VortexResult<()> {
     let res = SearchSortedPrimitiveArray::<i32>::new(&array, &mut SESSION.create_execution_ctx())
-        .search_sorted(&value, side)
-        .vortex_expect("Failed to search sorted");
+        .search_sorted(&value, side)?;
     assert_eq!(res, expected);
+    Ok(())
 }
 
 #[test]

vortex-array/src/patches.rs

@@ -1009,7 +1009,6 @@ fn search_index_binary_search_scalar(
     match_each_unsigned_integer_ptype!(indices.dtype().as_ptype(), |T| {
         SearchSortedPrimitiveArray::<T>::new(indices, &mut LEGACY_SESSION.create_execution_ctx())
             .search_sorted(&needle, SearchSortedSide::Left)
-            .map_err(|_| vortex_err!("indices must be a primitive array"))
     })
 }
 

vortex-array/src/search_sorted/primitive.rs

@@ -12,13 +12,21 @@ use crate::ExecutionCtx;
 use crate::dtype::NativePType;
 use crate::search_sorted::IndexOrd;
 
+/// A [`SearchSorted`](crate::search_sorted::SearchSorted) adapter over a sorted primitive-typed
+/// array, comparing elements as the native type `T`.
+///
+/// Values can be searched as `T`, `Option<T>`, or `usize`. Searching as `T` or `usize` treats
+/// null elements as `T::zero()`; use `Option<T>` when the array may contain nulls, in which case
+/// nulls sort before all non-null values.
 pub struct SearchSortedPrimitiveArray<'a, T>(
     &'a ArrayRef,
     RefCell<&'a mut ExecutionCtx>,
     PhantomData<T>,
 );
 
 impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> {
+    /// Wraps `array` for searching, panicking if the array's [`PType`](crate::dtype::PType) is
+    /// not `T::PTYPE`.
     pub fn new(array: &'a ArrayRef, ctx: &'a mut ExecutionCtx) -> Self {
         assert_eq!(
             array.dtype().as_ptype(),
@@ -28,11 +36,11 @@ impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> {
         Self(array, RefCell::new(ctx), PhantomData)
     }
 
+    /// Returns the value at `idx`, with nulls mapped to `T::zero()`.
     fn value(&self, idx: usize) -> VortexResult<T> {
-        let ctx_mut = &mut self.1.borrow_mut();
         Ok(self
             .0
-            .execute_scalar(idx, ctx_mut)?
+            .execute_scalar(idx, &mut self.1.borrow_mut())?
             .as_primitive()
             .typed_value::<T>()
             .unwrap_or_else(|| T::zero()))
@@ -52,12 +60,9 @@ impl<T: NativePType> IndexOrd<T> for SearchSortedPrimitiveArray<'_, T> {
 
 impl<T: NativePType> IndexOrd<Option<T>> for SearchSortedPrimitiveArray<'_, T> {
     fn index_cmp(&self, idx: usize, elem: &Option<T>) -> VortexResult<Option<Ordering>> {
-        let ctx_mut = &mut self.1.borrow_mut();
-        let value = self
-            .0
-            .is_valid(idx, ctx_mut)?
-            .then(|| self.value(idx))
-            .transpose()?;
+        // The borrow must end before `self.value` re-borrows the ctx.
+        let valid = self.0.is_valid(idx, &mut self.1.borrow_mut())?;
+        let value = valid.then(|| self.value(idx)).transpose()?;
 
         Ok(match (value, elem.as_ref()) {
             (Some(l), Some(r)) => Some(l.total_compare(*r)),
@@ -87,3 +92,46 @@ impl<T: NativePType> IndexOrd<usize> for SearchSortedPrimitiveArray<'_, T> {
         self.0.len()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
+
+    use crate::IntoArray;
+    use crate::array_session;
+    use crate::arrays::PrimitiveArray;
+    use crate::executor::VortexSessionExecute;
+    use crate::search_sorted::SearchResult;
+    use crate::search_sorted::SearchSorted;
+    use crate::search_sorted::SearchSortedPrimitiveArray;
+    use crate::search_sorted::SearchSortedSide;
+    use crate::validity::Validity;
+
+    #[test]
+    fn search_sorted_optional_value() -> VortexResult<()> {
+        let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::AllValid).into_array();
+        let mut ctx = array_session().create_execution_ctx();
+        let res = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx)
+            .search_sorted(&Some(3i32), SearchSortedSide::Left)?;
+        assert_eq!(res, SearchResult::Found(2));
+        Ok(())
+    }
+
+    #[test]
+    fn search_sorted_optional_value_with_nulls() -> VortexResult<()> {
+        let array =
+            PrimitiveArray::from_option_iter([None, None, Some(2i32), Some(3)]).into_array();
+        let mut ctx = array_session().create_execution_ctx();
+        let searcher = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx);
+        assert_eq!(
+            searcher.search_sorted(&Some(2i32), SearchSortedSide::Left)?,
+            SearchResult::Found(2)
+        );
+        assert_eq!(
+            searcher.search_sorted(&None, SearchSortedSide::Right)?,
+            SearchResult::Found(2)
+        );
+        Ok(())
+    }
+}