JVector: 革命性的嵌入式向量搜索引擎

JVector简介

JVector是一款先进的嵌入式向量搜索引擎,由Jonathan Ellis开发并开源。它采用基于图的索引结构,实现了高效的近似最近邻(ANN)搜索。相比传统的精确最近邻搜索,JVector能够在对数时间内返回近似结果,大大提高了搜索效率。

JVector的核心特性包括:

• 基于DiskANN设计的单层图索引结构
• 支持非阻塞并发控制,构建速度随CPU核心数线性扩展
• 两阶段搜索设计,平衡内存使用、延迟和准确性
• 支持产品量化(PQ)和二值量化(BQ)压缩向量
• 可构建超大规模索引,突破内存限制
• 支持增量构建和更新索引

这些特性使JVector成为目前最先进的嵌入式向量搜索引擎之一,能够满足各种大规模向量搜索的需求。

JVector架构设计

JVector采用基于图的索引结构,继承了DiskANN的设计思想并进行了多项创新。其核心架构包括:

1. 单层图索引:每个节点存储邻近节点列表,形成高效的搜索图结构。
2. 非阻塞并发控制:允许多线程并行构建索引,性能随CPU核心数线性扩展。
3. 两阶段搜索:
   • 第一阶段使用压缩向量在内存中快速搜索
   • 第二阶段使用磁盘上的精确向量重新排序
4. 向量压缩:
   • 产品量化(PQ)
   • 二值量化(BQ)
   • Fused ADC
5. 磁盘索引:支持构建超大规模索引,突破内存限制

这种架构设计使JVector能够在保持高准确率的同时,显著降低内存使用和搜索延迟。特别是支持构建超大规模索引的能力,让JVector可以充分利用对数时间复杂度的搜索优势,而不是退化为多索引间的线性时间合并。

使用JVector步骤详解

接下来,我们将通过一系列步骤,详细介绍如何使用JVector构建和查询向量索引。所有代码示例均来自JVector源代码仓库中的SiftSmall示例。

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

首先,我们来看如何在内存中构建和查询一个基本的索引:

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

    // 使用原始向量创建BuildScoreProvider
    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);
        }
    }
}

这段代码展示了JVector的基本用法:

1. 创建一个RandomAccessVectorValues对象来包装原始向量。
2. 使用BuildScoreProvider和GraphIndexBuilder构建索引。
3. 使用GraphSearcher.search方法执行搜索。

值得注意的是:

• 所有索引都假定向量源具有一致的固定维度。
• GraphIndexBuilder的溢出参数对于内存构建通常设为1.2,对于磁盘构建设为1.5。
• alpha参数控制边缘距离和多样性之间的权衡,高维向量通常使用1.2,2D或3D数据集推荐使用2.0。

步骤2:更精细地控制GraphSearcher

为了更好地控制搜索过程,我们可以直接使用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);
    }
}

这种方法允许我们:

• 重用GraphSearcher实例,避免重复初始化开销。
• 自定义SearchScoreProvider,控制如何计算查询向量与索引中其他节点的相似度。

步骤3:测量召回率

为了确保索引的准确性,我们可以使用testRecall方法测量召回率:

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

这将输出类似以下的结果:

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

召回率的微小波动是正常的,这是由于索引的近似性质和构建过程中的多线程并发导致的。

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

JVector支持将索引持久化到磁盘,以及从磁盘加载索引:

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 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而不是简单的Reader,因为搜索时可能需要打开额外的读取器。
• JVector提供了两种内置的RandomAccessReader实现:
  1. SimpleMappedReader: 使用FileChannel.map实现,兼容性好但限制文件大小为2GB。
  2. MemorySegmentReader: 使用较新的MemorySegment API实现,无文件大小限制但要求Java 22+。

步骤5:使用压缩向量进行搜索

为了进一步优化内存使用和搜索性能,JVector支持使用压缩向量:

// 计算并将压缩向量写入磁盘
Path pqPath = Files.createTempFile("siftsmall", ".pq");
try (DataOutputStream out = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(pqPath)))) {
    // 使用PQ压缩原始向量,压缩比为 128 * 4 / 16 = 32x
    ProductQuantization pq = ProductQuantization.compute(ravv,
                                                         16, // 子空