Skip to main content

vortex_array/arrays/varbinview/compute/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_error::VortexResult;
7
8use crate::ArrayRef;
9use crate::IntoArray;
10use crate::array::ArrayView;
11use crate::arrays::VarBinView;
12use crate::arrays::VarBinViewArray;
13use crate::scalar_fn::fns::mask::MaskReduce;
14use crate::validity::Validity;
15
16impl MaskReduce for VarBinView {
17    fn mask(array: ArrayView<'_, VarBinView>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
18        // SAFETY: masking the validity does not affect the invariants
19        unsafe {
20            Ok(Some(
21                VarBinViewArray::new_handle_unchecked(
22                    array.views_handle().clone(),
23                    Arc::clone(array.data_buffers()),
24                    array.dtype().as_nullable(),
25                    array.validity()?.and(Validity::Array(mask.clone()))?,
26                )
27                .into_array(),
28            ))
29        }
30    }
31}
32
33#[cfg(test)]
34mod tests {
35    use crate::IntoArray;
36    use crate::VortexSessionExecute;
37    use crate::array_session;
38    use crate::arrays::VarBinViewArray;
39    use crate::compute::conformance::mask::test_mask_conformance;
40
41    #[test]
42    fn take_mask_var_bin_view_array() {
43        test_mask_conformance(
44            &VarBinViewArray::from_iter_str(["one", "two", "three", "four", "five"]).into_array(),
45            &mut array_session().create_execution_ctx(),
46        );
47
48        test_mask_conformance(
49            &VarBinViewArray::from_iter_nullable_str([
50                Some("one"),
51                None,
52                Some("three"),
53                Some("four"),
54                Some("five"),
55            ])
56            .into_array(),
57            &mut array_session().create_execution_ctx(),
58        );
59    }
60}