Implement JSTable compaction with tombstone support for deletes

Author: michaelmior

ROADMAP.md

@@ -34,7 +34,7 @@ This stage will allow the data stored in memory to be dumped to disk to allow st
 
 This stage will implement the compaction algorithm for the LSM tree.
 
-- [] Implement compaction for JSTables by merging their schemas and updating the compressed documents
+- [x] Implement compaction for JSTables by merging their schemas and updating the compressed documents
 
 # Stage 6 - Internal query API
 

specs/storage.md

@@ -15,11 +15,13 @@ Example:
 {"type":"object","properties":{"a":{"type":["integer"]}}}
 ["01H4J3J4J3J4J3J4J3J4J3J4J3", {"a": 1}]
 ["01H4J3J4J3J4J3J4J3J4J3J4J4", {"a": 2}]
+["01H4J3J4J3J4J3J4J3J4J3J4J5", null]
 ```
 
 ## Compression
 
 Documents are stored as full JSON objects. This means that no compression is currently applied.
+A document can be `null` to indicate that it has been deleted.
 Future versions of the format might include schema-based compression. For example, if a field must be an integer, the type information need not be stored with each document. If a field can have multiple possible types, then the type information must be stored with each document.
 
 

src/db.rs

@@ -1,9 +1,12 @@
+use crate::jstable;
 use crate::log::{LogEntry, Logger, Operation};
 use crate::storage::MemTable;
-use serde_json::Value;
+use serde_json::{json, Value};
+use std::fs;
 use uuid::Uuid;
 
 const MEMTABLE_THRESHOLD: usize = 10;
+const JSTABLE_THRESHOLD: u64 = 5;
 
 pub struct DB {
     memtable: MemTable,
@@ -14,7 +17,7 @@ pub struct DB {
 
 impl DB {
     pub fn new(jstable_dir: &str) -> Self {
-        std::fs::create_dir_all(jstable_dir).unwrap();
+        fs::create_dir_all(jstable_dir).unwrap();
         let log_path = format!("{}/argus.log", jstable_dir);
         let logger = Logger::new(&log_path, 1024 * 1024).unwrap();
         let mut memtable = MemTable::new();
@@ -84,6 +87,30 @@ impl DB {
         self.jstable_count += 1;
         self.memtable = MemTable::new();
         self.logger.rotate().unwrap();
+
+        if self.jstable_count >= JSTABLE_THRESHOLD {
+            self.compact();
+        }
+    }
+
+    fn compact(&mut self) {
+        let mut tables = Vec::new();
+        for i in 0..self.jstable_count {
+            let path = format!("{}/jstable-{}", self.jstable_dir, i);
+            tables.push(jstable::read_jstable(&path).unwrap());
+        }
+
+        let merged_table = jstable::merge_jstables(&tables);
+
+        for i in 0..self.jstable_count {
+            let path = format!("{}/jstable-{}", self.jstable_dir, i);
+            fs::remove_file(path).unwrap();
+        }
+
+        let new_path = format!("{}/jstable-0", self.jstable_dir);
+        merged_table.write(&new_path).unwrap();
+
+        self.jstable_count = 1;
     }
 }
 
@@ -170,4 +197,85 @@ mod tests {
         assert_eq!(db2.memtable.len(), 1);
         assert_eq!(*db2.memtable.documents.get(&id2).unwrap(), doc2);
     }
+
+    #[test]
+    fn test_db_compaction() {
+        let dir = tempdir().unwrap();
+        let mut db = DB::new(dir.path().to_str().unwrap());
+
+        for i in 0..(MEMTABLE_THRESHOLD * JSTABLE_THRESHOLD as usize) {
+            db.insert(json!({ "a": i }));
+        }
+
+        assert_eq!(db.jstable_count, JSTABLE_THRESHOLD - 1);
+        db.insert(json!({ "a": 999 })); // Trigger another flush, which triggers compaction
+        assert_eq!(db.jstable_count, 1);
+
+        // Verify data is preserved
+        // We inserted 0..50 (50 items) + 1 (999). 51 items total.
+        // Item "0" should be in the compacted table.
+        // We can't easily query DB yet (no read path implemented in DB), 
+        // so we manually check the file.
+        let jstable_path = format!("{}/jstable-0", dir.path().to_str().unwrap());
+        let table = jstable::read_jstable(&jstable_path).unwrap();
+        assert!(table.documents.len() >= 50); // It should have the 50 items from the first 5 flushes (0-49)
+        // The last inserted item (999) is in memtable, not in jstable-0 yet.
+    }
+
+    #[test]
+    fn test_db_compaction_with_delete() {
+        let dir = tempdir().unwrap();
+        let mut db = DB::new(dir.path().to_str().unwrap());
+
+        // 1. Insert doc to be deleted
+        let id_to_delete = db.insert(json!({ "a": 100 }));
+        
+        // Fill memtable to force flush 1 (jstable-0)
+        // 1 item already inserted. Insert 9 more to fill (total 10).
+        for i in 0..9 {
+            db.insert(json!({ "fill": i }));
+        }
+        // 11th insert triggers flush of the first 10
+        db.insert(json!({ "trigger_1": 1 }));
+        assert_eq!(db.jstable_count, 1); 
+
+        // 2. Delete the doc
+        // id_to_delete is in jstable-0.
+        // delete adds tombstone to memtable.
+        // memtable currently has "trigger_1" (1 item).
+        // delete adds 1 item. len = 2.
+        db.delete(&id_to_delete);
+        
+        // Fill memtable to force flush 2 (jstable-1)
+        // Memtable len is 2. Need 8 more to fill (total 10).
+        for i in 0..8 {
+            db.insert(json!({ "fill_2": i }));
+        }
+        // 11th insert (relative to this batch) triggers flush
+        db.insert(json!({ "trigger_2": 1 }));
+        assert_eq!(db.jstable_count, 2); 
+
+        // 3. Create 3 more tables to reach threshold 5
+        for t in 0..3 {
+            // Fill memtable (10 items)
+            // Memtable has 1 item ("trigger_2" or previous trigger).
+            // Need 9 more.
+            for i in 0..9 {
+                db.insert(json!({ "fill_more": t, "i": i }));
+            }
+            // Trigger flush
+            db.insert(json!({ "trigger_more": t }));
+        }
+        // After 3rd iteration (total 5th flush), compaction triggers.
+        // jstable_count goes 4 -> 5 -> 1.
+        assert_eq!(db.jstable_count, 1);
+        
+        // 4. Verify id_to_delete is NOT in jstable-0
+        let jstable_path = format!("{}/jstable-0", dir.path().to_str().unwrap());
+        let table = jstable::read_jstable(&jstable_path).unwrap();
+        assert!(!table.documents.contains_key(&id_to_delete));
+        
+        // Verify other documents exist (e.g. from flush 1)
+        assert!(table.documents.len() > 40);
+    }
 }

src/jstable.rs

@@ -1,14 +1,32 @@
 use crate::schema::Schema;
 use serde_json::Value;
 use std::collections::BTreeMap;
-use std::io::{self, BufReader, BufRead};
+use std::io::{self, BufReader, BufRead, Write};
 use std::fs::File;
 
 pub struct JSTable {
     pub schema: Schema,
     pub documents: BTreeMap<String, Value>,
 }
 
+impl JSTable {
+    pub fn new(schema: Schema, documents: BTreeMap<String, Value>) -> Self {
+        JSTable { schema, documents }
+    }
+
+    pub fn write(&self, path: &str) -> io::Result<()> {
+        let mut file = File::create(path)?;
+        let schema_json = serde_json::to_string(&self.schema)?;
+        writeln!(file, "{}", schema_json)?;
+        for (id, doc) in &self.documents {
+            let record: (String, &Value) = (id.clone(), doc);
+            let record_json = serde_json::to_string(&record)?;
+            writeln!(file, "{}", record_json)?;
+        }
+        Ok(())
+    }
+}
+
 pub fn read_jstable(path: &str) -> io::Result<JSTable> {
     let file = File::open(path)?;
     let reader = BufReader::new(file);
@@ -20,13 +38,32 @@ pub fn read_jstable(path: &str) -> io::Result<JSTable> {
     let mut documents = BTreeMap::new();
     for line_result in lines {
         let line = line_result?;
+        if line.is_empty() {
+            continue;
+        }
         let record: (String, Value) = serde_json::from_str(&line)?;
         documents.insert(record.0, record.1);
     }
 
     Ok(JSTable { schema, documents })
 }
 
+pub fn merge_jstables(tables: &[JSTable]) -> JSTable {
+    let mut merged_schema = Schema::new(crate::schema::SchemaType::Object);
+    let mut merged_documents = BTreeMap::new();
+
+    for table in tables {
+        merged_schema.merge(table.schema.clone());
+        for (id, doc) in &table.documents {
+            merged_documents.insert(id.clone(), doc.clone());
+        }
+    }
+
+    merged_documents.retain(|_, v| !v.is_null());
+
+    JSTable::new(merged_schema, merged_documents)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -46,8 +83,8 @@ mod tests {
             items: None,
         }).unwrap();
         writeln!(file, "{}", schema_json).unwrap();
-        writeln!(file, "{}", serde_json::to_string(&json!(["id1", {"a": 1}]))?).unwrap();
-        writeln!(file, "{}", serde_json::to_string(&json!(["id2", {"a": 2}]))?).unwrap();
+        writeln!(file, "{}", serde_json::to_string(&json!(["id1", {"a": 1}])).unwrap()).unwrap();
+        writeln!(file, "{}", serde_json::to_string(&json!(["id2", {"a": 2}])).unwrap()).unwrap();
 
         let jstable = read_jstable(file.path().to_str().unwrap()).unwrap();
 
@@ -57,4 +94,29 @@ mod tests {
         assert_eq!(*jstable.documents.get("id2").unwrap(), json!({"a": 2}));
         Ok(())
     }
+
+    #[test]
+    fn test_write_jstable() -> Result<(), Box<dyn std::error::Error>> {
+        let schema = Schema {
+            types: vec![SchemaType::Object],
+            properties: Some(BTreeMap::from([
+                ("a".to_string(), Schema::new(SchemaType::Integer)),
+            ])),
+            items: None,
+        };
+        let mut documents = BTreeMap::new();
+        documents.insert("id1".to_string(), json!({"a": 1}));
+        documents.insert("id2".to_string(), json!({"a": 2}));
+        let jstable = JSTable::new(schema, documents);
+
+        let file = NamedTempFile::new().unwrap();
+        jstable.write(file.path().to_str().unwrap()).unwrap();
+
+        let content = std::fs::read_to_string(file.path()).unwrap();
+        assert!(content.contains("{\"type\":\"object\",\"properties\":{\"a\":{\"type\":[\"integer\"]}}}"));
+        assert!(content.contains("[\"id1\",{\"a\":1}]"));
+        assert!(content.contains("[\"id2\",{\"a\":2}]"));
+
+        Ok(())
+    }
 }

src/storage.rs

@@ -1,8 +1,7 @@
+use crate::jstable::JSTable;
 use crate::schema::{infer_schema, Schema};
 use serde_json::Value;
 use std::collections::BTreeMap;
-use std::io::Write;
-use serde_json::json;
 
 pub struct MemTable {
     pub documents: BTreeMap<String, Value>,
@@ -26,15 +25,8 @@ impl MemTable {
     }
 
     pub fn flush(&self, path: &str) -> std::io::Result<()> {
-        let mut file = std::fs::File::create(path)?;
-        let schema_json = serde_json::to_string(&self.schema)?;
-        writeln!(file, "{}", schema_json)?;
-        for (id, doc) in &self.documents {
-            let record = json!([id, doc]);
-            let record_json = serde_json::to_string(&record)?;
-            writeln!(file, "{}", record_json)?;
-        }
-        Ok(())
+        let jstable = JSTable::new(self.schema.clone(), self.documents.clone());
+        jstable.write(path)
     }
 
     pub fn insert(&mut self, id: String, doc: Value) {
@@ -43,15 +35,15 @@ impl MemTable {
         self.documents.insert(id, doc);
     }
 
-    pub fn delete(&mut self, id: &str) {
-        self.documents.remove(id);
-    }
-
     pub fn update(&mut self, id: &str, doc: Value) {
         let doc_schema = infer_schema(&doc);
         self.schema.merge(doc_schema);
         self.documents.insert(id.to_string(), doc);
     }
+
+    pub fn delete(&mut self, id: &str) {
+        self.documents.insert(id.to_string(), Value::Null);
+    }
 }
 
 #[cfg(test)]
@@ -75,15 +67,6 @@ mod tests {
         assert_eq!(props.get("b").unwrap().types, vec![SchemaType::String]);
     }
 
-    #[test]
-    fn test_memtable_delete() {
-        let mut memtable = MemTable::new();
-        memtable.insert("test-id".to_string(), json!({"a": 1}));
-        assert_eq!(memtable.len(), 1);
-        memtable.delete("test-id");
-        assert_eq!(memtable.len(), 0);
-    }
-
     #[test]
     fn test_memtable_update() {
         let mut memtable = MemTable::new();