Zenn
🌏

DuckDBで1024次元×10万レコードから類似検索してみた

2024/12/21に公開
SQL
機械学習
Database
numpy
DuckDB
tech

はじめに

先日、「NumPy構造体配列をDuckDBのテーブルに格納する」という記事を書きました。
DuckDBにはarray_cosine_similarityという関数があり、これを使えば類似検索を簡単に実装できそうだったので、試してみました。

https://zenn.dev/yuyakato/articles/af67b301d37e17

環境

今回は以下の環境で実験しました。

  • ハードウェア: MacBook Air M2 2022
    • メモリ: 24GB
  • Python: 3.11.6

Pythonには以下のライブラリをインストールしています。

requirements.txt
base58==2.1.1
contourpy==1.3.1
cycler==0.12.1
duckdb==1.1.3
fonttools==4.55.3
joblib==1.4.2
kiwisolver==1.4.7
matplotlib==3.10.0
numpy==2.2.0
packaging==24.2
pillow==11.0.0
polars==1.17.1
pyarrow==18.1.0
pyparsing==3.2.0
python-dateutil==2.9.0.post0
ruff==0.8.4
scikit-learn==1.6.0
scipy==1.14.1
six==1.17.0
threadpoolctl==3.5.0

コード

コード一式は、以下のGitHubリポジトリにも格納しています。

https://github.com/nayutaya/202412-duckdb-similarity-search/tree/main/vector-similarity-search

予備実験: 3次元×10,000レコード

まずは予備実験として、小さな次元数、少ないレコード数での類似検索を行ってみたいと思います。
なお、本記事における「類似」の比較には「コサイン類似度」を用います。

実験用データを作る

可視化しやすいので次元数は3、データ数は10,000としてランダムな実験用データを作ります。
次元数とデータ数を指定してランダムなデータ（L2正規化済み）を生成するPythonスクリプトは以下の通りです。

make_random.py
import os
import sys

import base58
import duckdb
import numpy as np
import polars as pl
from sklearn.preprocessing import normalize

# 次元数、データ数をコマンドライン引数から取得する
n_dims = int(sys.argv[1])
n_records = int(sys.argv[2])

# ランダムなデータを含むNumPy構造体配列を生成する
records_numpy = np.zeros(
    (n_records,), dtype=[("id", "U11"), ("feature", np.float32, (n_dims,))]
)
records_numpy["id"] = np.array(
    [base58.b58encode(os.urandom(8)).decode() for _ in range(n_records)]
)
records_numpy["feature"] = normalize(
    np.random.rand(n_records, n_dims).astype(np.float32) - 0.5
)

# Polarsのデータフレームに変換する
records_df = pl.DataFrame(records_numpy)

# DuckDBに書き込む
db_file_name = f"random_dim{n_dims}.duckdb"
print(db_file_name)
con = duckdb.connect(db_file_name)
con.sql(f"CREATE TABLE records(id VARCHAR PRIMARY KEY, feature FLOAT4[{n_dims}])")
con.sql("INSERT INTO records SELECT * FROM records_df")
print(con.sql("SELECT COUNT(*) FROM records"))

このPythonスクリプトを使って3次元、10,000レコードの実験用データを生成しましょう。

$ python make_random.py 3 10000
random_dim3.duckdb
┌──────────────┐
│ count_star() │
│    int64     │
├──────────────┤
│        10000 │
└──────────────┘

生成されたrandom_dim3.duckdbファイルのサイズは1323008（約1.3MB）バイトでした。
純粋なデータ量は、1レコードが(11+4*3)=23バイト、10,000レコードが(11+4*3)*10000=230000バイトなので、1323008/230000=5.75と約6倍のストレージが必要そうです。容量効率の面では微妙ですね。

実験用データを可視化する

実験用データが想定通りに生成されているかどうかを確認するために可視化してみます。

plot_random_dim3.py
import duckdb
import matplotlib.pyplot as plt
import numpy as np

# DuckDBからfeature列を取得する
con = duckdb.connect("random_dim3.duckdb")
features = np.vstack(
    con.sql("SELECT feature FROM records").fetchnumpy()["feature"].tolist()
)
print((features.dtype, features.shape))

# 散布図を描画する
fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(111, projection="3d")
ax.scatter(xs=features[:, 0], ys=features[:, 1], zs=features[:, 2], s=1, alpha=0.7)
plt.tight_layout()
plt.savefig("random_dim3.png", dpi=300, bbox_inches="tight")
plt.show()
plt.close()

$ python plot_random_dim3.py
(dtype('float32'), (10000, 3))

いい感じに3次元球が見えます。3次元の実験用データは想定通りに生成されているようです。
上記は出力された静止画を貼っていますが、実際にPythonスクリプトを動かすとマウスで視点をグリグリできます。

クエリベクトルを生成する

コサイン類似度を比較する対象のクエリベクトルを作りましょう。

make_query.py
import sys

import numpy as np
from sklearn.preprocessing import normalize

# 次元数をコマンドライン引数から取得する
n_dims = int(sys.argv[1])

# ランダムなクエリベクトルを生成する
query = normalize(np.random.rand(1, n_dims).astype(np.float32) - 0.5)[0]
print((query.dtype, query.shape))

# クエリベクトルを保存する
query_file_name = f"query_dim{n_dims}.npy"
print(query_file_name)
np.save(query_file_name, query)

$ python make_query.py 3
(dtype('float32'), (3,))
query_dim3.npy

クエリベクトルが生成されました。

類似検索＆可視化する

次に、DuckDBでコサイン類似度を算出し、可視化してみましょう。

plot_similarity_dim3.py
import duckdb
import matplotlib.pyplot as plt
import numpy as np

# クエリベクトルを読み込む
query = np.load("query_dim3.npy")

# DuckDBでコサイン類似度を算出する
con = duckdb.connect("random_dim3.duckdb")
result = con.execute(
    "SELECT feature, array_cosine_similarity(feature, $query::FLOAT4[3]) AS similarity FROM records",
    {"query": query},
).fetchnumpy()

features = np.vstack(result["feature"])
similarities = result["similarity"]
similarity_mask = similarities >= 0.9

# 結果を可視化する
# * 原点からクエリベクトル点を赤い線として描画する
# * コサイン類似度0.9未満の点を黒い点として描画する
# * コサイン類似度0.9以上の点を類似度で着色し描画する
fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(111, projection="3d")
ax.plot(xs=[0, query[0]], ys=[0, query[1]], zs=[0, query[2]], color="red", linewidth=1)
ax.scatter(
    xs=features[~similarity_mask][:, 0],
    ys=features[~similarity_mask][:, 1],
    zs=features[~similarity_mask][:, 2],
    c="black",
    s=1,
    alpha=0.1,
)
ax.scatter(
    xs=features[similarity_mask][:, 0],
    ys=features[similarity_mask][:, 1],
    zs=features[similarity_mask][:, 2],
    c=similarities[similarity_mask],
    s=2,
    alpha=0.7,
)
plt.tight_layout()
plt.savefig("similarity_dim3.png", dpi=300, bbox_inches="tight")
plt.show()
plt.close()

$ python plot_similarity_dim3.py

ちゃんと類似検索を行えているようですね。
今回はクエリベクトルを赤い線で示し、コサイン類似度が0.9以上の点を距離に応じて色づけしてみました。

本実験: 1024次元×100,000レコード

予備実験も上手くいったので、実際のユースケースに近い実験用データを生成してみましょう。
今回はJina CLIP v2の特徴量を使うことを想定し、1024次元、100,000レコードで試してみました。
なお、Jina CLIP v2については、別途記事を書きました。

https://zenn.dev/yuyakato/articles/175d6d590da13a

実験用データを作る

予備実験と同じようにランダムなデータを作ります。

$ python make_random.py 1024 100000
random_dim1024.duckdb
┌──────────────┐
│ count_star() │
│    int64     │
├──────────────┤
│       100000 │
└──────────────┘

生成されたrandom_dim1024.duckdbファイルのサイズは662188032（約622MB）バイトでした。
純粋なデータ量は(11+4*1024)*100000=410700000バイトなので、662188032/410700000=1.61倍の大きさ。レコード数が多ければ容量効率は上がりそうです。

ちなみに今回はduckdbファイルとして保存していますが、600MB程度であれば余裕でオンメモリ処理できそうですね。

クエリベクトルを生成する

予備実験と同じように、クエリベクトルを生成します。

$ python make_query.py 1024
(dtype('float32'), (1024,))
query_dim1024.npy

類似検索する

1024次元、100,000レコードからSQL文を使って類似検索してみます。
今回はコサイン類似度の高い上位10レコードを出力してみました。

search_similar_dim1024.py
import time

import duckdb
import numpy as np

# クエリベクトルを読み込む
n_dims = 1024
query = np.load(f"query_dim{n_dims}.npy")

# DuckDBでコサイン類似度を算出し、上位10レコードを取得する
con = duckdb.connect(f"random_dim{n_dims}.duckdb")
start_time = time.perf_counter()
similar_records = con.execute(
    f"""
    SELECT
      id,
      array_cosine_similarity(feature, $query::FLOAT4[{n_dims}]) AS similarity
    FROM records
    ORDER BY similarity DESC
    LIMIT 10
    """,
    {"query": query},
).fetchall()
end_time = time.perf_counter()

# 検索結果を表示する
for id, similarity in similar_records:
    print(f"{id}: {similarity:.6f}")

print(f"time: {end_time - start_time:.3f} sec")

$ python search_similar_dim1024.py
Cv6286ob1Qs: 0.133982
fWFsfiu1cua: 0.132817
VbzppVTpQ85: 0.130720
QNCMKn4JQjH: 0.125984
BRdc5nA9Er7: 0.121714
TeJWRacEFKE: 0.120769
N8SeBAZQezx: 0.118648
SdB5uVdxpRB: 0.114718
CNWMbYNRE4y: 0.114493
fzFajJPn9d8: 0.114473
time: 0.584 sec

1024次元と高次元なのでコサイン類似度はだいぶ小さな値になっていますが、確かに値が大きいものが抽出できてそうですね。
しかも検索時間も600ミリ秒程度。MacBook Airでこの速度なら、十分実用的かと思います。

おわりに

DuckDBを使えばとても簡単に類似検索を行うことができました。
今回はid、featureという2つの列を持つ1つのテーブルしかなったのでありがたみは薄い（NumPyで十分と感じる）かもしれないですが、列が増え、テーブルが増えたときには、DuckDBの良さがより活きてくる気がしています。
また、DuckDBにはベクトルに対してインデックスを生成する機能もあるみたいなので、より素早い処理も期待できそうです。

本記事が何らかの参考になれば幸いです。

Discussion