pgbench自定义压测脚本如何选取随机UUID主键且不计入延迟

问题背景 #

• 压测场景：基于pgbench对自定义数据集开展基准测试，目标待查询记录均以UUID作为主键
• 现有方案局限：官方及公开示例均通过random()函数生成连续整型值匹配顺序主键，无法直接适配UUID主键场景
  默认pgbench测试脚本示例如下：

\set bid random(1, 1 * :scale)
\set tid random(1, 10 * :scale)
\set delta random(-5000, 5000)

BEGIN;
UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;
SELECT abalance FROM pgbench_accounts WHERE aid = :aid;
UPDATE pgbench_tellers SET tbalance = tbalance + :delta WHERE tid = :tid;
UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid;
INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);
END;

• 核心需求：实现从目标表选取随机UUID存入pgbench变量，且该选取操作的耗时不计入压测延迟统计。

实现原理 #

pgbench默认仅统计BEGIN;到END;事务块内部的操作耗时，所有事务块外的变量赋值、外部命令调用、预查询逻辑均属于压测准备动作，耗时不会被纳入最终的延迟、TPS统计结果，不需要额外配置计时排除规则。

可选实现方案 #

方案1：预导出UUID列表随机抽取（适合千万级以内数据量场景） #

1. 压测前将目标表的所有UUID主键导出为纯文本文件，每行存储1个UUID
2. 编写压测脚本时，在事务块外通过shell命令随机抽取文件中的UUID赋值给变量，示例脚本：

-- 以下为事务块外逻辑，耗时不计入压测统计
\set aid `shuf -n 1 /data/pgbench/uuid_list.txt`
\set delta random(-5000, 5000)

-- 以下为事务块内压测逻辑，耗时全额计入统计
BEGIN;
UPDATE your_accounts SET abalance = abalance + :delta WHERE aid = ':aid'::uuid;
SELECT abalance FROM your_accounts WHERE aid = ':aid'::uuid;
-- 可按需替换为自定义压测SQL
INSERT INTO your_history (aid, delta, mtime) VALUES (':aid'::uuid, :delta, CURRENT_TIMESTAMP);
END;

注意事项：

• UUID为字符串类型，传入SQL时需要包裹单引号并显式转换为uuid类型，避免类型不匹配报错
• 该方案随机选取的记录完全均匀，只要UUID列表文件存储在本地SSD，shuf抽取单行的开销在微秒级，不会干扰压测节奏
• 若表数据量超过1千万，单文件体积过大会导致shuf扫描开销上升，不建议使用该方案

方案2：块采样实时查询随机UUID（适合亿级以上大数据量场景） #

数据量过大不适合全量导出UUID时，可在事务块外通过调用psql执行块级采样SQL获取随机UUID，这部分调用耗时不会计入压测统计，示例脚本：

-- 事务块外执行，耗时不计入统计
\set aid `psql -d your_db -U your_user -t -c "SELECT aid FROM your_accounts TABLESAMPLE SYSTEM (0.001) LIMIT 1;" | tr -d '[:space:]'`
\set delta random(-5000, 5000)

-- 事务块内压测逻辑
BEGIN;
UPDATE your_accounts SET abalance = abalance + :delta WHERE aid = ':aid'::uuid;
SELECT abalance FROM your_accounts WHERE aid = ':aid'::uuid;
INSERT INTO your_history (aid, delta, mtime) VALUES (':aid'::uuid, :delta, CURRENT_TIMESTAMP);
END;

注意事项：

• 禁止使用ORDER BY random()方式取随机记录，该写法会触发全表扫描性能极差；TABLESAMPLE SYSTEM为块级采样，亿级表也可在毫秒级返回结果
• 采样百分比参数（示例中为0.001）可根据表的实际总数据量调整，保证每次采样至少返回1条记录即可
• 调用psql时使用和pgbench相同的连接参数，避免出现权限不足、连接失败问题

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