Skip to main content

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