vortex_expr/exprs/
is_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Not;
5
6use vortex_array::arrays::{BoolArray, ConstantArray};
7use vortex_array::{Array, ArrayRef, DeserializeMetadata, EmptyMetadata, IntoArray};
8use vortex_dtype::{DType, Nullability};
9use vortex_error::{VortexResult, vortex_bail};
10use vortex_mask::Mask;
11
12use crate::display::{DisplayAs, DisplayFormat};
13use crate::{AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable};
14
15vtable!(IsNull);
16
17#[allow(clippy::derived_hash_with_manual_eq)]
18#[derive(Clone, Debug, Hash, Eq)]
19pub struct IsNullExpr {
20    child: ExprRef,
21}
22
23impl PartialEq for IsNullExpr {
24    fn eq(&self, other: &Self) -> bool {
25        self.child.eq(&other.child)
26    }
27}
28
29pub struct IsNullExprEncoding;
30
31impl VTable for IsNullVTable {
32    type Expr = IsNullExpr;
33    type Encoding = IsNullExprEncoding;
34    type Metadata = EmptyMetadata;
35
36    fn id(_encoding: &Self::Encoding) -> ExprId {
37        ExprId::new_ref("is_null")
38    }
39
40    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
41        ExprEncodingRef::new_ref(IsNullExprEncoding.as_ref())
42    }
43
44    fn metadata(_expr: &Self::Expr) -> Option<Self::Metadata> {
45        Some(EmptyMetadata)
46    }
47
48    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
49        vec![&expr.child]
50    }
51
52    fn with_children(_expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
53        Ok(IsNullExpr::new(children[0].clone()))
54    }
55
56    fn build(
57        _encoding: &Self::Encoding,
58        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
59        children: Vec<ExprRef>,
60    ) -> VortexResult<Self::Expr> {
61        if children.len() != 1 {
62            vortex_bail!("IsNull expects exactly one child, got {}", children.len());
63        }
64        Ok(IsNullExpr::new(children[0].clone()))
65    }
66
67    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
68        let array = expr.child.unchecked_evaluate(scope)?;
69        match array.validity_mask() {
70            Mask::AllTrue(len) => Ok(ConstantArray::new(false, len).into_array()),
71            Mask::AllFalse(len) => Ok(ConstantArray::new(true, len).into_array()),
72            Mask::Values(mask) => Ok(BoolArray::from(mask.boolean_buffer().not()).into_array()),
73        }
74    }
75
76    fn return_dtype(_expr: &Self::Expr, _scope: &DType) -> VortexResult<DType> {
77        Ok(DType::Bool(Nullability::NonNullable))
78    }
79}
80
81impl IsNullExpr {
82    pub fn new(child: ExprRef) -> Self {
83        Self { child }
84    }
85
86    pub fn new_expr(child: ExprRef) -> ExprRef {
87        Self::new(child).into_expr()
88    }
89}
90
91impl DisplayAs for IsNullExpr {
92    fn fmt_as(&self, df: DisplayFormat, f: &mut std::fmt::Formatter) -> std::fmt::Result {
93        match df {
94            DisplayFormat::Compact => {
95                write!(f, "is_null({})", self.child)
96            }
97            DisplayFormat::Tree => {
98                write!(f, "IsNull")
99            }
100        }
101    }
102}
103
104impl AnalysisExpr for IsNullExpr {}
105
106/// Creates an expression that checks for null values.
107///
108/// Returns a boolean array indicating which positions contain null values.
109///
110/// ```rust
111/// # use vortex_expr::{is_null, root};
112/// let expr = is_null(root());
113/// ```
114pub fn is_null(child: ExprRef) -> ExprRef {
115    IsNullExpr::new(child).into_expr()
116}
117
118#[cfg(test)]
119mod tests {
120    use vortex_array::IntoArray;
121    use vortex_array::arrays::{PrimitiveArray, StructArray};
122    use vortex_dtype::{DType, Nullability};
123    use vortex_scalar::Scalar;
124
125    use crate::is_null::is_null;
126    use crate::{Scope, get_item, root, test_harness};
127
128    #[test]
129    fn dtype() {
130        let dtype = test_harness::struct_dtype();
131        assert_eq!(
132            is_null(root()).return_dtype(&dtype).unwrap(),
133            DType::Bool(Nullability::NonNullable)
134        );
135    }
136
137    #[test]
138    fn replace_children() {
139        let expr = is_null(root());
140        let _ = expr.with_children(vec![root()]);
141    }
142
143    #[test]
144    fn evaluate_mask() {
145        let test_array =
146            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
147                .into_array();
148        let expected = [false, true, false, true, false];
149
150        let result = is_null(root())
151            .evaluate(&Scope::new(test_array.clone()))
152            .unwrap();
153
154        assert_eq!(result.len(), test_array.len());
155        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
156
157        for (i, expected_value) in expected.iter().enumerate() {
158            assert_eq!(
159                result.scalar_at(i),
160                Scalar::bool(*expected_value, Nullability::NonNullable)
161            );
162        }
163    }
164
165    #[test]
166    fn evaluate_all_false() {
167        let test_array = PrimitiveArray::from_iter(vec![1, 2, 3, 4, 5]).into_array();
168
169        let result = is_null(root())
170            .evaluate(&Scope::new(test_array.clone()))
171            .unwrap();
172
173        assert_eq!(result.len(), test_array.len());
174        assert_eq!(
175            result.as_constant().unwrap(),
176            Scalar::bool(false, Nullability::NonNullable)
177        );
178    }
179
180    #[test]
181    fn evaluate_all_true() {
182        let test_array =
183            PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
184                .into_array();
185
186        let result = is_null(root())
187            .evaluate(&Scope::new(test_array.clone()))
188            .unwrap();
189
190        assert_eq!(result.len(), test_array.len());
191        assert_eq!(
192            result.as_constant().unwrap(),
193            Scalar::bool(true, Nullability::NonNullable)
194        );
195    }
196
197    #[test]
198    fn evaluate_struct() {
199        let test_array = StructArray::from_fields(&[(
200            "a",
201            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
202                .into_array(),
203        )])
204        .unwrap()
205        .into_array();
206        let expected = [false, true, false, true, false];
207
208        let result = is_null(get_item("a", root()))
209            .evaluate(&Scope::new(test_array.clone()))
210            .unwrap();
211
212        assert_eq!(result.len(), test_array.len());
213        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
214
215        for (i, expected_value) in expected.iter().enumerate() {
216            assert_eq!(
217                result.scalar_at(i),
218                Scalar::bool(*expected_value, Nullability::NonNullable)
219            );
220        }
221    }
222
223    #[test]
224    fn test_display() {
225        let expr = is_null(get_item("name", root()));
226        assert_eq!(expr.to_string(), "is_null($.name)");
227
228        let expr2 = is_null(root());
229        assert_eq!(expr2.to_string(), "is_null($)");
230    }
231}
vortex_expr/exprs/is_null.rs

vortex_expr/exprs/
is_null.rs
