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::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#[derive(Clone, Debug)]
24pub struct AllNull;
25
26impl AggregateFnVTable for AllNull {
27 type Options = EmptyOptions;
28 type Partial = bool;
29
30 fn id(&self) -> AggregateFnId {
31 static ID: CachedId = CachedId::new("vortex.all_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)? == batch.len();
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 *partial &= match batch {
116 Columnar::Constant(c) => c.is_empty() || c.scalar().is_null(),
117 Columnar::Canonical(c) => {
118 let array = c.clone().into_array();
119 array.invalid_count(ctx)? == array.len()
120 }
121 };
122 Ok(())
123 }
124
125 fn finalize(
126 &self,
127 _args: AggregateArgs<'_, Self::Options>,
128 partials: ArrayRef,
129 ) -> VortexResult<ArrayRef> {
130 Ok(partials)
131 }
132
133 fn finalize_scalar(
134 &self,
135 args: AggregateArgs<'_, Self::Options>,
136 partial: &Self::Partial,
137 ) -> VortexResult<Scalar> {
138 self.to_scalar(args, partial)
139 }
140}
141
142#[cfg(test)]
143mod tests {
144 use vortex_error::VortexResult;
145
146 use crate::IntoArray;
147 use crate::VortexSessionExecute;
148 use crate::aggregate_fn::Accumulator;
149 use crate::aggregate_fn::DynAccumulator;
150 use crate::aggregate_fn::EmptyOptions;
151 use crate::aggregate_fn::fns::all_null::AllNull;
152 use crate::array_session;
153 use crate::arrays::PrimitiveArray;
154 use crate::dtype::DType;
155 use crate::dtype::Nullability;
156 use crate::dtype::PType;
157
158 #[test]
159 fn all_null_aggregate_fn() -> 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::from_option_iter::<i32, _>([None, None, None]).into_array();
165 acc.accumulate(&batch, &mut ctx)?;
166
167 assert!(bool::try_from(&acc.finish()?)?);
168 Ok(())
169 }
170
171 #[test]
172 fn all_null_false_with_non_nulls() -> VortexResult<()> {
173 let mut ctx = array_session().create_execution_ctx();
174 let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
175 let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?;
176
177 let batch = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).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_null_true_for_empty_input() -> VortexResult<()> {
186 let mut ctx = array_session().create_execution_ctx();
187 let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
188 let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?;
189
190 let batch = PrimitiveArray::empty::<i32>(Nullability::Nullable).into_array();
191 acc.accumulate(&batch, &mut ctx)?;
192
193 assert!(bool::try_from(&acc.finish()?)?);
194 Ok(())
195 }
196}