fix

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

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/encodings/fastlanes/src/bitpacking/compute/search_sorted.rs b/encodings/fastlanes/src/bitpacking/compute/search_sorted.rs
@@ -20,7 +20,7 @@ impl SearchSortedFn for BitPackedArray {
             let unwrapped_value: $P = value.cast(self.dtype())?.try_into().unwrap();
                 if let Some(patches_array) = self.patches() {
                 if unwrapped_value as usize >= self.max_packed_value() {
-                    search_sorted(&patches_array, value.clone(), side)
+                    Ok(search_sorted(&patches_array, value.clone(), side)?.map(|i| i - self.offset()))
                 } else {
                     Ok(SearchSorted::search_sorted(&BitPackedSearch::new(self), &unwrapped_value, side))
                 }
@@ -127,13 +127,17 @@ mod test {
             )
             .unwrap()
             .into_array(),
-            3,
-            5,
+            2,
+            4,
         )
         .unwrap();
         assert_eq!(
-            search_sorted(&bitpacked, 4, SearchSortedSide::Left).unwrap(),
+            search_sorted(&bitpacked, 3, SearchSortedSide::Left).unwrap(),
             SearchResult::Found(0)
         );
+        assert_eq!(
+            search_sorted(&bitpacked, 4, SearchSortedSide::Left).unwrap(),
+            SearchResult::Found(1)
+        );
     }
 }
diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs
@@ -58,7 +58,12 @@ impl BitPackedArray {
         }
 
         let ptype: PType = (&dtype).try_into()?;
-        let expected_packed_size = (((length + 1023) / 1024) + if offset == 0 { 0 } else { 1 })
+        let expected_packed_size = (((length + 1023) / 1024)
+            + if offset > 0 && offset + length > 1024 {
+                1
+            } else {
+                0
+            })
             * (128 * bit_width / ptype.byte_width());
         if packed.len() != expected_packed_size {
             return Err(vortex_err!(

diff --git a/vortex-array/src/compute/search_sorted.rs b/vortex-array/src/compute/search_sorted.rs
@@ -33,6 +33,13 @@ impl SearchResult {
             Self::NotFound(i) => i,
         }
     }
+
+    pub fn map<F: FnOnce(usize) -> usize>(self, f: F) -> Self {
+        match self {
+            Self::Found(i) => Self::Found(f(i)),
+            Self::NotFound(i) => Self::NotFound(f(i)),
+        }
+    }
 }
 
 pub trait SearchSortedFn {