Skip to main content

lance_index/vector/
hnsw.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! HNSW graph implementation.
5//!
6//! Hierarchical Navigable Small World (HNSW).
7//!
8
9use arrow_schema::{DataType, Field};
10use itertools::Itertools;
11use lance_core::deepsize::DeepSizeOf;
12use serde::{Deserialize, Serialize};
13
14use self::builder::HnswBuildParams;
15use super::graph::OrderedNode;
16use super::storage::VectorStore;
17
18pub mod builder;
19pub mod index;
20pub mod online;
21
22pub use builder::HNSW;
23pub use index::HNSWIndex;
24pub use online::OnlineHnswBuilder;
25
26const HNSW_TYPE: &str = "HNSW";
27const VECTOR_ID_COL: &str = "__vector_id";
28const POINTER_COL: &str = "__pointer";
29
30use std::sync::LazyLock;
31
32/// POINTER field.
33///
34pub static POINTER_FIELD: LazyLock<Field> =
35    LazyLock::new(|| Field::new(POINTER_COL, DataType::UInt32, true));
36
37/// Id of the vector in the `VectorStorage`.
38pub static VECTOR_ID_FIELD: LazyLock<Field> =
39    LazyLock::new(|| Field::new(VECTOR_ID_COL, DataType::UInt32, true));
40
41#[derive(Debug, Clone, Serialize, Deserialize, DeepSizeOf)]
42pub struct HnswMetadata {
43    pub entry_point: u32,
44    pub params: HnswBuildParams,
45    pub level_offsets: Vec<usize>,
46}
47
48impl Default for HnswMetadata {
49    fn default() -> Self {
50        let params = HnswBuildParams::default();
51        let level_offsets = vec![0; params.max_level as usize];
52        Self {
53            entry_point: 0,
54            params,
55            level_offsets,
56        }
57    }
58}
59
60/// Algorithm 4 in the HNSW paper.
61///
62/// # NOTE
63/// The results are not ordered.
64pub(crate) fn select_neighbors_heuristic(
65    storage: &impl VectorStore,
66    candidates: &[OrderedNode],
67    k: usize,
68) -> Vec<OrderedNode> {
69    if candidates.len() <= k {
70        return candidates.iter().cloned().collect_vec();
71    }
72
73    select_neighbors_heuristic_owned(storage, candidates.to_vec(), k)
74}
75
76pub(crate) fn select_neighbors_heuristic_owned(
77    storage: &impl VectorStore,
78    mut candidates: Vec<OrderedNode>,
79    k: usize,
80) -> Vec<OrderedNode> {
81    if candidates.len() <= k {
82        return candidates;
83    }
84
85    candidates.sort_unstable();
86
87    let mut results: Vec<OrderedNode> = Vec::with_capacity(k);
88    for u in candidates.iter() {
89        if results.len() >= k {
90            break;
91        }
92
93        if results.is_empty() || storage.prefers_candidate(u, &results) {
94            results.push(u.clone());
95        }
96    }
97    results
98}