Skip to main content

vortex_array/arrays/dict/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::IntoArray;
8use crate::arrays::DictArray;
9use crate::arrays::DictVTable;
10use crate::arrays::ScalarFnArrayExt;
11use crate::compute::MaskReduce;
12use crate::expr::EmptyOptions;
13use crate::expr::mask::Mask as MaskExpr;
14
15impl MaskReduce for DictVTable {
16    fn mask(array: &DictArray, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
17        let masked_codes = MaskExpr.try_new_array(
18            array.codes().len(),
19            EmptyOptions,
20            [array.codes().clone(), mask.clone()],
21        )?;
22        // SAFETY: masking codes doesn't change dict invariants
23        Ok(Some(unsafe {
24            DictArray::new_unchecked(masked_codes, array.values().clone()).into_array()
25        }))
26    }
27}