vortex_array/arrays/null/
compute.rs

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