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)]
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
99pub fn is_null(child: ExprRef) -> ExprRef {
100    IsNullExpr::new(child).into_expr()
101}
102
103#[cfg(test)]
104mod tests {
105    use vortex_array::IntoArray;
106    use vortex_array::arrays::{PrimitiveArray, StructArray};
107    use vortex_dtype::{DType, Nullability};
108    use vortex_scalar::Scalar;
109
110    use crate::is_null::is_null;
111    use crate::{Scope, get_item, root, test_harness};
112
113    #[test]
114    fn dtype() {
115        let dtype = test_harness::struct_dtype();
116        assert_eq!(
117            is_null(root()).return_dtype(&dtype).unwrap(),
118            DType::Bool(Nullability::NonNullable)
119        );
120    }
121
122    #[test]
123    fn replace_children() {
124        let expr = is_null(root());
125        let _ = expr.with_children(vec![root()]);
126    }
127
128    #[test]
129    fn evaluate_mask() {
130        let test_array =
131            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
132                .into_array();
133        let expected = [false, true, false, true, false];
134
135        let result = is_null(root())
136            .evaluate(&Scope::new(test_array.clone()))
137            .unwrap();
138
139        assert_eq!(result.len(), test_array.len());
140        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
141
142        for (i, expected_value) in expected.iter().enumerate() {
143            assert_eq!(
144                result.scalar_at(i).unwrap(),
145                Scalar::bool(*expected_value, Nullability::NonNullable)
146            );
147        }
148    }
149
150    #[test]
151    fn evaluate_all_false() {
152        let test_array = PrimitiveArray::from_iter(vec![1, 2, 3, 4, 5]).into_array();
153
154        let result = is_null(root())
155            .evaluate(&Scope::new(test_array.clone()))
156            .unwrap();
157
158        assert_eq!(result.len(), test_array.len());
159        assert_eq!(
160            result.as_constant().unwrap(),
161            Scalar::bool(false, Nullability::NonNullable)
162        );
163    }
164
165    #[test]
166    fn evaluate_all_true() {
167        let test_array =
168            PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
169                .into_array();
170
171        let result = is_null(root())
172            .evaluate(&Scope::new(test_array.clone()))
173            .unwrap();
174
175        assert_eq!(result.len(), test_array.len());
176        assert_eq!(
177            result.as_constant().unwrap(),
178            Scalar::bool(true, Nullability::NonNullable)
179        );
180    }
181
182    #[test]
183    fn evaluate_struct() {
184        let test_array = StructArray::from_fields(&[(
185            "a",
186            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
187                .into_array(),
188        )])
189        .unwrap()
190        .into_array();
191        let expected = [false, true, false, true, false];
192
193        let result = is_null(get_item("a", root()))
194            .evaluate(&Scope::new(test_array.clone()))
195            .unwrap();
196
197        assert_eq!(result.len(), test_array.len());
198        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
199
200        for (i, expected_value) in expected.iter().enumerate() {
201            assert_eq!(
202                result.scalar_at(i).unwrap(),
203                Scalar::bool(*expected_value, Nullability::NonNullable)
204            );
205        }
206    }
207}
vortex_expr/exprs/is_null.rs

vortex_expr/exprs/
is_null.rs
