vortex_array/arrays/null/
compute.rs

1use arrow_array::{ArrayRef as ArrowArrayRef, new_null_array};
2use arrow_schema::DataType;
3use vortex_dtype::{DType, match_each_integer_ptype};
4use vortex_error::{VortexResult, vortex_bail};
5use vortex_mask::Mask;
6use vortex_scalar::Scalar;
7
8use crate::arrays::NullEncoding;
9use crate::arrays::null::NullArray;
10use crate::compute::{
11    MaskFn, MinMaxFn, MinMaxResult, ScalarAtFn, SliceFn, TakeFn, ToArrowFn, UncompressedSizeFn,
12};
13use crate::nbytes::NBytes;
14use crate::variants::PrimitiveArrayTrait;
15use crate::vtable::ComputeVTable;
16use crate::{Array, ArrayRef, ToCanonical};
17
18impl ComputeVTable for NullEncoding {
19    fn mask_fn(&self) -> Option<&dyn MaskFn<&dyn Array>> {
20        Some(self)
21    }
22
23    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
24        Some(self)
25    }
26
27    fn slice_fn(&self) -> Option<&dyn SliceFn<&dyn Array>> {
28        Some(self)
29    }
30
31    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
32        Some(self)
33    }
34
35    fn to_arrow_fn(&self) -> Option<&dyn ToArrowFn<&dyn Array>> {
36        Some(self)
37    }
38
39    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
40        Some(self)
41    }
42
43    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
44        Some(self)
45    }
46}
47
48impl MaskFn<&NullArray> for NullEncoding {
49    fn mask(&self, array: &NullArray, _mask: Mask) -> VortexResult<ArrayRef> {
50        Ok(array.to_array().into_array())
51    }
52}
53
54impl SliceFn<&NullArray> for NullEncoding {
55    fn slice(&self, _array: &NullArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
56        Ok(NullArray::new(stop - start).into_array())
57    }
58}
59
60impl ScalarAtFn<&NullArray> for NullEncoding {
61    fn scalar_at(&self, _array: &NullArray, _index: usize) -> VortexResult<Scalar> {
62        Ok(Scalar::null(DType::Null))
63    }
64}
65
66impl TakeFn<&NullArray> for NullEncoding {
67    fn take(&self, array: &NullArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
68        let indices = indices.to_primitive()?;
69
70        // Enforce all indices are valid
71        match_each_integer_ptype!(indices.ptype(), |$T| {
72            for index in indices.as_slice::<$T>() {
73                if !((*index as usize) < array.len()) {
74                    vortex_bail!(OutOfBounds: *index as usize, 0, array.len());
75                }
76            }
77        });
78
79        Ok(NullArray::new(indices.len()).into_array())
80    }
81}
82
83impl ToArrowFn<&NullArray> for NullEncoding {
84    fn to_arrow(
85        &self,
86        array: &NullArray,
87        data_type: &DataType,
88    ) -> VortexResult<Option<ArrowArrayRef>> {
89        if data_type != &DataType::Null {
90            vortex_bail!("Unsupported data type: {data_type}");
91        }
92        Ok(Some(new_null_array(data_type, array.len())))
93    }
94}
95
96impl MinMaxFn<&NullArray> for NullEncoding {
97    fn min_max(&self, _array: &NullArray) -> VortexResult<Option<MinMaxResult>> {
98        Ok(None)
99    }
100}
101
102impl UncompressedSizeFn<&NullArray> for NullEncoding {
103    fn uncompressed_size(&self, array: &NullArray) -> VortexResult<usize> {
104        Ok(array.nbytes())
105    }
106}
107
108#[cfg(test)]
109mod test {
110    use vortex_buffer::buffer;
111    use vortex_dtype::DType;
112    use vortex_mask::Mask;
113
114    use crate::array::Array;
115    use crate::arrays::null::NullArray;
116    use crate::compute::{scalar_at, slice, take};
117    use crate::{ArrayExt, IntoArray};
118
119    #[test]
120    fn test_slice_nulls() {
121        let nulls = NullArray::new(10);
122        let sliced = slice(&nulls, 0, 4).unwrap().as_::<NullArray>().clone();
123
124        assert_eq!(sliced.len(), 4);
125        assert!(matches!(sliced.validity_mask().unwrap(), Mask::AllFalse(4)));
126    }
127
128    #[test]
129    fn test_take_nulls() {
130        let nulls = NullArray::new(10);
131        let taken = take(&nulls, &buffer![0u64, 2, 4, 6, 8].into_array())
132            .unwrap()
133            .as_::<NullArray>()
134            .clone();
135
136        assert_eq!(taken.len(), 5);
137        assert!(matches!(taken.validity_mask().unwrap(), Mask::AllFalse(5)));
138    }
139
140    #[test]
141    fn test_scalar_at_nulls() {
142        let nulls = NullArray::new(10);
143
144        let scalar = scalar_at(&nulls, 0).unwrap();
145        assert!(scalar.is_null());
146        assert_eq!(scalar.dtype().clone(), DType::Null);
147    }
148}
vortex_array/arrays/null/compute.rs

vortex_array/arrays/null/
compute.rs
