sqlite-vec vs pgvector: Tìm kiếm Vector Nhúng Cục bộ cho Ứng dụng Desktop & Edge

import numpy as np
import sqlite3
import time
from pgvector.psycopg import register_vector
import psycopg
import os

# Configuration
NUM_EMBEDDINGS = 100_000
DIMENSIONS = 384
SQLITE_DB_PATH = "embeddings.db"
PG_DB_NAME = "vector_benchmark"
PG_USER = "postgres"
PG_PASSWORD = "password" # Use a strong password in production
PG_HOST = "localhost"
PG_PORT = 5432

# Generate synthetic embeddings
print(f"Generating {NUM_EMBEDDINGS} synthetic embeddings...")
embeddings = np.random.rand(NUM_EMBEDDINGS, DIMENSIONS).astype(np.float32)
# Normalize embeddings to unit length, common practice for cosine similarity
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
print("Embeddings generated.")

# --- SQLite-vec Setup ---
print("\n--- Setting up sqlite-vec ---")
try:
    # Connect to SQLite database
    sqlite_conn = sqlite3.connect(SQLITE_DB_PATH)
    sqlite_cursor = sqlite_conn.cursor()

    # Load sqlite-vec extension (adjust path as necessary for your OS)
    # On Linux: libsqlite_vec.so, macOS: libsqlite_vec.dylib, Windows: sqlite_vec.dll
    # Ensure you have compiled/downloaded the correct extension for your system.
    # Example for a pre-compiled Linux extension:
    sqlite_conn.enable_load_extension(True)
    try:
        sqlite_cursor.execute("SELECT load_extension('./libsqlite_vec.so');")
    except sqlite3.OperationalError as e:
        print(f"Error loading sqlite-vec extension: {e}")
        print("Please ensure 'libsqlite_vec.so' (or .dylib/.dll) is in the current directory or specified path.")
        print("You might need to compile it from source: https://github.com/asg017/sqlite-vec")
        exit(1)
    sqlite_conn.enable_load_extension(False)

    # Create table for embeddings
    sqlite_cursor.execute(f"""
        CREATE TABLE IF NOT EXISTS items (
            id INTEGER PRIMARY KEY,
            embedding BLOB
        );
    """)
    sqlite_conn.commit()

    # Insert embeddings
    print(f"Inserting {NUM_EMBEDDINGS} embeddings into sqlite-vec...")
    insert_start_time = time.time()
    batch_size = 1000
    for i in range(0, NUM_EMBEDDINGS, batch_size):
        batch = embeddings[i:i+batch_size]
        sqlite_cursor.executemany(
            "INSERT INTO items (id, embedding) VALUES (?, ?)",
            [(i + j, vec.tobytes()) for j, vec in enumerate(batch)]
        )
        sqlite_conn.commit()
    insert_end_time = time.time()
    print(f"sqlite-vec insertion time: {insert_end_time - insert_start_time:.2f} seconds")

except Exception as e:
    print(f"Error during sqlite-vec setup: {e}")
    if 'sqlite_vec' in str(e):
        print("Ensure the sqlite-vec extension is correctly compiled and loaded.")
    exit(1)


# --- pgvector Setup ---
print("\n--- Setting up pgvector ---")
try:
    # Connect to PostgreSQL
    pg_conn = psycopg.connect(f"dbname={PG_DB_NAME} user={PG_USER} password={PG_PASSWORD} host={PG_HOST} port={PG_PORT}")
    pg_conn.autocommit = True # For CREATE DATABASE/EXTENSION
    pg_cursor = pg_conn.cursor()

    # Create database if not exists
    try:
        pg_cursor.execute(f"CREATE DATABASE {PG_DB_NAME};")
        print(f"Database '{PG_DB_NAME}' created.")
    except psycopg.errors.DuplicateDatabase:
        print(f"Database '{PG_DB_NAME}' already exists.")
    except Exception as e:
        print(f"Error creating database: {e}")
        print("Ensure PostgreSQL is running and user has permissions.")
        exit(1)

    pg_conn.close() # Close and reconnect to the new database

    pg_conn = psycopg.connect(f"dbname={PG_DB_NAME} user={PG_USER} password={PG_PASSWORD} host={PG_HOST} port={PG_PORT}")
    pg_cursor = pg_conn.cursor()

    # Enable pgvector extension
    pg_cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
    print("pgvector extension enabled.")

    # Register vector type for psycopg
    register_vector(pg_conn)

    # Create table for embeddings
    pg_cursor.execute(f"""
        CREATE TABLE IF NOT EXISTS items (
            id SERIAL PRIMARY KEY,
            embedding vector({DIMENSIONS})
        );
    """)
    pg_conn.commit()

    # Insert embeddings
    print(f"Inserting {NUM_EMBEDDINGS} embeddings into pgvector...")
    insert_start_time = time.time()
    batch_size = 1000
    for i in range(0, NUM_EMBEDDINGS, batch_size):
        batch = embeddings[i:i+batch_size]
        pg_cursor.executemany(
            "INSERT INTO items (embedding) VALUES (%s)",
            [(vec,) for vec in batch]
        )
        pg_conn.commit()
    insert_end_time = time.time()
    print(f"pgvector insertion time: {insert_end_time - insert_start_time:.2f} seconds")

except Exception as e:
    print(f"Error during pgvector setup: {e}")
    print("Ensure PostgreSQL is running and accessible.")
    exit(1)

print("\nData setup complete for both databases.")

import numpy as np
import sqlite3
import time
from pgvector.psycopg import register_vector
import psycopg
import os
import faiss # For ground truth and recall calculation
from sklearn.metrics import top_k_accuracy_score

# Configuration (same as above, ensure NUM_EMBEDDINGS, DIMENSIONS match)
NUM_EMBEDDINGS = 100_000
DIMENSIONS = 384
SQLITE_DB_PATH = "embeddings.db"
PG_DB_NAME = "vector_benchmark"
PG_USER = "postgres"
PG_PASSWORD = "password"
PG_HOST = "localhost"
PG_PORT = 5432
K_SEARCH = 10 # Number of neighbors to retrieve
NUM_QUERIES = 100 # Number of queries for benchmarking

# Reconnect to databases
sqlite_conn = sqlite3.connect(SQLITE_DB_PATH)
sqlite_conn.enable_load_extension(True)
sqlite_cursor = sqlite_conn.cursor()
try:
    sqlite_cursor.execute("SELECT load_extension('./libsqlite_vec.so');")
except sqlite3.OperationalError as e:
    print(f"Error loading sqlite-vec extension: {e}")
    exit(1)
sqlite_conn.enable_load_extension(False)

pg_conn = psycopg.connect(f"dbname={PG_DB_NAME} user={PG_USER} password={PG_PASSWORD} host={PG_HOST} port={PG_PORT}")
pg_cursor = pg_conn.cursor()
register_vector(pg_conn)

# Load all embeddings for ground truth and query generation
print("Loading all embeddings for ground truth...")
sqlite_cursor.execute("SELECT embedding FROM items ORDER BY id")
all_embeddings_bytes = sqlite_cursor.fetchall()
all_embeddings = np.array([np.frombuffer(e[0], dtype=np.float32) for e in all_embeddings_bytes])
print(f"Loaded {len(all_embeddings)} embeddings.")

# Generate query embeddings (random subset of existing embeddings)
query_indices = np.random.choice(NUM_EMBEDDINGS, NUM_QUERIES, replace=False)
query_embeddings = all_embeddings[query_indices]

# --- Ground Truth (using FAISS for exact nearest neighbors) ---
print("\nCalculating ground truth using FAISS...")
faiss_index = faiss.IndexFlatL2(DIMENSIONS) # L2 for cosine similarity on normalized vectors
faiss_index.add(all_embeddings)
D_gt, I_gt = faiss_index.search(query_embeddings, K_SEARCH)
print("Ground truth calculated.")

# Function to calculate recall@K
def calculate_recall(retrieved_ids, ground_truth_ids, k):
    hits = 0
    for i in range(len(retrieved_ids)):
        # Check if any of the top K retrieved IDs are in the ground truth top K
        # Note: For exact matches, we expect the query itself to be in the top K.
        # For approximate search, we check if the retrieved set overlaps with the ground truth set.
        retrieved_set = set(retrieved_ids[i])
        ground_truth_set = set(ground_truth_ids[i])
        if not ground_truth_set.isdisjoint(retrieved_set):
            hits += 1
    return hits / len(retrieved_ids)

# --- Benchmarking Function ---
def benchmark_search(db_type, index_type, index_params, query_sql, conn, cursor, is_sqlite_vec=False, is_quantized=False):
    print(f"\n--- Benchmarking {db_type} with {index_type} ({index_params}) ---")

    # Drop existing index if any
    if db_type == "sqlite-vec":
        try:
            cursor.execute(f"DROP TABLE IF EXISTS items_vec_idx;")
            conn.commit()
        except sqlite3.OperationalError:
            pass # Index might not exist
    elif db_type == "pgvector":
        try:
            cursor.execute(f"DROP INDEX IF EXISTS items_embedding_idx;")
            conn.commit()
        except psycopg.errors.UndefinedObject:
            pass # Index might not exist

    # Build Index
    print(f"Building {index_type} index...")
    build_start_time = time.time()
    if db_type == "sqlite-vec":
        if is_quantized:
            cursor.execute(f"CREATE VIRTUAL TABLE items_vec_idx USING vec0(items, embedding, {DIMENSIONS}, 'quantizer=int8', '{index_type}={index_params}');")
        else:
            cursor.execute(f"CREATE VIRTUAL TABLE items_vec_idx USING vec0(items, embedding, {DIMENSIONS}, '{index_type}={index_params}');")
        conn.commit()
    elif db_type == "pgvector":
        cursor.execute(f"CREATE INDEX items_embedding_idx ON items USING {index_type}(embedding {index_params});")
        conn.commit()
    build_end_time = time.time()
    print(f"Index build time: {build_end_time - build_start_time:.2f} seconds")

    # Measure RAM usage (approximate, depends on OS/tooling)
    # For sqlite-vec, this is the process's RAM. For pgvector, it's the PostgreSQL server's RAM.
    # This is hard to measure programmatically in a cross-platform way.
    # For a real benchmark, use OS-specific tools (e.g., `ps aux` on Linux, Activity Monitor on macOS).
    print("RAM usage measurement requires external OS tools.")

    # Perform searches
    search_latencies = []
    retrieved_ids_list = []

    print(f"Performing {NUM_QUERIES} searches...")
    for i, query_vec in enumerate(query_embeddings):
        search_start_time = time.time()
        if db_type == "sqlite-vec":
            # For sqlite-vec, the query uses the virtual table
            cursor.execute(query_sql, (query_vec.tobytes(), K_SEARCH))
        elif db_type == "pgvector":
            cursor.execute(query_sql, (query_vec, K_SEARCH))

        results = cursor.fetchall()
        search_end_time = time.time()
        search_latencies.append(search_end_time - search_start_time)
        retrieved_ids_list.append([r[0] for r in results]) # Assuming ID is the first column

    avg_latency = np.mean(search_latencies)
    p95_latency = np.percentile(search_latencies, 95)
    print(f"Average search latency: {avg_latency * 1000:.2f} ms")
    print(f"P95 search latency: {p95_latency * 1000:.2f} ms")

    # Calculate recall
    # We need to map the ground truth IDs to the retrieved IDs.
    # For simplicity, we assume the IDs are 0-indexed and correspond to the original `all_embeddings` array.
    # The ground truth `I_gt` contains indices, which are effectively IDs.
    recall = calculate_recall(retrieved_ids_list, I_gt, K_SEARCH)
    print(f"Recall@{K_SEARCH}: {recall:.4f}")

    return {
        "db_type": db_type,
        "index_type": index_type,
        "index_params": index_params,
        "build_time": build_end_time - build_start_time,
        "avg_latency_ms": avg_latency * 1000,
        "p95_latency_ms": p95_latency * 1000,
        "recall_at_k": recall
    }

# --- Run Benchmarks ---
results = []

# sqlite-vec: HNSW (default float32)
results.append(benchmark_search(
    "sqlite-vec", "hnsw", "M=16,ef_construction=100",
    "SELECT id FROM items_vec_idx WHERE embedding MATCH ? ORDER BY distance LIMIT ?",
    sqlite_conn, sqlite_cursor, is_sqlite_vec=True
))

# sqlite-vec: HNSW with int8 quantization
results.append(benchmark_search(
    "sqlite-vec", "hnsw", "M=16,ef_construction=100",
    "SELECT id FROM items_vec_idx WHERE embedding MATCH ? ORDER BY distance LIMIT ?",
    sqlite_conn, sqlite_cursor, is_sqlite_vec=True, is_quantized=True
))

# pgvector: IVF (lists=100)
# For 100k embeddings, lists=100 is a reasonable starting point.
# The `vector_l2_ops` is for L2 distance, which is equivalent to cosine for normalized vectors.
results.append(benchmark_search(
    "pgvector", "ivfflat", "vector_l2_ops, lists=100",
    f"SELECT id FROM items ORDER BY embedding <-> %s LIMIT %s",
    pg_conn, pg_cursor
))

# pgvector: HNSW (m=16, ef_construction=64)
# Note: pgvector's HNSW parameters are slightly different from sqlite-vec's.
# `m` is the number of neighbors, `ef_construction` is for build time.
results.append(benchmark_search(
    "pgvector", "hnsw", "vector_l2_ops, m=16, ef_construction=64",
    f"SELECT id FROM items ORDER BY embedding <-> %s LIMIT %s",
    pg_conn, pg_cursor
))

# Print results table
print("\n--- Benchmark Summary ---")
print("| DB | Index Type | Params | Build Time (s) | Avg Latency (ms) | P95 Latency (ms) | Recall@10 |")
print("|---|---|---|---|---|---|---|")
for r in results:
    print(f"| {r['db_type']} | {r['index_type']} | {r['index_params']} | {r['build_time']:.2f} | {r['avg_latency_ms']:.2f} | {r['p95_latency_ms']:.2f} | {r['recall_at_k']:.4f} |")

# Clean up
sqlite_conn.close()
pg_conn.close()