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

vortex_array/arrays/masked/
array.rs
