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::bit::get_bit;
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            let ptr = bits.inner().as_ptr();
83            let offset = bits.offset();
84            BitBuffer::collect_bool(indices.len(), |i| unsafe {
85                // SAFETY: we're already iterating over every index, a bound
86                // is excessive
87                let idx = indices.get_unchecked(i).as_();
88                // SAFETY: we've asserted index is in bounds before
89                get_bit_unchecked(ptr, offset + idx)
90            })
91        };
92    };
93
94    let ptr = validity.inner().as_ptr();
95    if bits.len() <= COLLECT_TO_VEC {
96        let offset = validity.offset();
97        let bools: Vec<bool> = bits.iter().collect();
98        return BitBuffer::collect_bool(indices.len(), |i| {
99            // SAFETY: we're already iterating over every index, a bound
100            // is excessive
101            let idx = unsafe { indices.get_unchecked(i).as_() };
102            // SAFETY: we're verified validity has same size as indices before
103            let mask_idx: usize = unsafe { get_bit_unchecked(ptr, offset + i) }.as_();
104            bools[idx & mask_idx.wrapping_neg()]
105        });
106    }
107
108    let validity_offset = validity.offset();
109    let buf = bits.inner();
110    let bits_offset = bits.offset();
111    BitBuffer::collect_bool(indices.len(), |i| {
112        // SAFETY: we're already iterating over every index, a bound
113        // is excessive
114        let idx = unsafe { indices.get_unchecked(i).as_() };
115        // SAFETY: we're verified validity has same size as indices before
116        let mask_idx: usize = unsafe { get_bit_unchecked(ptr, validity_offset + i) }.as_();
117        let masked_idx = idx & mask_idx.wrapping_neg();
118        get_bit(buf, bits_offset + masked_idx)
119    })
120}
121
122fn first_unset(bits: BitBufferView<'_>) -> Option<usize> {
123    let chunks = bits.chunks();
124    for (word_idx, word) in chunks.iter().enumerate() {
125        if word != u64::MAX {
126            return Some(word_idx * 64 + (!word).trailing_zeros() as usize);
127        }
128    }
129
130    let remainder_start = bits.len() / 64 * 64;
131    let target = remainder_start + (!chunks.remainder_bits()).trailing_zeros() as usize;
132    (target < bits.len()).then_some(target)
133}
134
135#[cfg(test)]
136mod tests {
137    use rand::RngExt;
138    use rand::SeedableRng;
139    use rand::rngs::StdRng;
140    use rstest::rstest;
141
142    use super::take_bits;
143    use crate::BitBuffer;
144
145    fn random_bits(len: usize) -> BitBuffer {
146        let mut rng = StdRng::seed_from_u64(42);
147        BitBuffer::collect_bool(len, |_| rng.random())
148    }
149
150    fn random_indices(len: usize, bound: usize) -> Vec<usize> {
151        let mut rng = StdRng::seed_from_u64(43);
152        (0..len).map(|_| rng.random_range(0..bound)).collect()
153    }
154
155    #[rstest]
156    #[case(300)]
157    #[case(5000)]
158    fn gathers(#[case] bits_len: usize) {
159        let bits = random_bits(bits_len + 11);
160        let view = bits.as_view().slice(11..);
161        let len = 400;
162        let indices: Vec<i64> = random_indices(len, bits_len)
163            .into_iter()
164            .map(|idx| i64::try_from(idx).unwrap())
165            .collect();
166
167        let taken = take_bits(view, &indices, None);
168        assert_eq!(taken.len(), len);
169        for (i, &idx) in indices.iter().enumerate() {
170            assert_eq!(
171                taken.value(i),
172                bits.value(usize::try_from(idx).unwrap() + 11)
173            );
174        }
175
176        let validity = BitBuffer::collect_bool(len, |i| i % 3 != 0);
177        let with_garbage: Vec<i64> = indices
178            .iter()
179            .enumerate()
180            .map(|(i, &idx)| match i % 6 {
181                0 => i64::MAX,
182                3 => -1,
183                _ => idx,
184            })
185            .collect();
186
187        let taken = take_bits(view, &with_garbage, Some(validity.as_view()));
188        assert_eq!(taken.len(), len);
189        for (i, &idx) in with_garbage.iter().enumerate() {
190            if validity.value(i) {
191                assert_eq!(
192                    taken.value(i),
193                    bits.value(usize::try_from(idx).unwrap() + 11)
194                );
195            }
196        }
197    }
198
199    #[rstest]
200    #[case(BitBuffer::new_set(100), |_: usize| true)]
201    #[case(BitBuffer::new_unset(100), |_: usize| false)]
202    #[case(BitBuffer::collect_bool(100, |i| i == 37), |idx: usize| idx == 37)]
203    #[case(BitBuffer::collect_bool(100, |i| i != 37), |idx: usize| idx != 37)]
204    fn nearly_full(#[case] bits: BitBuffer, #[case] expected: fn(usize) -> bool) {
205        let indices: Vec<u8> = random_indices(300, bits.len())
206            .into_iter()
207            .map(|idx| u8::try_from(idx).unwrap())
208            .collect();
209
210        let taken = take_bits(bits.as_view(), &indices, None);
211        for (i, &idx) in indices.iter().enumerate() {
212            assert_eq!(taken.value(i), expected(usize::from(idx)));
213        }
214    }
215
216    #[test]
217    fn small_indices() {
218        let bits = BitBuffer::collect_bool(10_000, |i| i != 37);
219        let taken = take_bits(bits.as_view(), &[36u32, 37, 38], None);
220        assert!(taken.value(0));
221        assert!(!taken.value(1));
222        assert!(taken.value(2));
223    }
224
225    #[test]
226    fn empty() {
227        let bits = random_bits(64);
228        assert_eq!(take_bits::<u16>(bits.as_view(), &[], None).len(), 0);
229
230        let bits = BitBuffer::new_unset(0);
231        let validity = BitBuffer::new_unset(3);
232        let taken = take_bits(bits.as_view(), &[7u32, 8, 9], Some(validity.as_view()));
233        assert_eq!(taken, BitBuffer::new_unset(3));
234    }
235}
vortex_buffer/bit/take.rs

vortex_buffer/bit/
take.rs
