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