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

vortex_array/arrays/masked/
array.rs
