vortex_array/arrays/masked/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::DType;
5use vortex_error::VortexResult;
6use vortex_error::vortex_bail;
7
8use crate::ArrayRef;
9use crate::compute::mask;
10use crate::stats::ArrayStats;
11use crate::validity::Validity;
12
13#[derive(Clone, Debug)]
14pub struct MaskedArray {
15    pub(super) child: ArrayRef,
16    pub(super) validity: Validity,
17    pub(super) dtype: DType,
18    pub(super) stats: ArrayStats,
19}
20
21impl MaskedArray {
22    pub fn try_new(child: ArrayRef, validity: Validity) -> VortexResult<Self> {
23        if matches!(validity, Validity::NonNullable) {
24            vortex_bail!("MaskedArray must have nullable validity, got {validity:?}")
25        }
26
27        if !child.all_valid() {
28            vortex_bail!("MaskedArray children must not have nulls");
29        }
30
31        if let Some(validity_len) = validity.maybe_len()
32            && validity_len != child.len()
33        {
34            vortex_bail!("Validity must be the same length as a MaskedArray's child");
35        }
36
37        // MaskedArray's nullability is determined solely by its validity, not the child's dtype.
38        // The child can have nullable dtype but must not have any actual null values.
39        let dtype = child.dtype().as_nullable();
40
41        Ok(Self {
42            child,
43            validity,
44            dtype,
45            stats: ArrayStats::default(),
46        })
47    }
48
49    pub fn child(&self) -> &ArrayRef {
50        &self.child
51    }
52
53    pub(crate) fn masked_child(&self) -> VortexResult<ArrayRef> {
54        // Invert the validity mask - we want to set values to null where validity is false.
55        let inverted_mask = !self.validity.to_mask(self.len());
56        mask(&self.child, &inverted_mask)
57    }
58}