feat: add XOR filter to JSTable for optimized lookups

Added a BinaryFuse8 XOR filter to JSTable headers to skip unnecessary disk reads for non-existent keys. Replaced bincode with serde_json for serialization as requested. Integrated the filter into the Collection struct and added a get() method to DB.

Author: michaelmior

Cargo.lock

@@ -19,6 +19,7 @@ config = { version = "0.15.19", default-features = false, features = ["toml"] }
 clap = { version = "4.5.54", features = ["derive"] }
 tracing = "0.1.44"
 tracing-subscriber = "0.3.22"
+xorf = { version = "0.12.0", features = ["binary-fuse", "serde"] }
 
 [dev-dependencies]
 criterion = "0.8.1"

Cargo.toml

@@ -13,7 +13,8 @@ The file consists of a sequence of entries. Each entry is encoded as:
 1.  **Header Entry**: The first entry in the file. It is a JSONB-encoded object containing:
     *   `timestamp`: The time the table was created (Unix timestamp in milliseconds).
     *   `schema`: The JSON Schema for the documents.
-2.  **Record Entries**: All subsequent entries. Each is a JSONB-encoded array `[id, document]`:
+2.  **Filter Entry**: The second entry in the file. It is a [Binary Fuse8](https://github.com/ayazhafiz/xorf) filter of the record IDs in the table, serialized as a JSON byte vector.
+3.  **Record Entries**: All subsequent entries. Each is a JSONB-encoded array `[id, document]`:
     *   `id`: String.
     *   `document`: The document object (or `null` for tombstone).
 

specs/storage.md

@@ -8,6 +8,7 @@ use std::fs;
 use std::iter::Peekable;
 use std::path::PathBuf;
 use uuid::Uuid;
+use xorf::{BinaryFuse8, Filter};
 
 struct MergedIterator<'a> {
     sources: Vec<Peekable<Box<dyn Iterator<Item = (String, Value)> + 'a>>>,
@@ -86,6 +87,7 @@ struct Collection {
     logger: Box<dyn Log>,
     memtable_threshold: usize,
     jstable_threshold: u64,
+    filters: Vec<BinaryFuse8>,
 }
 
 impl Collection {
@@ -104,9 +106,21 @@ impl Collection {
             Box::new(NullLogger)
         };
         let memtable = MemTable::new();
-        // Count existing JSTables
+        // Count existing JSTables and load filters
         let mut jstable_count = 0;
+        let mut filters = Vec::new();
         while dir.join(format!("jstable-{}", jstable_count)).exists() {
+            let path = dir.join(format!("jstable-{}", jstable_count));
+            if let Ok(filter) = jstable::read_filter(path.to_str().unwrap()) {
+                filters.push(filter);
+            } else {
+                // Should not happen if file exists and is valid, but handle gracefully?
+                // For now, if we can't read the filter, we might just panic or log error.
+                // Given this is a simple DB, let's assume valid state or fail.
+                // However, we need to push *something* or fail the whole load.
+                // Let's panic to signal corruption.
+                panic!("Failed to read filter for jstable-{}", jstable_count);
+            }
             jstable_count += 1;
         }
 
@@ -118,6 +132,7 @@ impl Collection {
             logger,
             memtable_threshold,
             jstable_threshold,
+            filters,
         }
     }
 
@@ -161,6 +176,11 @@ impl Collection {
         self.memtable
             .flush(jstable_path.to_str().unwrap(), self.name.clone())
             .unwrap();
+
+        // Load the new filter
+        let filter = jstable::read_filter(jstable_path.to_str().unwrap()).unwrap();
+        self.filters.push(filter);
+
         self.jstable_count += 1;
         self.memtable = MemTable::new();
         self.logger.rotate().unwrap();
@@ -187,6 +207,11 @@ impl Collection {
         let new_path = self.dir.join("jstable-0");
         merged_table.write(new_path.to_str().unwrap()).unwrap();
 
+        // Reset filters
+        self.filters.clear();
+        let filter = jstable::read_filter(new_path.to_str().unwrap()).unwrap();
+        self.filters.push(filter);
+
         self.jstable_count = 1;
     }
 
@@ -213,6 +238,47 @@ impl Collection {
 
         MergedIterator { sources }
     }
+
+    fn get(&self, id: &str) -> Option<Value> {
+        // 1. Check MemTable
+        if let Some(doc) = self.memtable.documents.get(id) {
+            if doc.is_null() {
+                return None; // Tombstone
+            }
+            return Some(doc.clone());
+        }
+
+        // 2. Check JSTables (Newer to Older)
+        let hash = {
+            use std::hash::{Hash, Hasher};
+            let mut hasher = std::collections::hash_map::DefaultHasher::new();
+            id.hash(&mut hasher);
+            hasher.finish()
+        };
+
+        for i in (0..self.jstable_count).rev() {
+            if let Some(filter) = self.filters.get(i as usize) {
+                if filter.contains(&hash) {
+                    // Possible match, scan the table
+                    let path = self.dir.join(format!("jstable-{}", i));
+                    if let Ok(iter) = jstable::JSTableIterator::new(path.to_str().unwrap()) {
+                        for res in iter {
+                            if let Ok((rid, doc)) = res {
+                                if rid == id {
+                                    if doc.is_null() {
+                                        return None; // Tombstone
+                                    }
+                                    return Some(doc);
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        None
+    }
 }
 
 impl Debug for Collection {
@@ -373,6 +439,10 @@ impl DB {
         self.get_collection(collection)
             .map(|c| Box::new(c.scan()) as Box<dyn Iterator<Item = (String, Value)> + '_>)
     }
+
+    pub fn get(&self, collection: &str, id: &str) -> Result<Option<Value>, String> {
+        self.get_collection(collection).map(|c| c.get(id))
+    }
 }
 
 #[cfg(test)]
@@ -694,4 +764,30 @@ mod tests {
         );
         assert!(db2.collections.contains_key("test"));
     }
+
+    #[test]
+    fn test_db_get() {
+        let dir = tempdir().unwrap();
+        let mut db = DB::new(
+            dir.path().to_str().unwrap(),
+            MEMTABLE_THRESHOLD,
+            JSTABLE_THRESHOLD,
+            Some(1024 * 1024),
+        );
+        db.create_collection("test").unwrap();
+        let id = db.insert("test", json!({ "a": 1 })).unwrap();
+
+        let doc = db.get("test", &id).unwrap().unwrap();
+        assert_eq!(doc, json!({ "a": 1 }));
+
+        // Flush to force creation of JSTable
+        for i in 0..MEMTABLE_THRESHOLD {
+            db.insert("test", json!({ "fill": i })).unwrap();
+        }
+
+        let doc = db.get("test", &id).unwrap().unwrap();
+        assert_eq!(doc, json!({ "a": 1 }));
+
+        assert!(db.get("test", "non-existent").unwrap().is_none());
+    }
 }

src/db.rs

@@ -4,6 +4,7 @@ use serde_json::Value;
 use std::collections::BTreeMap;
 use std::fs::File;
 use std::io::{self, BufReader, Read, Write};
+use xorf::BinaryFuse8;
 
 pub struct JSTable {
     pub timestamp: u64,
@@ -51,6 +52,27 @@ impl JSTable {
         file.write_all(&header_len.to_le_bytes())?;
         file.write_all(&header_bytes)?;
 
+        // Write Filter
+        let keys: Vec<u64> = self
+            .documents
+            .keys()
+            .map(|k| {
+                use std::hash::{Hash, Hasher};
+                let mut hasher = std::collections::hash_map::DefaultHasher::new();
+                k.hash(&mut hasher);
+                hasher.finish()
+            })
+            .collect();
+        let filter = BinaryFuse8::try_from(&keys).map_err(|_| {
+            io::Error::new(io::ErrorKind::InvalidData, "Failed to create XOR filter")
+        })?;
+        // Use serde_json for filter serialization as fallback
+        let filter_bytes = serde_json::to_vec(&filter)
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+        let filter_len = filter_bytes.len() as u32;
+        file.write_all(&filter_len.to_le_bytes())?;
+        file.write_all(&filter_bytes)?;
+
         // Write Documents
         for (id, doc) in &self.documents {
             let record: (String, &Value) = (id.clone(), doc);
@@ -93,6 +115,16 @@ impl JSTableIterator {
         let header: JSTableHeader = serde_json::from_str(&header_str)
             .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
 
+        // Read/Skip Filter
+        let mut len_buf = [0u8; 4];
+        reader.read_exact(&mut len_buf)?;
+        let filter_len = u32::from_le_bytes(len_buf) as usize;
+        // Skip filter bytes
+        io::copy(
+            &mut reader.by_ref().take(filter_len as u64),
+            &mut io::sink(),
+        )?;
+
         Ok(Self {
             reader,
             timestamp: header.timestamp,
@@ -157,6 +189,37 @@ pub fn read_jstable(path: &str) -> io::Result<JSTable> {
     })
 }
 
+pub fn read_filter(path: &str) -> io::Result<BinaryFuse8> {
+    let file = File::open(path)?;
+    let mut reader = BufReader::new(file);
+
+    // Read Header Length
+    let mut len_buf = [0u8; 4];
+    reader.read_exact(&mut len_buf)?;
+    let header_len = u32::from_le_bytes(len_buf) as usize;
+
+    // Skip Header Blob
+    io::copy(
+        &mut reader.by_ref().take(header_len as u64),
+        &mut io::sink(),
+    )?;
+
+    // Read Filter Length
+    let mut len_buf = [0u8; 4];
+    reader.read_exact(&mut len_buf)?;
+    let filter_len = u32::from_le_bytes(len_buf) as usize;
+
+    // Read Filter Blob
+    let mut filter_blob = vec![0u8; filter_len];
+    reader.read_exact(&mut filter_blob)?;
+
+    // Deserialize using serde_json
+    let filter: BinaryFuse8 = serde_json::from_slice(&filter_blob)
+        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+
+    Ok(filter)
+}
+
 pub fn merge_jstables(tables: &[JSTable]) -> JSTable {
     let mut sorted_tables: Vec<&JSTable> = tables.iter().collect();
     sorted_tables.sort_by_key(|t| t.timestamp);
@@ -192,6 +255,7 @@ mod tests {
     use crate::schema::SchemaType;
     use serde_json::json;
     use tempfile::NamedTempFile;
+    use xorf::Filter;
 
     #[test]
     fn test_read_jstable() -> Result<(), Box<dyn std::error::Error>> {
@@ -269,6 +333,46 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_read_filter() -> 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(
+            12345,
+            "test_col".to_string(),
+            schema.clone(),
+            documents.clone(),
+        );
+
+        let file = NamedTempFile::new().unwrap();
+        jstable.write(file.path().to_str().unwrap()).unwrap();
+
+        let filter = read_filter(file.path().to_str().unwrap())?;
+
+        // Helper to hash string for filter check
+        let hash = |s: &str| {
+            use std::hash::{Hash, Hasher};
+            let mut hasher = std::collections::hash_map::DefaultHasher::new();
+            s.hash(&mut hasher);
+            hasher.finish()
+        };
+
+        assert!(filter.contains(&hash("id1")));
+        assert!(filter.contains(&hash("id2")));
+        assert!(!filter.contains(&hash("id3")));
+
+        Ok(())
+    }
+
     #[test]
     fn test_merge_jstables_conflict_resolution() {
         let schema = Schema::new(SchemaType::Object);