Fix delete compact bug

2 files changed, 9 insertions, 10 deletions

diff --git a/log_db/benches/benchmark.rs b/log_db/benches/benchmark.rs
index 0f72aed..dc853b6 100644
--- a/log_db/benches/benchmark.rs
+++ b/log_db/benches/benchmark.rs
@@ -1,5 +1,7 @@
 mod utils;
 
+use std::collections::HashSet;
+
 use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
 use log_db::*;
 use tempfile;
@@ -48,14 +50,16 @@ pub fn delete_existing_compacted(c: &mut Criterion) {
     for size in [100_000, 1_000_000, 10_000_000] {
         println!("Prefilling DB to {} entries", size);
         prefill_db(&mut db, &mut insts, size, true).expect("Failed to prefill DB");
-        let mut inst_it = insts.iter();
+
+        let pk_set: HashSet<i64> = insts.iter().map(|inst| inst.id).collect();
+        let mut pk_set_it = pk_set.into_iter();
 
         group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| {
             b.iter(|| {
                 let result = db
-                    .delete(black_box(&Value::Int(inst_it.next().unwrap().id)))
+                    .delete(black_box(&Value::Int(pk_set_it.next().unwrap())))
                     .unwrap();
-                assert!(result.is_some())
+                assert!(result.is_some());
             });
         });
     }
diff --git a/log_db/src/lib.rs b/log_db/src/lib.rs
index 918f5e1..928da8d 100644
--- a/log_db/src/lib.rs
+++ b/log_db/src/lib.rs
@@ -813,7 +813,7 @@ impl<R: Recordable> DB<R> {
         let active_num = parse_segment_number(&active_target)?;
 
         debug!("Reading segment data into a BTreeMap");
-        let mut pk_to_item_map = BTreeMap::new();
+        let mut pk_to_item_map: BTreeMap<&IndexableValue, &Record> = BTreeMap::new();
         let forward_read_items: Vec<(IndexableValue, Record)> = ForwardLogReader::new(
             self.active_metadata_file.try_clone()?,
             self.active_data_file.try_clone()?,
@@ -832,12 +832,7 @@ impl<R: Recordable> DB<R> {
         self.active_data_file.unlock()?;
 
         for (pk, record) in forward_read_items.iter() {
-            // If the record is a tombstone, remove the PK from the map
-            if record.tombstone {
-                pk_to_item_map.remove(pk);
-            } else {
-                pk_to_item_map.insert(pk, record);
-            }
+            pk_to_item_map.insert(pk, record);
         }
 
         debug!(