wip

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

encodings/fsst/src/tests.rs

@@ -12,6 +12,7 @@ use vortex_array::compute::take;
 use vortex_buffer::buffer;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability;
+use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::FSSTVTable;
@@ -39,7 +40,7 @@ pub(crate) fn build_fsst_array() -> ArrayRef {
 }
 
 #[test]
-fn test_fsst_array_ops() {
+fn test_fsst_array_ops() -> VortexResult<()> {
     // first test the scalar_at values
     let fsst_array = build_fsst_array();
     assert_nth_scalar!(
@@ -75,7 +76,7 @@ fn test_fsst_array_ops() {
 
     // test take
     let indices = buffer![0, 2].into_array();
-    let fsst_taken = take(&fsst_array, &indices).unwrap();
+    let fsst_taken = take(&fsst_array, &indices)?;
     assert!(fsst_taken.is::<FSSTVTable>());
     assert_eq!(fsst_taken.len(), 2);
     assert_nth_scalar!(
@@ -92,7 +93,7 @@ fn test_fsst_array_ops() {
     // test filter
     let mask = Mask::from_iter([false, true, true]);
 
-    let fsst_filtered = filter(&fsst_array, &mask).unwrap();
+    let fsst_filtered = filter(&fsst_array, &mask)?;
     assert!(fsst_filtered.is::<FSSTVTable>());
     assert_eq!(fsst_filtered.len(), 2);
     assert_nth_scalar!(
@@ -102,7 +103,9 @@ fn test_fsst_array_ops() {
     );
 
     // test to_canonical
-    let canonical_array = fsst_array.to_varbinview().unwrap().into_array();
+    let canonical_array = fsst_array.to_varbinview()?.into_array();
 
     assert_arrays_eq!(fsst_array.to_array(), canonical_array);
+
+    Ok(())
 }

encodings/pco/src/test.rs

@@ -20,6 +20,7 @@ use vortex_buffer::BufferMut;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability;
 use vortex_dtype::PType;
+use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::PcoArray;
@@ -32,7 +33,7 @@ macro_rules! assert_nth_scalar {
 }
 
 #[test]
-fn test_compress_decompress() {
+fn test_compress_decompress() -> VortexResult<()> {
     let data: Vec<i32> = (0..200).collect();
     let array = PrimitiveArray::from_iter(data.clone());
     let compressed = PcoArray::from_primitive(&array, 3, 0).unwrap();
@@ -48,15 +49,17 @@ fn test_compress_decompress() {
     for i in 0_i32..5 {
         assert_nth_scalar!(slice, i as usize, 100 + i);
     }
-    let primitive = slice.to_primitive().unwrap();
+    let primitive = slice.to_primitive()?;
     assert_arrays_eq!(
         primitive,
         PrimitiveArray::from_iter([100, 101, 102, 103, 104])
     );
 
     let slice = compressed.slice(200..200);
-    let primitive = slice.to_primitive().unwrap();
+    let primitive = slice.to_primitive()?;
     assert_arrays_eq!(primitive, PrimitiveArray::from_iter(Vec::<i32>::new()));
+
+    Ok(())
 }
 
 #[test]
@@ -81,7 +84,7 @@ fn test_empty() {
 }
 
 #[test]
-fn test_validity_and_multiple_chunks_and_pages() {
+fn test_validity_and_multiple_chunks_and_pages() -> VortexResult<()> {
     let data: Vec<i32> = (0..200).collect();
     let mut validity: Vec<bool> = vec![true; 200];
     validity[7..15].fill(false);
@@ -115,15 +118,17 @@ fn test_validity_and_multiple_chunks_and_pages() {
     let slice = compressed.slice(100..103);
     assert_nth_scalar!(slice, 0, 100);
     assert_nth_scalar!(slice, 2, 102);
-    let primitive = slice.to_primitive().unwrap();
+    let primitive = slice.to_primitive()?;
     assert_eq!(
         primitive.validity(),
         &Validity::Array(BoolArray::from_iter(vec![true, false, true]).to_array())
     );
+
+    Ok(())
 }
 
 #[test]
-fn test_validity_vtable() {
+fn test_validity_vtable() -> VortexResult<()> {
     let data: Vec<i32> = (0..5).collect();
     let mask_bools = vec![false, true, true, false, true];
     let array = PrimitiveArray::new(
@@ -132,13 +137,15 @@ fn test_validity_vtable() {
     );
     let compressed = PcoArray::from_primitive(&array, 3, 0).unwrap();
     assert_eq!(
-        compressed.validity_mask().unwrap(),
+        compressed.validity_mask()?,
         Mask::from_iter(mask_bools)
     );
     assert_eq!(
-        compressed.slice(1..4).validity_mask().unwrap(),
+        compressed.slice(1..4).validity_mask()?,
         Mask::from_iter(vec![true, true, false])
     );
+
+    Ok(())
 }
 
 #[test]

encodings/zstd/src/test.rs

@@ -12,6 +12,7 @@ use vortex_array::vtable::ValidityHelper;
 use vortex_buffer::Buffer;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability;
+use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
 use crate::ZstdArray;
@@ -23,7 +24,7 @@ macro_rules! assert_nth_scalar {
 }
 
 #[test]
-fn test_zstd_compress_decompress() {
+fn test_zstd_compress_decompress() -> VortexResult<()> {
     let data: Vec<i32> = (0..200).collect();
     let array = PrimitiveArray::from_iter(data.clone());
 
@@ -33,27 +34,29 @@ fn test_zstd_compress_decompress() {
     assert!(compressed.dictionary.is_none());
 
     // check full decompression works
-    let decompressed = compressed.decompress().unwrap().to_primitive().unwrap();
+    let decompressed = compressed.decompress()?.to_primitive()?;
     assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data));
 
     // check slicing works
     let slice = compressed.slice(100..105);
     for i in 0_i32..5 {
         assert_nth_scalar!(slice, i as usize, 100 + i);
     }
-    let primitive = slice.to_primitive().unwrap();
+    let primitive = slice.to_primitive()?;
     assert_arrays_eq!(
         primitive,
         PrimitiveArray::from_iter([100, 101, 102, 103, 104])
     );
 
     let slice = compressed.slice(200..200);
-    let primitive = slice.to_primitive().unwrap();
+    let primitive = slice.to_primitive()?;
     assert_arrays_eq!(primitive, PrimitiveArray::from_iter(Vec::<i32>::new()));
+
+    Ok(())
 }
 
 #[test]
-fn test_zstd_empty() {
+fn test_zstd_empty() -> VortexResult<()> {
     let data: Vec<i32> = vec![];
     let array = PrimitiveArray::new(
         data.iter().cloned().collect::<Buffer<_>>(),
@@ -62,12 +65,14 @@ fn test_zstd_empty() {
 
     let compressed = ZstdArray::from_primitive(&array, 3, 100).unwrap();
 
-    let primitive = compressed.to_primitive().unwrap();
+    let primitive = compressed.to_primitive()?;
     assert_arrays_eq!(primitive, PrimitiveArray::from_iter(data));
+
+    Ok(())
 }
 
 #[test]
-fn test_zstd_with_validity_and_multi_frame() {
+fn test_zstd_with_validity_and_multi_frame() -> VortexResult<()> {
     let data: Vec<i32> = (0..200).collect();
     let mut validity: Vec<bool> = vec![false; 200];
     validity[3] = true;
@@ -84,15 +89,15 @@ fn test_zstd_with_validity_and_multi_frame() {
     assert_nth_scalar!(compressed, 10, None::<i32>);
     assert_nth_scalar!(compressed, 177, 177);
 
-    let decompressed = compressed.decompress().unwrap().to_primitive().unwrap();
+    let decompressed = compressed.decompress()?.to_primitive()?;
     let decompressed_values = decompressed.as_slice::<i32>();
     assert_eq!(decompressed_values[3], 3);
     assert_eq!(decompressed_values[177], 177);
     assert_eq!(decompressed.validity(), array.validity());
 
     // check slicing works
     let slice = compressed.slice(176..179);
-    let primitive = slice.to_primitive().unwrap();
+    let primitive = slice.to_primitive()?;
     assert_eq!(
         TryInto::<i32>::try_into(primitive.scalar_at(1).as_ref())
             .ok()
@@ -103,10 +108,12 @@ fn test_zstd_with_validity_and_multi_frame() {
         primitive.validity(),
         &Validity::Array(BoolArray::from_iter(vec![false, true, false]).to_array())
     );
+
+    Ok(())
 }
 
 #[test]
-fn test_zstd_with_dict() {
+fn test_zstd_with_dict() -> VortexResult<()> {
     let data: Vec<i32> = (0..200).collect();
     let array = PrimitiveArray::new(
         data.iter().cloned().collect::<Buffer<_>>(),
@@ -118,32 +125,36 @@ fn test_zstd_with_dict() {
     assert_nth_scalar!(compressed, 0, 0);
     assert_nth_scalar!(compressed, 199, 199);
 
-    let decompressed = compressed.decompress().unwrap().to_primitive().unwrap();
+    let decompressed = compressed.decompress()?.to_primitive()?;
     assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data));
     assert_eq!(decompressed.validity(), array.validity());
 
     // check slicing works
     let slice = compressed.slice(176..179);
-    let primitive = slice.to_primitive().unwrap();
+    let primitive = slice.to_primitive()?;
     assert_arrays_eq!(primitive, PrimitiveArray::from_iter([176, 177, 178]));
+
+    Ok(())
 }
 
 #[test]
-fn test_validity_vtable() {
+fn test_validity_vtable() -> VortexResult<()> {
     let mask_bools = vec![false, true, true, false, true];
     let array = PrimitiveArray::new(
         (0..5).collect::<Buffer<_>>(),
         Validity::Array(BoolArray::from_iter(mask_bools.clone()).to_array()),
     );
     let compressed = ZstdArray::from_primitive(&array, 3, 0).unwrap();
     assert_eq!(
-        compressed.validity_mask().unwrap(),
+        compressed.validity_mask()?,
         Mask::from_iter(mask_bools)
     );
     assert_eq!(
-        compressed.slice(1..4).validity_mask().unwrap(),
+        compressed.slice(1..4).validity_mask()?,
         Mask::from_iter(vec![true, true, false])
     );
+
+    Ok(())
 }
 
 #[test]
@@ -172,7 +183,7 @@ fn test_zstd_var_bin_view() {
 }
 
 #[test]
-fn test_zstd_decompress_var_bin_view() {
+fn test_zstd_decompress_var_bin_view() -> VortexResult<()> {
     let data: [Option<&'static [u8]>; 5] = [
         Some(b"foo"),
         Some(b"bar"),
@@ -189,12 +200,14 @@ fn test_zstd_decompress_var_bin_view() {
     assert_nth_scalar!(compressed, 2, None::<String>);
     assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet");
     assert_nth_scalar!(compressed, 4, "baz");
-    let decompressed = compressed.decompress().unwrap().to_varbinview().unwrap();
+    let decompressed = compressed.decompress()?.to_varbinview()?;
     assert_nth_scalar!(decompressed, 0, "foo");
     assert_nth_scalar!(decompressed, 1, "bar");
     assert_nth_scalar!(decompressed, 2, None::<String>);
     assert_nth_scalar!(decompressed, 3, "Lorem ipsum dolor sit amet");
     assert_nth_scalar!(decompressed, 4, "baz");
+
+    Ok(())
 }
 
 #[test]

vortex-array/src/arrays/datetime/test.rs

@@ -5,6 +5,7 @@ use rstest::rstest;
 use vortex_buffer::buffer;
 use vortex_dtype::datetime::TemporalMetadata;
 use vortex_dtype::datetime::TimeUnit;
+use vortex_error::VortexResult;
 
 use crate::IntoArray;
 use crate::ToCanonical;
@@ -15,6 +16,18 @@ use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 
 macro_rules! test_temporal_roundtrip {
+    ($prim:ty, $constructor:expr, $unit:expr) => {{
+        let array = buffer![100 as $prim].into_array();
+        let temporal: TemporalArray = $constructor(array, $unit);
+        let prims = temporal.temporal_values().to_primitive()?;
+
+        assert_eq!(prims.as_slice::<$prim>(), vec![100 as $prim].as_slice(),);
+        assert_eq!(temporal.temporal_metadata().time_unit(), $unit);
+        Ok(())
+    }};
+}
+
+macro_rules! test_temporal_roundtrip_panic {
     ($prim:ty, $constructor:expr, $unit:expr) => {{
         let array = buffer![100 as $prim].into_array();
         let temporal: TemporalArray = $constructor(array, $unit);
@@ -28,8 +41,8 @@ macro_rules! test_temporal_roundtrip {
 macro_rules! test_success_case {
     ($name:ident, $prim:ty, $constructor:expr, $unit:expr) => {
         #[test]
-        fn $name() {
-            test_temporal_roundtrip!($prim, $constructor, $unit);
+        fn $name() -> VortexResult<()> {
+            test_temporal_roundtrip!($prim, $constructor, $unit)
         }
     };
 }
@@ -39,7 +52,7 @@ macro_rules! test_fail_case {
         #[test]
         #[should_panic]
         fn $name() {
-            test_temporal_roundtrip!($prim, $constructor, $unit)
+            test_temporal_roundtrip_panic!($prim, $constructor, $unit)
         }
     };
 }
@@ -132,7 +145,7 @@ test_fail_case!(
 
 // We test Timestamp explicitly to avoid the macro getting too complex.
 #[test]
-fn test_timestamp() {
+fn test_timestamp() -> VortexResult<()> {
     let ts = buffer![100i64].into_array();
     let ts_array = ts.into_array();
 
@@ -146,14 +159,15 @@ fn test_timestamp() {
             let temporal_array =
                 TemporalArray::new_timestamp(ts_array.to_array(), unit, tz.clone());
 
-            let values = temporal_array.temporal_values().to_primitive().unwrap();
+            let values = temporal_array.temporal_values().to_primitive()?;
             assert_eq!(values.as_slice::<i64>(), vec![100i64].as_slice());
             assert_eq!(
                 temporal_array.temporal_metadata(),
                 &TemporalMetadata::Timestamp(unit, tz)
             );
         }
     }
+    Ok(())
 }
 
 #[test]
@@ -170,7 +184,7 @@ fn test_timestamp_fails_i32() {
 #[case(Validity::AllValid)]
 #[case(Validity::AllInvalid)]
 #[case(Validity::from_iter([true, false, true]))]
-fn test_validity_preservation(#[case] validity: Validity) {
+fn test_validity_preservation(#[case] validity: Validity) -> VortexResult<()> {
     let milliseconds = PrimitiveArray::new(
         buffer![
             86_400i64,            // element with only day component
@@ -188,9 +202,9 @@ fn test_validity_preservation(#[case] validity: Validity) {
     assert_eq!(
         temporal_array
             .temporal_values()
-            .to_primitive()
-            .unwrap()
+            .to_primitive()?
             .validity(),
         &validity
     );
+    Ok(())
 }