Optimizing Retrieval-Augmented Generation (RAG) with Hierarchical Navigable Small World (HNSW) Index Quantization in pgvector — code snippets

snippet-1.sql

-- Step 1: Add the half-precision vector column allowing NULLs initially
ALTER TABLE documents ADD COLUMN embedding_half halfvec(1536);

-- Step 2: Populate the new column in batches to prevent transaction log bloat
-- (For a live database, execute this in chunks of 50,000 rows using a background job)
UPDATE documents 
SET embedding_half = embedding::halfvec(1536) 
WHERE embedding_half IS NULL;

-- Step 3: Enforce the NOT NULL constraint after verification
ALTER TABLE documents ALTER COLUMN embedding_half SET NOT NULL;

-- Step 4: Build the HNSW index using halfvec_cosine_ops
-- We use CONCURRENTLY to avoid blocking read/write queries during index construction
CREATE INDEX CONCURRENTLY idx_documents_hnsw_half 
ON documents 
USING hnsw (embedding_half halfvec_cosine_ops)
WITH (m = 16, ef_construction = 64);

snippet-2.sql

-- Create an HNSW index on the binary-quantized representation of your float32 vector.
-- The cast to bit(1536) is mandatory so the database can allocate the correct bit-length.
CREATE INDEX CONCURRENTLY idx_documents_bq_hnsw 
ON documents
USING hnsw ((binary_quantize(embedding)::bit(1536)) bit_hamming_ops)
WITH (m = 16, ef_construction = 128);

snippet-3.sql

-- Query the binary expression index using Hamming distance (<~>)
-- The query vector must be explicitly cast to vector, then binary-quantized and cast to bit
SELECT id, title
FROM documents
ORDER BY (binary_quantize(embedding)::bit(1536)) <~> (binary_quantize(:query_vector::vector(1536))::bit(1536))
LIMIT 10;

snippet-4.sql

-- Single-query execution of two-stage retrieval using a Common Table Expression (CTE)
WITH candidates AS (
    SELECT 
        id, 
        embedding,
        -- Force the planner to use the BQ HNSW index by matching the index expression
        (binary_quantize(embedding)::bit(1536)) <~> (binary_quantize(:query_vector::vector(1536))::bit(1536)) AS hamming_dist
    FROM documents
    ORDER BY (binary_quantize(embedding)::bit(1536)) <~> (binary_quantize(:query_vector::vector(1536))::bit(1536))
    -- Fetch 100 candidates to ensure high recall
    LIMIT 100
)
SELECT 
    id, 
    hamming_dist,
    -- Compute exact cosine distance (operator <=>) only on the 100 candidates
    (embedding <=> :query_vector::vector(1536)) AS exact_cosine_dist
FROM candidates
ORDER BY exact_cosine_dist ASC
-- Return the final top 10 items to the LLM context window
LIMIT 10;

snippet-5.sql

-- Allocate temporary memory for the index build (recommend 25% of system RAM, up to 16GB-32GB)
SET maintenance_work_mem = '16GB';

-- Maximize parallel worker utilization (match physical CPU cores allocated to Postgres)
SET max_parallel_maintenance_workers = 8;

-- Build the index with slightly relaxed HNSW properties for speed
CREATE INDEX CONCURRENTLY idx_documents_hnsw_optimized 
ON documents 
USING hnsw (embedding_half halfvec_cosine_ops)
WITH (m = 24, ef_construction = 100);

-- Reset configuration variables back to system defaults for this session
RESET maintenance_work_mem;
RESET max_parallel_maintenance_workers;

snippet-6.sql

-- Force update of table statistics and histograms for the index expression
ANALYZE documents;

-- Adjust the cost model for NVMe storage (default random_page_cost is 4.0, which assumes slow HDDs)
-- Setting this to 1.1 or 1.0 tells the planner that random access is almost as fast as sequential access
SET random_page_cost = 1.1;

-- Run EXPLAIN ANALYZE to verify index execution paths
EXPLAIN ANALYZE
SELECT id
FROM documents
ORDER BY (binary_quantize(embedding)::bit(1536)) <~> (binary_quantize('[-0.01, 0.05, 0.12]::vector(3)'::vector(1536))::bit(1536))
LIMIT 10;

snippet-7.py

import logging
import psycopg
from psycopg.rows import dict_row

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("rag_vector_search")

def search_documents_two_stage(
    conn: psycopg.Connection, 
    query_vector: list[float], 
    limit: int = 10, 
    candidate_multiplier: int = 10
) -> list[dict]:
    """
    Executes a high-performance two-stage vector search.
    
    1. First stage uses Binary Quantization (BQ) HNSW index with Hamming distance.
    2. Second stage performs exact cosine distance re-ranking on the candidates.
    """
    vector_dim = len(query_vector)
    if vector_dim != 1536:
        raise ValueError(f"Expected 1536-dimensional vector, got {vector_dim}")
        
    # Serialize the float list to pgvector literal format: '[0.012,-0.432,...]'
    query_vector_literal = f"[{','.join(map(str, query_vector))}]"
    
    # Calculate candidate pool size (e.g., limit 10 * multiplier 10 = 100 candidates)
    candidate_limit = limit * candidate_multiplier
    
    # Single query combining CTE candidate retrieval and outer re-ranking
    query = """
        WITH candidates AS (
            SELECT 
                id, 
                title,
                embedding,
                (binary_quantize(embedding)::bit(1536)) <~> (binary_quantize(%s::vector(1536))::bit(1536)) AS hamming_dist
            FROM documents
            ORDER BY (binary_quantize(embedding)::bit(1536)) <~> (binary_quantize(%s::vector(1536))::bit(1536))
            LIMIT %s
        )
        SELECT 
            id, 
            title,
            hamming_dist,
            (embedding <=> %s::vector(1536)) AS cosine_dist
        FROM candidates
        ORDER BY cosine_dist ASC
        LIMIT %s;
    """
    
    try:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(
                query, 
                (
                    query_vector_literal, 
                    query_vector_literal, 
                    candidate_limit, 
                    query_vector_literal, 
                    limit
                )
            )
            return cur.fetchall()
    except psycopg.Error as e:
        logger.error("Failed to execute two-stage vector search: %s", e)
        raise