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

vortex_array/aggregate_fn/fns/null_count/
mod.rs
