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), vec![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::arrays::bool::BoolArrayExt;
131    use crate::assert_arrays_eq;
132    use crate::dtype::DType;
133    use crate::dtype::Nullability;
134    use crate::expr::col;
135    use crate::expr::get_item;
136    use crate::expr::not;
137    use crate::expr::root;
138    use crate::expr::test_harness;
139    use crate::scalar_fn::fns::not::BoolArray;
140
141    #[test]
142    fn is_strict() {
143        assert!(
144            not(root())
145                .as_scalar()
146                .is_some_and(|f| f.signature().is_strict())
147        );
148    }
149
150    #[test]
151    fn preserves_nulls() -> VortexResult<()> {
152        let mut ctx = array_session().create_execution_ctx();
153        let input = BoolArray::from_iter([Some(false), None, Some(true)]).into_array();
154
155        let result = input.apply(&not(root()))?;
156
157        assert_arrays_eq!(
158            result,
159            BoolArray::from_iter([Some(true), None, Some(false)]),
160            &mut ctx
161        );
162        Ok(())
163    }
164
165    #[test]
166    fn invert_booleans() {
167        let mut ctx = array_session().create_execution_ctx();
168        let not_expr = not(root());
169        let bools = BoolArray::from_iter([false, true, false, false, true, true]);
170        let result = bools
171            .into_array()
172            .apply(&not_expr)
173            .unwrap()
174            .execute::<BoolArray>(&mut ctx)
175            .unwrap();
176        assert_eq!(
177            result.to_bit_buffer().iter().collect::<Vec<_>>(),
178            vec![true, false, true, true, false, false]
179        );
180    }
181
182    #[test]
183    fn test_display_order_of_operations() {
184        let a = not(get_item("a", root()));
185        let b = get_item("a", not(root()));
186        assert_ne!(a.to_string(), b.to_string());
187        assert_eq!(a.to_string(), "vortex.not($.a)");
188        assert_eq!(b.to_string(), "vortex.not($).a");
189    }
190
191    #[test]
192    fn dtype() {
193        let not_expr = not(root());
194        let dtype = DType::Bool(Nullability::NonNullable);
195        assert_eq!(
196            not_expr.return_dtype(&dtype).unwrap(),
197            DType::Bool(Nullability::NonNullable)
198        );
199
200        let dtype = test_harness::struct_dtype();
201        assert_eq!(
202            not(col("bool1")).return_dtype(&dtype).unwrap(),
203            DType::Bool(Nullability::NonNullable)
204        );
205    }
206}
vortex_array/scalar_fn/fns/not/mod.rs

vortex_array/scalar_fn/fns/not/
mod.rs
