DuckDBのVector Searchを試してみる

この記事はDuckDB Advent Calendar 2024の9日目です。

DuckDBとは

DuckDBはデータ分析に特化した高速な組み込み型SQLデータベースです。
SQLiteが一般的なトランザクション処理向けであるのに対して、DuckDBは大量データの集計や分析処理(OLAP)に特化しており、より高速な処理が可能になっています。

面白い機能が色々備えられているのですが、例えば以下の特徴があります。

• メモリ内で動作する
• 多様な言語から利用可能
• Pythonの場合、Pandasと簡単に連携できる
• 地理空間データ処理ができるようになるspatialやS3からSELECTできるようになるhttpfsなど、標準機能を拡張する仕組みを有している

Vector Searchとは

Vector Searchとは、多次元ベクトル空間において与えられたベクトル(数値の配列)に類似するベクトルを検索する技術です。
テキスト、画像、音声などのデータをベクトルに変換して、多次元ベクトル空間にマッピングしていくことで、データを意味的な類似性で検索することができるようになります。
言うまでもなく、LLMの文脈で非常に注目されていますね。

標準機能で試してみる

DuckDBはVector型に対応しているので、Vector Searchすることが可能です。

SQL

DuckDBのコマンドラインからは以下のように試すことができます。

CREATE TABLE my_vector_table (i INTEGER, v FLOAT[3]);
INSERT INTO my_vector_table VALUES
    (11, array_value(1.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT)),
    (12, array_value(1.0::FLOAT, 1.0::FLOAT, 2.0::FLOAT)),
    (13, array_value(1.0::FLOAT, 1.0::FLOAT, 3.0::FLOAT)),
    (14, array_value(1.0::FLOAT, 2.0::FLOAT, 1.0::FLOAT)),
    (15, array_value(1.0::FLOAT, 2.0::FLOAT, 2.0::FLOAT)),
    (16, array_value(1.0::FLOAT, 2.0::FLOAT, 3.0::FLOAT)),
    (17, array_value(1.0::FLOAT, 3.0::FLOAT, 1.0::FLOAT)),
    (18, array_value(1.0::FLOAT, 3.0::FLOAT, 2.0::FLOAT)),
    (19, array_value(1.0::FLOAT, 3.0::FLOAT, 3.0::FLOAT)),
    (21, array_value(2.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT)),
    (22, array_value(2.0::FLOAT, 1.0::FLOAT, 2.0::FLOAT)),
    (23, array_value(2.0::FLOAT, 1.0::FLOAT, 3.0::FLOAT)),
    (24, array_value(2.0::FLOAT, 2.0::FLOAT, 1.0::FLOAT)),
    (25, array_value(2.0::FLOAT, 2.0::FLOAT, 2.0::FLOAT)),
    (26, array_value(2.0::FLOAT, 2.0::FLOAT, 3.0::FLOAT)),
    (27, array_value(2.0::FLOAT, 3.0::FLOAT, 1.0::FLOAT)),
    (28, array_value(2.0::FLOAT, 3.0::FLOAT, 2.0::FLOAT)),
    (29, array_value(2.0::FLOAT, 3.0::FLOAT, 3.0::FLOAT)),
    (31, array_value(3.0::FLOAT, 1.0::FLOAT, 1.0::FLOAT)),
    (32, array_value(3.0::FLOAT, 1.0::FLOAT, 2.0::FLOAT)),
    (33, array_value(3.0::FLOAT, 1.0::FLOAT, 3.0::FLOAT)),
    (34, array_value(3.0::FLOAT, 2.0::FLOAT, 1.0::FLOAT)),
    (35, array_value(3.0::FLOAT, 2.0::FLOAT, 2.0::FLOAT)),
    (36, array_value(3.0::FLOAT, 2.0::FLOAT, 3.0::FLOAT)),
    (37, array_value(3.0::FLOAT, 3.0::FLOAT, 1.0::FLOAT)),
    (38, array_value(3.0::FLOAT, 3.0::FLOAT, 2.0::FLOAT)),
    (39, array_value(3.0::FLOAT, 3.0::FLOAT, 3.0::FLOAT)),
;
SELECT * FROM my_vector_table;
SELECT *, array_distance(v, [1, 2, 3]::FLOAT[3]) as d FROM my_vector_table ORDER BY d LIMIT 10;

今回は試したデータの次元数が低いですが、LLMの文脈で使うような768次元や1024次元でも対応してくれそうです。

Python

pip install duckdbでPythonのプログラムに組み込む形でDuckDBを使うことができます。
この場合も同様にVector Searchすることができます。
今回はUIとしてStreamlitを使っています。

import duckdb
import streamlit as st

# タイトルを出力
st.title("DuckDB test")

# バージョンを出力
st.write(f"DuckDB version: {duckdb.__version__}")

# Vectorを含むテーブルを作成
duckdb.query("CREATE OR REPLACE TABLE my_vector_table (i INTEGER, v FLOAT[3])")

# Vectorデータを挿入する
idx = 0
for x in range(11):
    for y in range(11):
        for z in range(11):
            duckdb.query(f"INSERT INTO my_vector_table VALUES ({idx}, [{x}, {y}, {z}])")
            idx += 1

# Vectorデータを取得する
result = duckdb.sql("SELECT * FROM my_vector_table").df()
st.dataframe(result)

# 類似度検索
x = st.slider("x", 0.0, 10.0, 5.0, step=0.1)
y = st.slider("y", 0.0, 10.0, 5.0, step=0.1)
z = st.slider("z", 0.0, 10.0, 5.0, step=0.1)
result = duckdb.sql(f"SELECT *, array_distance(v, [{x}, {y}, {z}]::FLOAT[3]) as d FROM my_vector_table ORDER BY d LIMIT 30").df()
st.dataframe(result)

補足

ここまで、array_distanceのみ試しましたが、以下のようなものも使えそうです。

Array Function	意味
array_distance	ベクトル間のユークリッド距離
array_cosine_similarity	ベクトルのコサイン類似度
array_cosine_distance	1からコサイン類似度を引いたメトリクス、0から2の範囲になる
array_negative_inner_product	2つのベクトルの内積の負の値。ベクトルの長さが正規化されていればコサイン類似度と同じ挙動になる

詳細はこちらのドキュメントをご確認ください。

拡張機能

VSSという拡張機能をインストールすることで、HNSW(Hierarchical Navigable Small World)といういい感じのアルゴリズムでVector Searchするためのインデックスを張れるそうです。
引き続き、こちらも試してみたい！