Vector Search・Vectorインデックスの作成と、ドキュメント内の類似検索

はじめに
このブログで使用したノートブック
前提条件
ノートブックの概要
検証結果
類似検索
おわりに

はじめに

GLB事業部Lakehouse部の佐藤です。この記事では、DatabricksでのVectorインデックスの作り方と、Vectorインデックスを使っての類似検索についてご紹介します。

Vector Search（ベクトル検索）は、データポイントをベクトルとして表現し、そのベクトル間の類似度に基づいて検索を行う手法です。 Vectorインデックスは、大量のベクトルデータを効率的に検索するためのデータ構造です。通常のデータベースインデックスと同様に、検索のスピードを向上させるために使用されますが、高次元ベクトルデータに特化しています。ベクトル検索とインデックスの組み合わせにより、膨大なデータセットの中から類似するデータを高速かつ効率的に見つけることができます。

Databricksでは統合された環境が提供されているため、データ処理からベクトル検索までの一貫したワークフローをノートブックひとつで実現できます。なお、DatabricksのVector Searchは正式にはMosaic AI Vector Searchという名称になります。Databricks Mosaic AIは、Databricksの、特にAIおよび機械学習に特化した機能を強化したツールセットです。

このブログで使用したノートブック

この記事では、こちらのサンプルコードを実行しています。

docs.databricks.com

前提条件

AWSに構築されたDatabricksワークスペース
リージョンはus-east1（tokyoリージョンではまだモデルサービングがリリースされていないため）
15.1 MLクラスターを使用
Unity Catalogを有効化
サーバレスコンピュートを有効化
Vectorインデックスを作成するカタログ・スキーマは既存のものを使用、もしくは事前に作成しておく
エンべディングモデルのモデルサービングエンドポイントも事前に作成しておく

ノートブックの概要

このノートブックでは、Vector Searchを操作する際の主要なAPIであるVectorSearchClientを提供するVector Search Python SDKを使って、Vector SearchエンドポイントやVectorインデックスの作成、インデックスの取得までを行っています。（Vector SearchエンドポイントやVectorインデックスはUIを使っての作成も可能です。）

また、Vectorインデックス作成後、類似検索を行い、類似検索の結果をLangChainドキュメントに変換しています。

ノートブックの大まかな流れは以下の通りです。

VectorSearchClientのインストール
トイデータセットをソースDeltaテーブルにロード
Vector Searchエンドポイントの作成
Vectorインデックスの作成、取得
類似検索
類似検索した結果をLangChainドキュメントに変換
Vectorインデックスの削除

若干長めのコードになるため、このブログでは、

Vector Searchエンドポイントの作成
Vectorインデックスの作成
類似検索の実行

の実行結果に絞って共有させていただきます。

検証結果

Vector Searchエンドポイントの作成

vector_search_endpoint_name = "mihsato-vector-search-demo-endpoint"

vsc.create_endpoint(
    name=vector_search_endpoint_name,
    endpoint_type="STANDARD"
)

endpoint = vsc.get_endpoint(
  name=vector_search_endpoint_name)
endpoint

コードを実行した結果、このようにVector Searchエンドポイントが出来ました。

Vectorインデックスの作成、取得

 'mihsato-e5-small-v2'

vs_index = "en_wiki_index"
vs_index_fullname = f"{source_catalog}.{source_schema}.{vs_index}"

embedding_model_endpoint = "mihsato-e5-small-v2"

index = vsc.create_delta_sync_index(
  endpoint_name=vector_search_endpoint_name,
  source_table_name=source_table_fullname,
  index_name=vs_index_fullname,
  pipeline_type='TRIGGERED',
  primary_key="id",
  embedding_source_column="text",
  embedding_model_endpoint_name=embedding_model_endpoint
)
index.describe()

このコードでは、事前に作成済みのモデルサービングエンドポイントを指定する必要があります。私は、前回のブログで作成したモデルサービングエンドポイントを使用しました。

コードを実行すると、このようにVectorインデックスが作成されます。

Vectorインデックスを取得するコード

index = vsc.get_index(endpoint_name=vector_search_endpoint_name, index_name=vs_index_fullname)

index.describe()

# Wait for index to come online. Expect this command to take several minutes.
import time
while not index.describe().get('status').get('detailed_state').startswith('ONLINE'):

 print("Waiting for index to be ONLINE...")
  time.sleep(5)
print("Index is ONLINE")
index.describe()

Vectorインデックスを取得するコードを実行すると、このように取得されたベクトルのカラムが作成されます。

＊Vectorインデックスの取得には時間が掛かります。 Endpoint StatusがReadyになる前に次の類似検索のコードブロックを実行してしまうと、エラーが返ってきます。その場合は、VectorインデックスのStatusを確認し、Readyになったことが確認出来てから、再度類似検索を実行するようにしてください。