Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

pyvsag

pyvsag 是 VSAG 的 Python 绑定包,接口封装基于 pybind11 实现,源代码位于仓库 python_bindings/ 目录,打包脚本位于 python/

安装

从 PyPI 安装最新发布版本:

pip install pyvsag

需要在 Linux 环境下使用(manylinux2014 wheel)。如果希望构建本地 wheel,可以运行:

# 构建特定 Python 版本的 wheel
make pyvsag PY_VERSION=3.11

# 或一次构建所有受支持版本
make pyvsag-all

快速开始

pyvsag 暴露一个与 C++ Index 对象对应的 Index 类,构建与搜索参数使用 JSON 字符串传递:

import json
import numpy as np
import pyvsag

dim = 128
num_elements = 1000

ids = np.arange(num_elements, dtype=np.int64)
data = np.float32(np.random.random((num_elements, dim)))

index_params = json.dumps({
    "dtype": "float32",
    "metric_type": "l2",
    "dim": dim,
    "index_param": {
        "base_quantization_type": "fp32",
        "max_degree": 16,
        "ef_construction": 100,
    },
})

index = pyvsag.Index("hgraph", index_params)
index.build(vectors=data, ids=ids, num_elements=num_elements, dim=dim)

query = np.float32(np.random.random(dim))
search_params = json.dumps({"hgraph": {"ef_search": 100}})
result_ids, result_dists = index.knn_search(
    vector=query, k=10, parameters=search_params,
)
for rid, rdist in zip(result_ids, result_dists):
    print(f"{rid}: {rdist}")

完整示例请查阅仓库中的 examples/python/ 目录,建议从 103_index_hgraph.py 开始。

保存与加载

index.save("index.bin")

new_index = pyvsag.Index("hgraph", index_params)
new_index.load("index.bin")

当前 saveload 包装不会传播底层序列化和反序列化操作返回的错误。请确认输出文件 已经生成,并对重要索引执行保存后重新加载的 round-trip 校验,再依赖持久化结果。

索引操作

除构建、搜索、保存与加载外,Python 绑定还提供以下索引管理方法:

方法返回值说明
get_num_elements()int索引当前元素数量
get_memory_usage()int索引占用的内存字节数
check_id_exist(id)bool指定 ID 当前是否存在
get_min_max_id()(min_id, max_id)操作失败时返回 (-1, -1)
add(vectors, ids, num_elements, dim)None添加连续的稠密向量矩阵;dtype 必须与索引一致
remove(ids)int删除一维 int64 ID 数组;没有 ID 被删除和底层操作失败时都会返回 0
cal_distance_by_id(query, ids)numpy.ndarray每个 ID 返回一个 float32 距离;无效 ID 为 -1,底层操作失败时整个结果保持为 -1
# 添加两条稠密向量
new_vectors = np.random.random((2, dim)).astype(np.float32)
new_ids = np.array([10_001, 10_002], dtype=np.int64)
index.add(new_vectors, new_ids, len(new_ids), dim)

assert index.check_id_exist(10_001)
print(index.get_num_elements(), index.get_memory_usage())
print(index.get_min_max_id())

# 计算一个 query 到候选 ID 列表的距离
candidate_ids = np.array([10_001, 10_002], dtype=np.int64)
distances = index.cal_distance_by_id(query, candidate_ids)

# 删除 ID;返回值为实际删除数量
removed = index.remove(candidate_ids)

add 接受展开的连续数组或 row-major 二维矩阵。dtype: "float16" 时传入 numpy.float16dtype: "bfloat16" 时传入包含 BF16 位模式的 numpy.uint16 数组。 具体操作是否受支持仍取决于索引类型。输入校验和 add 等操作会抛出 Python 异常; removecal_distance_by_idsaveload 则使用上述 sentinel 或未检查行为, 不会统一传播底层操作错误。

与 C++ 库的关系

pyvsag 绑定的是同一份核心 C++ 实现,行为和性能特征与 C++ 版本保持一致。因此:

  • 大多数 C++ 参数在 Python 中以相同的 JSON 字段传递;
  • C++ 版本新增的索引类型、量化方式、距离度量会随 pyvsag 的下一个 release 一同发布;
  • 构建 wheel 时所使用的依赖项与发布版 C++ 库相同(OpenBLAS、libaio 等)。

关于可用参数和索引类型,请参考 创建索引索引参数