vortex_array/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::BitAnd;
5
6use vortex_error::VortexResult;
7use vortex_error::vortex_bail;
8use vortex_mask::Mask;
9
10use crate::ArrayRef;
11use crate::Executable;
12use crate::ExecutionCtx;
13use crate::IntoArray;
14use crate::arrays::BoolArray;
15use crate::arrays::Constant;
16use crate::columnar::Columnar;
17use crate::dtype::DType;
18
19impl Executable for Mask {
20    fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
21        if !matches!(array.dtype(), DType::Bool(_)) {
22            vortex_bail!("Mask array must have boolean dtype, not {}", array.dtype());
23        }
24
25        if let Some(constant) = array.as_opt::<Constant>() {
26            let mask_value = constant.scalar().as_bool().value().unwrap_or(false);
27            return Ok(Mask::new(array.len(), mask_value));
28        }
29
30        let array_len = array.len();
31        Ok(match array.execute(ctx)? {
32            Columnar::Constant(s) => {
33                Mask::new(array_len, s.scalar().as_bool().value().unwrap_or(false))
34            }
35            Columnar::Canonical(a) => {
36                let bool = a.into_array().execute::<BoolArray>(ctx)?;
37                let mask = bool.validity_mask()?;
38                let bits = bool.into_bit_buffer();
39                // To handle nullable boolean arrays, we treat nulls as false in the mask.
40                // TODO(ngates): is this correct? Feels like we should just force the caller to
41                //  pass non-nullable boolean arrays.
42                mask.bitand(&Mask::from(bits))
43            }
44        })
45    }
46}
vortex_array/mask.rs

vortex_array/
mask.rs
