vortex_array/aggregate_fn/fns/all_non_null/
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::dtype::DType;
17use crate::dtype::Nullability;
18use crate::scalar::Scalar;
19
20/// Compute whether every value in an array is non-null.
21///
22/// Like other `all` aggregates, this is vacuously true for empty input.
23#[derive(Clone, Debug)]
24pub struct AllNonNull;
25
26impl AggregateFnVTable for AllNonNull {
27    type Options = EmptyOptions;
28    type Partial = bool;
29
30    fn id(&self) -> AggregateFnId {
31        static ID: CachedId = CachedId::new("vortex.all_non_null");
32        *ID
33    }
34
35    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
36        Ok(Some(vec![]))
37    }
38
39    fn deserialize(
40        &self,
41        _metadata: &[u8],
42        _session: &VortexSession,
43    ) -> VortexResult<Self::Options> {
44        Ok(EmptyOptions)
45    }
46
47    fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {
48        Some(DType::Bool(Nullability::NonNullable))
49    }
50
51    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
52        self.return_dtype(options, input_dtype)
53    }
54
55    fn empty_partial(
56        &self,
57        _args: AggregateArgs<'_, Self::Options>,
58    ) -> VortexResult<Self::Partial> {
59        Ok(true)
60    }
61
62    fn partial_from_scalar(
63        &self,
64        _args: AggregateArgs<'_, Self::Options>,
65        scalar: &Scalar,
66    ) -> VortexResult<Self::Partial> {
67        bool::try_from(scalar)
68    }
69
70    fn merge_partials(
71        &self,
72        _args: AggregateArgs<'_, Self::Options>,
73        first: Self::Partial,
74        second: Self::Partial,
75    ) -> VortexResult<Self::Partial> {
76        Ok(first && second)
77    }
78
79    fn to_scalar(
80        &self,
81        _args: AggregateArgs<'_, Self::Options>,
82        partial: &Self::Partial,
83    ) -> VortexResult<Scalar> {
84        Ok(Scalar::bool(*partial, Nullability::NonNullable))
85    }
86
87    fn is_saturated(
88        &self,
89        _args: AggregateArgs<'_, Self::Options>,
90        partial: &Self::Partial,
91    ) -> bool {
92        !*partial
93    }
94
95    fn try_accumulate(
96        &self,
97        _args: AggregateArgs<'_, Self::Options>,
98        state: &mut Self::Partial,
99        batch: &ArrayRef,
100        ctx: &mut ExecutionCtx,
101    ) -> VortexResult<bool> {
102        *state &= batch.invalid_count(ctx)? == 0;
103        Ok(true)
104    }
105
106    fn accumulate(
107        &self,
108        _args: AggregateArgs<'_, Self::Options>,
109        partial: &mut Self::Partial,
110        batch: &Columnar,
111        ctx: &mut ExecutionCtx,
112    ) -> VortexResult<()> {
113        // Normal array dispatch is handled by `try_accumulate`, which always short-circuits.
114        // Keep this fallback in sync for direct Columnar accumulation paths.
115        *partial &= match batch {
116            Columnar::Constant(c) => c.is_empty() || !c.scalar().is_null(),
117            Columnar::Canonical(c) => c.clone().into_array().invalid_count(ctx)? == 0,
118        };
119        Ok(())
120    }
121
122    fn finalize(
123        &self,
124        _args: AggregateArgs<'_, Self::Options>,
125        partials: ArrayRef,
126    ) -> VortexResult<ArrayRef> {
127        Ok(partials)
128    }
129
130    fn finalize_scalar(
131        &self,
132        args: AggregateArgs<'_, Self::Options>,
133        partial: &Self::Partial,
134    ) -> VortexResult<Scalar> {
135        self.to_scalar(args, partial)
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use vortex_error::VortexResult;
142
143    use crate::IntoArray;
144    use crate::VortexSessionExecute;
145    use crate::aggregate_fn::Accumulator;
146    use crate::aggregate_fn::DynAccumulator;
147    use crate::aggregate_fn::EmptyOptions;
148    use crate::aggregate_fn::fns::all_non_null::AllNonNull;
149    use crate::array_session;
150    use crate::arrays::PrimitiveArray;
151    use crate::dtype::DType;
152    use crate::dtype::Nullability;
153    use crate::dtype::PType;
154
155    #[test]
156    fn all_non_null_aggregate_fn() -> VortexResult<()> {
157        let mut ctx = array_session().create_execution_ctx();
158        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
159        let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?;
160
161        let batch = PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3)]).into_array();
162        acc.accumulate(&batch, &mut ctx)?;
163
164        assert!(bool::try_from(&acc.finish()?)?);
165        Ok(())
166    }
167
168    #[test]
169    fn all_non_null_false_with_nulls() -> VortexResult<()> {
170        let mut ctx = array_session().create_execution_ctx();
171        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
172        let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?;
173
174        let batch = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
175        acc.accumulate(&batch, &mut ctx)?;
176
177        assert!(!bool::try_from(&acc.finish()?)?);
178        Ok(())
179    }
180
181    #[test]
182    fn all_non_null_true_for_empty_input() -> VortexResult<()> {
183        let mut ctx = array_session().create_execution_ctx();
184        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
185        let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?;
186
187        let batch = PrimitiveArray::empty::<i32>(Nullability::Nullable).into_array();
188        acc.accumulate(&batch, &mut ctx)?;
189
190        assert!(bool::try_from(&acc.finish()?)?);
191        Ok(())
192    }
193}
vortex_array/aggregate_fn/fns/all_non_null/mod.rs

vortex_array/aggregate_fn/fns/all_non_null/
mod.rs
