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

vortex_array/scalar_fn/fns/not/
mod.rs
