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    pub(crate) fn new(_id: u32, max_level: usize) -> Self {
23        let bottom_neighbors = Arc::new(Vec::new());
24        let level_neighbors = (0..max_level).map(|_| Arc::new(Vec::new())).collect();
25        let level_neighbors_ranked = (0..max_level).map(|_| Vec::new()).collect();
26        Self {
27            bottom_neighbors,
28            level_neighbors,
29            level_neighbors_ranked,
30        }
31    }
32
33    pub(crate) fn add_neighbor(&mut self, v: u32, dist: OrderedFloat, level: u16) {
34        self.level_neighbors_ranked[level as usize].push(OrderedNode { dist, id: v });
35    }
36
37    pub(crate) fn update_from_ranked_neighbors(&mut self, level: u16) {
38        let level_index = level as usize;
39        self.level_neighbors[level_index] = Arc::new(
40            self.level_neighbors_ranked[level_index]
41                .iter()
42                .map(|ordered_node| ordered_node.id)
43                .collect(),
44        );
45        if level == 0 {
46            self.bottom_neighbors = self.level_neighbors[0].clone();
47        }
48    }
49
50    pub(crate) fn cutoff(&self, level: u16, max_size: usize) -> OrderedFloat {
51        let neighbors = &self.level_neighbors_ranked[level as usize];
52        if neighbors.len() < max_size {
53            OrderedFloat(f32::INFINITY)
54        } else {
55            neighbors.last().unwrap().dist
56        }
57    }
58}
59
60#[derive(Debug)]
61pub struct GraphBuilderStats {
62    #[allow(dead_code)]
63    pub num_nodes: usize,
64    #[allow(dead_code)]
65    pub max_edges: usize,
66    #[allow(dead_code)]
67    pub mean_edges: f32,
68    #[allow(dead_code)]
69    pub mean_distance: f32,
70}