如何使用pgbench结合自定义函数对pgvector向量表进行压测？

错误原因 #

pgbench的\set命令是客户端侧的变量赋值逻辑，仅支持pgbench内置函数（如random()、random_gaussian()）或直接常量赋值，无法直接调用PostgreSQL数据库中定义的自定义函数，这就是你遇到语法错误的原因。

解决方案 #

方案1：使用pgbench的\gset从SQL查询结果赋值变量 #

通过SELECT ... \gset将自定义函数的执行结果赋值给pgbench变量，确保后续SQL使用同一组变量值：

修改simple-read-write.sql如下：

-- 获取随机ID并赋值给:id变量
SELECT random_between(1, 100000) AS id \gset
-- 获取随机vector并赋值给:v变量
SELECT random_array(1536)::VECTOR(1536) AS v \gset

BEGIN;
-- UPDATE（修正原脚本表名错误：vtet → vtest）
UPDATE vtest SET v = :v WHERE id = :id;
-- SELECT
SELECT v FROM vtest WHERE id = :id;
-- INSERT
INSERT INTO vtest (id,v,created_at) VALUES (:id, :v, CURRENT_TIMESTAMP);
END;

方案2：在SQL事务块中直接调用自定义函数（或用PL/pgSQL统一变量） #

如果希望所有操作使用同一组随机值，可在PL/pgSQL块中定义变量并调用函数，避免pgbench侧变量赋值：

BEGIN;
DECLARE
  target_id INT := random_between(1, 100000);
  target_v VECTOR(1536) := random_array(1536)::VECTOR(1536);
BEGIN
  UPDATE vtest SET v = target_v WHERE id = target_id;
  SELECT v FROM vtest WHERE id = target_id;
  INSERT INTO vtest (id,v,created_at) VALUES (target_id, target_v, CURRENT_TIMESTAMP);
END;
END;

如果允许各操作使用不同随机值，可简化为直接在SQL语句中调用函数：

BEGIN;
UPDATE vtest SET v = random_array(1536)::VECTOR(1536) WHERE id = random_between(1, 100000);
SELECT v FROM vtest WHERE id = random_between(1, 100000);
INSERT INTO vtest (id,v,created_at) VALUES (random_between(1, 100000), random_array(1536)::VECTOR(1536), CURRENT_TIMESTAMP);
END;

前置检查 #

1. 确保random_between和random_array函数已在pgvector数据库中创建完成
2. 确保vtest表存在，且提前插入一定量的测试数据（否则UPDATE操作可能无匹配行）
3. 修正原脚本中的表名错误：将SELECT v FROM vtet改为SELECT v FROM vtest

执行压测 #

保持原pgbench命令不变，直接运行即可：

/usr/pgsql-14/bin/pgbench -f ./simple-read-write.sql -c2 -T10 -P1 -U postgres -p 5432 -d pgvector

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