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::AggregateDTypes;
21use crate::aggregate_fn::AggregateFnId;
22use crate::aggregate_fn::AggregateFnVTable;
23use crate::aggregate_fn::DynAccumulator;
24use crate::aggregate_fn::EmptyOptions;
25use crate::dtype::DType;
26use crate::dtype::Nullability::NonNullable;
27use crate::dtype::PType;
28use crate::expr::stats::Precision;
29use crate::expr::stats::Stat;
30use crate::expr::stats::StatsProvider;
31use crate::scalar::Scalar;
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    // The accumulator caches the NaN count on `array` as a statistic.
68    Ok(usize::try_from(count).vortex_expect("Cannot be more nans than usize::MAX"))
69}
70
71/// Count the NaN values of an already-executed batch without touching any statistics.
72///
73/// Aggregates that delegate to [`NanCount`] use this on the columnar form of their input, so that
74/// results are only ever cached on the array the aggregate was requested for, never on a
75/// temporary canonical array.
76pub(crate) fn columnar_nan_count(batch: &Columnar, ctx: &mut ExecutionCtx) -> VortexResult<u64> {
77    // Non-float types have no NaN values.
78    if NanCount
79        .return_dtype(&EmptyOptions, batch.dtype())
80        .is_none()
81    {
82        return Ok(0);
83    }
84
85    let dtypes = AggregateDTypes::try_new(&NanCount, &EmptyOptions, batch.dtype().clone())?;
86    let args = dtypes.args(&EmptyOptions);
87    let mut count = NanCount.empty_partial(args)?;
88    NanCount.accumulate(args, &mut count, batch, ctx)?;
89    Ok(count)
90}
91
92/// Count the number of NaN values in an array.
93///
94/// Only applies to floating-point primitive types. Returns a `u64` count.
95/// Null values are not NaN, so if the array is all-invalid, the NaN count is zero.
96#[derive(Clone, Debug)]
97pub struct NanCount;
98
99impl AggregateFnVTable for NanCount {
100    type Options = EmptyOptions;
101    type Partial = u64;
102
103    fn id(&self) -> AggregateFnId {
104        static ID: CachedId = CachedId::new("vortex.nan_count");
105        *ID
106    }
107
108    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
109        Ok(Some(vec![]))
110    }
111
112    fn deserialize(
113        &self,
114        _metadata: &[u8],
115        _session: &VortexSession,
116    ) -> VortexResult<Self::Options> {
117        Ok(EmptyOptions)
118    }
119
120    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
121        if let DType::Primitive(ptype, ..) = input_dtype
122            && ptype.is_float()
123        {
124            Some(DType::Primitive(PType::U64, NonNullable))
125        } else {
126            None
127        }
128    }
129
130    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
131        self.return_dtype(options, input_dtype)
132    }
133
134    fn empty_partial(
135        &self,
136        _args: AggregateArgs<'_, Self::Options>,
137    ) -> VortexResult<Self::Partial> {
138        Ok(0)
139    }
140
141    fn partial_from_scalar(
142        &self,
143        _args: AggregateArgs<'_, Self::Options>,
144        scalar: &Scalar,
145    ) -> VortexResult<Self::Partial> {
146        Ok(scalar
147            .as_primitive()
148            .typed_value::<u64>()
149            .vortex_expect("nan_count partial should not be null"))
150    }
151
152    fn merge_partials(
153        &self,
154        _args: AggregateArgs<'_, Self::Options>,
155        first: Self::Partial,
156        second: Self::Partial,
157    ) -> VortexResult<Self::Partial> {
158        Ok(first + second)
159    }
160
161    fn to_scalar(
162        &self,
163        _args: AggregateArgs<'_, Self::Options>,
164        partial: &Self::Partial,
165    ) -> VortexResult<Scalar> {
166        Ok(Scalar::primitive(*partial, NonNullable))
167    }
168
169    #[inline]
170    fn is_saturated(
171        &self,
172        _args: AggregateArgs<'_, Self::Options>,
173        _partial: &Self::Partial,
174    ) -> bool {
175        false
176    }
177
178    fn accumulate(
179        &self,
180        _args: AggregateArgs<'_, Self::Options>,
181        partial: &mut Self::Partial,
182        batch: &Columnar,
183        ctx: &mut ExecutionCtx,
184    ) -> VortexResult<()> {
185        match batch {
186            Columnar::Constant(c) => {
187                if c.scalar().is_null() {
188                    // Null values are not NaN.
189                    return Ok(());
190                }
191                if c.scalar().as_primitive().is_nan() {
192                    *partial += c.len() as u64;
193                }
194                Ok(())
195            }
196            Columnar::Canonical(c) => match c {
197                Canonical::Primitive(p) => accumulate_primitive(partial, p, ctx),
198                _ => vortex_bail!(
199                    "Unsupported canonical type for nan_count: {}",
200                    batch.dtype()
201                ),
202            },
203        }
204    }
205
206    fn finalize(
207        &self,
208        _args: AggregateArgs<'_, Self::Options>,
209        partials: ArrayRef,
210    ) -> VortexResult<ArrayRef> {
211        Ok(partials)
212    }
213
214    fn finalize_scalar(
215        &self,
216        args: AggregateArgs<'_, Self::Options>,
217        partial: &Self::Partial,
218    ) -> VortexResult<Scalar> {
219        self.to_scalar(args, partial)
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use vortex_buffer::buffer;
226    use vortex_error::VortexResult;
227
228    use crate::IntoArray;
229    use crate::VortexSessionExecute;
230    use crate::aggregate_fn::Accumulator;
231    use crate::aggregate_fn::AggregateDTypes;
232    use crate::aggregate_fn::AggregateFnVTable;
233    use crate::aggregate_fn::DynAccumulator;
234    use crate::aggregate_fn::EmptyOptions;
235    use crate::aggregate_fn::fns::nan_count::NanCount;
236    use crate::aggregate_fn::fns::nan_count::nan_count;
237    use crate::array_session;
238    use crate::arrays::ChunkedArray;
239    use crate::arrays::ConstantArray;
240    use crate::arrays::PrimitiveArray;
241    use crate::dtype::DType;
242    use crate::dtype::Nullability;
243    use crate::dtype::PType;
244    use crate::validity::Validity;
245
246    #[test]
247    fn nan_count_multi_batch() -> VortexResult<()> {
248        let mut ctx = array_session().create_execution_ctx();
249        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
250        let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?;
251
252        let batch1 =
253            PrimitiveArray::new(buffer![f64::NAN, 1.0f64, f64::NAN], Validity::NonNullable)
254                .into_array();
255        acc.accumulate(&batch1, &mut ctx)?;
256
257        let batch2 =
258            PrimitiveArray::new(buffer![2.0f64, f64::NAN], Validity::NonNullable).into_array();
259        acc.accumulate(&batch2, &mut ctx)?;
260
261        let result = acc.finish()?;
262        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(3));
263        Ok(())
264    }
265
266    #[test]
267    fn nan_count_finish_resets_state() -> VortexResult<()> {
268        let mut ctx = array_session().create_execution_ctx();
269        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
270        let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?;
271
272        let batch1 =
273            PrimitiveArray::new(buffer![f64::NAN, 1.0f64], Validity::NonNullable).into_array();
274        acc.accumulate(&batch1, &mut ctx)?;
275        let result1 = acc.finish()?;
276        assert_eq!(result1.as_primitive().typed_value::<u64>(), Some(1));
277
278        let batch2 = PrimitiveArray::new(buffer![f64::NAN, f64::NAN, 2.0], Validity::NonNullable)
279            .into_array();
280        acc.accumulate(&batch2, &mut ctx)?;
281        let result2 = acc.finish()?;
282        assert_eq!(result2.as_primitive().typed_value::<u64>(), Some(2));
283        Ok(())
284    }
285
286    #[test]
287    fn nan_count_state_merge() -> VortexResult<()> {
288        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
289        let dtypes = AggregateDTypes::try_new(&NanCount, &EmptyOptions, dtype)?;
290
291        let state = NanCount.merge_partials(dtypes.args(&EmptyOptions), 5, 3)?;
292
293        let result = NanCount.to_scalar(dtypes.args(&EmptyOptions), &state)?;
294        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(8));
295        Ok(())
296    }
297
298    #[test]
299    fn nan_count_constant_nan() -> VortexResult<()> {
300        let array = ConstantArray::new(f64::NAN, 10);
301        let mut ctx = array_session().create_execution_ctx();
302        assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 10);
303        Ok(())
304    }
305
306    #[test]
307    fn nan_count_constant_non_nan() -> VortexResult<()> {
308        let array = ConstantArray::new(1.0f64, 10);
309        let mut ctx = array_session().create_execution_ctx();
310        assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 0);
311        Ok(())
312    }
313
314    #[test]
315    fn nan_count_empty() -> VortexResult<()> {
316        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
317        let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?;
318        let result = acc.finish()?;
319        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(0));
320        Ok(())
321    }
322
323    #[test]
324    fn nan_count_chunked() -> VortexResult<()> {
325        let chunk1 = PrimitiveArray::from_option_iter([Some(f64::NAN), None, Some(1.0)]);
326        let chunk2 = PrimitiveArray::from_option_iter([Some(f64::NAN), Some(f64::NAN), None]);
327        let dtype = chunk1.dtype().clone();
328        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
329        let mut ctx = array_session().create_execution_ctx();
330        assert_eq!(nan_count(&chunked.into_array(), &mut ctx)?, 3);
331        Ok(())
332    }
333
334    #[test]
335    fn nan_count_all_null() -> VortexResult<()> {
336        let p = PrimitiveArray::new(buffer![f64::NAN, f64::NAN, f64::NAN], Validity::AllInvalid);
337        let mut ctx = array_session().create_execution_ctx();
338        assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 0);
339        Ok(())
340    }
341}
vortex_array/aggregate_fn/fns/nan_count/mod.rs

vortex_array/aggregate_fn/fns/nan_count/
mod.rs
