vortex_roaring/integer/
compress.rs

1use croaring::Bitmap;
2use num_traits::NumCast;
3use vortex_array::array::PrimitiveArray;
4use vortex_array::variants::PrimitiveArrayTrait;
5use vortex_dtype::{NativePType, PType};
6use vortex_error::{vortex_bail, vortex_err, VortexResult};
7
8use crate::RoaringIntArray;
9
10pub fn roaring_int_encode(parray: PrimitiveArray) -> VortexResult<RoaringIntArray> {
11    match parray.ptype() {
12        PType::U8 => roaring_encode_primitive::<u8>(parray.maybe_null_slice()),
13        PType::U16 => roaring_encode_primitive::<u16>(parray.maybe_null_slice()),
14        PType::U32 => roaring_encode_primitive::<u32>(parray.maybe_null_slice()),
15        PType::U64 => roaring_encode_primitive::<u64>(parray.maybe_null_slice()),
16        _ => vortex_bail!("Unsupported PType {}", parray.ptype()),
17    }
18}
19
20fn roaring_encode_primitive<T: NumCast + NativePType>(
21    values: &[T],
22) -> VortexResult<RoaringIntArray> {
23    let mut bitmap = Bitmap::new();
24    bitmap.extend(
25        values
26            .iter()
27            .map(|i| {
28                i.to_u32()
29                    .ok_or_else(|| vortex_err!("Failed to cast value {} to u32", i))
30            })
31            .collect::<VortexResult<Vec<u32>>>()?,
32    );
33    bitmap.run_optimize();
34    bitmap.shrink_to_fit();
35    RoaringIntArray::try_new(bitmap, T::PTYPE)
36}
vortex_roaring/integer/compress.rs

vortex_roaring/integer/
compress.rs
