Skip to main content

Command Palette

Search for a command to run...

pgvector for RAG Search: Schema, Storage, and Retrieval, With Real Results

Updated
6 min readView as Markdown
pgvector for RAG Search: Schema, Storage, and Retrieval, With Real Results
G
Backend engineer building ShiftMailer, an AI-powered email tool. Sharing lessons from shipping AI agents and scalable systems. Follow for real-world AI, product, and engineering insights—no fluff.

RAG search has four moving parts: split the document, turn text into vectors, store those vectors somewhere queryable, and pull the right ones back out at query time. This post covers the last two — configuring Postgres to store vectors and structuring retrieval around them

Introduction

This post assumes you've already decided how to chunk your documents — if not, see the companion post, "Chunking Strategies in RAG: A Deep Dive" which measured fixed-size, recursive, and semantic chunking against real blog content and found recursive chunking the strongest default (100% clean sentence-ending rate vs. 5% for fixed-size, and structure-aware since it respects markdown headers). This post picks up from there: given chunks and a way to embed them, how do you actually get them into Postgres and query them back correctly?

Big Picture

💡 Key Point: embed_chunks() and embed_query() are the same underlying function, called at two different times. This single point of truth matters — if ingestion and query use different embedding logic (different model, different preprocessing, different normalization), the vectors won't be comparable and similarity search silently degrades without throwing any error.

Step 1: A Single Embedding Function, Used Twice

class EmbeddingService:
    def __init__(self, fitting_corpus):
        self._backend, self.dim = self._load_backend(fitting_corpus)

    def embed_chunk(self, text):
        return self._backend.embed_documents([text])[0]

    def embed_chunks(self, texts):
        return self._backend.embed_documents(texts)

    def embed_query(self, text):
        return self._backend.embed_query(text)

Wrapping the embedding backend in one class with embed_chunks() and embed_query() methods means the rest of the pipeline never touches the backend directly — it calls this interface. Swapping BAAI/bge-small-en-v1.5 for text-embedding-3-small later is a one-line change inside _load_backend(), and nothing in the schema, ingestion, or retrieval code needs to know or care.

⚠️ This only works if self.dim (the embedding dimension) is read from the actual backend after it loads, not hardcoded. Different models produce different dimensions — 384 for BGE-small, 1536 for text-embedding-3-small — and the Postgres schema in the next section is built with this exact number.

Step 2: Postgres Configuration for Vector Storage

The only mandatory step is enabling the extension:

CREATE EXTENSION IF NOT EXISTS vector;

📝 Terminology: this needs to run once per database, by a role with sufficient privilege (superuser, or a role granted CREATE on the database). It's a standard extension install — no changes to postgresql.conf, no shared_preload_libraries entry, and no server restart required, unlike some Postgres extensions.

Step 3: Schema Design

CREATE TABLE rag_chunks (
    id SERIAL PRIMARY KEY,
    document_id TEXT NOT NULL,
    document_title TEXT NOT NULL,
    chunk_index INT NOT NULL,
    content TEXT NOT NULL,
    embedding VECTOR(384) NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
    UNIQUE (document_id, chunk_index)
);

Each column exists for a specific reason tied to how retrieval will use it later:

Column Purpose
document_id Traces a chunk back to its source document — needed to cite results or fetch neighboring chunks
chunk_index This chunk's position within its document — enables ordered reconstruction and neighbor lookups
content The raw chunk text — this is what actually gets shown to the user or passed to an LLM, the embedding itself is not human-readable
embedding VECTOR(384) — dimension must match EmbeddingService.dim exactly
created_at Standard audit column — useful once you're re-ingesting updated documents and need to know which rows are stale
UNIQUE (document_id, chunk_index) Makes ingestion idempotent — re-running the pipeline on the same document fails loudly on duplicate insert rather than silently doubling the corpus

Step 4: Ingestion

# -- for each document, chunk it, embed the chunks, insert them
for document in documents: # each document has id and text
    chunks = chunk_text(document.text) # from the chunking
    vectors = embed_chunks(chunks) # single embedding fun

    for chunk_index, (chunk, vector) in enumerate(zip(chunks, vectors)):
        insert_row(
            document_id=document.document_id,
            document_title=document.title,
            chunk_index=chunk_index,
            content=chunk,
            embedding=vector,
        )

    commit()

Measured result, running this against 5 real blog posts (gauravbytes.dev):

Ingesting corpus:
  chunking-strategies-in-langchain-rag-a-deep-dive: 37 chunks
  how-i-created-an-mcp-server-for-postgresql-to-power-ai-agents-components-architecture-and-real-testing: 33 chunks
  how-to-pick-the-perfect-database-without-losing-your-mind: 51 chunks
  these-ai-memory-types-decide-whether-your-agent-is-smart-or-useless: 19 chunks
  upgrading-my-postgresql-ai-agent-3-architecture-decisions-i-made-and-why: 38 chunks
Total chunks stored: 178

Chunk count scales with document length and structure, not a fixed number per document — the database-picking post produced 51 chunks (it's the longest, most heavily ##-sectioned post in the corpus) while the AI memory post produced 19 (shorter, flatter structure).

Step 5: Retrieval

def search(conn, embedder, query, k=3):
    qvec = embedder.embed_query(query)
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT document_title, chunk_index, content, 1 - (embedding <=> %s::vector) AS similarity
            FROM rag_chunks
            ORDER BY embedding <=> %s::vector
            LIMIT %s;
            """,
            (qvec, qvec, k),
        )
        return cur.fetchall()

<=> is pgvector's cosine distance operator — smaller means more similar, which is why 1 - distance is computed separately to display an intuitive similarity score while the raw distance still drives the ORDER BY.

Measured results, three real queries against the 178 stored chunks:


Query: 'how should I structure memory for an AI agent'
-------------------------------------------------------------

1. similarity=0.819  [# These AI Memory Types Decide Whether Your A... chunk #18]
   If you truly want to build production-grade AI agents…  start with memory architecture first....
2. similarity=0.812  [# These AI Memory Types Decide Whether Your A... chunk #3]
   ## Why Memory Matters in AI Agents  Imagine hiring a human assistant who:  *   forgets your name eve...


Query: 'what commands should be blocked in a database agent for safety'
-------------------------------------------------------------

1. similarity=0.742  [# Upgrading My PostgreSQL AI Agent: 3 Archite... chunk #18]
   Two things worth noting here. First, this goes through `execute_sql()` — the single gateway from Dec...
2. similarity=0.737  [# Upgrading My PostgreSQL AI Agent: 3 Archite... chunk #11]
   The private connection API is about making the right path the only path. A future contributor should...


Query: 'why do I need to clean up log files on a server'
-------------------------------------------------------------

1. similarity=0.696  [# Upgrading My PostgreSQL AI Agent: 3 Archite... chunk #34]
   **The interface abstraction is what matters.** Because `save_session()` and `load_session()` are the...
2. similarity=0.657  [# Upgrading My PostgreSQL AI Agent: 3 Archite... chunk #12]
   ## Decision 2: Schema Caching at Startup...

Final Thoughts

Postgres, with the vector extension enabled, can store embeddings in a normal table and rank them by cosine similarity against a query using nothing more exotic than ORDER BY embedding <=> query_vector. No separate vector database, no new infrastructure to run — a schema, an insert, and a query, all standard SQL plus one operator.

Some of the results still wasn't clearly relevant, even with a real embedding model in place as we can see in the last query result. That's not a reason to distrust pgvector or the embedding — it's a sign that plain cosine similarity over chunk embeddings is a starting point, not the finished system.