nodedb_lite/engine/vector/
state.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Shared runtime state for HNSW vector search on Lite.
4//!
5//! Held as `Arc<VectorState<S>>` on both `NodeDbLite<S>` (user-facing
6//! entry points) and `LiteQueryEngine<S>` (PhysicalPlan executor) so
7//! the visitor pipeline can run vector ops without re-architecting the
8//! engine boundary.
9
10use std::collections::{HashMap, HashSet};
11use std::sync::{Arc, Mutex};
12
13use nodedb_mem::ScopedMemory;
14use nodedb_types::collection_config::VectorPrimaryConfig;
15use nodedb_vector::rerank::CodecSidecar;
16
17use crate::engine::vector::HnswIndex;
18use crate::engine::vector::id_map::VectorIdMap;
19use crate::storage::engine::StorageEngine;
20
21pub struct VectorState<S: StorageEngine> {
22    pub(crate) hnsw_indices: Mutex<HashMap<String, HnswIndex>>,
23    /// Per-index node ↔ document id bindings. Lock order: `hnsw_indices`
24    /// first, then this.
25    pub(crate) vector_id_map: Mutex<VectorIdMap>,
26    pub(crate) search_ef: usize,
27    pub(crate) storage: Arc<S>,
28    /// index_key → trained codec sidecar (populated by S2.a.11).
29    pub(crate) codec_sidecars: Arc<Mutex<HashMap<String, CodecSidecar>>>,
30    /// Per-(index_key) collection config — populated when a collection is
31    /// registered via DDL (C2c will wire that). Lookup is best-effort:
32    /// callers that don't find an entry default to F32 storage, matching
33    /// the previous behavior.
34    pub(crate) per_index_config: Arc<Mutex<HashMap<String, VectorPrimaryConfig>>>,
35    /// Index keys whose stored checkpoint exists but cannot be turned into a
36    /// usable index — unreadable checkpoint, or a segment that cannot serve the
37    /// graph with no durable vectors to rebuild from.
38    ///
39    /// Without this, `ensure_index_loaded` gives up WITHOUT caching anything, so
40    /// every later search on that collection repeats the entire cost — read the
41    /// checkpoint, deserialize the full graph, open and validate the segment,
42    /// scan the durable rows — and still finds nothing. On a collection with
43    /// thousands of nodes that turns one unusable segment into an operation that
44    /// pins a core indefinitely while reporting no progress.
45    ///
46    /// This is a NEGATIVE cache for the load path only. It is not consulted once
47    /// the collection is present in `hnsw_indices`, so a later insert (which
48    /// creates the index through `resident::lock_resident_or_create`) resolves
49    /// the collection normally without anything here needing to be cleared.
50    pub(crate) unloadable: Mutex<HashSet<String>>,
51    /// Index keys eviction dropped from `hnsw_indices` after writing their
52    /// checkpoint. An index absent from memory and named here must be loaded
53    /// back before any use: creating an empty one in its place hides every
54    /// vector it held. The lazy loader clears the mark once it has loaded
55    /// the index or found it unloadable. Lock order: `hnsw_indices` first.
56    pub(crate) evicted: Mutex<HashSet<String>>,
57    /// Memory scope for codec sidecar and rerank allocations owned by this state.
58    pub(crate) memory: ScopedMemory,
59}
60
61/// Restored state for [`VectorState::from_restored`].
62///
63/// Grouped into a struct because the constructor otherwise carries five
64/// unrelated arguments.
65pub struct RestoredVectorState<S: StorageEngine> {
66    pub storage: Arc<S>,
67    pub search_ef: usize,
68    pub indices: HashMap<String, HnswIndex>,
69    pub id_map: VectorIdMap,
70    pub memory: ScopedMemory,
71}
72
73impl<S: StorageEngine> VectorState<S> {
74    pub fn new(storage: Arc<S>, search_ef: usize, memory: ScopedMemory) -> Self {
75        Self {
76            hnsw_indices: Mutex::new(HashMap::new()),
77            vector_id_map: Mutex::new(VectorIdMap::new()),
78            search_ef,
79            storage,
80            codec_sidecars: Arc::new(Mutex::new(HashMap::new())),
81            per_index_config: Arc::new(Mutex::new(HashMap::new())),
82            unloadable: Mutex::new(HashSet::new()),
83            evicted: Mutex::new(HashSet::new()),
84            memory,
85        }
86    }
87
88    pub fn from_restored(restored: RestoredVectorState<S>) -> Self {
89        Self {
90            hnsw_indices: Mutex::new(restored.indices),
91            vector_id_map: Mutex::new(restored.id_map),
92            search_ef: restored.search_ef,
93            storage: restored.storage,
94            codec_sidecars: Arc::new(Mutex::new(HashMap::new())),
95            per_index_config: Arc::new(Mutex::new(HashMap::new())),
96            unloadable: Mutex::new(HashSet::new()),
97            evicted: Mutex::new(HashSet::new()),
98            memory: restored.memory,
99        }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::query::engine::{test_governor, test_scoped_memory};
107    use crate::storage::pagedb_storage::PagedbStorageMem;
108
109    #[tokio::test]
110    async fn per_index_config_starts_empty() {
111        let storage = Arc::new(
112            PagedbStorageMem::open_in_memory()
113                .await
114                .expect("in-memory pagedb"),
115        );
116        let memory = test_scoped_memory(&test_governor(), nodedb_mem::EngineId::Vector);
117        let state = VectorState::new(storage, 100, memory);
118        let configs = state.per_index_config.lock().expect("lock");
119        assert!(
120            configs.is_empty(),
121            "per_index_config must be empty on construction"
122        );
123    }
124}
nodedb_lite/engine/vector/state.rs

nodedb_lite/engine/vector/
state.rs
