vortex_array/compute/conformance/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_mask::Mask;
6
7use crate::ArrayRef;
8use crate::IntoArray;
9use crate::LEGACY_SESSION;
10use crate::VortexSessionExecute;
11use crate::arrays::BoolArray;
12use crate::arrays::bool::BoolArrayExt;
13use crate::builtins::ArrayBuiltins;
14
15/// Test mask compute function with various array sizes and patterns.
16/// The mask operation sets elements to null where the mask is true.
17pub fn test_mask_conformance(array: &ArrayRef) {
18    let len = array.len();
19
20    if len > 0 {
21        test_heterogenous_mask(array);
22        test_empty_mask(array);
23        test_full_mask(array);
24        test_alternating_mask(array);
25        test_sparse_mask(array);
26        test_single_element_mask(array);
27    }
28
29    if len >= 5 {
30        test_double_mask(array);
31    }
32
33    if len > 0 {
34        test_nullable_mask_input(array);
35    }
36}
37
38/// Tests masking with a heterogeneous pattern
39fn test_heterogenous_mask(array: &ArrayRef) {
40    let len = array.len();
41
42    // Create a pattern where roughly half the values are masked
43    let mask_pattern: Vec<bool> = (0..len).map(|i| i % 3 != 1).collect();
44    let mask_array = Mask::from_iter(mask_pattern.clone());
45
46    let masked = array
47        .clone()
48        .mask((!&mask_array).into_array())
49        .vortex_expect("mask should succeed in conformance test");
50    assert_eq!(masked.len(), array.len());
51
52    // Verify masked elements are null and unmasked elements are preserved
53    for (i, &masked_out) in mask_pattern.iter().enumerate() {
54        if masked_out {
55            assert!(
56                !masked
57                    .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx())
58                    .vortex_expect("is_valid should succeed in conformance test")
59            );
60        } else {
61            assert_eq!(
62                masked
63                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
64                    .vortex_expect("scalar_at should succeed in conformance test"),
65                array
66                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
67                    .vortex_expect("scalar_at should succeed in conformance test")
68                    .into_nullable()
69            );
70        }
71    }
72}
73
74/// Tests that an empty mask (all false) preserves all elements
75fn test_empty_mask(array: &ArrayRef) {
76    let len = array.len();
77    let all_unmasked = vec![false; len];
78    let mask_array = Mask::from_iter(all_unmasked);
79
80    let masked = array
81        .clone()
82        .mask((!&mask_array).into_array())
83        .vortex_expect("mask should succeed in conformance test");
84    assert_eq!(masked.len(), array.len());
85
86    // All elements should be preserved
87    for i in 0..len {
88        assert_eq!(
89            masked
90                .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
91                .vortex_expect("scalar_at should succeed in conformance test"),
92            array
93                .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
94                .vortex_expect("scalar_at should succeed in conformance test")
95                .into_nullable()
96        );
97    }
98}
99
100/// Tests that a full mask (all true) makes all elements null
101fn test_full_mask(array: &ArrayRef) {
102    let len = array.len();
103    let all_masked = vec![true; len];
104    let mask_array = Mask::from_iter(all_masked);
105
106    let masked = array
107        .clone()
108        .mask((!&mask_array).into_array())
109        .vortex_expect("mask should succeed in conformance test");
110    assert_eq!(masked.len(), array.len());
111
112    // All elements should be null
113    for i in 0..len {
114        assert!(
115            !masked
116                .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx())
117                .vortex_expect("is_valid should succeed in conformance test")
118        );
119    }
120}
121
122/// Tests alternating mask pattern
123fn test_alternating_mask(array: &ArrayRef) {
124    let len = array.len();
125    let pattern: Vec<bool> = (0..len).map(|i| i % 2 == 0).collect();
126    let mask_array = Mask::from_iter(pattern);
127
128    let masked = array
129        .clone()
130        .mask((!&mask_array).into_array())
131        .vortex_expect("mask should succeed in conformance test");
132    assert_eq!(masked.len(), array.len());
133
134    for i in 0..len {
135        if i % 2 == 0 {
136            assert!(
137                !masked
138                    .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx())
139                    .vortex_expect("is_valid should succeed in conformance test")
140            );
141        } else {
142            assert_eq!(
143                masked
144                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
145                    .vortex_expect("scalar_at should succeed in conformance test"),
146                array
147                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
148                    .vortex_expect("scalar_at should succeed in conformance test")
149                    .into_nullable()
150            );
151        }
152    }
153}
154
155/// Tests sparse mask (only a few elements masked)
156fn test_sparse_mask(array: &ArrayRef) {
157    let len = array.len();
158    if len < 10 {
159        return; // Skip for small arrays
160    }
161
162    // Mask only about 10% of elements
163    let pattern: Vec<bool> = (0..len).map(|i| i % 10 == 0).collect();
164    let mask_array = Mask::from_iter(pattern.clone());
165
166    let masked = array
167        .clone()
168        .mask((!&mask_array).into_array())
169        .vortex_expect("mask should succeed in conformance test");
170    assert_eq!(masked.len(), array.len());
171
172    // Count how many elements are valid after masking
173    let valid_count = (0..len)
174        .filter(|&i| {
175            masked
176                .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx())
177                .vortex_expect("is_valid should succeed in conformance test")
178        })
179        .count();
180
181    // Count how many elements should be invalid:
182    // - Elements that were masked (pattern[i] == true)
183    // - Elements that were already invalid in the original array
184    let expected_invalid_count = (0..len)
185        .filter(|&i| {
186            pattern[i]
187                || !array
188                    .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx())
189                    .vortex_expect("is_valid should succeed in conformance test")
190        })
191        .count();
192
193    assert_eq!(valid_count, len - expected_invalid_count);
194}
195
196/// Tests masking a single element
197fn test_single_element_mask(array: &ArrayRef) {
198    let len = array.len();
199
200    // Mask only the first element
201    let mut pattern = vec![false; len];
202    pattern[0] = true;
203    let mask_array = Mask::from_iter(pattern);
204
205    let masked = array
206        .clone()
207        .mask((!&mask_array).into_array())
208        .vortex_expect("mask should succeed in conformance test");
209    assert!(
210        !masked
211            .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx())
212            .vortex_expect("is_valid should succeed in conformance test")
213    );
214
215    for i in 1..len {
216        assert_eq!(
217            masked
218                .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
219                .vortex_expect("scalar_at should succeed in conformance test"),
220            array
221                .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
222                .vortex_expect("scalar_at should succeed in conformance test")
223                .into_nullable()
224        );
225    }
226}
227
228/// Tests double masking operations
229fn test_double_mask(array: &ArrayRef) {
230    let len = array.len();
231
232    // Create two different mask patterns
233    let mask1_pattern: Vec<bool> = (0..len).map(|i| i % 3 == 0).collect();
234    let mask2_pattern: Vec<bool> = (0..len).map(|i| i % 2 == 0).collect();
235
236    let mask1 = Mask::from_iter(mask1_pattern.clone());
237    let mask2 = Mask::from_iter(mask2_pattern.clone());
238
239    let first_masked = array
240        .clone()
241        .mask((!&mask1).into_array())
242        .vortex_expect("mask should succeed in conformance test");
243    let double_masked = first_masked
244        .mask((!&mask2).into_array())
245        .vortex_expect("mask should succeed in conformance test");
246
247    // Elements should be null if either mask is true
248    for i in 0..len {
249        if mask1_pattern[i] || mask2_pattern[i] {
250            assert!(
251                !double_masked
252                    .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx())
253                    .vortex_expect("is_valid should succeed in conformance test")
254            );
255        } else {
256            assert_eq!(
257                double_masked
258                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
259                    .vortex_expect("scalar_at should succeed in conformance test"),
260                array
261                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
262                    .vortex_expect("scalar_at should succeed in conformance test")
263                    .into_nullable()
264            );
265        }
266    }
267}
268
269/// Tests masking with nullable mask (nulls treated as false)
270fn test_nullable_mask_input(array: &ArrayRef) {
271    let len = array.len();
272    if len < 3 {
273        return; // Skip for very small arrays
274    }
275
276    // Create a nullable mask
277    let bool_values: Vec<bool> = (0..len).map(|i| i % 2 == 0).collect();
278    let validity_values: Vec<bool> = (0..len).map(|i| i % 3 != 0).collect();
279
280    let bool_array = BoolArray::from_iter(bool_values.clone());
281    let validity = crate::validity::Validity::from_iter(validity_values.clone());
282    let nullable_mask = BoolArray::new(bool_array.to_bit_buffer(), validity);
283
284    let mask_array =
285        nullable_mask.to_mask_fill_null_false(&mut LEGACY_SESSION.create_execution_ctx());
286    let masked = array
287        .clone()
288        .mask((!&mask_array).into_array())
289        .vortex_expect("mask should succeed in conformance test");
290
291    // Elements are masked only if the mask is true AND valid
292    for i in 0..len {
293        if bool_values[i] && validity_values[i] {
294            assert!(
295                !masked
296                    .is_valid(i, &mut LEGACY_SESSION.create_execution_ctx())
297                    .vortex_expect("is_valid should succeed in conformance test")
298            );
299        } else {
300            assert_eq!(
301                masked
302                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
303                    .vortex_expect("scalar_at should succeed in conformance test"),
304                array
305                    .execute_scalar(i, &mut LEGACY_SESSION.create_execution_ctx())
306                    .vortex_expect("scalar_at should succeed in conformance test")
307                    .into_nullable()
308            );
309        }
310    }
311}
vortex_array/compute/conformance/mask.rs

vortex_array/compute/conformance/
mask.rs
