vortex_fastlanes/bitpacking/compute/
slice.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::cmp::max;
5use std::ops::Range;
6
7use vortex_array::ArrayRef;
8use vortex_array::ArrayView;
9use vortex_array::ExecutionCtx;
10use vortex_array::IntoArray;
11use vortex_array::arrays::slice::SliceKernel;
12use vortex_array::arrays::slice::SliceReduce;
13use vortex_array::patches::Patches;
14use vortex_error::VortexResult;
15
16use crate::BitPacked;
17use crate::bitpacking::array::BitPackedArrayExt;
18
19impl SliceReduce for BitPacked {
20    fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
21        // We cannot access buffers (to slice the patches).
22        if array.patches().is_some() {
23            return Ok(None);
24        }
25
26        Ok(Some(slice_bitpacked(array, range, None)?))
27    }
28}
29
30impl SliceKernel for BitPacked {
31    fn slice(
32        array: ArrayView<'_, Self>,
33        range: Range<usize>,
34        _ctx: &mut ExecutionCtx,
35    ) -> VortexResult<Option<ArrayRef>> {
36        let patches = array
37            .patches()
38            .map(|p| p.slice(range.clone()))
39            .transpose()?
40            .flatten();
41
42        Ok(Some(slice_bitpacked(array, range, patches)?))
43    }
44}
45
46fn slice_bitpacked(
47    array: ArrayView<'_, BitPacked>,
48    range: Range<usize>,
49    patches: Option<Patches>,
50) -> VortexResult<ArrayRef> {
51    let offset_start = range.start + array.offset() as usize;
52    let offset_stop = range.end + array.offset() as usize;
53    let offset = offset_start % 1024;
54    let block_start = max(0, offset_start - offset);
55    let block_stop = offset_stop.div_ceil(1024) * 1024;
56
57    let encoded_start = (block_start / 8) * array.bit_width() as usize;
58    let encoded_stop = (block_stop / 8) * array.bit_width() as usize;
59
60    Ok(BitPacked::try_new(
61        array.packed().slice(encoded_start..encoded_stop),
62        array.dtype().as_ptype(),
63        array.validity()?.slice(range.clone())?,
64        patches,
65        array.bit_width(),
66        range.len(),
67        offset as u16,
68    )?
69    .into_array())
70}
71
72#[cfg(test)]
73mod tests {
74    use vortex_array::IntoArray;
75    use vortex_array::VortexSessionExecute;
76    use vortex_array::array_session;
77    use vortex_array::arrays::PrimitiveArray;
78    use vortex_array::arrays::SliceArray;
79    use vortex_error::VortexResult;
80
81    use crate::BitPacked;
82    use crate::bitpack_compress::bitpack_encode;
83
84    #[test]
85    fn test_reduce_parent_returns_bitpacked_slice() -> VortexResult<()> {
86        let mut ctx = array_session().create_execution_ctx();
87        let values = PrimitiveArray::from_iter(0u32..2048);
88        let bitpacked = bitpack_encode(&values, 11, None, &mut ctx)?;
89
90        let slice_array = SliceArray::new(bitpacked.clone().into_array(), 500..1500);
91
92        let bitpacked_ref = bitpacked.into_array();
93        let reduced = bitpacked_ref
94            .reduce_parent(&slice_array.into_array(), 0)?
95            .expect("expected slice kernel to execute");
96
97        assert!(reduced.is::<BitPacked>());
98        let reduced_bp = reduced.as_::<BitPacked>();
99        assert_eq!(reduced_bp.offset(), 500);
100        assert_eq!(reduced.len(), 1000);
101
102        Ok(())
103    }
104}
vortex_fastlanes/bitpacking/compute/slice.rs

vortex_fastlanes/bitpacking/compute/
slice.rs
