首批通过分布式安全可靠测评,为关键业务系统打造
使用 Python 快速进行向量搜索
更新时间:2026-04-11 17:12:13
OceanBase 支持使用 pyobvector 进行向量存储和搜索,本文介绍了快速上手方法。pyobvector 是 OceanBase 向量存储的 Python SDK,基于 SQLAlchemy,基本兼容 Milvus API。
前提条件
请确保你已部署了 OceanBase 集群并创建了 MySQL 租户。
请确保你的环境已安装 Python 3.9 及以上版本。
快速上手
安装 pyobvector
首先,需要将 pyobvector 先安装到我们的本地环境中。可参考如下命令:
pip install -U pyobvector
-U 参数表示如果你的环境已经安装了 pyobvector,则自动升级到最新版本;如果未安装,则直接安装最新版本。
用法
pyobvector 支持两种模式:
Milvus 兼容模式:用 MilvusLikeClient 类提供的类似于 Milvus API 的方式来使用向量存储。
SQLAlchemy 混合模式:使用 ObVecClient 类提供的向量存储功能,并用 SQLAlchemy 库执行关系数据库语句。在这种模式下,可以将 pyobvector 视为 SQLAlchemy 的扩展。
Milvus 兼容模式
建立客户端连接
pyobvector 提供了 MilvusLikeClient 让用户能够以 Milvus 兼容的方式使用 OceanBase 的向量存储和搜索能力,通过下面的语句能够创建一个客户端对象:
from pyobvector import *
# 请将下面的数据库连接信息参数修改为您的数据库实例的信息
client = MilvusLikeClient(uri="127.0.0.1:2881", user="root@test", db_name="test")
创建具有向量索引的集合
为了兼容 Milvus 的 API,pyobvector 的 MilvusClient 也提供了 create_collection 等方法。虽然方法名是创建集合,但是映射到 OceanBase 中实际上是创建了一张表,通过下面的例子能够创建一张具有 id、embedding、metadata 三列的表,并且针对 embedding 列创建一个 HNSW 类型的向量索引。其中 embedding 列就是一个 64 维的向量类型。
fields = [
FieldSchema(
name="id",
dtype=DataType.INT64,
is_primary=True,
auto_id=True,
),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=64),
FieldSchema(name="metadata", dtype=DataType.JSON),
]
index_params = MilvusLikeClient.prepare_index_params()
index_params.add_index(
field_name="embedding",
index_name="embedding_idx",
index_type=VecIndexType.HNSW,
distance="l2",
m=16,
ef_construction=256,
)
schema = CollectionSchema(fields)
table_name = "vector_search"
client.create_collection(table_name, schema=schema, index_params=index_params)
构造并写入数据
为了模拟在大量向量数据中进行搜索的场景,在这一步中先构造一些向量数据。方式是使用 python 中的 random 模块来构造随机的浮点数列表。
import random
random.seed(20241023)
batch_size = 100
batch = []
for i in range(1000):
batch.append(
{
"embedding": [random.uniform(-1, 1) for _ in range(64)],
"metadata": {"idx": i},
}
)
if len(batch) == batch_size:
client.insert(collection_name=table_name, data=batch)
batch = []
if len(batch) > 0:
client.insert(collection_name=table_name, data=batch)
执行相似向量搜索
通过 random.uniform 构造一个目标向量数据作为输入,在刚才插入数据的集合中进行向量搜索:
target_data = [random.uniform(-1, 1) for _ in range(64)]
res = client.search(
collection_name=table_name,
data=target_data,
anns_field="embedding",
limit=5,
output_fields=["id", "metadata"],
)
print(res)
# 预期返回结果如下
# [{'id': 63, 'metadata': {'idx': 62}}, {'id': 796, 'metadata': {'idx': 795}}, {'id': 187, 'metadata': {'idx': 186}}, {'id': 784, 'metadata': {'idx': 783}}, {'id': 880, 'metadata': {'idx': 879}}]
SQLAlchemy 混合模式
建立客户端连接
pyobvector 提供了 ObVecClient 让用户能够以 SQLAlchemy 混合模式使用 OceanBase 的向量存储和搜索能力,通过下面的语句能够创建一个客户端对象:
from pyobvector import *
# 请将下面的数据库连接信息参数修改为您的数据库实例的信息
client = ObVecClient(uri="127.0.0.1:2881", user="root@test", db_name="test")
创建表和向量索引
通过下面的例子能够创建一张具有 id、embedding、metadata 三列的表,并且针对 embedding 列创建一个 HNSW 类型的向量索引。其中 embedding 列就是一个 64 维的向量类型。
from sqlalchemy import Column, Integer, JSON
from sqlalchemy import func
cols = [
Column("id", Integer, primary_key=True, autoincrement=True),
Column("embedding", VECTOR(64)),
Column("metadata", JSON),
]
table_name = "vector_test3"
client.create_table(table_name, columns=cols)
print(f"Table {table_name} created")
client.create_index(
table_name,
is_vec_index=True,
index_name="embedding_idx",
column_names=["embedding"],
vidx_params="distance=l2, type=hnsw, lib=vsag", # m=16, ef_construction=256
)
print(f"Index {table_name}.embedding_idx created")
构造并写入数据
为了模拟在大量向量数据中进行搜索的场景,在这一步中先构造一些向量数据。方式是使用 python 中的 random 模块来构造随机的浮点数列表。
import random
random.seed(20241023)
batch_size = 100
batch = []
for i in range(1000):
batch.append(
{
"embedding": [random.uniform(-1, 1) for _ in range(64)],
"metadata": {"idx": i},
}
)
if len(batch) == batch_size:
client.insert(table_name, data=batch)
batch = []
if len(batch) > 0:
client.insert(table_name, data=batch)
执行相似向量搜索
通过 random.uniform 构造一个目标向量数据作为输入,在刚才插入数据的集合中进行向量搜索:
target_data = [random.uniform(-1, 1) for _ in range(64)]
res = client.ann_search(
table_name,
vec_data=target_data,
vec_column_name="embedding",
distance_func=func.l2_distance,
topk=5,
output_column_names=["id", "metadata"],
)
for r in res:
print(r)
# 预期返回结果如下
# (63, '{"idx": 62}')
# (796, '{"idx": 795}')
# (187, '{"idx": 186}')
# (784, '{"idx": 783}')
# (880, '{"idx": 879}')