polars_utils/
bloom_filter.rs

1//! Split-block bloom filter, the layout of parquet's bloom filters: 32-byte
2//! blocks, and a key sets one bit in each of the eight words of one block.
3//! See <https://github.com/apache/parquet-format/blob/master/BloomFilter.md>.
4
5use std::mem::ManuallyDrop;
6use std::sync::atomic::{AtomicU64, Ordering};
7
8use bytemuck::Zeroable;
9use bytemuck::allocation::zeroed_vec;
10
11use crate::mem::prefetch::prefetch_l1;
12
13const SALT: [u32; 8] = [
14    1203114875, 1150766481, 2284105051, 2729912477, 1884591559, 770785867, 2667333959, 1550580529,
15];
16
17const BLOCK_BYTES: usize = 32;
18
19/// Always below `num_blocks`.
20#[inline]
21fn block_index(hash: u64, num_blocks: usize) -> usize {
22    (((hash >> 32) * num_blocks as u64) >> 32) as usize
23}
24
25#[inline]
26fn block_mask(hash: u64) -> [u32; 8] {
27    let key = hash as u32;
28    std::array::from_fn(|i| 1u32 << (key.wrapping_mul(SALT[i]) >> 27))
29}
30
31#[inline]
32fn load_block(bytes: &[u8; BLOCK_BYTES]) -> [u32; 8] {
33    // SAFETY: the block is 8 words of 4 bytes.
34    let words: [u32; 8] = unsafe { std::ptr::read_unaligned(bytes.as_ptr().cast()) };
35    words.map(u32::from_le)
36}
37
38#[inline]
39fn store_block(block: [u32; 8], bytes: &mut [u8; BLOCK_BYTES]) {
40    // SAFETY: the block is 8 words of 4 bytes.
41    unsafe { std::ptr::write_unaligned(bytes.as_mut_ptr().cast(), block.map(u32::to_le)) }
42}
43
44/// The byte offset of the block of `hash` in `bitset`, which holds at least
45/// one block.
46#[inline]
47fn block_offset(bitset: &[u8], hash: u64) -> usize {
48    let num_blocks = bitset.len() / BLOCK_BYTES;
49    assert!(num_blocks > 0);
50    block_index(hash, num_blocks) * BLOCK_BYTES
51}
52
53#[inline]
54fn block_bytes(bitset: &[u8], hash: u64) -> &[u8; BLOCK_BYTES] {
55    let offset = block_offset(bitset, hash);
56    // SAFETY: `offset + BLOCK_BYTES <= bitset.len()`.
57    unsafe { &*(bitset.as_ptr().add(offset) as *const [u8; BLOCK_BYTES]) }
58}
59
60#[inline]
61fn block_bytes_mut(bitset: &mut [u8], hash: u64) -> &mut [u8; BLOCK_BYTES] {
62    let offset = block_offset(bitset, hash);
63    // SAFETY: `offset + BLOCK_BYTES <= bitset.len()`.
64    unsafe { &mut *(bitset.as_mut_ptr().add(offset) as *mut [u8; BLOCK_BYTES]) }
65}
66
67#[inline]
68fn block_contains(block: [u32; 8], mask: [u32; 8]) -> bool {
69    let mut found = true;
70    for i in 0..8 {
71        found &= block[i] & mask[i] != 0;
72    }
73    found
74}
75
76/// Whether `hash` is in the filter held by `bitset`, at least one block.
77pub fn is_in_set(bitset: &[u8], hash: u64) -> bool {
78    let block = load_block(block_bytes(bitset, hash));
79    block_contains(block, block_mask(hash))
80}
81
82/// Add `hash` to the filter held by `bitset`, at least one block.
83pub fn insert(bitset: &mut [u8], hash: u64) {
84    let bytes = block_bytes_mut(bitset, hash);
85    let mut block = load_block(bytes);
86    let mask = block_mask(hash);
87    for i in 0..8 {
88        block[i] |= mask[i];
89    }
90    store_block(block, bytes);
91}
92
93/// Eight words of a filter. Aligned so that a block never spans two cache
94/// lines.
95#[derive(Clone, Copy, Zeroable)]
96#[repr(C, align(32))]
97struct Block([u32; 8]);
98
99/// A split-block bloom filter over 64-bit hashes.
100#[derive(Clone)]
101pub struct SplitBlockBloom {
102    blocks: Vec<Block>,
103}
104
105impl SplitBlockBloom {
106    fn num_blocks_for(num_keys: usize, bits_per_key: usize) -> usize {
107        num_keys
108            .saturating_mul(bits_per_key)
109            .div_ceil(BLOCK_BYTES * 8)
110            .max(1)
111            .next_power_of_two()
112    }
113
114    /// The bytes `with_capacity` allocates.
115    pub fn size_for(num_keys: usize, bits_per_key: usize) -> usize {
116        Self::num_blocks_for(num_keys, bits_per_key).saturating_mul(BLOCK_BYTES)
117    }
118
119    /// A filter sized for `num_keys` keys at `bits_per_key` bits each, rounded
120    /// up to a power of two blocks. Its pages are only touched once a key is
121    /// inserted in them.
122    pub fn with_capacity(num_keys: usize, bits_per_key: usize) -> Self {
123        Self {
124            blocks: zeroed_vec(Self::num_blocks_for(num_keys, bits_per_key)),
125        }
126    }
127
128    pub fn size_bytes(&self) -> usize {
129        self.blocks.len() * BLOCK_BYTES
130    }
131
132    /// The number of bits the filter holds.
133    pub fn num_bits(&self) -> usize {
134        self.size_bytes() * 8
135    }
136
137    pub fn insert_many(&mut self, hashes: &[u64]) {
138        let blocks = &mut self.blocks;
139        for_each_block_mask(hashes, blocks.as_ptr().cast(), blocks.len(), |b, mask| {
140            // SAFETY: `b` is below the number of blocks.
141            let block = unsafe { &mut blocks.get_unchecked_mut(b).0 };
142            for i in 0..8 {
143                block[i] |= mask[i];
144            }
145        });
146    }
147
148    #[inline]
149    pub fn contains(&self, hash: u64) -> bool {
150        let b = block_index(hash, self.blocks.len());
151        // SAFETY: `b` is below the number of blocks.
152        let block = unsafe { self.blocks.get_unchecked(b).0 };
153        block_contains(block, block_mask(hash))
154    }
155}
156
157/// Hashes are inserted in groups of this many.
158const INSERT_GROUP: usize = 64;
159/// Filters with fewer blocks than this stay in the cache and aren't prefetched.
160const MIN_PREFETCH_BLOCKS: usize = 1 << 15;
161
162/// Call `f` with the block index and mask of every hash, for a filter of
163/// `num_blocks` blocks starting at `blocks`. The masks of a group of hashes are
164/// computed together, which vectorizes, and on a large filter the blocks of the
165/// group are fetched before any is written, so that their cache misses overlap.
166#[inline(always)]
167fn for_each_block_mask(
168    hashes: &[u64],
169    blocks: *const u8,
170    num_blocks: usize,
171    mut f: impl FnMut(usize, [u32; 8]),
172) {
173    let prefetch = num_blocks >= MIN_PREFETCH_BLOCKS;
174    let mut masks = [[0u32; 8]; INSERT_GROUP];
175    let mut idxs = [0usize; INSERT_GROUP];
176    for group in hashes.chunks(INSERT_GROUP) {
177        for ((hash, mask), b) in group.iter().zip(&mut masks).zip(&mut idxs) {
178            *mask = block_mask(*hash);
179            *b = block_index(*hash, num_blocks);
180        }
181        let idxs = &idxs[..group.len()];
182        if prefetch {
183            // A read prefetch, so that the blocks other threads share stay
184            // shared when the key's bits are set already.
185            for b in idxs {
186                prefetch_l1(blocks.wrapping_add(b * BLOCK_BYTES));
187            }
188        }
189        for (mask, b) in masks.iter().zip(idxs) {
190            f(*b, *mask);
191        }
192    }
193}
194
195/// The words of a [`Block`], two at a time, so fewer atomic operations are
196/// needed per key.
197#[repr(C, align(32))]
198struct AtomicBlock([AtomicU64; 4]);
199
200const _: () = assert!(
201    size_of::<AtomicBlock>() == size_of::<Block>()
202        && align_of::<AtomicBlock>() == align_of::<Block>()
203);
204
205/// A [`SplitBlockBloom`] that many threads insert into at once. An atomic
206/// insert costs several times a plain one, so a filter only one thread fills
207/// is better built as a [`SplitBlockBloom`].
208pub struct AtomicSplitBlockBloom {
209    blocks: Vec<AtomicBlock>,
210}
211
212impl From<SplitBlockBloom> for AtomicSplitBlockBloom {
213    fn from(bloom: SplitBlockBloom) -> Self {
214        // SAFETY: the blocks have the same size and alignment, and any bits
215        // are valid in both.
216        Self {
217            blocks: unsafe { cast_vec(bloom.blocks) },
218        }
219    }
220}
221
222impl AtomicSplitBlockBloom {
223    pub fn insert_many(&self, hashes: &[u64]) {
224        let blocks = &self.blocks;
225        for_each_block_mask(hashes, blocks.as_ptr().cast(), blocks.len(), |b, mask| {
226            // SAFETY: `b` is below the number of blocks.
227            let block = unsafe { &blocks.get_unchecked(b).0 };
228            let mask: [u64; 4] = bytemuck::cast(mask);
229            let words: [u64; 4] = std::array::from_fn(|i| block[i].load(Ordering::Relaxed));
230            let mut missing = 0;
231            for i in 0..4 {
232                missing |= mask[i] & !words[i];
233            }
234            // A block that has the bits of the key already is only read, so
235            // the blocks of keys seen before stay shared between threads.
236            if missing != 0 {
237                for i in 0..4 {
238                    if mask[i] & !words[i] != 0 {
239                        block[i].fetch_or(mask[i], Ordering::Relaxed);
240                    }
241                }
242            }
243        });
244    }
245
246    /// The filter, once no thread inserts anymore.
247    pub fn into_inner(self) -> SplitBlockBloom {
248        // SAFETY: as in `from`.
249        SplitBlockBloom {
250            blocks: unsafe { cast_vec(self.blocks) },
251        }
252    }
253}
254
255/// The allocation of `v` as a `Vec<U>`.
256///
257/// # Safety
258/// `U` must have the size and alignment of `T`, and every value of `T` must be
259/// a valid `U`.
260unsafe fn cast_vec<T, U>(v: Vec<T>) -> Vec<U> {
261    let mut v = ManuallyDrop::new(v);
262    unsafe { Vec::from_raw_parts(v.as_mut_ptr().cast(), v.len(), v.capacity()) }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    #[should_panic]
271    fn is_in_set_needs_a_block() {
272        is_in_set(&[], 0);
273    }
274
275    #[test]
276    #[should_panic]
277    fn insert_needs_a_block() {
278        insert(&mut [0u8; 31], 0);
279    }
280
281    fn spread_hashes(range: std::ops::Range<u64>) -> Vec<u64> {
282        range.map(|i| i.wrapping_mul(0x9E3779B97F4A7C15)).collect()
283    }
284
285    fn words(bloom: &SplitBlockBloom) -> Vec<[u32; 8]> {
286        bloom.blocks.iter().map(|b| b.0).collect()
287    }
288
289    /// The words of the filter of `hashes` that the bitset functions build.
290    fn bitset_words(hashes: &[u64], size_bytes: usize) -> Vec<[u32; 8]> {
291        let mut bitset = vec![0u8; size_bytes];
292        hashes.iter().for_each(|h| insert(&mut bitset, *h));
293        bitset
294            .chunks_exact(BLOCK_BYTES)
295            .map(|b| load_block(b.try_into().unwrap()))
296            .collect()
297    }
298
299    #[test]
300    fn matches_bitset() {
301        // Below and above the size that is prefetched, in batches that do and
302        // don't fill a group.
303        for keys in [1000, 1 << 20] {
304            let hashes = spread_hashes(0..keys as u64);
305            let expected = |bloom: &SplitBlockBloom| bitset_words(&hashes, bloom.size_bytes());
306            for batch in [7, 100, keys] {
307                let mut plain = SplitBlockBloom::with_capacity(keys, 8);
308                hashes.chunks(batch).for_each(|c| plain.insert_many(c));
309                assert_eq!(plain.size_bytes(), SplitBlockBloom::size_for(keys, 8));
310                assert!(words(&plain) == expected(&plain));
311
312                let atomic = AtomicSplitBlockBloom::from(SplitBlockBloom::with_capacity(keys, 8));
313                hashes.chunks(batch).for_each(|c| atomic.insert_many(c));
314                let atomic = atomic.into_inner();
315                assert!(words(&atomic) == expected(&atomic));
316            }
317        }
318    }
319
320    #[test]
321    fn contains_like_bitset() {
322        let hashes = spread_hashes(0..1000);
323        let mut bloom = SplitBlockBloom::with_capacity(1000, 8);
324        bloom.insert_many(&hashes);
325        let mut bitset = vec![0u8; bloom.size_bytes()];
326        hashes.iter().for_each(|h| insert(&mut bitset, *h));
327        assert!(hashes.iter().all(|h| bloom.contains(*h)));
328        for h in spread_hashes(1000..2000) {
329            assert_eq!(bloom.contains(h), is_in_set(&bitset, h));
330        }
331    }
332
333    #[test]
334    fn insert_from_several_threads() {
335        let hashes = spread_hashes(0..100_000);
336        let shared = AtomicSplitBlockBloom::from(SplitBlockBloom::with_capacity(hashes.len(), 8));
337        std::thread::scope(|s| {
338            for part in hashes.chunks(10_000) {
339                let shared = &shared;
340                s.spawn(move || shared.insert_many(part));
341            }
342        });
343        let bloom = shared.into_inner();
344        assert!(words(&bloom) == bitset_words(&hashes, bloom.size_bytes()));
345    }
346}
polars_utils/bloom_filter.rs

polars_utils/
bloom_filter.rs
