vortex_expr/exprs/
not.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::hash::Hash;
6
7use vortex_array::compute::invert;
8use vortex_array::{ArrayRef, DeserializeMetadata, EmptyMetadata};
9use vortex_dtype::DType;
10use vortex_error::{VortexResult, vortex_bail};
11
12use crate::{AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable};
13
14vtable!(Not);
15
16#[allow(clippy::derived_hash_with_manual_eq)]
17#[derive(Clone, Debug, Hash, Eq)]
18pub struct NotExpr {
19    child: ExprRef,
20}
21
22impl PartialEq for NotExpr {
23    fn eq(&self, other: &Self) -> bool {
24        self.child.eq(&other.child)
25    }
26}
27
28pub struct NotExprEncoding;
29
30impl VTable for NotVTable {
31    type Expr = NotExpr;
32    type Encoding = NotExprEncoding;
33    type Metadata = EmptyMetadata;
34
35    fn id(_encoding: &Self::Encoding) -> ExprId {
36        ExprId::new_ref("not")
37    }
38
39    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
40        ExprEncodingRef::new_ref(NotExprEncoding.as_ref())
41    }
42
43    fn metadata(_expr: &Self::Expr) -> Option<Self::Metadata> {
44        Some(EmptyMetadata)
45    }
46
47    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
48        vec![&expr.child]
49    }
50
51    fn with_children(_expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
52        Ok(NotExpr::new(children[0].clone()))
53    }
54
55    fn build(
56        _encoding: &Self::Encoding,
57        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
58        children: Vec<ExprRef>,
59    ) -> VortexResult<Self::Expr> {
60        if children.len() != 1 {
61            vortex_bail!(
62                "Not expression expects exactly one child, got {}",
63                children.len()
64            );
65        }
66        Ok(NotExpr::new(children[0].clone()))
67    }
68
69    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
70        let child_result = expr.child.unchecked_evaluate(scope)?;
71        invert(&child_result)
72    }
73
74    fn return_dtype(expr: &Self::Expr, scope: &DType) -> VortexResult<DType> {
75        let child = expr.child.return_dtype(scope)?;
76        if !matches!(child, DType::Bool(_)) {
77            vortex_bail!("Not expression expects a boolean child, got: {}", child);
78        }
79        Ok(child)
80    }
81}
82
83impl NotExpr {
84    pub fn new(child: ExprRef) -> Self {
85        Self { child }
86    }
87
88    pub fn new_expr(child: ExprRef) -> ExprRef {
89        Self::new(child).into_expr()
90    }
91
92    pub fn child(&self) -> &ExprRef {
93        &self.child
94    }
95}
96
97impl Display for NotExpr {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        write!(f, "(!{})", self.child)
100    }
101}
102
103impl AnalysisExpr for NotExpr {}
104
105pub fn not(operand: ExprRef) -> ExprRef {
106    NotExpr::new(operand).into_expr()
107}
108
109#[cfg(test)]
110mod tests {
111    use vortex_array::ToCanonical;
112    use vortex_array::arrays::BoolArray;
113    use vortex_dtype::{DType, Nullability};
114
115    use crate::{Scope, col, get_item, not, root, test_harness};
116
117    #[test]
118    fn invert_booleans() {
119        let not_expr = not(root());
120        let bools = BoolArray::from_iter([false, true, false, false, true, true]);
121        assert_eq!(
122            not_expr
123                .evaluate(&Scope::new(bools.to_array()))
124                .unwrap()
125                .to_bool()
126                .unwrap()
127                .boolean_buffer()
128                .iter()
129                .collect::<Vec<_>>(),
130            vec![true, false, true, true, false, false]
131        );
132    }
133
134    #[test]
135    fn test_display_order_of_operations() {
136        let a = not(get_item("a", root()));
137        let b = get_item("a", not(root()));
138        assert_ne!(a.to_string(), b.to_string());
139        assert_eq!(a.to_string(), "(!$.a)");
140        assert_eq!(b.to_string(), "(!$).a");
141    }
142
143    #[test]
144    fn dtype() {
145        let not_expr = not(root());
146        let dtype = DType::Bool(Nullability::NonNullable);
147        assert_eq!(
148            not_expr.return_dtype(&dtype).unwrap(),
149            DType::Bool(Nullability::NonNullable)
150        );
151
152        let dtype = test_harness::struct_dtype();
153        assert_eq!(
154            not(col("bool1")).return_dtype(&dtype).unwrap(),
155            DType::Bool(Nullability::NonNullable)
156        );
157    }
158}