vortex_expr/exprs/
is_null.rs1use std::any::Any;
2use std::fmt::Display;
3use std::ops::Not;
4use std::sync::Arc;
5
6use vortex_array::arrays::{BoolArray, ConstantArray};
7use vortex_array::{Array, ArrayRef, IntoArray};
8use vortex_dtype::{DType, Nullability};
9use vortex_error::{VortexExpect, VortexResult};
10use vortex_mask::Mask;
11
12use crate::{AnalysisExpr, ExprRef, Scope, ScopeDType, VortexExpr};
13
14#[derive(Debug, Eq, Hash)]
15#[allow(clippy::derived_hash_with_manual_eq)]
16pub struct IsNull {
17 child: ExprRef,
18}
19
20impl IsNull {
21 pub fn new_expr(child: ExprRef) -> ExprRef {
22 Arc::new(Self { child })
23 }
24}
25
26impl PartialEq for IsNull {
27 fn eq(&self, other: &Self) -> bool {
28 self.child.eq(&other.child)
29 }
30}
31
32impl Display for IsNull {
33 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34 write!(f, "is_null({})", self.child)
35 }
36}
37
38#[cfg(feature = "proto")]
39pub(crate) mod proto {
40 use vortex_error::{VortexResult, vortex_bail};
41 use vortex_proto::expr::kind;
42 use vortex_proto::expr::kind::Kind;
43
44 use crate::is_null::IsNull;
45 use crate::{ExprDeserialize, ExprRef, ExprSerializable, Id};
46
47 pub(crate) struct IsNullSerde;
48
49 impl Id for IsNullSerde {
50 fn id(&self) -> &'static str {
51 "is_null"
52 }
53 }
54
55 impl ExprDeserialize for IsNullSerde {
56 fn deserialize(&self, kind: &Kind, children: Vec<ExprRef>) -> VortexResult<ExprRef> {
57 let Kind::IsNull(kind::IsNull {}) = kind else {
58 vortex_bail!("wrong kind {:?}, want is_null", kind)
59 };
60
61 Ok(IsNull::new_expr(children[0].clone()))
62 }
63 }
64
65 impl ExprSerializable for IsNull {
66 fn id(&self) -> &'static str {
67 IsNullSerde.id()
68 }
69
70 fn serialize_kind(&self) -> VortexResult<Kind> {
71 Ok(Kind::IsNull(kind::IsNull {}))
72 }
73 }
74}
75
76impl AnalysisExpr for IsNull {}
77
78impl VortexExpr for IsNull {
79 fn as_any(&self) -> &dyn Any {
80 self
81 }
82
83 fn unchecked_evaluate(&self, ctx: &Scope) -> VortexResult<ArrayRef> {
84 let array = self.child.unchecked_evaluate(ctx)?;
85 match array.validity_mask()? {
86 Mask::AllTrue(len) => Ok(ConstantArray::new(false, len).into_array()),
87 Mask::AllFalse(len) => Ok(ConstantArray::new(true, len).into_array()),
88 Mask::Values(mask) => Ok(BoolArray::from(mask.boolean_buffer().not()).into_array()),
89 }
90 }
91
92 fn children(&self) -> Vec<&ExprRef> {
93 vec![&self.child]
94 }
95
96 fn replacing_children(self: Arc<Self>, mut children: Vec<ExprRef>) -> ExprRef {
97 Self::new_expr(
98 children
99 .pop()
100 .vortex_expect("IsNull::replacing_children should have one child"),
101 )
102 }
103
104 fn return_dtype(&self, _ctx: &ScopeDType) -> VortexResult<DType> {
105 Ok(DType::Bool(Nullability::NonNullable))
106 }
107}
108
109pub fn is_null(child: ExprRef) -> ExprRef {
110 IsNull::new_expr(child)
111}
112
113#[cfg(test)]
114mod tests {
115 use vortex_array::IntoArray;
116 use vortex_array::arrays::{PrimitiveArray, StructArray};
117 use vortex_dtype::{DType, Nullability};
118 use vortex_scalar::Scalar;
119
120 use crate::is_null::is_null;
121 use crate::{Scope, ScopeDType, get_item, root, test_harness};
122
123 #[test]
124 fn dtype() {
125 let dtype = test_harness::struct_dtype();
126 assert_eq!(
127 is_null(root())
128 .return_dtype(&ScopeDType::new(dtype))
129 .unwrap(),
130 DType::Bool(Nullability::NonNullable)
131 );
132 }
133
134 #[test]
135 fn replace_children() {
136 let expr = is_null(root());
137 let _ = expr.replacing_children(vec![root()]);
138 }
139
140 #[test]
141 fn evaluate_mask() {
142 let test_array =
143 PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
144 .into_array();
145 let expected = [false, true, false, true, false];
146
147 let result = is_null(root())
148 .evaluate(&Scope::new(test_array.clone()))
149 .unwrap();
150
151 assert_eq!(result.len(), test_array.len());
152 assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
153
154 for (i, expected_value) in expected.iter().enumerate() {
155 assert_eq!(
156 result.scalar_at(i).unwrap(),
157 Scalar::bool(*expected_value, Nullability::NonNullable)
158 );
159 }
160 }
161
162 #[test]
163 fn evaluate_all_false() {
164 let test_array = PrimitiveArray::from_iter(vec![1, 2, 3, 4, 5]).into_array();
165
166 let result = is_null(root())
167 .evaluate(&Scope::new(test_array.clone()))
168 .unwrap();
169
170 assert_eq!(result.len(), test_array.len());
171 assert_eq!(
172 result.as_constant().unwrap(),
173 Scalar::bool(false, Nullability::NonNullable)
174 );
175 }
176
177 #[test]
178 fn evaluate_all_true() {
179 let test_array =
180 PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
181 .into_array();
182
183 let result = is_null(root())
184 .evaluate(&Scope::new(test_array.clone()))
185 .unwrap();
186
187 assert_eq!(result.len(), test_array.len());
188 assert_eq!(
189 result.as_constant().unwrap(),
190 Scalar::bool(true, Nullability::NonNullable)
191 );
192 }
193
194 #[test]
195 fn evaluate_struct() {
196 let test_array = StructArray::from_fields(&[(
197 "a",
198 PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
199 .into_array(),
200 )])
201 .unwrap()
202 .into_array();
203 let expected = [false, true, false, true, false];
204
205 let result = is_null(get_item("a", root()))
206 .evaluate(&Scope::new(test_array.clone()))
207 .unwrap();
208
209 assert_eq!(result.len(), test_array.len());
210 assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
211
212 for (i, expected_value) in expected.iter().enumerate() {
213 assert_eq!(
214 result.scalar_at(i).unwrap(),
215 Scalar::bool(*expected_value, Nullability::NonNullable)
216 );
217 }
218 }
219}