vortex_array/aggregate_fn/fns/all_nan/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::Columnar;
8use crate::ExecutionCtx;
9use crate::IntoArray;
10use crate::aggregate_fn::AggregateFnId;
11use crate::aggregate_fn::AggregateFnVTable;
12use crate::aggregate_fn::EmptyOptions;
13use crate::aggregate_fn::fns::nan_count::nan_count;
14use crate::dtype::DType;
15use crate::dtype::Nullability;
16use crate::scalar::Scalar;
17
18/// Compute whether every value in an array is NaN.
19///
20/// Like other `all` aggregates, this is vacuously true for empty input.
21///
22/// This is a pruning aggregate, not just a convenience wrapper around
23/// [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount]. Pruning aggregates must prove a
24/// row-wise fact for every value in the scope, so their partials remain valid when a stats column is
25/// sliced or concatenated alongside the data. [`NanCount`][crate::aggregate_fn::fns::nan_count::NanCount]
26/// carries cross-row count information instead, so it is useful as a legacy storage format but not
27/// as the pruning expression itself.
28#[derive(Clone, Debug)]
29pub struct AllNan;
30
31impl AggregateFnVTable for AllNan {
32    type Options = EmptyOptions;
33    type Partial = bool;
34
35    fn id(&self) -> AggregateFnId {
36        AggregateFnId::new("vortex.all_nan")
37    }
38
39    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
40        Ok(None)
41    }
42
43    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
44        matches!(input_dtype, DType::Primitive(ptype, _) if ptype.is_float())
45            .then_some(DType::Bool(Nullability::Nullable))
46    }
47
48    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
49        self.return_dtype(options, input_dtype)
50    }
51
52    fn empty_partial(
53        &self,
54        _options: &Self::Options,
55        _input_dtype: &DType,
56    ) -> VortexResult<Self::Partial> {
57        Ok(true)
58    }
59
60    fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> {
61        *partial &= bool::try_from(&other)?;
62        Ok(())
63    }
64
65    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
66        Ok(Scalar::bool(*partial, Nullability::Nullable))
67    }
68
69    fn reset(&self, partial: &mut Self::Partial) {
70        *partial = true;
71    }
72
73    fn is_saturated(&self, partial: &Self::Partial) -> bool {
74        !*partial
75    }
76
77    fn try_accumulate(
78        &self,
79        state: &mut Self::Partial,
80        batch: &ArrayRef,
81        ctx: &mut ExecutionCtx,
82    ) -> VortexResult<bool> {
83        if !matches!(batch.dtype(), DType::Primitive(ptype, _) if ptype.is_float()) {
84            *state = false;
85            return Ok(true);
86        }
87
88        *state &= nan_count(batch, ctx)? == batch.len();
89        Ok(true)
90    }
91
92    fn accumulate(
93        &self,
94        partial: &mut Self::Partial,
95        batch: &Columnar,
96        ctx: &mut ExecutionCtx,
97    ) -> VortexResult<()> {
98        // Normal array dispatch is handled by `try_accumulate`, which always short-circuits.
99        // Keep this fallback in sync for direct Columnar accumulation paths.
100        let array = match batch {
101            Columnar::Constant(c) => c.clone().into_array(),
102            Columnar::Canonical(c) => c.clone().into_array(),
103        };
104        if !matches!(array.dtype(), DType::Primitive(ptype, _) if ptype.is_float()) {
105            *partial = false;
106            return Ok(());
107        }
108
109        *partial &= nan_count(&array, ctx)? == array.len();
110        Ok(())
111    }
112
113    fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
114        Ok(partials)
115    }
116
117    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
118        self.to_scalar(partial)
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use vortex_error::VortexResult;
125
126    use crate::IntoArray;
127    use crate::LEGACY_SESSION;
128    use crate::VortexSessionExecute;
129    use crate::aggregate_fn::Accumulator;
130    use crate::aggregate_fn::DynAccumulator;
131    use crate::aggregate_fn::EmptyOptions;
132    use crate::aggregate_fn::fns::all_nan::AllNan;
133    use crate::arrays::PrimitiveArray;
134    use crate::dtype::DType;
135    use crate::dtype::Nullability;
136    use crate::dtype::PType;
137
138    #[test]
139    fn all_nan_aggregate_fn() -> VortexResult<()> {
140        let mut ctx = LEGACY_SESSION.create_execution_ctx();
141        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
142        let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
143
144        let batch = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(f32::NAN)]).into_array();
145        acc.accumulate(&batch, &mut ctx)?;
146
147        assert!(bool::try_from(&acc.finish()?)?);
148        Ok(())
149    }
150
151    #[test]
152    fn all_nan_false_with_non_nan() -> VortexResult<()> {
153        let mut ctx = LEGACY_SESSION.create_execution_ctx();
154        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
155        let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
156
157        let batch = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32)]).into_array();
158        acc.accumulate(&batch, &mut ctx)?;
159
160        assert!(!bool::try_from(&acc.finish()?)?);
161        Ok(())
162    }
163
164    #[test]
165    fn all_nan_unsupported_for_non_float_values() -> VortexResult<()> {
166        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
167        assert!(Accumulator::try_new(AllNan, EmptyOptions, dtype).is_err());
168        Ok(())
169    }
170
171    #[test]
172    fn all_nan_false_with_null() -> VortexResult<()> {
173        let mut ctx = LEGACY_SESSION.create_execution_ctx();
174        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
175        let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
176
177        let batch = PrimitiveArray::from_option_iter([Some(f32::NAN), None]).into_array();
178        acc.accumulate(&batch, &mut ctx)?;
179
180        assert!(!bool::try_from(&acc.finish()?)?);
181        Ok(())
182    }
183
184    #[test]
185    fn all_nan_true_for_empty_float_values() -> VortexResult<()> {
186        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
187        let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
188
189        assert!(bool::try_from(&acc.finish()?)?);
190        Ok(())
191    }
192}
vortex_array/aggregate_fn/fns/all_nan/mod.rs

vortex_array/aggregate_fn/fns/all_nan/
mod.rs
