vortex_expr/exprs/
is_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::ops::Not;
6
7use vortex_array::arrays::{BoolArray, ConstantArray};
8use vortex_array::{Array, ArrayRef, DeserializeMetadata, EmptyMetadata, IntoArray};
9use vortex_dtype::{DType, Nullability};
10use vortex_error::{VortexResult, vortex_bail};
11use vortex_mask::Mask;
12
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 Display for IsNullExpr {
92    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
93        write!(f, "is_null({})", self.child)
94    }
95}
96
97impl AnalysisExpr for IsNullExpr {}
98
99/// Creates an expression that checks for null values.
100///
101/// Returns a boolean array indicating which positions contain null values.
102///
103/// ```rust
104/// # use vortex_expr::{is_null, root};
105/// let expr = is_null(root());
106/// ```
107pub fn is_null(child: ExprRef) -> ExprRef {
108    IsNullExpr::new(child).into_expr()
109}
110
111#[cfg(test)]
112mod tests {
113    use vortex_array::IntoArray;
114    use vortex_array::arrays::{PrimitiveArray, StructArray};
115    use vortex_dtype::{DType, Nullability};
116    use vortex_scalar::Scalar;
117
118    use crate::is_null::is_null;
119    use crate::{Scope, get_item, root, test_harness};
120
121    #[test]
122    fn dtype() {
123        let dtype = test_harness::struct_dtype();
124        assert_eq!(
125            is_null(root()).return_dtype(&dtype).unwrap(),
126            DType::Bool(Nullability::NonNullable)
127        );
128    }
129
130    #[test]
131    fn replace_children() {
132        let expr = is_null(root());
133        let _ = expr.with_children(vec![root()]);
134    }
135
136    #[test]
137    fn evaluate_mask() {
138        let test_array =
139            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
140                .into_array();
141        let expected = [false, true, false, true, false];
142
143        let result = is_null(root())
144            .evaluate(&Scope::new(test_array.clone()))
145            .unwrap();
146
147        assert_eq!(result.len(), test_array.len());
148        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
149
150        for (i, expected_value) in expected.iter().enumerate() {
151            assert_eq!(
152                result.scalar_at(i),
153                Scalar::bool(*expected_value, Nullability::NonNullable)
154            );
155        }
156    }
157
158    #[test]
159    fn evaluate_all_false() {
160        let test_array = PrimitiveArray::from_iter(vec![1, 2, 3, 4, 5]).into_array();
161
162        let result = is_null(root())
163            .evaluate(&Scope::new(test_array.clone()))
164            .unwrap();
165
166        assert_eq!(result.len(), test_array.len());
167        assert_eq!(
168            result.as_constant().unwrap(),
169            Scalar::bool(false, Nullability::NonNullable)
170        );
171    }
172
173    #[test]
174    fn evaluate_all_true() {
175        let test_array =
176            PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
177                .into_array();
178
179        let result = is_null(root())
180            .evaluate(&Scope::new(test_array.clone()))
181            .unwrap();
182
183        assert_eq!(result.len(), test_array.len());
184        assert_eq!(
185            result.as_constant().unwrap(),
186            Scalar::bool(true, Nullability::NonNullable)
187        );
188    }
189
190    #[test]
191    fn evaluate_struct() {
192        let test_array = StructArray::from_fields(&[(
193            "a",
194            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
195                .into_array(),
196        )])
197        .unwrap()
198        .into_array();
199        let expected = [false, true, false, true, false];
200
201        let result = is_null(get_item("a", root()))
202            .evaluate(&Scope::new(test_array.clone()))
203            .unwrap();
204
205        assert_eq!(result.len(), test_array.len());
206        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
207
208        for (i, expected_value) in expected.iter().enumerate() {
209            assert_eq!(
210                result.scalar_at(i),
211                Scalar::bool(*expected_value, Nullability::NonNullable)
212            );
213        }
214    }
215}
vortex_expr/exprs/is_null.rs

vortex_expr/exprs/
is_null.rs
