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;
10use vortex_session::VortexSession;
11use vortex_vector::Datum;
12
13use crate::Array;
14use crate::ArrayRef;
15use crate::executor::VectorExecutor;
16
17/// Executor for exporting a Vortex [`Mask`] from an [`ArrayRef`].
18pub trait MaskExecutor {
19    /// Execute the array to produce a mask.
20    fn execute_mask(&self, session: &VortexSession) -> VortexResult<Mask>;
21}
22
23impl MaskExecutor for ArrayRef {
24    fn execute_mask(&self, session: &VortexSession) -> VortexResult<Mask> {
25        if !matches!(self.dtype(), DType::Bool(_)) {
26            vortex_bail!("Mask array must have boolean dtype, not {}", self.dtype());
27        }
28
29        Ok(match self.execute_datum(session)? {
30            Datum::Scalar(s) => Mask::new(self.len(), s.as_bool().value().unwrap_or(false)),
31            Datum::Vector(v) => {
32                let (bits, mask) = v.into_bool().into_parts();
33                // To handle nullable boolean arrays, we treat nulls as false in the mask.
34                // TODO(ngates): is this correct? Feels like we should just force the caller to
35                //  pass non-nullable boolean arrays.
36                mask.bitand(&Mask::from(bits))
37            }
38        })
39    }
40}