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

vortex_array/aggregate_fn/fns/all_nan/
mod.rs
