diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index 84dab9dcf68..ad883eb5dc7 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -58,6 +58,7 @@ mod test { use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::assert_nth_scalar; use vortex_array::compute::take; use vortex_array::patches::Patches; use vortex_array::validity::Validity; @@ -81,8 +82,8 @@ mod test { ) .unwrap(); let sliced = arr.slice(1024..2048).as_::().clone(); - assert_eq!(sliced.scalar_at(0), (1024u32 % 64).into()); - assert_eq!(sliced.scalar_at(1023), (2047u32 % 64).into()); + assert_nth_scalar!(sliced, 0, 1024u32 % 64); + assert_nth_scalar!(sliced, 1023, 2047u32 % 64); assert_eq!(sliced.offset(), 0); assert_eq!(sliced.len(), 1024); } @@ -96,8 +97,8 @@ mod test { .unwrap() .into_array(); let sliced = arr.slice(512..1434).as_::().clone(); - assert_eq!(sliced.scalar_at(0), (512u32 % 64).into()); - assert_eq!(sliced.scalar_at(921), (1433u32 % 64).into()); + assert_nth_scalar!(sliced, 0, 512u32 % 64); + assert_nth_scalar!(sliced, 921, 1433u32 % 64); assert_eq!(sliced.offset(), 512); assert_eq!(sliced.len(), 922); } @@ -111,11 +112,8 @@ mod test { .unwrap(); let compressed = packed.slice(768..9999); - assert_eq!(compressed.scalar_at(0), ((768 % 63) as u8).into()); - assert_eq!( - compressed.scalar_at(compressed.len() - 1), - ((9998 % 63) as u8).into() - ); + assert_nth_scalar!(compressed, 0, (768 % 63) as u8); + assert_nth_scalar!(compressed, compressed.len() - 1, (9998 % 63) as u8); } #[test] @@ -127,11 +125,8 @@ mod test { .unwrap(); let compressed = packed.slice(7168..9216); - assert_eq!(compressed.scalar_at(0), ((7168 % 63) as u8).into()); - assert_eq!( - compressed.scalar_at(compressed.len() - 1), - ((9215 % 63) as u8).into() - ); + assert_nth_scalar!(compressed, 0, (7168 % 63) as u8); + assert_nth_scalar!(compressed, compressed.len() - 1, (9215 % 63) as u8); } #[test] @@ -143,13 +138,13 @@ mod test { .unwrap() .into_array(); let sliced = arr.slice(512..1434).as_::().clone(); - assert_eq!(sliced.scalar_at(0), (512u32 % 64).into()); - assert_eq!(sliced.scalar_at(921), (1433u32 % 64).into()); + assert_nth_scalar!(sliced, 0, 512u32 % 64); + assert_nth_scalar!(sliced, 921, 1433u32 % 64); assert_eq!(sliced.offset(), 512); assert_eq!(sliced.len(), 922); let doubly_sliced = sliced.slice(127..911).as_::().clone(); - assert_eq!(doubly_sliced.scalar_at(0), ((512u32 + 127) % 64).into()); - assert_eq!(doubly_sliced.scalar_at(783), ((512u32 + 910) % 64).into()); + assert_nth_scalar!(doubly_sliced, 0, (512u32 + 127) % 64); + assert_nth_scalar!(doubly_sliced, 783, (512u32 + 910) % 64); assert_eq!(doubly_sliced.offset(), 639); assert_eq!(doubly_sliced.len(), 784); } diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index d8e7955f8f3..0d8a24d52b6 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -93,38 +93,32 @@ fn compress_primitive((0u32..10_000).collect()); + do_roundtrip_test((0u32..10_000).collect()); } #[test] fn test_compress_nullable() { - do_roundtrip_test::(PrimitiveArray::from_option_iter( + do_roundtrip_test(PrimitiveArray::from_option_iter( (0u32..10_000).map(|i| (i % 2 == 0).then_some(i)), )); } #[test] fn test_compress_overflow() { - do_roundtrip_test::((0..10_000).map(|i| (i % (u8::MAX as i32)) as u8).collect()); + do_roundtrip_test((0..10_000).map(|i| (i % (u8::MAX as i32)) as u8).collect()); } - fn do_roundtrip_test(input: PrimitiveArray) { + fn do_roundtrip_test(input: PrimitiveArray) { let delta = DeltaArray::try_from_primitive_array(&input).unwrap(); assert_eq!(delta.len(), input.len()); let decompressed = delta_decompress(&delta); - let decompressed_slice = decompressed.as_slice::(); - assert_eq!(decompressed_slice.len(), input.len()); - for (actual, expected) in decompressed_slice.iter().zip(input.as_slice()) { - assert_eq!(actual, expected); - } - assert_eq!(decompressed.validity(), input.validity()); + assert_arrays_eq!(decompressed, input); } } diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 95700b875b5..e46235560d0 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -7,6 +7,7 @@ use vortex_array::IntoArray; use vortex_array::ToCanonical; use vortex_array::arrays::builder::VarBinBuilder; use vortex_array::assert_arrays_eq; +use vortex_array::assert_nth_scalar; use vortex_array::compute::filter; use vortex_array::compute::take; use vortex_buffer::buffer; @@ -18,12 +19,6 @@ use crate::FSSTVTable; use crate::fsst_compress; use crate::fsst_train_compressor; -macro_rules! assert_nth_scalar { - ($arr:expr, $n:expr, $expected:expr) => { - assert_eq!($arr.scalar_at($n), $expected.try_into().unwrap()); - }; -} - /// this function is VERY slow on miri, so we only want to run it once pub(crate) fn build_fsst_array() -> ArrayRef { let mut input_array = VarBinBuilder::::with_capacity(3); diff --git a/encodings/pco/src/test.rs b/encodings/pco/src/test.rs index 8af9ec4e159..0a4936dd136 100644 --- a/encodings/pco/src/test.rs +++ b/encodings/pco/src/test.rs @@ -9,6 +9,7 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrow::compute::to_arrow_preferred; use vortex_array::assert_arrays_eq; +use vortex_array::assert_nth_scalar; use vortex_array::serde::ArrayParts; use vortex_array::serde::SerializeOptions; use vortex_array::session::ArraySession; @@ -25,12 +26,6 @@ use vortex_mask::Mask; use crate::PcoArray; use crate::PcoVTable; -macro_rules! assert_nth_scalar { - ($arr:expr, $n:expr, $expected:expr) => { - assert_eq!($arr.scalar_at($n), $expected.try_into().unwrap()); - }; -} - #[test] fn test_compress_decompress() { let data: Vec = (0..200).collect(); diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index 250b336c61b..060c0bae72c 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -7,6 +7,7 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; +use vortex_array::assert_nth_scalar; use vortex_array::validity::Validity; use vortex_array::vtable::ValidityHelper; use vortex_buffer::Alignment; @@ -17,12 +18,6 @@ use vortex_mask::Mask; use crate::ZstdArray; -macro_rules! assert_nth_scalar { - ($arr:expr, $n:expr, $expected:expr) => { - assert_eq!($arr.scalar_at($n), $expected.try_into().unwrap()); - }; -} - #[test] fn test_zstd_compress_decompress() { let data: Vec = (0..200).collect(); @@ -115,7 +110,6 @@ fn test_zstd_with_dict() { 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); diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index fcff725c954..9aa30b03f3f 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -131,11 +131,11 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; use vortex_dtype::DecimalDType; use vortex_dtype::FieldNames; use vortex_dtype::Nullability; use vortex_mask::Mask; - use vortex_scalar::Scalar; use super::mask_canonical_array; @@ -160,12 +160,8 @@ mod tests { let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); - assert_eq!(result.len(), 5); - assert!(!result.is_valid(0)); - assert_eq!(result.scalar_at(1), Scalar::from(Some(false))); - assert_eq!(result.scalar_at(2), Scalar::from(Some(true))); - assert!(!result.is_valid(3)); - assert_eq!(result.scalar_at(4), Scalar::from(Some(true))); + let expected = BoolArray::from_iter([None, Some(false), Some(true), None, Some(true)]); + assert_arrays_eq!(result, expected); } #[test] @@ -175,12 +171,8 @@ mod tests { let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); - assert_eq!(result.len(), 5); - assert_eq!(result.scalar_at(0), Scalar::from(Some(1))); - assert!(!result.is_valid(1)); - assert_eq!(result.scalar_at(2), Scalar::from(Some(3))); - assert!(!result.is_valid(3)); - assert_eq!(result.scalar_at(4), Scalar::from(Some(5))); + let expected = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]); + assert_arrays_eq!(result, expected); } #[test] @@ -190,30 +182,24 @@ mod tests { let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); - assert_eq!(result.len(), 5); - assert!(!result.is_valid(0)); - assert!(!result.is_valid(1)); // was already null - assert_eq!(result.scalar_at(2), Scalar::from(Some(3))); - assert!(!result.is_valid(3)); - assert!(!result.is_valid(4)); // was already null + let expected = PrimitiveArray::from_option_iter([None, None, Some(3i32), None, None]); + assert_arrays_eq!(result, expected); } #[test] fn test_mask_decimal_array() { + let dtype = DecimalDType::new(10, 2); let array = DecimalArray::from_option_iter( [Some(1i128), Some(2), Some(3), Some(4), Some(5)], - DecimalDType::new(10, 2), + dtype, ); let mask = Mask::from_iter([false, false, true, false, false]); let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); - assert_eq!(result.len(), 5); - assert!(result.is_valid(0)); - assert!(result.is_valid(1)); - assert!(!result.is_valid(2)); - assert!(result.is_valid(3)); - assert!(result.is_valid(4)); + let expected = + DecimalArray::from_option_iter([Some(1i128), Some(2), None, Some(4), Some(5)], dtype); + assert_arrays_eq!(result, expected); } #[test] @@ -223,18 +209,9 @@ mod tests { let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); - assert_eq!(result.len(), 5); - assert!(!result.is_valid(0)); - assert_eq!( - result.scalar_at(1), - Scalar::utf8("two", Nullability::Nullable) - ); - assert!(!result.is_valid(2)); - assert_eq!( - result.scalar_at(3), - Scalar::utf8("four", Nullability::Nullable) - ); - assert!(!result.is_valid(4)); + let expected = + VarBinViewArray::from_iter_nullable_str([None, Some("two"), None, Some("four"), None]); + assert_arrays_eq!(result, expected); } #[test] @@ -304,11 +281,8 @@ mod tests { let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); - assert_eq!(result.len(), 5); - // All values should be masked out (null) - for i in 0..5 { - assert!(!result.is_valid(i)); - } + let expected = PrimitiveArray::from_option_iter([None, None, None, None, None::]); + assert_arrays_eq!(result, expected); } #[test] @@ -318,14 +292,9 @@ mod tests { let result = mask_canonical_array(array.to_canonical(), &mask).unwrap(); - assert_eq!(result.len(), 5); - // No values should be masked out - for i in 0..5 { - assert!(result.is_valid(i)); - #[allow(clippy::cast_possible_truncation)] - let expected = (i + 1) as i32; - assert_eq!(result.scalar_at(i), Scalar::from(Some(expected))); - } + let expected = + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4), Some(5)]); + assert_arrays_eq!(result, expected); } #[test] diff --git a/vortex-array/src/arrays/assertions.rs b/vortex-array/src/arrays/assertions.rs index ccc53d1b88a..3f452e1b575 100644 --- a/vortex-array/src/arrays/assertions.rs +++ b/vortex-array/src/arrays/assertions.rs @@ -9,6 +9,23 @@ pub fn format_indices>(indices: I) -> impl Display indices.into_iter().format(",") } +/// Asserts that the scalar at position `$n` in array `$arr` equals `$expected`. +/// +/// This is a convenience macro for testing that avoids verbose scalar comparison code. +/// +/// # Example +/// ```ignore +/// let arr = PrimitiveArray::from_iter([1, 2, 3]); +/// assert_nth_scalar!(arr, 0, 1); +/// assert_nth_scalar!(arr, 1, 2); +/// ``` +#[macro_export] +macro_rules! assert_nth_scalar { + ($arr:expr, $n:expr, $expected:expr) => { + assert_eq!($arr.scalar_at($n), $expected.try_into().unwrap()); + }; +} + #[macro_export] macro_rules! assert_arrays_eq { ($left:expr, $right:expr) => {{ diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index 1b7957c5e98..0e856ee68b0 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -102,8 +102,7 @@ fn scalar_at_empty_children_both_sides() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_eq!(array.scalar_at(0), 1u64.into()); - assert_eq!(array.scalar_at(1), 2u64.into()); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2])); } #[test] @@ -118,10 +117,7 @@ fn scalar_at_empty_children_trailing() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_eq!(array.scalar_at(0), 1u64.into()); - assert_eq!(array.scalar_at(1), 2u64.into()); - assert_eq!(array.scalar_at(2), 3u64.into()); - assert_eq!(array.scalar_at(3), 4u64.into()); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); } #[test] @@ -136,10 +132,7 @@ fn scalar_at_empty_children_leading() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_eq!(array.scalar_at(0), 1u64.into()); - assert_eq!(array.scalar_at(1), 2u64.into()); - assert_eq!(array.scalar_at(2), 3u64.into()); - assert_eq!(array.scalar_at(3), 4u64.into()); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); } #[test] diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 6eaf62127ef..55a5b311e83 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -16,6 +16,7 @@ use super::*; use crate::IntoArray; use crate::arrays::ListVTable; use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::ListBuilder; use crate::compute::filter; @@ -114,11 +115,11 @@ fn test_list_filter_dense_mask() { // Should have 5 lists remaining. assert_eq!(filtered_list.len(), 5); - // Verify first remaining list (originally index 1) has correct elements. - let first_list = filtered_list.list_elements_at(0); - assert_eq!(first_list.len(), 15); // 25 - 10 - assert_eq!(first_list.scalar_at(0), 10.into()); - assert_eq!(first_list.scalar_at(14), 24.into()); + // Verify first remaining list (originally index 1) has correct elements [10..25]. + assert_arrays_eq!( + filtered_list.list_elements_at(0), + PrimitiveArray::from_iter(10..25) + ); } #[test] @@ -143,17 +144,17 @@ fn test_list_filter_sparse_mask() { // Should have 2 lists remaining. assert_eq!(filtered_list.len(), 2); - // Verify first list (originally index 0). - let first_list = filtered_list.list_elements_at(0); - assert_eq!(first_list.len(), 10); - assert_eq!(first_list.scalar_at(0), 0.into()); - assert_eq!(first_list.scalar_at(9), 9.into()); - - // Verify second list (originally index 5). - let second_list = filtered_list.list_elements_at(1); - assert_eq!(second_list.len(), 15); // 100 - 85 - assert_eq!(second_list.scalar_at(0), 85.into()); - assert_eq!(second_list.scalar_at(14), 99.into()); + // Verify first list (originally index 0) has elements [0..10]. + assert_arrays_eq!( + filtered_list.list_elements_at(0), + PrimitiveArray::from_iter(0..10) + ); + + // Verify second list (originally index 5) has elements [85..100]. + assert_arrays_eq!( + filtered_list.list_elements_at(1), + PrimitiveArray::from_iter(85..100) + ); } #[test] @@ -177,10 +178,11 @@ fn test_list_filter_empty_lists() { // First list is empty. assert_eq!(filtered_list.list_elements_at(0).len(), 0); - // Second list has 3 elements. - let second_list = filtered_list.list_elements_at(1); - assert_eq!(second_list.len(), 3); - assert_eq!(second_list.scalar_at(0), 0.into()); + // Second list has elements [0..3]. + assert_arrays_eq!( + filtered_list.list_elements_at(1), + PrimitiveArray::from_iter(0..3) + ); // Third list is empty. assert_eq!(filtered_list.list_elements_at(2).len(), 0); @@ -236,11 +238,12 @@ fn test_list_filter_all_true() { // All lists should be preserved. assert_eq!(filtered_list.len(), 4); - // Verify all lists are intact. - for i in 0..4i32 { - let list_at_i = filtered_list.list_elements_at(i as usize); - assert_eq!(list_at_i.len(), 5); - assert_eq!(list_at_i.scalar_at(0), (i * 5).into()); + // Verify all lists are intact: [0..5], [5..10], [10..15], [15..20]. + for i in 0..4 { + assert_arrays_eq!( + filtered_list.list_elements_at(i), + PrimitiveArray::from_iter((i * 5) as i32..((i + 1) * 5) as i32) + ); } } @@ -282,10 +285,11 @@ fn test_list_filter_single_element() { assert_eq!(filtered_list.len(), 1); - let single_list = filtered_list.list_elements_at(0); - assert_eq!(single_list.len(), 10); - assert_eq!(single_list.scalar_at(0), 20.into()); - assert_eq!(single_list.scalar_at(9), 29.into()); + // The single remaining list has elements [20..30]. + assert_arrays_eq!( + filtered_list.list_elements_at(0), + PrimitiveArray::from_iter(20..30) + ); } #[test] @@ -309,11 +313,12 @@ fn test_list_filter_alternating_pattern() { assert_eq!(filtered_list.len(), 6); - // Verify that we have the correct lists (0, 2, 4, 6, 8, 10). + // Verify we have the correct lists: [0..5], [10..15], [20..25], [30..35], [40..45], [50..55]. for (i, expected_start) in [0, 10, 20, 30, 40, 50].iter().enumerate() { - let list_at_i = filtered_list.list_elements_at(i); - assert_eq!(list_at_i.len(), 5); - assert_eq!(list_at_i.scalar_at(0), (*expected_start).into()); + assert_arrays_eq!( + filtered_list.list_elements_at(i), + PrimitiveArray::from_iter(*expected_start..*expected_start + 5) + ); } } @@ -574,14 +579,11 @@ fn test_list_of_lists() { // Check first inner list [1, 2]. let first_inner = first_outer_list.list_elements_at(0); - assert_eq!(first_inner.len(), 2); - assert_eq!(first_inner.scalar_at(0), 1.into()); - assert_eq!(first_inner.scalar_at(1), 2.into()); + assert_arrays_eq!(first_inner, PrimitiveArray::from_iter([1, 2])); // Check second inner list [3]. let second_inner = first_outer_list.list_elements_at(1); - assert_eq!(second_inner.len(), 1); - assert_eq!(second_inner.scalar_at(0), 3.into()); + assert_arrays_eq!(second_inner, PrimitiveArray::from_iter([3])); // Check the second list of lists [[4, 5, 6]]. let second_outer = list_of_lists.list_elements_at(1); @@ -589,10 +591,7 @@ fn test_list_of_lists() { assert_eq!(second_outer_list.len(), 1); let inner = second_outer_list.list_elements_at(0); - assert_eq!(inner.len(), 3); - assert_eq!(inner.scalar_at(0), 4.into()); - assert_eq!(inner.scalar_at(1), 5.into()); - assert_eq!(inner.scalar_at(2), 6.into()); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([4, 5, 6])); // Check the third list of lists (empty). let third_outer = list_of_lists.list_elements_at(2); @@ -605,8 +604,7 @@ fn test_list_of_lists() { assert_eq!(fourth_outer_list.len(), 1); let inner = fourth_outer_list.list_elements_at(0); - assert_eq!(inner.len(), 1); - assert_eq!(inner.scalar_at(0), 7.into()); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([7])); // Test scalar conversion. let scalar = list_of_lists.scalar_at(0); @@ -762,8 +760,7 @@ fn test_list_of_lists_both_nullable() { let third_list = third_outer.as_::(); assert_eq!(third_list.len(), 1); let inner = third_list.list_elements_at(0); - assert_eq!(inner.len(), 1); - assert_eq!(inner.scalar_at(0), 3.into()); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([3])); // Fourth outer list should have a null inner list. let fourth_outer = list_of_lists.list_elements_at(3); diff --git a/vortex-array/src/arrays/varbin/tests.rs b/vortex-array/src/arrays/varbin/tests.rs index 51595369b91..0180605ac0a 100644 --- a/vortex-array/src/arrays/varbin/tests.rs +++ b/vortex-array/src/arrays/varbin/tests.rs @@ -11,7 +11,9 @@ use vortex_dtype::Nullability; use crate::Array; use crate::ArrayRef; use crate::IntoArray; +use crate::arrays::VarBinViewArray; use crate::arrays::varbin::VarBinArray; +use crate::assert_arrays_eq; use crate::validity::Validity; #[fixture] @@ -31,19 +33,17 @@ fn binary_array() -> ArrayRef { #[rstest] pub fn test_scalar_at(binary_array: ArrayRef) { - assert_eq!(binary_array.len(), 2); - assert_eq!(binary_array.scalar_at(0), "hello world".into()); - assert_eq!( - binary_array.scalar_at(1), - "hello world this is a long string".into() - ) + assert_arrays_eq!( + binary_array, + VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) + ); } #[rstest] pub fn slice_array(binary_array: ArrayRef) { let binary_arr = binary_array.slice(1..2); - assert_eq!( - binary_arr.scalar_at(0), - "hello world this is a long string".into() + assert_arrays_eq!( + binary_arr, + VarBinViewArray::from_iter_str(["hello world this is a long string"]) ); } diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index f33beba8a16..0e83fdae2cb 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -1,22 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_scalar::Scalar; use vortex_vector::binaryview::BinaryView; -use crate::Array; use crate::ToCanonical; use crate::arrays::VarBinViewArray; +use crate::assert_arrays_eq; #[test] pub fn varbin_view() { let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]); - assert_eq!(binary_arr.len(), 2); - assert_eq!(binary_arr.scalar_at(0), Scalar::from("hello world")); - assert_eq!( - binary_arr.scalar_at(1), - Scalar::from("hello world this is a long string") + assert_arrays_eq!( + binary_arr, + VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) ); } @@ -25,9 +22,9 @@ pub fn slice_array() { let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) .slice(1..2); - assert_eq!( - binary_arr.scalar_at(0), - Scalar::from("hello world this is a long string") + assert_arrays_eq!( + binary_arr, + VarBinViewArray::from_iter_str(["hello world this is a long string"]) ); } @@ -35,8 +32,10 @@ pub fn slice_array() { pub fn flatten_array() { let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); let var_bin = binary_arr.to_varbinview(); - assert_eq!(var_bin.scalar_at(0), Scalar::from("string1")); - assert_eq!(var_bin.scalar_at(1), Scalar::from("string2")); + assert_arrays_eq!( + var_bin, + VarBinViewArray::from_iter_str(["string1", "string2"]) + ); } #[test] diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index d434cb5ae51..4c39a1c924d 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -429,6 +429,7 @@ mod tests { use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::assert_nth_scalar; use crate::builders::ArrayBuilder; use crate::vtable::ValidityHelper; @@ -478,23 +479,17 @@ mod tests { // Check first list: [1, 2, 3]. let first_list = listview.list_elements_at(0); - assert_eq!(first_list.len(), 3); - assert_eq!(first_list.scalar_at(0), 1i32.into()); - assert_eq!(first_list.scalar_at(1), 2i32.into()); - assert_eq!(first_list.scalar_at(2), 3i32.into()); + assert_arrays_eq!(first_list, PrimitiveArray::from_iter([1i32, 2, 3])); // Check empty list. - let empty_list = listview.list_elements_at(1); - assert_eq!(empty_list.len(), 0); + assert_eq!(listview.list_elements_at(1).len(), 0); // Check null list. assert!(!listview.validity().is_valid(2)); // Check last list: [4, 5]. let last_list = listview.list_elements_at(3); - assert_eq!(last_list.len(), 2); - assert_eq!(last_list.scalar_at(0), 4i32.into()); - assert_eq!(last_list.scalar_at(1), 5i32.into()); + assert_arrays_eq!(last_list, PrimitiveArray::from_iter([4i32, 5])); } #[test] @@ -525,15 +520,16 @@ mod tests { assert_eq!(listview.len(), 2); // Verify first list: [1, 2]. - let first = listview.list_elements_at(0); - assert_eq!(first.scalar_at(0), 1i32.into()); - assert_eq!(first.scalar_at(1), 2i32.into()); + assert_arrays_eq!( + listview.list_elements_at(0), + PrimitiveArray::from_iter([1i32, 2]) + ); // Verify second list: [3, 4, 5]. - let second = listview.list_elements_at(1); - assert_eq!(second.scalar_at(0), 3i32.into()); - assert_eq!(second.scalar_at(1), 4i32.into()); - assert_eq!(second.scalar_at(2), 5i32.into()); + assert_arrays_eq!( + listview.list_elements_at(1), + PrimitiveArray::from_iter([3i32, 4, 5]) + ); // Test u64 offsets with u16 sizes. let dtype2: Arc = Arc::new(I32.into()); @@ -555,7 +551,7 @@ mod tests { for i in 0..5i32 { let list = listview2.list_elements_at(i as usize); assert_eq!(list.len(), 1); - assert_eq!(list.scalar_at(0), (i * 10).into()); + assert_nth_scalar!(list, 0, i * 10); } } @@ -591,10 +587,10 @@ mod tests { assert!(!listview.validity().is_valid(3)); // Last is the regular list: [10, 20]. - let last_list = listview.list_elements_at(4); - assert_eq!(last_list.len(), 2); - assert_eq!(last_list.scalar_at(0), 10i32.into()); - assert_eq!(last_list.scalar_at(1), 20i32.into()); + assert_arrays_eq!( + listview.list_elements_at(4), + PrimitiveArray::from_iter([10i32, 20]) + ); } #[test] @@ -631,25 +627,25 @@ mod tests { // Check the extended data. // First list: [0] (initial data). - let first = listview.list_elements_at(0); - assert_eq!(first.len(), 1); - assert_eq!(first.scalar_at(0), 0i32.into()); + assert_arrays_eq!( + listview.list_elements_at(0), + PrimitiveArray::from_iter([0i32]) + ); // Second list: [1, 2, 3] (from source). - let second = listview.list_elements_at(1); - assert_eq!(second.len(), 3); - assert_eq!(second.scalar_at(0), 1i32.into()); - assert_eq!(second.scalar_at(1), 2i32.into()); - assert_eq!(second.scalar_at(2), 3i32.into()); + assert_arrays_eq!( + listview.list_elements_at(1), + PrimitiveArray::from_iter([1i32, 2, 3]) + ); // Third list: null (from source). assert!(!listview.validity().is_valid(2)); // Fourth list: [4, 5] (from source). - let fourth = listview.list_elements_at(3); - assert_eq!(fourth.len(), 2); - assert_eq!(fourth.scalar_at(0), 4i32.into()); - assert_eq!(fourth.scalar_at(1), 5i32.into()); + assert_arrays_eq!( + listview.list_elements_at(3), + PrimitiveArray::from_iter([4i32, 5]) + ); } #[test] diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 87f2fce0748..a14495334b1 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -387,6 +387,7 @@ mod tests { use vortex_array::ToCanonical; use vortex_array::arrays::BoolArray; use vortex_array::arrays::StructArray; + use vortex_array::assert_nth_scalar; use vortex_array::expr::Expression; use vortex_array::expr::col; use vortex_array::expr::eq; @@ -702,8 +703,8 @@ mod tests { // ...and the result is masked with the validity of the parent StructArray assert_eq!(result.scalar_at(0), Scalar::null(result.dtype().clone()),); - assert_eq!(result.scalar_at(1), 2.into()); - assert_eq!(result.scalar_at(2), 3.into()); + assert_nth_scalar!(result, 1, 2); + assert_nth_scalar!(result, 2, 3); } #[rstest]