-- Step 1: Install the pgvector extension (skip this step if it is already installed)
CREATE EXTENSION IF NOT EXISTS vector;

-- Step 2: Install the pgvectorscale extension
CREATE EXTENSION IF NOT EXISTS vectorscale;

SELECT extname, extversion FROM pg_extension 
WHERE extname IN ('vector', 'vectorscale');

-- Create a document table containing 128-dimensional vectors.
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT,
    embedding vector(128)  -- Adjust the vector dimension according to your actual model.
);

-- Cosine distance (most commonly used, suitable for semantic search)
CREATE INDEX idx_doc_cosine ON documents 
USING diskann (embedding vector_cosine_ops);

-- Euclidean distance (L2)
CREATE INDEX idx_doc_l2 ON documents 
USING diskann (embedding vector_l2_ops);

-- Inner product distance
CREATE INDEX idx_doc_ip ON documents 
USING diskann (embedding vector_ip_ops);

-- Insert a single data record
INSERT INTO documents (title, content, embedding) 
VALUES (
    'TencentDB for PostgreSQL Introduction',
    'PostgreSQL is a powerful, open-source relational database...',
    '[0.1, 0.2, -0.3, ...]'::vector
);

-- Batch insertion (recommended to use COPY or multi-row INSERT)
INSERT INTO documents (title, content, embedding) VALUES
    ('Document 1', 'Content 1', '[0.1, 0.2, ...]'::vector),
    ('Document 2', 'Content 2', '[0.3, 0.4, ...]'::vector);

-- Query the Top 5 most similar documents using cosine distance
SELECT id, title, embedding <=> '[0.1, 0.2, ...]'::vector AS distance
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, ...]'::vector
LIMIT 5;

-- L2 distance query
SELECT id, title, embedding <-> '[0.1, 0.2, ...]'::vector AS distance
FROM documents
ORDER BY embedding <-> '[0.1, 0.2, ...]'::vector
LIMIT 5;

-- Inner product query
SELECT id, title, embedding <#> '[0.1, 0.2, ...]'::vector AS distance
FROM documents
ORDER BY embedding <#> '[0.1, 0.2, ...]'::vector
LIMIT 5;

-- High Recall configuration (suitable for scenarios requiring high accuracy)
CREATE INDEX idx_high_recall ON documents 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 64, search_list_size = 200);

-- High Throughput configuration (suitable for scenarios requiring low latency)
CREATE INDEX idx_low_latency ON documents 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 32, search_list_size = 80);

-- Large Dataset configuration (millions of vectors or more)
CREATE INDEX idx_large_scale ON documents 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 50, search_list_size = 150);

-- Create a knowledge base table
CREATE TABLE knowledge_base (
    id SERIAL PRIMARY KEY,
    doc_title TEXT,
    chunk_text TEXT,
    embedding vector(1024),  -- Using the BGE-M3 model as an example
    metadata JSONB
);

-- Create a DiskANN index
CREATE INDEX idx_kb_diskann ON knowledge_base 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 50, search_list_size = 150);

-- Search for the 5 document fragments most relevant to the question
SELECT chunk_text, metadata,
       embedding <=> $1::vector AS distance
FROM knowledge_base
ORDER BY embedding <=> $1::vector
LIMIT 5;

-- Product vector table
CREATE TABLE product_embeddings (
    product_id BIGINT PRIMARY KEY,
    product_name TEXT,
    category TEXT,
    embedding vector(256)
);

-- Similar product search filtered by category
SELECT product_id, product_name,
       embedding <=> $1::vector AS distance
FROM product_embeddings
WHERE category = 'Electronics'
ORDER BY embedding <=> $1::vector
LIMIT 20;

-- Image vector table
CREATE TABLE image_vectors (
    image_id BIGINT PRIMARY KEY,
    image_url TEXT,
    description TEXT,
    embedding vector(512)  -- CLIP model output
);

CREATE INDEX idx_img_diskann ON image_vectors 
USING diskann (embedding vector_cosine_ops);

-- Search for images using an image
SELECT image_id, image_url, description,
       embedding <=> $1::vector AS distance
FROM image_vectors
ORDER BY embedding <=> $1::vector
LIMIT 10;

-- Reindex
REINDEX INDEX idx_doc_cosine;

-- Or perform concurrent rebuilding (non-blocking to queries)
REINDEX INDEX CONCURRENTLY idx_doc_cosine;

-- View index size
SELECT pg_size_pretty(pg_relation_size('idx_doc_cosine')) AS index_size;

-- View index usage
SELECT indexrelname, idx_scan, idx_tup_read, idx_tup_fetch
FROM pg_stat_user_indexes
WHERE indexrelname LIKE '%diskann%';