vortex_array/aggregate_fn/fns/all_non_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 not 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 AllNonNan;
30
31impl AggregateFnVTable for AllNonNan {
32    type Options = EmptyOptions;
33    type Partial = bool;
34
35    fn id(&self) -> AggregateFnId {
36        AggregateFnId::new("vortex.all_non_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        *state &= nan_count(batch, ctx)? == 0;
84        Ok(true)
85    }
86
87    fn accumulate(
88        &self,
89        partial: &mut Self::Partial,
90        batch: &Columnar,
91        ctx: &mut ExecutionCtx,
92    ) -> VortexResult<()> {
93        // Normal array dispatch is handled by `try_accumulate`, which always short-circuits.
94        // Keep this fallback in sync for direct Columnar accumulation paths.
95        let array = match batch {
96            Columnar::Constant(c) => c.clone().into_array(),
97            Columnar::Canonical(c) => c.clone().into_array(),
98        };
99        *partial &= nan_count(&array, ctx)? == 0;
100        Ok(())
101    }
102
103    fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
104        Ok(partials)
105    }
106
107    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
108        self.to_scalar(partial)
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use vortex_error::VortexResult;
115
116    use crate::IntoArray;
117    use crate::LEGACY_SESSION;
118    use crate::VortexSessionExecute;
119    use crate::aggregate_fn::Accumulator;
120    use crate::aggregate_fn::DynAccumulator;
121    use crate::aggregate_fn::EmptyOptions;
122    use crate::aggregate_fn::fns::all_non_nan::AllNonNan;
123    use crate::arrays::PrimitiveArray;
124    use crate::dtype::DType;
125    use crate::dtype::Nullability;
126    use crate::dtype::PType;
127
128    #[test]
129    fn all_non_nan_aggregate_fn() -> VortexResult<()> {
130        let mut ctx = LEGACY_SESSION.create_execution_ctx();
131        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
132        let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?;
133
134        let batch = PrimitiveArray::from_option_iter([Some(1.0f32), None, Some(3.0)]).into_array();
135        acc.accumulate(&batch, &mut ctx)?;
136
137        assert!(bool::try_from(&acc.finish()?)?);
138        Ok(())
139    }
140
141    #[test]
142    fn all_non_nan_false_with_nan() -> VortexResult<()> {
143        let mut ctx = LEGACY_SESSION.create_execution_ctx();
144        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
145        let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?;
146
147        let batch = PrimitiveArray::from_option_iter([Some(1.0f32), Some(f32::NAN)]).into_array();
148        acc.accumulate(&batch, &mut ctx)?;
149
150        assert!(!bool::try_from(&acc.finish()?)?);
151        Ok(())
152    }
153
154    #[test]
155    fn all_non_nan_unsupported_for_non_float() -> VortexResult<()> {
156        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
157        assert!(Accumulator::try_new(AllNonNan, EmptyOptions, dtype).is_err());
158        Ok(())
159    }
160
161    #[test]
162    fn all_non_nan_true_for_empty_float() -> VortexResult<()> {
163        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
164        let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?;
165
166        assert!(bool::try_from(&acc.finish()?)?);
167        Ok(())
168    }
169
170    #[test]
171    fn all_non_nan_true_with_nulls() -> VortexResult<()> {
172        let mut ctx = LEGACY_SESSION.create_execution_ctx();
173        let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
174        let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?;
175
176        let batch = PrimitiveArray::from_option_iter([Some(1.0f32), None]).into_array();
177        acc.accumulate(&batch, &mut ctx)?;
178
179        assert!(bool::try_from(&acc.finish()?)?);
180        Ok(())
181    }
182}
vortex_array/aggregate_fn/fns/all_non_nan/mod.rs

vortex_array/aggregate_fn/fns/all_non_nan/
mod.rs
