vortex_fastlanes/bitpacking/vtable/
operations.rs

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

vortex_fastlanes/bitpacking/vtable/
operations.rs
