vortex_expr/exprs/
is_null.rs

1use 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}