跳到内容

使用 Lance 进行向量索引和向量搜索

Lance 提供了具有 ANN(近似最近邻)索引的高性能向量搜索功能。

通过本教程,您将能够构建和使用 ANN 索引,从而在保持高精度的同时显著加快向量搜索操作。您还将学习如何调整搜索参数以获得最佳性能,并将向量搜索与元数据查询结合在一个操作中。

安装 Python SDK

pip install pylance

设置您的环境

首先,导入必要的库

import shutil
import lance
import numpy as np
import pandas as pd
import pyarrow as pa
import duckdb

准备您的向量嵌入

在本教程中,下载并准备 SIFT 1M 数据集用于向量搜索实验。

  • 从以下链接下载 ANN_SIFT1M:http://corpus-texmex.irisa.fr/
  • 直接链接:ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz

您可以使用 wget

rm -rf sift* vec_data.lance
wget ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz
tar -xzf sift.tar.gz

将您的数据转换为 Lance 格式

然后,将原始向量数据转换为 Lance 格式,以实现高效存储和查询。

from lance.vector import vec_to_table
import struct

uri = "vec_data.lance"

with open("sift/sift_base.fvecs", mode="rb") as fobj:
    buf = fobj.read()
    data = np.array(struct.unpack("<128000000f", buf[4 : 4 + 4 * 1000000 * 128])).reshape((1000000, 128))
    dd = dict(zip(range(1000000), data))

table = vec_to_table(dd)
lance.write_dataset(table, uri, max_rows_per_group=8192, max_rows_per_file=1024*1024)

现在您可以加载数据集了

uri = "vec_data.lance"
sift1m = lance.dataset(uri)

无索引搜索

您将执行无索引的向量搜索以查看基线性能,然后将其与索引搜索进行比较。

首先,让我们采样一些查询向量

import duckdb
# Make sure DuckDB v0.7+ is installed
samples = duckdb.query("SELECT vector FROM sift1m USING SAMPLE 100").to_df().vector
0     [29.0, 10.0, 1.0, 50.0, 7.0, 89.0, 95.0, 51.0,...
1     [7.0, 5.0, 39.0, 49.0, 17.0, 12.0, 83.0, 117.0...
2     [0.0, 0.0, 0.0, 10.0, 12.0, 31.0, 6.0, 0.0, 0....
3     [0.0, 2.0, 9.0, 1.793662034335766e-43, 30.0, 1...
4     [54.0, 112.0, 16.0, 0.0, 0.0, 7.0, 112.0, 44.0...
                            ...
95    [1.793662034335766e-43, 33.0, 47.0, 28.0, 0.0,...
96    [1.0, 4.0, 2.0, 32.0, 3.0, 7.0, 119.0, 116.0, ...
97    [17.0, 46.0, 12.0, 0.0, 0.0, 3.0, 23.0, 58.0, ...
98    [0.0, 11.0, 30.0, 14.0, 34.0, 7.0, 0.0, 0.0, 1...
99    [20.0, 8.0, 121.0, 98.0, 37.0, 77.0, 9.0, 18.0...
Name: vector, Length: 100, dtype: object

现在,执行无索引的最近邻搜索

import time

start = time.time()
tbl = sift1m.to_table(columns=["id"], nearest={"column": "vector", "q": samples[0], "k": 10})
end = time.time()

print(f"Time(sec): {end-start}")
print(tbl.to_pandas())

预期输出

Time(sec): 0.10735273361206055
       id                                             vector    score
0  144678  [29.0, 10.0, 1.0, 50.0, 7.0, 89.0, 95.0, 51.0,...      0.0
1  575538  [2.0, 0.0, 1.0, 42.0, 3.0, 38.0, 152.0, 27.0, ...  76908.0
2  241428  [11.0, 0.0, 2.0, 118.0, 11.0, 108.0, 116.0, 21...  92877.0
...

如果没有索引,搜索将遍历整个数据集来计算每个数据点之间的距离。为了获得实用的实时性能,使用 ANN 索引将获得更好的性能。

构建搜索索引

如果您构建 ANN 索引,您可以显著加快向量搜索操作的速度,同时保持高精度。在此示例中,我们将构建 IVF_PQ 索引

sift1m.create_index(
    "vector",
    index_type="IVF_PQ", # specify the IVF_PQ index type
    num_partitions=256,  # IVF
    num_sub_vectors=16,  # PQ
)

示例响应应如下所示

Building vector index: IVF256,PQ16
CPU times: user 2min 23s, sys: 2.77 s, total: 2min 26s
Wall time: 22.7 s
Sample 65536 out of 1000000 to train kmeans of 128 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters
Sample 65536 out of 1000000 to train kmeans of 8 dim, 256 clusters

索引创建性能

如果您在自己的数据上尝试此操作,请确保您的向量(维度/子向量数)% 8 == 0,否则由于 SIMD 未对齐,索引创建时间将比预期长得多。

使用 ANN 索引进行向量搜索

您现在可以使用新创建的索引执行相同的搜索操作,并查看显著的性能提升。

sift1m = lance.dataset(uri)

import time

tot = 0
for q in samples:
    start = time.time()
    tbl = sift1m.to_table(nearest={"column": "vector", "q": q, "k": 10})
    end = time.time()
    tot += (end - start)

print(f"Avg(sec): {tot / len(samples)}")
print(tbl.to_pandas())

预期输出

Avg(sec): 0.0009334301948547364
       id                                             vector         score
0  378825  [20.0, 8.0, 121.0, 98.0, 37.0, 77.0, 9.0, 18.0...  16560.197266
1  143787  [11.0, 24.0, 122.0, 122.0, 53.0, 4.0, 0.0, 3.0...  61714.941406
2  356895  [0.0, 14.0, 67.0, 122.0, 83.0, 23.0, 1.0, 0.0,...  64147.218750
3  535431  [9.0, 22.0, 118.0, 118.0, 4.0, 5.0, 4.0, 4.0, ...  69092.593750
4  308778  [1.0, 7.0, 48.0, 123.0, 73.0, 36.0, 8.0, 4.0, ...  69131.812500
5  222477  [14.0, 73.0, 39.0, 4.0, 16.0, 94.0, 19.0, 8.0,...  69244.195312
6  672558  [2.0, 1.0, 0.0, 11.0, 36.0, 23.0, 7.0, 10.0, 0...  70264.828125
7  365538  [54.0, 43.0, 97.0, 59.0, 34.0, 17.0, 10.0, 15....  70273.710938
8  659787  [10.0, 9.0, 23.0, 121.0, 38.0, 26.0, 38.0, 9.0...  70374.703125
9  603930  [32.0, 32.0, 122.0, 122.0, 70.0, 4.0, 15.0, 12...  70583.375000

性能说明

您的实际数字将因您的存储而异。这些数字来自 M2 MacBook Air 上的本地磁盘。如果您直接查询 S3、HDD 或网络驱动器,性能会较慢。

调整搜索参数

您需要调整搜索参数以平衡速度和精度,为您的用例找到最佳设置。

延迟与召回率可通过以下参数进行调整:- nprobes:要搜索的 IVF 分区数 - refine_factor:确定在重新排序期间检索的向量数

%%time

sift1m.to_table(
    nearest={
        "column": "vector",
        "q": samples[0],
        "k": 10,
        "nprobes": 10,
        "refine_factor": 5,
    }
).to_pandas()

参数解释: - q => 示例向量 - k => 返回的邻居数 - nprobes => 要探测的分区数(在粗量化器中) - refine_factor => 控制“重新排序”。如果 k=10 且 refine_factor=5,则通过 ANN 检索 50 个最近邻居,并使用实际距离重新排序,然后返回前 10 个。这可以在不过多牺牲性能的情况下提高召回率

内存使用

上述延迟包括文件 I/O,因为 Lance 目前不在内存中保存任何内容。除了索引构建速度之外,创建纯内存版本的数据集将对性能产生最大影响。

结合特征和向量

您可以将元数据列添加到您的向量数据集中,并在单个操作中同时查询向量和特征。

在实际情况中,用户有其他需要一起存储和获取的特征或元数据列。如果您单独管理数据和索引,您需要做很多烦人的连接工作才能将它们组合在一起。

使用 Lance 只需一个调用

tbl = sift1m.to_table()
tbl = tbl.append_column("item_id", pa.array(range(len(tbl))))
tbl = tbl.append_column("revenue", pa.array((np.random.randn(len(tbl))+5)*1000))

然后您可以同时查询向量和元数据

# Get vectors and metadata together
result = sift1m.to_table(
    columns=["item_id", "revenue"],
    nearest={"column": "vector", "q": samples[0], "k": 10}
)
print(result.to_pandas())

后续步骤

您应该查看使用 Lance 对您的数据集进行版本控制。我们将向您展示如何对您的向量数据集进行版本控制并跟踪随时间的变化。