vortex_array/arrays/primitive/compute/
nan_count.rs1use 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}