Hybrid Search in Postgres: RRF for pgvector & FTS

RRF_Score(d) = Σ (1 / (k + rank_i(d)))

-- Ensure the necessary extensions are enabled
CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pg_trgm;

-- The main table for our documents
CREATE TABLE documents (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    content TEXT NOT NULL,
    embedding VECTOR(1536), -- Example dimension for OpenAI embeddings
    content_tsv TSVECTOR,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Create an HNSW index for vector search. HNSW is generally preferred for its
-- balance of speed and recall over IVFFlat for many use cases.
-- m: max number of connections per layer (higher means better recall, slower build)
-- ef_construction: size of the dynamic list for neighbors during build (higher means better index quality)
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);

-- Create a GIN index for FTS. GIN is ideal for tsvector columns.
CREATE INDEX ON documents USING GIN (content_tsv);

-- A function to update the tsvector column automatically
CREATE OR REPLACE FUNCTION update_content_tsv()
RETURNS TRIGGER AS $$
BEGIN
    NEW.content_tsv := to_tsvector('english', NEW.content);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- A trigger to execute the function on insert or update
CREATE TRIGGER tsvector_update
BEFORE INSERT OR UPDATE ON documents
FOR EACH ROW EXECUTE FUNCTION update_content_tsv();

// src/database.ts
import { Pool } from 'pg';

export const pool = new Pool({
  connectionString: process.env.DATABASE_URL,
});

// src/embedding-client.ts - A mock embedding client
// In a real application, this would call an API like OpenAI, Cohere, etc.
export class EmbeddingClient {
  async createEmbedding(text: string): Promise<number[]> {
    console.log(`Generating embedding for: "${text.substring(0, 30)}..."`);
    // In a real implementation, this would be a network call.
    // Here, we generate a normalized random vector for demonstration.
    const vector = Array.from({ length: 1536 }, () => Math.random() * 2 - 1);
    const magnitude = Math.sqrt(vector.reduce((sum, val) => sum + val * val, 0));
    return vector.map(v => v / magnitude);
  }
}

// src/hybrid-search-service.ts
import { pool } from './database';
import { EmbeddingClient } from './embedding-client';

// Type definitions for clarity
interface SearchResult {
  id: string;
  score: number; // The original score from the search engine
}

interface FusedResult {
  id: string;
  rrfScore: number;
}

interface Document {
  id: string;
  content: string;
  createdAt: Date;
}

export class HybridSearchService {
  private embeddingClient: EmbeddingClient;
  private readonly RRF_K = 60; // The 'k' constant for RRF

  constructor() {
    this.embeddingClient = new EmbeddingClient();
  }

  private async vectorSearch(queryVector: number[], limit: number): Promise<SearchResult[]> {
    const vectorString = `[${queryVector.join(',')}]`;
    const query = {
      text: `
        SELECT id, 1 - (embedding <=> $1) AS score
        FROM documents
        ORDER BY embedding <=> $1
        LIMIT $2;
      `,
      values: [vectorString, limit],
    };
    
    // pgvector HNSW requires setting ef_search for performance/accuracy tuning
    // This should be done per-transaction or per-session
    const client = await pool.connect();
    try {
      await client.query('SET LOCAL hnsw.ef_search = 100;'); // Tune this value
      const res = await client.query(query);
      return res.rows.map(row => ({ id: row.id, score: parseFloat(row.score) }));
    } finally {
      client.release();
    }
  }

  private async fullTextSearch(query: string, limit: number): Promise<SearchResult[]> {
    const queryText = {
      text: `
        SELECT id, ts_rank_cd(content_tsv, websearch_to_tsquery('english', $1)) AS score
        FROM documents
        WHERE content_tsv @@ websearch_to_tsquery('english', $1)
        ORDER BY score DESC
        LIMIT $2;
      `,
      values: [query, limit],
    };
    const res = await pool.query(queryText);
    return res.rows.map(row => ({ id: row.id, score: parseFloat(row.score) }));
  }

  private reciprocalRankFusion(results: SearchResult[][]): FusedResult[] {
    const rankedResults: { [docId: string]: number } = {};

    for (const resultSet of results) {
      for (let i = 0; i < resultSet.length; i++) {
        const docId = resultSet[i].id;
        const rank = i + 1;
        const score = 1 / (this.RRF_K + rank);

        if (!rankedResults[docId]) {
          rankedResults[docId] = 0;
        }
        rankedResults[docId] += score;
      }
    }

    const fused = Object.entries(rankedResults).map(([id, rrfScore]) => ({ id, rrfScore }));
    return fused.sort((a, b) => b.rrfScore - a.rrfScore);
  }

  private async hydrateDocuments(fusedResults: FusedResult[], limit: number): Promise<Document[]> {
    if (fusedResults.length === 0) {
      return [];
    }

    const topResults = fusedResults.slice(0, limit);
    const docIds = topResults.map(r => r.id);

    // We need to fetch the documents and preserve the RRF order.
    // A common pattern is using a JOIN with a VALUES list or a CASE statement.
    const query = {
      text: `
        SELECT id, content, created_at
        FROM documents
        JOIN unnest($1::uuid[]) WITH ORDINALITY t(id, ord) USING (id)
        ORDER BY t.ord;
      `,
      values: [docIds],
    };

    const res = await pool.query(query);
    return res.rows.map(row => ({
      id: row.id,
      content: row.content,
      createdAt: new Date(row.created_at),
    }));
  }

  public async search(query: string, limit: number = 10): Promise<Document[]> {
    // For production, fetch more results than needed for better fusion quality.
    // This is a critical tuning parameter.
    const searchLimit = limit * 5;

    // 1. Generate query embedding
    const queryVector = await this.embeddingClient.createEmbedding(query);

    // 2. Execute searches in parallel for performance
    const [vectorResults, ftsResults] = await Promise.all([
      this.vectorSearch(queryVector, searchLimit),
      this.fullTextSearch(query, searchLimit),
    ]);

    // 3. Fuse the results using RRF
    const fusedResults = this.reciprocalRankFusion([vectorResults, ftsResults]);

    // 4. Hydrate the final, ordered list of documents
    const finalDocuments = await this.hydrateDocuments(fusedResults, limit);

    return finalDocuments;
  }
}

// Example usage
async function main() {
  // Assume you have populated the 'documents' table with some data and embeddings
  const searchService = new HybridSearchService();
  const query = "advanced strategies for database indexing in postgres";

  console.log(`Executing hybrid search for: "${query}"`);
  const results = await searchService.search(query, 5);

  console.log("\n--- Hybrid Search Results ---");
  if (results.length > 0) {
    results.forEach((doc, index) => {
      console.log(`${index + 1}. [ID: ${doc.id}] ${doc.content.substring(0, 100)}...`);
    });
  } else {
    console.log("No results found.");
  }

  // Example of a query where FTS is likely to dominate
  const specificQuery = "PostgreSQL HNSW index with m=16";
  console.log(`\nExecuting hybrid search for: "${specificQuery}"`);
  const specificResults = await searchService.search(specificQuery, 5);

  console.log("\n--- Hybrid Search Results (Specific) ---");
  if (specificResults.length > 0) {
    specificResults.forEach((doc, index) => {
      console.log(`${index + 1}. [ID: ${doc.id}] ${doc.content.substring(0, 100)}...`);
    });
  } else {
    console.log("No results found.");
  }

  await pool.end();
}

main().catch(console.error);