Skip to main content

vortex_array/arrays/map/compute/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::array::ArrayView;
8use crate::arrays::ListView;
9use crate::arrays::map::Map;
10use crate::arrays::map::MapArrayExt;
11use crate::arrays::map::MapArraySlotsExt;
12use crate::arrays::map::compute::rebuild_map_from_array;
13use crate::executor::ExecutionCtx;
14use crate::scalar_fn::fns::mask::MaskKernel;
15use crate::scalar_fn::fns::mask::MaskReduce;
16
17impl MaskReduce for Map {
18    fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
19        let Some(entries) =
20            <ListView as MaskReduce>::mask(array.entries().as_::<ListView>(), mask)?
21        else {
22            return Ok(None);
23        };
24
25        rebuild_map_from_array(array.map_dtype().clone(), entries).map(Some)
26    }
27}
28
29impl MaskKernel for Map {
30    fn mask(
31        array: ArrayView<'_, Self>,
32        mask: &ArrayRef,
33        _ctx: &mut ExecutionCtx,
34    ) -> VortexResult<Option<ArrayRef>> {
35        <Self as MaskReduce>::mask(array, mask)
36    }
37}