FNNTW: 西部最快的最近邻

最快的最近邻搜索库（西部）是一个正在开发的 kD-tree 库，目标是成为现有性能最高的并行 kNN 库之一。

kD-trees 构建和查询过程中的几个关键组件使其性能得以提升。

1. 构建树

许多 kD-Tree 构建实现虽然很容易并行化，但并不是并行的。这里我们讨论了针对单个子树以及跨子树的并行化策略。

a. 使用 Quickselect

通过在 Rust 的 core::slice 中使用类似 select_nth_unstable_by 的 quickselect，而不是使用平均或其他中值算法，我们可以在相同的 O(N) 算法中免费获得 left 和 right 子集，而不必再进行 O(N) 的比较以获得 left 和 right 箱。因此，构建仍然是 O(N log(N))，但去除了在许多其他库中每个分割时都会出现的整个 O(N) 操作。对于节点数 >=100,000 的树，使用自制的并行近似中值查找器来找到子树的分割点。这极大地加快了大型树的构建速度，因为 ≈95% 的顺序树构建时间都花费在查找中值上。

b. 并行构建

kD-tree 的每个子树都是一个独立的 kD-tree，因此在每个分割时，可以并行构建每个子树直到某个最小子树大小。这在这里通过用户指定的 par_split_level 来实现，它是并行开始时的树深度。参见 与其他代码的基准测试 中的注释，了解推荐值。

2. 不安全访问

因为我们知道所有数组的形状（即树的维度）在编译时，并且我们在运行时知道树的大小和拓扑结构，所以代码中广泛使用了unsafe方法get_unchecked和get_unchecked_mut。这意味着几乎不做边界检查。

3. 内存分配器

我们测试了许多内存分配器，包括默认分配器jemalloc、tcmalloc、snmalloc、mimalloc、rpmalloc等，并发现tcmalloc的性能最佳。

与其他代码的基准测试

这个库的目的是由作者用来在宇宙学中计算总结统计量。因此，选择的基准测试参数接近于分析宇宙学模拟输出时可能使用的参数。在这种情况下，通常使用许多子样本或模拟盒子。因此，组合构建+查询时间是重要的，因为在分析中可能会构建多个不同的树。我们使用了

• 一个包含100,000个均匀随机点在单位立方体中的模拟数据集。
• 一个包含1,000,000个均匀随机点在单位立方体中的查询集。

在超过100个数据集和查询点的实现中，使用AMD EPYC 7502P和48个线程，我们得到了以下平均构建（串行）和1NN查询时间的结果。结果按照组合构建和查询时间排序。

由于FNNTW具有并行构建功能，在AMD EPYC 7502P上（在split_level = 2时）构建时间可以低至5毫秒（单精度下小于5毫秒）。由于并行和原子操作的开销在数据点数量少时减慢了构建速度，因此通过Tree:new(..)和Tree::new_parallel(..)提供了并行构建和非并行构建。后者接受上述参数par_split_level，即停止并行的分割深度。尽管对于我们的O(1e5)点应用，我们发现par_split_level = 2具有最大的改进，但我们预计最佳par_split_level将随着数据集大小的增加而增加。例如，对于O(1e8)点的树大小，我们发现在par_split_level = 4（16线程）时达到峰值性能。