vortex_fastlanes/bitpacking/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::mem;
5use std::mem::MaybeUninit;
6
7use fastlanes::BitPacking;
8use vortex_array::ArrayRef;
9use vortex_array::ArrayView;
10use vortex_array::ExecutionCtx;
11use vortex_array::IntoArray;
12use vortex_array::arrays::PrimitiveArray;
13use vortex_array::arrays::dict::TakeExecute;
14use vortex_array::dtype::IntegerPType;
15use vortex_array::dtype::NativePType;
16use vortex_array::dtype::PType;
17use vortex_array::match_each_integer_ptype;
18use vortex_array::match_each_unsigned_integer_ptype;
19use vortex_array::validity::Validity;
20use vortex_buffer::Buffer;
21use vortex_buffer::BufferMut;
22use vortex_error::VortexExpect as _;
23use vortex_error::VortexResult;
24
25use super::chunked_indices;
26use crate::BitPacked;
27use crate::BitPackedArrayExt;
28use crate::bitpack_decompress;
29
30// TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`.
31/// assuming the buffer is already allocated (which will happen at most once) then unpacking
32/// all 1024 elements takes ~8.8x as long as unpacking a single element on an M2 Macbook Air.
33/// see <https://github.com/vortex-data/vortex/pull/190#issue-2223752833>
34pub(super) const UNPACK_CHUNK_THRESHOLD: usize = 8;
35
36impl TakeExecute for BitPacked {
37    fn take(
38        array: ArrayView<'_, Self>,
39        indices: &ArrayRef,
40        ctx: &mut ExecutionCtx,
41    ) -> VortexResult<Option<ArrayRef>> {
42        // If the indices are large enough, it's faster to flatten and take the primitive array.
43        if indices.len() * UNPACK_CHUNK_THRESHOLD > array.len() {
44            let prim = array.array().clone().execute::<PrimitiveArray>(ctx)?;
45            return prim.into_array().take(indices.clone()).map(Some);
46        }
47
48        // NOTE: we use the unsigned PType because all values in the BitPackedArray must
49        //  be non-negative (pre-condition of creating the BitPackedArray).
50        let ptype: PType = PType::try_from(array.dtype())?;
51        let validity = array.validity()?;
52        let taken_validity = validity.take(indices)?;
53
54        let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
55        let taken = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| {
56            match_each_integer_ptype!(indices.ptype(), |I| {
57                take_primitive::<T, I>(array, &indices, taken_validity, ctx)?
58            })
59        });
60        let taken = if ptype.is_signed_int() {
61            PrimitiveArray::from_buffer_handle(
62                taken.buffer_handle().clone(),
63                ptype,
64                taken.validity()?,
65            )
66        } else {
67            taken
68        };
69        Ok(Some(taken.into_array()))
70    }
71}
72
73fn take_primitive<T: NativePType + BitPacking, I: IntegerPType>(
74    array: ArrayView<'_, BitPacked>,
75    indices: &PrimitiveArray,
76    taken_validity: Validity,
77    ctx: &mut ExecutionCtx,
78) -> VortexResult<PrimitiveArray> {
79    if indices.is_empty() {
80        return Ok(PrimitiveArray::new(Buffer::<T>::empty(), taken_validity));
81    }
82
83    let offset = array.offset() as usize;
84    let bit_width = array.bit_width() as usize;
85
86    let packed = array.packed_slice::<T>();
87
88    // Group indices by 1024-element chunk, *without* allocating on the heap
89    let indices_iter = indices.as_slice::<I>().iter().map(|i| {
90        i.to_usize()
91            .vortex_expect("index must be expressible as usize")
92    });
93
94    let mut output = BufferMut::<T>::with_capacity(indices.len());
95    let mut unpacked = [const { MaybeUninit::uninit() }; 1024];
96    let chunk_len = 128 * bit_width / size_of::<T>();
97
98    chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| {
99        let packed = &packed[chunk_idx * chunk_len..][..chunk_len];
100
101        let mut have_unpacked = false;
102        let mut offset_chunk_iter = indices_within_chunk.chunks_exact(UNPACK_CHUNK_THRESHOLD);
103
104        // this loop only runs if we have at least UNPACK_CHUNK_THRESHOLD offsets
105        for offset_chunk in &mut offset_chunk_iter {
106            assert_eq!(offset_chunk.len(), UNPACK_CHUNK_THRESHOLD); // let compiler know slice length
107            if !have_unpacked {
108                unsafe {
109                    let dst: &mut [MaybeUninit<T>] = &mut unpacked;
110                    let dst: &mut [T] = mem::transmute(dst);
111                    BitPacking::unchecked_unpack(bit_width, packed, dst);
112                }
113                have_unpacked = true;
114            }
115
116            for &index in offset_chunk {
117                output.push(unsafe { unpacked[index].assume_init() });
118            }
119        }
120
121        // if we have a remainder (i.e., < UNPACK_CHUNK_THRESHOLD leftover offsets), we need to handle it
122        if !offset_chunk_iter.remainder().is_empty() {
123            if have_unpacked {
124                // we already bulk unpacked this chunk, so we can just push the remaining elements
125                for &index in offset_chunk_iter.remainder() {
126                    output.push(unsafe { unpacked[index].assume_init() });
127                }
128            } else {
129                // we had fewer than UNPACK_CHUNK_THRESHOLD offsets in the first place,
130                // so we need to unpack each one individually
131                for &index in offset_chunk_iter.remainder() {
132                    output.push(unsafe {
133                        bitpack_decompress::unpack_single_primitive::<T>(packed, bit_width, index)
134                    });
135                }
136            }
137        }
138    });
139
140    let unpatched_taken = if array.dtype().as_ptype().is_signed_int() {
141        let primitive = PrimitiveArray::new(output, taken_validity);
142        PrimitiveArray::from_buffer_handle(
143            primitive.buffer_handle().clone(),
144            array.dtype().as_ptype(),
145            primitive.validity()?,
146        )
147    } else {
148        PrimitiveArray::new(output, taken_validity)
149    };
150    if let Some(patches) = array.patches()
151        && let Some(patches) = patches.take(&indices.clone().into_array(), ctx)?
152    {
153        return unpatched_taken.patch(&patches, ctx);
154    }
155
156    Ok(unpatched_taken)
157}
158
159#[cfg(test)]
160#[expect(clippy::cast_possible_truncation)]
161mod test {
162    use std::sync::LazyLock;
163
164    use rand::RngExt;
165    use rand::distr::Uniform;
166    use rand::rng;
167    use rstest::rstest;
168    use vortex_array::IntoArray;
169    use vortex_array::VortexSessionExecute;
170    use vortex_array::arrays::PrimitiveArray;
171    use vortex_array::assert_arrays_eq;
172    use vortex_array::compute::conformance::take::test_take_conformance;
173    use vortex_array::validity::Validity;
174    use vortex_buffer::Buffer;
175    use vortex_buffer::buffer;
176    use vortex_session::VortexSession;
177
178    use crate::BitPackedArray;
179    use crate::BitPackedData;
180    use crate::bitpacking::array::BitPackedArrayExt;
181    use crate::bitpacking::compute::take::take_primitive;
182
183    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
184        let session = vortex_array::array_session();
185        crate::initialize(&session);
186        session
187    });
188
189    #[test]
190    fn take_indices() {
191        let mut ctx = SESSION.create_execution_ctx();
192        let indices = buffer![0, 125, 2047, 2049, 2151, 2790].into_array();
193
194        // Create a u8 array modulo 63.
195        let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8));
196        let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap();
197
198        let primitive_result = bitpacked.take(indices).unwrap();
199        assert_arrays_eq!(
200            primitive_result,
201            PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]),
202            &mut ctx
203        );
204    }
205
206    #[test]
207    fn take_with_patches() {
208        let mut ctx = SESSION.create_execution_ctx();
209        let unpacked = Buffer::from_iter(0u32..1024).into_array();
210        let bitpacked = BitPackedData::encode(&unpacked, 2, &mut ctx).unwrap();
211
212        let indices = buffer![0, 2, 4, 6].into_array();
213
214        let primitive_result = bitpacked.take(indices).unwrap();
215        assert_arrays_eq!(
216            primitive_result,
217            PrimitiveArray::from_iter([0u32, 2, 4, 6]),
218            &mut ctx
219        );
220    }
221
222    #[test]
223    fn take_sliced_indices() {
224        let mut ctx = SESSION.create_execution_ctx();
225        let indices = buffer![1919, 1921].into_array();
226
227        // Create a u8 array modulo 63.
228        let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8));
229        let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap();
230        let sliced = bitpacked.slice(128..2050).unwrap();
231
232        let primitive_result = sliced.take(indices).unwrap();
233        assert_arrays_eq!(
234            primitive_result,
235            PrimitiveArray::from_iter([31u8, 33]),
236            &mut ctx
237        );
238    }
239
240    #[test]
241    #[cfg_attr(miri, ignore)] // This test is too slow on miri
242    fn take_random_indices() {
243        let mut ctx = SESSION.create_execution_ctx();
244        let num_patches: usize = 128;
245        let values = (0..u16::MAX as u32 + num_patches as u32).collect::<Buffer<_>>();
246        let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable);
247        let packed = BitPackedData::encode(&uncompressed.into_array(), 16, &mut ctx).unwrap();
248        assert!(packed.patches().is_some());
249
250        let rng = rng();
251        let range = Uniform::new(0, values.len()).unwrap();
252        let random_indices =
253            PrimitiveArray::from_iter(rng.sample_iter(range).take(10_000).map(|i| i as u32));
254        let taken = packed.take(random_indices.clone().into_array()).unwrap();
255
256        // sanity check
257        random_indices
258            .as_slice::<u32>()
259            .iter()
260            .enumerate()
261            .for_each(|(ti, i)| {
262                assert_eq!(
263                    u32::try_from(&packed.execute_scalar(*i as usize, &mut ctx).unwrap()).unwrap(),
264                    values[*i as usize]
265                );
266                assert_eq!(
267                    u32::try_from(&taken.execute_scalar(ti, &mut ctx).unwrap()).unwrap(),
268                    values[*i as usize]
269                );
270            });
271    }
272
273    #[test]
274    #[cfg_attr(miri, ignore)]
275    fn take_signed_with_patches() {
276        let mut ctx = SESSION.create_execution_ctx();
277        let start =
278            BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx)
279                .unwrap();
280
281        let taken_primitive = take_primitive::<u32, u64>(
282            start.as_view(),
283            &PrimitiveArray::from_iter([0u64, 1, 2, 3]),
284            Validity::NonNullable,
285            &mut ctx,
286        )
287        .unwrap();
288        assert_arrays_eq!(
289            taken_primitive,
290            PrimitiveArray::from_iter([1i32, 2, 3, 4]),
291            &mut ctx
292        );
293    }
294
295    #[test]
296    fn take_nullable_with_nullables() {
297        let mut ctx = SESSION.create_execution_ctx();
298        let start =
299            BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx)
300                .unwrap();
301
302        let taken_primitive = start
303            .take(
304                PrimitiveArray::from_option_iter([Some(0u64), Some(1), None, Some(3)]).into_array(),
305            )
306            .unwrap();
307        assert_arrays_eq!(
308            taken_primitive,
309            PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]),
310            &mut ctx
311        );
312        let taken_primitive_prim = taken_primitive.execute::<PrimitiveArray>(&mut ctx).unwrap();
313        assert_eq!(taken_primitive_prim.invalid_count(&mut ctx).unwrap(), 1);
314    }
315
316    fn bp(array: vortex_array::ArrayRef, bit_width: u8) -> BitPackedArray {
317        BitPackedData::encode(&array, bit_width, &mut SESSION.create_execution_ctx()).unwrap()
318    }
319
320    #[rstest]
321    #[case(bp(PrimitiveArray::from_iter((0..100).map(|i| (i % 63) as u8)).into_array(), 6))]
322    #[case(bp(PrimitiveArray::from_iter((0..256).map(|i| i as u32)).into_array(), 8))]
323    #[case(bp(buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(), 3))]
324    #[case(bp(
325        PrimitiveArray::from_option_iter([Some(10u16), None, Some(20), Some(30), None]).into_array(),
326        5
327    ))]
328    #[case(bp(buffer![42u32].into_array(), 6))]
329    #[case(bp(PrimitiveArray::from_iter((0..1024).map(|i| i as u32)).into_array(), 8))]
330    fn test_take_bitpacked_conformance(#[case] bitpacked: BitPackedArray) {
331        test_take_conformance(&bitpacked.into_array(), &mut SESSION.create_execution_ctx());
332    }
333}
vortex_fastlanes/bitpacking/compute/take.rs

vortex_fastlanes/bitpacking/compute/
take.rs
