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::ExecutionCtx;
9use vortex_array::IntoArray;
10use vortex_array::arrays::SliceKernel;
11use vortex_error::VortexResult;
12
13use crate::BitPackedArray;
14use crate::BitPackedVTable;
15
16impl SliceKernel for BitPackedVTable {
17    fn slice(
18        array: &BitPackedArray,
19        range: Range<usize>,
20        _ctx: &mut ExecutionCtx,
21    ) -> VortexResult<Option<ArrayRef>> {
22        let offset_start = range.start + array.offset() as usize;
23        let offset_stop = range.end + array.offset() as usize;
24        let offset = offset_start % 1024;
25        let block_start = max(0, offset_start - offset);
26        let block_stop = offset_stop.div_ceil(1024) * 1024;
27
28        let encoded_start = (block_start / 8) * array.bit_width() as usize;
29        let encoded_stop = (block_stop / 8) * array.bit_width() as usize;
30
31        // slice the buffer using the encoded start/stop values
32        // SAFETY: slicing packed values without decoding preserves invariants
33        Ok(Some(unsafe {
34            BitPackedArray::new_unchecked(
35                array.packed().slice(encoded_start..encoded_stop),
36                array.dtype().clone(),
37                array.validity()?.slice(range.clone())?,
38                array
39                    .patches()
40                    .map(|p| p.slice(range.clone()))
41                    .transpose()?
42                    .flatten(),
43                array.bit_width(),
44                range.len(),
45                offset as u16,
46            )
47            .into_array()
48        }))
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use std::sync::LazyLock;
55
56    use vortex_array::Array;
57    use vortex_array::IntoArray;
58    use vortex_array::VortexSessionExecute;
59    use vortex_array::arrays::SliceArray;
60    use vortex_array::session::ArraySession;
61    use vortex_array::vtable::VTable;
62    use vortex_error::VortexResult;
63    use vortex_session::VortexSession;
64
65    use crate::BitPackedVTable;
66    use crate::bitpack_compress::bitpack_encode;
67
68    static SESSION: LazyLock<VortexSession> =
69        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
70
71    #[test]
72    fn test_execute_parent_returns_bitpacked_slice() -> VortexResult<()> {
73        let values = vortex_array::arrays::PrimitiveArray::from_iter(0u32..2048);
74        let bitpacked = bitpack_encode(&values, 11, None)?;
75
76        let slice_array = SliceArray::new(bitpacked.clone().into_array(), 500..1500);
77
78        let mut ctx = SESSION.create_execution_ctx();
79        let reduced = <BitPackedVTable as VTable>::execute_parent(
80            &bitpacked,
81            &slice_array.into_array(),
82            0,
83            &mut ctx,
84        )?
85        .expect("expected slice kernel to execute");
86
87        assert!(reduced.is::<BitPackedVTable>());
88        let reduced_bp = reduced.as_::<BitPackedVTable>();
89        assert_eq!(reduced_bp.offset(), 500);
90        assert_eq!(reduced.len(), 1000);
91
92        Ok(())
93    }
94}
vortex_fastlanes/bitpacking/compute/slice.rs

vortex_fastlanes/bitpacking/compute/
slice.rs
