From aed6b0a1c213e8e8b7f1fc4ee16681d326fd3a4d Mon Sep 17 00:00:00 2001 From: Jan Tuomi Date: Mon, 10 Feb 2025 13:44:47 +0200 Subject: Fix secondary memtable by using a map instead of a set --- log_db/src/common.rs | 82 +++++++++++-------------- log_db/src/engine.rs | 126 +++++++++++++++++++++++++++++++++++++-- log_db/src/lib.rs | 28 ++++----- log_db/src/memtable_secondary.rs | 44 +++++++------- log_db/tests/integration.rs | 2 +- 5 files changed, 190 insertions(+), 92 deletions(-) (limited to 'log_db') diff --git a/log_db/src/common.rs b/log_db/src/common.rs index 3e06db5..a450c28 100644 --- a/log_db/src/common.rs +++ b/log_db/src/common.rs @@ -1,5 +1,6 @@ use super::*; +use std::collections::btree_map::Values; // For Unix-like systems #[cfg(unix)] use std::os::unix::fs::MetadataExt; @@ -48,10 +49,10 @@ pub enum DBError { } #[derive(Debug, Error)] -pub enum LogKeySetError { - #[error("log key not found in set")] +pub enum LogKeyMapError { + #[error("log key not found in map")] NotFoundError, - #[error("attempted to remove last element of non-empty set")] + #[error("attempted to remove last element of non-empty map")] RemovingLastElementError, } @@ -76,74 +77,59 @@ impl LogKey { } } -/// LogKeySet is a non-empty set of LogKeys. +/// LogKeyMap is a non-empty map of PK => LogKey mappings. #[derive(Debug, Clone, Eq, PartialEq)] -pub struct LogKeySet { - set: HashSet, +pub struct LogKeyMap { + map: BTreeMap, } -impl PartialOrd for LogKeySet { - fn partial_cmp(&self, other: &Self) -> Option { - let self_max_elem = self.set.iter().max()?; - let other_max_elem = other.set.iter().max()?; - Some(self_max_elem.cmp(other_max_elem)) - } -} - -impl LogKeySet { - /// Create a new LogKeySet with an initial LogKey. - /// The initial LogKey is required since LogKeySet must be non-empty. - pub fn new_with_initial(key: LogKey) -> Self { - let mut set = HashSet::new(); - set.insert(key); - LogKeySet { set } +impl LogKeyMap { + /// Create a new LogKeyMap with an initial mapping. + /// The initial mapping is required since LogKeyMap must be non-empty. + pub fn new_with_initial(pk: IndexableValue, log_key: LogKey) -> Self { + let mut map = BTreeMap::new(); + map.insert(pk, log_key); + LogKeyMap { map } } - pub fn contains(&self, key: &LogKey) -> bool { - self.set.contains(key) + pub fn contains_pk(&self, key: &IndexableValue) -> bool { + self.map.contains_key(key) } - /// The number of LogKeys in the set. + /// The number of LogKeys in the map. pub fn len(&self) -> usize { - self.set.len() + self.map.len() } - /// Insert a LogKey into the set. - pub fn insert(&mut self, key: LogKey) { - self.set.insert(key); + /// Insert a PK -> LogKey mapping into the map. + pub fn insert(&mut self, key: IndexableValue, log_key: LogKey) { + self.map.insert(key, log_key); } - /// Remove a LogKey from the set. Return Ok(()) if the key was found and removed. - /// Return `LogKeySetError::RemovingLastElementError` if trying to remove the last element. - /// Return `LogKeySetError::NotFoundError` if the key was not found. - pub fn remove(&mut self, key: &LogKey) -> Result<(), LogKeySetError> { - if self.set.len() == 1 { - return Err(LogKeySetError::RemovingLastElementError); + /// Remove a mapping from the map. Return Ok(()) if the key was found and removed. + /// Return `LogKeyMapError::RemovingLastElementError` if trying to remove the last element. + /// Return `LogKeyMapError::NotFoundError` if the key was not found. + pub fn remove_pk(&mut self, key: &IndexableValue) -> Result<(), LogKeyMapError> { + if self.map.len() == 1 { + return Err(LogKeyMapError::RemovingLastElementError); } - let removed = self.set.remove(key); + let removed = self.map.remove(key); - if !removed { - return Err(LogKeySetError::NotFoundError); + if removed.is_none() { + return Err(LogKeyMapError::NotFoundError); } assert!( - self.set.len() > 0, - "LogKeySet should not be empty after removal" + self.map.len() > 0, + "LogKeyMap should not be empty after removal" ); Ok(()) } /// Get a reference to the set of LogKeys. - pub fn log_keys(&self) -> &HashSet { - &self.set - } -} - -impl Ord for LogKeySet { - fn cmp(&self, other: &Self) -> Ordering { - self.partial_cmp(other) - .expect("LogKeySet comparison failed, possibly due to empty set") + pub fn log_keys(&self) -> Values { + self.map.values() } } diff --git a/log_db/src/engine.rs b/log_db/src/engine.rs index b8ac6b9..62a88cb 100644 --- a/log_db/src/engine.rs +++ b/log_db/src/engine.rs @@ -192,6 +192,8 @@ impl Engine { } fn insert_record_to_memtables(&mut self, log_key: LogKey, record: Record) { + let pk = record.at(self.primary_key_index).as_indexable().unwrap(); + for (sk_index, sk_field) in self.config.secondary_keys.iter().enumerate() { let secondary_memtable = &mut self.secondary_memtables[sk_index]; let sk_field_index = self @@ -202,19 +204,17 @@ impl Engine { .unwrap(); let sk = record.at(sk_field_index).as_indexable().unwrap(); - secondary_memtable.set(sk, log_key.clone()); + secondary_memtable.set(pk.clone(), sk, log_key.clone()); } // Doing this last because this moves log_key - let pk = record.at(self.primary_key_index).as_indexable().unwrap(); self.primary_memtable.set(pk, log_key); } fn remove_record_from_memtables(&mut self, record: &Record) { let pk = record.at(self.primary_key_index).as_indexable().unwrap(); - if let Some(plk) = self.primary_memtable.remove(&pk) { - // TODO this does not work (test_delete_by_multiple_indexes) + if let Some(_) = self.primary_memtable.remove(&pk) { for (sk_index, sk_field) in self.config.secondary_keys.iter_mut().enumerate() { let secondary_memtable = &mut self.secondary_memtables[sk_index]; let sk_field_index = self @@ -225,7 +225,7 @@ impl Engine { .unwrap(); let sk = record.at(sk_field_index).as_indexable().unwrap(); - secondary_memtable.remove(&sk, &plk); + secondary_memtable.remove(&pk, &sk); } } } @@ -746,3 +746,119 @@ impl Engine { result } } + +#[cfg(test)] +mod tests { + use ctor::ctor; + use env_logger; + + use super::*; + + #[ctor] + fn init_logger() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + #[derive(Eq, PartialEq, Clone, Debug)] + enum Field { + Id, + Name, + } + + #[derive(PartialEq, Eq, Debug, Clone)] + struct TestInst2 { + id: i64, + name: String, + } + + impl TestInst2 { + fn into_record(self) -> Vec { + vec![Value::Int(self.id), Value::String(self.name)] + } + + fn from_record(record: Vec) -> Self { + let mut it = record.into_iter(); + TestInst2 { + id: match it.next().unwrap() { + Value::Int(i) => i, + _ => panic!("Expected int"), + }, + name: match it.next().unwrap() { + Value::String(s) => s, + _ => panic!("Expected string"), + }, + } + } + } + + #[test] + fn test_memtable_insert_and_delete() { + let temp_dir = tempfile::tempdir().unwrap(); + let data_dir = temp_dir.path(); + + let capacity = 5; + let segment_size = capacity * 2 * 8 + METADATA_FILE_HEADER_SIZE; + + let mut db = DB::configure() + .data_dir(data_dir.to_str().unwrap()) + .schema(vec![ + (Field::Id, Type::int()), + (Field::Name, Type::string()), + ]) + .primary_key(Field::Id) + .secondary_keys(vec![Field::Name]) + .from_record(TestInst2::from_record) + .into_record(TestInst2::into_record) + .segment_size(segment_size) + .initialize() + .expect("Failed to create DB"); + + let engine = &mut db.engine; + + let inst = TestInst2 { + id: 0, + name: "foo".to_owned(), + }; + let id = IndexableValue::Int(0); + + assert_eq!(engine.primary_memtable.get(&id), None); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 0 + ); + engine.insert_record_to_memtables( + LogKey::new(1, 0), + Record::from(&inst.clone().into_record()), + ); + assert_eq!(engine.primary_memtable.get(&id), Some(&LogKey::new(1, 0))); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 1 + ); + + engine.insert_record_to_memtables( + LogKey::new(1, 1), + Record::from(&inst.clone().into_record()), + ); + assert_eq!(engine.primary_memtable.get(&id), Some(&LogKey::new(1, 1))); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 1 + ); + + engine.remove_record_from_memtables(&Record::from(&inst.into_record())); + assert_eq!(engine.primary_memtable.get(&id), None); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 0 + ); + } +} diff --git a/log_db/src/lib.rs b/log_db/src/lib.rs index 0317320..e92b68d 100644 --- a/log_db/src/lib.rs +++ b/log_db/src/lib.rs @@ -3,9 +3,7 @@ extern crate log; use once_cell::sync::Lazy; use rust_decimal::Decimal; -use std::cmp::Ordering; use std::collections::BTreeMap; -use std::collections::HashSet; use std::fmt::Debug; use std::fmt::Display; use std::fs::{self, metadata, File}; @@ -215,7 +213,7 @@ impl DB { pub fn tx_rollback(&mut self) -> DBResult<()> { if !self.engine.tx_active { return Err(DBError::TransactionError( - "No active transaction to rollback".to_string(), + "No active transaction to roll back".to_string(), )); } @@ -230,7 +228,6 @@ impl DB { mod tests { use ctor::ctor; use env_logger; - use std::collections::HashSet; use super::*; @@ -292,7 +289,6 @@ mod tests { #[test] fn test_compaction() { - let _ = env_logger::builder().is_test(true).try_init(); let temp_dir = tempfile::tempdir().unwrap(); let data_dir = temp_dir.path(); @@ -383,7 +379,6 @@ mod tests { #[test] fn test_repair() { - let _ = env_logger::builder().is_test(true).try_init(); let temp_dir = tempfile::tempdir().unwrap(); let data_dir = temp_dir.path(); @@ -438,7 +433,6 @@ mod tests { #[test] fn test_memtables_updated_on_write() { - let _ = env_logger::builder().is_test(true).try_init(); let temp_dir = tempfile::tempdir().unwrap(); let data_dir = temp_dir.path(); @@ -461,8 +455,10 @@ mod tests { None ); assert_eq!( - db.engine.secondary_memtables[0].find_by(&IndexableValue::String("John".to_string())), - &HashSet::new() + db.engine.secondary_memtables[0] + .find_by(&IndexableValue::String("John".to_string())) + .len(), + 0 ); // Insert record @@ -474,15 +470,15 @@ mod tests { // Check that the key is now indexed let expected_log_key = LogKey::new(1, 0); + let expected_pk = IndexableValue::Int(0); assert_eq!( - db.engine.primary_memtable.get(&IndexableValue::Int(0)), + db.engine.primary_memtable.get(&expected_pk), Some(&expected_log_key) ); - let mut expected_set: HashSet = HashSet::new(); - expected_set.insert(expected_log_key); - assert_eq!( - db.engine.secondary_memtables[0].find_by(&IndexableValue::String("John".to_string())), - &expected_set, - ); + let expected_vals = vec![&expected_log_key]; + let actual_vals = db.engine.secondary_memtables[0] + .find_by(&IndexableValue::String("John".to_string())) + .collect::>(); + assert_eq!(actual_vals, expected_vals); } } diff --git a/log_db/src/memtable_secondary.rs b/log_db/src/memtable_secondary.rs index cd6c252..e1ca835 100644 --- a/log_db/src/memtable_secondary.rs +++ b/log_db/src/memtable_secondary.rs @@ -1,16 +1,15 @@ use once_cell::sync::Lazy; use super::*; -use std::collections::{BTreeMap, HashSet}; +use std::collections::{btree_map::Values, BTreeMap}; pub struct SecondaryMemtable { - /// 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, + /// A 2-layer map of records indexed by SK => PK => LogKey. + /// The PK information is required to tell two records apart. + records: BTreeMap, } -static EMPTY_SET: Lazy> = Lazy::new(|| HashSet::new()); +static EMPTY_MAP: Lazy> = Lazy::new(|| BTreeMap::new()); impl SecondaryMemtable { pub fn new() -> SecondaryMemtable { @@ -19,38 +18,39 @@ impl SecondaryMemtable { } } - pub fn set(&mut self, key: IndexableValue, value: LogKey) { - match self.records.get_mut(&key) { - Some(set) => { - set.insert(value); + pub fn set(&mut self, pk: IndexableValue, sk: IndexableValue, value: LogKey) { + match self.records.get_mut(&sk) { + Some(map) => { + map.insert(pk, value); } None => { - self.records.insert(key, LogKeySet::new_with_initial(value)); + self.records + .insert(sk, LogKeyMap::new_with_initial(pk, value)); } }; } - pub fn find_by(&self, key: &IndexableValue) -> &HashSet { + pub fn find_by(&self, key: &IndexableValue) -> Values { match self.records.get(key) { Some(set) => set.log_keys(), - None => &EMPTY_SET, + None => EMPTY_MAP.values(), } } - // Remove a single log key associated with the given key. Returns `true` + // Remove a single mapping associated with the given PK and SK. Returns `true` // if the log key existed and was removed, `false` otherwise. - pub fn remove(&mut self, key: &IndexableValue, log_key: &LogKey) -> bool { - let set = match self.records.get_mut(key) { + pub fn remove(&mut self, pk: &IndexableValue, sk: &IndexableValue) -> bool { + let map = match self.records.get_mut(sk) { Some(set) => set, None => return false, }; - if set.len() == 1 && set.contains(log_key) { - self.records.remove(key); + if map.len() == 1 && map.contains_pk(pk) { + self.records.remove(sk); true } else { - return match set.remove(log_key) { + return match map.remove_pk(pk) { Ok(_) => true, - Err(LogKeySetError::NotFoundError) => false, + Err(LogKeyMapError::NotFoundError) => false, Err(e) => panic!("{:?}", e), }; } @@ -58,8 +58,8 @@ impl SecondaryMemtable { pub fn range>(&self, range: B) -> Vec<&LogKey> { let mut keys = Vec::new(); - for (_, set) in self.records.range(range) { - keys.extend(set.log_keys().iter()); + for (_, map) in self.records.range(range) { + keys.extend(map.log_keys()); } keys } diff --git a/log_db/tests/integration.rs b/log_db/tests/integration.rs index 4672c06..b09bff3 100644 --- a/log_db/tests/integration.rs +++ b/log_db/tests/integration.rs @@ -808,7 +808,7 @@ fn test_rollback_transaction() { }) .unwrap(); - db.tx_rollback().expect("Failed to rollback transaction"); + db.tx_rollback().expect("Failed to roll back transaction"); let johns = db .find_by(&Field::Name, &Value::String("John".to_string())) -- cgit v1.3