Практический пример использования pgvector

9 октября 2023

PGVector — это открытое (лицензия MIT) расширение PostgreSQL, решающее задачу поиска схожих векторов. Что еще за вектора такие, и зачем кому-то искать среди них похожие? Попробуем разобраться на конкретном примере.

Пусть имеется интернет-магазин ноутбуков. Пользователь открывает страницу с заданным ноутбуком. Задача — показать помимо самого ноутбука список рекомендуемых похожих ноутбуков.

Как определить, что два объекта (товара, картинки, аудио- или видео-файла, …) являются похожими? Практическое решение заключается в том, чтобы научиться выделять из объектов вектор признаков. Первому объекту будет соответствовать какой-то вектор (или точка) [1.5, 0, 42, 8], а второму — [2.1, -0.9, 314, 0]. Тогда похожесть объектов можно определить, как евклидово расстояние между точками, или как косинус угла между векторами, в зависимости от того, что лучше работает в конкретном случае. Соответственно, задача поиска похожих объектов сводится к поиску близких векторов в N-мерном пространстве.

Поиск схожих векторов pgvector берет на себя. От нас же требуется извлекать относительно осмысленные вектора.

Расширение устанавливается так:

git clone https://github.com/pgvector/pgvector.git
cd pgvector
git checkout v0.5.0
make
make install
make installcheck

Подсоединяемся к базе данных и говорим:

CREATE EXTENSION vector;

Добавим таблицу с ноутбуками:

CREATE TABLE laptops(
id SERIAL PRIMARY KEY,
manufacturer TEXT NOT NULL,
cpu TEXT NOT NULL,
harddrive_type TEXT NOT NULL,
harddrive_space TEXT NOT NULL,
ram_type TEXT NOT NULL,
ram TEXT NOT NULL,
wifi TEXT NOT NULL,
bluetooth TEXT NOT NULL,
ethernet TEXT NOT NULL,
webcam TEXT NOT NULL,
cardreader TEXT NOT NULL,
graphics TEXT NOT NULL,
display TEXT NOT NULL,
usb TEXT NOT NULL,
battery TEXT NOT NULL,
embedding VECTOR(50) NOT NULL
);

Все столбцы с полезной нагрузкой имеют тип TEXT для простоты. В реальных проектах не надо так делать. Столбец embedding хранит вектор признаков, по которому мы будем искать. Признаки было решено извлекать очень простым способом. Если ноутбук произведен производителем A, тогда первое значение в векторе равно 1, иначе оно равно 0. Если ноутбук произведен производителем Б, тогда второе значение равно 1, иначе оно равно 0. И так далее для N известных производителей. Затем, если используется процессор x86, то (N+1)-ое значение равно 1, иначе оно равно 0. Если используется ARM, (N+2)-ое значение равно 1, иначе оно равно 0. Продолжаем в том же духе для типа жесткого диска, наличия или отсутствия Bluetooth, и так далее. Надеюсь, идея ясна.

PGVector поддерживает два типа индексов — IVFFlat и HNSW. Первый индекс быстрее строится и потребляет меньше памяти, но медленнее ищет. Второй тип строится медленнее и потребляет больше памяти, но быстрее ищет. Запросов на чтение обычно больше, чем запросов на запись, поэтому воспользуемся HNSW:

CREATE INDEX ON laptops USING hnsw (embedding vector_l2_ops);

Здесь vector_l2_ops говорит о том, что используется евклидово расстояние. В PostgreSQL это называется класс операторов, или opclass.

Для заполнения таблицы случайными данными был написан небольшой скрипт на Python. Вставка 1000 строк на моем железе занимает порядка 5-7 секунд и замедляется по мере роста индекса. Для сравнения, без HNSW индекса та же операция занимает 80 мс. Данный бенчмарк непоказателен, поскольку я никак не оптимизировал СУБД, и вообще использовал debug сборки с включенными Assert’ами. Но стоит быть морально готовым к тому, что обновление HNSW — операция небыстрая. Может потребоваться добавлять данные фоново и/или во время минимальной нагрузки. Вставка 1 млн строк транзакциями по 1000 строк заняла 2 часа.

Попробуем найти ноутбук, максимально похожий на ноутбук с id равным 123:

SELECT * FROM laptops WHERE id != 123
ORDER BY embedding <-> (
SELECT embedding FROM laptops WHERE id = 123
) LIMIT 1;

Результат:

Запрос отрабатывает за ~8 мс. Если с LIMIT 10, то примерно за столько же.

Если без индекса:

BEGIN;
DROP INDEX laptops_embedding_idx;
-- повторяем запрос
ROLLBACK;

… то получается 1300 мс. Ускорение в 150+ раз, неплохо.

А ноутбуки точно похожие? Проверим:

SELECT * FROM laptops WHERE id = 123;

Результат:

Ноутбуки имеют разных производителей, а во всем остальном они одинаковые. Для ноутбука с id = 456 был найден такой же, но с большим жестким диском. У вас результаты будут другие, поскольку данные случайные. Но в общем и целом, выглядит так, как если бы все работало.

Больше информации о pgvector — в официальной документации. К моменту, когда вы будете читать эти строки, наверняка будет доступна более новая версия расширения, с новыми возможностями, улучшениями производительности, исправлениями ошибок, и так далее.

Метки: PostgreSQL, СУБД.

Вы можете прислать свой комментарий мне на почту, или воспользоваться комментариями в Telegram-группе.

Записки программиста

Практический пример использования pgvector

Коротко о себе

Основные рубрики