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            .chunks_exact(std::mem::size_of::<u64>())
274            .map(|chunk| u64::from_le_bytes(chunk.try_into().unwrap()))
275            .collect()
276    }
277
278    #[tokio::test]
279    async fn test_duplicate_hashes_do_not_grow_segment_set() {
280        let mut creator = BloomFilterCreator::new(
281            4,
282            0.01,
283            Arc::new(MockExternalTempFileProvider::new()),
284            Arc::new(AtomicUsize::new(0)),
285            None,
286        );
287        creator
288            .push_row_hashes(std::iter::repeat_n(7, 1_000_000))
289            .await
290            .unwrap();
291        assert_eq!(creator.cur_seg_distinct_elems.len(), 1);
292        assert!(creator.cur_seg_distinct_elems.capacity() < 16);
293    }
294
295    #[tokio::test]
296    async fn test_bloom_filter_creator() {
297        let mut writer = Cursor::new(Vec::new());
298        let mut creator = BloomFilterCreator::new(
299            2,
300            0.01,
301            Arc::new(MockExternalTempFileProvider::new()),
302            Arc::new(AtomicUsize::new(0)),
303            None,
304        );
305
306        creator
307            .push_row_elems(vec![b"a".to_vec(), b"b".to_vec()])
308            .await
309            .unwrap();
310        assert!(creator.cur_seg_distinct_elems_mem_usage > 0);
311        assert!(creator.memory_usage() > 0);
312
313        creator
314            .push_row_elems(vec![b"c".to_vec(), b"d".to_vec()])
315            .await
316            .unwrap();
317        // Finalize the first segment
318        assert_eq!(creator.cur_seg_distinct_elems_mem_usage, 0);
319        assert!(creator.memory_usage() > 0);
320
321        creator
322            .push_row_elems(vec![b"e".to_vec(), b"f".to_vec()])
323            .await
324            .unwrap();
325        assert!(creator.cur_seg_distinct_elems_mem_usage > 0);
326        assert!(creator.memory_usage() > 0);
327
328        creator.finish(&mut writer).await.unwrap();
329
330        let bytes = writer.into_inner();
331        let total_size = bytes.len();
332        let meta_size_offset = total_size - 4;
333        let meta_size = u32::from_le_bytes((&bytes[meta_size_offset..]).try_into().unwrap());
334
335        let meta_bytes = &bytes[total_size - meta_size as usize - 4..total_size - 4];
336        let meta = BloomFilterMeta::decode(meta_bytes).unwrap();
337
338        assert_eq!(meta.rows_per_segment, 2);
339        assert_eq!(meta.segment_count, 2);
340        assert_eq!(meta.row_count, 3);
341        assert_eq!(
342            meta.bloom_filter_size as usize + meta_bytes.len() + 4,
343            total_size
344        );
345
346        let mut bfs = Vec::new();
347        for segment in meta.bloom_filter_locs {
348            let bloom_filter_bytes =
349                &bytes[segment.offset as usize..(segment.offset + segment.size) as usize];
350            let v = u64_vec_from_bytes(bloom_filter_bytes);
351            let bloom_filter = BloomFilter::from_vec(v)
352                .seed(&SEED)
353                .expected_items(segment.element_count as usize);
354            bfs.push(bloom_filter);
355        }
356
357        assert_eq!(meta.segment_loc_indices.len(), 2);
358
359        let bf0 = &bfs[meta.segment_loc_indices[0] as usize];
360        assert!(bf0.contains(&b"a"));
361        assert!(bf0.contains(&b"b"));
362        assert!(bf0.contains(&b"c"));
363        assert!(bf0.contains(&b"d"));
364
365        let bf1 = &bfs[meta.segment_loc_indices[1] as usize];
366        assert!(bf1.contains(&b"e"));
367        assert!(bf1.contains(&b"f"));
368    }
369
370    #[tokio::test]
371    async fn test_bloom_filter_creator_batch_push() {
372        let mut writer = Cursor::new(Vec::new());
373        let mut creator: BloomFilterCreator = BloomFilterCreator::new(
374            2,
375            0.01,
376            Arc::new(MockExternalTempFileProvider::new()),
377            Arc::new(AtomicUsize::new(0)),
378            None,
379        );
380
381        creator
382            .push_n_row_elems(5, vec![b"a".to_vec(), b"b".to_vec()])
383            .await
384            .unwrap();
385        assert!(creator.cur_seg_distinct_elems_mem_usage > 0);
386        assert!(creator.memory_usage() > 0);
387
388        creator
389            .push_n_row_elems(5, vec![b"c".to_vec(), b"d".to_vec()])
390            .await
391            .unwrap();
392        assert_eq!(creator.cur_seg_distinct_elems_mem_usage, 0);
393        assert!(creator.memory_usage() > 0);
394
395        creator
396            .push_n_row_elems(10, vec![b"e".to_vec(), b"f".to_vec()])
397            .await
398            .unwrap();
399        assert_eq!(creator.cur_seg_distinct_elems_mem_usage, 0);
400        assert!(creator.memory_usage() > 0);
401
402        creator.finish(&mut writer).await.unwrap();
403
404        let bytes = writer.into_inner();
405        let total_size = bytes.len();
406        let meta_size_offset = total_size - 4;
407        let meta_size = u32::from_le_bytes((&bytes[meta_size_offset..]).try_into().unwrap());
408
409        let meta_bytes = &bytes[total_size - meta_size as usize - 4..total_size - 4];
410        let meta = BloomFilterMeta::decode(meta_bytes).unwrap();
411
412        assert_eq!(meta.rows_per_segment, 2);
413        assert_eq!(meta.segment_count, 10);
414        assert_eq!(meta.row_count, 20);
415        assert_eq!(
416            meta.bloom_filter_size as usize + meta_bytes.len() + 4,
417            total_size
418        );
419
420        let mut bfs = Vec::new();
421        for segment in meta.bloom_filter_locs {
422            let bloom_filter_bytes =
423                &bytes[segment.offset as usize..(segment.offset + segment.size) as usize];
424            let v = u64_vec_from_bytes(bloom_filter_bytes);
425            let bloom_filter = BloomFilter::from_vec(v)
426                .seed(&SEED)
427                .expected_items(segment.element_count as _);
428            bfs.push(bloom_filter);
429        }
430
431        // 4 bloom filters to serve 10 segments
432        assert_eq!(bfs.len(), 4);
433        assert_eq!(meta.segment_loc_indices.len(), 10);
434
435        for idx in meta.segment_loc_indices.iter().take(3) {
436            let bf = &bfs[*idx as usize];
437            assert!(bf.contains(&b"a"));
438            assert!(bf.contains(&b"b"));
439        }
440        for idx in meta.segment_loc_indices.iter().take(5).skip(2) {
441            let bf = &bfs[*idx as usize];
442            assert!(bf.contains(&b"c"));
443            assert!(bf.contains(&b"d"));
444        }
445        for idx in meta.segment_loc_indices.iter().take(10).skip(5) {
446            let bf = &bfs[*idx as usize];
447            assert!(bf.contains(&b"e"));
448            assert!(bf.contains(&b"f"));
449        }
450    }
451
452    #[tokio::test]
453    async fn borrowed_single_value_matches_owned_rows_across_segments() {
454        let make_creator = || {
455            BloomFilterCreator::new(
456                3,
457                0.01,
458                Arc::new(MockExternalTempFileProvider::new()),
459                Arc::new(AtomicUsize::new(0)),
460                None,
461            )
462        };
463        let mut borrowed = make_creator();
464        let mut owned = make_creator();
465        // Zero rows, nulls, empty values, duplicates and runs crossing segment boundaries.
466        for (rows, elem) in [
467            (0, Some(b"ignored".as_slice())),
468            (1, None),
469            (5, Some(b"".as_slice())),
470            (1, Some(b"".as_slice())),
471            (8, Some(b"label".as_slice())),
472            (1, None),
473        ] {
474            borrowed.push_n_row_elem(rows, elem).await.unwrap();
475            owned
476                .push_n_row_elems(rows, elem.map(<[u8]>::to_vec))
477                .await
478                .unwrap();
479            assert_eq!(borrowed.memory_usage(), owned.memory_usage());
480        }
481        let mut borrowed_blob = Cursor::new(Vec::new());
482        let mut owned_blob = Cursor::new(Vec::new());
483        borrowed.finish(&mut borrowed_blob).await.unwrap();
484        owned.finish(&mut owned_blob).await.unwrap();
485        assert_eq!(borrowed_blob.into_inner(), owned_blob.into_inner());
486    }
487
488    #[tokio::test]
489    async fn test_final_seg_all_null() {
490        let mut writer = Cursor::new(Vec::new());
491        let mut creator = BloomFilterCreator::new(
492            2,
493            0.01,
494            Arc::new(MockExternalTempFileProvider::new()),
495            Arc::new(AtomicUsize::new(0)),
496            None,
497        );
498
499        creator
500            .push_n_row_elems(4, vec![b"a".to_vec(), b"b".to_vec()])
501            .await
502            .unwrap();
503        creator.push_row_elems(Vec::new()).await.unwrap();
504
505        creator.finish(&mut writer).await.unwrap();
506
507        let bytes = writer.into_inner();
508        let total_size = bytes.len();
509        let meta_size_offset = total_size - 4;
510        let meta_size = u32::from_le_bytes((&bytes[meta_size_offset..]).try_into().unwrap());
511
512        let meta_bytes = &bytes[total_size - meta_size as usize - 4..total_size - 4];
513        let meta = BloomFilterMeta::decode(meta_bytes).unwrap();
514
515        assert_eq!(meta.rows_per_segment, 2);
516        assert_eq!(meta.segment_count, 3);
517        assert_eq!(meta.row_count, 5);
518    }
519}
index/bloom_filter/creator.rs

index/bloom_filter/
creator.rs
