feat: run end bool compressor (#1355)

Co-authored-by: Robert Kruszewski <[email protected]>

Author: lwwmanning

bench-vortex/benches/compressor_throughput.rs

@@ -5,8 +5,9 @@ use rand::distributions::Alphanumeric;
 use rand::seq::SliceRandom as _;
 use rand::{thread_rng, Rng, SeedableRng as _};
 use vortex::aliases::hash_set::HashSet;
-use vortex::array::{PrimitiveArray, VarBinViewArray};
+use vortex::array::{ConstantArray, PrimitiveArray, VarBinViewArray};
 use vortex::compute::unary::try_cast;
+use vortex::compute::{compare, Operator};
 use vortex::dict::{dict_encode_varbinview, DictArray};
 use vortex::dtype::PType;
 use vortex::fsst::{fsst_compress, fsst_train_compressor};
@@ -18,11 +19,14 @@ use vortex::sampling_compressor::compressors::bitpacked::{
 use vortex::sampling_compressor::compressors::delta::DeltaCompressor;
 use vortex::sampling_compressor::compressors::dict::DictCompressor;
 use vortex::sampling_compressor::compressors::r#for::FoRCompressor;
+use vortex::sampling_compressor::compressors::roaring_bool::RoaringBoolCompressor;
 use vortex::sampling_compressor::compressors::roaring_int::RoaringIntCompressor;
 use vortex::sampling_compressor::compressors::runend::DEFAULT_RUN_END_COMPRESSOR;
+use vortex::sampling_compressor::compressors::runend_bool::RunEndBoolCompressor;
 use vortex::sampling_compressor::compressors::zigzag::ZigZagCompressor;
 use vortex::sampling_compressor::compressors::CompressorRef;
 use vortex::sampling_compressor::SamplingCompressor;
+use vortex::scalar::Scalar;
 use vortex::validity::Validity;
 use vortex::{IntoArrayData as _, IntoCanonical};
 
@@ -44,11 +48,20 @@ fn primitive(c: &mut Criterion) {
     )
     .into_array();
     let int_array = try_cast(uint_array.clone(), PType::I32.into()).unwrap();
+
+    let bool_array = compare(
+        &uint_array,
+        ConstantArray::new(Scalar::from(0u32), uint_array.len()),
+        Operator::Eq,
+    )
+    .unwrap();
+
     let index_array = PrimitiveArray::from_vec(
         (0..num_values).map(|i| (i * 2) as u32 + 42).collect_vec(),
         Validity::NonNullable,
     )
     .into_array();
+
     let float_array = try_cast(uint_array.clone(), PType::F32.into()).unwrap();
 
     let compressors_names_and_arrays = [
@@ -61,7 +74,9 @@ fn primitive(c: &mut Criterion) {
         (&DEFAULT_RUN_END_COMPRESSOR, "runend", &uint_array),
         (&DeltaCompressor, "delta", &uint_array),
         (&DictCompressor, "dict", &uint_array),
+        (&RoaringBoolCompressor, "roaring_bool", &bool_array),
         (&RoaringIntCompressor, "roaring_int", &index_array),
+        (&RunEndBoolCompressor, "runend_bool", &bool_array),
         (&FoRCompressor, "frame_of_reference", &int_array),
         (&ZigZagCompressor, "zigzag", &int_array),
         (&ALPCompressor, "alp", &float_array),

encodings/runend-bool/src/array.rs

@@ -2,6 +2,7 @@ use std::fmt::{Debug, Display};
 
 use serde::{Deserialize, Serialize};
 use vortex_array::array::visitor::{AcceptArrayVisitor, ArrayVisitor};
+use vortex_array::array::{BoolArray, PrimitiveArray};
 use vortex_array::compute::unary::scalar_at;
 use vortex_array::compute::{search_sorted, SearchSortedSide};
 use vortex_array::encoding::ids;
@@ -12,11 +13,11 @@ use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayTrait, Canonical, IntoArrayData, IntoArrayVariant,
     IntoCanonical,
 };
-use vortex_dtype::{match_each_unsigned_integer_ptype, DType, PType};
+use vortex_dtype::{match_each_integer_ptype, match_each_unsigned_integer_ptype, DType, PType};
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 use vortex_scalar::Scalar;
 
-use crate::compress::runend_bool_decode;
+use crate::compress::{runend_bool_decode_slice, runend_bool_encode_slice};
 
 impl_encoding!("vortex.runendbool", ids::RUN_END_BOOL, RunEndBool);
 
@@ -146,6 +147,36 @@ impl RunEndBoolArray {
     }
 }
 
+pub fn encode_runend_bool(array: &BoolArray) -> VortexResult<RunEndBoolArray> {
+    let (ends, start) = runend_bool_encode_slice(&array.boolean_buffer());
+    RunEndBoolArray::try_new(
+        PrimitiveArray::from(ends).into_array(),
+        start,
+        array.validity(),
+    )
+}
+
+pub(crate) fn decode_runend_bool(
+    run_ends: &PrimitiveArray,
+    start: bool,
+    validity: Validity,
+    offset: usize,
+    length: usize,
+) -> VortexResult<BoolArray> {
+    match_each_integer_ptype!(run_ends.ptype(), |$E| {
+        let bools = runend_bool_decode_slice::<$E>(run_ends.maybe_null_slice(), start, offset, length);
+        Ok(BoolArray::try_new(bools, validity)?)
+    })
+}
+
+pub(crate) fn value_at_index(idx: usize, start: bool) -> bool {
+    if idx % 2 == 0 {
+        start
+    } else {
+        !start
+    }
+}
+
 impl BoolArrayTrait for RunEndBoolArray {
     fn invert(&self) -> VortexResult<ArrayData> {
         RunEndBoolArray::try_new(self.ends(), !self.start(), self.validity())
@@ -174,7 +205,7 @@ impl ArrayValidity for RunEndBoolArray {
 impl IntoCanonical for RunEndBoolArray {
     fn into_canonical(self) -> VortexResult<Canonical> {
         let pends = self.ends().into_primitive()?;
-        runend_bool_decode(
+        decode_runend_bool(
             &pends,
             self.start(),
             self.validity(),

encodings/runend-bool/src/compress.rs

@@ -3,16 +3,10 @@ use std::cmp::min;
 use arrow_buffer::buffer::BooleanBuffer;
 use arrow_buffer::BooleanBufferBuilder;
 use num_traits::{AsPrimitive, FromPrimitive};
-use vortex_array::array::{BoolArray, PrimitiveArray};
-use vortex_array::validity::Validity;
-use vortex_array::variants::PrimitiveArrayTrait;
-use vortex_dtype::{match_each_integer_ptype, NativePType};
-use vortex_error::{vortex_panic, VortexExpect as _, VortexResult};
-
-pub fn runend_bool_encode(elements: &BoolArray) -> (PrimitiveArray, bool) {
-    let (arr, start) = runend_bool_encode_slice(&elements.boolean_buffer());
-    (arr.into(), start)
-}
+use vortex_dtype::NativePType;
+use vortex_error::{vortex_panic, VortexExpect as _};
+
+use crate::value_at_index;
 
 pub fn runend_bool_encode_slice(elements: &BooleanBuffer) -> (Vec<u64>, bool) {
     let mut iter = elements.set_slices();
@@ -41,19 +35,6 @@ pub fn runend_bool_encode_slice(elements: &BooleanBuffer) -> (Vec<u64>, bool) {
     (ends, first_bool)
 }
 
-pub fn runend_bool_decode(
-    run_ends: &PrimitiveArray,
-    start: bool,
-    validity: Validity,
-    offset: usize,
-    length: usize,
-) -> VortexResult<BoolArray> {
-    match_each_integer_ptype!(run_ends.ptype(), |$E| {
-        let bools = runend_bool_decode_slice::<$E>(run_ends.maybe_null_slice(), start, offset, length);
-        Ok(BoolArray::try_new(bools, validity)?)
-    })
-}
-
 pub fn runend_bool_decode_slice<E: NativePType + AsPrimitive<usize> + FromPrimitive + Ord>(
     run_ends: &[E],
     start: bool,
@@ -86,28 +67,19 @@ pub fn runend_bool_decode_slice<E: NativePType + AsPrimitive<usize> + FromPrimit
     BooleanBuffer::from(decoded)
 }
 
-pub fn value_at_index(idx: usize, start: bool) -> bool {
-    if idx % 2 == 0 {
-        start
-    } else {
-        !start
-    }
-}
-
 #[cfg(test)]
 mod test {
     use arrow_buffer::BooleanBuffer;
     use itertools::Itertools;
     use rand::prelude::StdRng;
     use rand::{Rng, SeedableRng};
-    use vortex_array::array::BoolArray;
+    use vortex_array::array::{BoolArray, PrimitiveArray};
     use vortex_array::compute::SliceFn;
     use vortex_array::validity::Validity;
     use vortex_array::IntoArrayVariant;
 
-    use crate::compress::{
-        runend_bool_decode, runend_bool_decode_slice, runend_bool_encode, runend_bool_encode_slice,
-    };
+    use crate::compress::{runend_bool_decode_slice, runend_bool_encode_slice};
+    use crate::decode_runend_bool;
 
     #[test]
     fn encode_bool() {
@@ -175,9 +147,10 @@ mod test {
         let input = (0..1024 * 8 - 61).map(|_x| rng.gen::<bool>()).collect_vec();
         let b = BoolArray::from_iter(input.clone());
         let b = b.slice(3, 1024 * 8 - 66).unwrap().into_bool().unwrap();
-        let (ends, start) = runend_bool_encode(&b);
+        let (ends, start) = runend_bool_encode_slice(&b.boolean_buffer());
+        let ends = PrimitiveArray::from(ends);
 
-        let decoded = runend_bool_decode(&ends, start, Validity::NonNullable, 0, 1024 * 8 - 69)
+        let decoded = decode_runend_bool(&ends, start, Validity::NonNullable, 0, 1024 * 8 - 69)
             .unwrap()
             .into_bool()
             .unwrap()

encodings/runend-bool/src/compute.rs

@@ -7,8 +7,7 @@ use vortex_dtype::match_each_integer_ptype;
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 use vortex_scalar::Scalar;
 
-use crate::compress::value_at_index;
-use crate::RunEndBoolArray;
+use crate::{value_at_index, RunEndBoolArray};
 
 impl ArrayCompute for RunEndBoolArray {
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {

vortex-sampling-compressor/src/compressors/mod.rs

@@ -24,6 +24,7 @@ pub mod fsst;
 pub mod roaring_bool;
 pub mod roaring_int;
 pub mod runend;
+pub mod runend_bool;
 pub mod sparse;
 pub mod struct_;
 pub mod zigzag;

vortex-sampling-compressor/src/compressors/runend_bool.rs

@@ -0,0 +1,59 @@
+use vortex_array::aliases::hash_set::HashSet;
+use vortex_array::array::{Bool, PrimitiveArray};
+use vortex_array::encoding::EncodingRef;
+use vortex_array::stats::ArrayStatistics as _;
+use vortex_array::{ArrayData, ArrayDef, IntoArrayData, IntoArrayVariant};
+use vortex_error::VortexResult;
+use vortex_runend_bool::compress::runend_bool_encode_slice;
+use vortex_runend_bool::{RunEndBool, RunEndBoolArray, RunEndBoolEncoding};
+
+use crate::compressors::{CompressedArray, CompressionTree, EncodingCompressor};
+use crate::{constants, SamplingCompressor};
+
+#[derive(Debug)]
+pub struct RunEndBoolCompressor;
+
+impl EncodingCompressor for RunEndBoolCompressor {
+    fn id(&self) -> &str {
+        RunEndBool::ID.as_ref()
+    }
+
+    fn cost(&self) -> u8 {
+        constants::RUN_END_BOOL_COST
+    }
+
+    fn can_compress(&self, array: &ArrayData) -> Option<&dyn EncodingCompressor> {
+        // Only support bool arrays
+        if !array.is_encoding(Bool::ID) {
+            return None;
+        }
+
+        Some(self)
+    }
+
+    fn compress<'a>(
+        &'a self,
+        array: &ArrayData,
+        like: Option<CompressionTree<'a>>,
+        ctx: SamplingCompressor<'a>,
+    ) -> VortexResult<CompressedArray<'a>> {
+        let bool_array = array.clone().into_bool()?;
+        let (ends, start) = runend_bool_encode_slice(&bool_array.boolean_buffer());
+        let ends = PrimitiveArray::from(ends);
+
+        let compressed_ends = ctx
+            .auxiliary("ends")
+            .compress(&ends.into_array(), like.as_ref().and_then(|l| l.child(0)))?;
+
+        Ok(CompressedArray::compressed(
+            RunEndBoolArray::try_new(compressed_ends.array, start, bool_array.validity())?
+                .into_array(),
+            Some(CompressionTree::new(self, vec![compressed_ends.path])),
+            Some(array.statistics()),
+        ))
+    }
+
+    fn used_encodings(&self) -> HashSet<EncodingRef> {
+        HashSet::from([&RunEndBoolEncoding as EncodingRef])
+    }
+}

vortex-sampling-compressor/src/compressors/struct_.rs

@@ -5,7 +5,8 @@ use vortex_array::compress::compute_precompression_stats;
 use vortex_array::encoding::EncodingRef;
 use vortex_array::stats::ArrayStatistics as _;
 use vortex_array::variants::StructArrayTrait;
-use vortex_array::{ArrayData, ArrayDef, IntoArrayData};
+use vortex_array::{ArrayDType, ArrayData, ArrayDef, IntoArrayData};
+use vortex_dtype::DType;
 use vortex_error::VortexResult;
 
 use crate::compressors::{CompressedArray, CompressionTree, EncodingCompressor};
@@ -24,9 +25,8 @@ impl EncodingCompressor for StructCompressor {
     }
 
     fn can_compress(&self, array: &ArrayData) -> Option<&dyn EncodingCompressor> {
-        StructArray::try_from(array)
-            .ok()
-            .map(|_| self as &dyn EncodingCompressor)
+        let is_struct = matches!(array.dtype(), DType::Struct(..)) && array.is_encoding(Struct::ID);
+        is_struct.then_some(self)
     }
 
     fn compress<'a>(

vortex-sampling-compressor/src/constants.rs

@@ -1,8 +1,8 @@
 #![allow(dead_code)]
 
 // structural pass-throughs have no cost
-pub const SPARSE_COST: u8 = 0;
 pub const CHUNKED_COST: u8 = 0;
+pub const SPARSE_COST: u8 = 0;
 pub const STRUCT_COST: u8 = 0;
 
 // so fast that we can ignore the cost
@@ -20,6 +20,7 @@ pub const FOR_COST: u8 = 1;
 pub const FSST_COST: u8 = 1;
 pub const ROARING_BOOL_COST: u8 = 1;
 pub const ROARING_INT_COST: u8 = 1;
+pub const RUN_END_BOOL_COST: u8 = 1;
 pub const RUN_END_COST: u8 = 1;
 
 // "expensive" encodings