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 (offset_chunks, remainder) = indices_within_chunk.as_chunks::<UNPACK_CHUNK_THRESHOLD>();
103
104        // this loop only runs if we have at least UNPACK_CHUNK_THRESHOLD offsets
105        for offset_chunk in offset_chunks {
106            if !have_unpacked {
107                unsafe {
108                    let dst: &mut [MaybeUninit<T>] = &mut unpacked;
109                    let dst: &mut [T] = mem::transmute(dst);
110                    BitPacking::unchecked_unpack(bit_width, packed, dst);
111                }
112                have_unpacked = true;
113            }
114
115            for &index in offset_chunk {
116                output.push(unsafe { unpacked[index].assume_init() });
117            }
118        }
119
120        // if we have a remainder (i.e., < UNPACK_CHUNK_THRESHOLD leftover offsets), we need to handle it
121        if !remainder.is_empty() {
122            if have_unpacked {
123                // we already bulk unpacked this chunk, so we can just push the remaining elements
124                for &index in remainder {
125                    output.push(unsafe { unpacked[index].assume_init() });
126                }
127            } else {
128                // we had fewer than UNPACK_CHUNK_THRESHOLD offsets in the first place,
129                // so we need to unpack each one individually
130                for &index in remainder {
131                    output.push(unsafe {
132                        bitpack_decompress::unpack_single_primitive::<T>(packed, bit_width, index)
133                    });
134                }
135            }
136        }
137    });
138
139    let unpatched_taken = if array.dtype().as_ptype().is_signed_int() {
140        let primitive = PrimitiveArray::new(output, taken_validity);
141        PrimitiveArray::from_buffer_handle(
142            primitive.buffer_handle().clone(),
143            array.dtype().as_ptype(),
144            primitive.validity()?,
145        )
146    } else {
147        PrimitiveArray::new(output, taken_validity)
148    };
149    if let Some(patches) = array.patches()
150        && let Some(patches) = patches.take(&indices.clone().into_array(), ctx)?
151    {
152        return unpatched_taken.patch(&patches, ctx);
153    }
154
155    Ok(unpatched_taken)
156}
157
158#[cfg(test)]
159#[expect(clippy::cast_possible_truncation)]
160mod test {
161    use std::sync::LazyLock;
162
163    use rand::RngExt;
164    use rand::distr::Uniform;
165    use rand::rng;
166    use rstest::rstest;
167    use vortex_array::IntoArray;
168    use vortex_array::VortexSessionExecute;
169    use vortex_array::arrays::PrimitiveArray;
170    use vortex_array::assert_arrays_eq;
171    use vortex_array::compute::conformance::take::test_take_conformance;
172    use vortex_array::validity::Validity;
173    use vortex_buffer::Buffer;
174    use vortex_buffer::buffer;
175    use vortex_session::VortexSession;
176
177    use crate::BitPackedArray;
178    use crate::BitPackedData;
179    use crate::bitpacking::array::BitPackedArrayExt;
180    use crate::bitpacking::compute::take::take_primitive;
181
182    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
183        let session = vortex_array::array_session();
184        crate::initialize(&session);
185        session
186    });
187
188    #[test]
189    fn take_indices() {
190        let mut ctx = SESSION.create_execution_ctx();
191        let indices = buffer![0, 125, 2047, 2049, 2151, 2790].into_array();
192
193        // Create a u8 array modulo 63.
194        let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8));
195        let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap();
196
197        let primitive_result = bitpacked.take(indices).unwrap();
198        assert_arrays_eq!(
199            primitive_result,
200            PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]),
201            &mut ctx
202        );
203    }
204
205    #[test]
206    fn take_with_patches() {
207        let mut ctx = SESSION.create_execution_ctx();
208        let unpacked = Buffer::from_iter(0u32..1024).into_array();
209        let bitpacked = BitPackedData::encode(&unpacked, 2, &mut ctx).unwrap();
210
211        let indices = buffer![0, 2, 4, 6].into_array();
212
213        let primitive_result = bitpacked.take(indices).unwrap();
214        assert_arrays_eq!(
215            primitive_result,
216            PrimitiveArray::from_iter([0u32, 2, 4, 6]),
217            &mut ctx
218        );
219    }
220
221    #[test]
222    fn take_sliced_indices() {
223        let mut ctx = SESSION.create_execution_ctx();
224        let indices = buffer![1919, 1921].into_array();
225
226        // Create a u8 array modulo 63.
227        let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8));
228        let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap();
229        let sliced = bitpacked.slice(128..2050).unwrap();
230
231        let primitive_result = sliced.take(indices).unwrap();
232        assert_arrays_eq!(
233            primitive_result,
234            PrimitiveArray::from_iter([31u8, 33]),
235            &mut ctx
236        );
237    }
238
239    #[test]
240    #[cfg_attr(miri, ignore)] // This test is too slow on miri
241    fn take_random_indices() {
242        let mut ctx = SESSION.create_execution_ctx();
243        let num_patches: usize = 128;
244        let values = (0..u16::MAX as u32 + num_patches as u32).collect::<Buffer<_>>();
245        let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable);
246        let packed = BitPackedData::encode(&uncompressed.into_array(), 16, &mut ctx).unwrap();
247        assert!(packed.patches().is_some());
248
249        let rng = rng();
250        let range = Uniform::new(0, values.len()).unwrap();
251        let random_indices =
252            PrimitiveArray::from_iter(rng.sample_iter(range).take(10_000).map(|i| i as u32));
253        let taken = packed.take(random_indices.clone().into_array()).unwrap();
254
255        // sanity check
256        random_indices
257            .as_slice::<u32>()
258            .iter()
259            .enumerate()
260            .for_each(|(ti, i)| {
261                assert_eq!(
262                    u32::try_from(&packed.execute_scalar(*i as usize, &mut ctx).unwrap()).unwrap(),
263                    values[*i as usize]
264                );
265                assert_eq!(
266                    u32::try_from(&taken.execute_scalar(ti, &mut ctx).unwrap()).unwrap(),
267                    values[*i as usize]
268                );
269            });
270    }
271
272    #[test]
273    #[cfg_attr(miri, ignore)]
274    fn take_signed_with_patches() {
275        let mut ctx = SESSION.create_execution_ctx();
276        let start =
277            BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx)
278                .unwrap();
279
280        let taken_primitive = take_primitive::<u32, u64>(
281            start.as_view(),
282            &PrimitiveArray::from_iter([0u64, 1, 2, 3]),
283            Validity::NonNullable,
284            &mut ctx,
285        )
286        .unwrap();
287        assert_arrays_eq!(
288            taken_primitive,
289            PrimitiveArray::from_iter([1i32, 2, 3, 4]),
290            &mut ctx
291        );
292    }
293
294    #[test]
295    fn take_nullable_with_nullables() {
296        let mut ctx = SESSION.create_execution_ctx();
297        let start =
298            BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx)
299                .unwrap();
300
301        let taken_primitive = start
302            .take(
303                PrimitiveArray::from_option_iter([Some(0u64), Some(1), None, Some(3)]).into_array(),
304            )
305            .unwrap();
306        assert_arrays_eq!(
307            taken_primitive,
308            PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]),
309            &mut ctx
310        );
311        let taken_primitive_prim = taken_primitive.execute::<PrimitiveArray>(&mut ctx).unwrap();
312        assert_eq!(taken_primitive_prim.invalid_count(&mut ctx).unwrap(), 1);
313    }
314
315    fn bp(array: vortex_array::ArrayRef, bit_width: u8) -> BitPackedArray {
316        BitPackedData::encode(&array, bit_width, &mut SESSION.create_execution_ctx()).unwrap()
317    }
318
319    #[rstest]
320    #[case(bp(PrimitiveArray::from_iter((0..100).map(|i| (i % 63) as u8)).into_array(), 6))]
321    #[case(bp(PrimitiveArray::from_iter((0..256).map(|i| i as u32)).into_array(), 8))]
322    #[case(bp(buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(), 3))]
323    #[case(bp(
324        PrimitiveArray::from_option_iter([Some(10u16), None, Some(20), Some(30), None]).into_array(),
325        5
326    ))]
327    #[case(bp(buffer![42u32].into_array(), 6))]
328    #[case(bp(PrimitiveArray::from_iter((0..1024).map(|i| i as u32)).into_array(), 8))]
329    fn test_take_bitpacked_conformance(#[case] bitpacked: BitPackedArray) {
330        test_take_conformance(&bitpacked.into_array(), &mut SESSION.create_execution_ctx());
331    }
332}
vortex_fastlanes/bitpacking/compute/take.rs

vortex_fastlanes/bitpacking/compute/
take.rs
