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

vortex_array/compute/
invert.rs
