Fix out ouf bounds when taking from run end arrays (#501)

spiraldb · Jul 22, 2024 · 6a9873e · 6a9873e
1 parent b083633
commit 6a9873e
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 16 deletions.
diff --git a/encodings/runend-bool/src/array.rs b/encodings/runend-bool/src/array.rs
@@ -71,7 +71,7 @@ impl RunEndBoolArray {
             vortex_bail!("Index must be in array slice",);
         }
         search_sorted(&self.ends(), index + self.offset(), SearchSortedSide::Right)
-            .map(|s| s.to_index())
+            .map(|s| s.to_non_zero_offset_index(self.ends().len()))
     }
 
     pub fn validity(&self) -> Validity {

diff --git a/encodings/runend/src/runend.rs b/encodings/runend/src/runend.rs
@@ -73,7 +73,7 @@ impl RunEndArray {
 
     pub fn find_physical_index(&self, index: usize) -> VortexResult<usize> {
         search_sorted(&self.ends(), index + self.offset(), SearchSortedSide::Right)
-            .map(|s| s.to_index())
+            .map(|s| s.to_non_zero_offset_index(self.ends().len()))
     }
 
     pub fn encode(array: Array) -> VortexResult<Self> {

diff --git a/vortex-array/src/array/chunked/compute/take.rs b/vortex-array/src/array/chunked/compute/take.rs
@@ -79,7 +79,7 @@ fn take_strict_sorted(chunked: &ChunkedArray, indices: &Array) -> VortexResult<A
         let chunk_end = usize::try_from(&scalar_at(&chunked.chunk_ends(), chunk_idx + 1)?).unwrap();
         let chunk_end_pos = search_sorted(indices, chunk_end, SearchSortedSide::Left)
             .unwrap()
-            .to_index();
+            .to_zero_offset_index();
 
         // Now we can say the slice of indices belonging to this chunk is [pos, chunk_end_pos)
         let chunk_indices = slice(indices, pos, chunk_end_pos)?;

diff --git a/vortex-array/src/array/sparse/compute/slice.rs b/vortex-array/src/array/sparse/compute/slice.rs
@@ -8,9 +8,9 @@ impl SliceFn for SparseArray {
     fn slice(&self, start: usize, stop: usize) -> VortexResult<Array> {
         // Find the index of the first patch index that is greater than or equal to the offset of this array
         let index_start_index =
-            search_sorted(&self.indices(), start, SearchSortedSide::Left)?.to_index();
+            search_sorted(&self.indices(), start, SearchSortedSide::Left)?.to_zero_offset_index();
         let index_end_index =
-            search_sorted(&self.indices(), stop, SearchSortedSide::Left)?.to_index();
+            search_sorted(&self.indices(), stop, SearchSortedSide::Left)?.to_zero_offset_index();
 
         Ok(Self::try_new_with_offset(
             slice(&self.indices(), index_start_index, index_end_index)?,

diff --git a/vortex-array/src/compute/search_sorted.rs b/vortex-array/src/compute/search_sorted.rs
@@ -27,7 +27,20 @@ impl SearchResult {
         }
     }
 
-    pub fn to_index(self) -> usize {
+    pub fn to_non_zero_offset_index(self, len: usize) -> usize {
+        match self {
+            SearchResult::Found(i) => i,
+            SearchResult::NotFound(i) => {
+                if i == len {
+                    i - 1
+                } else {
+                    i
+                }
+            }
+        }
+    }
+
+    pub fn to_zero_offset_index(self) -> usize {
         match self {
             Self::Found(i) => i,
             Self::NotFound(i) => i,
@@ -221,47 +234,47 @@ mod test {
     fn left_side_equal() {
         let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9];
         let res = arr.search_sorted(&2, SearchSortedSide::Left);
-        assert_eq!(arr[res.to_index()], 2);
+        assert_eq!(arr[res.to_zero_offset_index()], 2);
         assert_eq!(res, SearchResult::Found(2));
     }
 
     #[test]
     fn right_side_equal() {
         let arr = [0, 1, 2, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9];
         let res = arr.search_sorted(&2, SearchSortedSide::Right);
-        assert_eq!(arr[res.to_index() - 1], 2);
+        assert_eq!(arr[res.to_zero_offset_index() - 1], 2);
         assert_eq!(res, SearchResult::Found(6));
     }
 
     #[test]
     fn left_side_equal_beginning() {
         let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
         let res = arr.search_sorted(&0, SearchSortedSide::Left);
-        assert_eq!(arr[res.to_index()], 0);
+        assert_eq!(arr[res.to_zero_offset_index()], 0);
         assert_eq!(res, SearchResult::Found(0));
     }
 
     #[test]
     fn right_side_equal_beginning() {
         let arr = [0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
         let res = arr.search_sorted(&0, SearchSortedSide::Right);
-        assert_eq!(arr[res.to_index() - 1], 0);
+        assert_eq!(arr[res.to_zero_offset_index() - 1], 0);
         assert_eq!(res, SearchResult::Found(4));
     }
 
     #[test]
     fn left_side_equal_end() {
         let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9];
         let res = arr.search_sorted(&9, SearchSortedSide::Left);
-        assert_eq!(arr[res.to_index()], 9);
+        assert_eq!(arr[res.to_zero_offset_index()], 9);
         assert_eq!(res, SearchResult::Found(9));
     }
 
     #[test]
     fn right_side_equal_end() {
         let arr = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 9];
         let res = arr.search_sorted(&9, SearchSortedSide::Right);
-        assert_eq!(arr[res.to_index() - 1], 9);
+        assert_eq!(arr[res.to_zero_offset_index() - 1], 9);
         assert_eq!(res, SearchResult::Found(13));
     }
 }
diff --git a/vortex-array/src/stream/take_rows.rs b/vortex-array/src/stream/take_rows.rs
@@ -72,13 +72,14 @@ impl<R: ArrayStream> Stream for TakeRows<R> {
 
         while let Some(batch) = ready!(this.reader.as_mut().poll_next(cx)?) {
             let curr_offset = *this.row_offset;
-            let left = search_sorted(this.indices, curr_offset, SearchSortedSide::Left)?.to_index();
+            let left = search_sorted(this.indices, curr_offset, SearchSortedSide::Left)?
+                .to_zero_offset_index();
             let right = search_sorted(
                 this.indices,
                 curr_offset + batch.len(),
                 SearchSortedSide::Left,
             )?
-            .to_index();
+            .to_zero_offset_index();
 
             *this.row_offset += batch.len();
 

diff --git a/vortex-ipc/src/chunked_reader/take_rows.rs b/vortex-ipc/src/chunked_reader/take_rows.rs
@@ -122,9 +122,9 @@ impl<R: VortexReadAt> ChunkedArrayReader<R> {
 
         // Relativize the indices to these chunks
         let indices_start =
-            search_sorted(indices, row_range.start, SearchSortedSide::Left)?.to_index();
+            search_sorted(indices, row_range.start, SearchSortedSide::Left)?.to_zero_offset_index();
         let indices_stop =
-            search_sorted(indices, row_range.end, SearchSortedSide::Right)?.to_index();
+            search_sorted(indices, row_range.end, SearchSortedSide::Right)?.to_zero_offset_index();
         let relative_indices = slice(indices, indices_start, indices_stop)?;
         let row_start_scalar = Scalar::from(row_range.start).cast(relative_indices.dtype())?;
         let relative_indices = subtract_scalar(&relative_indices, &row_start_scalar)?;