chore[encodings]: use vx array assert eq macro (#5076)

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

Author: joseph-isaacs

Cargo.lock

@@ -154,30 +154,26 @@ mod tests {
         let array = PrimitiveArray::new(buffer![1.234f32; 1025], Validity::NonNullable);
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_none());
-        assert_eq!(
-            encoded.encoded().to_primitive().as_slice::<i32>(),
-            vec![1234; 1025]
-        );
+        let expected_encoded = PrimitiveArray::from_iter(vec![1234i32; 1025]);
+        assert_arrays_eq!(encoded.encoded(), expected_encoded);
         assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 });
 
         let decoded = decompress(&encoded);
-        assert_eq!(array.as_slice::<f32>(), decoded.as_slice::<f32>());
+        assert_arrays_eq!(decoded, array);
     }
 
     #[test]
     fn test_nullable_compress() {
         let array = PrimitiveArray::from_option_iter([None, Some(1.234f32), None]);
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_none());
-        assert_eq!(
-            encoded.encoded().to_primitive().as_slice::<i32>(),
-            vec![0, 1234, 0]
-        );
+        let expected_encoded = PrimitiveArray::from_option_iter([None, Some(1234i32), None]);
+        assert_arrays_eq!(encoded.encoded(), expected_encoded);
         assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 });
 
         let decoded = decompress(&encoded);
-        let expected = vec![0f32, 1.234f32, 0f32];
-        assert_eq!(decoded.as_slice::<f32>(), expected.as_slice());
+        let expected = PrimitiveArray::from_option_iter(vec![None, Some(1.234f32), None]);
+        assert_arrays_eq!(decoded, expected);
     }
 
     #[test]
@@ -187,14 +183,13 @@ mod tests {
         let array = PrimitiveArray::new(values.clone(), Validity::NonNullable);
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_some());
-        assert_eq!(
-            encoded.encoded().to_primitive().as_slice::<i64>(),
-            vec![1234i64, 2718, 1234, 4000]
-        );
+        let expected_encoded = PrimitiveArray::from_iter(vec![1234i64, 2718, 1234, 4000]);
+        assert_arrays_eq!(encoded.encoded(), expected_encoded);
         assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
 
         let decoded = decompress(&encoded);
-        assert_eq!(values.as_slice(), decoded.as_slice::<f64>());
+        let expected_decoded = PrimitiveArray::new(values, Validity::NonNullable);
+        assert_arrays_eq!(decoded, expected_decoded);
     }
 
     #[test]
@@ -204,10 +199,9 @@ mod tests {
         let array = PrimitiveArray::new(values, Validity::from_iter([true, true, false, true]));
         let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_none());
-        assert_eq!(
-            encoded.encoded().to_primitive().as_slice::<i64>(),
-            vec![1234i64, 2718, 1234, 4000]
-        );
+        let expected_encoded =
+            PrimitiveArray::from_option_iter(buffer![Some(1234i64), Some(2718), None, Some(4000)]);
+        assert_arrays_eq!(encoded.encoded(), expected_encoded);
         assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
 
         let decoded = decompress(&encoded);
@@ -239,7 +233,7 @@ mod tests {
         let original = PrimitiveArray::from_iter([195.26274f32, 195.27837, -48.815685]);
         let alp_arr = alp_encode(&original, None).unwrap();
         let decompressed = alp_arr.to_primitive();
-        assert_eq!(original.as_slice::<f32>(), decompressed.as_slice::<f32>());
+        assert_arrays_eq!(decompressed, original);
     }
 
     #[test]
@@ -289,13 +283,16 @@ mod tests {
         let patches = encoded.patches().unwrap();
 
         let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive();
-        assert_eq!(chunk_offsets.as_slice::<u64>(), &[0, 1, 3]);
+        let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 3]);
+        assert_arrays_eq!(chunk_offsets, expected_offsets);
 
         let patch_indices = patches.indices().to_primitive();
-        assert_eq!(patch_indices.as_slice::<u64>(), &[1023, 1024, 1025]);
+        let expected_indices = PrimitiveArray::from_iter(vec![1023u64, 1024, 1025]);
+        assert_arrays_eq!(patch_indices, expected_indices);
 
         let patch_values = patches.values().to_primitive();
-        assert_eq!(patch_values.as_slice::<f64>(), &[PI, E, PI]);
+        let expected_values = PrimitiveArray::from_iter(vec![PI, E, PI]);
+        assert_arrays_eq!(patch_values, expected_values);
     }
 
     #[test]
@@ -309,13 +306,16 @@ mod tests {
         let patches = encoded.patches().unwrap();
 
         let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive();
-        assert_eq!(chunk_offsets.as_slice::<u64>(), &[0, 1, 1]);
+        let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]);
+        assert_arrays_eq!(chunk_offsets, expected_offsets);
 
         let patch_indices = patches.indices().to_primitive();
-        assert_eq!(patch_indices.as_slice::<u64>(), &[0, 2048]);
+        let expected_indices = PrimitiveArray::from_iter(vec![0u64, 2048]);
+        assert_arrays_eq!(patch_indices, expected_indices);
 
         let patch_values = patches.values().to_primitive();
-        assert_eq!(patch_values.as_slice::<f64>(), &[PI, E]);
+        let expected_values = PrimitiveArray::from_iter(vec![PI, E]);
+        assert_arrays_eq!(patch_values, expected_values);
     }
 
     #[test]
@@ -328,13 +328,16 @@ mod tests {
         let patches = encoded.patches().unwrap();
 
         let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive();
-        assert_eq!(chunk_offsets.as_slice::<u64>(), &[0, 1, 1]);
+        let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]);
+        assert_arrays_eq!(chunk_offsets, expected_offsets);
 
         let patch_indices = patches.indices().to_primitive();
-        assert_eq!(patch_indices.as_slice::<u64>(), &[0]);
+        let expected_indices = PrimitiveArray::from_iter(vec![0u64]);
+        assert_arrays_eq!(patch_indices, expected_indices);
 
         let patch_values = patches.values().to_primitive();
-        assert_eq!(patch_values.as_slice::<f64>(), &[PI]);
+        let expected_values = PrimitiveArray::from_iter(vec![PI]);
+        assert_arrays_eq!(patch_values, expected_values);
     }
 
     #[test]
@@ -348,12 +351,15 @@ mod tests {
         let patches = encoded.patches().unwrap();
 
         let chunk_offsets = patches.chunk_offsets().clone().unwrap().to_primitive();
-        assert_eq!(chunk_offsets.as_slice::<u64>(), &[0]);
+        let expected_offsets = PrimitiveArray::from_iter(vec![0u64]);
+        assert_arrays_eq!(chunk_offsets, expected_offsets);
 
         let patch_indices = patches.indices().to_primitive();
-        assert_eq!(patch_indices.as_slice::<u64>(), &[0, 100]);
+        let expected_indices = PrimitiveArray::from_iter(vec![0u64, 100]);
+        assert_arrays_eq!(patch_indices, expected_indices);
 
         let patch_values = patches.values().to_primitive();
-        assert_eq!(patch_values.as_slice::<f64>(), &[PI, E]);
+        let expected_values = PrimitiveArray::from_iter(vec![PI, E]);
+        assert_arrays_eq!(patch_values, expected_values);
     }
 }

encodings/alp/src/alp/compress.rs

@@ -45,10 +45,10 @@ register_kernel!(CastKernelAdapter(ALPVTable).lift());
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
-    use vortex_array::IntoArray;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::compute::cast;
     use vortex_array::compute::conformance::cast::test_cast_conformance;
+    use vortex_array::{IntoArray, assert_arrays_eq};
     use vortex_buffer::buffer;
     use vortex_dtype::{DType, Nullability, PType};
 
@@ -98,7 +98,7 @@ mod tests {
         );
 
         let decoded = casted.to_canonical().into_primitive();
-        assert_eq!(decoded.as_slice::<i32>(), &[1i32, 2, 3, 4]);
+        assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i32, 2, 3, 4]));
     }
 
     #[rstest]

encodings/alp/src/alp/compute/cast.rs

@@ -265,8 +265,8 @@ impl CanonicalVTable<ALPRDVTable> for ALPRDVTable {
 #[cfg(test)]
 mod test {
     use rstest::rstest;
-    use vortex_array::ToCanonical;
     use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::{ToCanonical, assert_arrays_eq};
 
     use crate::{ALPRDFloat, alp_rd};
 
@@ -294,7 +294,6 @@ mod test {
 
         let decoded = rd_array.to_primitive();
 
-        let maybe_null_reals: Vec<T> = reals.into_iter().map(|v| v.unwrap_or_default()).collect();
-        assert_eq!(decoded.as_slice::<T>(), &maybe_null_reals);
+        assert_arrays_eq!(decoded, PrimitiveArray::from_option_iter(reals));
     }
 }

encodings/alp/src/alp_rd/array.rs

@@ -37,7 +37,7 @@ mod test {
     use vortex_array::compute::conformance::filter::test_filter_conformance;
     use vortex_array::compute::filter;
     use vortex_array::validity::Validity;
-    use vortex_array::{IntoArray, ToCanonical};
+    use vortex_array::{IntoArray, ToCanonical, assert_arrays_eq};
     use vortex_buffer::buffer;
     use vortex_mask::Mask;
 
@@ -57,7 +57,7 @@ mod test {
         let filtered = filter(encoded.as_ref(), &Mask::from_iter([true, false, true]))
             .unwrap()
             .to_primitive();
-        assert_eq!(filtered.as_slice::<T>(), &[a, outlier]);
+        assert_arrays_eq!(filtered, PrimitiveArray::from_iter([a, outlier]));
     }
 
     #[rstest]

encodings/alp/src/alp_rd/compute/filter.rs

@@ -48,10 +48,10 @@ register_kernel!(TakeKernelAdapter(ALPRDVTable).lift());
 #[cfg(test)]
 mod test {
     use rstest::rstest;
-    use vortex_array::ToCanonical;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::compute::conformance::take::test_take_conformance;
     use vortex_array::compute::take;
+    use vortex_array::{ToCanonical, assert_arrays_eq};
 
     use crate::{ALPRDFloat, RDEncoder};
 
@@ -78,7 +78,7 @@ mod test {
             .unwrap()
             .to_primitive();
 
-        assert_eq!(taken.as_slice::<T>(), &[a, outlier]);
+        assert_arrays_eq!(taken, PrimitiveArray::from_iter([a, outlier]));
     }
 
     #[rstest]
@@ -104,9 +104,10 @@ mod test {
         .unwrap()
         .to_primitive();
 
-        assert_eq!(taken.as_slice::<T>()[0], a);
-        assert_eq!(taken.as_slice::<T>()[1], outlier);
-        assert!(!taken.validity_mask().value(2));
+        assert_arrays_eq!(
+            taken,
+            PrimitiveArray::from_option_iter([Some(a), Some(outlier), None])
+        );
     }
 
     #[rstest]

encodings/alp/src/alp_rd/compute/take.rs

@@ -26,5 +26,6 @@ vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true }
 
 [dev-dependencies]
+itertools = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { workspace = true, features = ["test-harness"] }

encodings/bytebool/Cargo.toml

@@ -74,6 +74,7 @@ register_kernel!(TakeKernelAdapter(ByteBoolVTable).lift());
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
+    use vortex_array::assert_arrays_eq;
     use vortex_array::compute::conformance::filter::test_filter_conformance;
     use vortex_array::compute::conformance::mask::test_mask_conformance;
     use vortex_array::compute::conformance::take::test_take_conformance;
@@ -87,18 +88,9 @@ mod tests {
         let vortex_arr = ByteBoolArray::from(original);
 
         let sliced_arr = vortex_arr.slice(1..4);
-        let sliced_arr = sliced_arr.as_::<ByteBoolVTable>();
 
-        let s = sliced_arr.scalar_at(0);
-        assert_eq!(s.as_bool().value(), Some(true));
-
-        let s = sliced_arr.scalar_at(1);
-        assert!(!sliced_arr.is_valid(1));
-        assert!(s.is_null());
-        assert_eq!(s.as_bool().value(), None);
-
-        let s = sliced_arr.scalar_at(2);
-        assert_eq!(s.as_bool().value(), Some(false));
+        let expected = ByteBoolArray::from(vec![Some(true), None, Some(false)]);
+        assert_arrays_eq!(sliced_arr, expected.to_array());
     }
 
     #[test]
@@ -108,11 +100,8 @@ mod tests {
 
         let arr = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
 
-        for i in 0..arr.len() {
-            let s = arr.scalar_at(i);
-            assert!(s.is_valid());
-            assert_eq!(s.as_bool().value(), Some(true));
-        }
+        let expected = ByteBoolArray::from(vec![true; 5]);
+        assert_arrays_eq!(arr, expected.to_array());
     }
 
     #[test]
@@ -122,11 +111,8 @@ mod tests {
 
         let arr = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
 
-        for i in 0..arr.len() {
-            let s = arr.scalar_at(i);
-            assert!(s.is_valid());
-            assert_eq!(s.as_bool().value(), Some(false));
-        }
+        let expected = ByteBoolArray::from(vec![false; 5]);
+        assert_arrays_eq!(arr, expected.to_array());
     }
 
     #[test]
@@ -136,18 +122,9 @@ mod tests {
 
         let arr = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
 
-        for i in 0..3 {
-            let s = arr.scalar_at(i);
-            assert!(s.is_valid());
-            assert_eq!(s.as_bool().value(), Some(true));
-        }
-
-        let s = arr.scalar_at(3);
-        assert!(s.is_valid());
-        assert_eq!(s.as_bool().value(), Some(false));
-
-        let s = arr.scalar_at(4);
-        assert!(s.is_null());
+        let expected =
+            ByteBoolArray::from(vec![Some(true), Some(true), Some(true), Some(false), None]);
+        assert_arrays_eq!(arr, expected.to_array());
     }
 
     #[test]