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

vortex_array/compute/
invert.rs
