Databricks Vector Searchのハイブリッド検索を試す

導入

以下の記事で、Databricks Vector Searchにハイブリッド検索が導入されるかも？という内容を書きました。

そして今日リリースノートをたまたま見たら、ハイブリッドサーチのサポートが記載されていました。

また、公式ドキュメントにも説明が載っていました。
2024/6月時点でパブリックプレビューのようです。

オフィシャルでサポートがアナウンスされましたので、公式ドキュメントの内容確認および試行してみます。

Databricks Vector Searchのハイブリッド検索とは？

公式ドキュメントから抜粋：

Mosaic AI Vector Search は、ベクターベースの埋め込み検索と従来のキーワードベースの検索手法を組み合わせたハイブリッド キーワード類似性検索もサポートしています。 このアプローチでは、クエリ内の単語と完全に一致すると同時に、ベクトルベースの類似性検索を使用して、クエリのセマンティックな関係とコンテキストをキャプチャします。

これら 2 つの手法を統合することで、ハイブリッド キーワード類似検索は、正確なキーワードだけでなく、概念的に類似したキーワードを含むドキュメントも取得し、より包括的で関連性の高い検索結果を提供します。 この方法は、純粋な類似性検索には適さない SKU や識別子などの固有のキーワードがソース データに含まれている RAG アプリケーションで特に役立ちます。

ハイブリッド検索は パブリック プレビュー段階です。

また、アルゴリズムの詳細も記載されています。

キーワード検索アルゴリズム
関連性スコアは、 Okapi BM25 を使用して計算されます。 テキストまたは文字列形式のソース テキスト埋め込みおよびメタデータ列を含む、すべてのテキストまたは文字列列が検索されます。 トークン化関数は、単語の境界で分割し、句読点を削除し、すべてのテキストを小文字に変換します。

類似検索とキーワード検索の組み合わせ
類似検索とキーワード検索結果は、相互順位融合(RRF)機能を使用して結合されます。
RRF は、スコアを使用して、各メソッドから各ドキュメントを再スコアリングします。

というわけで、ドキュメントを読んだ範囲での理解ですが、Databricks Vector Searchのハイブリッド検索は(近似最近傍探索による)ベクトル検索と、BM25によるキーワード検索のハイブリッドであり、各結果をReciprocal Rank Fusion (RRF)でマージ・再ランク付けしているようです。

※　手前みそですが、以前こちらの記事で類似の実装をしていました。

オーソドックスなハイブリッド検索の実現手法かなという印象ですが、キーワード検索においては形態素解析(単語分割)やストップワード除去といった前処理が大事という認識で、このあたりがどうなっているかは（特に日本語を扱う上で）気になるところです。

日本語がダメだとハイブリッド検索は性能低下を引き起こす可能性があるので、実際に試してみます。
検証環境はDatabricks on AWS、DBRは15.2MLです。

Step1. パッケージインストール

databricks-vectorsearchを最新バージョンに更新します。
また、念のためLangChainのパッケージを更新しておきます。

%pip install --upgrade --force-reinstall databricks-vectorsearch
%pip install -U langchain langchain_community

dbutils.library.restartPython()

検証時点でdatabricks-vectorsearchのバージョンは0.38でした。

Step2. インデックスの取得

検証に使うDatabricks Vector Searchのインデックスを用意します。
新しくインデックスを作り直すのも煩雑だったので、以下の記事で作成したインデックスを流用します。

from databricks.vector_search.client import VectorSearchClient

# 既に存在しているエンドポイントとインデックスを利用する
vector_search_endpoint_name = "default_vector_search_endpoint"
index_name = "training.databricks_qa_jp.databricks_documentation_vs_index"

# Vector Searchのクライアントを取得
client = VectorSearchClient()

# 既存のインデックスを取得
index = client.get_index(
    index_name=index_name,
    endpoint_name=vector_search_endpoint_name,
)

Step3. 類似文書検索

取得したインデックスを利用してい、類似文書検索を実行します。

Databricks Vector SearchのPython APIを利用する場合、類似検索はsimilarity_searchメソッドを利用します。
改めてヘルプを確認してみます。

help(index.similarity_search)

Help on method similarity_search in module databricks.vector_search.index:

similarity_search(columns, query_text=None, query_vector=None, filters=None, num_results=5, debug_level=1, score_threshold=None, query_type=None) method of databricks.vector_search.index.VectorSearchIndex instance
    Perform a similarity search on the index. This returns the top K results that are most similar to the query.
    
    :param columns: List of column names to return in the results.
    :param query_text: Query text to search for.
    :param query_vector: Query vector to search for.
    :param filters: Filters to apply to the query.
    :param num_results: Number of results to return.
    :param debug_level: Debug level to use for the query.
    :param score_threshold: Score threshold to use for the query.
    :param query_type: Query type of this query. Choices are "ANN" and "HYBRID".

どうもquery_typeパラメータを指定することでハイブリッド検索を利用できそうですね。
(バージョン0.36の時とは仕様が変わっています)
ANNはApproximate Nearest Neighbor: 近似最近傍探索のことかな。

では、まず通常のベクトル検索を実行。

import pandas as pd
result = index.similarity_search(
    query_text="Unity Catalogとは何?",
    columns=["id", "content", "url"],
    num_results=5,
)

columns = [c.get("name") for c in result["manifest"]["columns"]]
display(pd.DataFrame(result["result"]["data_array"], columns=columns))

結果はこちら。

では、ハイブリッド検索を同じクエリで実行してみます。

import pandas as pd
result = index.similarity_search(
    query_text="Unity Catalogとは何?",
    columns=["id", "content", "url"],
    num_results=5,
    query_type="HYBRID",
)

columns = [c.get("name") for c in result["manifest"]["columns"]]
display(pd.DataFrame(result["result"]["data_array"], columns=columns))

結果はこちら。

得られた文章、そしてスコアの付き方が変わってますね。

なんとなくですが、ハイブリッド検索の方が適した文章を持ってきているように思います。

次に、日本語のみのクエリではどうでしょうか。

import pandas as pd
result = index.similarity_search(
    query_text="データウェアハウスに関係する代表的な機能には何がある？",
    columns=["id", "content", "url"],
    num_results=5,
    query_type="HYBRID",
)

columns = [c.get("name") for c in result["manifest"]["columns"]]
display(pd.DataFrame(result["result"]["data_array"], columns=columns))

結果はこちら。

なお、ベクトル検索単体だと、以下のようになります。

こちらでもハイブリッド検索の方がよい結果に見えます。

日本語の文書やクエリでも問題なく動作していそうですね。

まとめ

Databricks Vector Searchでもハイブリッド検索が(パブリックプレビューですが)提供されました。
ハイブリッド検索自体は既に広く利用されるものとなっており、競合と比較すると早い提供というわけではありませんが、Lakehouseと結びついた強力なデータ管理とセットであることを考えると、非常に使い勝手がよいなと思います。

Vector Search周りはクラスタの0スケーリングが出来ないとか(オンライン利用想定なので、レイテンシが高くなる対応はしないんでしょうが)、改善して欲しいポイントはいくつかありますが、プラットフォーム全体で見ると非常に有用な機能です。
これからもうまく活用していきたいと思います。