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
6pub use kernel::*;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_session::VortexSession;
10use vortex_session::registry::CachedId;
11
12use crate::ArrayRef;
13use crate::ExecutionCtx;
14use crate::IntoArray;
15use crate::arrays::Bool;
16use crate::arrays::BoolArray;
17use crate::arrays::ConstantArray;
18use crate::arrays::ScalarFnArray;
19use crate::arrays::bool::BoolArrayExt;
20use crate::builtins::ArrayBuiltins;
21use crate::dtype::DType;
22use crate::scalar::Scalar;
23use crate::scalar_fn::Arity;
24use crate::scalar_fn::ChildName;
25use crate::scalar_fn::EmptyOptions;
26use crate::scalar_fn::ExecutionArgs;
27use crate::scalar_fn::ScalarFnId;
28use crate::scalar_fn::ScalarFnVTable;
29use crate::scalar_fn::ScalarFnVTableExt;
30
31/// Expression that logically inverts boolean values.
32#[derive(Clone)]
33pub struct Not;
34
35impl Not {
36    /// Creates a lazy boolean inversion of `input`.
37    ///
38    /// # Errors
39    ///
40    /// Returns an error if `input` is not boolean data.
41    pub fn try_new(input: ArrayRef) -> VortexResult<ScalarFnArray> {
42        ScalarFnArray::try_new(Not.bind(EmptyOptions), [input])
43    }
44}
45
46impl ScalarFnVTable for Not {
47    type Options = EmptyOptions;
48
49    fn id(&self) -> ScalarFnId {
50        static ID: CachedId = CachedId::new("vortex.not");
51        *ID
52    }
53
54    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
55        Ok(Some(vec![]))
56    }
57
58    fn deserialize(
59        &self,
60        _metadata: &[u8],
61        _session: &VortexSession,
62    ) -> VortexResult<Self::Options> {
63        Ok(EmptyOptions)
64    }
65
66    fn arity(&self, _options: &Self::Options) -> Arity {
67        Arity::Exact(1)
68    }
69
70    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
71        match child_idx {
72            0 => ChildName::from("input"),
73            _ => unreachable!("Invalid child index {} for Not expression", child_idx),
74        }
75    }
76
77    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
78        let child_dtype = &arg_dtypes[0];
79        if !matches!(child_dtype, DType::Bool(_)) {
80            vortex_bail!(
81                "Not expression expects a boolean child, got: {}",
82                child_dtype
83            );
84        }
85        Ok(child_dtype.clone())
86    }
87
88    fn execute(
89        &self,
90        _data: &Self::Options,
91        args: &dyn ExecutionArgs,
92        ctx: &mut ExecutionCtx,
93    ) -> VortexResult<ArrayRef> {
94        let child = args.get(0)?;
95
96        // For constant boolean
97        if let Some(scalar) = child.as_constant() {
98            let value = match scalar.as_bool().value() {
99                Some(b) => Scalar::bool(!b, child.dtype().nullability()),
100                None => Scalar::null(child.dtype().clone()),
101            };
102            return Ok(ConstantArray::new(value, args.row_count()).into_array());
103        }
104
105        // For boolean array
106        if let Some(bool) = child.as_opt::<Bool>() {
107            return Ok(BoolArray::new(!bool.to_bit_buffer(), bool.validity()?).into_array());
108        }
109
110        // Otherwise, execute and try again
111        child.execute::<ArrayRef>(ctx)?.not()
112    }
113
114    fn is_strict(&self, _options: &Self::Options) -> bool {
115        true
116    }
117
118    fn is_infallible(&self, _options: &Self::Options) -> bool {
119        true
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use vortex_error::VortexResult;
126
127    use crate::IntoArray;
128    use crate::VortexSessionExecute;
129    use crate::array_session;
130    use crate::assert_arrays_eq;
131    use crate::dtype::DType;
132    use crate::dtype::Nullability;
133    use crate::expr::col;
134    use crate::expr::get_item;
135    use crate::expr::not;
136    use crate::expr::root;
137    use crate::expr::test_harness;
138    use crate::scalar_fn::fns::not::BoolArray;
139
140    #[test]
141    fn is_strict() {
142        assert!(
143            not(root())
144                .as_scalar()
145                .is_some_and(|f| f.signature().is_strict())
146        );
147    }
148
149    #[test]
150    fn preserves_nulls() -> VortexResult<()> {
151        let mut ctx = array_session().create_execution_ctx();
152        let input = BoolArray::from_iter([Some(false), None, Some(true)]).into_array();
153
154        let result = input.apply(&not(root()))?;
155
156        assert_arrays_eq!(
157            result,
158            BoolArray::from_iter([Some(true), None, Some(false)]),
159            &mut ctx
160        );
161        Ok(())
162    }
163
164    #[test]
165    fn invert_booleans() {
166        let mut ctx = array_session().create_execution_ctx();
167        let not_expr = not(root());
168        let bools = BoolArray::from_iter([false, true, false, false, true, true]);
169        let result = bools.into_array().apply(&not_expr).unwrap();
170        assert_arrays_eq!(
171            result,
172            BoolArray::from_iter([true, false, true, true, false, false]),
173            &mut ctx
174        );
175    }
176
177    #[test]
178    fn test_display_order_of_operations() {
179        let a = not(get_item("a", root()));
180        let b = get_item("a", not(root()));
181        assert_ne!(a.to_string(), b.to_string());
182        assert_eq!(a.to_string(), "vortex.not($.a)");
183        assert_eq!(b.to_string(), "vortex.not($).a");
184    }
185
186    #[test]
187    fn dtype() {
188        let not_expr = not(root());
189        let dtype = DType::Bool(Nullability::NonNullable);
190        assert_eq!(
191            not_expr.return_dtype(&dtype).unwrap(),
192            DType::Bool(Nullability::NonNullable)
193        );
194
195        let dtype = test_harness::struct_dtype();
196        assert_eq!(
197            not(col("bool1")).return_dtype(&dtype).unwrap(),
198            DType::Bool(Nullability::NonNullable)
199        );
200    }
201}
vortex_array/scalar_fn/fns/not/mod.rs

vortex_array/scalar_fn/fns/not/
mod.rs
