diff --git a/encodings/roaring/src/boolean/compute.rs b/encodings/roaring/src/boolean/compute.rs index a9d3e92ecf..fe21449acd 100644 --- a/encodings/roaring/src/boolean/compute.rs +++ b/encodings/roaring/src/boolean/compute.rs @@ -35,3 +35,47 @@ impl SliceFn for RoaringBoolArray { Self::try_new(bitmap, stop - start).map(IntoArray::into_array) } } + +#[cfg(test)] +mod tests { + use vortex::array::BoolArray; + use vortex::compute::slice; + use vortex::compute::unary::scalar_at; + use vortex::{IntoArray, IntoArrayVariant}; + use vortex_scalar::Scalar; + + use crate::RoaringBoolArray; + + #[test] + #[cfg_attr(miri, ignore)] + pub fn test_scalar_at() { + let bool = BoolArray::from(vec![true, false, true, true]); + let array = RoaringBoolArray::encode(bool.into_array()).unwrap(); + + let truthy: Scalar = true.into(); + let falsy: Scalar = false.into(); + + assert_eq!(scalar_at(&array, 0).unwrap(), truthy); + assert_eq!(scalar_at(&array, 1).unwrap(), falsy); + assert_eq!(scalar_at(&array, 2).unwrap(), truthy); + assert_eq!(scalar_at(&array, 3).unwrap(), truthy); + } + + #[test] + #[cfg_attr(miri, ignore)] + pub fn test_slice() { + let bool = BoolArray::from(vec![true, false, true, true]); + let array = RoaringBoolArray::encode(bool.into_array()).unwrap(); + let sliced = slice(&array, 1, 3).unwrap(); + + assert_eq!( + sliced + .into_bool() + .unwrap() + .boolean_buffer() + .iter() + .collect::>(), + &[false, true] + ); + } +} diff --git a/encodings/roaring/src/boolean/mod.rs b/encodings/roaring/src/boolean/mod.rs index 0424877ca0..72ce6ca5d8 100644 --- a/encodings/roaring/src/boolean/mod.rs +++ b/encodings/roaring/src/boolean/mod.rs @@ -10,12 +10,11 @@ use vortex::validity::{ArrayValidity, LogicalValidity, Validity}; use vortex::variants::{ArrayVariants, BoolArrayTrait}; use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor}; use vortex::{ - impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoArray, IntoCanonical, - TypedArray, + impl_encoding, Array, ArrayDef, ArrayTrait, Canonical, IntoArray, IntoCanonical, TypedArray, }; use vortex_buffer::Buffer; use vortex_dtype::DType; -use vortex_dtype::Nullability::{NonNullable, Nullable}; +use vortex_dtype::Nullability::NonNullable; use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult}; mod compress; @@ -48,12 +47,7 @@ impl RoaringBoolArray { pub fn bitmap(&self) -> Bitmap { //TODO(@jdcasale): figure out a way to avoid this deserialization per-call - Bitmap::deserialize::( - self.array() - .buffer() - .vortex_expect("RoaringBoolArray buffer is missing") - .as_ref(), - ) + Bitmap::deserialize::(self.buffer().as_ref()) } pub fn encode(array: Array) -> VortexResult { @@ -63,6 +57,12 @@ impl RoaringBoolArray { vortex_bail!("RoaringBool can only encode boolean arrays") } } + + pub fn buffer(&self) -> &Buffer { + self.array() + .buffer() + .vortex_expect("Missing buffer in PrimitiveArray") + } } impl ArrayTrait for RoaringBoolArray {} @@ -89,25 +89,20 @@ impl AcceptArrayVisitor for RoaringBoolArray { // Or serialize into metadata? The only reason we support buffers is so we can write to // the wire without copying into FlatBuffers. But if we need to allocate to serialize // the bitmap anyway, then may as well shove it into metadata. - visitor.visit_buffer( - self.array() - .buffer() - .ok_or(vortex_err!("roaring bool should have a buffer"))?, - )?; - Ok(()) + visitor.visit_buffer(self.buffer()) } } impl ArrayStatisticsCompute for RoaringBoolArray {} impl ArrayValidity for RoaringBoolArray { - fn logical_validity(&self) -> LogicalValidity { - LogicalValidity::AllValid(self.len()) - } - fn is_valid(&self, _index: usize) -> bool { true } + + fn logical_validity(&self) -> LogicalValidity { + LogicalValidity::AllValid(self.len()) + } } impl IntoCanonical for RoaringBoolArray { @@ -119,55 +114,29 @@ impl IntoCanonical for RoaringBoolArray { .to_bitset() .ok_or_else(|| vortex_err!("Failed to convert RoaringBitmap to Bitset"))?; - let bytes = &bitset.as_slice()[0..bitset.size_in_bytes()]; - let buffer = ArrowBuffer::from_slice_ref(bytes); - Ok(Canonical::Bool(BoolArray::try_new( - BooleanBuffer::new(buffer, 0, bitset.size_in_bits()), - match self.dtype().nullability() { - NonNullable => Validity::NonNullable, - Nullable => Validity::AllValid, - }, - )?)) + let buffer = ArrowBuffer::from_slice_ref(bitset.as_slice()); + BoolArray::try_new( + BooleanBuffer::new(buffer, 0, self.len()), + Validity::NonNullable, + ) + .map(Canonical::Bool) } } #[cfg(test)] -#[allow(clippy::panic_in_result_fn)] mod test { use vortex::array::BoolArray; - use vortex::compute::unary::scalar_at; use vortex::IntoArray; - use vortex_error::VortexResult; - use vortex_scalar::Scalar; use crate::RoaringBoolArray; #[test] #[cfg_attr(miri, ignore)] - pub fn iter() -> VortexResult<()> { + pub fn iter() { let bool: BoolArray = BoolArray::from(vec![true, false, true, true]); - let array = RoaringBoolArray::encode(bool.into_array())?; - let round_trip = RoaringBoolArray::try_from(array)?; + let array = RoaringBoolArray::encode(bool.into_array()).unwrap(); + let round_trip = RoaringBoolArray::try_from(array).unwrap(); let values = round_trip.bitmap().to_vec(); assert_eq!(values, vec![0, 2, 3]); - - Ok(()) - } - - #[test] - #[cfg_attr(miri, ignore)] - pub fn test_scalar_at() -> VortexResult<()> { - let bool: BoolArray = BoolArray::from(vec![true, false, true, true]); - let array = RoaringBoolArray::encode(bool.into_array())?; - - let truthy: Scalar = true.into(); - let falsy: Scalar = false.into(); - - assert_eq!(scalar_at(&array, 0)?, truthy); - assert_eq!(scalar_at(&array, 1)?, falsy); - assert_eq!(scalar_at(&array, 2)?, truthy); - assert_eq!(scalar_at(&array, 3)?, truthy); - - Ok(()) } }