Skip to main content

vortex_array/scalar_fn/fns/not/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod kernel;
5
6use std::fmt::Formatter;
7
8pub use kernel::*;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_session::VortexSession;
13
14use crate::Array;
15use crate::ArrayRef;
16use crate::IntoArray;
17use crate::arrays::BoolArray;
18use crate::arrays::BoolVTable;
19use crate::arrays::ConstantArray;
20use crate::builtins::ArrayBuiltins;
21use crate::dtype::DType;
22use crate::expr::Expression;
23use crate::scalar::Scalar;
24use crate::scalar_fn::Arity;
25use crate::scalar_fn::ChildName;
26use crate::scalar_fn::EmptyOptions;
27use crate::scalar_fn::ExecutionArgs;
28use crate::scalar_fn::ScalarFnId;
29use crate::scalar_fn::ScalarFnVTable;
30
31/// Expression that logically inverts boolean values.
32#[derive(Clone)]
33pub struct Not;
34
35impl ScalarFnVTable for Not {
36    type Options = EmptyOptions;
37
38    fn id(&self) -> ScalarFnId {
39        ScalarFnId::from("vortex.not")
40    }
41
42    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
43        Ok(Some(vec![]))
44    }
45
46    fn deserialize(
47        &self,
48        _metadata: &[u8],
49        _session: &VortexSession,
50    ) -> VortexResult<Self::Options> {
51        Ok(EmptyOptions)
52    }
53
54    fn arity(&self, _options: &Self::Options) -> Arity {
55        Arity::Exact(1)
56    }
57
58    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
59        match child_idx {
60            0 => ChildName::from("input"),
61            _ => unreachable!("Invalid child index {} for Not expression", child_idx),
62        }
63    }
64
65    fn fmt_sql(
66        &self,
67        _options: &Self::Options,
68        expr: &Expression,
69        f: &mut Formatter<'_>,
70    ) -> std::fmt::Result {
71        write!(f, "not(")?;
72        expr.child(0).fmt_sql(f)?;
73        write!(f, ")")
74    }
75
76    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
77        let child_dtype = &arg_dtypes[0];
78        if !matches!(child_dtype, DType::Bool(_)) {
79            vortex_bail!(
80                "Not expression expects a boolean child, got: {}",
81                child_dtype
82            );
83        }
84        Ok(child_dtype.clone())
85    }
86
87    fn execute(&self, _data: &Self::Options, mut args: ExecutionArgs) -> VortexResult<ArrayRef> {
88        let child = args.inputs.pop().vortex_expect("Missing input child");
89
90        // For constant boolean
91        if let Some(scalar) = child.as_constant() {
92            let value = match scalar.as_bool().value() {
93                Some(b) => Scalar::bool(!b, child.dtype().nullability()),
94                None => Scalar::null(child.dtype().clone()),
95            };
96            return Ok(ConstantArray::new(value, args.row_count).into_array());
97        }
98
99        // For boolean array
100        if let Some(bool) = child.as_opt::<BoolVTable>() {
101            return Ok(BoolArray::new(!bool.to_bit_buffer(), bool.validity()?).into_array());
102        }
103
104        // Otherwise, execute and try again
105        child.execute::<ArrayRef>(args.ctx)?.not()
106    }
107
108    fn is_null_sensitive(&self, _options: &Self::Options) -> bool {
109        false
110    }
111
112    fn is_fallible(&self, _options: &Self::Options) -> bool {
113        false
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use crate::ToCanonical;
120    use crate::arrays::BoolArray;
121    use crate::dtype::DType;
122    use crate::dtype::Nullability;
123    use crate::expr::col;
124    use crate::expr::get_item;
125    use crate::expr::not;
126    use crate::expr::root;
127    use crate::expr::test_harness;
128
129    #[test]
130    fn invert_booleans() {
131        let not_expr = not(root());
132        let bools = BoolArray::from_iter([false, true, false, false, true, true]);
133        assert_eq!(
134            bools
135                .to_array()
136                .apply(&not_expr)
137                .unwrap()
138                .to_bool()
139                .to_bit_buffer()
140                .iter()
141                .collect::<Vec<_>>(),
142            vec![true, false, true, true, false, false]
143        );
144    }
145
146    #[test]
147    fn test_display_order_of_operations() {
148        let a = not(get_item("a", root()));
149        let b = get_item("a", not(root()));
150        assert_ne!(a.to_string(), b.to_string());
151        assert_eq!(a.to_string(), "not($.a)");
152        assert_eq!(b.to_string(), "not($).a");
153    }
154
155    #[test]
156    fn dtype() {
157        let not_expr = not(root());
158        let dtype = DType::Bool(Nullability::NonNullable);
159        assert_eq!(
160            not_expr.return_dtype(&dtype).unwrap(),
161            DType::Bool(Nullability::NonNullable)
162        );
163
164        let dtype = test_harness::struct_dtype();
165        assert_eq!(
166            not(col("bool1")).return_dtype(&dtype).unwrap(),
167            DType::Bool(Nullability::NonNullable)
168        );
169    }
170}