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