Use BitBuffer instead of BooleanBuffer (#4940)

Replace usages of BooleanBuffer with our own BitBuffer. This lets us
have our own backing Buffer for BitBuffer and implement additional
optimisations for use in vortex.

previously #2456 

This pr requires resolving outstanding performance issues before merging
1. Make iterators cheaper to construct, likely can't be owned
2. Improve append_packed_range performance. Arrow has a function that is
faster than bitvec crate

---------

Signed-off-by: Robert Kruszewski <[email protected]>
Signed-off-by: Nicholas Gates <[email protected]>
Co-authored-by: Andrew Duffy <[email protected]>
Co-authored-by: Nicholas Gates <[email protected]>

Author: 3 people

Cargo.lock

@@ -104,7 +104,7 @@ mod tests {
         let res = between_impl(arr, lower, upper, Nullability::Nullable, options)
             .unwrap()
             .to_bool()
-            .boolean_buffer()
+            .bit_buffer()
             .iter()
             .collect_vec();
         assert_eq!(res.len(), 1);

encodings/alp/src/alp/compute/between.rs

@@ -148,7 +148,7 @@ mod tests {
     {
         alp_scalar_compare(alp, value, operator)
             .unwrap()
-            .map(|a| a.to_bool().boolean_buffer().iter().collect())
+            .map(|a| a.to_bool().bit_buffer().iter().collect())
     }
 
     #[test]
@@ -166,7 +166,7 @@ mod tests {
             .unwrap()
             .to_bool();
 
-        for v in r.boolean_buffer().iter() {
+        for v in r.bit_buffer().iter() {
             assert!(!v);
         }
 
@@ -175,7 +175,7 @@ mod tests {
             .unwrap()
             .to_bool();
 
-        for v in r.boolean_buffer().iter() {
+        for v in r.bit_buffer().iter() {
             assert!(v);
         }
     }
@@ -196,15 +196,15 @@ mod tests {
             .unwrap()
             .to_bool();
 
-        assert!(r_eq.boolean_buffer().iter().all(|v| !v));
+        assert!(r_eq.bit_buffer().iter().all(|v| !v));
 
         #[allow(clippy::excessive_precision)]
         let r_neq = alp_scalar_compare(&encoded, 1.234444f32, Operator::NotEq)
             .unwrap()
             .unwrap()
             .to_bool();
 
-        assert!(r_neq.boolean_buffer().iter().all(|v| v));
+        assert!(r_neq.bit_buffer().iter().all(|v| v));
     }
 
     #[test]
@@ -223,31 +223,31 @@ mod tests {
             .to_bool();
 
         // !(0.0605_f32 >= 0.06051_f32);
-        assert!(r_gte.boolean_buffer().iter().all(|v| !v));
+        assert!(r_gte.bit_buffer().iter().all(|v| !v));
 
         let r_gt = alp_scalar_compare(&encoded, 0.06051_f32, Operator::Gt)
             .unwrap()
             .unwrap()
             .to_bool();
 
         // (0.0605_f32 > 0.06051_f32);
-        assert!(r_gt.boolean_buffer().iter().all(|v| !v));
+        assert!(r_gt.bit_buffer().iter().all(|v| !v));
 
         let r_lte = alp_scalar_compare(&encoded, 0.06051_f32, Operator::Lte)
             .unwrap()
             .unwrap()
             .to_bool();
 
         // 0.0605_f32 <= 0.06051_f32;
-        assert!(r_lte.boolean_buffer().iter().all(|v| v));
+        assert!(r_lte.bit_buffer().iter().all(|v| v));
 
         let r_lt = alp_scalar_compare(&encoded, 0.06051_f32, Operator::Lt)
             .unwrap()
             .unwrap()
             .to_bool();
 
         //0.0605_f32 < 0.06051_f32;
-        assert!(r_lt.boolean_buffer().iter().all(|v| v));
+        assert!(r_lt.bit_buffer().iter().all(|v| v));
     }
 
     #[test]
@@ -311,7 +311,7 @@ mod tests {
             .unwrap()
             .to_bool();
 
-        for v in r.boolean_buffer().iter() {
+        for v in r.bit_buffer().iter() {
             assert!(!v);
         }
     }

encodings/alp/src/alp/compute/compare.rs

@@ -34,9 +34,9 @@ register_kernel!(FilterKernelAdapter(ALPVTable).lift());
 #[cfg(test)]
 mod test {
     use rstest::rstest;
-    use vortex_array::IntoArray;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::compute::conformance::filter::test_filter_conformance;
+    use vortex_array::{ArrayRef, IntoArray};
     use vortex_buffer::buffer;
 
     use crate::ALPEncoding;
@@ -50,7 +50,7 @@ mod test {
         1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0,
         11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0
     ].into_array())]
-    fn test_filter_alp_conformance(#[case] array: vortex_array::ArrayRef) {
+    fn test_filter_alp_conformance(#[case] array: ArrayRef) {
         let alp = ALPEncoding
             .encode(&array.to_canonical(), None)
             .unwrap()

encodings/alp/src/alp/compute/filter.rs

@@ -33,7 +33,6 @@ mod tests {
     // Arrays with patterns
     #[case::repeating_pattern(alp_encode(&PrimitiveArray::from_iter([1.1f32, 2.2, 3.3, 1.1, 2.2, 3.3, 1.1, 2.2, 3.3]), None).unwrap())]
     #[case::close_values(alp_encode(&PrimitiveArray::from_iter([100.001f64, 100.002, 100.003, 100.004, 100.005]), None).unwrap())]
-
     fn test_alp_consistency(#[case] array: ALPArray) {
         test_array_consistency(array.as_ref());
     }

encodings/alp/src/alp/compute/mod.rs

@@ -64,7 +64,6 @@ mod tests {
         let encoder = RDEncoder::new(&values);
         encoder.encode(&arr)
     })]
-
     fn test_alp_rd_consistency(#[case] array: ALPRDArray) {
         test_array_consistency(array.as_ref());
     }

encodings/alp/src/alp_rd/compute/mod.rs

@@ -17,7 +17,6 @@ version = { workspace = true }
 workspace = true
 
 [dependencies]
-arrow-buffer = { workspace = true }
 num-traits = { workspace = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }

encodings/bytebool/Cargo.toml

@@ -4,7 +4,6 @@
 use std::fmt::Debug;
 use std::ops::Range;
 
-use arrow_buffer::BooleanBuffer;
 use vortex_array::arrays::BoolArray;
 use vortex_array::stats::{ArrayStats, StatsSetRef};
 use vortex_array::validity::Validity;
@@ -13,7 +12,7 @@ use vortex_array::vtable::{
     ValidityVTableFromValidityHelper,
 };
 use vortex_array::{ArrayRef, Canonical, EncodingId, EncodingRef, IntoArray, vtable};
-use vortex_buffer::ByteBuffer;
+use vortex_buffer::{BitBuffer, ByteBuffer};
 use vortex_dtype::DType;
 use vortex_error::vortex_panic;
 use vortex_scalar::Scalar;
@@ -114,9 +113,9 @@ impl ArrayVTable<ByteBoolVTable> for ByteBoolVTable {
 
 impl CanonicalVTable<ByteBoolVTable> for ByteBoolVTable {
     fn canonicalize(array: &ByteBoolArray) -> Canonical {
-        let boolean_buffer = BooleanBuffer::from(array.as_slice());
+        let boolean_buffer = BitBuffer::from(array.as_slice());
         let validity = array.validity().clone();
-        Canonical::Bool(BoolArray::from_bool_buffer(boolean_buffer, validity))
+        Canonical::Bool(BoolArray::from_bit_buffer(boolean_buffer, validity))
     }
 }
 

encodings/bytebool/src/array.rs

@@ -149,7 +149,7 @@ mod tests {
         let res = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
 
         assert_eq!(
-            res.to_bool().boolean_buffer().iter().collect::<Vec<_>>(),
+            res.to_bool().bit_buffer().iter().collect::<Vec<_>>(),
             vec![false, false, true]
         );
     }
@@ -206,13 +206,13 @@ mod tests {
 
         let res = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
 
-        assert_eq!(res.to_bool().bool_vec().unwrap(), vec![false, false, false]);
+        assert_eq!(res.to_bool().bool_vec(), vec![false, false, false]);
 
         let res = compare(lhs.as_ref(), rhs.as_ref(), Operator::Gt).unwrap();
-        assert_eq!(res.to_bool().bool_vec().unwrap(), vec![true, true, true]);
+        assert_eq!(res.to_bool().bool_vec(), vec![true, true, true]);
 
         let res = compare(lhs.as_ref(), rhs.as_ref(), Operator::Lt).unwrap();
-        assert_eq!(res.to_bool().bool_vec().unwrap(), vec![false, false, false]);
+        assert_eq!(res.to_bool().bool_vec(), vec![false, false, false]);
 
         // This cannot be converted to a i32.
         let rhs = ConstantArray::new(
@@ -222,12 +222,12 @@ mod tests {
 
         let res = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
 
-        assert_eq!(res.to_bool().bool_vec().unwrap(), vec![false, false, false]);
+        assert_eq!(res.to_bool().bool_vec(), vec![false, false, false]);
 
         let res = compare(lhs.as_ref(), rhs.as_ref(), Operator::Gt).unwrap();
-        assert_eq!(res.to_bool().bool_vec().unwrap(), vec![false, false, false]);
+        assert_eq!(res.to_bool().bool_vec(), vec![false, false, false]);
 
         let res = compare(lhs.as_ref(), rhs.as_ref(), Operator::Lt).unwrap();
-        assert_eq!(res.to_bool().bool_vec().unwrap(), vec![true, true, true]);
+        assert_eq!(res.to_bool().bool_vec(), vec![true, true, true]);
     }
 }

encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs

@@ -3,10 +3,10 @@
 
 use std::fmt::Debug;
 
-use arrow_buffer::BooleanBuffer;
 use vortex_array::stats::{ArrayStats, StatsSetRef};
 use vortex_array::vtable::{ArrayVTable, NotSupported, VTable, ValidityVTable};
 use vortex_array::{Array, ArrayRef, EncodingId, EncodingRef, ToCanonical, vtable};
+use vortex_buffer::BitBuffer;
 use vortex_dtype::{DType, match_each_integer_ptype};
 use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
 use vortex_mask::{AllOr, Mask};
@@ -142,13 +142,13 @@ impl ValidityVTable<DictVTable> for DictVTable {
 
     fn validity_mask(array: &DictArray) -> Mask {
         let codes_validity = array.codes().validity_mask();
-        match codes_validity.boolean_buffer() {
+        match codes_validity.bit_buffer() {
             AllOr::All => {
                 let primitive_codes = array.codes().to_primitive();
                 let values_mask = array.values().validity_mask();
                 let is_valid_buffer = match_each_integer_ptype!(primitive_codes.ptype(), |P| {
                     let codes_slice = primitive_codes.as_slice::<P>();
-                    BooleanBuffer::collect_bool(array.len(), |idx| {
+                    BitBuffer::collect_bool(array.len(), |idx| {
                         #[allow(clippy::cast_possible_truncation)]
                         values_mask.value(codes_slice[idx] as usize)
                     })
@@ -162,7 +162,7 @@ impl ValidityVTable<DictVTable> for DictVTable {
                 let is_valid_buffer = match_each_integer_ptype!(primitive_codes.ptype(), |P| {
                     let codes_slice = primitive_codes.as_slice::<P>();
                     #[allow(clippy::cast_possible_truncation)]
-                    BooleanBuffer::collect_bool(array.len(), |idx| {
+                    BitBuffer::collect_bool(array.len(), |idx| {
                         validity_buff.value(idx) && values_mask.value(codes_slice[idx] as usize)
                     })
                 });
@@ -174,15 +174,14 @@ impl ValidityVTable<DictVTable> for DictVTable {
 
 #[cfg(test)]
 mod test {
-    use arrow_buffer::BooleanBuffer;
     use rand::distr::{Distribution, StandardUniform};
     use rand::prelude::StdRng;
     use rand::{Rng, SeedableRng};
     use vortex_array::arrays::{ChunkedArray, PrimitiveArray};
     use vortex_array::builders::builder_with_capacity;
     use vortex_array::validity::Validity;
     use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
-    use vortex_buffer::buffer;
+    use vortex_buffer::{BitBuffer, buffer};
     use vortex_dtype::Nullability::NonNullable;
     use vortex_dtype::{DType, NativePType, PType, UnsignedPType};
     use vortex_error::{VortexExpect, VortexUnwrap, vortex_panic};
@@ -195,7 +194,7 @@ mod test {
         let dict = DictArray::try_new(
             PrimitiveArray::new(
                 buffer![0u32, 1, 2, 2, 1],
-                Validity::from(BooleanBuffer::from(vec![true, false, true, false, true])),
+                Validity::from(BitBuffer::from(vec![true, false, true, false, true])),
             )
             .into_array(),
             PrimitiveArray::new(buffer![3, 6, 9], Validity::AllValid).into_array(),
@@ -214,7 +213,7 @@ mod test {
             buffer![0u32, 1, 2, 2, 1].into_array(),
             PrimitiveArray::new(
                 buffer![3, 6, 9],
-                Validity::from(BooleanBuffer::from(vec![true, false, false])),
+                Validity::from(BitBuffer::from(vec![true, false, false])),
             )
             .into_array(),
         )
@@ -231,12 +230,12 @@ mod test {
         let dict = DictArray::try_new(
             PrimitiveArray::new(
                 buffer![0u32, 1, 2, 2, 1],
-                Validity::from(BooleanBuffer::from(vec![true, false, true, false, true])),
+                Validity::from(BitBuffer::from(vec![true, false, true, false, true])),
             )
             .into_array(),
             PrimitiveArray::new(
                 buffer![3, 6, 9],
-                Validity::from(BooleanBuffer::from(vec![false, true, true])),
+                Validity::from(BitBuffer::from(vec![false, true, true])),
             )
             .into_array(),
         )
@@ -253,7 +252,7 @@ mod test {
         let dict = DictArray::try_new(
             PrimitiveArray::new(
                 buffer![0u32, 1, 2, 2, 1],
-                Validity::from(BooleanBuffer::from(vec![true, false, true, false, true])),
+                Validity::from(BitBuffer::from(vec![true, false, true, false, true])),
             )
             .into_array(),
             PrimitiveArray::new(buffer![3, 6, 9], Validity::NonNullable).into_array(),
@@ -312,8 +311,8 @@ mod test {
 
         assert_eq!(into_prim.as_slice::<u64>(), prim_into.as_slice::<u64>());
         assert_eq!(
-            into_prim.validity_mask().boolean_buffer(),
-            prim_into.validity_mask().boolean_buffer()
+            into_prim.validity_mask().bit_buffer(),
+            prim_into.validity_mask().bit_buffer()
         )
     }
 }