-- 第一步：安装 pgvector 扩展（如已安装可跳过）
CREATE EXTENSION IF NOT EXISTS vector;

-- 第二步：安装 pgvectorscale 扩展
CREATE EXTENSION IF NOT EXISTS vectorscale;

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

-- 创建包含128维向量的文档表
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT,
    embedding vector(128)  -- 向量维度根据实际模型调整
);

-- 余弦距离（最常用，适合语义搜索）
CREATE INDEX idx_doc_cosine ON documents 
USING diskann (embedding vector_cosine_ops);

-- 欧氏距离（L2）
CREATE INDEX idx_doc_l2 ON documents 
USING diskann (embedding vector_l2_ops);

-- 内积距离
CREATE INDEX idx_doc_ip ON documents 
USING diskann (embedding vector_ip_ops);

-- 插入单条数据
INSERT INTO documents (title, content, embedding) 
VALUES (
    '云数据库 PostgreSQL 简介',
    'PostgreSQL 是功能强大的开源关系数据库...',
    '[0.1, 0.2, -0.3, ...]'::vector
);

-- 批量插入（推荐使用 COPY 或多行 INSERT）
INSERT INTO documents (title, content, embedding) VALUES
    ('文档1', '内容1', '[0.1, 0.2, ...]'::vector),
    ('文档2', '内容2', '[0.3, 0.4, ...]'::vector);

-- 余弦距离查询 Top 5最相似文档
SELECT id, title, embedding <=> '[0.1, 0.2, ...]'::vector AS distance
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, ...]'::vector
LIMIT 5;

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

-- 内积查询
SELECT id, title, embedding <#> '[0.1, 0.2, ...]'::vector AS distance
FROM documents
ORDER BY embedding <#> '[0.1, 0.2, ...]'::vector
LIMIT 5;

-- 高 Recall 配置（适合对准确率要求高的场景）
CREATE INDEX idx_high_recall ON documents 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 64, search_list_size = 200);

-- 高吞吐配置（适合对延迟要求高的场景）
CREATE INDEX idx_low_latency ON documents 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 32, search_list_size = 80);

-- 大数据集配置（百万级以上）
CREATE INDEX idx_large_scale ON documents 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 50, search_list_size = 150);

-- 创建知识库表
CREATE TABLE knowledge_base (
    id SERIAL PRIMARY KEY,
    doc_title TEXT,
    chunk_text TEXT,
    embedding vector(1024),  -- 以 BGE-M3模型为例
    metadata JSONB
);

-- 创建 DiskANN 索引
CREATE INDEX idx_kb_diskann ON knowledge_base 
USING diskann (embedding vector_cosine_ops)
WITH (num_neighbors = 50, search_list_size = 150);

-- 检索与问题最相关的5个文档片段
SELECT chunk_text, metadata,
       embedding <=> $1::vector AS distance
FROM knowledge_base
ORDER BY embedding <=> $1::vector
LIMIT 5;

-- 商品向量表
CREATE TABLE product_embeddings (
    product_id BIGINT PRIMARY KEY,
    product_name TEXT,
    category TEXT,
    embedding vector(256)
);

-- 按类别过滤的相似商品搜索
SELECT product_id, product_name,
       embedding <=> $1::vector AS distance
FROM product_embeddings
WHERE category = '电子产品'
ORDER BY embedding <=> $1::vector
LIMIT 20;

-- 图片向量表
CREATE TABLE image_vectors (
    image_id BIGINT PRIMARY KEY,
    image_url TEXT,
    description TEXT,
    embedding vector(512)  -- CLIP 模型输出
);

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

-- 以图搜图
SELECT image_id, image_url, description,
       embedding <=> $1::vector AS distance
FROM image_vectors
ORDER BY embedding <=> $1::vector
LIMIT 10;

-- 重建索引
REINDEX INDEX idx_doc_cosine;

-- 或并发重建（不阻塞查询）
REINDEX INDEX CONCURRENTLY idx_doc_cosine;

-- 查看索引大小
SELECT pg_size_pretty(pg_relation_size('idx_doc_cosine')) AS index_size;

-- 查看索引使用情况
SELECT indexrelname, idx_scan, idx_tup_read, idx_tup_fetch
FROM pg_stat_user_indexes
WHERE indexrelname LIKE '%diskann%';