~/tools / 埋め込みと検索(RAG) / pgvector SQL Builder

pgvector SQL Builder — pgvector のテーブル・索引・検索の SQL を書く

PostgreSQLのpgvectorでembeddingを保存して検索するためのSQLを書きます。vector(n) または halfvec(n) の列を持つテーブル、距離に合った演算子クラスを使ったHNSWまたはIVFFlatのインデックス、読みやすいスコアを返す最近傍検索のクエリが得られます。pgvectorのドキュメントにある制限を確認し、ベクトルが使う容量も見積もります。データがブラウザの外に出ることはありません。

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
  id bigserial PRIMARY KEY,
  content text NOT NULL,
  embedding vector(384) NOT NULL
);

CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);

PREPARE search_documents(vector(384)) AS
  SELECT id, content, 1 - (embedding <=> $1) AS cosine_similarity
  FROM documents
  ORDER BY embedding <=> $1
  LIMIT 5;

-- EXECUTE search_documents('[0.0123, -0.0456, ...]');

pgvectorは、PostgreSQLにベクトルの類似度検索を加えるオープンソースの拡張機能です。Ollamaで作ったembeddingを、ほかのデータと同じデータベースに保存できます。間違えやすい点がいくつかあります。列の次元数は、embeddingモデルの次元数と同じでなければなりません。距離ごとに専用の演算子があり、コサイン距離は <=>、L2距離は <->、内積の符号を反転した値は <#>、L1距離は <+> です。近似インデックスが使われるのは、同じ距離の演算子クラス(<=> なら vector_cosine_ops など)で作られていて、さらにクエリがその演算子そのものを昇順で使って並べ替え、LIMIT を付けているときだけです。このツールは、1組の選択から3つの部分を書きます。CREATE EXTENSION と CREATE TABLE、CREATE INDEX、そして読みやすいスコア付きで最も近い行を返す PREPARE 文です(コサイン類似度は1からコサイン距離を引いた値、内積は <#> に -1 を掛けた値です)。pgvector 0.8.7のREADMEとソースコードにある制限を適用します。1つの列は最大16,000次元、インデックスは vector で2,000次元、halfvec で4,000次元まで、IVFFlatにはL1の演算子クラスがなく、HNSWの m は2から100、ef_construction は4から1,000で、かつ 2 × m 以上です。HNSWでは検索が返すのは最大でも hnsw.ef_search 行(デフォルトは40)なので、LIMIT がそれより大きいと警告します。IVFFlatでは、READMEの出発点である、100万行までは rows / 1000、それより多いときは sqrt(rows)、probes には sqrt(lists) を使って、lists と probes を入れます。生成したSQLは、Ollamaで作った実際の384次元のベクトルを使い、PostgreSQL 16上のpgvector 0.8.7で実行し、EXPLAIN でインデックスが使われていることを確認しました。

このツールの計算部分は、オープンソース(MIT ライセンス)で公開しています。 github.com/seigo5555/toolkitbox-lib/src/pgvector-sql.ts