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

vortex_array/arrays/primitive/compute/
nan_count.rs
