vortex_array/arrays/primitive/compute/
nan_count.rs

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

vortex_array/arrays/primitive/compute/
nan_count.rs
