vortex_array/aggregate_fn/fns/all_null/
mod.rs1use 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::AggregateFnId;
13use crate::aggregate_fn::AggregateFnVTable;
14use crate::aggregate_fn::EmptyOptions;
15use crate::dtype::DType;
16use crate::dtype::Nullability;
17use crate::scalar::Scalar;
18
19#[derive(Clone, Debug)]
23pub struct AllNull;
24
25impl AggregateFnVTable for AllNull {
26 type Options = EmptyOptions;
27 type Partial = bool;
28
29 fn id(&self) -> AggregateFnId {
30 static ID: CachedId = CachedId::new("vortex.all_null");
31 *ID
32 }
33
34 fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
35 Ok(Some(vec![]))
36 }
37
38 fn deserialize(
39 &self,
40 _metadata: &[u8],
41 _session: &VortexSession,
42 ) -> VortexResult<Self::Options> {
43 Ok(EmptyOptions)
44 }
45
46 fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {
47 Some(DType::Bool(Nullability::NonNullable))
48 }
49
50 fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
51 self.return_dtype(options, input_dtype)
52 }
53
54 fn empty_partial(
55 &self,
56 _options: &Self::Options,
57 _input_dtype: &DType,
58 ) -> VortexResult<Self::Partial> {
59 Ok(true)
60 }
61
62 fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> {
63 *partial &= bool::try_from(&other)?;
64 Ok(())
65 }
66
67 fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
68 Ok(Scalar::bool(*partial, Nullability::NonNullable))
69 }
70
71 fn reset(&self, partial: &mut Self::Partial) {
72 *partial = true;
73 }
74
75 fn is_saturated(&self, partial: &Self::Partial) -> bool {
76 !*partial
77 }
78
79 fn try_accumulate(
80 &self,
81 state: &mut Self::Partial,
82 batch: &ArrayRef,
83 ctx: &mut ExecutionCtx,
84 ) -> VortexResult<bool> {
85 *state &= batch.invalid_count(ctx)? == batch.len();
86 Ok(true)
87 }
88
89 fn accumulate(
90 &self,
91 partial: &mut Self::Partial,
92 batch: &Columnar,
93 ctx: &mut ExecutionCtx,
94 ) -> VortexResult<()> {
95 *partial &= match batch {
98 Columnar::Constant(c) => c.is_empty() || c.scalar().is_null(),
99 Columnar::Canonical(c) => {
100 let array = c.clone().into_array();
101 array.invalid_count(ctx)? == array.len()
102 }
103 };
104 Ok(())
105 }
106
107 fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
108 Ok(partials)
109 }
110
111 fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
112 self.to_scalar(partial)
113 }
114}
115
116#[cfg(test)]
117mod tests {
118 use vortex_error::VortexResult;
119
120 use crate::IntoArray;
121 use crate::VortexSessionExecute;
122 use crate::aggregate_fn::Accumulator;
123 use crate::aggregate_fn::DynAccumulator;
124 use crate::aggregate_fn::EmptyOptions;
125 use crate::aggregate_fn::fns::all_null::AllNull;
126 use crate::array_session;
127 use crate::arrays::PrimitiveArray;
128 use crate::dtype::DType;
129 use crate::dtype::Nullability;
130 use crate::dtype::PType;
131
132 #[test]
133 fn all_null_aggregate_fn() -> VortexResult<()> {
134 let mut ctx = array_session().create_execution_ctx();
135 let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
136 let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?;
137
138 let batch = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]).into_array();
139 acc.accumulate(&batch, &mut ctx)?;
140
141 assert!(bool::try_from(&acc.finish()?)?);
142 Ok(())
143 }
144
145 #[test]
146 fn all_null_false_with_non_nulls() -> VortexResult<()> {
147 let mut ctx = array_session().create_execution_ctx();
148 let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
149 let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?;
150
151 let batch = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
152 acc.accumulate(&batch, &mut ctx)?;
153
154 assert!(!bool::try_from(&acc.finish()?)?);
155 Ok(())
156 }
157
158 #[test]
159 fn all_null_true_for_empty_input() -> VortexResult<()> {
160 let mut ctx = array_session().create_execution_ctx();
161 let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
162 let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?;
163
164 let batch = PrimitiveArray::empty::<i32>(Nullability::Nullable).into_array();
165 acc.accumulate(&batch, &mut ctx)?;
166
167 assert!(bool::try_from(&acc.finish()?)?);
168 Ok(())
169 }
170}