vortex_array/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_mask::Mask;
7
8use crate::ArrayRef;
9use crate::Executable;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::arrays::BoolArray;
13use crate::builtins::ArrayBuiltins;
14use crate::columnar::Columnar;
15use crate::dtype::DType;
16use crate::dtype::Nullability;
17
18impl Executable for Mask {
19    /// Executes a boolean array into a [`Mask`].
20    ///
21    /// The array must have a non-nullable boolean dtype. To execute a nullable boolean array,
22    /// coercing null elements to `false`, first call
23    /// [`ArrayRef::fill_null(false)`](crate::builtins::ArrayBuiltins::fill_null).
24    fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
25        if !matches!(array.dtype(), DType::Bool(Nullability::NonNullable)) {
26            vortex_bail!(
27                "Mask array must have boolean(NonNullable) dtype, not {}",
28                array.dtype()
29            );
30        }
31
32        let array_len = array.len();
33        Ok(match array.execute(ctx)? {
34            Columnar::Constant(s) => {
35                Mask::new(array_len, s.scalar().as_bool().value().unwrap_or(false))
36            }
37            Columnar::Canonical(a) => {
38                let bool = a.into_array().execute::<BoolArray>(ctx)?;
39                Mask::from(bool.into_bit_buffer())
40            }
41        })
42    }
43}
44
45#[deprecated = "Use array.fill_null(false)?.execute::<Mask>(ctx)"]
46pub struct NullAsFalse(ArrayRef);
47impl ArrayRef {
48    #[deprecated = "Use array.fill_null(false)?.execute::<Mask>(ctx)"]
49    #[allow(deprecated)]
50    pub fn null_as_false(self) -> NullAsFalse {
51        NullAsFalse(self)
52    }
53}
54#[allow(deprecated)]
55impl NullAsFalse {
56    #[deprecated = "Use array.fill_null(false)?.execute::<Mask>(ctx)"]
57    pub fn execute(self, ctx: &mut ExecutionCtx) -> VortexResult<Mask> {
58        self.0.fill_null(false)?.execute::<Mask>(ctx)
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use vortex_error::VortexResult;
65    use vortex_mask::Mask;
66
67    use crate::ExecutionCtx;
68    use crate::IntoArray;
69    use crate::VortexSessionExecute;
70    use crate::array_session;
71    use crate::arrays::BoolArray;
72    use crate::arrays::ConstantArray;
73    use crate::builtins::ArrayBuiltins;
74    use crate::dtype::DType;
75    use crate::dtype::Nullability;
76    use crate::scalar::Scalar;
77
78    fn ctx() -> ExecutionCtx {
79        array_session().create_execution_ctx()
80    }
81
82    #[test]
83    fn mask_non_nullable() -> VortexResult<()> {
84        let array = BoolArray::from_iter([true, false, true]).into_array();
85        let mask = array.execute::<Mask>(&mut ctx())?;
86        assert_eq!(mask, Mask::from_iter([true, false, true]));
87        Ok(())
88    }
89
90    #[test]
91    fn mask_rejects_nullable() {
92        let array = BoolArray::from_iter([Some(true), None]).into_array();
93        assert!(array.execute::<Mask>(&mut ctx()).is_err());
94    }
95
96    #[test]
97    fn fill_null_then_mask_coerces_nulls() -> VortexResult<()> {
98        let array = BoolArray::from_iter([Some(true), None, Some(false), None]).into_array();
99        let mask = array.fill_null(false)?.execute::<Mask>(&mut ctx())?;
100        assert_eq!(mask, Mask::from_iter([true, false, false, false]));
101        Ok(())
102    }
103
104    #[test]
105    fn fill_null_then_mask_null_constant() -> VortexResult<()> {
106        let array =
107            ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array();
108        let mask = array.fill_null(false)?.execute::<Mask>(&mut ctx())?;
109        assert_eq!(mask, Mask::new_false(4));
110        Ok(())
111    }
112}
vortex_array/mask.rs

vortex_array/
mask.rs
