vortex_array/arrays/masked/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10
11use crate::ArrayRef;
12use crate::array::Array;
13use crate::array::ArrayParts;
14use crate::array::TypedArrayRef;
15use crate::array::child_to_validity;
16use crate::array::validity_to_child;
17use crate::arrays::Masked;
18use crate::validity::Validity;
19
20/// The underlying child array being masked.
21pub(super) const CHILD_SLOT: usize = 0;
22/// The validity bitmap defining which elements are non-null.
23pub(super) const VALIDITY_SLOT: usize = 1;
24pub(super) const NUM_SLOTS: usize = 2;
25pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child", "validity"];
26
27#[derive(Clone, Debug)]
28pub struct MaskedData;
29
30impl Display for MaskedData {
31    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
32        Ok(())
33    }
34}
35
36pub trait MaskedArrayExt: TypedArrayRef<Masked> {
37    fn child(&self) -> &ArrayRef {
38        self.as_ref().slots()[CHILD_SLOT]
39            .as_ref()
40            .vortex_expect("validated masked child slot")
41    }
42
43    fn validity_child(&self) -> Option<&ArrayRef> {
44        self.as_ref().slots()[VALIDITY_SLOT].as_ref()
45    }
46
47    fn masked_validity(&self) -> Validity {
48        child_to_validity(
49            &self.as_ref().slots()[VALIDITY_SLOT],
50            self.as_ref().dtype().nullability(),
51        )
52    }
53
54    fn masked_validity_mask(&self) -> vortex_mask::Mask {
55        self.masked_validity().to_mask(self.as_ref().len())
56    }
57}
58impl<T: TypedArrayRef<Masked>> MaskedArrayExt for T {}
59
60impl MaskedData {
61    pub(crate) fn try_new(
62        child_len: usize,
63        child_all_valid: bool,
64        validity: Validity,
65    ) -> VortexResult<Self> {
66        if matches!(validity, Validity::NonNullable) {
67            vortex_bail!("MaskedArray must have nullable validity, got {validity:?}")
68        }
69
70        if !child_all_valid {
71            vortex_bail!("MaskedArray children must not have nulls");
72        }
73
74        if let Some(validity_len) = validity.maybe_len()
75            && validity_len != child_len
76        {
77            vortex_bail!("Validity must be the same length as a MaskedArray's child");
78        }
79
80        // MaskedArray's nullability is determined solely by its validity, not the child's dtype.
81        // The child can have nullable dtype but must not have any actual null values.
82        Ok(Self)
83    }
84}
85
86impl Array<Masked> {
87    /// Constructs a new `MaskedArray`.
88    pub fn try_new(child: ArrayRef, validity: Validity) -> VortexResult<Self> {
89        let dtype = child.dtype().as_nullable();
90        let len = child.len();
91        let validity_slot = validity_to_child(&validity, len);
92        let data = MaskedData::try_new(len, child.all_valid()?, validity)?;
93        Ok(unsafe {
94            Array::from_parts_unchecked(
95                ArrayParts::new(Masked, dtype, len, data)
96                    .with_slots(vec![Some(child), validity_slot]),
97            )
98        })
99    }
100}
vortex_array/arrays/masked/array.rs

vortex_array/arrays/masked/
array.rs
