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

vortex_array/arrays/masked/
array.rs
