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_buffer::buffer;
123    use vortex_dtype::{DType, Nullability};
124    use vortex_scalar::Scalar;
125
126    use crate::is_null::is_null;
127    use crate::{Scope, get_item, root, test_harness};
128
129    #[test]
130    fn dtype() {
131        let dtype = test_harness::struct_dtype();
132        assert_eq!(
133            is_null(root()).return_dtype(&dtype).unwrap(),
134            DType::Bool(Nullability::NonNullable)
135        );
136    }
137
138    #[test]
139    fn replace_children() {
140        let expr = is_null(root());
141        let _ = expr.with_children(vec![root()]);
142    }
143
144    #[test]
145    fn evaluate_mask() {
146        let test_array =
147            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
148                .into_array();
149        let expected = [false, true, false, true, false];
150
151        let result = is_null(root())
152            .evaluate(&Scope::new(test_array.clone()))
153            .unwrap();
154
155        assert_eq!(result.len(), test_array.len());
156        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
157
158        for (i, expected_value) in expected.iter().enumerate() {
159            assert_eq!(
160                result.scalar_at(i),
161                Scalar::bool(*expected_value, Nullability::NonNullable)
162            );
163        }
164    }
165
166    #[test]
167    fn evaluate_all_false() {
168        let test_array = buffer![1, 2, 3, 4, 5].into_array();
169
170        let result = is_null(root())
171            .evaluate(&Scope::new(test_array.clone()))
172            .unwrap();
173
174        assert_eq!(result.len(), test_array.len());
175        assert_eq!(
176            result.as_constant().unwrap(),
177            Scalar::bool(false, Nullability::NonNullable)
178        );
179    }
180
181    #[test]
182    fn evaluate_all_true() {
183        let test_array =
184            PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
185                .into_array();
186
187        let result = is_null(root())
188            .evaluate(&Scope::new(test_array.clone()))
189            .unwrap();
190
191        assert_eq!(result.len(), test_array.len());
192        assert_eq!(
193            result.as_constant().unwrap(),
194            Scalar::bool(true, Nullability::NonNullable)
195        );
196    }
197
198    #[test]
199    fn evaluate_struct() {
200        let test_array = StructArray::from_fields(&[(
201            "a",
202            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
203                .into_array(),
204        )])
205        .unwrap()
206        .into_array();
207        let expected = [false, true, false, true, false];
208
209        let result = is_null(get_item("a", root()))
210            .evaluate(&Scope::new(test_array.clone()))
211            .unwrap();
212
213        assert_eq!(result.len(), test_array.len());
214        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
215
216        for (i, expected_value) in expected.iter().enumerate() {
217            assert_eq!(
218                result.scalar_at(i),
219                Scalar::bool(*expected_value, Nullability::NonNullable)
220            );
221        }
222    }
223
224    #[test]
225    fn test_display() {
226        let expr = is_null(get_item("name", root()));
227        assert_eq!(expr.to_string(), "is_null($.name)");
228
229        let expr2 = is_null(root());
230        assert_eq!(expr2.to_string(), "is_null($)");
231    }
232}
vortex_expr/exprs/is_null.rs

vortex_expr/exprs/
is_null.rs
