vortex_array/aggregate_fn/fns/nan_count/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod primitive;
5
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_err;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use self::primitive::accumulate_primitive;
14use crate::ArrayRef;
15use crate::Canonical;
16use crate::Columnar;
17use crate::ExecutionCtx;
18use crate::aggregate_fn::Accumulator;
19use crate::aggregate_fn::AggregateArgs;
20use crate::aggregate_fn::AggregateFnId;
21use crate::aggregate_fn::AggregateFnVTable;
22use crate::aggregate_fn::DynAccumulator;
23use crate::aggregate_fn::EmptyOptions;
24use crate::dtype::DType;
25use crate::dtype::Nullability::NonNullable;
26use crate::dtype::PType;
27use crate::expr::stats::Precision;
28use crate::expr::stats::Stat;
29use crate::expr::stats::StatsProvider;
30use crate::scalar::Scalar;
31use crate::scalar::ScalarValue;
32
33/// Return the number of NaN values in an array.
34///
35/// Null values are not NaN and are not counted.
36///
37/// See [`NanCount`] for details.
38pub fn nan_count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
39    // Short-circuit using cached array statistics.
40    if let Precision::Exact(nan_count_scalar) = array.statistics().get(Stat::NaNCount) {
41        return usize::try_from(&nan_count_scalar)
42            .map_err(|e| vortex_err!("Failed to convert NaN count stat to usize: {e}"));
43    }
44
45    // Short-circuit for non-float types.
46    if NanCount
47        .return_dtype(&EmptyOptions, array.dtype())
48        .is_none()
49    {
50        return Ok(0);
51    }
52
53    // Short-circuit for empty arrays or all-null arrays.
54    if array.is_empty() || array.valid_count(ctx)? == 0 {
55        return Ok(0);
56    }
57
58    // Compute using Accumulator<NanCount>.
59    let mut acc = Accumulator::try_new(NanCount, EmptyOptions, array.dtype().clone())?;
60    acc.accumulate(array, ctx)?;
61    let result = acc.finish()?;
62
63    let count = result
64        .as_primitive()
65        .typed_value::<u64>()
66        .vortex_expect("nan_count result should not be null");
67    let count_usize = usize::try_from(count).vortex_expect("Cannot be more nans than usize::MAX");
68
69    // Cache the computed NaN count as a statistic.
70    array
71        .statistics()
72        .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(count)));
73
74    Ok(count_usize)
75}
76
77/// Count the number of NaN values in an array.
78///
79/// Only applies to floating-point primitive types. Returns a `u64` count.
80/// Null values are not NaN, so if the array is all-invalid, the NaN count is zero.
81#[derive(Clone, Debug)]
82pub struct NanCount;
83
84impl AggregateFnVTable for NanCount {
85    type Options = EmptyOptions;
86    type Partial = u64;
87
88    fn id(&self) -> AggregateFnId {
89        static ID: CachedId = CachedId::new("vortex.nan_count");
90        *ID
91    }
92
93    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
94        Ok(Some(vec![]))
95    }
96
97    fn deserialize(
98        &self,
99        _metadata: &[u8],
100        _session: &VortexSession,
101    ) -> VortexResult<Self::Options> {
102        Ok(EmptyOptions)
103    }
104
105    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
106        if let DType::Primitive(ptype, ..) = input_dtype
107            && ptype.is_float()
108        {
109            Some(DType::Primitive(PType::U64, NonNullable))
110        } else {
111            None
112        }
113    }
114
115    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
116        self.return_dtype(options, input_dtype)
117    }
118
119    fn empty_partial(
120        &self,
121        _args: AggregateArgs<'_, Self::Options>,
122    ) -> VortexResult<Self::Partial> {
123        Ok(0)
124    }
125
126    fn partial_from_scalar(
127        &self,
128        _args: AggregateArgs<'_, Self::Options>,
129        scalar: Scalar,
130    ) -> VortexResult<Self::Partial> {
131        Ok(scalar
132            .as_primitive()
133            .typed_value::<u64>()
134            .vortex_expect("nan_count partial should not be null"))
135    }
136
137    fn merge_partials(
138        &self,
139        _args: AggregateArgs<'_, Self::Options>,
140        first: Self::Partial,
141        second: Self::Partial,
142    ) -> VortexResult<Self::Partial> {
143        Ok(first + second)
144    }
145
146    fn to_scalar(
147        &self,
148        _args: AggregateArgs<'_, Self::Options>,
149        partial: &Self::Partial,
150    ) -> VortexResult<Scalar> {
151        Ok(Scalar::primitive(*partial, NonNullable))
152    }
153
154    #[inline]
155    fn is_saturated(
156        &self,
157        _args: AggregateArgs<'_, Self::Options>,
158        _partial: &Self::Partial,
159    ) -> bool {
160        false
161    }
162
163    fn accumulate(
164        &self,
165        _args: AggregateArgs<'_, Self::Options>,
166        partial: &mut Self::Partial,
167        batch: &Columnar,
168        ctx: &mut ExecutionCtx,
169    ) -> VortexResult<()> {
170        match batch {
171            Columnar::Constant(c) => {
172                if c.scalar().is_null() {
173                    // Null values are not NaN.
174                    return Ok(());
175                }
176                if c.scalar().as_primitive().is_nan() {
177                    *partial += c.len() as u64;
178                }
179                Ok(())
180            }
181            Columnar::Canonical(c) => match c {
182                Canonical::Primitive(p) => accumulate_primitive(partial, p, ctx),
183                _ => vortex_bail!(
184                    "Unsupported canonical type for nan_count: {}",
185                    batch.dtype()
186                ),
187            },
188        }
189    }
190
191    fn finalize(
192        &self,
193        _args: AggregateArgs<'_, Self::Options>,
194        partials: ArrayRef,
195    ) -> VortexResult<ArrayRef> {
196        Ok(partials)
197    }
198
199    fn finalize_scalar(
200        &self,
201        args: AggregateArgs<'_, Self::Options>,
202        partial: &Self::Partial,
203    ) -> VortexResult<Scalar> {
204        self.to_scalar(args, partial)
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use vortex_buffer::buffer;
211    use vortex_error::VortexResult;
212
213    use crate::IntoArray;
214    use crate::VortexSessionExecute;
215    use crate::aggregate_fn::Accumulator;
216    use crate::aggregate_fn::AggregateDTypes;
217    use crate::aggregate_fn::AggregateFnVTable;
218    use crate::aggregate_fn::DynAccumulator;
219    use crate::aggregate_fn::EmptyOptions;
220    use crate::aggregate_fn::fns::nan_count::NanCount;
221    use crate::aggregate_fn::fns::nan_count::nan_count;
222    use crate::array_session;
223    use crate::arrays::ChunkedArray;
224    use crate::arrays::ConstantArray;
225    use crate::arrays::PrimitiveArray;
226    use crate::dtype::DType;
227    use crate::dtype::Nullability;
228    use crate::dtype::PType;
229    use crate::validity::Validity;
230
231    #[test]
232    fn nan_count_multi_batch() -> VortexResult<()> {
233        let mut ctx = array_session().create_execution_ctx();
234        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
235        let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?;
236
237        let batch1 =
238            PrimitiveArray::new(buffer![f64::NAN, 1.0f64, f64::NAN], Validity::NonNullable)
239                .into_array();
240        acc.accumulate(&batch1, &mut ctx)?;
241
242        let batch2 =
243            PrimitiveArray::new(buffer![2.0f64, f64::NAN], Validity::NonNullable).into_array();
244        acc.accumulate(&batch2, &mut ctx)?;
245
246        let result = acc.finish()?;
247        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(3));
248        Ok(())
249    }
250
251    #[test]
252    fn nan_count_finish_resets_state() -> VortexResult<()> {
253        let mut ctx = array_session().create_execution_ctx();
254        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
255        let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?;
256
257        let batch1 =
258            PrimitiveArray::new(buffer![f64::NAN, 1.0f64], Validity::NonNullable).into_array();
259        acc.accumulate(&batch1, &mut ctx)?;
260        let result1 = acc.finish()?;
261        assert_eq!(result1.as_primitive().typed_value::<u64>(), Some(1));
262
263        let batch2 = PrimitiveArray::new(buffer![f64::NAN, f64::NAN, 2.0], Validity::NonNullable)
264            .into_array();
265        acc.accumulate(&batch2, &mut ctx)?;
266        let result2 = acc.finish()?;
267        assert_eq!(result2.as_primitive().typed_value::<u64>(), Some(2));
268        Ok(())
269    }
270
271    #[test]
272    fn nan_count_state_merge() -> VortexResult<()> {
273        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
274        let dtypes = AggregateDTypes::try_new(&NanCount, &EmptyOptions, dtype)?;
275
276        let state = NanCount.merge_partials(dtypes.args(&EmptyOptions), 5, 3)?;
277
278        let result = NanCount.to_scalar(dtypes.args(&EmptyOptions), &state)?;
279        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(8));
280        Ok(())
281    }
282
283    #[test]
284    fn nan_count_constant_nan() -> VortexResult<()> {
285        let array = ConstantArray::new(f64::NAN, 10);
286        let mut ctx = array_session().create_execution_ctx();
287        assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 10);
288        Ok(())
289    }
290
291    #[test]
292    fn nan_count_constant_non_nan() -> VortexResult<()> {
293        let array = ConstantArray::new(1.0f64, 10);
294        let mut ctx = array_session().create_execution_ctx();
295        assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 0);
296        Ok(())
297    }
298
299    #[test]
300    fn nan_count_empty() -> VortexResult<()> {
301        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
302        let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?;
303        let result = acc.finish()?;
304        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(0));
305        Ok(())
306    }
307
308    #[test]
309    fn nan_count_chunked() -> VortexResult<()> {
310        let chunk1 = PrimitiveArray::from_option_iter([Some(f64::NAN), None, Some(1.0)]);
311        let chunk2 = PrimitiveArray::from_option_iter([Some(f64::NAN), Some(f64::NAN), None]);
312        let dtype = chunk1.dtype().clone();
313        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
314        let mut ctx = array_session().create_execution_ctx();
315        assert_eq!(nan_count(&chunked.into_array(), &mut ctx)?, 3);
316        Ok(())
317    }
318
319    #[test]
320    fn nan_count_all_null() -> VortexResult<()> {
321        let p = PrimitiveArray::from_option_iter::<f64, _>([None, None, None]);
322        let mut ctx = array_session().create_execution_ctx();
323        assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 0);
324        Ok(())
325    }
326}
vortex_array/aggregate_fn/fns/nan_count/mod.rs

vortex_array/aggregate_fn/fns/nan_count/
mod.rs
