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::ArrayRef;
14use crate::DynArray;
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::IntoArray;
125    use crate::ToCanonical;
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    use crate::scalar_fn::fns::not::BoolArray;
134
135    #[test]
136    fn invert_booleans() {
137        let not_expr = not(root());
138        let bools = BoolArray::from_iter([false, true, false, false, true, true]);
139        assert_eq!(
140            bools
141                .into_array()
142                .apply(&not_expr)
143                .unwrap()
144                .to_bool()
145                .to_bit_buffer()
146                .iter()
147                .collect::<Vec<_>>(),
148            vec![true, false, true, true, false, false]
149        );
150    }
151
152    #[test]
153    fn test_display_order_of_operations() {
154        let a = not(get_item("a", root()));
155        let b = get_item("a", not(root()));
156        assert_ne!(a.to_string(), b.to_string());
157        assert_eq!(a.to_string(), "not($.a)");
158        assert_eq!(b.to_string(), "not($).a");
159    }
160
161    #[test]
162    fn dtype() {
163        let not_expr = not(root());
164        let dtype = DType::Bool(Nullability::NonNullable);
165        assert_eq!(
166            not_expr.return_dtype(&dtype).unwrap(),
167            DType::Bool(Nullability::NonNullable)
168        );
169
170        let dtype = test_harness::struct_dtype();
171        assert_eq!(
172            not(col("bool1")).return_dtype(&dtype).unwrap(),
173            DType::Bool(Nullability::NonNullable)
174        );
175    }
176}