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

vortex_fastlanes/bitpacking/vtable/
operations.rs
