PostgreSQLでベクトル検索を実装する方法

pgvectorはPostgreSQLにvector型のカラムを追加し、類似度クエリは距離演算子によるORDER BYとLIMITの組み合わせで表現されます。

この話題が持ち上がるのは、たいてい「cancel my subscription」でヘルプセンター検索をしても記事のタイトルが「Ending your plan」であるために何もヒットしなかった、といったケースです。そして誰かが、すでに運用中のPostgreSQLの隣にマネージドのベクトルデータベースを立ち上げようと提案します。しかし、その2つ目のサービスが必要になることはめったにありません。本記事の残りの部分では、その道筋すべてをSQLで示します。拡張機能の有効化、埋め込みモデルに合わせたカラムのサイズ設定、ベクトルの保存、コサイン距離によるクエリ、HNSWまたはIVFFlatによるインデックス作成、ベクトル検索結果と通常テーブルとの結合、そしてベクトル検索が適さずPostgreSQLの全文検索に任せるべきクエリについて解説します。

ベクトル検索は何に使うのか

ベクトル検索は意味に基づいてマッチングするため、「cancel my subscription」というクエリで「Ending your plan」というタイトルの文書を返すことができます。両者に共通する単語がなくてもです。各テキストは埋め込みモデルによって固定長の数値のリストに変換され、意味の近いテキストはその空間内で近い位置に配置されます。検索とは、クエリのベクトルに最も近い、保存済みのベクトルを見つけることです。埋め込みの仕組みについての背景は、Vector Databases Explainedを参照してください。

これが効果を発揮するのは主に3つの場面です。言い換えに強いサイト内検索、サポートチケットのマッチング(これに似た過去のチケットを探す)、そしてRAGのための検索です。RAGでは最も近い文書が言語モデルにコンテキストとして渡されます。詳しくはintroduction to RAG for web appsを参照してください。

pgvectorを有効化してベクトルカラムを追加する

pgvectorはデータベースごとに一度だけCREATE EXTENSION vectorで有効化し、カラム型はvector(n)です。ここでnは次元数を表します。バージョン0.8.6はPostgreSQL 13以降をサポートしていますが、PostgreSQL 13はすでにコミュニティサポートが終了しているため、実質的な下限は14以降となります。

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
  id        bigserial PRIMARY KEY,
  title     text NOT NULL,
  body      text NOT NULL,
  embedding vector(1536)   -- replace 1536 with your embedding model's output size
);

-- or, on a table you already have:
ALTER TABLE documents ADD COLUMN embedding vector(1536);

vector(n)の中の数値は、使用する埋め込みモデルが生成するベクトルの長さと一致していなければなりません。たとえばOpenAIのtext-embedding-3-smallはデフォルトで1536次元のベクトルを返します。異なるモデルのベクトル同士を比較しても意味はないため、モデルを切り替える場合はカラム全体を再度埋め込み直す必要があります。

任意のモデルの埋め込みを保存する

埋め込みは他のカラム値と同じように書き込みます。アプリケーションコードでベクトルを生成し、vectorにキャストしたバインドパラメータとして渡すだけです。pgvector自体がモデルを呼び出すことはなく、PostgreSQLドライバのある言語であれば何でも使えます。

INSERT INTO documents (title, body, embedding)
VALUES ($1, $2, $3::vector);

-- re-embed an existing row without creating a duplicate
INSERT INTO documents (id, title, body, embedding)
VALUES ($1, $2, $3, $4::vector)
ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding;

初回のバックフィルについて、pgvectorはCOPY ... FROM STDIN WITH (FORMAT BINARY)による一括ロードを推奨しており、インデックスはデータ投入前ではなく投入後に構築することを勧めています。

どのpgvector距離演算子を使うべきか

モデルのドキュメントに別段の記載がない限り、テキスト埋め込みには<=>を使ってください。この演算子はコサイン距離を返すため、昇順でソートすると最も近い一致が先頭に来ます。表示用にコサイン類似度が必要な場合は、結果を1から引いてください。

-- five closest documents; $1 is the query embedding
SELECT id, title, 1 - (embedding <=> $1::vector) AS similarity
FROM documents
ORDER BY embedding <=> $1::vector
LIMIT 5;

pgvectorは合計6つの距離演算子をサポートしており、それぞれに対応するインデックス演算子クラスが必要です。

<#>は内積の符号を反転した値を返します。PostgreSQLはインデックスを昇順にしかスキャンしないため、負の形にすることで最小の値が最も近い一致となります。素の内積を得るには-1を掛けてください。ベクトルがすでに長さ1に正規化されている場合、内積が厳密検索の最速の選択肢です。<~>(ハミング)と<%>(ジャッカード)はバイナリベクトル用で、本記事の範囲外です。

HNSWとIVFFlatのどちらでインデックスを作るべきか

インデックスがない場合、pgvectorはクエリベクトルをすべての行と比較し、厳密な最近傍を返します。HNSWまたはIVFFlatインデックスを追加すると、これが近似検索に切り替わります。近似検索ははるかに高速で、真の最近傍の大半を見つけますが、厳密クエリとは異なる行を返す可能性があります。どちらのインデックス型も近似です。厳密検索とは、単にカラムにベクトルインデックスがまったく存在しない場合に起こることです。

構築時間やメモリの制約でIVFFlatを選ばざるを得ない場合を除き、デフォルトはHNSWにしてください。演算子クラスはクエリで使う演算子と一致させる必要があります。

-- production: avoid blocking writes
CREATE INDEX CONCURRENTLY documents_embedding_hnsw
  ON documents USING hnsw (embedding vector_cosine_ops);

-- IVFFlat alternative; create only after the table has data
CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

HNSWにはトレーニング工程がないため、テーブルに1行も入っていない状態でもインデックスを構築できます。IVFFlatにはトレーニング工程があります。そのリストはインデックス構築時に存在するデータから作られるため、先に代表性のある行のセットを用意しておく必要があります。トランザクション内でSET LOCALを使ってef_searchやprobesを引き上げると、速度と引き換えにそのクエリの再現率が向上します。

プランナがベクトルインデックスを検討するのは、クエリが距離演算子に対して直接ORDER BYを昇順で指定し、かつLIMITを伴っている場合だけです。ORDER BY 1 - (embedding <=> $1) DESCではインデックスは使われません。小さなテーブルではプランナがシーケンシャルスキャンを選ぶこともあるため、EXPLAINで確認してください。インデックスによって再現率がどれだけ犠牲になったかを測るには、厳密検索を強制して2つの結果セットを比較します。

BEGIN;
SET LOCAL enable_indexscan = off;  -- exact search for comparison
SELECT id FROM documents ORDER BY embedding <=> $1::vector LIMIT 5;
COMMIT;

別途ベクトルデータベースは必要か

すでにPostgreSQLを運用しているなら、おそらく別途ベクトルデータベースは必要ありません。pgvectorを使えば、同じバックアップ、同じトランザクションのもとでベクトル検索が行え、ベクトル検索結果を通常のテーブルと1つのクエリで結合できます。文書とその埋め込みは1回のINSERTで書き込まれるため、2つのストア間の同期パイプラインも、ベクトルが孤立するという障害モードも存在しません。ベクトルは他のカラムと同様にWAL(先行書き込みログ)を通るので、レプリカやポイントインタイムリカバリでも追加の作業なしに反映されます。

これが具体的になるのが結合です。エンタープライズ顧客に限定して最も近いサポートチケットを探すのは、単一のステートメントで済みます。

SELECT t.id, t.subject, c.plan, t.embedding <=> $1::vector AS distance
FROM tickets t
JOIN customers c ON c.id = t.customer_id
WHERE c.plan = 'enterprise'
ORDER BY t.embedding <=> $1::vector
LIMIT 10;

近似インデックスには1つ落とし穴があります。先にインデックスがスキャンされ、その結果に対してWHERE句が適用されるため、選択性の高いフィルタではLIMITで要求した数より少ない行しか残らないことがあります。フィルタ対象カラムにB-treeを張れば高速な厳密結果が得られることが多く、残りは反復スキャン(iterative scan)や部分インデックスで対応できます。

専用のベクトルデータベースは、数十億規模のベクトルや極端な書き込みレートでは依然として価値があります。それ未満では、追加のサービスはユーザーから見えるメリットのない運用コストにしかなりません。

ベクトル検索の限界:全文検索とのハイブリッド検索

ベクトル検索は同じ意味を持つ行を見つけ、全文検索は同じ単語を含む行を見つけます。本番環境の検索では通常その両方を実行し、2つのランク付きリストをマージします。厳密な識別子、製品コード、エラー文字列、人名などは、埋め込みモデルにとって有用な「意味」を持ちません。SKU-4471というクエリは、そのトークンを含む行を返すべきであり、類似製品に関する行を返すべきではありません。こうしたクエリにはPostgreSQLの全文検索が適しており、pgvectorのドキュメントでもハイブリッドクエリのためにベクトル検索と組み合わせる方法が紹介されています。

まずは保存される生成カラムとしてのtsvectorと、その上のGINインデックスから始めます。

ALTER TABLE documents
  ADD COLUMN body_tsv tsvector
  GENERATED ALWAYS AS (to_tsvector('english', title || ' ' || body)) STORED;

CREATE INDEX ON documents USING gin (body_tsv);

次に、pgvector自身のRRFサンプルの形に従って、Reciprocal Rank Fusion(RRF)で2つのランキングを統合します。

-- $1 = query embedding, $2 = raw query text
WITH semantic AS (
  SELECT id, RANK() OVER (ORDER BY embedding <=> $1::vector) AS rnk
  FROM documents
  ORDER BY embedding <=> $1::vector
  LIMIT 20
),
lexical AS (
  SELECT id, RANK() OVER (ORDER BY ts_rank_cd(body_tsv, q) DESC) AS rnk
  FROM documents, plainto_tsquery('english', $2) AS q
  WHERE body_tsv @@ q
  ORDER BY ts_rank_cd(body_tsv, q) DESC
  LIMIT 20
)
SELECT COALESCE(s.id, l.id) AS id,
       COALESCE(1.0 / (60 + s.rnk), 0.0) + COALESCE(1.0 / (60 + l.rnk), 0.0) AS score
FROM semantic s
FULL OUTER JOIN lexical l ON l.id = s.id
ORDER BY score DESC
LIMIT 10;

各リストは1 / (k + rank)を寄与するため、どちらかのリストで上位にある文書は高いスコアを得て、両方に存在する文書が最も高いスコアになります。定数k = 60はCormack、Clarke、Büttcherに由来し、彼らはパイロット研究でこの値を固定し、正確な値は重要ではないと報告しています。このクエリは例示として扱い、自分のデータでHNSWとGINの両方のインデックスが使われていることをEXPLAIN (ANALYZE, BUFFERS)で確認してください。

どこから始めるか

実装全体は、カラム1つ、演算子1つ、インデックス1つに収まり、すべてがすでにバックアップとレプリケーションを行っているデータベースの中で完結します。モデルに合わせてサイズを決めたvector(n)カラムを追加し、まずはインデックスなしでORDER BY ... <=> ... LIMITのクエリを書き、その後HNSWを追加してenable_indexscan = offで再現率の差を確認してください。セマンティックな結果が妥当に見えるようになったら、全文検索側を組み込みましょう。そうしないと、最初に注文番号を貼り付けたユーザーがそのギャップを見つけることになります。

FAQ

pgvectorは2,000次元を超える埋め込みをインデックス化できますか?

標準のvector型だけでは不可能です。HNSWとIVFFlatは最大2,000次元までのvectorカラムをインデックス化できますが、カラム自体は最大16,000次元まで格納できます。それより大きな埋め込みには、式インデックス内でhalfvecにキャストする(4,000次元までインデックス可能)、再ランキングを伴うバイナリ量子化を使う(最大64,000次元)、サブベクトルをインデックス化する、あるいはモデルに出力次元数を減らすよう要求する、といった方法があります。最後の方法はOpenAIのtext-embedding-3系モデルがdimensionsパラメータでサポートしています。

新しい行を挿入した後、pgvectorのインデックスを再構築する必要がありますか?

HNSWでは不要です。新しい行は挿入されるたびにグラフに追加されます。だからこそ空のテーブルでもインデックスを構築できるのです。IVFFlatは挙動が異なります。そのリスト重心は構築時にk-meansで一度だけ計算され、その後移動しないため、テーブルが成長したり分布が変化したりすると再現率が低下することがあります。大量ロードや分布の変化の後は、REINDEX INDEX CONCURRENTLYでIVFFlatインデックスを再構築してください。

HNSWインデックスを追加したら、クエリがLIMITより少ない行しか返さないのはなぜですか?

HNSWスキャンが返す候補は最大でhnsw.ef_search件(デフォルト40)であり、WHEREフィルタはその候補に対して後から適用されるためです。40を超えるLIMIT、選択性の高いフィルタ、デッドタプルのいずれもが結果を不足させる原因になります。SET LOCALでhnsw.ef_searchを引き上げるか、pgvector 0.8.0以降であればhnsw.iterative_scanをstrict_orderに設定して反復インデックススキャンを有効にし、十分な行がマッチするまでスキャンを継続させてください。

2つの異なるモデルの埋め込みを同じテーブルに保存できますか?

可能ですが、共有された1つのインデックス付きカラムに保存することはできません。モデルごとに別々のvector(n)カラムを用意し、それぞれをそのモデルの出力サイズに合わせ、それぞれに独自のインデックスを張ってください。pgvectorは次元が混在する型指定なしのvectorカラムも許容しますが、インデックスは単一の次元の行しかカバーできず、キャストを伴う式インデックスとモデル識別子に対する部分WHEREを組み合わせて構築する必要があります。異なるモデル間の距離に意味はありません。