diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 3829b248df..e4518121c8 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -10,9 +10,7 @@ use vortex_array::stats::StatsSet; use vortex_array::validity::{LogicalValidity, Validity, ValidityMetadata, ValidityVTable}; use vortex_array::variants::{ArrayVariants, BoolArrayTrait}; use vortex_array::visitor::{ArrayVisitor, VisitorVTable}; -use vortex_array::{ - impl_encoding, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData, IntoCanonical, -}; +use vortex_array::{impl_encoding, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoCanonical}; use vortex_buffer::Buffer; use vortex_dtype::DType; use vortex_error::{VortexExpect as _, VortexResult}; @@ -92,15 +90,7 @@ impl ArrayVariants for ByteBoolArray { } } -impl BoolArrayTrait for ByteBoolArray { - fn invert(&self) -> VortexResult { - ByteBoolArray::try_from_vec( - self.maybe_null_slice().iter().map(|v| !v).collect(), - self.validity(), - ) - .map(|a| a.into_array()) - } -} +impl BoolArrayTrait for ByteBoolArray {} impl From> for ByteBoolArray { fn from(value: Vec) -> Self { diff --git a/encodings/dict/src/variants.rs b/encodings/dict/src/variants.rs index 05d1a1dbb3..5725a07370 100644 --- a/encodings/dict/src/variants.rs +++ b/encodings/dict/src/variants.rs @@ -1,9 +1,8 @@ use vortex_array::variants::{ ArrayVariants, BinaryArrayTrait, BoolArrayTrait, PrimitiveArrayTrait, Utf8ArrayTrait, }; -use vortex_array::{ArrayDType, ArrayData}; +use vortex_array::ArrayDType; use vortex_dtype::DType; -use vortex_error::VortexResult; use crate::DictArray; @@ -25,11 +24,7 @@ impl ArrayVariants for DictArray { } } -impl BoolArrayTrait for DictArray { - fn invert(&self) -> VortexResult { - todo!() - } -} +impl BoolArrayTrait for DictArray {} impl PrimitiveArrayTrait for DictArray {} diff --git a/encodings/roaring/src/boolean/compute.rs b/encodings/roaring/src/boolean/compute.rs index dc57986c89..c79a55551f 100644 --- a/encodings/roaring/src/boolean/compute.rs +++ b/encodings/roaring/src/boolean/compute.rs @@ -1,12 +1,16 @@ use croaring::Bitmap; -use vortex_array::compute::{ComputeVTable, ScalarAtFn, SliceFn}; -use vortex_array::{ArrayData, IntoArrayData}; +use vortex_array::compute::{ComputeVTable, InvertFn, ScalarAtFn, SliceFn}; +use vortex_array::{ArrayData, ArrayLen, IntoArrayData}; use vortex_error::VortexResult; use vortex_scalar::Scalar; use crate::{RoaringBoolArray, RoaringBoolEncoding}; impl ComputeVTable for RoaringBoolEncoding { + fn invert_fn(&self) -> Option<&dyn InvertFn> { + Some(self) + } + fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn> { Some(self) } @@ -16,6 +20,13 @@ impl ComputeVTable for RoaringBoolEncoding { } } +impl InvertFn for RoaringBoolEncoding { + fn invert(&self, array: &RoaringBoolArray) -> VortexResult { + RoaringBoolArray::try_new(array.bitmap().flip(0..(array.len() as u32)), array.len()) + .map(|a| a.into_array()) + } +} + impl ScalarAtFn for RoaringBoolEncoding { fn scalar_at(&self, array: &RoaringBoolArray, index: usize) -> VortexResult { Ok(array.bitmap().contains(index as u32).into()) diff --git a/encodings/roaring/src/boolean/mod.rs b/encodings/roaring/src/boolean/mod.rs index 234efc92eb..3b82db257b 100644 --- a/encodings/roaring/src/boolean/mod.rs +++ b/encodings/roaring/src/boolean/mod.rs @@ -90,12 +90,7 @@ impl ArrayVariants for RoaringBoolArray { } } -impl BoolArrayTrait for RoaringBoolArray { - fn invert(&self) -> VortexResult { - RoaringBoolArray::try_new(self.bitmap().flip(0..(self.len() as u32)), self.len()) - .map(|a| a.into_array()) - } -} +impl BoolArrayTrait for RoaringBoolArray {} impl VisitorVTable for RoaringBoolEncoding { fn accept(&self, array: &RoaringBoolArray, visitor: &mut dyn ArrayVisitor) -> VortexResult<()> { diff --git a/encodings/runend-bool/src/array.rs b/encodings/runend-bool/src/array.rs index 96f752d4b6..552821721d 100644 --- a/encodings/runend-bool/src/array.rs +++ b/encodings/runend-bool/src/array.rs @@ -176,18 +176,7 @@ pub(crate) fn value_at_index(idx: usize, start: bool) -> bool { } } -impl BoolArrayTrait for RunEndBoolArray { - fn invert(&self) -> VortexResult { - RunEndBoolArray::with_offset_and_size( - self.ends(), - !self.start(), - self.validity(), - self.len(), - self.offset(), - ) - .map(|a| a.into_array()) - } -} +impl BoolArrayTrait for RunEndBoolArray {} impl ArrayVariants for RunEndBoolArray { fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> { diff --git a/encodings/runend-bool/src/compute/invert.rs b/encodings/runend-bool/src/compute/invert.rs new file mode 100644 index 0000000000..2a943ddaaa --- /dev/null +++ b/encodings/runend-bool/src/compute/invert.rs @@ -0,0 +1,19 @@ +use vortex_array::compute::InvertFn; +use vortex_array::{ArrayData, ArrayLen, IntoArrayData}; +use vortex_error::VortexResult; + +use crate::{RunEndBoolArray, RunEndBoolEncoding}; + +impl InvertFn for RunEndBoolEncoding { + fn invert(&self, array: &RunEndBoolArray) -> VortexResult { + RunEndBoolArray::with_offset_and_size( + array.ends(), + // We only need to invert the starting bool + !array.start(), + array.validity(), + array.len(), + array.offset(), + ) + .map(|a| a.into_array()) + } +} diff --git a/encodings/runend-bool/src/compute.rs b/encodings/runend-bool/src/compute/mod.rs similarity index 92% rename from encodings/runend-bool/src/compute.rs rename to encodings/runend-bool/src/compute/mod.rs index 5ff787ade0..57a6bf6cd9 100644 --- a/encodings/runend-bool/src/compute.rs +++ b/encodings/runend-bool/src/compute/mod.rs @@ -1,5 +1,9 @@ +mod invert; + use vortex_array::array::BoolArray; -use vortex_array::compute::{slice, ComputeVTable, ScalarAtFn, SliceFn, TakeFn, TakeOptions}; +use vortex_array::compute::{ + slice, ComputeVTable, InvertFn, ScalarAtFn, SliceFn, TakeFn, TakeOptions, +}; use vortex_array::variants::PrimitiveArrayTrait; use vortex_array::{ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant, ToArrayData}; use vortex_dtype::match_each_integer_ptype; @@ -9,6 +13,10 @@ use vortex_scalar::Scalar; use crate::{value_at_index, RunEndBoolArray, RunEndBoolEncoding}; impl ComputeVTable for RunEndBoolEncoding { + fn invert_fn(&self) -> Option<&dyn InvertFn> { + Some(self) + } + fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn> { Some(self) } diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index ff81721c35..415ba99109 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -15,7 +15,7 @@ use vortex_array::{ IntoArrayVariant, IntoCanonical, }; use vortex_dtype::{DType, PType}; -use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult}; +use vortex_error::{vortex_bail, VortexExpect as _, VortexResult}; use vortex_scalar::Scalar; use crate::compress::{runend_decode_bools, runend_decode_primitive, runend_encode}; @@ -195,22 +195,7 @@ impl ArrayVariants for RunEndArray { impl PrimitiveArrayTrait for RunEndArray {} -impl BoolArrayTrait for RunEndArray { - fn invert(&self) -> VortexResult { - RunEndArray::with_offset_and_length( - self.ends(), - self.values().with_dyn(|v| { - v.as_bool_array() - .ok_or_else(|| vortex_err!("Values were not a bool dtype array"))? - .invert() - })?, - self.validity(), - self.len(), - self.offset(), - ) - .map(|a| a.into_array()) - } -} +impl BoolArrayTrait for RunEndArray {} impl ValidityVTable for RunEndEncoding { fn is_valid(&self, array: &RunEndArray, index: usize) -> bool { diff --git a/encodings/runend/src/compute/invert.rs b/encodings/runend/src/compute/invert.rs new file mode 100644 index 0000000000..d239e3805a --- /dev/null +++ b/encodings/runend/src/compute/invert.rs @@ -0,0 +1,18 @@ +use vortex_array::compute::{invert, InvertFn}; +use vortex_array::{ArrayData, ArrayLen, IntoArrayData}; +use vortex_error::VortexResult; + +use crate::{RunEndArray, RunEndEncoding}; + +impl InvertFn for RunEndEncoding { + fn invert(&self, array: &RunEndArray) -> VortexResult { + RunEndArray::with_offset_and_length( + array.ends(), + invert(&array.values())?, + array.validity(), + array.len(), + array.offset(), + ) + .map(|a| a.into_array()) + } +} diff --git a/encodings/runend/src/compute/mod.rs b/encodings/runend/src/compute/mod.rs index 1a31a2859a..9885e80a38 100644 --- a/encodings/runend/src/compute/mod.rs +++ b/encodings/runend/src/compute/mod.rs @@ -1,4 +1,5 @@ mod compare; +mod invert; use std::cmp::min; use std::ops::AddAssign; @@ -6,8 +7,8 @@ use std::ops::AddAssign; use num_traits::AsPrimitive; use vortex_array::array::{BooleanBuffer, ConstantArray, PrimitiveArray, SparseArray}; use vortex_array::compute::{ - filter, scalar_at, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, ScalarAtFn, - SliceFn, TakeFn, TakeOptions, + filter, scalar_at, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, InvertFn, + ScalarAtFn, SliceFn, TakeFn, TakeOptions, }; use vortex_array::validity::Validity; use vortex_array::variants::PrimitiveArrayTrait; @@ -27,6 +28,10 @@ impl ComputeVTable for RunEndEncoding { Some(self) } + fn invert_fn(&self) -> Option<&dyn InvertFn> { + Some(self) + } + fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn> { Some(self) } diff --git a/vortex-array/src/array/bool/compute/invert.rs b/vortex-array/src/array/bool/compute/invert.rs new file mode 100644 index 0000000000..db5bb53c00 --- /dev/null +++ b/vortex-array/src/array/bool/compute/invert.rs @@ -0,0 +1,13 @@ +use std::ops::Not; + +use vortex_error::VortexResult; + +use crate::array::{BoolArray, BoolEncoding}; +use crate::compute::InvertFn; +use crate::{ArrayData, IntoArrayData}; + +impl InvertFn for BoolEncoding { + fn invert(&self, array: &BoolArray) -> VortexResult { + Ok(BoolArray::try_new(array.boolean_buffer().not(), array.validity())?.into_array()) + } +} diff --git a/vortex-array/src/array/bool/compute/mod.rs b/vortex-array/src/array/bool/compute/mod.rs index a280510d2e..0111c82ca5 100644 --- a/vortex-array/src/array/bool/compute/mod.rs +++ b/vortex-array/src/array/bool/compute/mod.rs @@ -1,12 +1,13 @@ use crate::array::BoolEncoding; use crate::compute::{ - BinaryBooleanFn, ComputeVTable, FillForwardFn, FilterFn, ScalarAtFn, SliceFn, TakeFn, + BinaryBooleanFn, ComputeVTable, FillForwardFn, FilterFn, InvertFn, ScalarAtFn, SliceFn, TakeFn, }; use crate::ArrayData; mod fill; pub mod filter; mod flatten; +mod invert; mod scalar_at; mod slice; mod take; @@ -28,6 +29,10 @@ impl ComputeVTable for BoolEncoding { Some(self) } + fn invert_fn(&self) -> Option<&dyn InvertFn> { + Some(self) + } + fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn> { Some(self) } diff --git a/vortex-array/src/array/bool/mod.rs b/vortex-array/src/array/bool/mod.rs index b1fc057d9a..5e2db6a519 100644 --- a/vortex-array/src/array/bool/mod.rs +++ b/vortex-array/src/array/bool/mod.rs @@ -184,11 +184,7 @@ impl ArrayVariants for BoolArray { } } -impl BoolArrayTrait for BoolArray { - fn invert(&self) -> VortexResult { - Ok(BoolArray::try_new(!&self.boolean_buffer(), self.validity())?.into_array()) - } -} +impl BoolArrayTrait for BoolArray {} impl From for BoolArray { fn from(value: BooleanBuffer) -> Self { diff --git a/vortex-array/src/array/chunked/compute/invert.rs b/vortex-array/src/array/chunked/compute/invert.rs new file mode 100644 index 0000000000..303fa9e016 --- /dev/null +++ b/vortex-array/src/array/chunked/compute/invert.rs @@ -0,0 +1,13 @@ +use itertools::Itertools; +use vortex_error::VortexResult; + +use crate::array::{ChunkedArray, ChunkedEncoding}; +use crate::compute::{invert, InvertFn}; +use crate::{ArrayDType, ArrayData, IntoArrayData}; + +impl InvertFn for ChunkedEncoding { + fn invert(&self, array: &ChunkedArray) -> VortexResult { + let chunks = array.chunks().map(|c| invert(&c)).try_collect()?; + ChunkedArray::try_new(chunks, array.dtype().clone()).map(|a| a.into_array()) + } +} diff --git a/vortex-array/src/array/chunked/compute/mod.rs b/vortex-array/src/array/chunked/compute/mod.rs index c66833a924..1d1e5be899 100644 --- a/vortex-array/src/array/chunked/compute/mod.rs +++ b/vortex-array/src/array/chunked/compute/mod.rs @@ -4,12 +4,13 @@ use vortex_error::VortexResult; use crate::array::chunked::ChunkedArray; use crate::array::ChunkedEncoding; use crate::compute::{ - compare, slice, try_cast, CastFn, CompareFn, ComputeVTable, FilterFn, Operator, ScalarAtFn, - SliceFn, SubtractScalarFn, TakeFn, + compare, slice, try_cast, CastFn, CompareFn, ComputeVTable, FilterFn, InvertFn, Operator, + ScalarAtFn, SliceFn, SubtractScalarFn, TakeFn, }; use crate::{ArrayData, IntoArrayData}; mod filter; +mod invert; mod scalar_at; mod slice; mod take; @@ -27,6 +28,10 @@ impl ComputeVTable for ChunkedEncoding { Some(self) } + fn invert_fn(&self) -> Option<&dyn InvertFn> { + Some(self) + } + fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn> { Some(self) } diff --git a/vortex-array/src/array/chunked/variants.rs b/vortex-array/src/array/chunked/variants.rs index d364ca92cb..0efd383d8d 100644 --- a/vortex-array/src/array/chunked/variants.rs +++ b/vortex-array/src/array/chunked/variants.rs @@ -46,21 +46,7 @@ impl ArrayVariants for ChunkedArray { impl NullArrayTrait for ChunkedArray {} -impl BoolArrayTrait for ChunkedArray { - fn invert(&self) -> VortexResult { - let chunks = self - .chunks() - .map(|c| { - c.with_dyn(|a| { - a.as_bool_array() - .ok_or_else(|| vortex_err!("Child was not a bool array")) - .and_then(|b| b.invert()) - }) - }) - .collect::>>()?; - ChunkedArray::try_new(chunks, self.dtype().clone()).map(|a| a.into_array()) - } -} +impl BoolArrayTrait for ChunkedArray {} impl PrimitiveArrayTrait for ChunkedArray {} diff --git a/vortex-array/src/array/constant/compute/invert.rs b/vortex-array/src/array/constant/compute/invert.rs new file mode 100644 index 0000000000..0a376e354c --- /dev/null +++ b/vortex-array/src/array/constant/compute/invert.rs @@ -0,0 +1,14 @@ +use vortex_error::VortexResult; + +use crate::array::{ConstantArray, ConstantEncoding}; +use crate::compute::InvertFn; +use crate::{ArrayData, ArrayLen, IntoArrayData, ToArrayData}; + +impl InvertFn for ConstantEncoding { + fn invert(&self, array: &ConstantArray) -> VortexResult { + match array.scalar().as_bool().value() { + None => Ok(array.to_array()), + Some(b) => Ok(ConstantArray::new(!b, array.len()).into_array()), + } + } +} diff --git a/vortex-array/src/array/constant/compute/mod.rs b/vortex-array/src/array/constant/compute/mod.rs index 2384bd34d0..c3bd9994fb 100644 --- a/vortex-array/src/array/constant/compute/mod.rs +++ b/vortex-array/src/array/constant/compute/mod.rs @@ -1,5 +1,6 @@ mod boolean; mod compare; +mod invert; mod search_sorted; use vortex_error::VortexResult; @@ -8,8 +9,8 @@ use vortex_scalar::Scalar; use crate::array::constant::ConstantArray; use crate::array::ConstantEncoding; use crate::compute::{ - BinaryBooleanFn, CompareFn, ComputeVTable, FilterFn, FilterMask, ScalarAtFn, SearchSortedFn, - SliceFn, TakeFn, TakeOptions, + BinaryBooleanFn, CompareFn, ComputeVTable, FilterFn, FilterMask, InvertFn, ScalarAtFn, + SearchSortedFn, SliceFn, TakeFn, TakeOptions, }; use crate::{ArrayData, IntoArrayData}; @@ -26,6 +27,10 @@ impl ComputeVTable for ConstantEncoding { Some(self) } + fn invert_fn(&self) -> Option<&dyn InvertFn> { + Some(self) + } + fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn> { Some(self) } diff --git a/vortex-array/src/array/constant/variants.rs b/vortex-array/src/array/constant/variants.rs index d37af26f4c..42f926d6a5 100644 --- a/vortex-array/src/array/constant/variants.rs +++ b/vortex-array/src/array/constant/variants.rs @@ -10,7 +10,7 @@ use crate::variants::{ ArrayVariants, BinaryArrayTrait, BoolArrayTrait, ExtensionArrayTrait, ListArrayTrait, NullArrayTrait, PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait, }; -use crate::{ArrayDType, ArrayData, ArrayLen, IntoArrayData, ToArrayData}; +use crate::{ArrayDType, ArrayData, ArrayLen, IntoArrayData}; /// Constant arrays support all DTypes impl ArrayVariants for ConstantArray { @@ -49,14 +49,7 @@ impl ArrayVariants for ConstantArray { impl NullArrayTrait for ConstantArray {} -impl BoolArrayTrait for ConstantArray { - fn invert(&self) -> VortexResult { - match self.scalar().as_bool().value() { - None => Ok(self.to_array()), - Some(b) => Ok(ConstantArray::new(!b, self.len()).into_array()), - } - } -} +impl BoolArrayTrait for ConstantArray {} impl Accessor for ConstantArray where diff --git a/vortex-array/src/array/sparse/compute/invert.rs b/vortex-array/src/array/sparse/compute/invert.rs new file mode 100644 index 0000000000..cdb6e4be80 --- /dev/null +++ b/vortex-array/src/array/sparse/compute/invert.rs @@ -0,0 +1,18 @@ +use vortex_error::VortexResult; + +use crate::array::{SparseArray, SparseEncoding}; +use crate::compute::{invert, InvertFn}; +use crate::{ArrayData, ArrayLen, IntoArrayData}; + +impl InvertFn for SparseEncoding { + fn invert(&self, array: &SparseArray) -> VortexResult { + let inverted_fill = array.fill_scalar().as_bool().invert().into_scalar(); + SparseArray::try_new( + array.indices(), + invert(&array.values())?, + array.len(), + inverted_fill, + ) + .map(|a| a.into_array()) + } +} diff --git a/vortex-array/src/array/sparse/compute/mod.rs b/vortex-array/src/array/sparse/compute/mod.rs index 17c6951e8c..3433ed85f3 100644 --- a/vortex-array/src/array/sparse/compute/mod.rs +++ b/vortex-array/src/array/sparse/compute/mod.rs @@ -5,12 +5,13 @@ use vortex_scalar::Scalar; use crate::array::sparse::SparseArray; use crate::array::{PrimitiveArray, SparseEncoding}; use crate::compute::{ - scalar_at, search_sorted, take, ComputeVTable, FilterFn, FilterMask, ScalarAtFn, SearchResult, - SearchSortedFn, SearchSortedSide, SliceFn, TakeFn, TakeOptions, + scalar_at, search_sorted, take, ComputeVTable, FilterFn, FilterMask, InvertFn, ScalarAtFn, + SearchResult, SearchSortedFn, SearchSortedSide, SliceFn, TakeFn, TakeOptions, }; use crate::variants::PrimitiveArrayTrait; use crate::{ArrayData, IntoArrayData, IntoArrayVariant}; +mod invert; mod slice; mod take; @@ -19,6 +20,10 @@ impl ComputeVTable for SparseEncoding { Some(self) } + fn invert_fn(&self) -> Option<&dyn InvertFn> { + Some(self) + } + fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn> { Some(self) } diff --git a/vortex-array/src/array/sparse/variants.rs b/vortex-array/src/array/sparse/variants.rs index 88af892f09..62f65d55ba 100644 --- a/vortex-array/src/array/sparse/variants.rs +++ b/vortex-array/src/array/sparse/variants.rs @@ -47,22 +47,7 @@ impl ArrayVariants for SparseArray { impl NullArrayTrait for SparseArray {} -impl BoolArrayTrait for SparseArray { - fn invert(&self) -> VortexResult { - let inverted_fill = self.fill_scalar().as_bool().invert().into_scalar(); - SparseArray::try_new( - self.indices(), - self.values().with_dyn(|a| { - a.as_bool_array() - .ok_or_else(|| vortex_err!("Not a bool array")) - .and_then(|b| b.invert()) - })?, - self.len(), - inverted_fill, - ) - .map(|a| a.into_array()) - } -} +impl BoolArrayTrait for SparseArray {} impl PrimitiveArrayTrait for SparseArray {} @@ -133,6 +118,7 @@ mod tests { use vortex_scalar::Scalar; use crate::array::{BoolArray, PrimitiveArray, SparseArray}; + use crate::compute::invert; use crate::{IntoArrayData, IntoArrayVariant}; #[test] @@ -145,9 +131,7 @@ mod tests { ) .unwrap() .into_array(); - let inverted = sparse_bools - .with_dyn(|a| a.as_bool_array_unchecked().invert()) - .unwrap(); + let inverted = invert(&sparse_bools).unwrap(); assert_eq!( inverted .into_bool() diff --git a/vortex-array/src/compute/invert.rs b/vortex-array/src/compute/invert.rs new file mode 100644 index 0000000000..0a25c3a373 --- /dev/null +++ b/vortex-array/src/compute/invert.rs @@ -0,0 +1,44 @@ +use vortex_dtype::DType; +use vortex_error::{vortex_bail, vortex_err, VortexError, VortexResult}; + +use crate::encoding::Encoding; +use crate::{ArrayDType, ArrayData, IntoArrayData, IntoArrayVariant}; + +pub trait InvertFn { + /// Logically invert a boolean array. Converts true -> false, false -> true, null -> null. + fn invert(&self, array: &Array) -> VortexResult; +} + +impl InvertFn for E +where + E: InvertFn, + for<'a> &'a E::Array: TryFrom<&'a ArrayData, Error = VortexError>, +{ + fn invert(&self, array: &ArrayData) -> VortexResult { + let array_ref = <&E::Array>::try_from(array)?; + let encoding = array + .encoding() + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Mismatched encoding"))?; + InvertFn::invert(encoding, array_ref) + } +} + +/// Logically invert a boolean array. +pub fn invert(array: &ArrayData) -> VortexResult { + if !matches!(array.dtype(), DType::Bool(..)) { + vortex_bail!("Expected boolean array, got {}", array.dtype()); + } + + if let Some(f) = array.encoding().invert_fn() { + return f.invert(array); + } + + // Otherwise, we canonicalize into a boolean array and invert. + log::debug!( + "No invert implementation found for encoding {}", + array.encoding().id(), + ); + invert(&array.clone().into_bool()?.into_array()) +} diff --git a/vortex-array/src/compute/mod.rs b/vortex-array/src/compute/mod.rs index 711d3c19df..1bfe86a507 100644 --- a/vortex-array/src/compute/mod.rs +++ b/vortex-array/src/compute/mod.rs @@ -12,6 +12,7 @@ pub use cast::{try_cast, CastFn}; pub use compare::{compare, scalar_cmp, CompareFn, Operator}; pub use fill_forward::{fill_forward, FillForwardFn}; pub use filter::*; +pub use invert::{invert, InvertFn}; pub use scalar_at::{scalar_at, ScalarAtFn}; pub use scalar_subtract::{subtract_scalar, SubtractScalarFn}; pub use search_sorted::*; @@ -25,6 +26,7 @@ mod cast; mod compare; mod fill_forward; mod filter; +mod invert; mod scalar_at; mod scalar_subtract; mod search_sorted; @@ -68,6 +70,13 @@ pub trait ComputeVTable { None } + /// Invert a boolean array. Converts true -> false, false -> true, null -> null. + /// + /// See [InvertFn] + fn invert_fn(&self) -> Option<&dyn InvertFn> { + None + } + /// Single item indexing on Vortex arrays. /// /// See: [ScalarAtFn]. diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs index ee6a8af742..74be72b40e 100644 --- a/vortex-array/src/variants.rs +++ b/vortex-array/src/variants.rs @@ -81,14 +81,7 @@ pub trait ArrayVariants { pub trait NullArrayTrait: ArrayTrait {} -pub trait BoolArrayTrait: ArrayTrait { - /// Return a new inverted version of this array. - /// - /// True -> False - /// False -> True - /// Null -> Null - fn invert(&self) -> VortexResult; -} +pub trait BoolArrayTrait: ArrayTrait {} /// Iterate over an array of primitives by dispatching at run-time on the array type. #[macro_export] diff --git a/vortex-expr/src/not.rs b/vortex-expr/src/not.rs index 2ed0383c45..9a7e2acf43 100644 --- a/vortex-expr/src/not.rs +++ b/vortex-expr/src/not.rs @@ -3,9 +3,10 @@ use std::fmt::Display; use std::sync::Arc; use vortex_array::aliases::hash_set::HashSet; +use vortex_array::compute::invert; use vortex_array::ArrayData; use vortex_dtype::field::Field; -use vortex_error::{vortex_err, VortexResult}; +use vortex_error::VortexResult; use crate::{unbox_any, ExprRef, VortexExpr}; @@ -38,11 +39,7 @@ impl VortexExpr for Not { fn evaluate(&self, batch: &ArrayData) -> VortexResult { let child_result = self.child.evaluate(batch)?; - child_result.with_dyn(|a| { - a.as_bool_array() - .ok_or_else(|| vortex_err!("Child was not a bool array")) - .and_then(|b| b.invert()) - }) + invert(&child_result) } fn collect_references<'a>(&'a self, references: &mut HashSet<&'a Field>) {