vortex_roaring/integer/
compute.rs

1use croaring::Bitmap;
2use vortex_array::compute::{ComputeVTable, ScalarAtFn, SliceFn};
3use vortex_array::{ArrayData, ArrayLen, IntoArrayData};
4use vortex_dtype::PType;
5use vortex_error::{vortex_err, VortexResult};
6use vortex_scalar::Scalar;
7
8use crate::{RoaringIntArray, RoaringIntEncoding};
9
10impl ComputeVTable for RoaringIntEncoding {
11    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
12        Some(self)
13    }
14
15    fn slice_fn(&self) -> Option<&dyn SliceFn<ArrayData>> {
16        Some(self)
17    }
18}
19
20impl ScalarAtFn<RoaringIntArray> for RoaringIntEncoding {
21    fn scalar_at(&self, array: &RoaringIntArray, index: usize) -> VortexResult<Scalar> {
22        let bitmap_value = array
23            .owned_bitmap()
24            .select(index as u32)
25            .ok_or_else(|| vortex_err!(OutOfBounds: index, 0, array.len()))?;
26        let scalar: Scalar = match array.metadata().ptype {
27            PType::U8 => (bitmap_value as u8).into(),
28            PType::U16 => (bitmap_value as u16).into(),
29            PType::U32 => bitmap_value.into(),
30            PType::U64 => (bitmap_value as u64).into(),
31            _ => unreachable!("RoaringIntArray constructor should have disallowed this type"),
32        };
33        Ok(scalar)
34    }
35}
36
37impl SliceFn<RoaringIntArray> for RoaringIntEncoding {
38    fn slice(&self, array: &RoaringIntArray, start: usize, stop: usize) -> VortexResult<ArrayData> {
39        let mut bitmap = array.owned_bitmap();
40        let start = bitmap
41            .select(start as u32)
42            .ok_or_else(|| vortex_err!(OutOfBounds: start, 0, array.len()))?;
43        let stop_inclusive = if stop == array.len() {
44            bitmap.maximum().unwrap_or(0)
45        } else {
46            bitmap
47                .select(stop.saturating_sub(1) as u32)
48                .ok_or_else(|| vortex_err!(OutOfBounds: stop, 0, array.len()))?
49        };
50
51        bitmap.and_inplace(&Bitmap::from_range(start..=stop_inclusive));
52        RoaringIntArray::try_new(bitmap, array.cached_ptype()).map(IntoArrayData::into_array)
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use vortex_array::array::PrimitiveArray;
59    use vortex_array::compute::{scalar_at, slice};
60
61    use super::*;
62
63    #[test]
64    #[cfg_attr(miri, ignore)]
65    pub fn test_scalar_at() {
66        let ints = PrimitiveArray::from(vec![2u32, 12, 22, 32]).into_array();
67        let array = RoaringIntArray::encode(ints).unwrap();
68
69        assert_eq!(scalar_at(&array, 0).unwrap(), 2u32.into());
70        assert_eq!(scalar_at(&array, 1).unwrap(), 12u32.into());
71    }
72
73    #[test]
74    #[cfg_attr(miri, ignore)]
75    fn test_slice() {
76        let array = RoaringIntArray::try_new(Bitmap::from_range(10..20), PType::U32).unwrap();
77
78        let sliced = slice(&array, 0, 5).unwrap();
79        assert_eq!(sliced.len(), 5);
80        assert_eq!(scalar_at(&sliced, 0).unwrap(), 10u32.into());
81        assert_eq!(scalar_at(&sliced, 4).unwrap(), 14u32.into());
82
83        let sliced = slice(&array, 5, 10).unwrap();
84        assert_eq!(sliced.len(), 5);
85        assert_eq!(scalar_at(&sliced, 0).unwrap(), 15u32.into());
86        assert_eq!(scalar_at(&sliced, 4).unwrap(), 19u32.into());
87
88        let sliced = slice(&sliced, 3, 5).unwrap();
89        assert_eq!(sliced.len(), 2);
90        assert_eq!(scalar_at(&sliced, 0).unwrap(), 18u32.into());
91        assert_eq!(scalar_at(&sliced, 1).unwrap(), 19u32.into());
92    }
93}
vortex_roaring/integer/compute.rs

vortex_roaring/integer/
compute.rs
