vector_search_optimization.py

# =============================================================================
# Vector Search Optimization - Python pgvector Examples
# pip install: psycopg[binary] pgvector numpy
# =============================================================================

import numpy as np
import psycopg
from pgvector.psycopg import register_vector


# =============================================================================
# Example 1: Create and Compare Index Types
# =============================================================================
def create_indexes():
    """Create HNSW and IVFFlat indexes and compare build times."""
    conninfo = (
        "host=myserver.postgres.database.azure.com "
        "port=5432 dbname=vectordb user=adminuser "
        "password=secret sslmode=require"
    )

    with psycopg.connect(conninfo) as conn:
        conn.autocommit = True
        register_vector(conn)

        with conn.cursor() as cur:
            # Create table
            cur.execute("""
                CREATE TABLE IF NOT EXISTS vec_items (
                    id SERIAL PRIMARY KEY,
                    content TEXT,
                    embedding vector(256)
                )
            """)

            # Insert random vectors if table is empty
            cur.execute("SELECT count(*) FROM vec_items")
            count = cur.fetchone()[0]
            if count == 0:
                print("Inserting 50,000 random vectors...")
                vectors = np.random.rand(50000, 256).astype(np.float32)
                for i, vec in enumerate(vectors):
                    cur.execute(
                        "INSERT INTO vec_items (content, embedding) VALUES (%s, %s)",
                        (f"document_{i}", vec.tolist()),
                    )
                print("Done inserting vectors.")

            # Build HNSW index
            cur.execute("DROP INDEX IF EXISTS idx_hnsw")
            cur.execute("SET maintenance_work_mem = '512MB'")
            print("Building HNSW index...")
            cur.execute("""
                CREATE INDEX idx_hnsw ON vec_items
                USING hnsw (embedding vector_cosine_ops)
                WITH (m = 16, ef_construction = 64)
            """)
            print("HNSW index created.")

            # Build IVFFlat index (lists = 50 for 50K rows)
            cur.execute("DROP INDEX IF EXISTS idx_ivfflat")
            print("Building IVFFlat index...")
            cur.execute("""
                CREATE INDEX idx_ivfflat ON vec_items
                USING ivfflat (embedding vector_cosine_ops)
                WITH (lists = 50)
            """)
            print("IVFFlat index created.")

            # Show index sizes
            cur.execute("""
                SELECT indexname,
                       pg_size_pretty(pg_relation_size(indexname::regclass)) AS size
                FROM pg_indexes
                WHERE tablename = 'vec_items' AND indexname LIKE 'idx_%'
            """)
            print("\nIndex sizes:")
            for row in cur.fetchall():
                print(f"  {row[0]}: {row[1]}")

            cur.execute("RESET maintenance_work_mem")


# =============================================================================
# Example 2: Measure Recall at Different Parameter Settings
# =============================================================================
def measure_recall():
    """Compare recall between exact search and ANN search at various settings."""
    conninfo = (
        "host=myserver.postgres.database.azure.com "
        "port=5432 dbname=vectordb user=adminuser "
        "password=secret sslmode=require"
    )

    with psycopg.connect(conninfo) as conn:
        register_vector(conn)
        k = 20  # number of nearest neighbors

        with conn.cursor() as cur:
            # Get a query vector
            cur.execute("SELECT embedding FROM vec_items WHERE id = 1")
            query_vec = cur.fetchone()[0]

            # Ground truth: exact nearest neighbors (sequential scan)
            cur.execute("SET enable_indexscan = off")
            cur.execute("SET enable_bitmapscan = off")
            cur.execute(
                "SELECT id FROM vec_items WHERE id != 1 "
                "ORDER BY embedding <=> %s LIMIT %s",
                (query_vec, k),
            )
            ground_truth = {row[0] for row in cur.fetchall()}
            cur.execute("SET enable_indexscan = on")
            cur.execute("SET enable_bitmapscan = on")

            print(f"Ground truth ({k} nearest neighbors): {sorted(ground_truth)}")

            # Test HNSW with different ef_search values
            # Ensure HNSW index exists
            cur.execute("DROP INDEX IF EXISTS idx_ivfflat")
            print("\n--- HNSW Recall ---")
            for ef in [10, 40, 100, 200]:
                cur.execute(f"SET hnsw.ef_search = {ef}")
                cur.execute(
                    "SELECT id FROM vec_items WHERE id != 1 "
                    "ORDER BY embedding <=> %s LIMIT %s",
                    (query_vec, k),
                )
                ann_results = {row[0] for row in cur.fetchall()}
                recall = len(ground_truth & ann_results) / len(ground_truth)
                print(f"  ef_search={ef:3d}: recall={recall:.2f} "
                      f"({len(ground_truth & ann_results)}/{k} correct)")

            # Rebuild IVFFlat and test with different probes
            cur.execute("""
                CREATE INDEX idx_ivfflat ON vec_items
                USING ivfflat (embedding vector_cosine_ops)
                WITH (lists = 50)
            """)
            cur.execute("DROP INDEX idx_hnsw")
            print("\n--- IVFFlat Recall ---")
            for probes in [1, 5, 10, 25]:
                cur.execute(f"SET ivfflat.probes = {probes}")
                cur.execute(
                    "SELECT id FROM vec_items WHERE id != 1 "
                    "ORDER BY embedding <=> %s LIMIT %s",
                    (query_vec, k),
                )
                ann_results = {row[0] for row in cur.fetchall()}
                recall = len(ground_truth & ann_results) / len(ground_truth)
                print(f"  probes={probes:3d}: recall={recall:.2f} "
                      f"({len(ground_truth & ann_results)}/{k} correct)")

            # Recreate HNSW for other examples
            cur.execute("""
                CREATE INDEX idx_hnsw ON vec_items
                USING hnsw (embedding vector_cosine_ops)
            """)


# =============================================================================
# Example 3: EXPLAIN ANALYZE to Verify Index Usage
# =============================================================================
def explain_analyze():
    """Use EXPLAIN ANALYZE to check whether a vector index is being used."""
    conninfo = (
        "host=myserver.postgres.database.azure.com "
        "port=5432 dbname=vectordb user=adminuser "
        "password=secret sslmode=require"
    )

    with psycopg.connect(conninfo) as conn:
        register_vector(conn)

        with conn.cursor() as cur:
            cur.execute("SELECT embedding FROM vec_items WHERE id = 1")
            query_vec = cur.fetchone()[0]

            # Check query plan
            cur.execute(
                "EXPLAIN ANALYZE "
                "SELECT id, embedding <=> %s AS distance "
                "FROM vec_items "
                "ORDER BY embedding <=> %s LIMIT 10",
                (query_vec, query_vec),
            )
            plan = cur.fetchall()

            print("Query plan:")
            for line in plan:
                text = line[0]
                print(f"  {text}")
                # Flag whether index is used
                if "Index Scan" in text:
                    print("  >>> INDEX IS BEING USED")
                elif "Seq Scan" in text:
                    print("  >>> WARNING: SEQUENTIAL SCAN (no index used)")


# =============================================================================
# Example 4: Operator Class Mismatch Detection
# =============================================================================
def detect_operator_mismatch():
    """Demonstrate and detect operator class mismatches."""
    conninfo = (
        "host=myserver.postgres.database.azure.com "
        "port=5432 dbname=vectordb user=adminuser "
        "password=secret sslmode=require"
    )

    with psycopg.connect(conninfo) as conn:
        register_vector(conn)

        with conn.cursor() as cur:
            cur.execute("SELECT embedding FROM vec_items WHERE id = 1")
            query_vec = cur.fetchone()[0]

            # Check which indexes exist and their operator classes
            cur.execute("""
                SELECT indexname, indexdef
                FROM pg_indexes
                WHERE tablename = 'vec_items' AND indexname LIKE 'idx_%'
            """)
            print("Current indexes:")
            for row in cur.fetchall():
                print(f"  {row[0]}: {row[1]}")

            # Test cosine operator (should use index if vector_cosine_ops exists)
            cur.execute(
                "EXPLAIN (FORMAT TEXT) "
                "SELECT id FROM vec_items ORDER BY embedding <=> %s LIMIT 10",
                (query_vec,),
            )
            plan_cosine = [row[0] for row in cur.fetchall()]
            uses_index = any("Index Scan" in line for line in plan_cosine)
            print(f"\nCosine <=> query: {'USES INDEX' if uses_index else 'SEQ SCAN'}")

            # Test L2 operator (will seq scan if only cosine index exists)
            cur.execute(
                "EXPLAIN (FORMAT TEXT) "
                "SELECT id FROM vec_items ORDER BY embedding <-> %s LIMIT 10",
                (query_vec,),
            )
            plan_l2 = [row[0] for row in cur.fetchall()]
            uses_index = any("Index Scan" in line for line in plan_l2)
            print(f"L2 <-> query:     {'USES INDEX' if uses_index else 'SEQ SCAN (operator mismatch!)'}")


# =============================================================================
# Example 5: Hybrid Search — Vector + Full-Text
# =============================================================================
def hybrid_search():
    """Combine full-text keyword filtering with vector similarity ranking."""
    conninfo = (
        "host=myserver.postgres.database.azure.com "
        "port=5432 dbname=vectordb user=adminuser "
        "password=secret sslmode=require"
    )

    with psycopg.connect(conninfo) as conn:
        conn.autocommit = True
        register_vector(conn)

        with conn.cursor() as cur:
            # Ensure tsvector column and GIN index exist
            cur.execute("""
                DO $$
                BEGIN
                    IF NOT EXISTS (
                        SELECT 1 FROM information_schema.columns
                        WHERE table_name = 'vec_items' AND column_name = 'content_tsv'
                    ) THEN
                        ALTER TABLE vec_items ADD COLUMN content_tsv tsvector
                            GENERATED ALWAYS AS (
                                to_tsvector('english', coalesce(content, ''))
                            ) STORED;
                        CREATE INDEX idx_fts ON vec_items USING GIN (content_tsv);
                    END IF;
                END $$;
            """)

            # Update some rows with searchable content
            updates = [
                (1, "Azure OpenAI provides GPT models for text generation"),
                (2, "PostgreSQL pgvector enables vector similarity search"),
                (3, "Azure Cognitive Search supports hybrid search patterns"),
                (4, "Azure Functions runs serverless compute workloads"),
                (5, "Azure OpenAI embeddings convert text to vectors"),
            ]
            for doc_id, content in updates:
                cur.execute(
                    "UPDATE vec_items SET content = %s WHERE id = %s",
                    (content, doc_id),
                )

            # Get query vector
            cur.execute("SELECT embedding FROM vec_items WHERE id = 1")
            query_vec = cur.fetchone()[0]

            # Hybrid query: keyword filter + vector ranking
            cur.execute(
                "SELECT id, content, embedding <=> %s AS distance "
                "FROM vec_items "
                "WHERE content_tsv @@ plainto_tsquery('english', %s) "
                "ORDER BY embedding <=> %s "
                "LIMIT 5",
                (query_vec, "Azure OpenAI", query_vec),
            )
            results = cur.fetchall()

            print("Hybrid search results (keyword: 'Azure OpenAI', ranked by vector distance):")
            for row in results:
                print(f"  id={row[0]}, distance={row[2]:.4f}, content={row[1]}")


if __name__ == "__main__":
    print("=== Example 1: Create and Compare Indexes ===")
    create_indexes()

    print("\n=== Example 2: Measure Recall ===")
    measure_recall()

    print("\n=== Example 3: EXPLAIN ANALYZE ===")
    explain_analyze()

    print("\n=== Example 4: Operator Class Mismatch ===")
    detect_operator_mismatch()

    print("\n=== Example 5: Hybrid Search ===")
    hybrid_search()