🦆

ベクトル検索による文書検索にDuckDBを使ってみた

2025/02/01に公開1

はじめに

最近よく見るデータ分析(OLAP)向けのデータベース「DuckDB」は、標準でベクトル検索に対応しています。DuckDBはPandasと連携しやすいという特徴もあります。
https://qiita.com/ak-sakatoku/items/54ed6ab29708ed4a6bb9

そこで、今回はちょっと実践的なベクトル検索にDuckDBを使ってみました。
今回実装したものはこちらのリポジトリにアップしています。
https://github.com/Shakshi3104/cobalt-duckdb
また、Gradioを用いて作成したデモアプリをHugging Face Spaceにアップしています。
https://huggingface.co/spaces/Shakshi3104/Cobalt-DuckDB

ベクトル検索をDuckDBで実装する

ベクトル検索に使う埋め込みモデルは、以前ハイブリッド検索を実装してみたときに使用した「Ruri」を今回も使用します。
https://huggingface.co/cl-nagoya/ruri-large
https://qiita.com/Shakshi3104/items/6ca3882ba45a4924bf0d

インデックスの作成

検索文書のインデックス作成をするためには、文書を埋め込みベクトルに変換しDuckDBに入れる必要があります。DuckDBはDataFrameから簡単にテーブルを作成できるので、DataFrameを作成してDuckDBを作成してみます。
https://duckdb.org/docs/api/python/data_ingestion

以下のようなコードでインデックスを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以外にも距離を計算する関数があります。詳しくは、公式のドキュメントも見てみてください。
https://duckdb.org/docs/sql/functions/array.html

まとめ

DuckDBは、DataFrameからテーブルが作れたり、クエリの結果をDataFrameで得られたり、Pythonでのデータ分析と相性がめちゃくちゃ良いなと思いました。ベクトル検索も、コサイン類似度の計算で次元数の指定しないとエラーで計算できなかったところでちょっとつまづきましたが、割と簡単に実装できました。DuckDBはデータ分析で今後も活用していきたいです!

参考

Discussion

rob oudendijkrob oudendijk

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