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::ExecutionCtx;
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    #[slot(0)]
26    pub child: ArrayRef,
27    /// The validity bitmap defining which elements are non-null.
28    #[slot(1)]
29    pub validity: Option<ArrayRef>,
30}
31
32#[derive(Clone, Debug)]
33pub struct MaskedData;
34
35impl Display for MaskedData {
36    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
37        Ok(())
38    }
39}
40
41pub trait MaskedArrayExt: TypedArrayRef<Masked> + MaskedArraySlotsExt {
42    fn masked_validity(&self) -> Validity {
43        child_to_validity(
44            self.slots()[MaskedSlots::VALIDITY].as_ref(),
45            self.dtype().nullability(),
46        )
47    }
48}
49impl<T: TypedArrayRef<Masked>> MaskedArrayExt for T {}
50
51impl MaskedData {
52    pub(crate) fn try_new(
53        child_len: usize,
54        child_all_valid: bool,
55        validity: Validity,
56    ) -> VortexResult<Self> {
57        if matches!(validity, Validity::NonNullable) {
58            vortex_bail!("MaskedArray must have nullable validity, got {validity:?}")
59        }
60
61        if !child_all_valid {
62            vortex_bail!("MaskedArray children must not have nulls");
63        }
64
65        if let Some(validity_len) = validity.maybe_len()
66            && validity_len != child_len
67        {
68            vortex_bail!("Validity must be the same length as a MaskedArray's child");
69        }
70
71        // MaskedArray's nullability is determined solely by its validity, not the child's dtype.
72        // The child can have nullable dtype but must not have any actual null values.
73        Ok(Self)
74    }
75}
76
77impl Array<Masked> {
78    /// Constructs a new `MaskedArray`.
79    pub fn try_new(
80        child: ArrayRef,
81        validity: Validity,
82        ctx: &mut ExecutionCtx,
83    ) -> VortexResult<Self> {
84        let dtype = child.dtype().as_nullable();
85        let len = child.len();
86        let validity_slot = validity_to_child(&validity, len);
87        let data = MaskedData::try_new(len, child.all_valid(ctx)?, validity)?;
88        Ok(unsafe {
89            Array::from_parts_unchecked(ArrayParts::new(
90                Masked,
91                dtype,
92                len,
93                data,
94                smallvec![Some(child), validity_slot],
95            ))
96        })
97    }
98
99    /// Constructs a new `MaskedArray` without validating the child or validity.
100    ///
101    /// # Safety
102    ///
103    /// The caller must ensure that `child` contains no nulls, that `validity` is nullable, and
104    /// that an array-backed `validity` has the same length as `child`.
105    pub unsafe fn new_unchecked(child: ArrayRef, validity: Validity) -> Self {
106        let dtype = child.dtype().as_nullable();
107        let len = child.len();
108        let validity_slot = validity_to_child(&validity, len);
109        unsafe {
110            Array::from_parts_unchecked(ArrayParts::new(
111                Masked,
112                dtype,
113                len,
114                MaskedData,
115                smallvec![Some(child), validity_slot],
116            ))
117        }
118    }
119}
vortex_array/arrays/masked/array.rs

vortex_array/arrays/masked/
array.rs
