Merge pull request #78 from LevKruglyak/52-mvp-for-learned-components

MVP for learned components

Cargo.toml

@@ -9,8 +9,9 @@ members = [
   "utils/slice_search", 
   "utils/sorted_array", 
   "utils/id_allocator", 
-
-  "bench/runner",
+  "utils/learned_segment",
+  "utils/gapped_array",
+  "bench/runner", 
 ]
 
 exclude = [

core/Cargo.toml

@@ -23,7 +23,10 @@ bincode = { version = "1.3.3" }
 anyhow = "1.0.82"
 
 sorted_array = { path = "../utils/sorted_array", version = "0.1.3", features = ["serde"] }
+gapped_array = { path = "../utils/gapped_array", version = "0.1.0" }
 id_allocator = { path = "../utils/id_allocator", version = "0.1.0", features = ["serde"] }
+learned_index_segmentation = { path = "../utils/learned_segment", version = "0.1.0" }
+
 lazy_static = "1.4.0"
 
 [dev-dependencies]

core/src/classical/memory/mod.rs → core/src/classical/btree_memory/mod.rs

@@ -3,7 +3,7 @@ mod layer;
 use crate::common::list::memory::ArenaID;
 use crate::node_layer::{impl_node_layer, NodeLayer};
 use crate::traits::Address;
-use crate::{component::*, Key, StaticBounded, Value};
+use crate::{component::*, Key, Value};
 use layer::*;
 
 // -------------------------------------------------------
@@ -12,15 +12,15 @@ use layer::*;
 
 pub type BTreeInternalAddress = ArenaID;
 
-pub struct BTreeInternalComponent<K: Ord, X: 'static, const FANOUT: usize, BA, PA> {
+pub struct BTreeInternalComponent<K: Key, X: 'static, const FANOUT: usize, BA, PA> {
     inner: MemoryBTreeLayer<K, BA, FANOUT, PA>,
     _ph: std::marker::PhantomData<X>,
 }
 
 impl<K, X, const FANOUT: usize, BA, PA> NodeLayer<K, BTreeInternalAddress, PA>
     for BTreeInternalComponent<K, X, FANOUT, BA, PA>
 where
-    K: Clone + Ord + StaticBounded,
+    K: Key,
     BA: Address,
     PA: Address,
 {

core/src/classical/mod.rs

@@ -1,9 +1,9 @@
+pub mod btree_disk;
+pub mod btree_memory;
 pub mod btree_top;
-pub mod disk;
-pub mod memory;
 
 mod node;
 
+pub use btree_disk::*;
+pub use btree_memory::*;
 pub use btree_top::*;
-pub use disk::*;
-pub use memory::*;

core/src/common/list/memory.rs

@@ -45,10 +45,11 @@ where
         }
     }
 
-    pub fn insert_after(&mut self, inner: N, ptr: ArenaID) -> ArenaID {
+    #[must_use]
+    pub fn insert_after(&mut self, node: N, ptr: ArenaID) -> ArenaID {
         let next_ptr = self.arena[ptr].0.next;
 
-        let mut new_node = MemoryNode::new(inner);
+        let mut new_node = MemoryNode::new(node);
         new_node.previous = Some(ptr);
         new_node.next = next_ptr;
 
@@ -65,10 +66,11 @@ where
     }
 
     #[allow(unused)]
-    pub fn insert_before(&mut self, inner: N, ptr: ArenaID) -> ArenaID {
+    #[must_use]
+    pub fn insert_before(&mut self, node: N, ptr: ArenaID) -> ArenaID {
         let previous_ptr = self.arena[ptr].0.previous;
 
-        let mut new_node = MemoryNode::new(inner);
+        let mut new_node = MemoryNode::new(node);
         new_node.previous = previous_ptr;
         new_node.next = Some(ptr);
 
@@ -84,6 +86,7 @@ where
         new_node_ptr
     }
 
+    #[must_use]
     pub fn clear(&mut self) -> ArenaID {
         self.arena.clear();
         let ptr = self

core/src/learned/mod.rs

@@ -0,0 +1,5 @@
+pub mod pgm_memory;
+
+mod node;
+
+pub use pgm_memory::*;

core/src/learned/node.rs

@@ -0,0 +1,68 @@
+use learned_index_segmentation::LinearModel;
+
+use crate::{Key, KeyBounded, StaticBounded};
+use gapped_array::GappedKVArray;
+
+impl<K: StaticBounded, const EPSILON: usize> KeyBounded<K> for LinearModel<K, EPSILON> {
+    fn lower_bound(&self) -> &K {
+        self.min_key()
+    }
+}
+
+#[derive(Debug)]
+pub struct PGMNode<K: Key, V, const EPSILON: usize> {
+    gapped: GappedKVArray<K, V>,
+    model: LinearModel<K, EPSILON>,
+}
+
+impl<K: Key, V, const EPSILON: usize> KeyBounded<K> for PGMNode<K, V, EPSILON> {
+    fn lower_bound(&self) -> &K {
+        self.gapped.min().unwrap_or(&K::max_ref())
+    }
+}
+
+impl<K: Key, V, const EPSILON: usize> Default for PGMNode<K, V, EPSILON> {
+    fn default() -> Self {
+        Self {
+            gapped: GappedKVArray::new(0),
+            model: LinearModel::sentinel(),
+        }
+    }
+}
+
+impl<K: Key, V, const EPSILON: usize> PGMNode<K, V, EPSILON> {
+    pub fn from_trained(model: LinearModel<K, EPSILON>, entries: Vec<(K, V)>) -> Self {
+        // NOTE: Filling at 0.5 utilization is just a heuristic, eventually this should be a param
+        let mut gapped = GappedKVArray::new(entries.len() * 2);
+        for (key, value) in entries {
+            let hint = model.hint(&key).min(gapped.len() - 1);
+            gapped
+                .initial_model_based_insert((key, value), hint)
+                .unwrap();
+        }
+        Self { gapped, model }
+    }
+
+    pub fn search_exact(&self, key: &K) -> Option<&V> {
+        let hint = self.model.hint(key);
+        self.gapped.search_exact(key, Some(hint))
+    }
+
+    pub fn search_pir(&self, key: &K) -> &V {
+        let hint = self.model.hint(key);
+        match self.gapped.search_pir(key, Some(hint)) {
+            Some(val) => val,
+            None => self.gapped.min_val().unwrap(),
+        }
+    }
+
+    pub fn grow_insert(&mut self, entry: (K, V)) {
+        if self.gapped.density() >= 0.8 {
+            let scale_factor = 2.0;
+            self.gapped.rescale(scale_factor).unwrap();
+            self.model.rescale(scale_factor as f64);
+        }
+        let hint = self.model.hint(&entry.0);
+        self.gapped.upsert_with_hint(entry, hint).unwrap();
+    }
+}