PostgreSQL でのベクトル検索

PostgreSQL でベクトル検索を行う方法のメモ。

テーブル作成・データ登録

postgres ユーザで test_vector1 DB で vector を利用できるようにします。

$ psql -h 127.0.0.1 -U postgres
# create databse test_vector1;
# \connect test_vector1
# create extension vector;

テーブルを作成し、ベクトルの次元の制約を追加

# create table vector1 (
    id      integer not null,
    name    varchar(256) not null,
    vec     vector(2),
    primary key (id)
);

# alter table vector1 add check (vector_dims(vec::vector) = 2);

データを登録

# insert into vector1 (id, name, vec) values (0, 'name_000', array[0.0, 0.0]);
# insert into vector1 (id, name, vec) values (1, 'name_001', array[0.0, 1.0]);
# insert into vector1 (id, name, vec) values (2, 'name_002', array[1.0, 0.0]);
# insert into vector1 (id, name, vec) values (3, 'name_003', array[2.0, 1.0]);
# insert into vector1 (id, name, vec) values (4, 'name_004', array[1.0, 2.0]);
# insert into vector1 (id, name, vec) values (5, 'name_005', array[2.0, 2.0]);
# insert into vector1 (id, name, vec) values (6, 'name_006', array[3.0, 1.0]);
# insert into vector1 (id, name, vec) values (7, 'name_007', array[3.0, 2.0]);
# insert into vector1 (id, name, vec) values (8, 'name_008', array[1.0, 3.0]);
# insert into vector1 (id, name, vec) values (9, 'name_009', array[2.0, 3.0]);

検索

<->: 距離での検索

# select * from vector1 order by vec <-> '[1.5, 1.5]' limit 3;

 id |   name   |  vec
----+----------+-------
  3 | name_003 | [2,1]
  4 | name_004 | [1,2]
  5 | name_005 | [2,2]

<=>: コサインでの検索

#select * from vector1 order by vec <=> '[1.5, 1.5]' limit 3;

 id |   name   |  vec
----+----------+-------
  5 | name_005 | [2,2]
  7 | name_007 | [3,2]
  9 | name_009 | [2,3]

<#>: 内積 * -1での検索

select * from vector1 order by vec <#> '[1.5, 1.5]' limit 3;

 id |   name   |  vec
----+----------+-------
  7 | name_007 | [3,2]
  9 | name_009 | [2,3]
  6 | name_006 | [3,1]

条件付き検索

# select * from vector1 where id % 2 = 0 order by vec <-> '[1.5, 1.5]' limit 3;

 id |   name   |  vec
----+----------+-------
  4 | name_004 | [1,2]
  2 | name_002 | [1,0]
  6 | name_006 | [3,1]