Kishore Uppala
Kishore Uppala

Reputation: 349

pgvector stress test using pgbench with custom function or custom script

I am trying to run a stress test on a vector data type table. But not able to run pgbench with custom functions: random_between and random_array. Kindly suggest how to perform the stress test with custom functions.

random_between:

   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;

random_array:

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;

Table "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 command:

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

Error:

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

Expecting to run pgbench with custom functions and custom sql script.

Upvotes: 1

Views: 86

Answers (1)

Vladlen
Vladlen

Reputation: 138

The metacommand \set has the same syntax in psql and pgbench, it is easy mix them. Your commands \set id random_between(1, 100000) and \set v random_array(1536)::VECTOR(1536) are written for psql. pgbench has own command processor and can not call user defined functions in \set metacommand, it has the list of predefined metafunctions, like random(), sqrt(), double(), approximately 10 functions. You can look at the list in the documentation page https://postgrespro.com/docs/postgrespro/10/pgbench .

Anyway, you pgbench script has SQL commands and you can call your functions from SQL code directly. psql variables can be emulated in SQL code using WITH clause. For example , your code can be represented as:

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;

Upvotes: 0

Related Questions