use shared cache for HNSW vertices

Author: shisoft

src/apps/hnsw/partition/search.rs

@@ -1,7 +1,6 @@
 use std::{
     cell::RefCell,
     collections::{BTreeMap, BTreeSet, BinaryHeap, HashSet, VecDeque},
-    rc::Rc,
     sync::Arc,
     vec,
 };
@@ -16,7 +15,7 @@ use dovahkiin::{
 use lightning::map::Map;
 use lightning::map::PtrHashMap;
 use neb::{
-    index::{ranged::trees::Ordering, IndexerClients},
+    index::{ranged::trees::Ordering},
     query::data_client::IndexedDataClient,
     ram::{
         chunk::Chunks,
@@ -27,7 +26,7 @@ use neb::{
 use crate::{
     apps::hnsw::measurements::{vector_distance, HnswMetric, Metric, MetricEncoding},
     graph::{
-        edge::Edge, format_vertex_cell, id_list::IdList, partitioner::DefaultPartitioner,
+        format_vertex_cell, id_list::IdList, partitioner::DefaultPartitioner,
         vertex::Vertex, EdgeDirection, GraphEngine, NeighbourhoodError,
     },
     job::{
@@ -76,6 +75,9 @@ pub struct HnswOnlinePartition {
     pub server_id: u64,
     pub chunks: Arc<Chunks>,
     pub engine: Arc<GraphEngine>,
+    // Shared caches across searches in this partition
+    pub shared_vertex_cache: PtrHashMap<Id, (Id, Id)>,
+    pub shared_vertex_vector_cache: PtrHashMap<(Id, u64), Arc<OwnedPrimArray>>,
 }
 
 impl HnswOnlinePartition {
@@ -168,6 +170,8 @@ impl HnswOnlinePartition {
             server_id,
             chunks: chunks.clone(),
             engine: engine.clone(),
+            shared_vertex_cache: PtrHashMap::with_capacity(1 << 16),
+            shared_vertex_vector_cache: PtrHashMap::with_capacity(1 << 18),
         })
     }
 
@@ -226,14 +230,15 @@ impl HnswOnlinePartition {
         metadata: &SearchMetadata,
         job_id: JobId,
     ) -> Option<(Id, Id)> {
-        if let Some(v) = metadata.vertex_cache.borrow().get(vertex_id) {
+        // shared cache only
+        if let Some(v) = self.shared_vertex_cache.get(vertex_id) {
             append_job_log(
                 logger,
                 job_id,
                 JobLogLevel::Trace,
-                format!("Found vertex in cache for vertex {:?}", vertex_id),
+                format!("Found vertex in SHARED cache for vertex {:?}", vertex_id),
             );
-            return Some(*v);
+            return Some(v);
         }
         let direction = EdgeDirection::Undirected;
         let edge_list = direction.as_field();
@@ -261,10 +266,8 @@ impl HnswOnlinePartition {
             (SharedValue::Id(cell_id), SharedValue::Id(type_list_id)) => {
                 let cell_id = **cell_id;
                 let type_list_id = **type_list_id;
-                metadata
-                    .vertex_cache
-                    .borrow_mut()
-                    .insert(*vertex_id, (cell_id, type_list_id));
+                // write-through to shared cache
+                self.shared_vertex_cache.insert(*vertex_id, (cell_id, type_list_id));
                 // println!("Inserted cell id in cache for vertex {:?}", vertex_id);
                 Some((cell_id, type_list_id))
             }
@@ -300,11 +303,16 @@ impl HnswOnlinePartition {
         chunks: &'a Arc<Chunks>,
         logger: &Arc<JobLogger>,
         job_id: JobId,
-    ) -> Option<Rc<OwnedPrimArray>> {
+    ) -> Option<Arc<OwnedPrimArray>> {
+        // per-search 
         let mut vertex_vector_cache = metadata.vertex_vector_cache.borrow_mut();
-        if let Some(vector) = vertex_vector_cache.get(vertex_id) {
+        if let Some(vector) = vertex_vector_cache.get(&(*vertex_id, field_id)) {
             return Some(vector.clone());
         }
+        // shared cache fallback
+        if let Some(vector) = self.shared_vertex_vector_cache.get(&(*vertex_id, field_id)) {
+            return Some(vector);
+        }
         let (cell_id, _type_list_id) =
             self.get_vertex(vertex_id, chunks, logger, metadata, job_id)?;
         let cell = match chunks.read_selected(&cell_id, &[field_id], false) {
@@ -327,9 +335,11 @@ impl HnswOnlinePartition {
         match field {
             SharedValue::PrimArray(vector) => {
                 let vector = vector.copy_into_owned();
-                let rc_vector = Rc::new(vector);
-                vertex_vector_cache.insert(*vertex_id, rc_vector.clone());
-                Some(rc_vector)
+                let arc_vector = Arc::new(vector);
+                self.shared_vertex_vector_cache
+                    .insert((*vertex_id, field_id), arc_vector.clone());
+                vertex_vector_cache.insert((*vertex_id, field_id), arc_vector.clone());
+                Some(arc_vector)
             }
             _ => {
                 append_job_log(
@@ -356,7 +366,8 @@ impl HnswOnlinePartition {
     ) -> Option<Distance> {
         let mut vs = vec![vertex_a, vertex_b];
         vs.sort();
-        let distance_key = (*vs[0], *vs[1]);
+        let distance_key = (*vs[0], *vs[1], field_id);
+        // per-search cache
         let mut distance_cache = metadata.vertex_distance_cache.borrow_mut();
         if let Some(distance) = distance_cache.get(&distance_key) {
             return Some(*distance);
@@ -550,7 +561,6 @@ impl HnswOnlinePartition {
                 level_entries: HashMap::default(),
                 vertex_distance_cache: RefCell::new(BTreeMap::default()),
                 vertex_vector_cache: RefCell::new(BTreeMap::default()),
-                vertex_cache: RefCell::new(BTreeMap::default()),
                 ef,
                 ef_construction,
             },
@@ -562,7 +572,7 @@ impl HnswOnlinePartition {
             metric,
             max_level,
         };
-        partition_search.metadata.set_query_vector(query);
+        partition_search.metadata.set_query_vector(query, field_id);
         Ok(partition_search)
     }
 
@@ -1116,6 +1126,10 @@ impl HnswOnlinePartition {
                     }
                 }
             }
+            // Invalidate only affected entries
+            // Note: leaving distance cache entries intact for stability; vector values are immutable
+            self.shared_vertex_cache.remove(&vertex_id);
+            self.shared_vertex_vector_cache.remove(&(vertex_id.clone(), field_id));
         }
         // If this is a high-level vertex, add it to the index's top_level_vertices
         if random_level >= max_level {

src/apps/hnsw/partition/types.rs

@@ -1,6 +1,6 @@
 use std::cell::RefCell;
 use std::collections::{BTreeMap, HashSet};
-use std::rc::Rc;
+use std::sync::Arc;
 
 use ahash::HashMap;
 use bifrost::rpc::RPCError;
@@ -62,9 +62,8 @@ pub struct SearchMetadata {
     pub level_entries: ahash::HashMap<usize, Id>,
     pub remote_frontiers: HashSet<Id>,
     pub last_distance: Distance,
-    pub vertex_distance_cache: RefCell<BTreeMap<(Id, Id), Distance>>,
-    pub vertex_cache: RefCell<BTreeMap<Id, (Id, Id)>>,
-    pub vertex_vector_cache: RefCell<BTreeMap<Id, Rc<OwnedPrimArray>>>,
+    pub vertex_distance_cache: RefCell<BTreeMap<(Id, Id, u64), Distance>>,
+    pub vertex_vector_cache: RefCell<BTreeMap<(Id, u64), Arc<OwnedPrimArray>>>,
     pub ef: usize, // Extension factor for search, controls exploration vs. exploitation
     pub ef_construction: usize, // Extension factor for construction, controls how many neighbors to consider when building the graph
 }
@@ -73,9 +72,9 @@ unsafe impl Sync for SearchMetadata {}
 unsafe impl Send for SearchMetadata {}
 
 impl SearchMetadata {
-    pub fn set_query_vector(&self, query: OwnedPrimArray) {
+    pub fn set_query_vector(&self, query: OwnedPrimArray, field_id: u64) {
         self.vertex_vector_cache
             .borrow_mut()
-            .insert(Id::unit_id(), Rc::new(query));
+            .insert((Id::unit_id(), field_id), Arc::new(query));
     }
 }