vortex_array/aggregate_fn/fns/count/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod grouped;
5pub(crate) use grouped::CountGroupedKernel;
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_session::registry::CachedId;
10
11use crate::ArrayRef;
12use crate::Columnar;
13use crate::ExecutionCtx;
14use crate::aggregate_fn::AggregateArgs;
15use crate::aggregate_fn::AggregateFnId;
16use crate::aggregate_fn::AggregateFnVTable;
17use crate::aggregate_fn::NumericalAggregateOpts;
18use crate::aggregate_fn::fns::nan_count::nan_count;
19use crate::dtype::DType;
20use crate::dtype::Nullability;
21use crate::dtype::PType;
22use crate::scalar::Scalar;
23
24/// Count the number of non-null elements in an array.
25///
26/// Applies to all types. Returns a `u64` count.
27/// The identity value is zero.
28///
29/// For float inputs, NaN handling is controlled by [`NumericalAggregateOpts`]: with `skip_nans` (the
30/// default) NaN values are treated as missing and excluded from the count, otherwise they are
31/// counted like any other non-null value.
32#[derive(Clone, Debug)]
33pub struct Count;
34
35impl AggregateFnVTable for Count {
36    type Options = NumericalAggregateOpts;
37    type Partial = u64;
38
39    fn id(&self) -> AggregateFnId {
40        static ID: CachedId = CachedId::new("vortex.count");
41        *ID
42    }
43
44    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
45        vortex_bail!("Count is not yet serializable");
46    }
47
48    fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {
49        Some(DType::Primitive(PType::U64, Nullability::NonNullable))
50    }
51
52    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
53        self.return_dtype(options, input_dtype)
54    }
55
56    fn empty_partial(
57        &self,
58        _args: AggregateArgs<'_, Self::Options>,
59    ) -> VortexResult<Self::Partial> {
60        Ok(0)
61    }
62
63    fn partial_from_scalar(
64        &self,
65        _args: AggregateArgs<'_, Self::Options>,
66        scalar: Scalar,
67    ) -> VortexResult<Self::Partial> {
68        Ok(scalar
69            .as_primitive()
70            .typed_value::<u64>()
71            .vortex_expect("count partial should not be null"))
72    }
73
74    fn merge_partials(
75        &self,
76        _args: AggregateArgs<'_, Self::Options>,
77        first: Self::Partial,
78        second: Self::Partial,
79    ) -> VortexResult<Self::Partial> {
80        Ok(first + second)
81    }
82
83    fn to_scalar(
84        &self,
85        _args: AggregateArgs<'_, Self::Options>,
86        partial: &Self::Partial,
87    ) -> VortexResult<Scalar> {
88        Ok(Scalar::primitive(*partial, Nullability::NonNullable))
89    }
90
91    #[inline]
92    fn is_saturated(
93        &self,
94        _args: AggregateArgs<'_, Self::Options>,
95        _partial: &Self::Partial,
96    ) -> bool {
97        false
98    }
99
100    fn try_accumulate(
101        &self,
102        args: AggregateArgs<'_, Self::Options>,
103        state: &mut Self::Partial,
104        batch: &ArrayRef,
105        ctx: &mut ExecutionCtx,
106    ) -> VortexResult<bool> {
107        let mut count = batch.valid_count(ctx)? as u64;
108        // NaN values are excluded from the count of a float input when they are skipped.
109        if args.options.skip_nans && args.dtype.is_float() {
110            // `nan_count` shortcircuits on an exact `Stat::NaNCount` before scanning the batch.
111            count = count.saturating_sub(nan_count(batch, ctx)? as u64);
112        }
113        *state += count;
114        Ok(true)
115    }
116
117    fn accumulate(
118        &self,
119        _args: AggregateArgs<'_, Self::Options>,
120        _partial: &mut Self::Partial,
121        _batch: &Columnar,
122        _ctx: &mut ExecutionCtx,
123    ) -> VortexResult<()> {
124        unreachable!("Count::try_accumulate handles all arrays")
125    }
126
127    fn finalize(
128        &self,
129        _args: AggregateArgs<'_, Self::Options>,
130        partials: ArrayRef,
131    ) -> VortexResult<ArrayRef> {
132        Ok(partials)
133    }
134
135    fn finalize_scalar(
136        &self,
137        args: AggregateArgs<'_, Self::Options>,
138        partial: &Self::Partial,
139    ) -> VortexResult<Scalar> {
140        self.to_scalar(args, partial)
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use std::sync::LazyLock;
147
148    use vortex_buffer::buffer;
149    use vortex_error::VortexExpect;
150    use vortex_error::VortexResult;
151    use vortex_session::VortexSession;
152
153    use crate::ArrayRef;
154    use crate::ExecutionCtx;
155    use crate::IntoArray;
156    use crate::VortexSessionExecute;
157    use crate::aggregate_fn::Accumulator;
158    use crate::aggregate_fn::AggregateDTypes;
159    use crate::aggregate_fn::AggregateFnVTable;
160    use crate::aggregate_fn::DynAccumulator;
161    use crate::aggregate_fn::NumericalAggregateOpts;
162    use crate::aggregate_fn::fns::count::Count;
163    use crate::arrays::ChunkedArray;
164    use crate::arrays::ConstantArray;
165    use crate::arrays::PrimitiveArray;
166    use crate::dtype::DType;
167    use crate::dtype::Nullability;
168    use crate::dtype::PType;
169    use crate::expr::stats::Precision;
170    use crate::expr::stats::Stat;
171    use crate::scalar::Scalar;
172    use crate::scalar::ScalarValue;
173    use crate::validity::Validity;
174
175    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
176
177    pub fn count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
178        let mut acc = Accumulator::try_new(
179            Count,
180            NumericalAggregateOpts::default(),
181            array.dtype().clone(),
182        )?;
183        acc.accumulate(array, ctx)?;
184        let result = acc.finish()?;
185
186        Ok(usize::try_from(
187            result
188                .as_primitive()
189                .typed_value::<u64>()
190                .vortex_expect("count result should not be null"),
191        )?)
192    }
193
194    #[test]
195    fn count_all_valid() -> VortexResult<()> {
196        let array =
197            PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::NonNullable).into_array();
198        let mut ctx = SESSION.create_execution_ctx();
199        assert_eq!(count(&array, &mut ctx)?, 5);
200        Ok(())
201    }
202
203    #[test]
204    fn count_with_nulls() -> VortexResult<()> {
205        let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)])
206            .into_array();
207        let mut ctx = SESSION.create_execution_ctx();
208        assert_eq!(count(&array, &mut ctx)?, 3);
209        Ok(())
210    }
211
212    #[test]
213    fn count_all_null() -> VortexResult<()> {
214        let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]).into_array();
215        let mut ctx = SESSION.create_execution_ctx();
216        assert_eq!(count(&array, &mut ctx)?, 0);
217        Ok(())
218    }
219
220    #[test]
221    fn count_empty() -> VortexResult<()> {
222        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
223        let mut acc = Accumulator::try_new(Count, NumericalAggregateOpts::default(), dtype)?;
224        let result = acc.finish()?;
225        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(0));
226        Ok(())
227    }
228
229    #[test]
230    fn count_multi_batch() -> VortexResult<()> {
231        let mut ctx = SESSION.create_execution_ctx();
232        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
233        let mut acc = Accumulator::try_new(Count, NumericalAggregateOpts::default(), dtype)?;
234
235        let batch1 = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
236        acc.accumulate(&batch1, &mut ctx)?;
237
238        let batch2 = PrimitiveArray::from_option_iter([None, Some(5i32)]).into_array();
239        acc.accumulate(&batch2, &mut ctx)?;
240
241        let result = acc.finish()?;
242        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(3));
243        Ok(())
244    }
245
246    #[test]
247    fn count_finish_resets_state() -> VortexResult<()> {
248        let mut ctx = SESSION.create_execution_ctx();
249        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
250        let mut acc = Accumulator::try_new(Count, NumericalAggregateOpts::default(), dtype)?;
251
252        let batch1 = PrimitiveArray::from_option_iter([Some(1i32), None]).into_array();
253        acc.accumulate(&batch1, &mut ctx)?;
254        let result1 = acc.finish()?;
255        assert_eq!(result1.as_primitive().typed_value::<u64>(), Some(1));
256
257        let batch2 = PrimitiveArray::from_option_iter([Some(2i32), Some(3), None]).into_array();
258        acc.accumulate(&batch2, &mut ctx)?;
259        let result2 = acc.finish()?;
260        assert_eq!(result2.as_primitive().typed_value::<u64>(), Some(2));
261        Ok(())
262    }
263
264    #[test]
265    fn count_state_merge() -> VortexResult<()> {
266        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
267        let options = NumericalAggregateOpts::default();
268        let dtypes = AggregateDTypes::try_new(&Count, &options, dtype)?;
269
270        let state = Count.merge_partials(dtypes.args(&options), 5, 3)?;
271
272        let result = Count.to_scalar(dtypes.args(&options), &state)?;
273        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(8));
274        Ok(())
275    }
276
277    fn count_with_options(
278        array: &ArrayRef,
279        ctx: &mut ExecutionCtx,
280        options: NumericalAggregateOpts,
281    ) -> VortexResult<u64> {
282        let mut acc = Accumulator::try_new(Count, options, array.dtype().clone())?;
283        acc.accumulate(array, ctx)?;
284        Ok(acc
285            .finish()?
286            .as_primitive()
287            .typed_value::<u64>()
288            .vortex_expect("count result should not be null"))
289    }
290
291    #[test]
292    fn count_float_excludes_nans_by_default() -> VortexResult<()> {
293        let array =
294            PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(3.0)])
295                .into_array();
296        let mut ctx = SESSION.create_execution_ctx();
297        assert_eq!(count(&array, &mut ctx)?, 2);
298        Ok(())
299    }
300
301    #[test]
302    fn count_float_includes_nans_when_not_skipping() -> VortexResult<()> {
303        let array =
304            PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(3.0)])
305                .into_array();
306        let mut ctx = SESSION.create_execution_ctx();
307        assert_eq!(
308            count_with_options(&array, &mut ctx, NumericalAggregateOpts::include_nans())?,
309            3
310        );
311        Ok(())
312    }
313
314    #[test]
315    fn count_float_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> {
316        // The array has no NaNs; a planted exact NaNCount stat proves the count is derived from
317        // the stat rather than a scan.
318        let array =
319            PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0, 4.0], Validity::NonNullable).into_array();
320        array
321            .statistics()
322            .set(Stat::NaNCount, Precision::Exact(ScalarValue::from(3u64)));
323        let mut ctx = SESSION.create_execution_ctx();
324        assert_eq!(count(&array, &mut ctx)?, 1);
325        Ok(())
326    }
327
328    #[test]
329    fn count_constant_nan() -> VortexResult<()> {
330        let array = ConstantArray::new(f64::NAN, 5).into_array();
331        let mut ctx = SESSION.create_execution_ctx();
332        assert_eq!(count(&array, &mut ctx)?, 0);
333        assert_eq!(
334            count_with_options(&array, &mut ctx, NumericalAggregateOpts::include_nans())?,
335            5
336        );
337        Ok(())
338    }
339
340    #[test]
341    fn count_constant_non_null() -> VortexResult<()> {
342        let array = ConstantArray::new(42i32, 10);
343        let mut ctx = SESSION.create_execution_ctx();
344        assert_eq!(count(&array.into_array(), &mut ctx)?, 10);
345        Ok(())
346    }
347
348    #[test]
349    fn count_constant_null() -> VortexResult<()> {
350        let array = ConstantArray::new(
351            Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
352            10,
353        );
354        let mut ctx = SESSION.create_execution_ctx();
355        assert_eq!(count(&array.into_array(), &mut ctx)?, 0);
356        Ok(())
357    }
358
359    #[test]
360    fn count_chunked() -> VortexResult<()> {
361        let chunk1 = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]);
362        let chunk2 = PrimitiveArray::from_option_iter([None, Some(5i32), None]);
363        let dtype = chunk1.dtype().clone();
364        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
365        let mut ctx = SESSION.create_execution_ctx();
366        assert_eq!(count(&chunked.into_array(), &mut ctx)?, 3);
367        Ok(())
368    }
369}
vortex_array/aggregate_fn/fns/count/mod.rs

vortex_array/aggregate_fn/fns/count/
mod.rs
