💾
Faissによるベクトル検索の検証とその先の選択肢（SQLite / DuckDB / RDB / BigQuery）
はじめに
昨今の機械学習やLLMの発展の結果、ベクトル検索が多用されています。
たとえば、OpenAIで作成したエンベディング表現を用いて類似のテキストを検索する方法は、RAG（Retrieval-Augmented Generation）の前処理としてもよく活用されています。
ベクトル検索を行う場合、データ量が少なければFaissを使用してメモリ上で検索するだけで対応できますが、規模が大きくなると困難になります。
本記事ではまず Faiss を復習し、次に「メモリに載らない」ケースでどう扱うかを整理します。続いて、SQLite / DuckDB といった組み込み DB でのベクトル検索を試し、最後に PostgreSQL / MariaDB / BigQuery での実験結果を示します。
本記事は以下の知識を前提とします。
- pythonが実装できること
- Faissによるベクトル検索を使ったことがある
- dockerで環境を構築できること
- PostgreSQL, MariaDB, BigQueryの最低限の知識
実験環境
必要なライブラリ
以下のファイルのdependenciesで指定されているライブラリをuv addまたはpip installで実験環境にインストールする必要があります。
[project]
name = "vecdb"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"faiss-cpu>=1.13.2",
"numpy>=2.4.0",
"openai>=2.14.0",
"pydantic>=2.12.5",
"python-dotenv>=1.2.1",
"tqdm>=4.67.1",
"ipykernel",
"datasets>=4.4.2",
"sqlite-vec>=0.1.6",
"pgvector>=0.4.2",
"mysql-connector-python>=9.5.0",
"psycopg[binary]>=3.3.2",
"pymysql>=1.1.2",
"duckdb>=1.4.3",
"duckdb-extensions>=1.4.3",
"duckdb-extension-vss>=1.4.3",
"google-cloud-bigquery>=3.40.0",
]
[dependency-groups]
dev = [
"pyrefly>=0.47.0",
"pytest>=9.0.2",
"pytest-cov>=7.0.0",
"ruff>=0.14.10",
]
実験に使用するデータセット
ベクトル検索の実験を行うためのデータセットとして、英語wikipediaをOpenAIのtext-embedding-ada-002でエンベディングした結果を利用します。
import numpy as np
from datasets import load_dataset
# OpenAI埋め込みだけ読む
DATA_URL = "https://huggingface.co/datasets/Supabase/wikipedia-en-embeddings/resolve/main/wiki_openai.ndjson.gz"
# streaming=False: まとめてDLして処理（サンプル用途）
ds = load_dataset("json", data_files=DATA_URL, split="train")
#
ds.num_rows
列としては以下の３つがあります。
- id
- body: エンベディング元の文章
- text-embedding-ada-002: 1536次元のfloat64の配列
また、データ件数は224482件存在します。
Faissの復習
Facebook AI Similarity Search (Faiss)はベクトルの類似検索を行う場合に有力の候補となるライブラリです。まずはこのライブラリについて復習します。
ここでは実験環境で取得したデータセットを使用してベクトル検索を実験します。
単純なサンプル
今回のデータセットの次元数とデータ量が大きいため、メモリに乗せて動かすのは難しいです。そのため、今回はデータ量を間引いて実験します。
# データ量が大きいので間引く
N = 10_000
docs = ds.select(range(min(N, ds.num_rows)))
texts = docs["body"]
emb = np.array(docs["text-embedding-ada-002"], dtype=np.float32) # (N, dim)
dim = emb.shape[1]
print('dim=', emb.shape)
IndexFlatIPを使用して内積での類似検索を行うためのインデックスを作成します。
import faiss
import numpy as np
index = faiss.IndexFlatIP(dim) # exact search (inner product)
index.add(emb)
print("is_trained", index.is_trained)
検索対象のクエリー用のベクトルを作成します。今回はindex:123のドキュメントに類似するものを探すことにします。
q_idx = 123
query_emb = emb[q_idx : q_idx + 1] # shape (1, dim)
print('query_emb=', query_emb.shape, query_emb, texts[q_idx][:120].replace("\n", " "), "...")
query_emb= (1, 1536) [[-0.01161382 -0.00404291 0.02300174 ... 0.01169355 0.00800609
-0.00926182]] Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
作成したindexを使用してクエリ用のベクトルに近い文章・上位10件をindex.searchで検索します。
k = 10
scores, ids = index.search(query_emb, k)
print("QUERY:", texts[q_idx][:120].replace("\n", " "), "...")
print("--- top-k ---")
for rank, (doc_id, score) in enumerate(zip(ids[0], scores[0]), start=1):
print(f"{rank:02d} score={score:.4f} text={texts[doc_id][:120].replace('\n', ' ')} ...")
この関数は検索結果として、近さの度合いと対象のidを返します。
検索結果
QUERY: Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
--- top-k ---
01 score=1.0000 text=Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
02 score=0.8530 text=Title: Tom Graveney Content: thumb|right|200px|Graveney in October 1954 Thomas William "Tom" Graveney OBE (16 June 1927 ...
03 score=0.8361 text=Title: Viv Richards Content: Viv Richards (born in 1952) is a former West Indian cricketer. Richards is from Antigua. He ...
04 score=0.8358 text=Title: Don Shepherd Content: thumb|right|200px|Shepherd in May 2007 Donald John Shepherd (12 August 1927 - 18 August 201 ...
05 score=0.8300 text=Title: David English (cricketer) Content: David Stuart English, CBE (4 March 1946 – 12 November 2022) was a British char ...
06 score=0.8283 text=Title: Rex Whitehead Content: Rex Vernon Whitehead (26 October 1948 - 26 June 2014) was an Australian Test cricket match ...
07 score=0.8258 text=Title: David Allen (cricketer) Content: David Arthur Allen (29 October 1935 – 25 May 2014) was an English cricketer. He ...
08 score=0.8248 text=Title: Steve Waugh Content: thumb|200px|Waugh in 2002 Stephen Rodger Waugh (born 2 June 1965) is a former Australian cri ...
09 score=0.8243 text=Title: Hansie Cronje Content: right|200px Wessel Johannes "Hansie" Cronje (25 September 1969 – 1 June 2002) was a South ...
10 score=0.8227 text=Title: Barry Jarman Content: Barrington Noel Jarman OAM (17 February 1936 – 18 July 2020) was an Australian Test cricket ...
同一ベクトル（同一文書）なのでスコアが 1.0 になっています。以降は内積スコアが高いほど文章が類似していると解釈できます。
次項以降で別の検索方法を試します。そのため、今回の結果との差異を確認できるように、baselineに保存しておきます。
baseline = ids[0].tolist()
baseline
# [123, 6689, 5774, 814, 7407, 7396, 3226, 8757, 5998, 2685]
このbaselineを使用して、今回の結果と別の方法の検索結果について差を確認できます。
np.isin(baseline, ids)
# array([ True, True, True, True, True, True, True, True, True, True])
効率的な検索方法
IndexFlatIPとIndexFlatL2はクエリベクトルに対して インデックス中の全ベクトルとの内積または L2 距離を総当たりで計算します。
この方法は近似誤差が存在しない一方、インデックス中のデータやベクトルの次元数が大きくなると計算量が増加し、処理速度の要件を満たせなくなるケースがあります。
そういったケースでは、精度と速度のトレードオフを行います。
たとえば、IndexIVFFlatでは、クラスタで分割を行い、検索時に探索するクラスタ数を限定することで計算量を削減します。
Faissには他にもインデックスを作成する手法は存在しますが、実行速度、メモリ使用量、精度のバランスをうまくとる必要があります。
HNSWインデックス
HNSW(Hierarchical Navigable Small World) インデックスはPostgreSQL: pgvectorなどでもサポートされている近似近傍探索アルゴリズムになります。
ベクトル同士をグラフ構造で接続し、そのグラフを辿って近傍探索を行います。
詳細については以下の解説を参照してください。
実際の使用例は以下の通りです。
indexの構築方法
IndexHNSWFlatを使用してインデックスを作成します。
# HNSW
# 近いベクトル同士をグラフで結び、上位→下位の多層構造で高速に近傍探索する kNN インデックス
M = 8 # グラフの近傍数（大きいほど精度↑/メモリ↑）
index = faiss.IndexHNSWFlat(dim, M, faiss.METRIC_INNER_PRODUCT)
# 検索精度/速度に効くパラメータ
index.hnsw.efConstruction = 32 # 構築時（大きいほど精度↑/遅い）
index.hnsw.efSearch = 4 # 検索時（大きいほど精度↑/遅い）
index.add(emb)
検索方法
k = 10
scores, ids = index.search(query_emb, k)
これを実行した場合、インデックスの構築についてはIndexFlatIPのときより時間がかかります。
検索はデータ量によって変わりますが、データ量が十分に大きければ、IndexFlatIPより早くなっているかと思います。
私の実験環境ではIndexFlatIPとの結果の差は、9/10になります。
np.isin(baseline, ids)
# array([ True, True, True, True, True, True, True, True, False, True])
これは近似による処理のため、IndexFlatIPとまったく同じ結果になることが保証されていません。
もし、速度を犠牲にして、精度を上げたい場合はIndexHNSWFlatのパラメータを調整します。MやefSearchを増やすことで検索時の速度を犠牲にして精度を上げることが可能になります。
# パラメータを調整すると精度をあげることはできる
M = 16 # グラフの近傍数（大きいほど精度↑/メモリ↑）
index = faiss.IndexHNSWFlat(dim, M, faiss.METRIC_INNER_PRODUCT)
# 検索精度/速度に効くパラメータ
index.hnsw.efConstruction = 64 # 構築時（大きいほど精度↑/遅い）
index.hnsw.efSearch = 8 # 検索時（大きいほど精度↑/遅い）
IVFインデックスを使用した例
IVF(Inverted File)ではクラスタリングによる検索範囲を制限します。
以下がその例となります。
インデックスの構築
この例では128のクラスタに分割します。
nlist = 128
quantizer = faiss.IndexFlatIP(dim)
index = faiss.IndexIVFFlat(quantizer, dim, nlist)
IVFインデックスを使用するには、クラスタに割り当てるかを決定するためのindexであるquantizerが必要です。
また、クラスタ分割を行うためのトレーニングが必要です。
この時点ではindex.is_trainedはFalseとなっているので、以下のコードを実行してトレーニングを行います。
index.train(emb)
これにより、index.is_trainedはTrueとなります。
あとはIndexIVFFlatにベクトルを追加します。
index.add(emb)
IndexIVFFlatで検索を実施する場合はまず、index.nprobeを決定します。
以下の例ではクラスタを1つしかみないため、処理速度は向上しますが、精度は落ちるでしょう。
index.nprobe = 1
k = 10
scores, ids = index.search(query_emb, k)
もし、index.nprobeを増やせば検索するクラスタが増える分、処理速度は落ちますが、精度は向上します。
index.nprobe = 2
k = 10
scores, ids = index.search(query_emb, k)
私の環境ではnp.isin(baseline, ids)の結果について、index.nprobe = 1の場合は9/10, index.nprobe = 2の場合は10/10となりました。
ディスクを使用したインデックス
今までの実験ではメモリ使用量の制限により、データ件数を制限して実験を行っていました。
IndexIVFFlat の inverted lists を OnDiskInvertedLists としてディスク（.ivfdata）に置くことで、巨大なリスト本体を RAM に載せずに検索できます。
まず、データセットの内容をディスクに保存します。
- テキスト情報→sqliteに保存する
- エンベディング結果→np.memmapに保存する
具体的なコードは以下のようになります。
import sqlite3
from datasets import load_dataset
from tqdm import tqdm
import numpy as np
dim = 1536 # OpenAI ada-002 の場合
DATA_URL = "https://huggingface.co/datasets/Supabase/wikipedia-en-embeddings/resolve/main/wiki_openai.ndjson.gz"
ds = load_dataset("json", data_files=DATA_URL, split="train")
emb_path = "embeddings.f32"
embeddings = np.memmap(emb_path, dtype="float32", mode="w+", shape=(ds.num_rows, dim))
with sqlite3.connect("docs.sqlite") as conn:
cur = conn.cursor()
cur.execute("CREATE TABLE IF NOT EXISTS docs (id INTEGER PRIMARY KEY, body TEXT NOT NULL)")
conn.commit()
for i, row in tqdm(enumerate(ds)):
v = np.asarray(row["text-embedding-ada-002"], dtype=np.float32)
embeddings[i] = v
cur.execute("INSERT INTO docs(id, body) VALUES (?, ?)", (i, row["body"]))
if (i+1) % 10000 == 0:
conn.commit()
embeddings.flush()
conn.commit()
embeddings.flush()
np.memmapで保存したエンベディングの結果は以下のような実装で取得することができます。
tmp = np.memmap("embeddings.f32", dtype="float32", mode="r", shape=(ds.num_rows, dim))
print(tmp[0])
print(tmp[-1])
クラスタ分割のためのトレーニングを実施します。今回は全件ではなく2万件のみ実施します。訓練済みのインデックスはtrained.indexとして保存します。
import faiss
import numpy as np
import tracemalloc
tracemalloc.start()
tracemalloc.reset_peak()
dim = 1536
nlist = 128
quantizer = faiss.IndexFlatIP(dim)
index = faiss.IndexIVFFlat(quantizer, dim, nlist, faiss.METRIC_INNER_PRODUCT)
# train はサンプルでOK（例：先頭20k）
train_n = 20_000
xt = np.memmap("embeddings.f32", dtype="float32", mode="r", shape=(ds.num_rows, dim))[:train_n].copy()
index.train(xt)
faiss.write_index(index, "trained.index")
続いて、トレーニング済みのIndexIVFFlatにデータを追加していきます。この例では2万件ごとにデータ追加して、その結果を別ファイルに保存します。
import numpy as np
import faiss
import os
import tracemalloc
tracemalloc.start()
tracemalloc.reset_peak()
dim = 1536
chunk = 20_000 # 1チャンク何件ずつ add するか（RAMに合わせて調整）
os.makedirs("shards", exist_ok=True)
xb = np.memmap("embeddings.f32", dtype="float32", mode="r", shape=(ds.num_rows, dim))
for shard_id, start in tqdm(enumerate(range(0, ds.num_rows, chunk))):
end = min(start + chunk, ds.num_rows)
sub = faiss.read_index("trained.index")
x = np.array(xb[start:end], dtype=np.float32, copy=True)
ids = np.arange(start, end, dtype=np.int64)
sub.add_with_ids(x, ids)
out = f"shards/block_{shard_id:05d}.index"
faiss.write_index(sub, out)
print("wrote", out, "size", end-start)
merge_ondiskを使用して、複数シャード（block_.index）が持つ inverted listsを 1 つの on-disk ファイル（.ivfdata）に連結し、trained 側に “ディスク上の inverted lists” として接続します。
import faiss
from faiss.contrib.ondisk import merge_ondisk
import glob
trained = faiss.read_index("trained.index")
block_files = sorted(glob.glob("shards/block_*.index"))
ivfdata = "merged.ivfdata"
merge_ondisk(trained, block_files, ivfdata)
faiss.write_index(trained, "populated_ondisk.index")
これを検索する例は以下の通りです。
import faiss
import sqlite3
index = faiss.read_index("populated_ondisk.index")
index.nprobe = 32
D, I = index.search(query_emb, 10)
# 本文を外部ストアから取得
conn = sqlite3.connect("docs.sqlite")
cur = conn.cursor()
for doc_id, score in zip(I[0], D[0]):
cur.execute("SELECT body FROM docs WHERE id=?", (int(doc_id),))
body = cur.fetchone()[0]
print(score, doc_id, body[:120].replace("\n", " "), "...")
populated_ondisk.index だけでは検索できず、merged.ivfdata が参照可能なパスにある必要があります。
nprobe=32の場合
1.0000002 123 Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
0.85713804 99980 Title: Frank Worrell Content: Sir Frank Mortimer Maglinne Worrell (1 August 1924, Bank Hall, St Michael, Barbados – 13 M ...
0.8530156 6689 Title: Tom Graveney Content: thumb|right|200px|Graveney in October 1954 Thomas William "Tom" Graveney OBE (16 June 1927 ...
0.85217476 14905 Title: Easton McMorris Content: Easton Dudley Ashton St John McMorris (4 April 1935 – 1 February 2022) was a West Indian ...
0.8499083 208851 Title: West Indies cricket team Content: The West Indian cricket team, also known as the West Indies or the Windies, is ...
0.8475273 23438 Title: Shane Warne Content: Shane Keith Warne (13 September 1969 – 4 March 2022) was an Australian cricketer. He captain ...
0.84659135 142127 Title: Warren Wisneski Content: Warren Anthony Wisneski was born 19 February, 1969 in New Plymouth, Taranaki. He was a N ...
0.8457886 73128 Title: Bevan Congdon Content: Bevan Ernest Congdon OBE (11 February 1938 – 10 February 2018) was a New Zealand cricket a ...
0.8445078 120439 Title: Tony Cozier Content: Winston Anthony Lloyd "Tony" Cozier (10 July 1940 – 11 May 2016) was a Barbadian cricket wri ...
0.8432672 66095 Title: Warren Lees Content: Warren Kenneth Lees MBE was born 19 March, 1952 in Dunedin. He is a former New Zealand crick ...
nprobe=1の場合
1.0000002 123 Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
0.85217476 14905 Title: Easton McMorris Content: Easton Dudley Ashton St John McMorris (4 April 1935 – 1 February 2022) was a West Indian ...
0.8475273 23438 Title: Shane Warne Content: Shane Keith Warne (13 September 1969 – 4 March 2022) was an Australian cricketer. He captain ...
0.84659135 142127 Title: Warren Wisneski Content: Warren Anthony Wisneski was born 19 February, 1969 in New Plymouth, Taranaki. He was a N ...
0.8432672 66095 Title: Warren Lees Content: Warren Kenneth Lees MBE was born 19 March, 1952 in Dunedin. He is a former New Zealand crick ...
0.8426095 78127 Title: Peter Walker (cricketer, born 1936) Content: Peter Michael Walker (17 February 1936 – 5 April 2020) was an Englis ...
0.8410813 69403 Title: Eddie Stuart Content: Edward Albert "Eddie" Stuart (12 May 1931 – 4 November 2014) was a South African profession ...
0.83654416 38670 Title: John Watkins (South African cricketer) Content: John Cecil Watkins (10 April 1923 - 6 September 2021) was a South ...
0.83609366 5774 Title: Viv Richards Content: Viv Richards (born in 1952) is a former West Indian cricketer. Richards is from Antigua. He ...
0.8354683 90500 Title: Graeme Watson (cricketer) Content: Graeme Donald Watson (8 March 1945 – 24 April 2020) was an Australian crickete ...
ローカルDBでのベクトル検索(sqlite/DuckDB)
sqliteやDuckDBの拡張機能を使用することでベクトル列によるベクトル検索が可能になります。
これはメモリに乗らない量のベクトル検索を手軽に実験する有力な候補となります。
sqlite+sqlite-vec
sqliteの拡張機能である、sqlite-vecを使用することでベクトル検索が可能になります。
ただし、現時点では近似近傍探索アルゴリズムでの検索はサポートされていません。
では実際にsqlite-vecを使用してみます。
まず、sqliteの拡張機能としてのsqlite_vecを有効にします。
import sqlite3
import sqlite_vec
DB_PATH = "wiki.sqlite"
db = sqlite3.connect(DB_PATH)
db.enable_load_extension(True)
sqlite_vec.load(db)
db.enable_load_extension(False)
vec_version, = db.execute("select vec_version()").fetchone()
print(f"vec_version={vec_version}")
次にベクトルを格納するvec0 Virtual Tableを作成します。
db = sqlite3.connect(DB_PATH)
db.enable_load_extension(True)
sqlite_vec.load(db)
db.enable_load_extension(False)
cur = db.cursor()
# 本文（メタデータ）用テーブル
cur.execute("""
CREATE VIRTUAL TABLE IF NOT EXISTS vec_documents USING vec0(
id INTEGER PRIMARY KEY,
body TEXT NOT NULL,
embedding FLOAT[1536] distance_metric=cosine
)
""")
作成した仮想テーブルにデータを投入する方法は以下の通りです。
from tqdm import tqdm
batch_size = 1_000
total = ds.num_rows
for start in range(0, total, batch_size):
end = min(start + batch_size, total)
docs = ds.select(range(start, end))
print(f"Processing rows {start} to {end} ...")
texts = docs["body"]
emb = np.array(docs["text-embedding-ada-002"], dtype=np.float32) # (N, dim)
db.execute("BEGIN")
for body, vec in tqdm(zip(texts, emb)):
cur.execute("INSERT INTO vec_documents(body, embedding) VALUES (?, ?)", (body,vec.tobytes()))
db.commit()
ベクトルの検索はMATCHを使用して行います。
q = query_emb.astype(np.float32)
k = 10
rows = db.execute("""
SELECT id, body, distance, 1-distance
FROM vec_documents
WHERE embedding MATCH ?
AND k = ?
ORDER BY distance
""", (q.tobytes(), k)).fetchall()
for doc_id, body, dist, similar in rows:
print(doc_id, dist, similar, body[:120].replace("\n", " "), "...")
出力結果
124 0.0 1.0 Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
99981 0.14286094903945923 0.8571390509605408 Title: Frank Worrell Content: Sir Frank Mortimer Maglinne Worrell (1 August 1924, Bank Hall, St Michael, Barbados – 13 M ...
6690 0.14698246121406555 0.8530175387859344 Title: Tom Graveney Content: thumb|right|200px|Graveney in October 1954 Thomas William "Tom" Graveney OBE (16 June 1927 ...
14906 0.14782372117042542 0.8521762788295746 Title: Easton McMorris Content: Easton Dudley Ashton St John McMorris (4 April 1935 – 1 February 2022) was a West Indian ...
208852 0.15009190142154694 0.8499080985784531 Title: West Indies cricket team Content: The West Indian cricket team, also known as the West Indies or the Windies, is ...
23439 0.1524720937013626 0.8475279062986374 Title: Shane Warne Content: Shane Keith Warne (13 September 1969 – 4 March 2022) was an Australian cricketer. He captain ...
142128 0.15340891480445862 0.8465910851955414 Title: Warren Wisneski Content: Warren Anthony Wisneski was born 19 February, 1969 in New Plymouth, Taranaki. He was a N ...
73129 0.15421102941036224 0.8457889705896378 Title: Bevan Congdon Content: Bevan Ernest Congdon OBE (11 February 1938 – 10 February 2018) was a New Zealand cricket a ...
120440 0.1554921567440033 0.8445078432559967 Title: Tony Cozier Content: Winston Anthony Lloyd "Tony" Cozier (10 July 1940 – 11 May 2016) was a Barbadian cricket wri ...
66096 0.15673358738422394 0.8432664126157761 Title: Warren Lees Content: Warren Kenneth Lees MBE was born 19 March, 1952 in Dunedin. He is a former New Zealand crick ...
⚠️idについてDBに採番を任せているため、1スタートとなり、前述のFaissの結果とはずれています。
DuckDB+duckdb-extension-vss
DuckDBは列指向の組み込みで使用できるデータベースです。
Vector Similarity Search Extensionを用いることでベクトルの類似検索が行えます。
近似近傍探索アルゴリズムをサポートしておりインデックスとしてHNSWを付与することが可能です。
では実際にDuckDBによるベクトル検索を実験してみます。
まず、DuckDBのconnectionを作成、拡張機能を有効にします。
import duckdb
DB_PATH = "./wiki.duckdb" # DuckDBのDBファイル（sqliteとは別物）
con = duckdb.connect(DB_PATH)
# vss拡張をインストール＆ロード（core extension）
con.execute("INSTALL vss;")
con.execute("LOAD vss;")
次にテーブルを作成します。
con.execute("""
CREATE SEQUENCE IF NOT EXISTS vec_documents_id_seq START 1;
CREATE TABLE IF NOT EXISTS vec_documents (
id BIGINT PRIMARY KEY DEFAULT nextval('vec_documents_id_seq'),
body TEXT NOT NULL,
embedding FLOAT[1536]
);
""")
デフォルトではファイルで永続化している場合、HNSWインデックスを作成することはできません。
実験機能としては有効なので、もしHNSWインデックスを付与する場合はhnsw_enable_experimental_persistenceを有効にします。
# 永続DBでもHNSWを作れるようにする（experimental）
con.execute("SET hnsw_enable_experimental_persistence = true;")
# HNSW インデックス作成（例：cosine）
# https://duckdb.org/docs/stable/core_extensions/vss
con.execute("""
CREATE INDEX IF NOT EXISTS vec_documents_hnsw_cosine
ON vec_documents
USING HNSW (embedding)
WITH (
metric = 'cosine',
ef_construction = 128,
ef_search = 64,
M = 16,
M0 = 32
);
""")
テーブルを作成したのち、データの投入を行います。
insert文で逐次投入するか、pyarrowのテーブル経由でデータの投入が可能です。
insert文で逐次投入
from tqdm import tqdm
import numpy as np
batch_size = 1000
total = ds.num_rows
for start in range(0, total, batch_size):
con.begin()
end = min(start + batch_size, total)
docs = ds.select(range(start, end))
print(f"Processing rows {start} to {end} ...")
texts = docs["body"]
emb = np.array(docs["text-embedding-ada-002"], dtype=np.float32) # (N, 1536)
rows = [(body, vec.tolist()) for body, vec in zip(texts, emb)]
con.executemany("INSERT INTO vec_documents(body, embedding) VALUES (?, ?)", rows)
con.commit()
pyarrowのテーブル経由
import pyarrow as pa
import numpy as np
batch_size = 10_000
total = ds.num_rows
for start in tqdm(range(0, total, batch_size)):
con.execute("BEGIN TRANSACTION;")
end = min(start + batch_size, total)
batch = ds.select(range(start, end))
texts = batch["body"]
emb = np.asarray(batch["text-embedding-ada-002"], dtype=np.float32)
# DuckDBに渡す形：list[list[float]]（ただし行ループは最小限）
tbl = pa.Table.from_pydict({
"body": texts,
"embedding": [v.tolist() for v in emb],
})
con.register("tmp", tbl)
con.execute("INSERT INTO vec_documents(body, embedding) SELECT body, embedding FROM tmp")
con.commit()
検索を実施する場合は以下のようなコードになります。
q = query_emb.astype(np.float32).tolist()[0]
k = 10
rows = con.execute("""
SELECT
id,
body,
array_cosine_distance(embedding, ?::FLOAT[1536]) AS distance,
(1 - array_cosine_distance(embedding, ?::FLOAT[1536])) AS similarity
FROM vec_documents
ORDER BY distance
LIMIT ?;
""", [q, q, k]).fetchall()
for doc_id, body, dist, sim in rows:
print(doc_id, dist, sim, body[:120].replace("\n", " "), "...")
出力結果
124 0.0 1.0 Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
99981 0.14286094903945923 0.8571390509605408 Title: Frank Worrell Content: Sir Frank Mortimer Maglinne Worrell (1 August 1924, Bank Hall, St Michael, Barbados – 13 M ...
6690 0.14698243141174316 0.8530175685882568 Title: Tom Graveney Content: thumb|right|200px|Graveney in October 1954 Thomas William "Tom" Graveney OBE (16 June 1927 ...
14906 0.14782369136810303 0.852176308631897 Title: Easton McMorris Content: Easton Dudley Ashton St John McMorris (4 April 1935 – 1 February 2022) was a West Indian ...
208852 0.15009188652038574 0.8499081134796143 Title: West Indies cricket team Content: The West Indian cricket team, also known as the West Indies or the Windies, is ...
23439 0.15247207880020142 0.8475279211997986 Title: Shane Warne Content: Shane Keith Warne (13 September 1969 – 4 March 2022) was an Australian cricketer. He captain ...
142128 0.15340888500213623 0.8465911149978638 Title: Warren Wisneski Content: Warren Anthony Wisneski was born 19 February, 1969 in New Plymouth, Taranaki. He was a N ...
73129 0.15421104431152344 0.8457889556884766 Title: Bevan Congdon Content: Bevan Ernest Congdon OBE (11 February 1938 – 10 February 2018) was a New Zealand cricket a ...
120440 0.1554921269416809 0.8445078730583191 Title: Tony Cozier Content: Winston Anthony Lloyd "Tony" Cozier (10 July 1940 – 11 May 2016) was a Barbadian cricket wri ...
66096 0.15673351287841797 0.843266487121582 Title: Warren Lees Content: Warren Kenneth Lees MBE was born 19 March, 1952 in Dunedin. He is a former New Zealand crick ...
⚠️idについてDBに採番を任せているため、1スタートとなり、前述のFaissの結果とはずれています。
外部DBでのベクトル検索（PostgreSQL / MariaDB / BigQuery）
ファイルベースでは複数の環境からの操作が困難です。
ここではVectorデータを独立したサービスで管理する方法について確認します。
今回は以下の事例をみていみます。
- PostgreSQL
- MariaDB
- BigQuery
なお、MySQLについてはMySQL9.xからVector列のサポート自体は行われていますが、距離を計算するためのDISTANCE関数などはHeatWave MySQLユーザーのみが使用でき、コミュニティ版ディストリビューションには含まれていません。
PostgreSQL
pgvector拡張機能を使用することでVECTOR列の使用が可能になります。
近似近傍探索アルゴリズムをサポートしておりインデックスとしてはHNSWとIVFFlatをサポートしています。
また、AWS Aurora PostgreSQLでもpgvectorをサポートしています。そのため、既存の環境にVector検索を導入する場合には有力な候補となるでしょう。
ではPostgreSQLでベクトル検索を動かしてみます。
まず、実際に実験を行うため以下のようなコンテナを作成します。
PostgreSQLのコンテナ
docker-compose.yaml
services:
postgres:
image: pgvector/pgvector:pg16
environment:
POSTGRES_PASSWORD: postgres
POSTGRES_DB: vecdb
POSTGRES_USER: postgres
ports:
- "5432:5432"
volumes:
- ./postgres/init.sql:/docker-entrypoint-initdb.d/init.sql:ro
- ./pgdata:/var/lib/postgresql/data
./postgres/init.sql
CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE IF NOT EXISTS documents (
id BIGSERIAL PRIMARY KEY,
body TEXT NOT NULL,
embedding VECTOR(1536) NOT NULL
);
-- インデックスの付与（任意）
CREATE INDEX IF NOT EXISTS documents_embedding_hnsw
ON documents USING hnsw (embedding vector_cosine_ops);
コンテナを起動したのち、以下のコードでデータを投入します。
import psycopg
import os
from pgvector.psycopg import register_vector
DSN = os.getenv("PG_DSN", "postgresql://postgres:postgres@localhost:5432/vecdb")
batch_size = 1000
total = ds.num_rows
sql = "INSERT INTO documents(body, embedding) VALUES (%s, %s)"
with psycopg.connect(DSN) as conn:
register_vector(conn)
with conn.cursor() as cur:
cur.execute("TRUNCATE TABLE documents")
for start in tqdm(range(0, total, batch_size)):
end = min(start + batch_size, total)
docs = ds.select(range(start, end))
texts = docs["body"]
emb = np.array(docs["text-embedding-ada-002"], dtype=np.float32) # (N, dim)
conn.execute("BEGIN")
rows = [(t, e.tolist()) for t, e in zip(texts, emb)]
cur.executemany(sql, rows)
conn.commit()
なお、hnswインデックスを付与した場合、私の環境では投入に3.25h程度かかりました。
検索については以下のようになります。
import psycopg
import os
from pgvector.psycopg import register_vector
DSN = os.getenv("PG_DSN", "postgresql://postgres:postgres@localhost:5432/vecdb")
q = query_emb.astype(np.float32)
k = 10
# use a 1-D Python list and cast the parameter to pgvector in SQL
param = q[0].tolist()
with psycopg.connect(DSN) as conn:
register_vector(conn)
with conn.cursor() as cur:
cur.execute(
"""
SELECT id, body,
(embedding <=> %s::vector) AS cosine_distance,
(1 - (embedding <=> %s::vector)) AS cosine_similarity
FROM documents
ORDER BY embedding <=> %s::vector
LIMIT %s
""",
(param, param, param, k),
)
rows = cur.fetchall()
for row in rows:
print(row)
出力結果
(224606, 'Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February 1925 - 1 July 2020) was a cricket player for the West Indies. Along with Frank Worrell and Clyde Walcott, he formed what was known as "The Three Ws" of West Indian cricket. He was born in Saint Michael, Barbados. Weekes began his career in 1948. He scored 88 and 117 runs (retired) and was selected for the tour. He played his first match for Barbados on 24 February 1945. He was aged 19 years; it was one day before his 20th birthday. The match was against Trinidad and Tobago at Queen\'s Park Oval, Port of Spain. Batting at number six, he scored 0 and eight runs. Barbados lost by ten wickets.Cricket Archive, Trinidad v Barbados scorecard http://www.cricketarchive.co.uk/Archive/Scorecards/17/17467.html Accessed 24 April 2008. He retired in 1958. Weekes died on 1 July 2020 in Christ Church, Barbados at the age of 95.Cricket legend Sir Everton passes ', 0.0, 1.0)
(324463, 'Title: Frank Worrell Content: Sir Frank Mortimer Maglinne Worrell (1 August 1924, Bank Hall, St Michael, Barbados – 13 March 1967, Kingston, Jamaica) was a West Indies cricketer and Jamaican senator. He is sometimes called by his nickname of Tae. He was a stylish right-handed batsman and useful left-arm seam bowler. He became famous in the 1950s as the first black captain of the West Indies cricket team. He is the only batsman to be in two 500-run partnerships in first-class cricket. 500-run partnerships Since the 1960-61 series, the Frank Worrell Trophy is awarded to the winner of the Test series between Australia and West Indies.Report announcing the Frank Worrell Trophy, Glasgow Herald, 10 February 1961One of the two Halls of Residence at the University of the West Indies, Cave Hill Campus in Barbados is named after him. Caribbean journalist Ernest Eytle wrote a biography of Worrell titled Frank Worrell: The Career of a Great Cricketer in 1963. In June 1988 Worrell was celebrated on the $2 Barbadian stamp along with the Barbados Cricket Buckle. ', 0.142862090041876, 0.857137909958124)
(231172, 'Title: Tom Graveney Content: thumb|right|200px|Graveney in October 1954 Thomas William "Tom" Graveney OBE (16 June 1927 - 3 November 2015) was a leading English cricketer. He represented his country in 79 Test matches and scoring over 4,800 runs. In a career lasting from 1948 to 1972, he was the 15th person to score one hundred first-class centuries, and the first post-war player to do so. ==Early life== Graveney was born in Riding Mill, Hexham, Northumberland on 16 June 1927. His brother was Ken Graveney, also a cricketer. Graveney studied at the Bristol Grammar School. ==Career== He played for Gloucestershire and Worcestershire, and helped Worcestershire win the county championship for the first time in their history. His achievements for England after being recalled in 1966 have been described as "the stuff of legend." Graveney was a Wisden Cricketer of the Year in 1953, captained England on one occasion and was awarded the OBE while still playing. In later life he worked as a cricket commentator for BBC Television and was the first former professional to be President of the Marylebone Cricket Club. He was one of the first 55 players inducted to the ICC Cricket Hall of Fame in 2009. ==Personal life== Graveney was married to Jackie Graveney from 1952 until her death of Alzheimer\'s disease in 2013. Graveney\'s brother, Ken, died on 25 October 2015. A week after his brother\'s death, Graveney died on 3 November 2015 from Parkinson\'s disease at the age of 88. ', 0.14698460173971573, 0.8530153982602843)
(239388, 'Title: Easton McMorris Content: Easton Dudley Ashton St John McMorris (4 April 1935 – 1 February 2022) was a West Indian cricketer. He played in 13 Tests from 1958 to 1966.Wisden 1963, pp. 919–20. He toured England with the West Indies team in 1963 and 1966. He was a heavy scorer for Jamaica. He captained Jamaica from 1967-68 until he retired after the 1971-72 season. McMorris was born in Saint Andrew, Jamaica. McMorris died on 1 February 2022, at the age of 86. ', 0.14782573347515493, 0.8521742665248451)
(433334, 'Title: West Indies cricket team Content: The West Indian cricket team, also known as the West Indies or the Windies, is a multi-national cricket team. They are a sporting group of 15 mainly English- speaking Caribbean countries, British dependencies and non-British dependencies. From the mid 1970s to the early 1990s, the West Indies team was one of the strongest in the world in both Test and One Day International cricket. A number of cricketers among the best in the world have played for the West Indies. Sir Garfield Sobers, Lance Gibbs, Gordon Greenidge, George Headley, Clive Lloyd, Malcolm Marshall, Andy Roberts, Alvin Kallicharran, Rohan Kanhai, Frank Worrell, Everton Weekes, Curtly Ambrose, Michael Holding, Joel Garner and Sir Viv Richards are in the ICC Hall of Fame. World-record holder Brian Lara was a West Indies Test player. The West Indies won the ICC Cricket World Cup twice, in 1975 and 1979. They won the ICC Champions Trophy in 2004. They were runners up in the Under 19 Cricket World Cup in 2004. They were semi finalists in the ICC World Twenty20 in 2009. They were the first cricket team to win the World Cup twice. West Indies are also the first team to win back to back World Cups. West Indies is the first team to appear in three consecutive World Cup finals (1975, 1979 and 1983). ', 0.1500918850110703, 0.8499081149889297)
(247921, "Title: Shane Warne Content: Shane Keith Warne (13 September 1969 – 4 March 2022) was an Australian cricketer. He captained the Australian national team in One Day Internationals (ODI). He is often called the greatest bowlers in the history of the sport. He took 708 wickets from 40,704 balls bowled during his Test career. He has the second most wickets in Test history and most of any leg spin bowler. He retired from international cricket in January 2007, at the end of Australia's 5–0 Ashes series whitewash victory over England. After retirement from playing, he was a cricket commentator. Warne was born in Upper Ferntree Gully, an outer suburb of Melbourne, Victoria. He had three children with his ex-wife Simone Callahan. He had complete heterochromia, from which one of his eyes is blue, and the other green. Warne died from a possible heart attack on 4 March 2022 at his villa on the island of Ko Samui in Thailand. He was 52. ", 0.15247271627507553, 0.8475272837249245)
(366610, 'Title: Warren Wisneski Content: Warren Anthony Wisneski was born 19 February, 1969 in New Plymouth, Taranaki. He was a New Zealand cricketer who played three one-day internationals in 2000 but no tests. He also played for Central Districts and Canterbury in the Plunket Shield before retiring after the 2003/04 season. ', 0.15340908684511478, 0.8465909131548852)
(297611, 'Title: Bevan Congdon Content: Bevan Ernest Congdon OBE (11 February 1938 – 10 February 2018) was a New Zealand cricket all-rounder. He played 61 Test matches and 11 One Day Internationals from 1965 to 1978, which included a spell as captain. He was born in Motueka, New Zealand. He was captain of the New Zealand Test and ODI team from 1972 to 1974, and was the first New Zealand captain to record a victory over Australia. Congdon died in Auckland on 10 February 2018, one day before his 80th birthday. ', 0.15421166319566626, 0.8457883368043337)
(344922, 'Title: Tony Cozier Content: Winston Anthony Lloyd "Tony" Cozier (10 July 1940 – 11 May 2016) was a Barbadian cricket writer. He was a commentator on West Indian cricket from 1958 until his death. Cozier was described as having an "encyclopaedic" knowledge of cricket.Nice Cozier chat Mail Online, 28 April 1999 Cozier was born in Bridgetown, Barbados. He was known for his knowledge of statistics. Cozier died on 11 May 2016 in Bridgetown from a long-illness, aged 75. ', 0.1554922872194361, 0.8445077127805639)
(290578, 'Title: Warren Lees Content: Warren Kenneth Lees MBE was born 19 March, 1952 in Dunedin. He is a former New Zealand cricketer and coach. He played 21 tests and 31 one-day internationals from 1976 to 1983 and played for Otago in the Plunket Shield. Category:1952 births Category:Living people Category:New Zealand cricketers Category:Sportspeople from Dunedin', 0.15673293906854635, 0.8432670609314536)
MariaDB
MariaDB 11.7からVector列がサポートされています。
近似近傍探索アルゴリズムをサポートしておりインデックスとしてはmodified HNSWが使用できます。
まず実験を行うため以下のようなコンテナを作成します。
MariaDBのコンテナ
docker-compose.yaml
services:
mariadb:
image: mariadb:11.8
environment:
MARIADB_ROOT_PASSWORD: mariadb
MARIADB_DATABASE: vecdb
ports:
- "3307:3306"
volumes:
- ./mariadbdata:/var/lib/mysql
- ./mariadb/init.sql:/docker-entrypoint-initdb.d/init.sql:ro
mariadb/init.sql
CREATE TABLE IF NOT EXISTS documents (
id BIGINT AUTO_INCREMENT PRIMARY KEY,
body TEXT NOT NULL,
embedding VECTOR(1536) NOT NULL
);
コンテナに実験データを投入します。
import pymysql
conn = pymysql.connect(
host="127.0.0.1",
port=3307,
user="root",
password="mariadb",
database="vecdb",
autocommit=False,
)
cur = conn.cursor()
def vec_to_text(v: np.ndarray) -> str:
return "[" + ",".join(f"{x:.8g}" for x in v.tolist()) + "]"
batch_size = 1000
total = ds.num_rows
sql = "INSERT INTO documents(body, embedding) VALUES (%s, VEC_FromText(%s))"
#print('TRUNCATE....')
# 時間がかかる
#cur.execute("TRUNCATE TABLE documents")
for start in range(0, total, batch_size):
end = min(start + batch_size, total)
print("Processing rows", start, "to", end, "...")
docs = ds.select(range(start, end))
texts = docs["body"]
emb = np.array(docs["text-embedding-ada-002"], dtype=np.float32) # (N, dim)
cur.execute("START TRANSACTION")
rows = [(t, vec_to_text(e)) for t, e in zip(texts, emb)]
cur.executemany(sql, rows)
conn.commit()
cur.close()
検索については以下のようになります。
q = query_emb.astype(np.float32)
k = 10
q = vec_to_text(q[0])
cur.execute(
"""
SELECT id, body,
VEC_DISTANCE_COSINE(embedding, VEC_FromText(%s)) AS cosine_distance,
(1 - VEC_DISTANCE_COSINE(embedding, VEC_FromText(%s))) AS cosine_similarity
FROM documents
ORDER BY VEC_DISTANCE_COSINE(embedding, VEC_FromText(%s))
LIMIT %s
""",
(q, q, q, k)
)
rows = cur.fetchall()
for row in rows:
print(row)
cur.close()
conn.close()
出力結果
(1124, 'Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February 1925 - 1 July 2020) was a cricket player for the West Indies. Along with Frank Worrell and Clyde Walcott, he formed what was known as "The Three Ws" of West Indian cricket. He was born in Saint Michael, Barbados. Weekes began his career in 1948. He scored 88 and 117 runs (retired) and was selected for the tour. He played his first match for Barbados on 24 February 1945. He was aged 19 years; it was one day before his 20th birthday. The match was against Trinidad and Tobago at Queen\'s Park Oval, Port of Spain. Batting at number six, he scored 0 and eight runs. Barbados lost by ten wickets.Cricket Archive, Trinidad v Barbados scorecard http://www.cricketarchive.co.uk/Archive/Scorecards/17/17467.html Accessed 24 April 2008. He retired in 1958. Weekes died on 1 July 2020 in Christ Church, Barbados at the age of 95.Cricket legend Sir Everton passes ', 0.0, 1.0)
(100981, 'Title: Frank Worrell Content: Sir Frank Mortimer Maglinne Worrell (1 August 1924, Bank Hall, St Michael, Barbados – 13 March 1967, Kingston, Jamaica) was a West Indies cricketer and Jamaican senator. He is sometimes called by his nickname of Tae. He was a stylish right-handed batsman and useful left-arm seam bowler. He became famous in the 1950s as the first black captain of the West Indies cricket team. He is the only batsman to be in two 500-run partnerships in first-class cricket. 500-run partnerships Since the 1960-61 series, the Frank Worrell Trophy is awarded to the winner of the Test series between Australia and West Indies.Report announcing the Frank Worrell Trophy, Glasgow Herald, 10 February 1961One of the two Halls of Residence at the University of the West Indies, Cave Hill Campus in Barbados is named after him. Caribbean journalist Ernest Eytle wrote a biography of Worrell titled Frank Worrell: The Career of a Great Cricketer in 1963. In June 1988 Worrell was celebrated on the $2 Barbadian stamp along with the Barbados Cricket Buckle. ', 0.14286197073500184, 0.8571380292649982)
(15906, 'Title: Easton McMorris Content: Easton Dudley Ashton St John McMorris (4 April 1935 – 1 February 2022) was a West Indian cricketer. He played in 13 Tests from 1958 to 1966.Wisden 1963, pp. 919–20. He toured England with the West Indies team in 1963 and 1966. He was a heavy scorer for Jamaica. He captained Jamaica from 1967-68 until he retired after the 1971-72 season. McMorris was born in Saint Andrew, Jamaica. McMorris died on 1 February 2022, at the age of 86. ', 0.1478250809075755, 0.8521749190924245)
(24439, "Title: Shane Warne Content: Shane Keith Warne (13 September 1969 – 4 March 2022) was an Australian cricketer. He captained the Australian national team in One Day Internationals (ODI). He is often called the greatest bowlers in the history of the sport. He took 708 wickets from 40,704 balls bowled during his Test career. He has the second most wickets in Test history and most of any leg spin bowler. He retired from international cricket in January 2007, at the end of Australia's 5–0 Ashes series whitewash victory over England. After retirement from playing, he was a cricket commentator. Warne was born in Upper Ferntree Gully, an outer suburb of Melbourne, Victoria. He had three children with his ex-wife Simone Callahan. He had complete heterochromia, from which one of his eyes is blue, and the other green. Warne died from a possible heart attack on 4 March 2022 at his villa on the island of Ko Samui in Thailand. He was 52. ", 0.15247275850069164, 0.8475272414993084)
(143128, 'Title: Warren Wisneski Content: Warren Anthony Wisneski was born 19 February, 1969 in New Plymouth, Taranaki. He was a New Zealand cricketer who played three one-day internationals in 2000 but no tests. He also played for Central Districts and Canterbury in the Plunket Shield before retiring after the 2003/04 season. ', 0.15340871346813634, 0.8465912865318637)
(74129, 'Title: Bevan Congdon Content: Bevan Ernest Congdon OBE (11 February 1938 – 10 February 2018) was a New Zealand cricket all-rounder. He played 61 Test matches and 11 One Day Internationals from 1965 to 1978, which included a spell as captain. He was born in Motueka, New Zealand. He was captain of the New Zealand Test and ODI team from 1972 to 1974, and was the first New Zealand captain to record a victory over Australia. Congdon died in Auckland on 10 February 2018, one day before his 80th birthday. ', 0.15421141232634705, 0.845788587673653)
(67096, 'Title: Warren Lees Content: Warren Kenneth Lees MBE was born 19 March, 1952 in Dunedin. He is a former New Zealand cricketer and coach. He played 21 tests and 31 one-day internationals from 1976 to 1983 and played for Otago in the Plunket Shield. Category:1952 births Category:Living people Category:New Zealand cricketers Category:Sportspeople from Dunedin', 0.15673288684184894, 0.8432671131581511)
(79128, 'Title: Peter Walker (cricketer, born 1936) Content: Peter Michael Walker (17 February 1936 – 5 April 2020) was an English cricketer. He played for England in three Test matches against South Africa in 1960. He also played for Glamorgan. He was born in Clifton, Bristol. Walker died of a stroke on 5 April 2020. He was 84. ', 0.15739049143063144, 0.8426095085693686)
(70404, 'Title: Eddie Stuart Content: Edward Albert "Eddie" Stuart (12 May 1931 – 4 November 2014) was a South African professional footballer. He played as a defender. He played 322 times in England for Wolverhampton Wanderers between 1951 and 1962, winning three league titles and the FA Cup. He also served Stoke City, Tranmere Rovers and Stockport County. Stuart died on 4 November 2014 in Wrexham, Wales aged 83 following a long illness. ', 0.15891878575963125, 0.8410812142403687)
(159174, 'Title: Bob Willis Content: right|200px Robert George Dylan Willis (born Robert George Willis; 30 May 1949 – 4 December 2019) was an English cricketer. He played for Surrey, Warwickshire, Northern Transvaal and England. He was a Wisden Cricketer of the Year for 1978. He was born in Sunderland, Tyne and Wear, England. Willis died on 4 December 2019 at the age of 70. ', 0.1589504110844192, 0.8410495889155808)
BigQuery
BigQueryでもベクトル検索はサポートされています。
近似近傍探索アルゴリズムをサポートしておりインデックスとしてはIVFインデックスとTreeAHをサポートしています。
TreeAHは木構造で段階的にハッシュ空間を分割し、探索候補を急速に絞り込むアルゴリズムで大規模なクエリバッチの場合に推奨されます。
BigQueryをすでに使用している場合、BigQueryのベクトル検索を行うのは有力な候補となるでしょう。
ただし、基本的にBigQueryは検索量に応じて課金が行われるため、検索量が予算範囲内におさまるかが採用の是非の基準となるでしょう。
では実際にBigQueryでのベクトル検索を行います。
まず、bigqueryに接続します。
import os
from google.cloud import bigquery
# BigQueryに接続するためのファイルパスを指定
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = './mytestproject_key.json'
# BigQueryの環境に合わせて設定する
dataset_id = 'mydataset'
project_id = 'あなたのプロジェクトID'
client = bigquery.Client(project=project_id)
次にARRAY<FLOAT64>の列を含むテーブルを作成します。
create_table_query = f"""
-- 例: project.dataset を適宜置換
CREATE TABLE IF NOT EXISTS `{project_id}.{dataset_id}.documents` (
id INT64,
body STRING,
embedding ARRAY<FLOAT64>
);
"""
query_job = client.query(create_table_query)
# ジョブの完了を待機
query_job.result()
load_table_from_dataframeを使用してデータを投入します。
import pandas as pd
from tqdm import tqdm
batch_size = 10_000
total = ds.num_rows
table_id = f"{project_id}.{dataset_id}.documents"
for start in tqdm(range(0, total, batch_size)):
end = min(start + batch_size, total)
docs = ds.select(range(start, end))
# BigQuery推奨に寄せて FLOAT64 配列へ（float32でも持てるが、最終的にfloat64に上げるのが無難）
emb64 = np.asarray(docs["text-embedding-ada-002"], dtype=np.float64)
df = pd.DataFrame({
"id": list(range(start, end)),
"body": docs["body"],
"embedding": [v.tolist() for v in emb64], # ARRAY<FLOAT64>
})
job = client.load_table_from_dataframe(df, table_id)
job.result()
データ投入後にインデックスを作成する場合は以下を実行してください。
# データを入れてからじゃないとエラーとなる
# Failed to calculate array_min_len since the entries in corresponding column 'embedding' are all NULLs
# https://docs.cloud.google.com/bigquery/docs/vector-index#understanding_when_data_is_indexed
create_index_query = f"""
CREATE OR REPLACE VECTOR INDEX documents_embedding_idx
ON `{project_id}.{dataset_id}.documents`(embedding)
STORING(id, body)
OPTIONS(distance_type='COSINE', index_type='IVF');
"""
query_job = client.query(create_index_query)
# ジョブの完了を待機
query_job.result()
検索については以下のようなコードになります。
q_float_list = query_emb.astype(np.float64).tolist()[0] # ARRAY<FLOAT64>
k = 10
sql = f"""
SELECT
base.id,
base.body,
distance,
1.0 - distance AS similarity
FROM VECTOR_SEARCH(
TABLE `{project_id}.{dataset_id}.documents`,
'embedding',
(SELECT @q AS embedding),
query_column_to_search => 'embedding',
top_k => @k,
distance_type => 'COSINE'
)
ORDER BY distance
"""
job_config = bigquery.QueryJobConfig(
query_parameters=[
bigquery.ArrayQueryParameter("q", "FLOAT64", q_float_list), # len=1536 の list[float]
bigquery.ScalarQueryParameter("k", "INT64", k),
]
)
rows = client.query(sql, job_config=job_config).result()
for r in rows:
print(r["id"], r["distance"], r["similarity"], r["body"][:120].replace("\n", " "), "...")
出力
123 2.220446049250313e-16 0.9999999999999998 Title: Everton Weekes Content: thumb|200px|Weekes in the 1940s Sir Everton DeCourcy Weekes, KCMG, GCM, OBE (26 February ...
99980 0.1428619778369229 0.8571380221630771 Title: Frank Worrell Content: Sir Frank Mortimer Maglinne Worrell (1 August 1924, Bank Hall, St Michael, Barbados – 13 M ...
6689 0.14698436851924224 0.8530156314807578 Title: Tom Graveney Content: thumb|right|200px|Graveney in October 1954 Thomas William "Tom" Graveney OBE (16 June 1927 ...
14905 0.14782507925730715 0.8521749207426929 Title: Easton McMorris Content: Easton Dudley Ashton St John McMorris (4 April 1935 – 1 February 2022) was a West Indian ...
208851 0.15009166333064838 0.8499083366693516 Title: West Indies cricket team Content: The West Indian cricket team, also known as the West Indies or the Windies, is ...
23438 0.1524727466263729 0.8475272533736271 Title: Shane Warne Content: Shane Keith Warne (13 September 1969 – 4 March 2022) was an Australian cricketer. He captain ...
142127 0.15340872122224558 0.8465912787777544 Title: Warren Wisneski Content: Warren Anthony Wisneski was born 19 February, 1969 in New Plymouth, Taranaki. He was a N ...
73128 0.15421143002805848 0.8457885699719415 Title: Bevan Congdon Content: Bevan Ernest Congdon OBE (11 February 1938 – 10 February 2018) was a New Zealand cricket a ...
120439 0.15549208508716972 0.8445079149128303 Title: Tony Cozier Content: Winston Anthony Lloyd "Tony" Cozier (10 July 1940 – 11 May 2016) was a Barbadian cricket wri ...
66095 0.15673289138508706 0.8432671086149129 Title: Warren Lees Content: Warren Kenneth Lees MBE was born 19 March, 1952 in Dunedin. He is a former New Zealand crick ...
まとめ
今回はFaissでのベクトル検索について確認しました。
この際、ベクトル検索をする場合は全数チェックする方法と、近似近傍探索アルゴリズムにより検索する範囲を絞り込む方法があることを確認しました。
インデックスをどのように永続化して検索するかを確認しました。
Faissの機能だけでも永続化はできますが、DuckDBなどを使用した方がシンプルにできることが確認できます。
続いて、Vector情報を独立して管理するデータベースの動作を確認しました。
今回触れた範囲では、多くの製品は近似近傍探索アルゴリズムをサポートしているのでデータ量が増えた場合でも対応は可能であると考えられます。
他にもベクトル検索を行うサービスは存在するので今回、実験した内容を元に検証してみるのをお勧めします。
参考
-
Faiss: The Missing Manual
- Faissの解説としてはわかりやすいマニュアル。動画もついている。
-
Guidelines to choose an index
- Faissでどのインデックスを選択すべきかの参考
-
Indexes that do not fit in RAM
- FaissでRAMに収まらないインデックスをどう取り扱うか？
- THE FAISS LIBRARY
- sqlite-vec
- DuckDB Docs Vector Similarity Search Extension
- pgvector
- MariaDB Vector Overview
- BigQuery - ベクトル検索を使用してエンベディングを検索する
- BigQuery - ベクトル インデックスを管理する
Discussion