На грани ИИ: пример поиска и обработки векторов в PostgreSQL + pgvector

На Хабре было много упоминаний pgvector в обзорах Postgresso. И каждый раз новость была про место которое где‑то за границей и далеко. Многие коммерческие решения для хранения и поиска векторов в базе данных нынче не доступны, а pgvector доступен любому, тем более в самой популярной базе в России. Применим pgvector для задачи поиска похожих домов по инфраструктуре для детей в Москве.

В этой статье покажу на этом практическом примере как хранить, кластеризовать алгоритмом DBSCANвекторы и искать по ним в базе данных. В примере задача с векторами на грани типичного хранения и обработки результатов работы нейросетевых моделей в базе данных.

Прежде всего надо установить pgvector в PostgreSQL, он доступен в виде расширения. Поскольку я работаю с базой данных из Docker, то могу просто добавить в Dockerfile строчки и пересобрать образ:

RUN git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git RUN cd pgvector && make && make install

А в самой базе данных, нужно загрузить расширение:

osmworld=# CREATE EXTENSION vector; CREATE EXTENSION Time: 32,606 ms

Данные для векторов можно получить, например, из модели машинного обучения в python скрипте или ML модели в spark и вставить в таблицу с колонкой типа vector. А можно создать в SQL как гистограмму определенных категорий. В этом случае можно значения в массивах real[], integer[] или double precision[], numeric[] привести к типу ::vector.

Данными для примера послужат гистограммы числа объектов детской инфраструктуры в окрестностях жилых домов в Москве. Про то как рассчитать эти данные я рассказывал здесь, а в этой публикации я просто возьму готовые данные и создам из них таблицу с колонкой типа одинадцатимерный vector:

create table infrastructure_for_children_features as    select (row_number() over ())::integer id, null::integer cluster,           district, street, housenumber,           ARRAY[kindergarten::integer, school::integer,college::integer, university::integer, language_school::integer, music_school::integer,training::integer,sports_centre::integer,community_centre::integer,playground::integer,clinic::integer]    ::vector(11) feature  from infrastructure_for_children; alter table infrastructure_for_children_features add primary key(id);

Так в базе создал таблицу на 30237 записей со структурой:

osmworld=# \d infrastructure_for_children_features    Table "public.infrastructure_for_children_features"    Column    |    Type    |  -------------+------------|  id          | integer    |   cluster     | integer    |   district    | text       |   street      | text       |   housenumber | text       |   feature     | vector(11) | Indexes:     "infrastructure_for_children_features_pkey" PRIMARY KEY, btree (id)

Теперь хотелось бы объединить их в группы по близости векторов. Опять же можно использовать нейросети, а можно использовать классические алгоритмы кластеризации — метод k‑средних(k‑means) или основанную на плотности пространственную кластеризацию для приложений с шумами (DBSCAN). Для метрики близости использую Евклидово расстояние между векторами из колонки feature.

Поскольку число кластеров мне не известно, то я выберу метод DBSCAN и прогоню этот набор данных через него чтобы посмотреть зависимость от epsilon числа групп и число элементов не попавших в группы:

eps|clusters|not_in_cluster 0.07529667 0.57529667 1.020228648 1.547526904 2.092822630 2.5122717620 3.0117311778 3.58567605 4.06014562 4.53642760 5.02321574 5.5138972 6.077604 6.551377 7.029265 7.514168 8.01096 8.5454 9.0437 9.5230

На свой субъективный взгляд выберу eps=5.5 и запущу Java программу, которая заполнит колонку cluster значениями алгоритма DBSCAN для minPoints=3 и eps=5.5:

package com.github.isuhorukov;  import com.pgvector.PGvector; import org.apache.commons.math3.ml.clustering.Cluster; import org.apache.commons.math3.ml.clustering.Clusterable; import org.apache.commons.math3.ml.clustering.DBSCANClusterer; import org.apache.commons.math3.ml.distance.EuclideanDistance;  import java.sql.*; import java.util.ArrayList; import java.util.List;  public class Main {     public static void main(String[] args) throws Exception {         try (Connection connection = DriverManager.getConnection(                 System.getenv("jdbc_url"), System.getenv("user"), System.getenv("password"))) {             connection.setAutoCommit(false);             PGvector.addVectorType(connection);             float eps = Float.parseFloat(System.getenv("eps"));             int minPoints = Integer.parseInt(System.getenv("minPoints"));             DBSCANClusterer<Feature> dbscanClusterer = new DBSCANClusterer<>(eps,minPoints,new EuclideanDistance());             List<Feature> features = fetchFeatures(connection,                                                     "select id,feature from infrastructure_for_children_features");             List<Cluster<Feature>> cluster = dbscanClusterer.cluster(features);             saveClusters(connection, cluster);         }     }      private static void saveClusters(Connection connection, List<Cluster<Feature>> cluster) throws SQLException {         try (PreparedStatement clusterPs = connection.prepareStatement(              "update infrastructure_for_children_features set cluster = ? where id = ?")){             for (int idx = 0; idx < cluster.size(); idx++) {                 List<Feature> featureCluster = cluster.get(idx).getPoints();                 for (Feature feature : featureCluster) {                     clusterPs.setInt(1, idx);                     clusterPs.setInt(2, feature.id);                     clusterPs.addBatch();                 }                 clusterPs.executeBatch();             }             connection.commit();         } catch (Exception e) {             connection.rollback();             throw new RuntimeException(e);         }     }      private static List<Feature> fetchFeatures(Connection connection, String query) {         List<Feature> features = new ArrayList<>();         try (Statement statement = connection.createStatement();              ResultSet resultSet = statement.executeQuery(query))         {             while (resultSet.next()) {                 int id = resultSet.getInt(1);                 float[] feature = ((PGvector) resultSet.getObject(2)).toArray();                 features.add(new Feature(id, feature));             }         } catch (Exception e) {             throw new RuntimeException(e);         }         return features;     }      static class Feature implements Clusterable {         public int id;         public double[] feature;         public Feature(int id, float[] feature) {             this.id = id;             this.feature = new double[feature.length];             for (int i = 0; i < feature.length; i++) {                 this.feature[i] = feature[i];             }         }          @Override         public double[] getPoint() {             return feature;         }     } }

Для компиляции программы нужен pom.xml для maven:

<?xml version="1.0" encoding="UTF-8"?> <project xmlns="http://maven.apache.org/POM/4.0.0"          xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"          xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">     <modelVersion>4.0.0</modelVersion>      <groupId>com.github.igor-suhorukov</groupId>     <artifactId>vectors</artifactId>     <version>1.0-SNAPSHOT</version>      <properties>         <maven.compiler.source>11</maven.compiler.source>         <maven.compiler.target>11</maven.compiler.target>         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>     </properties>     <dependencies>         <dependency>             <groupId>org.apache.commons</groupId>             <artifactId>commons-math3</artifactId>             <version>3.6.1</version>         </dependency>         <dependency>             <groupId>com.pgvector</groupId>             <artifactId>pgvector</artifactId>             <version>0.1.3</version>         </dependency>         <dependency>             <groupId>org.postgresql</groupId>             <artifactId>postgresql</artifactId>             <version>42.6.0</version>         </dependency>     </dependencies> </project>

Внутри этой программы происходит следующее:

• Устанавливается соединение с базой данных используя URL соединения, имя пользователя и пароль, полученные из переменных окружения. Отключается режим автоматической фиксации транзакций соединения.

• Регистрируется пользовательский тип вектора PGvector в драйвере данных.

• Вызывается метод fetchFeatures для извлечения списка характеристик из базы данных из созданной нами ранее таблицы infrastructure_for_children_features.

• Значения eps и minPoints извлекаются из переменных окружения и преобразуются в типы float и int соответственно.

• Создается экземпляр DBSCANClusterer из commons-math3 с использованием eps, minPoints и метрики расстояния (EuclideanDistance). Вызывается метод cluster объекта DBSCANClusterer для кластеризации записей, полученных ранее из базы в методеfetchFeatures.

• Вызывается метод saveClusters для сохранения кластеров в базе данных в поле cluster таблицы infrastructure_for_children_features, учитывая номер кластера и id каждой записи из результатов алгоритма кластеризации. Если данные успешно записались, то транзакция фиксируется, в случае ошибки — транзакция откатывается.

Найдем в базе данных дом со скриншота по идентификатору:

osmworld=# select * from infrastructure_for_children_features where id=831;  id  | cluster |     district      |        street         | housenumber |             feature              -----+---------+-------------+--------------------+-------------------+-----------------------+-------------+---------------------------------  831 |       4 | Пресненский район | Большая Бронная улица | 19          | [37,28,3,23,6,4,17,23,9,148,92] (1 row)

А теперь в PostgreSQL найдем 10 ближайших к нему домов по значению вектора:

osmworld=# select id, cluster, district,street,housenumber from infrastructure_for_children_features order by  feature<-> '[37,28,3,23,6,4,17,23,9,148,92]' limit 10;   id  | cluster |     district      |            street            | housenumber  ------+---------+-------------------+------------------------------+-------------   831 |       4 | Пресненский район | Большая Бронная улица        | 19  1011 |       4 | Пресненский район | Большая Бронная улица        | 17   897 |       4 | Пресненский район | Сытинский переулок           | 5/10 с3   827 |       4 | Пресненский район | Богословский переулок        | 8/15  1019 |       4 | Пресненский район | Большая Бронная улица        | 16   823 |       4 | Тверской район    | Малый Палашёвский переулок   | 4   893 |       4 | Пресненский район | Большой Козихинский переулок | 4   631 |       4 | Пресненский район | Большая Бронная улица        | 25 с3   821 |       4 | Пресненский район | Сытинский переулок           | 5/10 с4  1117 |       4 | Пресненский район | Богословский переулок        | 5 (10 rows)  Time: 25,777 ms  osmworld=# explain select id, cluster, district,street,housenumber from infrastructure_for_children_features order by  feature<-> '[37,28,3,23,6,4,17,23,9,148,92]' limit 10;                                                QUERY PLAN                                                --------------------------------------------------------------------------------------------------------  Limit  (cost=2375.37..2375.40 rows=10 width=90)    ->  Sort  (cost=2375.37..2450.97 rows=30237 width=90)          Sort Key: ((feature <-> '[37,28,3,23,6,4,17,23,9,148,92]'::vector))          ->  Seq Scan on infrastructure_for_children_features  (cost=0.00..1721.96 rows=30237 width=90) (4 rows)

То есть, можно найти в городе похожие дома на указанный в запросе, по доступной в окрестностях инфраструктуре.

Можно ли ускорить поиск по векторам? Да, расширение поддерживает индексы IVFFlat и HNSW. Попробуем HNSW он более быстрый и точный, если верить научным публикациям:

osmworld=# CREATE INDEX ON infrastructure_for_children_features USING hnsw (feature vector_l2_ops) WITH (m = 16, ef_construction = 64); CREATE INDEX Time: 12705,611 ms (00:12,706)  osmworld=# \d infrastructure_for_children_features        Table "public.infrastructure_for_children_features"    Column    |       Type       | Collation | Nullable | Default  -------------+------------------+-----------+----------+---------  id          | integer          |           | not null |   cluster     | integer          |           |          |   lon         | double precision |           |          |   lat         | double precision |           |          |   district    | text             |           |          |   street      | text             |           |          |   housenumber | text             |           |          |   feature     | vector(11)       |           |          |  Indexes:     "infrastructure_for_children_features_pkey" PRIMARY KEY, btree (id)     "infrastructure_for_children_features_feature_idx" hnsw (feature vector_l2_ops) WITH (m='16', ef_construction='64')  osmworld=# explain select id, cluster, district,street,housenumber from infrastructure_for_children_features order by  feature<-> '[37,28,3,23,6,4,17,23,9,148,92]' limit 10;                                                                         QUERY PLAN                                                                          -----------------------------------------------------------------------------------------------------------------------------------------------------------  Limit  (cost=5.00..5.61 rows=10 width=90)    ->  Index Scan using infrastructure_for_children_features_feature_idx on infrastructure_for_children_features  (cost=5.00..1861.36 rows=30237 width=90)          Order By: (feature <-> '[37,28,3,23,6,4,17,23,9,148,92]'::vector) (3 rows) Time: 1,022 ms

Запрос стал использовать этот индекс и поиск по векторам стал быстрее на порядок по сравнению с seqscan, и в этом конкретном примере угадал с параметрами для DBSCAN:

osmworld=# select id, cluster, district,street,housenumber from infrastructure_for_children_features order by  feature<-> '[37,28,3,23,6,4,17,23,9,148,92]' limit 10;   id  | cluster |     district      |            street            | housenumber  ------+---------+-------------------+------------------------------+-------------   831 |       4 | Пресненский район | Большая Бронная улица        | 19  1011 |       4 | Пресненский район | Большая Бронная улица        | 17   827 |       4 | Пресненский район | Богословский переулок        | 8/15   897 |       4 | Пресненский район | Сытинский переулок           | 5/10 с3  1019 |       4 | Пресненский район | Большая Бронная улица        | 16   823 |       4 | Тверской район    | Малый Палашёвский переулок   | 4   893 |       4 | Пресненский район | Большой Козихинский переулок | 4   631 |       4 | Пресненский район | Большая Бронная улица        | 25 с3   821 |       4 | Пресненский район | Сытинский переулок           | 5/10 с4  1117 |       4 | Пресненский район | Богословский переулок        | 5 (10 rows)  Time: 1,644 ms

Вывод

Расширение pgvector PostgreSQL оказалось простым в использовании и с ним можно работать не только алгоритмами машинного обучения, но и классическими алгоритмами кластеризации из программы по JDBC, а так же быстро искать используя поиск по близости векторов и специализированный индекс HNSW.