PostgreSQL向量检索之pgvector入门实战指南

CREATE EXTENSION IF NOT EXISTS vector;

SELECT * FROM pg_extension WHERE extname='vector';

CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    embedding VECTOR(3) -- 假设嵌入向量是 3 维
);

INSERT INTO documents (embedding) VALUES ('[1,2,3]'), ('[4,5,6]'),('[4,0,5]');

-- 使用 L2 距离
SELECT id,embedding <-> '[1,2,3]' as score,embedding FROM documents ORDER BY embedding <-> '[1,2,3]' LIMIT 5;

# 向现有表中添加向量列
ALTER TABLE items ADD COLUMN embedding vector(3);
# 更新向量
UPDATE items SET embedding = '[1,2,3]' WHERE id = 1;
# 删除向量
UPDATE items SET embedding = '[1,2,3]' WHERE id = 1;

import numpy as np
import psycopg2
from pgvector.psycopg2 import register_vector
# 连接数据库
conn = psycopg2.connect(host='192.168.1.101', dbname='test-agent', user='postgres', password='123456', port=5433)
register_vector(conn)  # 注册 vector 类型
cur = conn.cursor()
# 插入示例向量
embedding = np.random.rand(3).astype('float32')
print(embedding)
cur.execute(
    "INSERT INTO documents (embedding) VALUES (%s)",
    (embedding,)  # 注意这里的逗号，确保是元组格式
)
conn.commit()
# 查询最近邻
query_embedding = np.random.rand(3).astype('float32')
print(query_embedding)
cur.execute(
    "SELECT id, embedding,embedding <-> %s as score  FROM documents ORDER BY embedding <-> %s LIMIT 5",
    (query_embedding, query_embedding,)
)
rows = cur.fetchall()
for r in rows:
    print(r)
cur.close()
conn.close()