在PostgreSQL中计算并存储向量余弦相似度（点积）的方案咨询

方案一：直接在PostgreSQL内用向量扩展计算（最优解） #

把向量运算放在数据库里做，能彻底避免数据来回传输的开销，而且PostgreSQL有专门的向量优化工具：

1. 用pgvector扩展（强烈推荐） #

这是PostgreSQL官方生态里的向量处理扩展，用C实现的原生向量运算，效率比Python/numpy高得多：

• 先安装扩展（不同环境安装方式略有差异，比如Ubuntu用apt install postgresql-15-pgvector），然后在数据库中启用：

  CREATE EXTENSION IF NOT EXISTS vector;
  

• 把你的Vector字段从JSONB转换成vector类型（更高效的存储和运算）：

  ALTER TABLE your_table ADD COLUMN vector_col vector(5000);
  UPDATE your_table SET vector_col = vector(Vector); -- 把JSONB数组转成vector类型
  -- 之后可以删掉原来的JSONB字段（可选）
  

• 预存每个向量的模长（避免重复计算）：

  ALTER TABLE your_table ADD COLUMN vector_norm float8;
  UPDATE your_table SET vector_norm = norm(vector_col);
  

• 创建存储点积的表，插入所有两两组合的点积（利用交换律只存id1 < id2的组合，省一半空间）：

  CREATE TABLE vector_dot_products (
      id1 INT REFERENCES your_table(id),
      id2 INT REFERENCES your_table(id),
      dot_product float8,
      PRIMARY KEY (id1, id2)
  );
  
  INSERT INTO vector_dot_products (id1, id2, dot_product)
  SELECT a.id, b.id, a.vector_col <#> b.vector_col
  FROM your_table a
  CROSS JOIN your_table b
  WHERE a.id < b.id;
  

后续计算余弦相似度时，直接用dot_product / (a.vector_norm * b.vector_norm)即可，速度拉满。

2. 原生JSONB函数（无扩展场景） #

如果没法装pgvector，可以用PostgreSQL的原生JSON函数写一个点积计算函数：

CREATE OR REPLACE FUNCTION jsonb_array_dot(a jsonb, b jsonb)
RETURNS float8 AS $$
SELECT SUM((a->>i)::float8 * (b->>i)::float8)
FROM generate_series(0, jsonb_array_length(a)-1) i
WHERE jsonb_array_length(a) = jsonb_array_length(b); -- 确保向量长度一致
$$ LANGUAGE sql IMMUTABLE;

之后的步骤和上面类似：预存模长、插入点积。缺点是纯SQL实现的效率不如pgvector，但比把数据导出到Python处理快。

方案二：优化Python + numpy的实现 #

如果必须用Python处理，那可以从这几个点优化效率：

• 批量拉取+矩阵运算：不要逐行取数据，一次性拉取批量数据转成numpy矩阵，用vectors @ vectors.T直接计算全量点积矩阵，比循环两两计算快几个数量级。
• 分块处理：如果数据量极大（比如10w+条，全量矩阵占内存太多），可以把数据分成若干块，分块计算点积再批量插入。
• 避免重复计算：只计算i < j的组合，减少一半的计算量和存储空间。

示例代码大概是这样：

import psycopg2
import numpy as np

# 连接数据库
conn = psycopg2.connect("dbname=your_db user=your_user password=your_pwd")
cur = conn.cursor()

# 一次性拉取所有数据（数据量太大就分批拉）
cur.execute("SELECT id, Vector FROM your_table")
data = cur.fetchall()
ids = [row[0] for row in data]
# 把JSONB数组转成numpy矩阵
vectors = np.array([np.array(row[1], dtype=np.float32) for row in data])

# 计算全量点积矩阵
dot_matrix = vectors @ vectors.T

# 生成待插入的记录（只存id1 < id2的组合）
records = []
for i in range(len(ids)):
    for j in range(i+1, len(ids)):
        records.append((ids[i], ids[j], float(dot_matrix[i][j])))

# 批量插入数据库
cur.executemany(
    "INSERT INTO vector_dot_products (id1, id2, dot_product) VALUES (%s, %s, %s)",
    records
)
conn.commit()

# 关闭连接
cur.close()
conn.close()

额外建议 #

• 如果你的表会有新增/更新，不要每次全量重新计算点积，而是做增量更新：比如插入新向量时，计算它和所有现有向量的点积并插入到点积表中。
• 给点积表的id1和id2字段建索引，后续查询相似度时能更快定位。

内容的提问来源于stack exchange，提问作者Drxxd