vortex_array/arrays/
assertions.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5
6use itertools::Itertools;
7use vortex_error::VortexExpect;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::LEGACY_SESSION;
13use crate::RecursiveCanonical;
14use crate::VortexSessionExecute;
15
16fn format_indices<I: IntoIterator<Item = usize>>(indices: I) -> impl Display {
17    indices.into_iter().format(",")
18}
19
20/// Executes an array to recursive canonical form with the given execution context.
21fn execute_to_canonical(array: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef {
22    array
23        .execute::<RecursiveCanonical>(ctx)
24        .vortex_expect("failed to execute array to recursive canonical form")
25        .0
26        .into_array()
27}
28
29/// Finds indices where two arrays differ based on `scalar_at` comparison.
30#[expect(clippy::unwrap_used)]
31fn find_mismatched_indices(left: &ArrayRef, right: &ArrayRef) -> Vec<usize> {
32    assert_eq!(left.len(), right.len());
33    let mut ctx = LEGACY_SESSION.create_execution_ctx();
34    (0..left.len())
35        .filter(|i| {
36            left.execute_scalar(*i, &mut ctx).unwrap()
37                != right.execute_scalar(*i, &mut ctx).unwrap()
38        })
39        .collect()
40}
41
42/// Asserts that the scalar at position `$n` in array `$arr` equals `$expected`.
43///
44/// This is a convenience macro for testing that avoids verbose scalar comparison code.
45///
46/// # Example
47/// ```ignore
48/// let arr = PrimitiveArray::from_iter([1, 2, 3]);
49/// assert_nth_scalar!(arr, 0, 1);
50/// assert_nth_scalar!(arr, 1, 2);
51/// ```
52#[macro_export]
53macro_rules! assert_nth_scalar {
54    ($arr:expr, $n:expr, $expected:expr) => {{
55        use $crate::IntoArray as _;
56        use $crate::LEGACY_SESSION;
57        use $crate::VortexSessionExecute as _;
58        let arr_ref: $crate::ArrayRef = $crate::IntoArray::into_array($arr.clone());
59        assert_eq!(
60            arr_ref
61                .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx())
62                .unwrap(),
63            $expected.try_into().unwrap()
64        );
65    }};
66}
67
68/// Asserts that the scalar at position `$n` in array `$arr` is null.
69///
70/// # Example
71///
72/// ```ignore
73/// let arr = PrimitiveArray::from_option_iter([Some(1), None, Some(3)]);
74/// assert_nth_scalar_null!(arr, 1);
75/// ```
76#[macro_export]
77macro_rules! assert_nth_scalar_is_null {
78    ($arr:expr, $n:expr) => {{
79        use $crate::LEGACY_SESSION;
80        use $crate::VortexSessionExecute as _;
81        let arr_ref: $crate::ArrayRef = $crate::IntoArray::into_array($arr.clone());
82        assert!(
83            arr_ref
84                .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx())
85                .unwrap()
86                .is_null(),
87            "expected scalar at index {} to be null, but was {:?}",
88            $n,
89            arr_ref
90                .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx())
91                .unwrap()
92        );
93    }};
94}
95
96#[macro_export]
97macro_rules! assert_arrays_eq {
98    ($left:expr, $right:expr) => {{
99
100
101        let left: $crate::ArrayRef = $crate::IntoArray::into_array($left.clone());
102        let right: $crate::ArrayRef = $crate::IntoArray::into_array($right.clone());
103        if left.dtype() != right.dtype() {
104            panic!(
105                "assertion left == right failed: arrays differ in type: {} != {}.\n  left: {}\n right: {}",
106                left.dtype(),
107                right.dtype(),
108                left.display_values(),
109                right.display_values()
110            )
111        }
112
113        assert_eq!(
114            left.len(),
115            right.len(),
116            "assertion left == right failed: arrays differ in length: {} != {}.\n  left: {}\n right: {}",
117            left.len(),
118            right.len(),
119            left.display_values(),
120            right.display_values()
121        );
122
123        let left = left.clone();
124        let right = right.clone();
125        $crate::arrays::assert_arrays_eq_impl(&left, &right);
126    }};
127}
128
129/// Implementation of `assert_arrays_eq!` — called by the macro after converting inputs to
130/// `ArrayRef`.
131#[track_caller]
132#[expect(clippy::panic)]
133pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef) {
134    let executed = execute_to_canonical(left.clone(), &mut LEGACY_SESSION.create_execution_ctx());
135
136    let left_right = find_mismatched_indices(left, right);
137    let executed_right = find_mismatched_indices(&executed, right);
138
139    if !left_right.is_empty() || !executed_right.is_empty() {
140        let mut msg = String::new();
141        if !left_right.is_empty() {
142            msg.push_str(&format!(
143                "\n  left != right at indices: {}",
144                format_indices(left_right)
145            ));
146        }
147        if !executed_right.is_empty() {
148            msg.push_str(&format!(
149                "\n  executed != right at indices: {}",
150                format_indices(executed_right)
151            ));
152        }
153        panic!(
154            "assertion failed: arrays do not match:{}\n     left: {}\n    right: {}\n executed: {}",
155            msg,
156            left.display_values(),
157            right.display_values(),
158            executed.display_values()
159        )
160    }
161}
vortex_array/arrays/assertions.rs

vortex_array/arrays/
assertions.rs
