APC 技術ブログ

株式会社エーピーコミュニケーションズの技術ブログです。

株式会社 エーピーコミュニケーションズの技術ブログです。

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

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

Databricks
＜ techblogAPCについて（はじめにお読み下さい） ＞

はじめに

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ドキュメントに変換しています。

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

  1. VectorSearchClientのインストール
  2. トイデータセットをソースDeltaテーブルにロード
  3. Vector Searchエンドポイントの作成
  4. Vectorインデックスの作成、取得
  5. 類似検索
  6. 類似検索した結果をLangChainドキュメントに変換
  7. 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になったことが確認出来てから、再度類似検索を実行するようにしてください。

類似検索

類似検索のコードを実行した結果、このように結果が返ってきました。

＊Mosaic AI Vector Searchでは、類似度はL2 距離（ユークリッド距離）によって計算されています。コサイン類似度を使いたい場合は、データポイントの埋め込みを検索に入力する前に正規化する必要があります。

おわりに

Databricksのシームレスなプラットフォームだと、ノートブックひとつでベクトルインデックスを作り、そのまま類似検索が出来てしまうので、プラットフォームを行ったり来たりするストレスが無く、とても楽にタスクを完了出来ました。

最後までご覧いただきありがとうございます。 私たちはDatabricksを用いたデータ分析基盤の導入から内製化支援まで幅広く支援をしております。 もしご興味がある方は、お問い合わせ頂ければ幸いです。

www.ap-com.co.jp