diff --git a/encodings/fastlanes/src/bitpacking/compute/search_sorted.rs b/encodings/fastlanes/src/bitpacking/compute/search_sorted.rs index d5097378e..07775fcdb 100644 --- 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 index 641cdcd99..7dc8c0730 100644 --- 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 index 03f04bf3a..3855285ee 100644 --- 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 usize>(self, f: F) -> Self { + match self { + Self::Found(i) => Self::Found(f(i)), + Self::NotFound(i) => Self::NotFound(f(i)), + } + } } pub trait SearchSortedFn {