Fix Slice and SearchSorted on BitPackedArray

spiraldb · Jun 25, 2024 · be17537 · be17537
1 parent 2378cb6
commit be17537
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 6 deletions.
diff --git a/encodings/fastlanes/src/bitpacking/compute/search_sorted.rs b/encodings/fastlanes/src/bitpacking/compute/search_sorted.rs
@@ -35,6 +35,7 @@ impl SearchSortedFn for BitPackedArray {
 #[derive(Debug)]
 struct BitPackedSearch {
     packed: PrimitiveArray,
+    offset: usize,
     length: usize,
     bit_width: usize,
     min_patch_offset: Option<usize>,
@@ -44,6 +45,7 @@ impl BitPackedSearch {
     pub fn new(array: &BitPackedArray) -> Self {
         Self {
             packed: array.packed().into_primitive().unwrap(),
+            offset: array.offset(),
             length: array.len(),
             bit_width: array.bit_width(),
             min_patch_offset: array.patches().map(|p| {
@@ -64,8 +66,12 @@ impl<T: BitPacking + NativePType> IndexOrd<T> for BitPackedSearch {
         }
         // SAFETY: Used in search_sorted_by which ensures that idx is within bounds
         let val: T = unsafe {
-            unpack_single_primitive(self.packed.maybe_null_slice::<T>(), self.bit_width, idx)
-                .unwrap()
+            unpack_single_primitive(
+                self.packed.maybe_null_slice::<T>(),
+                self.bit_width,
+                idx + self.offset,
+            )
+            .unwrap()
         };
         val.partial_cmp(elem)
     }
@@ -81,6 +87,7 @@ impl Len for BitPackedSearch {
 mod test {
     use vortex::array::primitive::PrimitiveArray;
     use vortex::compute::search_sorted::{search_sorted, SearchResult, SearchSortedSide};
+    use vortex::compute::slice::slice;
     use vortex::IntoArray;
 
     use crate::BitPackedArray;
@@ -110,4 +117,23 @@ mod test {
             SearchResult::NotFound(0)
         );
     }
+
+    #[test]
+    fn search_sliced() {
+        let bitpacked = slice(
+            &BitPackedArray::encode(
+                &PrimitiveArray::from(vec![1u32, 2, 3, 4, 5]).into_array(),
+                2,
+            )
+            .unwrap()
+            .into_array(),
+            3,
+            5,
+        )
+        .unwrap();
+        assert_eq!(
+            search_sorted(&bitpacked, 4, SearchSortedSide::Left).unwrap(),
+            SearchResult::Found(0)
+        );
+    }
 }
diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs
@@ -12,8 +12,8 @@ impl SliceFn for BitPackedArray {
         let block_start = max(0, start - offset);
         let block_stop = ((stop + 1023) / 1024) * 1024;
 
-        let encoded_start = (block_start / 8) * self.bit_width();
-        let encoded_stop = (block_stop / 8) * self.bit_width();
+        let encoded_start = (block_start / 8) * self.bit_width() / self.ptype().byte_width();
+        let encoded_stop = (block_stop / 8) * self.bit_width() / self.ptype().byte_width();
         Self::try_new_from_offset(
             slice(&self.packed(), encoded_start, encoded_stop)?,
             self.validity().slice(start, stop)?,
@@ -25,3 +25,38 @@ impl SliceFn for BitPackedArray {
         .map(|a| a.into_array())
     }
 }
+
+#[cfg(test)]
+mod test {
+    use vortex::array::primitive::PrimitiveArray;
+    use vortex::compute::slice::slice;
+    use vortex::{ArrayTrait, IntoArray};
+
+    use crate::BitPackedArray;
+
+    #[test]
+    pub fn slice_block() {
+        let arr = BitPackedArray::encode(
+            &PrimitiveArray::from((0u32..2048).map(|v| v % 64).collect::<Vec<_>>()).into_array(),
+            6,
+        )
+        .unwrap()
+        .into_array();
+        let sliced = BitPackedArray::try_from(slice(&arr, 1024, 2048).unwrap()).unwrap();
+        assert_eq!(sliced.offset(), 0);
+        assert_eq!(sliced.len(), 1024);
+    }
+
+    #[test]
+    pub fn slice_within_block() {
+        let arr = BitPackedArray::encode(
+            &PrimitiveArray::from((0u32..2048).map(|v| v % 64).collect::<Vec<_>>()).into_array(),
+            6,
+        )
+        .unwrap()
+        .into_array();
+        let sliced = BitPackedArray::try_from(slice(&arr, 512, 1434).unwrap()).unwrap();
+        assert_eq!(sliced.offset(), 512);
+        assert_eq!(sliced.len(), 922);
+    }
+}
diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs
@@ -50,10 +50,16 @@ impl BitPackedArray {
         if bit_width > 64 {
             vortex_bail!("Unsupported bit width {}", bit_width);
         }
+        if offset > 1023 {
+            vortex_bail!(
+                "Offset must be less than full block, i.e. 1024, got {}",
+                offset
+            );
+        }
 
         let ptype: PType = (&dtype).try_into()?;
-        let expected_packed_size =
-            ((length + 1023) / 1024) * (128 * bit_width / ptype.byte_width());
+        let expected_packed_size = (((length + 1023) / 1024) + if offset == 0 { 0 } else { 1 })
+            * (128 * bit_width / ptype.byte_width());
         if packed.len() != expected_packed_size {
             return Err(vortex_err!(
                 "Expected {} packed bytes, got {}",