使用带有自定义函数或自定义脚本的 pgbench 进行 pgvector 压力测试

Question

我正在尝试对矢量数据类型表运行压力测试。但无法使用自定义函数运行 pgbench：random_ Between 和 random_array。请建议如何使用自定义函数进行压力测试。

随机_之间：

   CREATE OR REPLACE FUNCTION random_between(low INT ,high INT) 
    RETURNS INT AS
    $$
    BEGIN
    RETURN floor(random()* (high-low + 1) + low);
    END;
    $$ language 'plpgsql' STRICT;

随机数组：

CREATE OR REPLACE FUNCTION random_array(dim integer)
        RETURNS DOUBLE PRECISION[]
    AS $$
        SELECT array_agg(random())
        FROM generate_series(1, dim);
    $$
    LANGUAGE SQL
    VOLATILE
    COST 1;

表“public.vtest”

    +------------+-----------------------------+-----------+----------+---------+
    |   Column   |            Type             | Collation | Nullable | Default |
    +------------+-----------------------------+-----------+----------+---------+
    | id         | bigint                      |           |          |         |
    | v          | vector(1536)                |           |          |         |
    | created_at | timestamp without time zone |           |          | now()   |
    +------------+-----------------------------+-----------+----------+---------+

cat simple-read-write.sql

  \set id random_between(1, 100000)
  \set v random_array(1536)::VECTOR(1536)

    BEGIN;
    --UPDATE 
    UPDATE vtest SET v = :v WHERE id = :id;
    --SELECT 
    SELECT v FROM vtet WHERE id = :id;
    --INSERT 
    INSERT INTO vtest (id,v,created_at) VALUES (:id, :v, CURRENT_TIMESTAMP);
    END;

Pgbench 命令：

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

错误：

\set id random_between(1, 100000)
                     ^ error found here
\set v random_array(1536)::VECTOR(1536)
                         ^ error found here

期望使用自定义函数和自定义 sql 脚本运行 pgbench。

Answer 1

元命令 \set 在 psql 和 pgbench 中具有相同的语法，很容易混合它们。您的命令

\set id random_between(1, 100000)

和

\set v random_array(1536)::VECTOR(1536)

是为 psql 编写的。 pgbench 有自己的命令处理器，不能调用 \set metacommand 中用户定义的函数，它有预定义的元函数列表，如 random()、sqrt()、double() 等大约 10 个函数。您可以查看文档页面中的列表https://postgrespro.com/docs/postgrespro/10/pgbench.

无论如何，你的pgbench脚本有SQL命令，你可以直接从SQL代码调用你的函数。 psql 变量可以使用WITH 子句在SQL 代码中模拟。例如，您的代码可以表示为：

BEGIN;
    WITH var1 AS (select random_between(1, 100000) as id),
    var2 AS (select random_array(1536)::VECTOR(1536) as v),
    update1 AS (
    UPDATE vtest as vtest  SET v = var2.v
     FROM var1 as var1 JOIN var2 as var2 on 1=1
     WHERE vtest.id = var1.id RETURNING vtest.id),
    select1 AS (SELECT v FROM vtest as vtest 
     JOIN var1 as var1 on 1=1 WHERE vtest.id = var1.id)
    INSERT INTO vtest (id,v,created_at) SELECT var1.id, var2.v,
      CURRENT_TIMESTAMP FROM var1 as var1 JOIN var2 as var2 on 1=1;
    END;