Skip to main content

lance_index/vector/graph/
builder.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use deepsize::DeepSizeOf;
5
6use super::OrderedFloat;
7use super::OrderedNode;
8use std::sync::Arc;
9
10/// GraphNode during build.
11///
12/// WARNING: Internal API,  API stability is not guaranteed
13#[derive(Debug, Clone, DeepSizeOf)]
14pub struct GraphBuilderNode {
15    /// neighbors of each level of the node.
16    pub(crate) bottom_neighbors: Arc<Vec<u32>>,
17    pub(crate) level_neighbors: Vec<Arc<Vec<u32>>>,
18    pub(crate) level_neighbors_ranked: Vec<Vec<OrderedNode>>,
19}
20
21impl GraphBuilderNode {
22    /// Construct a node from already-computed parts. Used by the online
23    /// builder to convert its mutable state into an immutable `HNSW`.
24    pub(crate) fn from_parts(
25        level_neighbors: Vec<Arc<Vec<u32>>>,
26        level_neighbors_ranked: Vec<Vec<OrderedNode>>,
27        bottom_neighbors: Arc<Vec<u32>>,
28    ) -> Self {
29        Self {
30            bottom_neighbors,
31            level_neighbors,
32            level_neighbors_ranked,
33        }
34    }
35
36    pub(crate) fn new(_id: u32, max_level: usize) -> Self {
37        let bottom_neighbors = Arc::new(Vec::new());
38        let level_neighbors = (0..max_level).map(|_| Arc::new(Vec::new())).collect();
39        let level_neighbors_ranked = (0..max_level).map(|_| Vec::new()).collect();
40        Self {
41            bottom_neighbors,
42            level_neighbors,
43            level_neighbors_ranked,
44        }
45    }
46
47    pub(crate) fn add_neighbor(&mut self, v: u32, dist: OrderedFloat, level: u16) {
48        self.level_neighbors_ranked[level as usize].push(OrderedNode { dist, id: v });
49    }
50
51    pub(crate) fn update_from_ranked_neighbors(&mut self, level: u16) {
52        let level_index = level as usize;
53        self.level_neighbors[level_index] = Arc::new(
54            self.level_neighbors_ranked[level_index]
55                .iter()
56                .map(|ordered_node| ordered_node.id)
57                .collect(),
58        );
59        if level == 0 {
60            self.bottom_neighbors = self.level_neighbors[0].clone();
61        }
62    }
63
64    pub(crate) fn cutoff(&self, level: u16, max_size: usize) -> OrderedFloat {
65        let neighbors = &self.level_neighbors_ranked[level as usize];
66        if neighbors.len() < max_size {
67            OrderedFloat(f32::INFINITY)
68        } else {
69            neighbors.last().unwrap().dist
70        }
71    }
72}
73
74#[derive(Debug)]
75pub struct GraphBuilderStats {
76    #[allow(dead_code)]
77    pub num_nodes: usize,
78    #[allow(dead_code)]
79    pub max_edges: usize,
80    #[allow(dead_code)]
81    pub mean_edges: f32,
82    #[allow(dead_code)]
83    pub mean_distance: f32,
84}