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