vortex_fastlanes/bitpacking/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::scalar::Scalar;
5use vortex_array::vtable::OperationsVTable;
6use vortex_error::VortexResult;
7
8use crate::BitPackedArray;
9use crate::BitPackedVTable;
10use crate::bitpack_decompress;
11
12impl OperationsVTable<BitPackedVTable> for BitPackedVTable {
13    fn scalar_at(array: &BitPackedArray, index: usize) -> VortexResult<Scalar> {
14        Ok(
15            if let Some(patches) = array.patches()
16                && let Some(patch) = patches.get_patched(index)?
17            {
18                patch
19            } else {
20                bitpack_decompress::unpack_single(array, index)
21            },
22        )
23    }
24}
25
26#[cfg(test)]
27mod test {
28    use std::ops::Range;
29    use std::sync::LazyLock;
30
31    use vortex_array::Array;
32    use vortex_array::IntoArray;
33    use vortex_array::VortexSessionExecute;
34    use vortex_array::arrays::PrimitiveArray;
35    use vortex_array::arrays::SliceArray;
36    use vortex_array::assert_arrays_eq;
37    use vortex_array::assert_nth_scalar;
38    use vortex_array::buffer::BufferHandle;
39    use vortex_array::dtype::DType;
40    use vortex_array::dtype::Nullability;
41    use vortex_array::dtype::PType;
42    use vortex_array::patches::Patches;
43    use vortex_array::scalar::Scalar;
44    use vortex_array::session::ArraySession;
45    use vortex_array::validity::Validity;
46    use vortex_array::vtable::VTable;
47    use vortex_buffer::Alignment;
48    use vortex_buffer::Buffer;
49    use vortex_buffer::ByteBuffer;
50    use vortex_buffer::buffer;
51
52    use crate::BitPackedArray;
53    use crate::BitPackedVTable;
54
55    static SESSION: LazyLock<vortex_session::VortexSession> =
56        LazyLock::new(|| vortex_session::VortexSession::empty().with::<ArraySession>());
57
58    fn slice_via_kernel(array: &BitPackedArray, range: Range<usize>) -> BitPackedArray {
59        let slice_array = SliceArray::new(array.clone().into_array(), range);
60        let mut ctx = SESSION.create_execution_ctx();
61        let sliced = <BitPackedVTable as VTable>::execute_parent(
62            array,
63            &slice_array.into_array(),
64            0,
65            &mut ctx,
66        )
67        .expect("execute_parent failed")
68        .expect("expected slice kernel to execute");
69        sliced.as_::<BitPackedVTable>().clone()
70    }
71
72    #[test]
73    pub fn slice_block() {
74        let arr = BitPackedArray::encode(
75            PrimitiveArray::from_iter((0u32..2048).map(|v| v % 64)).as_ref(),
76            6,
77        )
78        .unwrap();
79        let sliced = slice_via_kernel(&arr, 1024..2048);
80        assert_nth_scalar!(sliced, 0, 1024u32 % 64);
81        assert_nth_scalar!(sliced, 1023, 2047u32 % 64);
82        assert_eq!(sliced.offset(), 0);
83        assert_eq!(sliced.len(), 1024);
84    }
85
86    #[test]
87    pub fn slice_within_block() {
88        let arr = BitPackedArray::encode(
89            PrimitiveArray::from_iter((0u32..2048).map(|v| v % 64)).as_ref(),
90            6,
91        )
92        .unwrap();
93        let sliced = slice_via_kernel(&arr, 512..1434);
94        assert_nth_scalar!(sliced, 0, 512u32 % 64);
95        assert_nth_scalar!(sliced, 921, 1433u32 % 64);
96        assert_eq!(sliced.offset(), 512);
97        assert_eq!(sliced.len(), 922);
98    }
99
100    #[test]
101    fn slice_within_block_u8s() {
102        let packed = BitPackedArray::encode(
103            PrimitiveArray::from_iter((0..10_000).map(|i| (i % 63) as u8)).as_ref(),
104            7,
105        )
106        .unwrap();
107
108        let compressed = packed.slice(768..9999).unwrap();
109        assert_nth_scalar!(compressed, 0, (768 % 63) as u8);
110        assert_nth_scalar!(compressed, compressed.len() - 1, (9998 % 63) as u8);
111    }
112
113    #[test]
114    fn slice_block_boundary_u8s() {
115        let packed = BitPackedArray::encode(
116            PrimitiveArray::from_iter((0..10_000).map(|i| (i % 63) as u8)).as_ref(),
117            7,
118        )
119        .unwrap();
120
121        let compressed = packed.slice(7168..9216).unwrap();
122        assert_nth_scalar!(compressed, 0, (7168 % 63) as u8);
123        assert_nth_scalar!(compressed, compressed.len() - 1, (9215 % 63) as u8);
124    }
125
126    #[test]
127    fn double_slice_within_block() {
128        let arr = BitPackedArray::encode(
129            PrimitiveArray::from_iter((0u32..2048).map(|v| v % 64)).as_ref(),
130            6,
131        )
132        .unwrap();
133        let sliced = slice_via_kernel(&arr, 512..1434);
134        assert_nth_scalar!(sliced, 0, 512u32 % 64);
135        assert_nth_scalar!(sliced, 921, 1433u32 % 64);
136        assert_eq!(sliced.offset(), 512);
137        assert_eq!(sliced.len(), 922);
138        let doubly_sliced = slice_via_kernel(&sliced, 127..911);
139        assert_nth_scalar!(doubly_sliced, 0, (512u32 + 127) % 64);
140        assert_nth_scalar!(doubly_sliced, 783, (512u32 + 910) % 64);
141        assert_eq!(doubly_sliced.offset(), 639);
142        assert_eq!(doubly_sliced.len(), 784);
143    }
144
145    #[test]
146    fn slice_empty_patches() {
147        // We create an array that has 1 element that does not fit in the 6-bit range.
148        let array = BitPackedArray::encode(&buffer![0u32..=64].into_array(), 6).unwrap();
149
150        assert!(array.patches().is_some());
151
152        let patch_indices = array.patches().unwrap().indices().clone();
153        assert_eq!(patch_indices.len(), 1);
154
155        // Slicing drops the empty patches array.
156        let sliced_bp = slice_via_kernel(&array, 0..64);
157        assert!(sliced_bp.patches().is_none());
158    }
159
160    #[test]
161    fn take_after_slice() {
162        // Check that our take implementation respects the offsets applied after slicing.
163
164        let array =
165            BitPackedArray::encode(PrimitiveArray::from_iter((63u32..).take(3072)).as_ref(), 6)
166                .unwrap();
167
168        // Slice the array.
169        // The resulting array will still have 3 1024-element chunks.
170        let sliced = array.slice(922..2061).unwrap();
171
172        // Take one element from each chunk.
173        // Chunk 1: physical indices  922-1023, logical indices    0-101
174        // Chunk 2: physical indices 1024-2047, logical indices  102-1125
175        // Chunk 3: physical indices 2048-2060, logical indices 1126-1138
176
177        let taken = sliced
178            .take(buffer![101i64, 1125, 1138].into_array())
179            .unwrap();
180        assert_eq!(taken.len(), 3);
181    }
182
183    #[test]
184    fn scalar_at_invalid_patches() {
185        let packed_array = unsafe {
186            BitPackedArray::new_unchecked(
187                BufferHandle::new_host(ByteBuffer::copy_from_aligned(
188                    [0u8; 128],
189                    Alignment::of::<u32>(),
190                )),
191                DType::Primitive(PType::U32, true.into()),
192                Validity::AllInvalid,
193                Some(
194                    Patches::new(
195                        8,
196                        0,
197                        buffer![1u32].into_array(),
198                        PrimitiveArray::new(buffer![999u32], Validity::AllValid).to_array(),
199                        None,
200                    )
201                    .unwrap(),
202                ),
203                1,
204                8,
205                0,
206            )
207            .into_array()
208        };
209        assert_eq!(
210            packed_array.scalar_at(1).unwrap(),
211            Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable))
212        );
213    }
214
215    #[test]
216    fn scalar_at() {
217        let values = (0u32..257).collect::<Buffer<_>>();
218        let uncompressed = values.clone().into_array();
219        let packed = BitPackedArray::encode(&uncompressed, 8).unwrap();
220        assert!(packed.patches().is_some());
221
222        let patches = packed.patches().unwrap().indices().clone();
223        assert_eq!(
224            usize::try_from(&patches.scalar_at(0).unwrap()).unwrap(),
225            256
226        );
227
228        let expected = PrimitiveArray::from_iter(values.iter().copied());
229        assert_arrays_eq!(packed, expected);
230    }
231}
vortex_fastlanes/bitpacking/vtable/operations.rs

vortex_fastlanes/bitpacking/vtable/
operations.rs
