pgvector 向量搜尋：從 Embedding 到 RAG 的完整實戰 | PostgreSQL

# Ubuntu / Debian
sudo apt install postgresql-17-pgvector

# macOS（Homebrew）
brew install pgvector

# 確認安裝
psql -c "SELECT * FROM pg_available_extensions WHERE name = 'vector';"

# docker-compose.yml
services:
  postgres:
    image: pgvector/pgvector:pg17
    environment:
      POSTGRES_PASSWORD: secret
    volumes:
      - pgdata:/var/lib/postgresql/data

-- 在目標資料庫中啟用
CREATE EXTENSION IF NOT EXISTS vector;

-- 確認版本
SELECT extversion FROM pg_extension WHERE extname = 'vector';
-- 結果：0.8.0

-- 建立含向量欄位的資料表
CREATE TABLE documents (
    id          BIGSERIAL PRIMARY KEY,
    content     TEXT NOT NULL,
    metadata    JSONB,
    embedding   vector(1536)  -- OpenAI text-embedding-3-small 的維度
);

-- 插入向量資料
INSERT INTO documents (content, embedding)
VALUES (
    '機器學習是人工智慧的子領域',
    '[0.021, -0.043, 0.187, ..., 0.034]'::vector
);

-- 基本相似性查詢（餘弦距離，取最近 5 筆）
SELECT id, content,
       embedding <=> '[0.019, -0.041, 0.191, ...]'::vector AS distance
FROM documents
ORDER BY distance
LIMIT 5;

-- 使用 16 位元半精度浮點數，儲存空間減少 50%
CREATE TABLE documents_half (
    id        BIGSERIAL PRIMARY KEY,
    content   TEXT,
    embedding halfvec(3072)  -- text-embedding-3-large 的維度
);

-- 查詢語法相同
SELECT content, embedding <=> '[0.019, ...]'::halfvec AS distance
FROM documents_half
ORDER BY distance
LIMIT 10;

-- 只儲存非零元素，適合 BM25/SPLADE 等稀疏表示
CREATE TABLE sparse_docs (
    id        BIGSERIAL PRIMARY KEY,
    content   TEXT,
    embedding sparsevec(100000)
);

-- 稀疏向量表示格式：{index:value, ...}/dimensions
INSERT INTO sparse_docs (content, embedding)
VALUES (
    'PostgreSQL 向量搜尋',
    '{1:0.8, 42:0.3, 1337:0.5}/100000'::sparsevec
);

-- 餘弦距離（最常用於文字 Embedding）
SELECT id, content, embedding <=> query_vec AS cosine_dist
FROM documents ORDER BY cosine_dist LIMIT 5;

-- L2 距離
SELECT id, content, embedding <-> query_vec AS l2_dist
FROM documents ORDER BY l2_dist LIMIT 5;

-- 內積（注意回傳負值，最小值 = 最相似）
SELECT id, content, embedding <#> query_vec AS neg_inner_product
FROM documents ORDER BY neg_inner_product LIMIT 5;

-- 建立 IVFFlat 索引（餘弦距離）
CREATE INDEX ON documents
  USING ivfflat (embedding vector_cosine_ops)
  WITH (lists = 100);

-- 查詢時調整 probes（搜尋更多群集 = 更高 Recall）
SET ivfflat.probes = 10;

-- 建立 HNSW 索引（餘弦距離）
CREATE INDEX ON documents
  USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);

-- 查詢時調整 ef_search（更大 = 更高 Recall、更慢）
SET hnsw.ef_search = 100;

-- 使用 EXPLAIN 確認查詢計畫
EXPLAIN (ANALYZE, BUFFERS)
SELECT id, content
FROM documents
ORDER BY embedding <=> '[0.021, -0.043, ...]'::vector
LIMIT 10;

-- 若看到 "Index Scan using xxx_idx" → 索引有效
-- 若看到 "Seq Scan" → 可能原因：
--   1. 資料量不足（優化器認為全表掃描更快）
--   2. LIMIT 值太大
--   3. 索引 operator class 與查詢運算子不符

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import PGVector
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 設定連線
CONNECTION_STRING = "postgresql+psycopg2://user:password@localhost:5432/mydb"

# 初始化 Embedding 模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 建立或連接向量資料庫
vectorstore = PGVector(
    collection_name="documents",
    connection_string=CONNECTION_STRING,
    embedding_function=embeddings,
)

# 新增文件
docs = [
    Document(page_content="PostgreSQL 是功能強大的關聯式資料庫", metadata={"source": "wiki"}),
    Document(page_content="pgvector 讓 PostgreSQL 支援向量搜尋", metadata={"source": "docs"}),
    Document(page_content="RAG 架構使用向量搜尋增強 LLM 回答", metadata={"source": "blog"}),
]
vectorstore.add_documents(docs)

# 相似性搜尋
results = vectorstore.similarity_search("PostgreSQL 如何支援 AI 應用？", k=3)

# 建立完整 RAG Chain
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
llm = ChatOpenAI(model="gpt-4o-mini")

prompt = ChatPromptTemplate.from_template("""
根據以下上下文回答問題：

上下文：{context}

問題：{question}
""")

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

answer = rag_chain.invoke("pgvector 有什麼優勢？")
print(answer)

import psycopg2
import psycopg2.extras
import openai

def get_embedding(text: str, model: str = "text-embedding-3-small") -> list:
    """呼叫 OpenAI API 取得 Embedding 向量"""
    client = openai.OpenAI()
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

def ingest_document(conn, content: str, metadata: dict = None):
    """將文件及其 Embedding 存入 PostgreSQL"""
    embedding = get_embedding(content)
    with conn.cursor() as cur:
        cur.execute(
            """INSERT INTO documents (content, metadata, embedding)
               VALUES (%s, %s, %s::vector) RETURNING id""",
            (content, psycopg2.extras.Json(metadata or {}), embedding)
        )
        doc_id = cur.fetchone()[0]
    conn.commit()
    return doc_id

def semantic_search(conn, query: str, k: int = 5, threshold: float = 0.8):
    """語意搜尋：找出與查詢最相似的文件"""
    query_embedding = get_embedding(query)
    with conn.cursor() as cur:
        cur.execute(
            """SELECT id, content, metadata,
                      1 - (embedding <=> %s::vector) AS similarity
               FROM documents
               WHERE 1 - (embedding <=> %s::vector) > %s
               ORDER BY embedding <=> %s::vector
               LIMIT %s""",
            (query_embedding, query_embedding, threshold, query_embedding, k)
        )
        return [
            {"id": r[0], "content": r[1], "metadata": r[2], "similarity": r[3]}
            for r in cur.fetchall()
        ]

def rag_answer(conn, question: str) -> str:
    """完整 RAG 流程：搜尋 → 組合提示詞 → LLM 回答"""
    # Step 1: 向量搜尋
    relevant_docs = semantic_search(conn, question, k=4)
    context = "\n\n".join(doc["content"] for doc in relevant_docs)

    # Step 2: 呼叫 LLM
    client = openai.OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": f"根據以下資料回答：\n{context}\n\n問題：{question}"}]
    )
    return response.choices[0].message.content

-- 混合搜尋：向量相似度 + BM25 全文搜尋
WITH vector_results AS (
    SELECT id, content,
           ROW_NUMBER() OVER (
               ORDER BY embedding <=> '[...]'::vector
           ) AS vector_rank
    FROM documents
    ORDER BY embedding <=> '[...]'::vector
    LIMIT 50
),
text_results AS (
    SELECT id, content,
           ROW_NUMBER() OVER (
               ORDER BY ts_rank(
                   to_tsvector('english', content),
                   plainto_tsquery('pgvector')
               ) DESC
           ) AS text_rank
    FROM documents
    WHERE to_tsvector('english', content) @@ plainto_tsquery('pgvector')
    LIMIT 50
)
-- RRF 融合排名（k=60 是常用常數）
SELECT
    COALESCE(v.id, t.id) AS id,
    COALESCE(v.content, t.content) AS content,
    COALESCE(1.0 / (60 + v.vector_rank), 0) +
    COALESCE(1.0 / (60 + t.text_rank), 0) AS rrf_score
FROM vector_results v
FULL OUTER JOIN text_results t ON v.id = t.id
ORDER BY rrf_score DESC
LIMIT 10;

-- postgresql.conf 關鍵設定

-- 共享記憶體（建議為 RAM 的 25%）
shared_buffers = '8GB'

-- 工作記憶體（影響排序、Hash Join）
work_mem = '256MB'

-- 維護記憶體（影響 HNSW 建立、VACUUM）
maintenance_work_mem = '2GB'

-- 平行查詢
max_parallel_workers_per_gather = 4

-- 平行建立 HNSW 索引
SET max_parallel_maintenance_workers = 7;

from psycopg2.extras import execute_values

def batch_insert(conn, records, batch_size=1000):
    """分批次插入向量，避免單次事務過大"""
    for i in range(0, len(records), batch_size):
        batch = records[i:i + batch_size]
        with conn.cursor() as cur:
            execute_values(
                cur,
                "INSERT INTO documents (content, embedding) VALUES %s",
                [(r["content"], r["embedding"]) for r in batch],
                template="(%s, %s::vector)"
            )
        conn.commit()
        print(f"已插入 {min(i + batch_size, len(records))}/{len(records)} 筆")

-- 先大量插入，再建立索引（最快）
DROP INDEX IF EXISTS documents_embedding_idx;
-- 大量 INSERT ...

-- HNSW 百萬級參數建議
CREATE INDEX CONCURRENTLY documents_embedding_idx
    ON documents USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 64);
-- m: 16（平衡效能與記憶體）
-- ef_construction: 64-128（品質，愈大愈慢但 Recall 更高）
-- ef_search: 40-100（視延遲容忍度調整）

-- Scalar Quantization：float32 → int8，節省 75% 索引記憶體
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64, quantization = sq);

-- Binary Quantization：最激進壓縮，轉為 1-bit
-- 適合維度 >= 1024，搭配重新排序使用
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64, quantization = bq);

-- 對百萬級向量使用範圍分割
CREATE TABLE documents (
    id          BIGSERIAL,
    created_at  TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    content     TEXT,
    embedding   vector(1536)
) PARTITION BY RANGE (created_at);

-- 按月建立分割
CREATE TABLE documents_2026_01 PARTITION OF documents
    FOR VALUES FROM ('2026-01-01') TO ('2026-02-01');

-- 在每個分割上建立 HNSW 索引
CREATE INDEX ON documents_2026_01 USING hnsw (embedding vector_cosine_ops);

-- 監控索引大小
SELECT indexname,
       pg_size_pretty(pg_relation_size(indexname::regclass)) AS index_size
FROM pg_indexes
WHERE tablename = 'documents';

-- 查看快取命中率
SELECT
    heap_blks_hit::float / NULLIF(heap_blks_hit + heap_blks_read, 0) AS table_cache,
    idx_blks_hit::float / NULLIF(idx_blks_hit + idx_blks_read, 0) AS index_cache
FROM pg_statio_user_tables
WHERE relname = 'documents';

-- 升級 pgvector Extension
ALTER EXTENSION vector UPDATE;

-- 升級後可能需要重建索引（尤其跨大版本）
REINDEX INDEX CONCURRENTLY documents_embedding_idx;