vortex_roaring/boolean/
compute.rs

1use croaring::Bitmap;
2use vortex_array::compute::{ComputeVTable, InvertFn, ScalarAtFn, SliceFn};
3use vortex_array::{ArrayData, ArrayLen, IntoArrayData};
4use vortex_error::VortexResult;
5use vortex_scalar::Scalar;
6
7use crate::{RoaringBoolArray, RoaringBoolEncoding};
8
9impl ComputeVTable for RoaringBoolEncoding {
10    fn invert_fn(&self) -> Option<&dyn InvertFn<ArrayData>> {
11        Some(self)
12    }
13
14    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
15        Some(self)
16    }
17
18    fn slice_fn(&self) -> Option<&dyn SliceFn<ArrayData>> {
19        Some(self)
20    }
21}
22
23impl InvertFn<RoaringBoolArray> for RoaringBoolEncoding {
24    fn invert(&self, array: &RoaringBoolArray) -> VortexResult<ArrayData> {
25        RoaringBoolArray::try_new(array.bitmap().flip(0..(array.len() as u32)), array.len())
26            .map(|a| a.into_array())
27    }
28}
29
30impl ScalarAtFn<RoaringBoolArray> for RoaringBoolEncoding {
31    fn scalar_at(&self, array: &RoaringBoolArray, index: usize) -> VortexResult<Scalar> {
32        Ok(array.bitmap().contains(index as u32).into())
33    }
34}
35
36impl SliceFn<RoaringBoolArray> for RoaringBoolEncoding {
37    fn slice(
38        &self,
39        array: &RoaringBoolArray,
40        start: usize,
41        stop: usize,
42    ) -> VortexResult<ArrayData> {
43        let slice_bitmap = Bitmap::from_range(start as u32..stop as u32);
44        let bitmap = array
45            .bitmap()
46            .and(&slice_bitmap)
47            .add_offset(-(start as i64));
48
49        RoaringBoolArray::try_new(bitmap, stop - start).map(IntoArrayData::into_array)
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use vortex_array::array::BoolArray;
56    use vortex_array::compute::{scalar_at, slice};
57    use vortex_array::{IntoArrayData, IntoArrayVariant};
58    use vortex_scalar::Scalar;
59
60    use crate::RoaringBoolArray;
61
62    #[test]
63    #[cfg_attr(miri, ignore)]
64    pub fn test_scalar_at() {
65        let bool = BoolArray::from_iter([true, false, true, true]);
66        let array = RoaringBoolArray::encode(bool.into_array()).unwrap();
67
68        let truthy: Scalar = true.into();
69        let falsy: Scalar = false.into();
70
71        assert_eq!(scalar_at(&array, 0).unwrap(), truthy);
72        assert_eq!(scalar_at(&array, 1).unwrap(), falsy);
73        assert_eq!(scalar_at(&array, 2).unwrap(), truthy);
74        assert_eq!(scalar_at(&array, 3).unwrap(), truthy);
75    }
76
77    #[test]
78    #[cfg_attr(miri, ignore)]
79    pub fn test_slice() {
80        let bool = BoolArray::from_iter([true, false, true, true]);
81        let array = RoaringBoolArray::encode(bool.into_array()).unwrap();
82        let sliced = slice(&array, 1, 3).unwrap();
83
84        assert_eq!(
85            sliced
86                .into_bool()
87                .unwrap()
88                .boolean_buffer()
89                .iter()
90                .collect::<Vec<_>>(),
91            &[false, true]
92        );
93    }
94}
vortex_roaring/boolean/compute.rs

vortex_roaring/boolean/
compute.rs
