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

vortex_array/arrays/null/
compute.rs
