ベクトル検索による文書検索にDuckDBを使ってみた
はじめに
最近よく見るデータ分析(OLAP)向けのデータベース「DuckDB」は、標準でベクトル検索に対応しています。DuckDBはPandasと連携しやすいという特徴もあります。
そこで、今回はちょっと実践的なベクトル検索にDuckDBを使ってみました。
今回実装したものはこちらのリポジトリにアップしています。
また、Gradioを用いて作成したデモアプリをHugging Face Spaceにアップしています。
ベクトル検索をDuckDBで実装する
ベクトル検索に使う埋め込みモデルは、以前ハイブリッド検索を実装してみたときに使用した「Ruri」を今回も使用します。
インデックスの作成
検索文書のインデックス作成をするためには、文書を埋め込みベクトルに変換しDuckDBに入れる必要があります。DuckDBはDataFrameから簡単にテーブルを作成できるので、DataFrameを作成してDuckDBを作成してみます。
以下のようなコードでインデックスをDuckDBに作成します。
import sentence_transformers as st
import pandas as pd
import duckdb
model = st.SentenceTransformer("cl-nagoya/ruri-large")
# 検索文書には「文書: 」と頭につける
search_fields = [f"文書: {text}" for text in df["content"].values.tolist()]
# ベクトル化
embeddings = model.encode(search_fields)
# DuckDBにテーブルを作成
# テーブル名
vector_store_name = "ruri_vector_index"
# DuckDBに作成するテーブルをDataFrameにする
vdb = pd.DataFrame({"index": range(len(embeddings)), "embedding": embeddings.tolist()})
# DuckDBにテーブルを作成
duckdb.register(vector_store_name, vdb)
文書の検索
作成したインデックスから検索クエリと類似する文書を検索する場合は、以下のようなコードを実行します。
# 検索クエリには「クエリ: 」と頭につける
text = "日本の首都"
embedding = model.encode(f"クエリ: {text}")
# 埋め込みベクトルの次元数
num_dim = len(embeddings_query)
# DuckDBのテーブルに対してクエリ実行
distance = duckdb.sql(f"""
select
index,
array_cosine_distance(embedding::DOUBLE[{num_dim}], {embeddings_query.tolist()}::DOUBLE[{num_dim}]) as distance
from {self.vector_store_name}
order by distance
limit 10
""").df()
# 検索結果
result = df.iloc[distance["index"].tolist()])
DuckDBのテーブルに対するクエリだけ抜き出すと以下のような形です。
コサイン類似度を計算する際に、インデックスと検索クエリの埋め込みベクトルの次元数を、DOUBLE[1024]
のように明示的に指定する必要がありました。
select
index,
-- コサイン類似度による距離を計算
-- ベクトルの次元数を明示的に指定する
array_cosine_distance(embedding::DOUBLE[{num_dim}], {embeddings_query.tolist()}::DOUBLE[{num_dim}) as distance
from {self.vector_store_name}
-- 距離が小さい順にソート
order by distance
limit 10
array_cosine_distance
以外にも距離を計算する関数があります。詳しくは、公式のドキュメントも見てみてください。
まとめ
DuckDBは、DataFrameからテーブルが作れたり、クエリの結果をDataFrameで得られたり、Pythonでのデータ分析と相性がめちゃくちゃ良いなと思いました。ベクトル検索も、コサイン類似度の計算で次元数の指定しないとエラーで計算できなかったところでちょっとつまづきましたが、割と簡単に実装できました。DuckDBはデータ分析で今後も活用していきたいです!
Discussion
Yes, DuckDB is great!! On 50mb and runs in memory only if needed.
I use it instead of Postgres and with Grafana.
Regards,
Rob Oudendijk