From 27c864f975aadc708b2977350f7485ee50f7fdfa Mon Sep 17 00:00:00 2001
From: Robert Kruszewski <github@robertk.io>
Date: Tue, 25 Jun 2024 13:38:04 +0100
Subject: [PATCH] more tests

---
 .../fastlanes/src/bitpacking/compute/slice.rs | 53 ++++++++++++++++++-
 encodings/fastlanes/src/bitpacking/mod.rs     | 49 +----------------
 2 files changed, 53 insertions(+), 49 deletions(-)

diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs
index 1fa7bdbc68..9b37cc1fea 100644
--- a/encodings/fastlanes/src/bitpacking/compute/slice.rs
+++ b/encodings/fastlanes/src/bitpacking/compute/slice.rs
@@ -30,6 +30,7 @@ impl SliceFn for BitPackedArray {
 mod test {
     use vortex::array::primitive::PrimitiveArray;
     use vortex::compute::slice::slice;
+    use vortex::compute::unary::scalar_at::scalar_at;
     use vortex::{ArrayTrait, IntoArray};
 
     use crate::BitPackedArray;
@@ -37,12 +38,17 @@ mod test {
     #[test]
     pub fn slice_block() {
         let arr = BitPackedArray::encode(
-            &PrimitiveArray::from((0u32..2048).map(|v| v % 64).collect::<Vec<_>>()).into_array(),
+            PrimitiveArray::from((0u32..2048).map(|v| v % 64).collect::<Vec<_>>()).array(),
             6,
         )
         .unwrap()
         .into_array();
         let sliced = BitPackedArray::try_from(slice(&arr, 1024, 2048).unwrap()).unwrap();
+        assert_eq!(scalar_at(sliced.array(), 0).unwrap(), (1024u32 % 64).into());
+        assert_eq!(
+            scalar_at(sliced.array(), 1023).unwrap(),
+            (2047u32 % 64).into()
+        );
         assert_eq!(sliced.offset(), 0);
         assert_eq!(sliced.len(), 1024);
     }
@@ -50,13 +56,56 @@ mod test {
     #[test]
     pub fn slice_within_block() {
         let arr = BitPackedArray::encode(
-            &PrimitiveArray::from((0u32..2048).map(|v| v % 64).collect::<Vec<_>>()).into_array(),
+            PrimitiveArray::from((0u32..2048).map(|v| v % 64).collect::<Vec<_>>()).array(),
             6,
         )
         .unwrap()
         .into_array();
         let sliced = BitPackedArray::try_from(slice(&arr, 512, 1434).unwrap()).unwrap();
+        assert_eq!(scalar_at(sliced.array(), 0).unwrap(), (512u32 % 64).into());
+        assert_eq!(
+            scalar_at(sliced.array(), 921).unwrap(),
+            (1433u32 % 64).into()
+        );
         assert_eq!(sliced.offset(), 512);
         assert_eq!(sliced.len(), 922);
     }
+
+    #[test]
+    fn slice_within_block_u8s() {
+        let packed = BitPackedArray::encode(
+            PrimitiveArray::from((0..10_000).map(|i| (i % 63) as u8).collect::<Vec<_>>()).array(),
+            7,
+        )
+        .unwrap();
+
+        let compressed = slice(packed.array(), 768, 9999).unwrap();
+        assert_eq!(
+            scalar_at(&compressed, 0).unwrap(),
+            ((768 % 63) as u8).into()
+        );
+        assert_eq!(
+            scalar_at(&compressed, compressed.len() - 1).unwrap(),
+            ((9998 % 63) as u8).into()
+        );
+    }
+
+    #[test]
+    fn slice_block_boundary_u8s() {
+        let packed = BitPackedArray::encode(
+            PrimitiveArray::from((0..10_000).map(|i| (i % 63) as u8).collect::<Vec<_>>()).array(),
+            7,
+        )
+        .unwrap();
+
+        let compressed = slice(packed.array(), 7168, 9216).unwrap();
+        assert_eq!(
+            scalar_at(&compressed, 0).unwrap(),
+            ((7168 % 63) as u8).into()
+        );
+        assert_eq!(
+            scalar_at(&compressed, compressed.len() - 1).unwrap(),
+            ((9215 % 63) as u8).into()
+        );
+    }
 }
diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs
index 7dc8c07301..2c02471e6b 100644
--- a/encodings/fastlanes/src/bitpacking/mod.rs
+++ b/encodings/fastlanes/src/bitpacking/mod.rs
@@ -58,13 +58,8 @@ impl BitPackedArray {
         }
 
         let ptype: PType = (&dtype).try_into()?;
-        let expected_packed_size = (((length + 1023) / 1024)
-            + if offset > 0 && offset + length > 1024 {
-                1
-            } else {
-                0
-            })
-            * (128 * bit_width / ptype.byte_width());
+        let expected_packed_size =
+            ((length + offset + 1023) / 1024) * (128 * bit_width / ptype.byte_width());
         if packed.len() != expected_packed_size {
             return Err(vortex_err!(
                 "Expected {} packed bytes, got {}",
@@ -191,50 +186,10 @@ impl ArrayTrait for BitPackedArray {
 #[cfg(test)]
 mod test {
     use vortex::array::primitive::PrimitiveArray;
-    use vortex::compute::slice::slice;
-    use vortex::compute::unary::scalar_at::scalar_at;
     use vortex::{IntoArray, IntoCanonical};
 
     use crate::BitPackedArray;
 
-    #[test]
-    fn slice_within_block() {
-        let packed = BitPackedArray::encode(
-            PrimitiveArray::from((0..10_000).map(|i| (i % 63) as u8).collect::<Vec<_>>()).array(),
-            7,
-        )
-        .unwrap();
-
-        let compressed = slice(packed.array(), 768, 9999).unwrap();
-        assert_eq!(
-            scalar_at(&compressed, 0).unwrap(),
-            ((768 % 63) as u8).into()
-        );
-        assert_eq!(
-            scalar_at(&compressed, compressed.len() - 1).unwrap(),
-            ((9998 % 63) as u8).into()
-        );
-    }
-
-    #[test]
-    fn slice_block_boundary() {
-        let packed = BitPackedArray::encode(
-            PrimitiveArray::from((0..10_000).map(|i| (i % 63) as u8).collect::<Vec<_>>()).array(),
-            7,
-        )
-        .unwrap();
-
-        let compressed = slice(packed.array(), 7168, 9216).unwrap();
-        assert_eq!(
-            scalar_at(&compressed, 0).unwrap(),
-            ((7168 % 63) as u8).into()
-        );
-        assert_eq!(
-            scalar_at(&compressed, compressed.len() - 1).unwrap(),
-            ((9215 % 63) as u8).into()
-        );
-    }
-
     #[test]
     fn test_encode() {
         let values = vec![Some(1), None, Some(1), None, Some(1), None, Some(u64::MAX)];