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::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_session::VortexSession;
12
13use crate::Array;
14use crate::ArrayRef;
15use crate::ExecutionCtx;
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(
88        &self,
89        _data: &Self::Options,
90        args: &dyn ExecutionArgs,
91        ctx: &mut ExecutionCtx,
92    ) -> VortexResult<ArrayRef> {
93        let child = args.get(0)?;
94
95        // For constant boolean
96        if let Some(scalar) = child.as_constant() {
97            let value = match scalar.as_bool().value() {
98                Some(b) => Scalar::bool(!b, child.dtype().nullability()),
99                None => Scalar::null(child.dtype().clone()),
100            };
101            return Ok(ConstantArray::new(value, args.row_count()).into_array());
102        }
103
104        // For boolean array
105        if let Some(bool) = child.as_opt::<BoolVTable>() {
106            return Ok(BoolArray::new(!bool.to_bit_buffer(), bool.validity()?).into_array());
107        }
108
109        // Otherwise, execute and try again
110        child.execute::<ArrayRef>(ctx)?.not()
111    }
112
113    fn is_null_sensitive(&self, _options: &Self::Options) -> bool {
114        false
115    }
116
117    fn is_fallible(&self, _options: &Self::Options) -> bool {
118        false
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use crate::ToCanonical;
125    use crate::arrays::BoolArray;
126    use crate::dtype::DType;
127    use crate::dtype::Nullability;
128    use crate::expr::col;
129    use crate::expr::get_item;
130    use crate::expr::not;
131    use crate::expr::root;
132    use crate::expr::test_harness;
133
134    #[test]
135    fn invert_booleans() {
136        let not_expr = not(root());
137        let bools = BoolArray::from_iter([false, true, false, false, true, true]);
138        assert_eq!(
139            bools
140                .to_array()
141                .apply(&not_expr)
142                .unwrap()
143                .to_bool()
144                .to_bit_buffer()
145                .iter()
146                .collect::<Vec<_>>(),
147            vec![true, false, true, true, false, false]
148        );
149    }
150
151    #[test]
152    fn test_display_order_of_operations() {
153        let a = not(get_item("a", root()));
154        let b = get_item("a", not(root()));
155        assert_ne!(a.to_string(), b.to_string());
156        assert_eq!(a.to_string(), "not($.a)");
157        assert_eq!(b.to_string(), "not($).a");
158    }
159
160    #[test]
161    fn dtype() {
162        let not_expr = not(root());
163        let dtype = DType::Bool(Nullability::NonNullable);
164        assert_eq!(
165            not_expr.return_dtype(&dtype).unwrap(),
166            DType::Bool(Nullability::NonNullable)
167        );
168
169        let dtype = test_harness::struct_dtype();
170        assert_eq!(
171            not(col("bool1")).return_dtype(&dtype).unwrap(),
172            DType::Bool(Nullability::NonNullable)
173        );
174    }
175}