Add compressor for constant nonnullable and all valid bool arrays

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-btrblocks/public-api.lock

@@ -1,5 +1,53 @@
 pub mod vortex_btrblocks
 
+pub enum vortex_btrblocks::BoolCode
+
+pub vortex_btrblocks::BoolCode::Constant
+
+pub vortex_btrblocks::BoolCode::Uncompressed
+
+impl core::clone::Clone for vortex_btrblocks::BoolCode
+
+pub fn vortex_btrblocks::BoolCode::clone(&self) -> vortex_btrblocks::BoolCode
+
+impl core::cmp::Eq for vortex_btrblocks::BoolCode
+
+impl core::cmp::Ord for vortex_btrblocks::BoolCode
+
+pub fn vortex_btrblocks::BoolCode::cmp(&self, other: &vortex_btrblocks::BoolCode) -> core::cmp::Ordering
+
+impl core::cmp::PartialEq for vortex_btrblocks::BoolCode
+
+pub fn vortex_btrblocks::BoolCode::eq(&self, other: &vortex_btrblocks::BoolCode) -> bool
+
+impl core::cmp::PartialOrd for vortex_btrblocks::BoolCode
+
+pub fn vortex_btrblocks::BoolCode::partial_cmp(&self, other: &vortex_btrblocks::BoolCode) -> core::option::Option<core::cmp::Ordering>
+
+impl core::fmt::Debug for vortex_btrblocks::BoolCode
+
+pub fn vortex_btrblocks::BoolCode::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::hash::Hash for vortex_btrblocks::BoolCode
+
+pub fn vortex_btrblocks::BoolCode::hash<__H: core::hash::Hasher>(&self, state: &mut __H)
+
+impl core::marker::Copy for vortex_btrblocks::BoolCode
+
+impl core::marker::StructuralPartialEq for vortex_btrblocks::BoolCode
+
+impl enum_iterator::Sequence for vortex_btrblocks::BoolCode
+
+pub const vortex_btrblocks::BoolCode::CARDINALITY: usize
+
+pub fn vortex_btrblocks::BoolCode::first() -> core::option::Option<Self>
+
+pub fn vortex_btrblocks::BoolCode::last() -> core::option::Option<Self>
+
+pub fn vortex_btrblocks::BoolCode::next(&self) -> core::option::Option<Self>
+
+pub fn vortex_btrblocks::BoolCode::previous(&self) -> core::option::Option<Self>
+
 pub enum vortex_btrblocks::FloatCode
 
 pub vortex_btrblocks::FloatCode::Alp
@@ -188,6 +236,8 @@ pub fn vortex_btrblocks::StringCode::previous(&self) -> core::option::Option<Sel
 
 pub struct vortex_btrblocks::BtrBlocksCompressor
 
+pub vortex_btrblocks::BtrBlocksCompressor::bool_schemes: alloc::vec::Vec<&'static dyn vortex_btrblocks::compressor::bool::BoolScheme>
+
 pub vortex_btrblocks::BtrBlocksCompressor::float_schemes: alloc::vec::Vec<&'static dyn vortex_btrblocks::compressor::float::FloatScheme>
 
 pub vortex_btrblocks::BtrBlocksCompressor::int_schemes: alloc::vec::Vec<&'static dyn vortex_btrblocks::compressor::integer::IntegerScheme>
@@ -208,6 +258,8 @@ pub fn vortex_btrblocks::BtrBlocksCompressor::default() -> Self
 
 impl vortex_btrblocks::CanonicalCompressor for vortex_btrblocks::BtrBlocksCompressor
 
+pub fn vortex_btrblocks::BtrBlocksCompressor::bool_schemes(&self) -> &[&'static dyn vortex_btrblocks::compressor::bool::BoolScheme]
+
 pub fn vortex_btrblocks::BtrBlocksCompressor::compress_canonical(&self, array: vortex_array::canonical::Canonical, ctx: vortex_btrblocks::ctx::CompressorContext, excludes: vortex_btrblocks::ctx::Excludes<'_>) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
 
 pub fn vortex_btrblocks::BtrBlocksCompressor::float_schemes(&self) -> &[&'static dyn vortex_btrblocks::compressor::float::FloatScheme]
@@ -224,12 +276,16 @@ pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::build(self) -> vortex_btrbl
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::empty() -> Self
 
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::exclude_bool(self, codes: impl core::iter::traits::collect::IntoIterator<Item = vortex_btrblocks::BoolCode>) -> Self
+
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::exclude_float(self, codes: impl core::iter::traits::collect::IntoIterator<Item = vortex_btrblocks::FloatCode>) -> Self
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::exclude_int(self, codes: impl core::iter::traits::collect::IntoIterator<Item = vortex_btrblocks::IntCode>) -> Self
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::exclude_string(self, codes: impl core::iter::traits::collect::IntoIterator<Item = vortex_btrblocks::StringCode>) -> Self
 
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::include_bool(self, codes: impl core::iter::traits::collect::IntoIterator<Item = vortex_btrblocks::BoolCode>) -> Self
+
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::include_float(self, codes: impl core::iter::traits::collect::IntoIterator<Item = vortex_btrblocks::FloatCode>) -> Self
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::include_int(self, codes: impl core::iter::traits::collect::IntoIterator<Item = vortex_btrblocks::IntCode>) -> Self
@@ -282,6 +338,8 @@ pub fn vortex_btrblocks::IntegerStats::source(&self) -> &vortex_array::arrays::p
 
 pub trait vortex_btrblocks::CanonicalCompressor
 
+pub fn vortex_btrblocks::CanonicalCompressor::bool_schemes(&self) -> &[&'static dyn vortex_btrblocks::compressor::bool::BoolScheme]
+
 pub fn vortex_btrblocks::CanonicalCompressor::compress_canonical(&self, array: vortex_array::canonical::Canonical, ctx: vortex_btrblocks::ctx::CompressorContext, excludes: vortex_btrblocks::ctx::Excludes<'_>) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
 
 pub fn vortex_btrblocks::CanonicalCompressor::float_schemes(&self) -> &[&'static dyn vortex_btrblocks::compressor::float::FloatScheme]
@@ -292,6 +350,8 @@ pub fn vortex_btrblocks::CanonicalCompressor::string_schemes(&self) -> &[&'stati
 
 impl vortex_btrblocks::CanonicalCompressor for vortex_btrblocks::BtrBlocksCompressor
 
+pub fn vortex_btrblocks::BtrBlocksCompressor::bool_schemes(&self) -> &[&'static dyn vortex_btrblocks::compressor::bool::BoolScheme]
+
 pub fn vortex_btrblocks::BtrBlocksCompressor::compress_canonical(&self, array: vortex_array::canonical::Canonical, ctx: vortex_btrblocks::ctx::CompressorContext, excludes: vortex_btrblocks::ctx::Excludes<'_>) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
 
 pub fn vortex_btrblocks::BtrBlocksCompressor::float_schemes(&self) -> &[&'static dyn vortex_btrblocks::compressor::float::FloatScheme]

vortex-btrblocks/src/builder.rs

@@ -6,10 +6,13 @@
 use itertools::Itertools;
 use vortex_utils::aliases::hash_set::HashSet;
 
+use crate::BoolCode;
 use crate::BtrBlocksCompressor;
 use crate::FloatCode;
 use crate::IntCode;
 use crate::StringCode;
+use crate::compressor::bool::ALL_BOOL_SCHEMES;
+use crate::compressor::bool::BoolScheme;
 use crate::compressor::float::ALL_FLOAT_SCHEMES;
 use crate::compressor::float::FloatScheme;
 use crate::compressor::integer::ALL_INT_SCHEMES;
@@ -43,6 +46,7 @@ use crate::compressor::string::StringScheme;
 /// ```
 #[derive(Debug, Clone)]
 pub struct BtrBlocksCompressorBuilder {
+    bool_schemes: HashSet<&'static dyn BoolScheme>,
     int_schemes: HashSet<&'static dyn IntegerScheme>,
     float_schemes: HashSet<&'static dyn FloatScheme>,
     string_schemes: HashSet<&'static dyn StringScheme>,
@@ -51,6 +55,7 @@ pub struct BtrBlocksCompressorBuilder {
 impl Default for BtrBlocksCompressorBuilder {
     fn default() -> Self {
         Self {
+            bool_schemes: ALL_BOOL_SCHEMES.iter().copied().collect(),
             int_schemes: ALL_INT_SCHEMES
                 .iter()
                 .copied()
@@ -74,12 +79,20 @@ impl BtrBlocksCompressorBuilder {
     /// Create a new builder with no encodings enabled.
     pub fn empty() -> Self {
         Self {
+            bool_schemes: Default::default(),
             int_schemes: Default::default(),
             float_schemes: Default::default(),
             string_schemes: Default::default(),
         }
     }
 
+    /// Excludes the specified bool compression schemes.
+    pub fn exclude_bool(mut self, codes: impl IntoIterator<Item = BoolCode>) -> Self {
+        let codes: HashSet<_> = codes.into_iter().collect();
+        self.bool_schemes.retain(|s| !codes.contains(&s.code()));
+        self
+    }
+
     /// Excludes the specified integer compression schemes.
     pub fn exclude_int(mut self, codes: impl IntoIterator<Item = IntCode>) -> Self {
         let codes: HashSet<_> = codes.into_iter().collect();
@@ -101,6 +114,17 @@ impl BtrBlocksCompressorBuilder {
         self
     }
 
+    /// Includes the specified bool compression schemes.
+    pub fn include_bool(mut self, codes: impl IntoIterator<Item = BoolCode>) -> Self {
+        let codes: HashSet<_> = codes.into_iter().collect();
+        for scheme in ALL_BOOL_SCHEMES {
+            if codes.contains(&scheme.code()) {
+                self.bool_schemes.insert(*scheme);
+            }
+        }
+        self
+    }
+
     /// Includes the specified integer compression schemes.
     pub fn include_int(mut self, codes: impl IntoIterator<Item = IntCode>) -> Self {
         let codes: HashSet<_> = codes.into_iter().collect();
@@ -138,6 +162,11 @@ impl BtrBlocksCompressorBuilder {
     pub fn build(self) -> BtrBlocksCompressor {
         // Note we should apply the schemes in the same order, in case try conflict.
         BtrBlocksCompressor {
+            bool_schemes: self
+                .bool_schemes
+                .into_iter()
+                .sorted_by_key(|s| s.code())
+                .collect_vec(),
             int_schemes: self
                 .int_schemes
                 .into_iter()

vortex-btrblocks/src/canonical_compressor.rs

@@ -27,6 +27,7 @@ use vortex_array::vtable::ValidityHelper;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
+use crate::BoolCompressor;
 use crate::BtrBlocksCompressorBuilder;
 use crate::CompressorContext;
 use crate::CompressorExt;
@@ -35,6 +36,7 @@ use crate::FloatCompressor;
 use crate::IntCode;
 use crate::IntCompressor;
 use crate::StringCompressor;
+use crate::compressor::bool::BoolScheme;
 use crate::compressor::decimal::compress_decimal;
 use crate::compressor::float::FloatScheme;
 use crate::compressor::integer::IntegerScheme;
@@ -54,6 +56,9 @@ pub trait CanonicalCompressor {
         excludes: Excludes,
     ) -> VortexResult<ArrayRef>;
 
+    /// Returns the enabled bool compression schemes.
+    fn bool_schemes(&self) -> &[&'static dyn BoolScheme];
+
     /// Returns the enabled integer compression schemes.
     fn int_schemes(&self) -> &[&'static dyn IntegerScheme];
 
@@ -93,6 +98,9 @@ pub trait CanonicalCompressor {
 /// ```
 #[derive(Clone)]
 pub struct BtrBlocksCompressor {
+    /// Bool compressor with configured schemes.
+    pub bool_schemes: Vec<&'static dyn BoolScheme>,
+
     /// Integer compressor with configured schemes.
     pub int_schemes: Vec<&'static dyn IntegerScheme>,
 
@@ -127,6 +135,12 @@ impl BtrBlocksCompressor {
         self.compress_canonical(compact, CompressorContext::default(), Excludes::none())
     }
 
+    pub(crate) fn bool_compressor(&self) -> BoolCompressor<'_> {
+        BoolCompressor {
+            btr_blocks_compressor: self,
+        }
+    }
+
     pub(crate) fn integer_compressor(&self) -> IntCompressor<'_> {
         IntCompressor {
             btr_blocks_compressor: self,
@@ -215,7 +229,10 @@ impl CanonicalCompressor for BtrBlocksCompressor {
     ) -> VortexResult<ArrayRef> {
         match array {
             Canonical::Null(null_array) => Ok(null_array.into_array()),
-            Canonical::Bool(bool_array) => Ok(bool_array.into_array()),
+            Canonical::Bool(bool_array) => {
+                self.bool_compressor()
+                    .compress(self, &bool_array, ctx, excludes.bool)
+            }
             Canonical::Primitive(primitive) => {
                 if primitive.ptype().is_int() {
                     self.integer_compressor()
@@ -304,6 +321,10 @@ impl CanonicalCompressor for BtrBlocksCompressor {
         }
     }
 
+    fn bool_schemes(&self) -> &[&'static dyn BoolScheme] {
+        &self.bool_schemes
+    }
+
     fn int_schemes(&self) -> &[&'static dyn IntegerScheme] {
         &self.int_schemes
     }