feat: encode and decode BlockedBloomFilter

Author: k-jingyang

benches/bloom.rs

@@ -63,7 +63,7 @@ fn blocked_filter_construction(c: &mut Criterion) {
 }
 
 fn blocked_filter_contains(c: &mut Criterion) {
-    use lsm_tree::segment::filter::blocked_bloom::Builder;
+    use lsm_tree::segment::filter::{blocked_bloom::Builder, AMQFilter};
 
     let keys = (0..100_000u128)
         .map(|x| x.to_be_bytes().to_vec())

src/segment/filter/blocked_bloom/mod.rs

@@ -3,8 +3,14 @@
 // (found in the LICENSE-* files in the repository)
 
 mod builder;
-use super::{bit_array::BitArrayReader, CACHE_LINE_BYTES};
+use super::{bit_array::BitArrayReader, AMQFilter, CACHE_LINE_BYTES};
+use crate::{
+    coding::{DecodeError, Encode, EncodeError},
+    file::MAGIC_BYTES,
+};
 pub use builder::Builder;
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use std::io::{Read, Write};
 
 /// Two hashes that are used for double hashing
 pub type CompositeHash = (u64, u64);
@@ -20,29 +26,22 @@ pub struct BlockedBloomFilter {
     num_blocks: usize,
 }
 
-// TODO: Implement Encode and Decode for BlockedBloomFilter
+impl AMQFilter for BlockedBloomFilter {
+    fn bytes(&self) -> &[u8] {
+        self.inner.bytes()
+    }
 
-impl BlockedBloomFilter {
     /// Size of bloom filter in bytes
     #[must_use]
-    pub fn len(&self) -> usize {
+    fn len(&self) -> usize {
         self.inner.bytes().len()
     }
 
-    fn from_raw(m: usize, k: usize, slice: crate::Slice) -> Self {
-        let num_blocks = m.div_ceil(CACHE_LINE_BYTES);
-        Self {
-            inner: BitArrayReader::new(slice),
-            k,
-            num_blocks,
-        }
-    }
-
     /// Returns `true` if the hash may be contained.
     ///
     /// Will never have a false negative.
     #[must_use]
-    pub fn contains_hash(&self, (mut h1, mut h2): CompositeHash) -> bool {
+    fn contains_hash(&self, (mut h1, mut h2): CompositeHash) -> bool {
         let block_idx = h1 % (self.num_blocks as u64);
 
         for i in 1..(self.k as u64) {
@@ -65,10 +64,56 @@ impl BlockedBloomFilter {
     ///
     /// Will never have a false negative.
     #[must_use]
-    pub fn contains(&self, key: &[u8]) -> bool {
+    fn contains(&self, key: &[u8]) -> bool {
         self.contains_hash(Self::get_hash(key))
     }
 
+    fn filter_type(&self) -> super::FilterType {
+        super::FilterType::BlockedBloom
+    }
+}
+
+impl Encode for BlockedBloomFilter {
+    fn encode_into<W: Write>(&self, writer: &mut W) -> Result<(), EncodeError> {
+        // Write header
+        writer.write_all(&MAGIC_BYTES)?;
+
+        writer.write_u8(super::FilterType::BlockedBloom as u8)?;
+
+        // NOTE: Hash type (unused)
+        writer.write_u8(0)?;
+
+        writer.write_u64::<LittleEndian>(self.num_blocks as u64)?;
+        writer.write_u64::<LittleEndian>(self.k as u64)?;
+        writer.write_all(self.inner.bytes())?;
+
+        Ok(())
+    }
+}
+
+impl BlockedBloomFilter {
+    // To be used by AMQFilter after magic bytes and filter type have been read and parsed
+    pub(super) fn decode_from<R: Read>(reader: &mut R) -> Result<Self, DecodeError> {
+        // NOTE: Hash type (unused)
+        let hash_type = reader.read_u8()?;
+        assert_eq!(0, hash_type, "Invalid bloom hash type");
+
+        let num_blocks = reader.read_u64::<LittleEndian>()? as usize;
+        let k = reader.read_u64::<LittleEndian>()? as usize;
+
+        let mut bytes = vec![0; num_blocks * CACHE_LINE_BYTES];
+        reader.read_exact(&mut bytes)?;
+
+        Ok(Self::from_raw(num_blocks, k, bytes.into()))
+    }
+
+    fn from_raw(num_blocks: usize, k: usize, slice: crate::Slice) -> Self {
+        Self {
+            inner: BitArrayReader::new(slice),
+            k,
+            num_blocks,
+        }
+    }
     /// Returns `true` if the bit at `idx` is `1`.
     fn has_bit(&self, block_idx: usize, idx_in_block: usize) -> bool {
         self.inner
@@ -84,6 +129,57 @@ impl BlockedBloomFilter {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::segment::filter::{AMQFilterBuilder, FilterType};
+
+    use std::fs::File;
+    use test_log::test;
+
+    #[test]
+    fn blocked_bloom_serde_round_trip() -> crate::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let path = dir.path().join("bf");
+        let mut file = File::create(&path)?;
+
+        let mut filter = Builder::with_fp_rate(10, 0.0001);
+
+        let keys = &[
+            b"item0", b"item1", b"item2", b"item3", b"item4", b"item5", b"item6", b"item7",
+            b"item8", b"item9",
+        ];
+
+        for key in keys {
+            filter.set_with_hash(BlockedBloomFilter::get_hash(*key));
+        }
+
+        let filter = filter.build();
+
+        for key in keys {
+            assert!(filter.contains(&**key));
+        }
+        assert!(!filter.contains(b"asdasads"));
+        assert!(!filter.contains(b"item10"));
+        assert!(!filter.contains(b"cxycxycxy"));
+
+        filter.encode_into(&mut file)?;
+        file.sync_all()?;
+        drop(file);
+
+        let mut file = File::open(&path)?;
+        let filter_copy = AMQFilterBuilder::decode_from(&mut file)?;
+
+        assert_eq!(filter.inner.bytes(), filter_copy.bytes());
+        assert_eq!(FilterType::BlockedBloom, filter_copy.filter_type());
+
+        for key in keys {
+            assert!(filter.contains(&**key));
+        }
+        assert!(!filter_copy.contains(b"asdasads"));
+        assert!(!filter_copy.contains(b"item10"));
+        assert!(!filter_copy.contains(b"cxycxycxy"));
+
+        Ok(())
+    }
 
     #[test]
     fn blocked_bloom_basic() {

src/segment/filter/mod.rs

@@ -7,6 +7,7 @@ pub mod blocked_bloom;
 pub mod standard_bloom;
 
 use crate::{coding::DecodeError, file::MAGIC_BYTES};
+use blocked_bloom::BlockedBloomFilter;
 use byteorder::ReadBytesExt;
 use std::io::Read;
 
@@ -46,7 +47,8 @@ impl BloomConstructionPolicy {
     }
 }
 
-enum FilterType {
+#[derive(PartialEq, Debug)]
+pub enum FilterType {
     StandardBloom = 0,
     BlockedBloom = 1,
 }
@@ -67,6 +69,7 @@ pub trait AMQFilter: Sync + Send {
     fn len(&self) -> usize;
     fn contains(&self, item: &[u8]) -> bool;
     fn contains_hash(&self, hash: (u64, u64)) -> bool;
+    fn filter_type(&self) -> FilterType;
 }
 
 pub struct AMQFilterBuilder {}
@@ -89,8 +92,9 @@ impl AMQFilterBuilder {
             FilterType::StandardBloom => StandardBloomFilter::decode_from(reader)
                 .map(Self::wrap_filter)
                 .map_err(|e| DecodeError::from(e)),
-            // TODO: Implement
-            FilterType::BlockedBloom => Err(DecodeError::InvalidHeader("BlockedBloom")),
+            FilterType::BlockedBloom => BlockedBloomFilter::decode_from(reader)
+                .map(Self::wrap_filter)
+                .map_err(|e| DecodeError::from(e)),
         }
     }
 

src/segment/filter/standard_bloom/mod.rs

@@ -72,15 +72,18 @@ impl AMQFilter for StandardBloomFilter {
 
         true
     }
+
+    fn filter_type(&self) -> super::FilterType {
+        super::FilterType::StandardBloom
+    }
 }
 
 impl Encode for StandardBloomFilter {
     fn encode_into<W: Write>(&self, writer: &mut W) -> Result<(), EncodeError> {
         // Write header
         writer.write_all(&MAGIC_BYTES)?;
 
-        // NOTE: Filter type (unused)
-        writer.write_u8(0)?;
+        writer.write_u8(super::FilterType::StandardBloom as u8)?;
 
         // NOTE: Hash type (unused)
         writer.write_u8(0)?;
@@ -131,7 +134,7 @@ impl StandardBloomFilter {
 
 #[cfg(test)]
 mod tests {
-    use crate::segment::filter::AMQFilterBuilder;
+    use crate::segment::filter::{AMQFilterBuilder, FilterType};
 
     use super::*;
     use std::fs::File;
@@ -172,6 +175,7 @@ mod tests {
         let filter_copy = AMQFilterBuilder::decode_from(&mut file)?;
 
         assert_eq!(filter.inner.bytes(), filter_copy.bytes());
+        assert_eq!(FilterType::StandardBloom, filter_copy.filter_type());
 
         for key in keys {
             assert!(filter.contains(&**key));