nodedb_lite/engine/vector/
nodes.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Binding a document's vector to its HNSW node, and releasing it.
4//!
5//! A document id holds at most one live node per index. Every insert path
6//! binds through [`bind_node`], which tombstones the node the id held before,
7//! so re-inserting an id replaces its vector instead of adding a second node.
8//! Every delete path releases through [`unbind_node`].
9
10use std::sync::Arc;
11
12use crate::engine::vector::resident::lock_resident_or_create;
13use crate::engine::vector::{HnswIndex, VectorState};
14use crate::error::LiteError;
15use crate::nodedb::lock_ext::LockExt;
16use crate::storage::engine::StorageEngine;
17
18/// Insert `embedding` into `index` as the node of `doc_id` in `index_key`,
19/// under the caller's lock on the index map. The node `doc_id` held before
20/// is tombstoned and its codec sidecar code dropped. Returns the new node id.
21///
22/// Fails when the index refuses the vector; the old node then stays live.
23pub(crate) fn bind_node<S: StorageEngine>(
24    state: &VectorState<S>,
25    index: &mut HnswIndex,
26    index_key: &str,
27    doc_id: &str,
28    embedding: Vec<f32>,
29) -> Result<u32, LiteError> {
30    let node = index.len() as u32;
31    index.insert(embedding)?;
32    let displaced = state
33        .vector_id_map
34        .lock_or_recover()
35        .bind(index_key, doc_id, node);
36    if let Some(old) = displaced {
37        index.delete(old);
38        drop_sidecar_code(state, index_key, old);
39    }
40    Ok(node)
41}
42
43/// Release the node of `doc_id` in `index_key`: unbind it, tombstone it in
44/// `index` when the index is resident, and drop its codec sidecar code.
45/// Returns the released node, or `None` when `doc_id` held none.
46pub(crate) fn unbind_node<S: StorageEngine>(
47    state: &VectorState<S>,
48    index: Option<&mut HnswIndex>,
49    index_key: &str,
50    doc_id: &str,
51) -> Option<u32> {
52    let node = state
53        .vector_id_map
54        .lock_or_recover()
55        .unbind_doc(index_key, doc_id)?;
56    if let Some(index) = index {
57        index.delete(node);
58    }
59    drop_sidecar_code(state, index_key, node);
60    Some(node)
61}
62
63/// Load or create `index_key`'s index and bind `embedding` as the node of
64/// `doc_id`, as [`bind_node`] does. Returns the new node id.
65///
66/// Fails when loading the index fails or the index refuses the vector.
67pub(crate) async fn upsert_node<S: StorageEngine>(
68    state: &Arc<VectorState<S>>,
69    index_key: &str,
70    doc_id: &str,
71    embedding: &[f32],
72) -> Result<u32, LiteError> {
73    let mut resident = lock_resident_or_create(state, index_key, embedding.len()).await?;
74    bind_node(
75        state,
76        resident.index(),
77        index_key,
78        doc_id,
79        embedding.to_vec(),
80    )
81}
82
83/// Encode `embedding` into `index_key`'s codec sidecar as `node`, installing
84/// the sidecar first when the index config calls for one.
85///
86/// Fails only when installing the sidecar fails. An encode error is logged
87/// and the row falls back to FP32 rerank. Call it without the index map
88/// locked: installing a sidecar trains on the index.
89pub(crate) fn encode_sidecar<S: StorageEngine>(
90    state: &VectorState<S>,
91    index_key: &str,
92    node: u32,
93    embedding: &[f32],
94) -> Result<(), LiteError> {
95    if !crate::engine::vector::sidecar::ensure_sidecar(state, index_key)? {
96        return Ok(());
97    }
98    let mut sidecars = state.codec_sidecars.lock_or_recover();
99    if let Some(sidecar) = sidecars.get_mut(index_key)
100        && let Err(e) = sidecar.encode_and_insert(node, embedding)
101    {
102        tracing::warn!(
103            index_key,
104            id = node,
105            error = %e,
106            "sidecar encode_and_insert failed; row falls back to FP32 rerank"
107        );
108    }
109    Ok(())
110}
111
112/// Drop `node`'s code from `index_key`'s codec sidecar, when one is installed.
113fn drop_sidecar_code<S: StorageEngine>(state: &VectorState<S>, index_key: &str, node: u32) {
114    if let Some(sidecar) = state.codec_sidecars.lock_or_recover().get_mut(index_key) {
115        sidecar.remove(node);
116    }
117}
nodedb_lite/engine/vector/nodes.rs

nodedb_lite/engine/vector/
nodes.rs
