vortex_roaring/integer/
compute.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
use croaring::Bitmap;
use vortex_array::compute::unary::ScalarAtFn;
use vortex_array::compute::{ArrayCompute, SliceFn};
use vortex_array::{Array, IntoArray};
use vortex_dtype::PType;
use vortex_error::{vortex_err, VortexResult, VortexUnwrap as _};
use vortex_scalar::Scalar;

use crate::RoaringIntArray;

impl ArrayCompute for RoaringIntArray {
    fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
        Some(self)
    }

    fn slice(&self) -> Option<&dyn SliceFn> {
        Some(self)
    }
}

impl ScalarAtFn for RoaringIntArray {
    fn scalar_at(&self, index: usize) -> VortexResult<Scalar> {
        let bitmap_value = self
            .owned_bitmap()
            .select(index as u32)
            .ok_or_else(|| vortex_err!(OutOfBounds: index, 0, self.len()))?;
        let scalar: Scalar = match self.metadata().ptype {
            PType::U8 => (bitmap_value as u8).into(),
            PType::U16 => (bitmap_value as u16).into(),
            PType::U32 => bitmap_value.into(),
            PType::U64 => (bitmap_value as u64).into(),
            _ => unreachable!("RoaringIntArray constructor should have disallowed this type"),
        };
        Ok(scalar)
    }

    fn scalar_at_unchecked(&self, index: usize) -> Scalar {
        <Self as ScalarAtFn>::scalar_at(self, index).vortex_unwrap()
    }
}

impl SliceFn for RoaringIntArray {
    fn slice(&self, start: usize, stop: usize) -> VortexResult<Array> {
        let mut bitmap = self.owned_bitmap();
        let start = bitmap
            .select(start as u32)
            .ok_or_else(|| vortex_err!(OutOfBounds: start, 0, self.len()))?;
        let stop_inclusive = if stop == self.len() {
            bitmap.maximum().unwrap_or(0)
        } else {
            bitmap
                .select(stop.saturating_sub(1) as u32)
                .ok_or_else(|| vortex_err!(OutOfBounds: stop, 0, self.len()))?
        };

        bitmap.and_inplace(&Bitmap::from_range(start..=stop_inclusive));
        Self::try_new(bitmap, self.ptype()).map(IntoArray::into_array)
    }
}

#[cfg(test)]
mod tests {
    use vortex_array::array::PrimitiveArray;
    use vortex_array::compute::slice;
    use vortex_array::compute::unary::scalar_at;

    use super::*;

    #[test]
    #[cfg_attr(miri, ignore)]
    pub fn test_scalar_at() {
        let ints = PrimitiveArray::from(vec![2u32, 12, 22, 32]).into_array();
        let array = RoaringIntArray::encode(ints).unwrap();

        assert_eq!(scalar_at(&array, 0).unwrap(), 2u32.into());
        assert_eq!(scalar_at(&array, 1).unwrap(), 12u32.into());
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn test_slice() {
        let array = RoaringIntArray::try_new(Bitmap::from_range(10..20), PType::U32).unwrap();

        let sliced = slice(&array, 0, 5).unwrap();
        assert_eq!(sliced.len(), 5);
        assert_eq!(scalar_at(&sliced, 0).unwrap(), 10u32.into());
        assert_eq!(scalar_at(&sliced, 4).unwrap(), 14u32.into());

        let sliced = slice(&array, 5, 10).unwrap();
        assert_eq!(sliced.len(), 5);
        assert_eq!(scalar_at(&sliced, 0).unwrap(), 15u32.into());
        assert_eq!(scalar_at(&sliced, 4).unwrap(), 19u32.into());

        let sliced = slice(&sliced, 3, 5).unwrap();
        assert_eq!(sliced.len(), 2);
        assert_eq!(scalar_at(&sliced, 0).unwrap(), 18u32.into());
        assert_eq!(scalar_at(&sliced, 1).unwrap(), 19u32.into());
    }
}