vortex_array/compute/
invert.rs

1use std::sync::LazyLock;
2
3use vortex_dtype::DType;
4use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err, vortex_panic};
5
6use crate::arcref::ArcRef;
7use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output};
8use crate::encoding::Encoding;
9use crate::{Array, ArrayRef, ToCanonical};
10
11/// Logically invert a boolean array.
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 InvertArgs { array } = InvertArgs::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(),
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 InvertArgs { array } = InvertArgs::try_from(args)?;
53        if !matches!(array.dtype(), DType::Bool(..)) {
54            vortex_bail!("Expected boolean array, got {}", array.dtype());
55        }
56        Ok(array.dtype().clone())
57    }
58
59    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
60        let InvertArgs { array } = InvertArgs::try_from(args)?;
61        Ok(array.len())
62    }
63
64    fn is_elementwise(&self) -> bool {
65        true
66    }
67}
68
69struct InvertArgs<'a> {
70    array: &'a dyn Array,
71}
72
73impl<'a> TryFrom<&InvocationArgs<'a>> for InvertArgs<'a> {
74    type Error = VortexError;
75
76    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
77        if value.inputs.len() != 1 {
78            vortex_bail!("Invert expects exactly one argument",);
79        }
80        let array = value.inputs[0]
81            .array()
82            .ok_or_else(|| vortex_err!("Invert expects an array argument"))?;
83        Ok(InvertArgs { array })
84    }
85}
86
87pub struct InvertKernelRef(ArcRef<dyn Kernel>);
88inventory::collect!(InvertKernelRef);
89
90pub trait InvertKernel: Encoding {
91    /// Logically invert a boolean array. Converts true -> false, false -> true, null -> null.
92    fn invert(&self, array: &Self::Array) -> VortexResult<ArrayRef>;
93}
94
95#[derive(Debug)]
96pub struct InvertKernelAdapter<E: Encoding>(pub E);
97
98impl<E: Encoding + InvertKernel> InvertKernelAdapter<E> {
99    pub const fn lift(&'static self) -> InvertKernelRef {
100        InvertKernelRef(ArcRef::new_ref(self))
101    }
102}
103
104impl<E: Encoding + InvertKernel> Kernel for InvertKernelAdapter<E> {
105    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
106        let args = InvertArgs::try_from(args)?;
107        let Some(array) = args.array.as_any().downcast_ref::<E::Array>() else {
108            return Ok(None);
109        };
110        Ok(Some(E::invert(&self.0, array)?.into()))
111    }
112}
113
114pub static INVERT_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
115    let compute = ComputeFn::new("invert".into(), ArcRef::new_ref(&Invert));
116    for kernel in inventory::iter::<InvertKernelRef> {
117        compute.register_kernel(kernel.0.clone());
118    }
119    compute
120});
vortex_array/compute/invert.rs

vortex_array/compute/
invert.rs
