vortex_array/aggregate_fn/fns/null_count/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6use vortex_error::vortex_err;
7use vortex_session::VortexSession;
8use vortex_session::registry::CachedId;
9
10use crate::ArrayRef;
11use crate::Columnar;
12use crate::ExecutionCtx;
13use crate::IntoArray;
14use crate::aggregate_fn::Accumulator;
15use crate::aggregate_fn::AggregateArgs;
16use crate::aggregate_fn::AggregateFnId;
17use crate::aggregate_fn::AggregateFnVTable;
18use crate::aggregate_fn::DynAccumulator;
19use crate::aggregate_fn::EmptyOptions;
20use crate::dtype::DType;
21use crate::dtype::Nullability::NonNullable;
22use crate::dtype::PType;
23use crate::expr::stats::Precision;
24use crate::expr::stats::Stat;
25use crate::expr::stats::StatsProvider;
26use crate::scalar::Scalar;
27
28/// Return the number of null values in an array.
29pub fn null_count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
30    if let Precision::Exact(null_count_scalar) = array.statistics().get(Stat::NullCount) {
31        return usize::try_from(&null_count_scalar)
32            .map_err(|e| vortex_err!("Failed to convert null count stat to usize: {e}"));
33    }
34
35    let mut acc = Accumulator::try_new(NullCount, EmptyOptions, array.dtype().clone())?;
36    acc.accumulate(array, ctx)?;
37    let result = acc.finish()?;
38
39    let count = result
40        .as_primitive()
41        .typed_value::<u64>()
42        .vortex_expect("null_count result should not be null");
43    // The accumulator caches the null count on `array` as a statistic.
44    Ok(usize::try_from(count).vortex_expect("Cannot be more nulls than usize::MAX"))
45}
46
47/// Count the number of null values in an array.
48///
49/// Applies to all types and returns a non-null `u64`.
50#[derive(Clone, Debug)]
51pub struct NullCount;
52
53impl AggregateFnVTable for NullCount {
54    type Options = EmptyOptions;
55    type Partial = u64;
56
57    fn id(&self) -> AggregateFnId {
58        static ID: CachedId = CachedId::new("vortex.null_count");
59        *ID
60    }
61
62    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
63        Ok(Some(vec![]))
64    }
65
66    fn deserialize(
67        &self,
68        _metadata: &[u8],
69        _session: &VortexSession,
70    ) -> VortexResult<Self::Options> {
71        Ok(EmptyOptions)
72    }
73
74    fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {
75        Some(DType::Primitive(PType::U64, NonNullable))
76    }
77
78    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
79        self.return_dtype(options, input_dtype)
80    }
81
82    fn empty_partial(
83        &self,
84        _args: AggregateArgs<'_, Self::Options>,
85    ) -> VortexResult<Self::Partial> {
86        Ok(0)
87    }
88
89    fn partial_from_scalar(
90        &self,
91        _args: AggregateArgs<'_, Self::Options>,
92        scalar: &Scalar,
93    ) -> VortexResult<Self::Partial> {
94        Ok(scalar
95            .as_primitive()
96            .typed_value::<u64>()
97            .vortex_expect("null_count partial should not be null"))
98    }
99
100    fn merge_partials(
101        &self,
102        _args: AggregateArgs<'_, Self::Options>,
103        first: Self::Partial,
104        second: Self::Partial,
105    ) -> VortexResult<Self::Partial> {
106        Ok(first + second)
107    }
108
109    fn to_scalar(
110        &self,
111        _args: AggregateArgs<'_, Self::Options>,
112        partial: &Self::Partial,
113    ) -> VortexResult<Scalar> {
114        Ok(Scalar::primitive(*partial, NonNullable))
115    }
116
117    #[inline]
118    fn is_saturated(
119        &self,
120        _args: AggregateArgs<'_, Self::Options>,
121        _partial: &Self::Partial,
122    ) -> bool {
123        false
124    }
125
126    fn try_accumulate(
127        &self,
128        _args: AggregateArgs<'_, Self::Options>,
129        state: &mut Self::Partial,
130        batch: &ArrayRef,
131        ctx: &mut ExecutionCtx,
132    ) -> VortexResult<bool> {
133        *state += batch.invalid_count(ctx)? as u64;
134        Ok(true)
135    }
136
137    fn accumulate(
138        &self,
139        _args: AggregateArgs<'_, Self::Options>,
140        partial: &mut Self::Partial,
141        batch: &Columnar,
142        ctx: &mut ExecutionCtx,
143    ) -> VortexResult<()> {
144        *partial += match batch {
145            Columnar::Constant(c) => {
146                if c.scalar().is_null() {
147                    c.len() as u64
148                } else {
149                    0
150                }
151            }
152            Columnar::Canonical(c) => c.clone().into_array().invalid_count(ctx)? as u64,
153        };
154        Ok(())
155    }
156
157    fn finalize(
158        &self,
159        _args: AggregateArgs<'_, Self::Options>,
160        partials: ArrayRef,
161    ) -> VortexResult<ArrayRef> {
162        Ok(partials)
163    }
164
165    fn finalize_scalar(
166        &self,
167        args: AggregateArgs<'_, Self::Options>,
168        partial: &Self::Partial,
169    ) -> VortexResult<Scalar> {
170        self.to_scalar(args, partial)
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use vortex_error::VortexResult;
177
178    use crate::IntoArray;
179    use crate::VortexSessionExecute;
180    use crate::aggregate_fn::Accumulator;
181    use crate::aggregate_fn::DynAccumulator;
182    use crate::aggregate_fn::EmptyOptions;
183    use crate::aggregate_fn::fns::null_count::NullCount;
184    use crate::aggregate_fn::fns::null_count::null_count;
185    use crate::array_session;
186    use crate::arrays::PrimitiveArray;
187    use crate::dtype::DType;
188    use crate::dtype::Nullability;
189    use crate::dtype::PType;
190    use crate::expr::stats::Precision;
191    use crate::expr::stats::Stat;
192    use crate::expr::stats::StatsProviderExt;
193
194    #[test]
195    fn null_count_with_nulls() -> VortexResult<()> {
196        let array =
197            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array();
198        let mut ctx = array_session().create_execution_ctx();
199
200        assert_eq!(null_count(&array, &mut ctx)?, 2);
201        assert_eq!(
202            array.statistics().get_as::<u64>(Stat::NullCount),
203            Precision::exact(2u64)
204        );
205        Ok(())
206    }
207
208    #[test]
209    fn null_count_multi_batch() -> VortexResult<()> {
210        let mut ctx = array_session().create_execution_ctx();
211        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
212        let mut acc = Accumulator::try_new(NullCount, EmptyOptions, dtype)?;
213
214        let batch1 = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
215        acc.accumulate(&batch1, &mut ctx)?;
216
217        let batch2 = PrimitiveArray::from_option_iter([None, Some(5i32), None]).into_array();
218        acc.accumulate(&batch2, &mut ctx)?;
219
220        let result = acc.finish()?;
221        assert_eq!(result.as_primitive().typed_value::<u64>(), Some(3));
222        Ok(())
223    }
224}
vortex_array/aggregate_fn/fns/null_count/mod.rs

vortex_array/aggregate_fn/fns/null_count/
mod.rs
