saves/loads elements_added on BloomFilters during export

barrust · Aug 6, 2020 · c64ddf3 · c64ddf3
1 parent ef4d8ab
commit c64ddf3
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/probables/blooms/expandingbloom.py b/probables/blooms/expandingbloom.py
@@ -150,8 +150,8 @@ def export(self, filepath):
         with open(filepath, 'wb') as fileobj:
             # add all the different Bloom bit arrays...
             for blm in self._blooms:
-                rep = 'B' * blm.bloom_length
-                fileobj.write(pack(rep, *blm.bloom))
+                rep = 'Q' + 'B' * blm.bloom_length
+                fileobj.write(pack(rep, blm.elements_added, *blm.bloom))
             fileobj.write(pack('QQQf', len(self._blooms),
                                self.estimated_elements,
                                self.elements_added,
@@ -175,9 +175,11 @@ def __load(self, filename):
                                   false_positive_rate=self.__fpr,
                                   hash_function=self.__hash_func)
                 # now we need to read in the correct number of bytes...
-                offset = calcsize('B') * blm.bloom_length
-                rep = 'B' * blm.bloom_length
-                blm._bloom = list(unpack(rep, fileobj.read(offset)))
+                offset = calcsize('Q') + calcsize('B') * blm.bloom_length
+                rep = 'Q' + 'B' * blm.bloom_length
+                elements_added, *bloom_bytes = list(unpack(rep, fileobj.read(offset)))
+                blm._bloom = bloom_bytes
+                blm.elements_added = elements_added
                 self._blooms.append(blm)
 
 

diff --git a/tests/expandingbloom_test.py b/tests/expandingbloom_test.py
@@ -81,14 +81,14 @@ def test_ebf_export(self):
         ''' basic expanding Bloom Filter export test '''
         blm = ExpandingBloomFilter(est_elements=25, false_positive_rate=0.05)
         blm.export('test.ebf')
-        self.assertEqual(calc_file_md5('test.ebf'), '1581beab91f83b7e5aaf0f059ee94eaf')
+        self.assertEqual(calc_file_md5('test.ebf'), 'eb5769ae9babdf7b37d6ce64d58812bc')
         os.remove('test.ebf')
 
     def test_ebf_import_empty(self):
         ''' test that expanding Bloom Filter is correct on import '''
         blm = ExpandingBloomFilter(est_elements=25, false_positive_rate=0.05)
         blm.export('test.ebf')
-        self.assertEqual(calc_file_md5('test.ebf'), '1581beab91f83b7e5aaf0f059ee94eaf')
+        self.assertEqual(calc_file_md5('test.ebf'), 'eb5769ae9babdf7b37d6ce64d58812bc')
 
         blm2 = ExpandingBloomFilter(filepath='test.ebf')
         for bloom in blm2._blooms:
@@ -217,14 +217,14 @@ def test_rfb_basic_export(self):
         ''' basic rotating Bloom Filter export test '''
         blm = RotatingBloomFilter(est_elements=25, false_positive_rate=0.05)
         blm.export('test.rbf')
-        self.assertEqual(calc_file_md5('test.rbf'), '1581beab91f83b7e5aaf0f059ee94eaf')
+        self.assertEqual(calc_file_md5('test.rbf'), 'eb5769ae9babdf7b37d6ce64d58812bc')
         os.remove('test.rbf')
 
     def test_rbf_import_empty(self):
         ''' test that rotating Bloom Filter is correct on import '''
         blm = RotatingBloomFilter(est_elements=25, false_positive_rate=0.05)
         blm.export('test.rbf')
-        self.assertEqual(calc_file_md5('test.rbf'), '1581beab91f83b7e5aaf0f059ee94eaf')
+        self.assertEqual(calc_file_md5('test.rbf'), 'eb5769ae9babdf7b37d6ce64d58812bc')
 
         blm2 = ExpandingBloomFilter(filepath='test.rbf')
         for bloom in blm2._blooms: