vortex_array/arrays/null/
compute.rs

1use vortex_dtype::match_each_integer_ptype;
2use vortex_error::{VortexResult, vortex_bail};
3use vortex_mask::Mask;
4
5use crate::arrays::NullVTable;
6use crate::arrays::null::NullArray;
7use crate::compute::{
8    FilterKernel, FilterKernelAdapter, MaskKernel, MaskKernelAdapter, MinMaxKernel,
9    MinMaxKernelAdapter, MinMaxResult, TakeKernel, TakeKernelAdapter,
10};
11use crate::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
12
13impl FilterKernel for NullVTable {
14    fn filter(&self, _array: &Self::Array, mask: &Mask) -> VortexResult<ArrayRef> {
15        Ok(NullArray::new(mask.true_count()).into_array())
16    }
17}
18
19register_kernel!(FilterKernelAdapter(NullVTable).lift());
20
21impl MaskKernel for NullVTable {
22    fn mask(&self, array: &NullArray, _mask: &Mask) -> VortexResult<ArrayRef> {
23        Ok(array.to_array())
24    }
25}
26
27register_kernel!(MaskKernelAdapter(NullVTable).lift());
28
29impl TakeKernel for NullVTable {
30    #[allow(clippy::cast_possible_truncation)]
31    fn take(&self, array: &NullArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
32        let indices = indices.to_primitive()?;
33
34        // Enforce all indices are valid
35        match_each_integer_ptype!(indices.ptype(), |T| {
36            for index in indices.as_slice::<T>() {
37                if (*index as usize) >= array.len() {
38                    vortex_bail!(OutOfBounds: *index as usize, 0, array.len());
39                }
40            }
41        });
42
43        Ok(NullArray::new(indices.len()).into_array())
44    }
45}
46
47register_kernel!(TakeKernelAdapter(NullVTable).lift());
48
49impl MinMaxKernel for NullVTable {
50    fn min_max(&self, _array: &NullArray) -> VortexResult<Option<MinMaxResult>> {
51        Ok(None)
52    }
53}
54
55register_kernel!(MinMaxKernelAdapter(NullVTable).lift());
56
57#[cfg(test)]
58mod test {
59    use vortex_buffer::buffer;
60    use vortex_dtype::DType;
61    use vortex_mask::Mask;
62
63    use crate::arrays::null::NullArray;
64    use crate::compute::take;
65    use crate::{IntoArray, ToCanonical};
66
67    #[test]
68    fn test_slice_nulls() {
69        let nulls = NullArray::new(10);
70        let sliced = nulls.slice(0, 4).unwrap().to_null().unwrap();
71
72        assert_eq!(sliced.len(), 4);
73        assert!(matches!(sliced.validity_mask().unwrap(), Mask::AllFalse(4)));
74    }
75
76    #[test]
77    fn test_take_nulls() {
78        let nulls = NullArray::new(10);
79        let taken = take(nulls.as_ref(), &buffer![0u64, 2, 4, 6, 8].into_array())
80            .unwrap()
81            .to_null()
82            .unwrap();
83
84        assert_eq!(taken.len(), 5);
85        assert!(matches!(taken.validity_mask().unwrap(), Mask::AllFalse(5)));
86    }
87
88    #[test]
89    fn test_scalar_at_nulls() {
90        let nulls = NullArray::new(10);
91
92        let scalar = nulls.scalar_at(0).unwrap();
93        assert!(scalar.is_null());
94        assert_eq!(scalar.dtype().clone(), DType::Null);
95    }
96}