Posted inデータベース

PostgreSQL pgvectorチュートリアル：Embeddingの保存と高速セマンティック検索の構築

Database tutorial - IT technology blog
Database tutorial - IT technology blog

先月の午前2時15分、オンコールのPagerDutyアラートが鳴り響きました。急激な4,000 RPMのトラフィックスパイクにより、スタンドアロンのベクトルデータベースクラスタが504 Gateway Timeoutを返し始めたのです。

さらに悪いことに、非同期CDC同期パイプラインに45分の遅延が発生し、何千人ものアクティブユーザーに古い破損した検索結果を提供してしまっていました。私たちは次のスプリントで、この脆弱な2つのデータベース構成を撤廃しました。pgvector拡張機能を使ってベクトル操作を直接PostgreSQLに移行したことで、データ同期の遅延を完全に解消することができました。

クイックスタート（5分）：Embeddingの保存とクエリ実行

ベクトルの類似度クエリを実行するには、拡張機能を有効化し、カラムにvectorデータ型を割り当てます。OpenAIのtext-embedding-3-smallの1536次元や、all-MiniLM-L6-v2のような軽量なHugging Faceモデルの384次元など、使用するEmbeddingモデルが出力する正確な次元数を指定してください。

1. 拡張機能の有効化とテーブルの作成

-- pgvector拡張機能を有効化
CREATE EXTENSION IF NOT EXISTS vector;

-- 1536次元のEmbeddingを持つドキュメント用テーブルを作成
CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    category VARCHAR(50),
    embedding vector(1536),
    created_at TIMESTAMPTZ DEFAULT NOW()
);

2. ベクトルデータの挿入

-- 対応するEmbeddingベクトルを持つレコードを挿入
INSERT INTO documents (title, content, category, embedding)
VALUES (
    'Postgresのコネクションプーリング',
    'PgBouncerを使用した接続飽和管理の徹底解説。',
    'database',
    '[-0.0023, 0.0142, -0.0311, ...]' -- 1536個の浮動小数点数値
);

3. 最近傍探索の実行

-- コサイン類似度を用いてセマンティックに関連性の高い上位5件のドキュメントを取得
SELECT 
    id,
    title,
    1 - (embedding <=> '[-0.0023, 0.0142, -0.0311, ...]') AS similarity_score
FROM documents
ORDER BY embedding <=> '[-0.0023, 0.0142, -0.0311, ...]'
LIMIT 5;

詳細解説：ベクトル演算とインデックスの選定

リレーショナルデータ用のPostgreSQLとベクトル用のPinecone、Qdrant、Milvusといった2つのデータベースを運用する構成は、運用上の摩擦を生み出します。インフラコストが倍増し、トランザクションの一貫性が損なわれ、ネットワークホップが増加します。pgvectorを使用すれば、ACID特性を維持したまま、Embeddingを顧客テーブルと直接結合し、厳格な外部キー制約を適用できます。

距離演算子の理解

Embeddingモデルの距離計算方法に基づいて演算子を選択します：

  • <=> コサイン距離（Cosine Distance）：ベクトルの大きさを考慮せず、角度の距離を計算します。OpenAI、Cohere、および最新のLLMテキストEmbeddingのほとんどにおいて最適な選択肢です。
  • <-> ユークリッド距離 / L2距離（Euclidean / L2 Distance）：幾何学的な直線距離を測定します。コンピュータビジョンや音声分類モデルで一般的です。
  • <#> 負の内積（Negative Inner Product）：負の内積を計算します。内積がコサイン類似度と等しくなる正規化済みの単位ベクトルに使用することで、計算コストを削減できます。

適切なインデックスの選択：HNSW vs. IVFFlat

厳密な最近傍探索（シーケンシャルスキャン）は1万行未満であれば問題なく動作します。しかしその閾値を超えると、クエリのレイテンシは15msから数秒へと跳ね上がります。本番環境のワークロードでは、近似最近傍探索（ANN）インデックスが不可欠です。

-- 1. HNSWインデックス（Hierarchical Navigable Small World）- 本番環境に推奨
CREATE INDEX idx_documents_embedding_hnsw ON documents 
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- 2. IVFFlatインデックス（Inverted File Flat）- 低メモリ消費、作成前にデータが必要
CREATE INDEX idx_documents_embedding_ivfflat ON documents 
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);
インデックスの種類 構築速度 クエリレイテンシ（100万件） メモリ消費量 再現率（Recall Rate）
HNSW 低速（約15〜30分 / 100万ベクトル） < 8ms 約1.5 GB / 100万ベクトル 97% – 99%
IVFFlat 高速（約2〜5分 / 100万ベクトル） 約35〜70ms 約300 MB / 100万ベクトル 85% – 92%

ユーザー向けアプリケーションでは、基本的にHNSWを採用してください。IVFFlatとは異なり、HNSWは空のテーブルでも確実に構築でき、継続的な書き込みがあっても再現率が低下せず、常に10ms未満の応答速度を実現します。

高度な活用法：Pythonパイプラインとハイブリッド検索

純粋なベクトル検索は、請求書番号、製品SKU、専門的な型番といった完全一致キーワードの検索を苦手とすることがよくあります。PostgreSQLの全文検索（BM25スタイルのtsvector）と、Reciprocal Rank Fusion（RRF）を用いた高密度なセマンティック検索を組み合わせることで、両方の利点を活かした最高の検索品質を実現できます。

本番環境向けハイブリッド検索クエリ（Reciprocal Rank Fusion）

WITH semantic_search AS (
    SELECT id, RANK() OVER (ORDER BY embedding <=> $1) AS rank
    FROM documents
    ORDER BY embedding <=> $1
    LIMIT 20
),
keyword_search AS (
    SELECT id, RANK() OVER (ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC) AS rank
    FROM documents, plainto_tsquery('english', $2) query
    WHERE to_tsvector('english', content) @@ query
    LIMIT 20
)
SELECT 
    d.id, 
    d.title,
    COALESCE(1.0 / (60 + s.rank), 0.0) + 
    COALESCE(1.0 / (60 + k.rank), 0.0) AS rrf_score
FROM documents d
LEFT JOIN semantic_search s ON d.id = s.id
LEFT JOIN keyword_search k ON d.id = k.id
WHERE s.id IS NOT NULL OR k.id IS NOT NULL
ORDER BY rrf_score DESC
LIMIT 10;

Psycopg 3を使用したPython連携

import psycopg
from pgvector.psycopg import register_vector
from openai import OpenAI

client = OpenAI(api_key="your-api-key")

def get_embedding(text: str) -> list[float]:
    response = client.embeddings.create(
        input=text,
        model="text-embedding-3-small"
    )
    return response.data[0].embedding

# PostgreSQLに接続し、1536次元のEmbeddingをクエリ
with psycopg.connect("postgresql://postgres:secret@localhost:5432/app_db") as conn:
    register_vector(conn)
    
    query_text = "PgBouncerの接続制限をチューニングするにはどうすればよいですか？"
    query_vector = get_embedding(query_text)
    
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, title, content
            FROM documents
            ORDER BY embedding <=> %s::vector
            LIMIT 3;
        """, (query_vector,))
        
        for doc_id, title, snippet in cur.fetchall():
            print(f"[{doc_id}] {title}: {snippet[:80]}...")

実践的なTips：本番環境のトラフィックを乗り切るために

1. インデックス作成時にmaintenance_work_memを割り当てる

HNSWインデックスの構築では、グラフ作成時に大量のメモリを消費します。デフォルトのmaintenance_work_mem（通常64MB）では、Postgresが一時バッファをディスクにスワップしてしまい、10分で終わるはずの構築が4時間も停滞する原因になります。マイグレーションセッション用に値を引き上げてください：

SET maintenance_work_mem = '4GB';
CREATE INDEX CONCURRENTLY idx_documents_embedding_hnsw 
ON documents USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

2. マルチテナント環境では部分インデックスを活用する

一般的なWHERE tenant_id = 42のようなフィルタでグローバルなベクトルインデックスをスキャンしないでください。PostgresはまずグローバルなHNSWグラフ全体をトラバースし、その後に一致しないレコードを除外するため、小規模テナントのクエリパフォーマンスが著しく低下します。

代わりに、テナント固有またはフィルタリングされた部分インデックスを作成します：

-- エンタープライズテナント専用のインデックスパーティション
CREATE INDEX idx_tenant_42_vectors ON documents 
USING hnsw (embedding vector_cosine_ops)
WHERE tenant_id = 42;

3. pg_prewarmでベクトルインデックスを共有バッファに常駐させる

HNSWグラフのトラバースには、ランダムなメモリページ参照が必要です。コールドページがRAMからNVMeストレージに退避されると、クエリレイテンシが6msから250msに跳ね上がる可能性があります。データベースの再起動やデプロイ直後に、アクティブなインデックスをキャッシュに事前ロードしてください：

CREATE EXTENSION IF NOT EXISTS pg_prewarm;
SELECT pg_prewarm('idx_documents_embedding_hnsw', 'buffer', 'main');
Share:
Tags: