Fix RunEnd take and scalar_at compute functions (#588)

When accessing the last element we would return out of bounds index
spiraldb · Aug 9, 2024 · c15cc74 · c15cc74
1 parent 6986386
commit c15cc74
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 25 deletions.
diff --git a/encodings/runend-bool/src/array.rs b/encodings/runend-bool/src/array.rs
@@ -67,11 +67,13 @@ impl RunEndBoolArray {
     }
 
     pub fn find_physical_index(&self, index: usize) -> VortexResult<usize> {
-        if index > self.len() {
-            vortex_bail!("Index must be in array slice",);
-        }
-        search_sorted(&self.ends(), index + self.offset(), SearchSortedSide::Right)
-            .map(|s| s.to_index())
+        let searched_index =
+            search_sorted(&self.ends(), index + self.offset(), SearchSortedSide::Right)?.to_index();
+        Ok(if searched_index == self.ends().len() {
+            searched_index - 1
+        } else {
+            searched_index
+        })
     }
 
     pub fn validity(&self) -> Validity {

diff --git a/encodings/runend-bool/src/compute.rs b/encodings/runend-bool/src/compute.rs
@@ -3,7 +3,7 @@ use vortex::compute::unary::ScalarAtFn;
 use vortex::compute::{slice, ArrayCompute, SliceFn, TakeFn};
 use vortex::{Array, IntoArray, IntoArrayVariant, ToArray};
 use vortex_dtype::match_each_integer_ptype;
-use vortex_error::VortexResult;
+use vortex_error::{vortex_bail, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::compress::value_at_index;
@@ -40,8 +40,12 @@ impl TakeFn for RunEndBoolArray {
             primitive_indices
                 .maybe_null_slice::<$P>()
                 .iter()
+                .map(|idx| *idx as usize)
                 .map(|idx| {
-                    self.find_physical_index(*idx as usize)
+                    if idx >= self.len() {
+                        vortex_bail!(OutOfBounds: idx, 0, self.len())
+                    }
+                    self.find_physical_index(idx)
                 })
                 .collect::<VortexResult<Vec<_>>>()?
         });
@@ -60,12 +64,11 @@ impl SliceFn for RunEndBoolArray {
     fn slice(&self, start: usize, stop: usize) -> VortexResult<Array> {
         let slice_begin = self.find_physical_index(start)?;
         let slice_end = self.find_physical_index(stop)?;
-        let bounded_slice_end = usize::min(slice_end + 1, self.ends().len());
 
         Ok(Self::with_offset_and_size(
-            slice(&self.ends(), slice_begin, bounded_slice_end)?,
+            slice(&self.ends(), slice_begin, slice_end + 1)?,
             value_at_index(slice_begin, self.start()),
-            self.validity().slice(slice_begin, bounded_slice_end)?,
+            self.validity().slice(slice_begin, slice_end + 1)?,
             stop - start,
             start,
         )?

diff --git a/encodings/runend/src/compute.rs b/encodings/runend/src/compute.rs
@@ -3,7 +3,7 @@ use vortex::compute::unary::{scalar_at, ScalarAtFn};
 use vortex::compute::{slice, take, ArrayCompute, SliceFn, TakeFn};
 use vortex::{Array, IntoArray, IntoArrayVariant};
 use vortex_dtype::match_each_integer_ptype;
-use vortex_error::VortexResult;
+use vortex_error::{vortex_bail, VortexResult};
 use vortex_scalar::Scalar;
 
 use crate::RunEndArray;
@@ -35,9 +35,12 @@ impl TakeFn for RunEndArray {
             primitive_indices
                 .maybe_null_slice::<$P>()
                 .iter()
+                .map(|idx| *idx as usize)
                 .map(|idx| {
-                    self.find_physical_index(*idx as usize)
-                        .map(|loc| loc as u64)
+                    if idx >= self.len() {
+                        vortex_bail!(OutOfBounds: idx, 0, self.len())
+                    }
+                    self.find_physical_index(idx).map(|loc| loc as u64)
                 })
                 .collect::<VortexResult<Vec<_>>>()?
         });
@@ -52,12 +55,11 @@ impl SliceFn for RunEndArray {
     fn slice(&self, start: usize, stop: usize) -> VortexResult<Array> {
         let slice_begin = self.find_physical_index(start)?;
         let slice_end = self.find_physical_index(stop)?;
-        let bounded_slice_end = usize::min(slice_end + 1, self.ends().len());
 
         Ok(Self::with_offset_and_size(
-            slice(&self.ends(), slice_begin, bounded_slice_end)?,
-            slice(&self.values(), slice_begin, bounded_slice_end)?,
-            self.validity().slice(slice_begin, bounded_slice_end)?,
+            slice(&self.ends(), slice_begin, slice_end + 1)?,
+            slice(&self.values(), slice_begin, slice_end + 1)?,
+            self.validity().slice(slice_begin, slice_end + 1)?,
             stop - start,
             start,
         )?
@@ -69,20 +71,46 @@ impl SliceFn for RunEndArray {
 mod test {
     use vortex::array::PrimitiveArray;
     use vortex::compute::take;
-    use vortex::{IntoArrayVariant, ToArray};
+    use vortex::compute::unary::scalar_at;
+    use vortex::{Array, IntoArray, IntoArrayVariant, ToArray};
 
     use crate::RunEndArray;
 
-    #[test]
-    fn ree_take() {
-        let ree = RunEndArray::encode(
+    fn ree_array() -> Array {
+        RunEndArray::encode(
             PrimitiveArray::from(vec![1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).to_array(),
         )
-        .unwrap();
-        let taken = take(ree.array(), PrimitiveArray::from(vec![9, 8, 1, 3]).array()).unwrap();
+        .unwrap()
+        .into_array()
+    }
+
+    #[test]
+    fn ree_take() {
+        let taken = take(&ree_array(), PrimitiveArray::from(vec![9, 8, 1, 3]).array()).unwrap();
         assert_eq!(
             taken.into_primitive().unwrap().maybe_null_slice::<i32>(),
             &[5, 5, 1, 4]
         );
     }
+
+    #[test]
+    fn ree_take_end() {
+        let taken = take(&ree_array(), PrimitiveArray::from(vec![11]).array()).unwrap();
+        assert_eq!(
+            taken.into_primitive().unwrap().maybe_null_slice::<i32>(),
+            &[5]
+        );
+    }
+
+    #[test]
+    #[should_panic]
+    fn ree_take_out_of_bounds() {
+        take(&ree_array(), PrimitiveArray::from(vec![12]).array()).unwrap();
+    }
+
+    #[test]
+    fn ree_scalar_at_end() {
+        let scalar = scalar_at(&ree_array(), 11).unwrap();
+        assert_eq!(scalar, 5.into());
+    }
 }
diff --git a/encodings/runend/src/runend.rs b/encodings/runend/src/runend.rs
@@ -72,8 +72,13 @@ 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())
+        let searched_index =
+            search_sorted(&self.ends(), index + self.offset(), SearchSortedSide::Right)?.to_index();
+        Ok(if searched_index == self.ends().len() {
+            searched_index - 1
+        } else {
+            searched_index
+        })
     }
 
     pub fn encode(array: Array) -> VortexResult<Self> {