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

vortex_array/arrays/null/
compute.rs
