Move invert from BoolArrayTrait to InvertFn (#1490)

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<ArrayData> {
-        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<Vec<bool>> for ByteBoolArray {
     fn from(value: Vec<bool>) -> Self {

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<ArrayData> {
-        todo!()
-    }
-}
+impl BoolArrayTrait for DictArray {}
 
 impl PrimitiveArrayTrait for DictArray {}
 

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<ArrayData>> {
+        Some(self)
+    }
+
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
         Some(self)
     }
@@ -16,6 +20,13 @@ impl ComputeVTable for RoaringBoolEncoding {
     }
 }
 
+impl InvertFn<RoaringBoolArray> for RoaringBoolEncoding {
+    fn invert(&self, array: &RoaringBoolArray) -> VortexResult<ArrayData> {
+        RoaringBoolArray::try_new(array.bitmap().flip(0..(array.len() as u32)), array.len())
+            .map(|a| a.into_array())
+    }
+}
+
 impl ScalarAtFn<RoaringBoolArray> for RoaringBoolEncoding {
     fn scalar_at(&self, array: &RoaringBoolArray, index: usize) -> VortexResult<Scalar> {
         Ok(array.bitmap().contains(index as u32).into())

encodings/roaring/src/boolean/mod.rs

@@ -90,12 +90,7 @@ impl ArrayVariants for RoaringBoolArray {
     }
 }
 
-impl BoolArrayTrait for RoaringBoolArray {
-    fn invert(&self) -> VortexResult<ArrayData> {
-        RoaringBoolArray::try_new(self.bitmap().flip(0..(self.len() as u32)), self.len())
-            .map(|a| a.into_array())
-    }
-}
+impl BoolArrayTrait for RoaringBoolArray {}
 
 impl VisitorVTable<RoaringBoolArray> for RoaringBoolEncoding {
     fn accept(&self, array: &RoaringBoolArray, visitor: &mut dyn ArrayVisitor) -> VortexResult<()> {

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<ArrayData> {
-        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> {

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<RunEndBoolArray> for RunEndBoolEncoding {
+    fn invert(&self, array: &RunEndBoolArray) -> VortexResult<ArrayData> {
+        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())
+    }
+}

encodings/runend-bool/src/compute.rs → 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<ArrayData>> {
+        Some(self)
+    }
+
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
         Some(self)
     }

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<ArrayData> {
-        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<RunEndArray> for RunEndEncoding {
     fn is_valid(&self, array: &RunEndArray, index: usize) -> bool {

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<RunEndArray> for RunEndEncoding {
+    fn invert(&self, array: &RunEndArray) -> VortexResult<ArrayData> {
+        RunEndArray::with_offset_and_length(
+            array.ends(),
+            invert(&array.values())?,
+            array.validity(),
+            array.len(),
+            array.offset(),
+        )
+        .map(|a| a.into_array())
+    }
+}

encodings/runend/src/compute/mod.rs

@@ -1,13 +1,14 @@
 mod compare;
+mod invert;
 
 use std::cmp::min;
 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<ArrayData>> {
+        Some(self)
+    }
+
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
         Some(self)
     }

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<BoolArray> for BoolEncoding {
+    fn invert(&self, array: &BoolArray) -> VortexResult<ArrayData> {
+        Ok(BoolArray::try_new(array.boolean_buffer().not(), array.validity())?.into_array())
+    }
+}

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<ArrayData>> {
+        Some(self)
+    }
+
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
         Some(self)
     }

vortex-array/src/array/bool/mod.rs

@@ -184,11 +184,7 @@ impl ArrayVariants for BoolArray {
     }
 }
 
-impl BoolArrayTrait for BoolArray {
-    fn invert(&self) -> VortexResult<ArrayData> {
-        Ok(BoolArray::try_new(!&self.boolean_buffer(), self.validity())?.into_array())
-    }
-}
+impl BoolArrayTrait for BoolArray {}
 
 impl From<BooleanBuffer> for BoolArray {
     fn from(value: BooleanBuffer) -> Self {

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<ChunkedArray> for ChunkedEncoding {
+    fn invert(&self, array: &ChunkedArray) -> VortexResult<ArrayData> {
+        let chunks = array.chunks().map(|c| invert(&c)).try_collect()?;
+        ChunkedArray::try_new(chunks, array.dtype().clone()).map(|a| a.into_array())
+    }
+}

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<ArrayData>> {
+        Some(self)
+    }
+
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
         Some(self)
     }

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<ArrayData> {
-        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::<VortexResult<Vec<_>>>()?;
-        ChunkedArray::try_new(chunks, self.dtype().clone()).map(|a| a.into_array())
-    }
-}
+impl BoolArrayTrait for ChunkedArray {}
 
 impl PrimitiveArrayTrait for ChunkedArray {}
 

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<ConstantArray> for ConstantEncoding {
+    fn invert(&self, array: &ConstantArray) -> VortexResult<ArrayData> {
+        match array.scalar().as_bool().value() {
+            None => Ok(array.to_array()),
+            Some(b) => Ok(ConstantArray::new(!b, array.len()).into_array()),
+        }
+    }
+}