index/bloom_filter/
creator.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod finalize_segment;
16mod intermediate_codec;
17
18use std::collections::HashSet;
19use std::sync::Arc;
20use std::sync::atomic::{AtomicUsize, Ordering};
21
22use finalize_segment::FinalizedBloomFilterStorage;
23use futures::{AsyncWrite, AsyncWriteExt, StreamExt};
24use greptime_proto::v1::index::{BloomFilterLoc, BloomFilterMeta};
25use prost::Message;
26use snafu::ResultExt;
27
28use crate::Bytes;
29use crate::bloom_filter::error::{IoSnafu, Result};
30use crate::bloom_filter::{PrehashedBuildHasher, element_hash};
31use crate::external_provider::ExternalTempFileProvider;
32
33/// `BloomFilterCreator` is responsible for creating and managing bloom filters
34/// for a set of elements. It divides the rows into segments and creates
35/// bloom filters for each segment.
36///
37/// # Format
38///
39/// The bloom filter creator writes the following format to the writer:
40///
41/// ```text
42/// +--------------------+--------------------+-----+----------------------+----------------------+
43/// | Bloom filter 0     | Bloom filter 1     | ... | BloomFilterMeta      | Meta size            |
44/// +--------------------+--------------------+-----+----------------------+----------------------+
45/// |<- bytes (size 0) ->|<- bytes (size 1) ->| ... |<- json (meta size) ->|<- u32 LE (4 bytes) ->|
46/// ```
47///
48pub struct BloomFilterCreator {
49    /// The number of rows per segment set by the user.
50    rows_per_segment: usize,
51
52    /// Row count that added to the bloom filter so far.
53    accumulated_row_count: usize,
54
55    /// Distinct element hashes (see [`element_hash`]) in the current segment.
56    ///
57    /// Elements with equal hashes set the same bits, so deduplicating by hash loses
58    /// nothing and avoids copying the values.
59    cur_seg_distinct_elems: HashSet<u64, PrehashedBuildHasher>,
60
61    /// The memory usage of the current segment's distinct elements.
62    cur_seg_distinct_elems_mem_usage: usize,
63
64    /// Storage for finalized Bloom filters.
65    finalized_bloom_filters: FinalizedBloomFilterStorage,
66
67    /// Row count that finalized so far.
68    finalized_row_count: usize,
69
70    /// Global memory usage of the bloom filter creator.
71    global_memory_usage: Arc<AtomicUsize>,
72}
73
74impl BloomFilterCreator {
75    /// Creates a new `BloomFilterCreator` with the specified number of rows per segment.
76    ///
77    /// # PANICS
78    ///
79    /// `rows_per_segment` <= 0
80    pub fn new(
81        rows_per_segment: usize,
82        false_positive_rate: f64,
83        intermediate_provider: Arc<dyn ExternalTempFileProvider>,
84        global_memory_usage: Arc<AtomicUsize>,
85        global_memory_usage_threshold: Option<usize>,
86    ) -> Self {
87        assert!(
88            rows_per_segment > 0,
89            "rows_per_segment must be greater than 0"
90        );
91
92        Self {
93            rows_per_segment,
94            accumulated_row_count: 0,
95            cur_seg_distinct_elems: HashSet::default(),
96            cur_seg_distinct_elems_mem_usage: 0,
97            global_memory_usage: global_memory_usage.clone(),
98            finalized_bloom_filters: FinalizedBloomFilterStorage::new(
99                false_positive_rate,
100                intermediate_provider,
101                global_memory_usage,
102                global_memory_usage_threshold,
103            ),
104            finalized_row_count: 0,
105        }
106    }
107
108    /// Adds multiple rows of elements to the bloom filter. If the number of accumulated rows
109    /// reaches `rows_per_segment`, it finalizes the current segment.
110    pub async fn push_n_row_elems(
111        &mut self,
112        nrows: usize,
113        elems: impl IntoIterator<Item = Bytes>,
114    ) -> Result<()> {
115        let hashes = elems
116            .into_iter()
117            .map(|e| element_hash(&e))
118            .collect::<Vec<_>>();
119        self.push_n_row_hashes(nrows, &hashes).await
120    }
121
122    /// Adds `nrows` copies of a single borrowed value (or null). Row counts advance for
123    /// nulls as well.
124    pub async fn push_n_row_elem(&mut self, nrows: usize, elem: Option<&[u8]>) -> Result<()> {
125        match elem {
126            Some(elem) => self.push_n_row_hashes(nrows, &[element_hash(elem)]).await,
127            None => self.push_n_row_hashes(nrows, &[]).await,
128        }
129    }
130
131    /// Adds a row of elements to the bloom filter. If the number of accumulated rows
132    /// reaches `rows_per_segment`, it finalizes the current segment.
133    pub async fn push_row_elems(&mut self, elems: impl IntoIterator<Item = Bytes>) -> Result<()> {
134        self.push_row_hashes(elems.into_iter().map(|e| element_hash(&e)))
135            .await
136    }
137
138    /// Adds a row of element hashes computed by [`element_hash`].
139    pub async fn push_row_hashes(&mut self, hashes: impl IntoIterator<Item = u64>) -> Result<()> {
140        self.accumulated_row_count += 1;
141        self.insert_hashes(hashes);
142
143        if self
144            .accumulated_row_count
145            .is_multiple_of(self.rows_per_segment)
146        {
147            self.finalize_segment().await?;
148            self.finalized_row_count = self.accumulated_row_count;
149        }
150
151        Ok(())
152    }
153
154    /// Adds `nrows` rows that all contain the element hashes in `hashes`.
155    pub async fn push_n_row_hashes(&mut self, mut nrows: usize, hashes: &[u64]) -> Result<()> {
156        while nrows > 0 {
157            let rows_to_seg_end =
158                self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment);
159            let rows_to_push = nrows.min(rows_to_seg_end);
160            nrows -= rows_to_push;
161            self.accumulated_row_count += rows_to_push;
162            self.insert_hashes(hashes.iter().copied());
163
164            if self
165                .accumulated_row_count
166                .is_multiple_of(self.rows_per_segment)
167            {
168                self.finalize_segment().await?;
169                self.finalized_row_count = self.accumulated_row_count;
170            }
171        }
172
173        Ok(())
174    }
175
176    fn insert_hashes(&mut self, hashes: impl IntoIterator<Item = u64>) {
177        let old_len = self.cur_seg_distinct_elems.len();
178        // Not `extend`: it reserves for the iterator's length, which counts duplicate
179        // tokens, and the capacity survives `drain` at segment boundaries.
180        for hash in hashes {
181            self.cur_seg_distinct_elems.insert(hash);
182        }
183        let mem_diff = (self.cur_seg_distinct_elems.len() - old_len) * size_of::<u64>();
184        if mem_diff > 0 {
185            self.cur_seg_distinct_elems_mem_usage += mem_diff;
186            self.global_memory_usage
187                .fetch_add(mem_diff, Ordering::Relaxed);
188        }
189    }
190
191    /// Finalizes any remaining segments and writes the bloom filters and metadata to the provided writer.
192    pub async fn finish(&mut self, mut writer: impl AsyncWrite + Unpin) -> Result<()> {
193        if self.accumulated_row_count > self.finalized_row_count {
194            self.finalize_segment().await?;
195        }
196
197        let mut meta = BloomFilterMeta {
198            rows_per_segment: self.rows_per_segment as _,
199            row_count: self.accumulated_row_count as _,
200            ..Default::default()
201        };
202
203        let (indices, mut segs) = self.finalized_bloom_filters.drain().await?;
204        meta.segment_loc_indices = indices.into_iter().map(|i| i as u64).collect();
205        meta.segment_count = meta.segment_loc_indices.len() as _;
206
207        while let Some(segment) = segs.next().await {
208            let segment = segment?;
209            writer
210                .write_all(&segment.bloom_filter_bytes)
211                .await
212                .context(IoSnafu)?;
213
214            let size = segment.bloom_filter_bytes.len() as u64;
215            meta.bloom_filter_locs.push(BloomFilterLoc {
216                offset: meta.bloom_filter_size as _,
217                size,
218                element_count: segment.element_count as _,
219            });
220            meta.bloom_filter_size += size;
221        }
222
223        let meta_bytes = meta.encode_to_vec();
224        writer.write_all(&meta_bytes).await.context(IoSnafu)?;
225
226        let meta_size = meta_bytes.len() as u32;
227        writer
228            .write_all(&meta_size.to_le_bytes())
229            .await
230            .context(IoSnafu)?;
231        writer.flush().await.unwrap();
232
233        Ok(())
234    }
235
236    /// Returns the memory usage of the creating bloom filter.
237    pub fn memory_usage(&self) -> usize {
238        self.cur_seg_distinct_elems_mem_usage + self.finalized_bloom_filters.memory_usage()
239    }
240
241    async fn finalize_segment(&mut self) -> Result<()> {
242        let elem_count = self.cur_seg_distinct_elems.len();
243        self.finalized_bloom_filters
244            .add(self.cur_seg_distinct_elems.drain(), elem_count)
245            .await?;
246
247        self.global_memory_usage
248            .fetch_sub(self.cur_seg_distinct_elems_mem_usage, Ordering::Relaxed);
249        self.cur_seg_distinct_elems_mem_usage = 0;
250        Ok(())
251    }
252}
253
254impl Drop for BloomFilterCreator {
255    fn drop(&mut self) {
256        self.global_memory_usage
257            .fetch_sub(self.cur_seg_distinct_elems_mem_usage, Ordering::Relaxed);
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use fastbloom::BloomFilter;
264    use futures::io::Cursor;
265
266    use super::*;
267    use crate::bloom_filter::SEED;
268    use crate::external_provider::MockExternalTempFileProvider;
269
270    /// Converts a slice of bytes to a vector of `u64`.
271    pub fn u64_vec_from_bytes(bytes: &[u8]) -> Vec<u64> {
272        bytes
273            .as_chunks::<{ std::mem::size_of::<u64>() }>()
274            .0
275            .iter()
276            .map(|chunk| u64::from_le_bytes(*chunk))
277            .collect()
278    }
279
280    #[tokio::test]
281    async fn test_duplicate_hashes_do_not_grow_segment_set() {
282        let mut creator = BloomFilterCreator::new(
283            4,
284            0.01,
285            Arc::new(MockExternalTempFileProvider::new()),
286            Arc::new(AtomicUsize::new(0)),
287            None,
288        );
289        creator
290            .push_row_hashes(std::iter::repeat_n(7, 1_000_000))
291            .await
292            .unwrap();
293        assert_eq!(creator.cur_seg_distinct_elems.len(), 1);
294        assert!(creator.cur_seg_distinct_elems.capacity() < 16);
295    }
296
297    #[tokio::test]
298    async fn test_bloom_filter_creator() {
299        let mut writer = Cursor::new(Vec::new());
300        let mut creator = BloomFilterCreator::new(
301            2,
302            0.01,
303            Arc::new(MockExternalTempFileProvider::new()),
304            Arc::new(AtomicUsize::new(0)),
305            None,
306        );
307
308        creator
309            .push_row_elems(vec![b"a".to_vec(), b"b".to_vec()])
310            .await
311            .unwrap();
312        assert!(creator.cur_seg_distinct_elems_mem_usage > 0);
313        assert!(creator.memory_usage() > 0);
314
315        creator
316            .push_row_elems(vec![b"c".to_vec(), b"d".to_vec()])
317            .await
318            .unwrap();
319        // Finalize the first segment
320        assert_eq!(creator.cur_seg_distinct_elems_mem_usage, 0);
321        assert!(creator.memory_usage() > 0);
322
323        creator
324            .push_row_elems(vec![b"e".to_vec(), b"f".to_vec()])
325            .await
326            .unwrap();
327        assert!(creator.cur_seg_distinct_elems_mem_usage > 0);
328        assert!(creator.memory_usage() > 0);
329
330        creator.finish(&mut writer).await.unwrap();
331
332        let bytes = writer.into_inner();
333        let total_size = bytes.len();
334        let meta_size_offset = total_size - 4;
335        let meta_size = u32::from_le_bytes((&bytes[meta_size_offset..]).try_into().unwrap());
336
337        let meta_bytes = &bytes[total_size - meta_size as usize - 4..total_size - 4];
338        let meta = BloomFilterMeta::decode(meta_bytes).unwrap();
339
340        assert_eq!(meta.rows_per_segment, 2);
341        assert_eq!(meta.segment_count, 2);
342        assert_eq!(meta.row_count, 3);
343        assert_eq!(
344            meta.bloom_filter_size as usize + meta_bytes.len() + 4,
345            total_size
346        );
347
348        let mut bfs = Vec::new();
349        for segment in meta.bloom_filter_locs {
350            let bloom_filter_bytes =
351                &bytes[segment.offset as usize..(segment.offset + segment.size) as usize];
352            let v = u64_vec_from_bytes(bloom_filter_bytes);
353            let bloom_filter = BloomFilter::from_vec(v)
354                .seed(&SEED)
355                .expected_items(segment.element_count as usize);
356            bfs.push(bloom_filter);
357        }
358
359        assert_eq!(meta.segment_loc_indices.len(), 2);
360
361        let bf0 = &bfs[meta.segment_loc_indices[0] as usize];
362        assert!(bf0.contains(&b"a"));
363        assert!(bf0.contains(&b"b"));
364        assert!(bf0.contains(&b"c"));
365        assert!(bf0.contains(&b"d"));
366
367        let bf1 = &bfs[meta.segment_loc_indices[1] as usize];
368        assert!(bf1.contains(&b"e"));
369        assert!(bf1.contains(&b"f"));
370    }
371
372    #[tokio::test]
373    async fn test_bloom_filter_creator_batch_push() {
374        let mut writer = Cursor::new(Vec::new());
375        let mut creator: BloomFilterCreator = BloomFilterCreator::new(
376            2,
377            0.01,
378            Arc::new(MockExternalTempFileProvider::new()),
379            Arc::new(AtomicUsize::new(0)),
380            None,
381        );
382
383        creator
384            .push_n_row_elems(5, vec![b"a".to_vec(), b"b".to_vec()])
385            .await
386            .unwrap();
387        assert!(creator.cur_seg_distinct_elems_mem_usage > 0);
388        assert!(creator.memory_usage() > 0);
389
390        creator
391            .push_n_row_elems(5, vec![b"c".to_vec(), b"d".to_vec()])
392            .await
393            .unwrap();
394        assert_eq!(creator.cur_seg_distinct_elems_mem_usage, 0);
395        assert!(creator.memory_usage() > 0);
396
397        creator
398            .push_n_row_elems(10, vec![b"e".to_vec(), b"f".to_vec()])
399            .await
400            .unwrap();
401        assert_eq!(creator.cur_seg_distinct_elems_mem_usage, 0);
402        assert!(creator.memory_usage() > 0);
403
404        creator.finish(&mut writer).await.unwrap();
405
406        let bytes = writer.into_inner();
407        let total_size = bytes.len();
408        let meta_size_offset = total_size - 4;
409        let meta_size = u32::from_le_bytes((&bytes[meta_size_offset..]).try_into().unwrap());
410
411        let meta_bytes = &bytes[total_size - meta_size as usize - 4..total_size - 4];
412        let meta = BloomFilterMeta::decode(meta_bytes).unwrap();
413
414        assert_eq!(meta.rows_per_segment, 2);
415        assert_eq!(meta.segment_count, 10);
416        assert_eq!(meta.row_count, 20);
417        assert_eq!(
418            meta.bloom_filter_size as usize + meta_bytes.len() + 4,
419            total_size
420        );
421
422        let mut bfs = Vec::new();
423        for segment in meta.bloom_filter_locs {
424            let bloom_filter_bytes =
425                &bytes[segment.offset as usize..(segment.offset + segment.size) as usize];
426            let v = u64_vec_from_bytes(bloom_filter_bytes);
427            let bloom_filter = BloomFilter::from_vec(v)
428                .seed(&SEED)
429                .expected_items(segment.element_count as _);
430            bfs.push(bloom_filter);
431        }
432
433        // 4 bloom filters to serve 10 segments
434        assert_eq!(bfs.len(), 4);
435        assert_eq!(meta.segment_loc_indices.len(), 10);
436
437        for idx in meta.segment_loc_indices.iter().take(3) {
438            let bf = &bfs[*idx as usize];
439            assert!(bf.contains(&b"a"));
440            assert!(bf.contains(&b"b"));
441        }
442        for idx in meta.segment_loc_indices.iter().take(5).skip(2) {
443            let bf = &bfs[*idx as usize];
444            assert!(bf.contains(&b"c"));
445            assert!(bf.contains(&b"d"));
446        }
447        for idx in meta.segment_loc_indices.iter().take(10).skip(5) {
448            let bf = &bfs[*idx as usize];
449            assert!(bf.contains(&b"e"));
450            assert!(bf.contains(&b"f"));
451        }
452    }
453
454    #[tokio::test]
455    async fn borrowed_single_value_matches_owned_rows_across_segments() {
456        let make_creator = || {
457            BloomFilterCreator::new(
458                3,
459                0.01,
460                Arc::new(MockExternalTempFileProvider::new()),
461                Arc::new(AtomicUsize::new(0)),
462                None,
463            )
464        };
465        let mut borrowed = make_creator();
466        let mut owned = make_creator();
467        // Zero rows, nulls, empty values, duplicates and runs crossing segment boundaries.
468        for (rows, elem) in [
469            (0, Some(b"ignored".as_slice())),
470            (1, None),
471            (5, Some(b"".as_slice())),
472            (1, Some(b"".as_slice())),
473            (8, Some(b"label".as_slice())),
474            (1, None),
475        ] {
476            borrowed.push_n_row_elem(rows, elem).await.unwrap();
477            owned
478                .push_n_row_elems(rows, elem.map(<[u8]>::to_vec))
479                .await
480                .unwrap();
481            assert_eq!(borrowed.memory_usage(), owned.memory_usage());
482        }
483        let mut borrowed_blob = Cursor::new(Vec::new());
484        let mut owned_blob = Cursor::new(Vec::new());
485        borrowed.finish(&mut borrowed_blob).await.unwrap();
486        owned.finish(&mut owned_blob).await.unwrap();
487        assert_eq!(borrowed_blob.into_inner(), owned_blob.into_inner());
488    }
489
490    #[tokio::test]
491    async fn test_final_seg_all_null() {
492        let mut writer = Cursor::new(Vec::new());
493        let mut creator = BloomFilterCreator::new(
494            2,
495            0.01,
496            Arc::new(MockExternalTempFileProvider::new()),
497            Arc::new(AtomicUsize::new(0)),
498            None,
499        );
500
501        creator
502            .push_n_row_elems(4, vec![b"a".to_vec(), b"b".to_vec()])
503            .await
504            .unwrap();
505        creator.push_row_elems(Vec::new()).await.unwrap();
506
507        creator.finish(&mut writer).await.unwrap();
508
509        let bytes = writer.into_inner();
510        let total_size = bytes.len();
511        let meta_size_offset = total_size - 4;
512        let meta_size = u32::from_le_bytes((&bytes[meta_size_offset..]).try_into().unwrap());
513
514        let meta_bytes = &bytes[total_size - meta_size as usize - 4..total_size - 4];
515        let meta = BloomFilterMeta::decode(meta_bytes).unwrap();
516
517        assert_eq!(meta.rows_per_segment, 2);
518        assert_eq!(meta.segment_count, 3);
519        assert_eq!(meta.row_count, 5);
520    }
521}
index/bloom_filter/creator.rs

index/bloom_filter/
creator.rs
