OceanBaseはpyobvectorを使用したベクトルストレージと検索をサポートしています。本記事では、クイックスタートの方法について説明します。pyobvectorはOceanBaseのベクトルストレージ用Python SDKであり、SQLAlchemyをベースとし、Milvus APIと基本的に互換性があります。
前提条件
OceanBaseクラスタのデプロイが完了し、MySQLテナントが作成されていることを確認してください。
環境にPython 3.9以降のバージョンがインストールされていることを確認してください。
クイックスタート
pyobvectorのインストール
まず、ローカル環境にpyobvectorをインストールする必要があります。以下のコマンドを参照してください:
pip install -U pyobvector
-U パラメータは、環境にpyobvectorが既にインストールされている場合、最新バージョンに自動的にアップグレードします。インストールされていない場合は、最新バージョンを直接インストールします。
使用方法
pyobvectorは2つのモードをサポートしています:
Milvus互換モード:MilvusLikeClientクラスが提供する、Milvus APIに類似した方法でベクトルストレージを使用します。
SQLAlchemyハイブリッドモード:ObVecClientクラスが提供するベクトルストレージ機能を使用し、SQLAlchemyライブラリでリレーショナルデータベースのステートメントを実行します。このモードでは、pyobvectorをSQLAlchemyの拡張機能と見なすことができます。
Milvus互換モード
クライアント接続の確立
pyobvectorは、Milvusと互換性のある方法でOceanBaseのベクトルストレージおよび検索機能を利用できるMilvusLikeClientを提供しています。以下のステートメントでクライアントオブジェクトを作成できます:
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 の3列を持つテーブルを作成し、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は、OceanBaseのベクトルストレージと検索機能をSQLAlchemyハイブリッドモードで利用できるObVecClientを提供します。以下のステートメントでクライアントオブジェクトを作成できます:
from pyobvector import *
# 以下のデータベース接続情報パラメータをご自身のデータベースインスタンスの情報に変更してください
client = ObVecClient(uri="127.0.0.1:2881", user="root@test", db_name="test")
テーブルとベクトルインデックスの作成
以下の例では、id、embedding、metadata の3列を持つテーブルを作成し、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}')
関連ドキュメント
- ベクトルデータの詳細については、ベクトルデータを参照してください。
- ベクトルインデックスの関連説明については、ベクトルインデックスを参照してください。