vortex_array/compute/
invert.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::LazyLock;
5
6use arcref::ArcRef;
7use vortex_dtype::DType;
8use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err, vortex_panic};
9
10use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, UnaryArgs};
11use crate::vtable::VTable;
12use crate::{Array, ArrayRef, IntoArray, ToCanonical};
13
14static INVERT_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
15    let compute = ComputeFn::new("invert".into(), ArcRef::new_ref(&Invert));
16    for kernel in inventory::iter::<InvertKernelRef> {
17        compute.register_kernel(kernel.0.clone());
18    }
19    compute
20});
21
22pub(crate) fn warm_up_vtable() -> usize {
23    INVERT_FN.kernels().len()
24}
25
26/// Logically invert a boolean array, preserving its validity.
27pub fn invert(array: &dyn Array) -> VortexResult<ArrayRef> {
28    INVERT_FN
29        .invoke(&InvocationArgs {
30            inputs: &[array.into()],
31            options: &(),
32        })?
33        .unwrap_array()
34}
35
36struct Invert;
37
38impl ComputeFnVTable for Invert {
39    fn invoke(
40        &self,
41        args: &InvocationArgs,
42        kernels: &[ArcRef<dyn Kernel>],
43    ) -> VortexResult<Output> {
44        let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;
45
46        for kernel in kernels {
47            if let Some(output) = kernel.invoke(args)? {
48                return Ok(output);
49            }
50        }
51        if let Some(output) = array.invoke(&INVERT_FN, args)? {
52            return Ok(output);
53        }
54
55        // Otherwise, we canonicalize into a boolean array and invert.
56        log::debug!(
57            "No invert implementation found for encoding {}",
58            array.encoding_id(),
59        );
60        if array.is_canonical() {
61            vortex_panic!("Canonical bool array does not implement invert");
62        }
63        Ok(invert(&array.to_bool().into_array())?.into())
64    }
65
66    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
67        let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;
68
69        if !matches!(array.dtype(), DType::Bool(..)) {
70            vortex_bail!("Expected boolean array, got {}", array.dtype());
71        }
72        Ok(array.dtype().clone())
73    }
74
75    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
76        let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;
77        Ok(array.len())
78    }
79
80    fn is_elementwise(&self) -> bool {
81        true
82    }
83}
84
85struct InvertArgs<'a> {
86    array: &'a dyn Array,
87}
88
89impl<'a> TryFrom<&InvocationArgs<'a>> for InvertArgs<'a> {
90    type Error = VortexError;
91
92    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
93        if value.inputs.len() != 1 {
94            vortex_bail!("Invert expects exactly one argument",);
95        }
96        let array = value.inputs[0]
97            .array()
98            .ok_or_else(|| vortex_err!("Invert expects an array argument"))?;
99        Ok(InvertArgs { array })
100    }
101}
102
103pub struct InvertKernelRef(ArcRef<dyn Kernel>);
104inventory::collect!(InvertKernelRef);
105
106pub trait InvertKernel: VTable {
107    /// Logically invert a boolean array. Converts true -> false, false -> true, null -> null.
108    fn invert(&self, array: &Self::Array) -> VortexResult<ArrayRef>;
109}
110
111#[derive(Debug)]
112pub struct InvertKernelAdapter<V: VTable>(pub V);
113
114impl<V: VTable + InvertKernel> InvertKernelAdapter<V> {
115    pub const fn lift(&'static self) -> InvertKernelRef {
116        InvertKernelRef(ArcRef::new_ref(self))
117    }
118}
119
120impl<V: VTable + InvertKernel> Kernel for InvertKernelAdapter<V> {
121    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
122        let args = InvertArgs::try_from(args)?;
123        let Some(array) = args.array.as_opt::<V>() else {
124            return Ok(None);
125        };
126        Ok(Some(V::invert(&self.0, array)?.into()))
127    }
128}
vortex_array/compute/invert.rs

vortex_array/compute/
invert.rs
