vortex_array/compute/
invert.rs

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

vortex_array/compute/
invert.rs
