vortex_array/arrays/primitive/compute/
nan_count.rs

1use vortex_dtype::{NativePType, match_each_float_ptype};
2use vortex_error::VortexResult;
3use vortex_mask::Mask;
4
5use crate::arrays::{PrimitiveArray, PrimitiveVTable};
6use crate::compute::{NaNCountKernel, NaNCountKernelAdapter};
7use crate::register_kernel;
8
9impl NaNCountKernel for PrimitiveVTable {
10    fn nan_count(&self, array: &PrimitiveArray) -> VortexResult<usize> {
11        Ok(match_each_float_ptype!(array.ptype(), |F| {
12            compute_nan_count_with_validity(array.as_slice::<F>(), array.validity_mask()?)
13        }))
14    }
15}
16
17register_kernel!(NaNCountKernelAdapter(PrimitiveVTable).lift());
18
19#[inline]
20fn compute_nan_count_with_validity<T: NativePType>(values: &[T], validity: Mask) -> usize {
21    match validity {
22        Mask::AllTrue(_) => values.iter().filter(|v| v.is_nan()).count(),
23        Mask::AllFalse(_) => 0,
24        Mask::Values(v) => values
25            .iter()
26            .zip(v.boolean_buffer().iter())
27            .filter_map(|(v, m)| m.then_some(v))
28            .filter(|v| v.is_nan())
29            .count(),
30    }
31}
32
33#[cfg(test)]
34mod tests {
35    use vortex_buffer::buffer;
36
37    use crate::arrays::PrimitiveArray;
38    use crate::compute::nan_count;
39    use crate::validity::Validity;
40
41    #[test]
42    fn primitive_nan_count() {
43        let p = PrimitiveArray::new(
44            buffer![
45                -f32::NAN,
46                f32::NAN,
47                0.1,
48                1.1,
49                -0.0,
50                f32::INFINITY,
51                f32::NEG_INFINITY
52            ],
53            Validity::NonNullable,
54        );
55        assert_eq!(nan_count(p.as_ref()).unwrap(), 2);
56    }
57}