weavatrix_search_vector/
graph.rs

1use crate::error::SearchError;
2use crate::hit::SearchHit;
3use crate::hnsw::VectorIndex;
4use crate::parallel;
5
6/// Directed nearest-neighbor candidate edge.
7#[derive(Debug, Clone, Copy, PartialEq)]
8pub struct NeighborEdge {
9    pub source: u64,
10    pub target: u64,
11    pub distance: f32,
12}
13
14/// Deterministic directed KNN candidate graph in compressed-row form.
15///
16/// This is deliberately not a semantic or domain graph: consumers decide how
17/// to interpret, threshold, symmetrize, or add provenance to candidate edges.
18#[derive(Debug, Clone, PartialEq)]
19pub struct KnnGraph {
20    nodes: Vec<u64>,
21    offsets: Vec<usize>,
22    edges: Vec<NeighborEdge>,
23}
24
25impl KnnGraph {
26    /// Builds an approximate directed KNN graph with bounded query workers.
27    ///
28    /// # Errors
29    ///
30    /// Returns a typed query, worker, capacity, or allocation error.
31    pub fn build(index: &VectorIndex, neighbors: usize) -> Result<Self, SearchError> {
32        Self::build_with(index, neighbors, false)
33    }
34
35    /// Builds an exact directed KNN graph for oracle and small-corpus use.
36    ///
37    /// # Errors
38    ///
39    /// Returns a typed query, worker, capacity, or allocation error.
40    pub fn build_exact(index: &VectorIndex, neighbors: usize) -> Result<Self, SearchError> {
41        Self::build_with(index, neighbors, true)
42    }
43
44    #[must_use]
45    pub fn len(&self) -> usize {
46        self.nodes.len()
47    }
48
49    #[must_use]
50    pub fn is_empty(&self) -> bool {
51        self.nodes.is_empty()
52    }
53
54    #[must_use]
55    pub fn edge_count(&self) -> usize {
56        self.edges.len()
57    }
58
59    #[must_use]
60    pub fn nodes(&self) -> impl ExactSizeIterator<Item = u64> + '_ {
61        self.nodes.iter().copied()
62    }
63
64    #[must_use]
65    pub fn neighbors(&self, source: u64) -> Option<&[NeighborEdge]> {
66        let index = self.nodes.binary_search(&source).ok()?;
67        Some(&self.edges[self.offsets[index]..self.offsets[index + 1]])
68    }
69
70    #[must_use]
71    pub fn edges(&self) -> impl ExactSizeIterator<Item = NeighborEdge> + '_ {
72        self.edges.iter().copied()
73    }
74
75    fn build_with(index: &VectorIndex, neighbors: usize, exact: bool) -> Result<Self, SearchError> {
76        let nodes = index.keys().collect::<Vec<_>>();
77        let per_node = neighbors.min(index.len().saturating_sub(1));
78        let queries = nodes
79            .iter()
80            .map(|key| {
81                index
82                    .vector(*key)
83                    .expect("index key iterator resolves to a vector")
84            })
85            .collect::<Vec<_>>();
86        let rows = if exact {
87            parallel::search_batch(&queries, index.config().query_threads, |query| {
88                index.search_exact(query, per_node.saturating_add(1))
89            })?
90        } else {
91            index.search_batch(&queries, per_node.saturating_add(1))?
92        };
93        Self::from_rows(nodes, rows, per_node)
94    }
95
96    fn from_rows(
97        nodes: Vec<u64>,
98        rows: Vec<Vec<SearchHit>>,
99        per_node: usize,
100    ) -> Result<Self, SearchError> {
101        let capacity = nodes
102            .len()
103            .checked_mul(per_node)
104            .ok_or(SearchError::CapacityOverflow)?;
105        let mut offsets = Vec::new();
106        offsets
107            .try_reserve_exact(nodes.len().saturating_add(1))
108            .map_err(|_| SearchError::AllocationFailed)?;
109        let mut edges = Vec::new();
110        edges
111            .try_reserve_exact(capacity)
112            .map_err(|_| SearchError::AllocationFailed)?;
113        offsets.push(0);
114        for (source, row) in nodes.iter().copied().zip(rows) {
115            edges.extend(
116                row.into_iter()
117                    .filter(|hit| hit.key != source)
118                    .take(per_node)
119                    .map(|hit| NeighborEdge {
120                        source,
121                        target: hit.key,
122                        distance: hit.distance,
123                    }),
124            );
125            offsets.push(edges.len());
126        }
127        Ok(Self {
128            nodes,
129            offsets,
130            edges,
131        })
132    }
133}
weavatrix_search_vector/graph.rs

weavatrix_search_vector/
graph.rs
