vortex_buffer/bit/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::AsPrimitive;
5use num_traits::Zero;
6use vortex_error::VortexExpect;
7
8use crate::BitBuffer;
9use crate::BitBufferView;
10use crate::BufferMut;
11use crate::bit::get_bit_unchecked;
12
13// If you have less than a Linux memory page of bits, it pays to convert to a
14// Vec
15const COLLECT_TO_VEC: usize = 4096;
16
17/// Select bits at indices into a new BitBuffer
18pub fn take_bits<I>(
19    bits: BitBufferView<'_>,
20    indices: &[I],
21    validity: Option<BitBufferView<'_>>,
22) -> BitBuffer
23where
24    I: AsPrimitive<usize> + Ord + Zero,
25{
26    if let Some(validity) = validity {
27        assert_eq!(validity.len(), indices.len());
28    }
29
30    if indices.is_empty() || bits.is_empty() {
31        return BitBuffer::new_unset(indices.len());
32    }
33
34    // If we don't have a lot of indices, mostly full/empty case is very fast.
35    // This is, for example, the case of comparing a dict to a constant.
36    if bits.len() <= indices.len() * 64 {
37        match bits.true_count() {
38            0 => return BitBuffer::new_unset(indices.len()),
39            count if count == bits.len() => return BitBuffer::new_set(indices.len()),
40            1 => {
41                let target = bits.select(0).vortex_expect("one set bit");
42                return BitBuffer::collect_bool_multiversioned(indices.len(), |i| {
43                    // SAFETY: we're already iterating over every index, a bound
44                    // is excessive
45                    unsafe { indices.get_unchecked(i) }.as_() == target
46                });
47            }
48            count if count == bits.len() - 1 => {
49                let target = first_unset(bits).vortex_expect("one unset bit");
50                return BitBuffer::collect_bool_multiversioned(indices.len(), |i| {
51                    // SAFETY: we're already iterating over every index, a bound
52                    // is excessive
53                    unsafe { indices.get_unchecked(i) }.as_() != target
54                });
55            }
56            _ => {}
57        }
58    }
59
60    let Some(validity) = validity else {
61        // iterating over indices here for a bounds check is faster than a
62        // per-element branch in [] operator.
63        let first = indices[0];
64        let (min, max) = indices
65            .iter()
66            .fold((first, first), |(lo, hi), &x| (lo.min(x), hi.max(x)));
67        assert!(
68            min >= I::zero() && max.as_() < bits.len(),
69            "take index out of bounds"
70        );
71
72        return if bits.len() <= COLLECT_TO_VEC {
73            let bools: Vec<bool> = bits.iter().collect();
74            BitBuffer::collect_bool_multiversioned(indices.len(), |i| unsafe {
75                // SAFETY: we're already iterating over every index, a bound
76                // is excessive
77                let idx = indices.get_unchecked(i).as_();
78                // SAFETY: we've asserted index is in bounds before
79                *bools.get_unchecked(idx)
80            })
81        } else {
82            #[cfg(target_arch = "x86_64")]
83            if let Some(taken) = avx2::take(bits, indices, None) {
84                return taken;
85            }
86
87            let ptr = bits.inner().as_ptr();
88            let offset = bits.offset();
89            BitBuffer::collect_bool_multiversioned(indices.len(), |i| unsafe {
90                // SAFETY: we're already iterating over every index, a bound
91                // is excessive
92                let idx = indices.get_unchecked(i).as_();
93                // SAFETY: we've asserted index is in bounds before
94                get_bit_unchecked(ptr, offset + idx)
95            })
96        };
97    };
98
99    #[cfg(target_arch = "x86_64")]
100    if let Some(taken) = avx2::take(bits, indices, Some(validity)) {
101        return taken;
102    }
103
104    if bits.len() <= COLLECT_TO_VEC {
105        let bools: Vec<bool> = bits.iter().collect();
106        return take_with_validity(indices, validity, bits.len() - 1, |idx| unsafe {
107            // SAFETY: we clamp index so it's never out of bounds
108            *bools.get_unchecked(idx)
109        });
110    }
111
112    let buf = bits.inner().as_ptr();
113    let bits_offset = bits.offset();
114    take_with_validity(indices, validity, bits.len() - 1, |idx| unsafe {
115        // SAFETY: we clamp index so it's never out of bounds
116        get_bit_unchecked(buf, bits_offset + idx)
117    })
118}
119
120fn take_with_validity<I>(
121    indices: &[I],
122    validity: BitBufferView<'_>,
123    max_index: usize,
124    read: impl Fn(usize) -> bool,
125) -> BitBuffer
126where
127    I: AsPrimitive<usize>,
128{
129    let total = indices.len();
130    let groups = total / 8;
131
132    let validity_bytes = validity.inner();
133    let validity_offset = validity.offset();
134    let validity_shift = validity_offset % 8;
135    let validity_first_byte = validity_offset / 8;
136    let has_non_byte_shift = usize::from(validity_shift != 0);
137
138    let mut out = BufferMut::<u8>::with_capacity(total.div_ceil(8));
139    let mut violations = 0usize;
140    for group in 0..groups {
141        let byte = validity_first_byte + group;
142        let next_byte = byte + has_non_byte_shift;
143        // SAFETY: caller verified validity and indices length are equal
144        let lo: u16 = unsafe { *validity_bytes.get_unchecked(byte) }.as_();
145        let hi: u16 = unsafe { *validity_bytes.get_unchecked(next_byte) }.as_();
146        let valid = (lo | hi << 8) >> validity_shift;
147
148        let mut packed = 0u8;
149        for bit in 0..8 {
150            // SAFETY: full groups stay within indices
151            let idx = unsafe { indices.get_unchecked(group * 8 + bit) }.as_();
152            let keep = (usize::from(valid >> bit) & 1).wrapping_neg();
153            let masked = idx & keep;
154            let clamped = masked.min(max_index);
155            violations += usize::from(masked != clamped);
156            packed |= (read(clamped) as u8) << bit;
157        }
158        // SAFETY: out has sufficient capacity
159        unsafe { out.push_unchecked(packed) };
160    }
161
162    assert_eq!(violations, 0, "take index out of bounds");
163
164    let tail = total % 8;
165    if tail == 0 {
166        return BitBuffer::new(out.freeze(), total);
167    }
168
169    let start = groups * 8;
170    let ptr = validity_bytes.as_ptr();
171    let mut packed = 0u8;
172    for bit in 0..tail {
173        let pos = start + bit;
174        // SAFETY: pos stays within indices
175        let idx = unsafe { indices.get_unchecked(pos) }.as_();
176        // SAFETY: validity has same length as indices
177        let keep =
178            (unsafe { get_bit_unchecked(ptr, validity_offset + pos) } as usize).wrapping_neg();
179        let masked = idx & keep;
180        assert!(masked <= max_index, "take index out of bounds");
181        packed |= (read(masked) as u8) << bit;
182    }
183    // SAFETY: out has sufficient capacity
184    unsafe { out.push_unchecked(packed) };
185    BitBuffer::new(out.freeze(), total)
186}
187
188fn first_unset(bits: BitBufferView<'_>) -> Option<usize> {
189    let chunks = bits.chunks();
190    for (word_idx, word) in chunks.iter().enumerate() {
191        if word != u64::MAX {
192            return Some(word_idx * 64 + (!word).trailing_zeros() as usize);
193        }
194    }
195
196    let remainder_start = bits.len() / 64 * 64;
197    let target = remainder_start + (!chunks.remainder_bits()).trailing_zeros() as usize;
198    (target < bits.len()).then_some(target)
199}
200
201#[cfg(target_arch = "x86_64")]
202#[path = "take_avx2.rs"]
203mod avx2;
204
205#[cfg(test)]
206mod tests {
207    use num_traits::AsPrimitive;
208    use num_traits::Zero;
209    use rand::RngExt;
210    use rand::SeedableRng;
211    use rand::rngs::StdRng;
212    use rstest::rstest;
213
214    use super::take_bits;
215    use crate::BitBuffer;
216
217    fn random_bits(len: usize) -> BitBuffer {
218        let mut rng = StdRng::seed_from_u64(42);
219        BitBuffer::collect_bool(len, |_| rng.random())
220    }
221
222    fn random_indices(len: usize, bound: usize) -> Vec<usize> {
223        let mut rng = StdRng::seed_from_u64(43);
224        (0..len).map(|_| rng.random_range(0..bound)).collect()
225    }
226
227    fn gathers<I>(bits_len: usize, garbage: [I; 2])
228    where
229        I: AsPrimitive<usize> + TryFrom<usize> + Ord + Zero,
230        <I as TryFrom<usize>>::Error: std::fmt::Debug,
231    {
232        let bits = random_bits(bits_len + 11);
233        let view = bits.as_view().slice(11..);
234        let len = 403;
235        let indices: Vec<I> = random_indices(len, bits_len)
236            .into_iter()
237            .map(|idx| I::try_from(idx).unwrap())
238            .collect();
239
240        let taken = take_bits(view, &indices, None);
241        assert_eq!(taken.len(), len);
242        for (i, &idx) in indices.iter().enumerate() {
243            assert_eq!(taken.value(i), bits.value(idx.as_() + 11));
244        }
245
246        let validity = BitBuffer::collect_bool(len, |i| i % 3 != 0);
247        let with_garbage: Vec<I> = indices
248            .iter()
249            .enumerate()
250            .map(|(i, &idx)| match i % 6 {
251                0 => garbage[0],
252                3 => garbage[1],
253                _ => idx,
254            })
255            .collect();
256
257        let taken = take_bits(view, &with_garbage, Some(validity.as_view()));
258        assert_eq!(taken.len(), len);
259        for (i, &idx) in with_garbage.iter().enumerate() {
260            if validity.value(i) {
261                assert_eq!(taken.value(i), bits.value(idx.as_() + 11));
262            }
263        }
264    }
265
266    #[rstest]
267    #[case(300)]
268    #[case(5000)]
269    fn gathers_i64(#[case] bits_len: usize) {
270        gathers::<i64>(bits_len, [i64::MAX, -1]);
271    }
272
273    #[rstest]
274    #[case(300)]
275    #[case(5000)]
276    fn gathers_u16(#[case] bits_len: usize) {
277        gathers::<u16>(bits_len, [u16::MAX, u16::MAX - 1]);
278    }
279
280    #[rstest]
281    #[case(300)]
282    #[case(5000)]
283    fn gathers_u32(#[case] bits_len: usize) {
284        gathers::<u32>(bits_len, [u32::MAX, u32::MAX - 1]);
285    }
286
287    #[rstest]
288    #[case(BitBuffer::new_set(100), |_: usize| true)]
289    #[case(BitBuffer::new_unset(100), |_: usize| false)]
290    #[case(BitBuffer::collect_bool(100, |i| i == 37), |idx: usize| idx == 37)]
291    #[case(BitBuffer::collect_bool(100, |i| i != 37), |idx: usize| idx != 37)]
292    fn nearly_full(#[case] bits: BitBuffer, #[case] expected: fn(usize) -> bool) {
293        let indices: Vec<u8> = random_indices(300, bits.len())
294            .into_iter()
295            .map(|idx| u8::try_from(idx).unwrap())
296            .collect();
297
298        let taken = take_bits(bits.as_view(), &indices, None);
299        for (i, &idx) in indices.iter().enumerate() {
300            assert_eq!(taken.value(i), expected(usize::from(idx)));
301        }
302    }
303
304    #[test]
305    fn small_indices() {
306        let bits = BitBuffer::collect_bool(10_000, |i| i != 37);
307        let taken = take_bits(bits.as_view(), &[36u32, 37, 38], None);
308        assert!(taken.value(0));
309        assert!(!taken.value(1));
310        assert!(taken.value(2));
311    }
312
313    #[test]
314    fn empty() {
315        let bits = random_bits(64);
316        assert_eq!(take_bits::<u16>(bits.as_view(), &[], None).len(), 0);
317
318        let bits = BitBuffer::new_unset(0);
319        let validity = BitBuffer::new_unset(3);
320        let taken = take_bits(bits.as_view(), &[7u32, 8, 9], Some(validity.as_view()));
321        assert_eq!(taken, BitBuffer::new_unset(3));
322    }
323
324    #[test]
325    #[should_panic(expected = "out of bounds")]
326    fn out_of_bounds() {
327        let bits = BitBuffer::collect_bool(10_000, |i| i % 2 == 0);
328        let view = bits.as_view().slice(3..);
329        let mut indices: Vec<u32> = random_indices(403, 9_000)
330            .into_iter()
331            .map(|idx| u32::try_from(idx).unwrap())
332            .collect();
333        indices[80] = u32::MAX;
334        let validity = BitBuffer::collect_bool(403, |i| i != 200);
335        let _taken = take_bits(view, &indices, Some(validity.as_view()));
336    }
337}
vortex_buffer/bit/take.rs

vortex_buffer/bit/
take.rs
