uqa-storage 0.1.11

Document store, inverted index, IVF/HNSW vectors, B-tree, R*Tree, catalog
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Diversity-aware neighbor selection and reciprocal degree pruning.

use std::collections::BTreeSet;

use super::metric::distance;
use super::search::Candidate;
use super::types::{HNSWIndex, NodeId};

impl HNSWIndex {
    pub(super) fn ensure_layer_zero_backbone(&self, node_id: NodeId, neighbors: &mut Vec<NodeId>) {
        let Some(previous_id) = node_id.checked_sub(1) else {
            return;
        };
        if self.nodes.contains_key(&previous_id) && !neighbors.contains(&previous_id) {
            neighbors.push(previous_id);
            neighbors.sort_unstable();
        }
    }

    pub(super) fn prune_node(&mut self, node_id: NodeId, layer: usize) {
        let Some(node) = self.nodes.get(&node_id) else {
            return;
        };
        let Some(current) = node.neighbors.get(layer).cloned() else {
            return;
        };
        let protected = current
            .iter()
            .copied()
            .filter(|neighbor_id| layer == 0 && node_id.abs_diff(*neighbor_id) == 1)
            .collect::<BTreeSet<_>>();
        let limit = self.max_connections(layer);
        let mut selected = protected.iter().copied().collect::<Vec<_>>();
        selected.extend(
            self.select_neighbors(
                &node.normalized_vector,
                current
                    .iter()
                    .copied()
                    .filter(|neighbor_id| !protected.contains(neighbor_id)),
                limit.saturating_sub(selected.len()),
                Some(node_id),
            ),
        );
        selected.sort_unstable();
        let selected_set = selected.iter().copied().collect::<BTreeSet<_>>();
        let removed = current
            .into_iter()
            .filter(|neighbor| !selected_set.contains(neighbor))
            .collect::<Vec<_>>();
        if let Some(node) = self.nodes.get_mut(&node_id) {
            node.neighbors[layer] = selected;
            self.dirty_nodes.insert(node_id);
        }
        for removed_id in removed {
            if let Some(neighbor) = self.nodes.get_mut(&removed_id) {
                if let Some(reverse) = neighbor.neighbors.get_mut(layer) {
                    reverse.retain(|candidate| *candidate != node_id);
                    self.dirty_nodes.insert(removed_id);
                }
            }
        }
    }

    pub(super) fn select_neighbors(
        &self,
        query: &[f32],
        candidates: impl IntoIterator<Item = NodeId>,
        limit: usize,
        exclude: Option<NodeId>,
    ) -> Vec<NodeId> {
        let mut candidates = candidates
            .into_iter()
            .filter(|node_id| Some(*node_id) != exclude)
            .filter_map(|node_id| {
                self.nodes.get(&node_id).map(|node| Candidate {
                    distance: distance(query, &node.normalized_vector),
                    node_id,
                })
            })
            .collect::<Vec<_>>();
        candidates.sort();
        candidates.dedup_by_key(|candidate| candidate.node_id);
        let mut selected = Vec::with_capacity(limit.min(candidates.len()));
        let mut rejected = Vec::new();
        for candidate in candidates {
            let Some(candidate_node) = self.nodes.get(&candidate.node_id) else {
                continue;
            };
            let diverse = selected.iter().all(|selected_id| {
                self.nodes.get(selected_id).is_none_or(|selected_node| {
                    distance(
                        &candidate_node.normalized_vector,
                        &selected_node.normalized_vector,
                    ) > candidate.distance
                })
            });
            if diverse && selected.len() < limit {
                selected.push(candidate.node_id);
            } else {
                rejected.push(candidate.node_id);
            }
        }
        for candidate in rejected {
            if selected.len() >= limit {
                break;
            }
            selected.push(candidate);
        }
        selected.sort_unstable();
        selected
    }
}