jvector - 高性能向量索引库支持图索引和大规模搜索

近似最近邻搜索简介

在高维空间中,精确的最近邻搜索(k-最近邻或KNN)的计算成本过高,因为在二维或三维中有效的搜索空间分割方法(如四叉树或k-d树)在高维中退化为线性扫描。这是所谓的"维度灾难"的一个方面。

对于大型数据集,以对数时间获得近似答案几乎总是比以线性时间获得精确答案更有用。这被简称为ANN(近似最近邻)搜索。

ANN索引主要分为两大类:

基于分区的索引,如LSH、IVF或SCANN
图索引,如HNSW或DiskANN

图索引通常更容易实现且速度更快,更重要的是它们可以增量构建和更新。这使得它们比只能在预先完全指定的静态数据集上工作的分区方法更适合作为通用索引。这就是为什么所有主要的商业向量索引都使用图方法的原因。

JVector是DiskANN家族中的一个图索引。

JVector架构

JVector是一个基于图的索引,在DiskANN设计的基础上增加了可组合的扩展。

JVector实现了一个具有非阻塞并发控制的单层图,允许构建过程随着核心数量的增加而线性扩展: 随着线程数增加,JVector呈线性扩展

图以每个节点的磁盘邻接表表示,内联存储额外数据以支持两遍搜索。第一遍由内存中保存的向量的有损压缩表示驱动,第二遍由从磁盘读取的更精确表示驱动。第一遍搜索可以通过以下方式执行:

乘积量化(PQ),可选择带有各向异性加权
二进制量化(BQ)
融合ADC,其中PQ码本被转置并与图邻接表内联写入

第二遍搜索可以使用:

全精度float32向量

这种两遍设计减少了内存使用并降低了延迟,同时保持了准确性。

此外,JVector独特地提供了使用两遍搜索构建索引本身的能力,允许构建大于内存的索引: 更大的索引

这一点很重要,因为它允许您在单个索引中利用对数级搜索,而不是溢出到多个索引的线性时间合并结果。

JVector步骤解析

所有代码示例都来自JVector源代码库中的SiftSmall,其中还包括siftsmall数据集。只需在IDE中导入项目并点击运行即可尝试!

步骤1:在内存中构建和查询索引

首先是代码:

public static void siftInMemory(ArrayList<VectorFloat<?>> baseVectors) throws IOException {
    // 从第一个向量推断维度
    int originalDimension = baseVectors.get(0).length();
    // 将原始向量包装在RandomAccessVectorValues中
    RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(baseVectors, originalDimension);

    // 使用原始的内存中向量的分数提供者
    BuildScoreProvider bsp = BuildScoreProvider.randomAccessScoreProvider(ravv, VectorSimilarityFunction.EUCLIDEAN);
    try (GraphIndexBuilder builder = new GraphIndexBuilder(bsp,
                                                           ravv.dimension(),
                                                           16, // 图度
                                                           100, // 构建搜索深度
                                                           1.2f, // 构建过程中允许度溢出的因子
                                                           1.2f)) // 放宽邻居多样性要求的因子
    {
        // 构建索引(在内存中)
        OnHeapGraphIndex index = builder.build(ravv);

        // 搜索随机向量
        VectorFloat<?> q = randomVector(originalDimension);
        SearchResult sr = GraphSearcher.search(q,
                                               10, // 结果数量
                                               ravv, // 我们搜索的向量,用于评分
                                               VectorSimilarityFunction.EUCLIDEAN, // 评分方式
                                               index,
                                               Bits.ALL); // 要考虑的有效序号
        for (SearchResult.NodeScore ns : sr.getNodes()) {
            System.out.println(ns);
        }
    }
}

注释:

所有索引都假设您有一个具有一致、固定维度(float32组件数量)的向量源。
向量源通常表示为RandomAccessVectorValues的子类,它围绕getVector / getVectorInto提供了一个简单的API。请注意在多线程上下文中的isValueShared();作为经验法则,内存中的RAVV不会使用共享值,而从磁盘读取的RAVV会(作为一种优化,避免为每次调用分配新的堆上向量)。
您不必一次性向索引提供所有向量,但由于这在原型设计时是一种常见场景,因此提供了一种便捷方法来执行此操作。我们稍后将介绍如何增量构建索引。
对于溢出Builder参数,内存构建的最佳值约为1.2,磁盘上构建约为1.5。(允许的溢出越多,需要重新计算最佳边的次数就越少,但每次搜索时需要咨询的邻居就越多。)
alpha参数控制边距距离和多样性之间的权衡;通常对于高维向量,1.2就足够了;对于2D或3D数据集,建议使用2.0。有关更多详细信息,请参阅DiskANN论文。
GraphSearcher的Bits参数旨在根据外部谓词控制结果集,在本教程中不会使用。

步骤2:对GraphSearcher的更多控制

保持Builder不变,更新后的搜索代码如下:

        // 使用GraphSearcher和SearchScoreProvider搜索随机向量
        VectorFloat<?> q = randomVector(originalDimension);
        try (GraphSearcher searcher = new GraphSearcher(index)) {
            SearchScoreProvider ssp = SearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, ravv);
            SearchResult sr = searcher.search(ssp, 10, Bits.ALL);
            for (SearchResult.NodeScore ns : sr.getNodes()) {
                System.out.println(ns);
            }
        }

注释:

Searcher的分配成本适中,因为在构造时需要初始化大量内部状态。因此,JVector支持搜索器池化(例如,使用ExplicitThreadLocal,下面会介绍)。
在管理GraphSearcher实例时,您还负责构造一个SearchScoreProvider,您可以将其视为:给定一个查询向量,告诉JVector如何计算索引中其他节点的分数。SearchScoreProvider可以是精确的,如这里所示,或者是ApproximateScoreFunction和Reranker的组合,下面会介绍。

步骤3:测量召回率

如果一个速度极快的向量索引无法返回准确的结果,那么它就没有多大用处。作为健全性检查,SiftSmall包含一个辅助方法_testRecall_。将其连接到我们的代码主要涉及将SearchScoreProvider转换为工厂lambda:

        Function<VectorFloat<?>, SearchScoreProvider> sspFactory = q -> SearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, ravv);
        testRecall(index, queryVectors, groundTruth, sspFactory);

如果运行代码,您会看到每次召回率略有不同(由testRecall打印):

(OnHeapGraphIndex) Recall: 0.9898
...
(OnHeapGraphIndex) Recall: 0.9890

考虑到所创建索引的近似性质以及build调用的多线程并发引入的扰动,这是预期的结果。

步骤4:将索引写入磁盘并从磁盘加载

代码:

    Path indexPath = Files.createTempFile("siftsmall", ".inline");
    try (GraphIndexBuilder builder = new GraphIndexBuilder(bsp, ravv.dimension(), 16, 100, 1.2f, 1.2f)) {
        // 构建索引(在内存中)
        OnHeapGraphIndex index = builder.build(ravv);
        // 使用默认选项将索引写入磁盘
        OnDiskGraphIndex.write(index, ravv, indexPath);
    }

    // 磁盘上的索引需要一个ReaderSupplier(而不仅仅是Reader),因为我们希望它为搜索打开额外的读取器
    ReaderSupplier rs = new SimpleMappedReaderSupplier(indexPath);
    OnDiskGraphIndex index = OnDiskGraphIndex.load(rs);
    // 根据(精确计算的)基准真值测量我们的召回率
    Function<VectorFloat<?>, SearchScoreProvider> sspFactory = q -> SearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, ravv);
    testRecall(index, queryVectors, groundTruth, sspFactory);

注释：

我们可以通过单个方法调用将内存中构建的索引（如本例）写入磁盘。
加载和搜索磁盘上的索引需要一个ReaderSupplier，它提供RandomAccessReader对象。RandomAccessReader接口旨在由使用项目进行扩展。例如，DataStax Astra实现了一个由Cassandra块缓存支持的RandomAccessReader。JVector提供了两种现成的实现。
- SimpleMappedReader：使用FileChannel.map实现，这意味着它与所有可运行JVector的Java版本兼容，但也意味着它限制为2GB文件大小。SimpleMappedReader主要用于示例代码。
- MemorySegmentReader：使用较新的MemorySegment API实现，没有文件大小限制，但仅限于Java 22+。（实际的MemorySegmentReader代码与Java 20+兼容，但为了方便，我们将其保留在22+模块中。有兴趣的读者可以重构构建以改进这一点。）如果您没有特殊需求，建议在生产环境中使用MemorySegmentReader。

第5步：在搜索中使用压缩向量

使用乘积量化压缩向量的方法如下：

        // 计算并将压缩向量写入磁盘
        Path pqPath = Files.createTempFile("siftsmall", ".pq");
        try (DataOutputStream out = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(pqPath)))) {
            // 使用PQ压缩原始向量。这代表了128 * 4 / 16 = 32倍的压缩比
            ProductQuantization pq = ProductQuantization.compute(ravv,
                                                                 16, // 子空间数量
                                                                 256, // 每个子空间的质心数量
                                                                 true); // 对数据集进行中心化
            ByteSequence<?>[] compressed = pq.encodeAll(ravv);
            // 将压缩向量写入磁盘
            PQVectors pqv = new PQVectors(pq, compressed);
            pqv.write(out);
        }

JVector也支持二进制量化，但由于BQ对搜索精度的影响很大，所以通常不如PQ有用。

然后我们可以通过从PQVectors获取快速ApproximateScoreFunction，并从索引View获取Reranker来连接压缩向量到两阶段搜索：

        ReaderSupplier rs = new MMapReaderSupplier(indexPath);
        OnDiskGraphIndex index = OnDiskGraphIndex.load(rs);
        // 加载我们刚刚写入磁盘的PQVectors
        try (RandomAccessReader in = new SimpleMappedReader(pqPath)) {
            PQVectors pqv = PQVectors.load(in);
            // SearchScoreProvider首先使用加载到内存中的PQVectors进行初步筛选，
            // 然后使用存储在索引磁盘中的精确向量进行重新排序
            Function<VectorFloat<?>, SearchScoreProvider> sspFactory = q -> {
                ApproximateScoreFunction asf = pqv.precomputedScoreFunctionFor(q, VectorSimilarityFunction.EUCLIDEAN);
                Reranker reranker = index.getView().rerankerFor(q, VectorSimilarityFunction.EUCLIDEAN);
                return new SearchScoreProvider(asf, reranker);
            };
            // 对照（精确计算的）真实结果测量我们的召回率
            testRecall(index, queryVectors, groundTruth, sspFactory);
        }

PQVectors提供precomputedScoreFunctionFor和scoreFunctionFor两种工厂方法。顾名思义，前者预先计算了ADC（非对称距离计算）所需的PQ码本距离片段。这对于搜索除最小索引之外的所有索引都更快，但如果您确实有一个微小的索引或需要在其他上下文中执行ADC，那么没有预计算的scoreFunctionFor版本会派上用场。

这组功能是经典的DiskANN设计。

第6步：构建大于内存的索引

JVector还可以应用PQ压缩来允许索引大于内存的数据集：只在内存中保留压缩向量。

首先，我们需要设置一个可以添加新向量的PQVectors实例，以及使用它的BuildScoreProvider：

        // 计算码本，但暂不编码任何向量
        ProductQuantization pq = ProductQuantization.compute(ravv, 16, 256, true);

        // 在构建索引时，我们将压缩新向量并将其添加到这个支持PQVectors的List中；
        // 这用于对构建搜索进行评分
        List<ByteSequence<?>> incrementallyCompressedVectors = new ArrayList<>();
        PQVectors pqv = new PQVectors(pq, incrementallyCompressedVectors);
        BuildScoreProvider bsp = BuildScoreProvider.pqBuildScoreProvider(VectorSimilarityFunction.EUCLIDEAN, pqv);

然后我们需要设置一个OnDiskGraphIndexWriter，以完全控制构建过程。

        Path indexPath = Files.createTempFile("siftsmall", ".inline");
        Path pqPath = Files.createTempFile("siftsmall", ".pq");
        // Builder创建看起来大致相同
        try (GraphIndexBuilder builder = new GraphIndexBuilder(bsp, ravv.dimension(), 16, 100, 1.2f, 1.2f);
             // 首次显式使用Writer，这是OnDiskGraphIndex.write背后的内容
             OnDiskGraphIndexWriter writer = new OnDiskGraphIndexWriter.Builder(builder.getGraph(), indexPath)
                     .with(new InlineVectors(ravv.dimension()))
                     .withMapper(new OnDiskGraphIndexWriter.IdentityMapper())
                     .build();
             // 压缩向量的输出
             DataOutputStream pqOut = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(pqPath))))

完成后，我们可以一次索引一个向量：

            // 逐个向量构建索引（在磁盘上）
            for (VectorFloat<?> v : baseVectors) {
                // 压缩新向量并将其添加到PQVectors（通过incrementallyCompressedVectors）
                int ordinal = incrementallyCompressedVectors.size();
                incrementallyCompressedVectors.add(pq.encode(v));
                // 将完整向量写入磁盘
                writer.writeInline(ordinal, Feature.singleState(FeatureId.INLINE_VECTORS, new InlineVectors.State(v)));
                // 现在将其添加到图中 -- 之前的步骤必须先完成，因为在addGraphNode构建过程中运行的搜索会使用PQVectors和InlineVectorValues
                builder.addGraphNode(ordinal, v);
            }

最后，我们需要运行cleanup()并将索引和PQVectors写入磁盘：

            // cleanup进行最终的maxDegree强制执行，并处理其他场景，如删除的节点
            // 在这里我们不需要担心这些
            builder.cleanup();

            // 完成索引写入（通过填充边列表）并写入我们完成的PQVectors
            writer.write(Map.of());
            pqv.write(pqOut);

注释：

切换到增量索引构建时，搜索代码不会改变 - 磁盘上的索引结构相同，只是（可能）大得多。
OnDiskGraphIndexWriter::writeInline通过同步实现线程安全，但如果您计划在多线程场景中使用它们，必须注意支持结构也是线程安全的（本例不是多线程的）。或者，您可以序列化对PQVectors的更新，只保留对GraphIndexBuilder::addGraphNode的并发调用。这代表了构建时间的大部分，所以使用这种方法您将看到良好的性能。

不太明显的要点

嵌入模型产生的输出来自一致的向量分布。这意味着您可以保存和重用ProductQuantization码本，即使是对不同的向量集，只要您第一次构建时有足够大的训练集。ProductQuantization.MAX_PQ_TRAINING_SET_SIZE（128,000个向量）已被证明足够大。
JDK ThreadLocal对象只能从创建它们的线程中引用。这是一个难以适应缓存Closeable对象（如GraphSearcher）的设计。JVector提供了ExplicitThreadLocal类来解决这个问题。
融合ADC仅与乘积量化兼容，不与二进制量化兼容。这并不是什么大损失，因为很少有模型生成最适合BQ的嵌入。尽管如此，非融合索引仍然支持BQ。
JVector大量利用Panama Vector API（SIMD）进行ANN索引和搜索。我们发现在索引和乘积量化过程中，内存带宽可能会饱和，导致进程变慢。为避免这种情况，索引和PQ构建的批处理方法使用PhysicalCoreExecutor将操作量限制在物理核心数量。默认值是Java看到的处理器数量的1/2。这可能并不适用于所有设置（例如，无超线程或混合架构），因此如果您希望覆盖默认值，请使用-Djvector.physical_core_count属性，或传入您自己的ForkJoinPool实例。

高级功能

融合ADC表示为在增量索引构建期间支持的Feature，类似于上面的InlineVectors。请参阅Grid类以获取示例代码。
各向异性PQ内置于ProductQuantization类中，可以提高召回率，但除了在每个模型的基础上进行实验外，没有人知道如何调整它（使用T/threshold参数），而选择错误的设置可能会使情况变糟。从论文的图3中可以看出：
JVector通过GraphIndexBuilder::markNodeDeleted支持原地删除。被删除的节点会在GraphIndexBuilder::cleanup期间被移除，连接会被替换，运行时间与被删除节点的数量成正比。
要checkpoint一个图并重新加载以继续编辑，请使用OnHeapGraphIndex::save和GraphIndexBuilder.load。

算法背后的研究

开发和测试

本项目组织为多模块Maven构建。目的是生成一个适合作为任何Java 11代码依赖的多版本jar包。当在启用Vector模块的Java 20+JVM上运行时，将使用优化的向量提供程序。一般来说，项目结构适合用JDK 20+构建，但当JAVA_HOME设置为Java 11到Java 19时，某些构建功能仍然可用。

基础代码在jvector-base中，将为Java 11版本构建，相应地限制语言特性和API。jvector-twenty中的代码将针对Java 20语言特性/API编译，并包含在最终的多版本jar中，面向支持的JVM。jvector-multirelease将jvector-base和jvector-twenty打包为发布用的多版本jar。jvector-examples是一个额外的同级模块，使用jvector-base/jvector-twenty的reactor表示来运行示例代码。jvector-tests包含项目的测试，能够在Java 11和Java 20+ JVM上运行。

要运行测试，使用mvn test。要在Java 20+上运行测试，使用mvn test。要在Java 11上运行测试，使用mvn -Pjdk11 test。要运行单个测试类，使用Maven Surefire测试过滤功能，例如，mvn -Dtest=TestNeighborArray test。你也可以使用方法级过滤和模式，例如，mvn -Dtest=TestNeighborArray#testRetain* test。

你可以直接运行SiftSmall和Bench来了解这里发生的一切。Bench将自动下载所需的数据集到fvec和hdf5目录。SiftSmall使用的文件可以在项目根目录的siftsmall目录中找到。

要运行这两个类，你可以使用Maven exec-plugin通过以下命令：