From ee0a81cf0277a7ff8683f94cc95a23f29996d980 Mon Sep 17 00:00:00 2001
From: Jan Tuomi <jan@jantuomi.fi>
Date: Thu, 10 Oct 2024 23:41:10 +0300
Subject: Add invariant that memtables should always have the same set of
 primary key references

---
 log_db/src/secondary_memtable.rs | 36 ++++++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

(limited to 'log_db/src/secondary_memtable.rs')

diff --git a/log_db/src/secondary_memtable.rs b/log_db/src/secondary_memtable.rs
index 194d4a0..994ee0d 100644
--- a/log_db/src/secondary_memtable.rs
+++ b/log_db/src/secondary_memtable.rs
@@ -5,17 +5,30 @@ use std::fmt::Debug;
 
 pub struct SecondaryMemtable<Field: Eq + Clone + Debug> {
     pub field: Field,
+    field_index: usize,
+    primary_key_index: usize,
 
     /// Map of records indexed by key. The value is the set of primary key values of records
     /// that have the secondary key value. The actual `Record` objects are stored in the
     /// primary memtable, which acts as the shared heap.
-    records: BTreeMap<IndexableValue, HashSet<IndexableValue>>,
+    pub records: BTreeMap<IndexableValue, HashSet<IndexableValue>>,
 }
 
 impl<Field: Eq + Clone + Debug> SecondaryMemtable<Field> {
-    pub fn new(field: &Field) -> SecondaryMemtable<Field> {
+    pub fn new(
+        field_schema: &Vec<(Field, RecordField)>,
+        field: &Field,
+        primary_key_index: usize,
+    ) -> SecondaryMemtable<Field> {
+        let field_index = field_schema
+            .iter()
+            .position(|(f, _)| f == field)
+            .expect("Field not found in schema");
+
         SecondaryMemtable {
             field: field.clone(),
+            field_index,
+            primary_key_index,
             records: BTreeMap::new(),
         }
     }
@@ -76,4 +89,23 @@ impl<Field: Eq + Clone + Debug> SecondaryMemtable<Field> {
                 .collect(),
         }
     }
+
+    pub fn remove(&mut self, record: &Record) {
+        let key = record.values[self.field_index]
+            .as_indexable()
+            .expect("Field is not indexable");
+
+        let primary_key = record.values[self.primary_key_index]
+            .as_indexable()
+            .expect("Primary key is not indexable");
+
+        match self.records.get_mut(&key) {
+            Some(set) => {
+                set.remove(&primary_key);
+            }
+            None => {
+                panic!("Record not found in secondary memtable");
+            }
+        }
+    }
 }
-- 
cgit v1.3