Fix SearchSorted for SparseArray when searching from Right (#770)

spiraldb · Sep 10, 2024 · a82eb07 · a82eb07
1 parent bcae104
commit a82eb07
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 6 deletions.
diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml
@@ -14,7 +14,7 @@ jobs:
       - name: Install cargo fuzz
         run: cargo install cargo-fuzz
       - name: Run fuzzing target
-        run: cargo fuzz run fuzz_target_1 -- -max_total_time=600
+        run: RUST_BACKTRACE=1 cargo fuzz run fuzz_target_1 -- -max_total_time=600
         continue-on-error: true
       - name: Archive crash artifacts
         uses: actions/upload-artifact@v4

diff --git a/fuzz/src/search_sorted.rs b/fuzz/src/search_sorted.rs
@@ -14,9 +14,8 @@ impl<T: PartialOrd> IndexOrd<Option<T>> for SearchNullableSlice<T> {
     fn index_cmp(&self, idx: usize, elem: &Option<T>) -> Option<Ordering> {
         match elem {
             None => unreachable!("Can't search for None"),
-            Some(v) =>
-            // SAFETY: Used in search_sorted_by same as the standard library. The search_sorted ensures idx is in bounds
-            {
+            Some(v) => {
+                // SAFETY: Used in search_sorted_by same as the standard library. The search_sorted ensures idx is in bounds
                 match unsafe { self.0.get_unchecked(idx) } {
                     None => Some(Ordering::Greater),
                     Some(i) => i.partial_cmp(v),

diff --git a/vortex-array/src/array/sparse/compute/mod.rs b/vortex-array/src/array/sparse/compute/mod.rs
@@ -54,8 +54,23 @@ impl SearchSortedFn for SparseArray {
     fn search_sorted(&self, value: &Scalar, side: SearchSortedSide) -> VortexResult<SearchResult> {
         search_sorted(&self.values(), value.clone(), side).and_then(|sr| match sr {
             SearchResult::Found(i) => {
-                let index: usize = scalar_at(&self.indices(), i)?.as_ref().try_into()?;
-                Ok(SearchResult::Found(index - self.indices_offset()))
+                let index: usize = scalar_at(
+                    &self.indices(),
+                    if i == self.metadata().indices_len {
+                        i - 1
+                    } else {
+                        i
+                    },
+                )?
+                .as_ref()
+                .try_into()?;
+                Ok(SearchResult::Found(
+                    if i == self.metadata().indices_len {
+                        index + 1
+                    } else {
+                        index
+                    } - self.indices_offset(),
+                ))
             }
             SearchResult::NotFound(i) => {
                 let index: usize = scalar_at(&self.indices(), if i == 0 { 0 } else { i - 1 })?
@@ -109,6 +124,12 @@ mod test {
         assert_eq!(res, SearchResult::Found(9));
     }
 
+    #[test]
+    pub fn search_not_found_right() {
+        let res = search_sorted(&array(), 56, SearchSortedSide::Right).unwrap();
+        assert_eq!(res, SearchResult::NotFound(16));
+    }
+
     #[test]
     pub fn search_sliced() {
         let array = slice(&array(), 7, 20).unwrap();
@@ -117,4 +138,25 @@ mod test {
             SearchResult::NotFound(2)
         );
     }
+
+    #[test]
+    pub fn search_right() {
+        let array = SparseArray::try_new(
+            PrimitiveArray::from(vec![0u64]).into_array(),
+            PrimitiveArray::from_vec(vec![0u8], Validity::AllValid).into_array(),
+            2,
+            Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)),
+        )
+        .unwrap()
+        .into_array();
+
+        assert_eq!(
+            search_sorted(&array, 0, SearchSortedSide::Right).unwrap(),
+            SearchResult::Found(1)
+        );
+        assert_eq!(
+            search_sorted(&array, 1, SearchSortedSide::Right).unwrap(),
+            SearchResult::NotFound(1)
+        );
+    }
 }