我们从索引中随机采样两个，分别对应于向量 A 和 B，计算 n = A - B，并找到 A 和 B 的中点（point_on_plane）。存储超平面使用的结构包含系数（向量 n）和常量（n 和 point_on_plane 的点积），形式为：n(x-x0) = nx - nx0。我们可以求向量和 n 之间的点积，然后减去常数，就可以将向量放在超平面的“上方”或“下方”。由于树中的内部节点包含超平面定义，叶节点包含向量 ID，因此我们可以使用 ADT 对树进行类型检查：
impl<const N: usize> ANNIndex<N> {
fn build_hyperplane(
indexes: &Vec<usize>,
all_vecs: &Vec<Vector<N>>,
) -> (HyperPlane<N>, Vec<usize>, Vec<usize>) {
let sample: Vec<_> = indexes
.choose_multiple(&mut rand::thread_rng(), 2)
.collect();
// cartesian eq for hyperplane n * (x - x_0) = 0
// n (normal vector) is the coefs x_1 to x_n
let (a, b) = (*sample[0], *sample[1]);
let coefficients = all_vecs[a].subtract_from(&all_vecs[b]);
let point_on_plane = all_vecs[a].avg(&all_vecs[b]);
let constant = -coefficients.dot_product(&point_on_plane);
let hyperplane = HyperPlane::<N> {
coefficients,
constant,
};
let (mut above, mut below) = (vec![], vec![]);
for &id in indexes.iter() {
if hyperplane.point_is_above(&all_vecs[id]) {
above.push(id)
} else {
below.push(id)
};
}
return (hyperplane, above, below);
}
}
impl<const N: usize> ANNIndex<N> {
fn deduplicate(
vectors: &Vec<Vector<N>>,
ids: &Vec<i32>,
dedup_vectors: &mut Vec<Vector<N>>,
ids_of_dedup_vectors: &mut Vec<i32>,
) {
let mut hashes_seen = HashSet::new();
for i in 1..vectors.len() {
let hash_key = vectors[i].to_hashkey();
if !hashes_seen.contains(&hash_key) {
hashes_seen.insert(hash_key);
dedup_vectors.push(vectors[i]);
ids_of_dedup_vectors.push(ids[i]);
}
}
}
pub fn build_index(
num_trees: i32,
max_size: i32,
vecs: &Vec<Vector<N>>,
vec_ids: &Vec<i32>,
) -> ANNIndex<N> {
let (mut unique_vecs, mut ids) = (vec![], vec![]);
Self::deduplicate(vecs, vec_ids, &mut unique_vecs, &mut ids);
// Trees hold an index into the [unique_vecs] list which is not
// necessarily its id, if duplicates existed
let all_indexes: Vec<usize> = (0..unique_vecs.len()).collect();
let trees: Vec<_> = (0..num_trees)
.map(|_| Self::build_a_tree(max_size, &all_indexes, &unique_vecs))
.collect();
return ANNIndex::<N> {
trees,
ids,
values: unique_vecs,
};
}
}
200 行 Rust 代码编写一个向量搜索库，代码已开源！
查询时间
索引建立之后，下一步我们如何使用它搜索某棵树中输入向量的 K 个近似最近邻？我们的超平面存储在非叶节点处，因此我们可以从树的根开始搜索：“这个向量是在这个超平面的上方还是下方？”我们可以用点积计算，复杂度为 O(D)。接下来，我们可以根据响应，递归搜索左子树或右子树，直到找到叶节点。请记住，叶节点中存储的向量数最大为“最大节点数”，这些向量位于输入向量的近似邻域中（它们将落入所有超平面下超空间的同一分区中）。如果这个叶节点的向量索引数超过 K，我们就需要根据到输入向量的 L2 距离，对所有这些向量进行排序，并返回最接近的 K 个向量。
假设我们的索引最后得到了一棵平衡树，对于维度 D、向量数量 N 和最大节点数 M << N，搜索耗费的时间为 O(Dlog(N) + DM + Mlog(M))，最坏的情况下我们需要比较 log(N) 个超平面（即树的高度）才能找到叶节点，每次比较耗费的时间为 O(D) 个点积，计算一个叶节点的所有候选向量的 L2 距离需要 O(DM)，最后在 O(Mlog(M)) 时间内对它们进行排序，并返回 前 K 个。
但是，如果我们找到的叶节点少于 K 个向量，会发生什么情况？如果最大节点数太小，或超平面分割不均匀，则会导致某个子树中的向量非常少。为了解决这个问题，我们可以在树搜索中添加一些简单的回溯功能。例如，如果返回的候选向量数量不足，我们可以在内部节点再递归调用一次，访问另一个分枝。代码如下：
1impl<const N: usize> ANNIndex<N> {
2 fn tree_result(
3 query: Vector<N>,
4 n: i32,
5 tree: &Node<N>,
6 candidates: &mut HashSet<usize>,
7 ) -> i32 {
8 // take everything in node, if still needed, take from alternate subtree
9 match tree {
10 Node::Leaf(box_leaf) => {
11 let leaf_values = &(box_leaf.0);
12 let num_candidates_found = min(n as usize, leaf_values.len());
13 for i in 0..num_candidates_found {
14 candidates.insert(leaf_values[i]);
15 }
16 return num_candidates_found as i32;
17 }
18 Node::Inner(inner) => {
19 let above = (*inner).hyperplane.point_is_above(&query);
20 let (main, backup) = match above {
21 true => (&(inner.right_node), &(inner.left_node)),
22 false => (&(inner.left_node), &(inner.right_node)),
23 };
24 match Self::tree_result(query, n, main, candidates) {
25 k if k < n => {
26 k + Self::tree_result(query, n - k, backup, candidates)
27 }
28 k => k,
29 }
30 }
31 }
32 }
33}
