Дизайн встраиваемой базы данных для ANN запросов: MusyaDB

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

Статья получилась размером с небольшую книгу и разбивается на две больших статьи. В первой части мы поймем с чего вообще начинается дизайн таких систем, выберем алгоритмы и модель вычислений. Во-второй части мы аккуратно разрежем систему на слои и опишем API их взаимодействия.

На всякий случай я решил сразу дать этой архитектуре название MusyaDB. Таким образом она будет как и Lucene, названа в честь жены. Только в данном случае используется прозвище, потому что моя остроумная идея назвать базу для ANN запросов ANNastasia, была отвергнута.

Постановка задачи

Как только вы сели за ноутбук с мыслью «сейчас я сделаю базу данных, которая всем подойдет», вы уже допустили фатальную ошибку. Эффективность даже такой тривиальной вещи как сортировка зависит от use case-ов (смотрите истории развития сортировок в стандартных библиотеках C++ и Java), а БД, пусть даже встроенная, это целая вселенная trade-off’-ов. Вы должны представить себя пользователем такой системы, понять чего вы хотите: какие use case-ы покрыть, а от каких отказаться. У Майкла Стоунбрейкера есть отличный доклад с весьма точным названием «One Size Fits None». Садиться за дизайн надо с этой мыслью.

Поскольку инициатива на моей стороне, я достану требования из шляпы. Но постараюсь дать некоторую мотивацию о том, почему такие требования вообще имеют смысл и какие для них есть альтернативы.

Чтобы было понятнее, мы оттолкнемся от следующей задачи: написать базу данных, которая поможет в разработке рекомендательной системы гетерогенных (разнородных) документов. Примером такой системы может быть алгоритмическая лента с разными типами документов (статьи, видео, твиты и прочее) или вкладка «рекомендуем для вас» в маркетплейсе, где продают самые разные товары.

То, что система рекомендательная, значит, что мы должны обеспечить быстрый приближенный поиск ближайших соседей (т.н. ANN запросы) по векторному полю. Для простоты такое поле у нас будет одно. Т.е. по заданному вектору в запросе найти в базе приближенный top документов наиболее близких запросу. Метрика расстояния может быть любой, но обычно используют косинус между векторами. Вектора нормализуют и тогда косинус равен скалярному произведению. Расчет расстояние ещё называют скорированием: это более общий термин означает вычисление релевантности документа.

А то, что документы разнородные, для нас, как для разработчиков базы данных, значит, что всего разных полей у документов много, но при этом у каждого отдельного документа заполнена лишь малая их часть. Или говоря в терминах SQL баз данных таблица разреженная, т.е. в ней много NULL-ов.

Это, с одной стороны, актуальный кейс, а, с другой стороны, на рынке не так много хороших решений (например qdrant). Во многом потому, что (как мы увидим дальше) нельзя просто добавить ANN индекс к существующей SQL базе данных, чтобы это работало эффективно. Так что есть пространство для фантазии.

Функционал или производительность

Для начала мы должны понять, как мы будем поступать, когда перед нами возникнет перспектива отказаться от каких-то сценариев использования в пользу быстродействия других. Будем ли мы отсекать сценарии использования или предпочтём поддержать больший функционал ценой эффективности системы. Никакого единственно верного ответа здесь, как и далее, нет: вы вполне можете позиционировать свой продукт как удобный швейцарский нож, который просто работает, и как узкоспециализированную бритву, которой нет равных в своём деле.

Мы пойдем по второму пути, т.к. мне кажется, что обычно разработчикам труднее отказаться от функционала, чем от производительности.

Асимметрия нагрузки между чтением и записью

Для начала мы подумаем о том, как будет выглядеть нагрузка на чтение. Будем строить систему из расчета, что именно нагрузка на чтение является основной. Это не только типичный сценарий, но и потребует от нас сложного представления данных.

Запросы и данные

Первая вещь, с которой нужно определиться: на какого вида запросы система будет ориентирована в первую очередь и на каких данных она будет их обрабатывать. Данными могут быть например графы, плоские таблицы или иерархические документы с произвольной схемой. Запросы могут быть например OLTP, OLAP, на полнотекстовый поиск и т.д.

Разнородные документы можно хранить либо без схемы в виде условных JSON документов, как в MongoDB. Либо в виде обычных плоских таблиц, как в SQL базах данных, только с эффективной поддержкой большого числа разреженных колонок, как например в Clickhouse или в Apache Lucene.

Я предлагаю взять второй вариант, т.к. его реализация будет эффективнее за счёт строгой схемы, если конечно полей меньше чем документов.

Проблемы, которые создают ANN запросы

Поиск ближайших соседей в пространствах большой размерности, даже приближенный, вычислительно очень тяжелая задача. На хабре есть обзорная статья разных алгоритмов. И чуть более старый доклад от разработчика Яндекс картинок на ту же тему. Так по моему опыту самым популярным решением является HNSW. Хотя по бенчмаркам самой быстрой считается библиотека ScaNN.

Если у вас есть вычислительно тяжелый индекс, такой как ANN или, скажем, полнотекстовый, то поиск в нём, скорее всего, будет горячим местом всего приложения. Но с такими индексами есть ещё одна проблема: они не updatable. Т.е. в такой индекс невозможно за небольшое число операций добавить документы, не деградировав при этом по чтению, в отличие от например B-деревьев.

Тяжелый индекс неизбежно будет набором иммутабельных (неизменяемых) файлов, а изменения будут происходить путем выпуска новых файлов на замену старым. При попытке обобщить эту схему получится что-то вроде LSM-деревьев. Похожим образом устроен Apache Lucene (Java-библиотека для построения поисковых движков), потому что её основной сценарий использования это полнотекстовый поиск. Хотя поиск по термам работает через обратный индекс, подсчет BM25 скоров документов требует вычисления тяжелых статистик. На всякий случай уточню, что никаких LSM-деревьев в Lucene нет, т.к. LSM-деревья это просто способ организовать индекс по ключу, аналог B-деревьев. Но идея использования иммутабельных файлов общая.

Ниже упрощенная схемы работы LSM и Lucene для наглядности:

Давайте сразу, немного забегая вперед, условимся о следующих обозначениях:

Слово таблица будет иметь тот же смысл, что и в SQL: набор документов с множеством полей.
Сегментами мы будем называть независимые части таблицы. В контексте LSM деревьев похожие вещи называют SSTable, но у нас будет другой формат хранения, поэтому я решил воспользоваться термином сегмент из Lucene. Один сегмент состоит из нескольких файлов. Все они будут иммутабельными, даже битсет живых документов т.н. liv файл. Об этом чуть ниже.
DocId — это порядковый номер документа в сегменте.

В этой статье мы будем подрезать много идей из Apache Lucene, т.к. эта библиотека тоже позволяет находить top релевантных разнородных документов. Однако Lucene используется для текстового поиска. И, несмотря на то, что и текстовый и ANN поиск скорируют документы, возвращая top, они существенно различаются по своей модели вычислений.

Полнотекстовый поиск с помощью обратного индекса сперва получает итератор по документам, содержащим необходимые термы (причём можно тривиально добиться того, чтобы итератор возвращал документы в порядке, в котором они лежат в сегменте), а затем уже производит вычисление скоров каждого подходящего документа, накапливая результаты. Есть оптимизация, позволяющая пропускать документы, у которых будет заведомо маленький скор, но в целом скорируются все документы, подходящие под запрос, по порядку нахождения в сегменте. В конце результаты разных сегментов объединяются вместе. Подробнее можно прочитать в этой статье.

Для таких вычислений подходит модель map-reduce (стадия map в Lucene называется collect).

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

К сожалению ANN индексы устроены иначе. Например HNSW пытается быстро попасть в область близких к искомому документов и потом начинает обходить её в ширину. Тут стоит оговориться, что речь идет о плотных векторах, т.к. приближенный поиск по разряженным векторам можно свести к поиску через обратный индекс, рассматривая отдельные координаты.

Так происходит потому, что умножение длинных вещественных векторов стоит существенно дороже, чем расчет BM25 скоров, из-за чего хочется просмотреть как можно меньше кандидатов.

Однако неясно как это совместить с обратными индексами, которые эффективно просматривают документы в порядке их расположения. Причем обратные индексы по другим полям нужны, т.к. позволяют существенно ускорить поиск если обрабатывать сразу нужное подпространство.

На первый взгляд может показаться, что векторный запрос и фильтрации должны быть сильно скоррелированы, однако вот вам контрпример. Представьте, что вашей рекомендательной системой пользуется подросток, которому контент 18+ интересен, но он ему запрещен, тогда мы получим такую ситуацию:

Позже мы обсудим как можно выкрутиться из этих трудностей.

Нагрузка на запись

Выбор структуры данных для разных сценариев нагрузки

LSM эффективнее на запись, чем B-деревья поэтому мы тоже будем писать эффективно. Однако если операция слияния перестает успевать поддерживать число sstable-ов маленьким, это приводит к существенному замедлению операций чтения. Подобная проблема есть и у Lucene индексов при росте числа сегментов. К сожалению все эффективные ANN индексы не updatable, поэтому с этим нам придется смириться. Это фундаментальная проблема тяжелых индексов, обойти которую мы не сможем.

Однако можно сгладить её для особого случая: когда чаще всего происходят update-ы полей, для которых нет тяжелых индексов. Тогда мы можем хранить их отдельно в updatable структуре.

У алгоритмического фида быстро меняющимся параметром может быть счетчик лайков, а у маркетплейса — цена и доступность товара.

Если вы хотите сделать production-ready систему, то об этом стоит подумать заранее.

Мы же пока обяжем пользователя вставлять документы большими батчами.

Соответственно, если кому-то такое неудобно: придется поставить перед системой очередь, например Kafka, которая будет накапливать документы и формировать батчи.

Throughput или latency

В дихотомии throughput / latency мы будем пытаться попасть примерно в золотую середину: в случае достаточного числа ресурсов иметь возможность их утилизировать снижая latency, и, наоборот, поддерживать большое число параллельных запросов в ущерб таймингам.

Чтобы не было иллюзии, будто это единственно возможный выбор, стоит обсудить альтернативы.

Скажем, для оптимизации throughput-а можно набирать запросы в батч, например по таймеру. Затем пробегаться по данным один раз собирая значения для всех запросов за один проход и выполнять все запросы одновременно. Таким образом можно добиться практически сколь угодно высокой пропускной способности, лишь бы памяти под запросы хватало. Однако latency при этом тоже будет огромным. Такое имеет смысл только для оффлайн обработки действительно огромного числа запросов, потому что обычно проблему throughput-а решают горизонтальным масштабированием.

Если же пытаться минимизировать latency, то нужно наоборот сторониться всякой батчёвости, блоковости и всего, что с этим связанно, скажем сжатия. Это создаёт большие накладные расходы на широкие запросы, но зато позволят быстрее доставать документы из узких запросов. Такой подход больше подходит под OLTP нагрузку.

Строчное хранение или колоночное

Документы можно хранить подряд и тогда это будет строчное хранение, а можно хранить отдельно каждую колонку и это будет колоночное (столбцовое) хранение.

Колоночное хранение лучше для OLAP нагрузки:

Лучше сжимаются данные.
Можно делать разреженные таблицы с тысячами полей, для хранения разнородных сущностей вместе.
Эффективно обрабатываются широкие запросы, где требуется заглянуть лишь в небольшое подмножество полей.

Строчное хранение лучше для OLTP:

Лучше для транзакций.
Лучше для узких запросов по первичному ключу.

В нашем же случае выбор не столь тривиален, т.к. наша нагрузка не OLTP, но и OLAP её назвать сложно.

В каждом запросе наши колонки делятся на разные виды:

Колонка с векторами. Грубо говоря участвует в ORDER BY секции.
Колонки по которым мы фильтруем. Условно, участвующие в WHERE секции.
Колонки, значения из которых мы возвращаем. Участвующие в SELECT секции.

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

Поля по которым фильтрация идёт без индексов (например по сложным или низкоселективным выражениям) лучше хранить отдельно от полей, которые участвуют в SELECT, что бы сократить чтения.

В то же время, поля, использующиеся вместе, лучше и хранить вместе.

Поэтому, видимо, оптимальной будет гибридная схема, где колонки представляют из себя плоские кортежи (возможно из одного элемента).

Соответственно, пользователь может сам решить, какие поля используются вместе и лучше хранятся вместе, а какие — используются и хранятся отдельно.

Поддерживать произвольно вложенные документы мы не будем, т.к. такого рода структуры в любом случае эффективно обрабатывать не получится.

Помещаются ли данные в память?

Может показаться что БД — это всегда про ситуацию, когда данные не помещаются в память, однако это не всегда так.

Ходить по тяжелым индексам на диске — гиблое дело: даже с быстрым случайным доступом приличное решение получить сложно.

Поэтому мы будем отталкиваться от того, что диск нужен для быстрых рестартов, данные mmap-ятся, и далее обращений к диску уже нет.

Нужно отметить, что поскольку CPU намного быстрее обрабатывает данные, чем они читаются из RAM, то алгоритм быстро работающий с файлом на диске скорее всего будет быстро работать и с данными в RAM. Поэтому часто удается добиться win-win решения.

При этом запись на диск нам обязательно нужна для durability.

Потенциальное использование в распределенной системе

Система, которую вы дизайнете, не только состоит из слоёв, которые вы комбинируете, но и сама может быть частью какой-то более сложной системы, которую строят другие.

Поэтому мы учтём возможность того, что поверх нашего хранилища может быть написано хранилище распределённое. Подобно тому как Elasticsearch написан поверх Apache Lucene.

Вообще построение распределенной БД поверх локальной, это не самый лучший способ построения распределенных БД. Если вы хотите узнать самый лучший способ читайте про YDB и аналогично устроенные системы. Но если наша цель реплицироваться на несколько тысяч машин и десяток шардов, то этот вариант вполне подходит.

Учтите однако, что шардировать ANN индекс очень сложно, и, насколько мне известно, в реальных системах так никто не делает, поэтому запрос всегда идет во все шарды.

Репликация

Запись в тяжелые индексы, как и поиск по ним, упирается главным образом в CPU и скорость памяти. Поэтому первый шаг на пути построения эффективной распределенной системы — это физическое разделение использования Write API и Read API, потому что они конкурируют за одни и те же ресурсы.

Это значит, что мы должны поддержать такой сценарий: на одной машине через Write API пишут таблицу на диск, потом тем или иным способом копируют дельту на другую машину, а та, в свою очередь, переключается на новую версию таблицы, желательно без существенной паузы в запросах.

Что бы не возникло конфликтов надо писать исключительно иммутабельные данные на диск. В нашей схеме надо ещё хранить файл живых документов, который на первый взгляд кажется естественным мутировать. Однако он очень маленький и можно просто писать новую его версию под другим именем. Тогда версия индекса будет просто списком нужных файлов.

Мы говорили ещё про updatable структуру для часто обновляемых данных. Забегая вперед, скажу, что мы не будем сейчас это детально прорабатывать, так как и без того работы по дизайну набежало порядочно. Локально такие файлы легко поддерживать хорошо известными алгоритмами, например B+деревьями. Однако, с эффективной репликацией возникает масса трудностей.

Шардирование

Если таблица ещё помещается на одну машину, но одна машина не справляется с потоком записей, можно собирать различные части таблицы на разных машинах, а потом скачивать их вместе во время репликации. Поскольку таблица в любом случае состоит из разных сегментов, то это в сущности ничего не меняет.

Если же таблица не умещается на одну машину, то требуется также ввести шарды именно на чтение (шардов на чтение и на запись может быть разное количество, т.к. это вообще говоря разные сервисы).

В таком случае пользователь будет вынужден ходить в несколько шардов при чтении, получать сырые данные, которые как-то блендить. Это, на самом деле, только упрощает нам задачу: мы можем просто сказать, что у нас нет высокоуровнего языка запросов, а есть только программное API для работы с данными, и пользователь может получать как сырые, так и не сырые данные на своё усмотрение.

Выбор фич для минимального прототипа

Разумеется embedded storage это задача, которую невозможно доделать, и фичи можно придумывать бесконечно.

Поэтому надо будет выбрать минимальный их набор, который позволит нам показать, что вся эта конструкция собирается во что-то полезное.

Поэтому, вместо того, чтобы реализовывать десятки методов для хранения целых чисел разной битности и знаковости, мы остановимся на следующем наборе.

Допустимые типы в схеме:

vector: нормализованные вещественные вектора фиксированной размерности
int64: знаковое 64 битное целое, в некоторых языках его называют long.
keyword: строки небольшой кардинальности.
blob: массивы байт

В int64 можно паковать любые примитивы почти во всех языках: для целей POC этого достаточно.
Хотя в будущем будет важно различать целые и вещественные числа, т.к. для них эффективны разные алгоритмы сжатия.

Массив байт нужен для не повторяющихся строк, к тому же это пример объекта переменной длины. Индексов для него не будет.

Keyword — нужны для повторяющихся строк, на их примере мы рассмотрим одну занимательную оптимизацию.

Для int64 и keyword будет поддержка обратного индекса и поиска диапазонов.

Для вещественных векторов фиксированной размерности будет доступен ANN индекс, который позволит искать top векторов наиболее близких к заданному.

Языка запросов у нас не будет — это задача какой-то системы сверху нашей. Весь наш API будет программным. Мы должны предоставлять способ фильтрации по индексам, доставать поля и исполнять на них произвольный код. К деталям мы вернёмся значительно позже.

Т.к. мы поддерживаем кортежи, то для простоты будем пока считать что индексы доступны только для кортежей размера 1.

Поиск

Наконец мы подошли к возможно самой интересной части статьи, где я расскажу как предлагаю решить проблему объединения ANN-индекса и обратных, чтобы при этом вычисления оставались эффективными.

Ключевая идея

У нас есть, казалось бы, фундаментальное противоречие:

Скорирование документов очень дорогое.
Поэтому мы хотим делать эту операцию как можно реже.
Поэтому нам надо обходить документы в порядке максимально приближённом к идеальному (т.е. такому, который бы возник, если бы мы отсортировали все документы по скору с запросом).
А так как все запросы разные, этот порядок для всех запросов тоже разный.
Следовательно не существует одной сортировки для документов такой, чтобы мы могли эффективно по ним ходить последовательно.
По этой же причине невозможно лениво обходить обратный индекс в порядке, в котором документы лежат физически.

Однако после изучения алгоритмов ScaNN и FAISS мне в голову пришел один трюк, и со временем я даже понял, как его можно элегантно объяснить.

Допустим, мы заранее знаем — размер топа документов, который мы будем возвращать. Допустим, у нас есть какой-то ANN индекс на n > k документов, и нас устраивает некоторая ошибка $\epsilon$ — процент документов в собранном top , которые на самом деле не вошли бы в top , если бы мы честно обошли вообще все документы (можно выбрать любую вычислимую метрику ошибки, на рассуждение это не повлияет). Очевидно существует функция $C(k, n, \epsilon)$ , которая говорит, сколько нам надо просмотреть кандидатов через наш индекс, чтобы на индексе размера сформировать top документов с погрешностью не больше $\epsilon$ . Скажем для идеального индекса $C(k, n, \epsilon) = k$ , а для худшего случая, когда индекса вообще нет и мы сканируем документы подряд $C(k, n, \epsilon) = n - e \cdot k$ . И для любого реального индекса функция будет принимать значения внутри этого диапазона. Нам важно, что для любого конкретного набора документов и конкретного индекса функция существует и ограничена. Её даже можно эмпирически построить, правда формально для этого потребуется перебрать все возможные запросы, но строить мы её всё равно не будем, нас здесь интересует только её существование.

А теперь внимание: нам неважно в каком порядке обходить эти документов. И поэтому мы можем их обходить в порядке физического расположения.

Т.е. вообще порядок обхода документов нам важен, если мы ищем какие-то произвольные top-ы, как очень маленькие, так и очень большие. Или если у нас есть лимит по времени, а ищем мы сколько получится: если есть свободные ресурсы просто минимизируем ошибку до нуля, а если нет то уж как получится. Но если мы заранее задали наши таргеты по ошибке, времени ответа, размер топа и число документов в индексе, хотя бы приблизительно, то мы можем так же приблизительно установить число кандидатов для просмотра. На практике это делается эмпирически в обратную сторону: крутят ручку числа кандидатов через АБ эксперимент, пока не достигнут максимума бизнесметрик, либо пока бюджет на железо не кончится. Так же нам важен порядок в том смысле, что хочется иметь как можно меньше. Но мы могли бы иметь индекс, в котором сперва как-то быстро находятся эти документов без скорирования, а затем мы обходим их один за другим. А это уже очень похоже на обратный индекс.

Да, в отличие от чистого обратного индекса, для которого можно конструировать ответ полностью лениво, здесь будет шаг, требующий предварительно выбрать кандидатов. Но это всё равно лучше, чем блуждать по индексу в случайном порядке.

Более того, идею можно докрутить. Если произвести кластеризацию документов, можно на первом шаге набирать кластера в порядке близости их центрального элемента, пока в сумме не наберётся кандидатов по ним, а потом обходить. Тут остается вопрос, надо ли ещё выбирать порядок обхода кластеров или лучше их обходить как они лежат на диске. Но главное, что в любом случае внутри кластера мы сканируем документы последовательно, а значит можем эффективно применять обратный индекс. Для удобства можно делать скиплисты на начало кластеров для длинных постинг листов.

Такая структура данных называется IVF (inverted file) в FAISS. Алгоритм ScANN тоже делает похожий трюк только кластера называет партициями. Мы также можем использовать первый слой HNSW алгоритма для аналогичных целей. Это очень хорошо, т.к. возникает независимость от алгоритма.

Отмечу, что имена «кластер» и «партиция» крайне неудачны, т.к. в будущем могут внести путаницу в контексте распределенной системы, которую кто-то будет строить поверх нашей, поэтому я предлагаю здесь и далее называть это neighbors set. Вообще мой вам совет, когда не уверены как назвать какую-то сущность, которая является набором чего-то, лучше не поддаваться первому порыву назвать это блоком или чанком, а придумать название подлиннее. Потому что, к сожалению, короткие имена очень быстро кончаются.

Таблица состоит из сегментов, сегмент из neighbors set, neighbors set из близких документов.

Тут надо немного пояснить, как организовывают работу с векторами. Исходно есть какая-то модель, в которую клеют все-все факторы, их могут быть тысячи. Такая модель выплёвывает сильно разреженные вектора огромной размерности порядка 10000. По ним, конечно, никто не ищет, т.к. даже хранить их трудно. Поэтому используют какую-то другую модель которая их сжимает в размерность порядка 100. Т.к. исходные вектора разреженные, эта операция не приводит к сильным потерям, и получается плотный вектор. Но если его сжимать сильнее, то это приведёт к большой потери в точности.

Однако умножать вектора размерности порядка 100 всё ещё довольно затратно, поэтому для минимизации тяжелого скорирования часто применяют трюк под названием легкая/тяжелая формулы.

Идея в том, что мы можем хранить наши вектора в целевом виде размерности порядка 100 и одновременно в сильно сжатом виде меньшей размерности порядка 10. И сперва набирать топ с запасом, скорируясь по сильно сжатому представлению — так называемая легкая формула, а в конце уже на этом топе выбирать истинный топ по целевым векторам, т.н. тяжелая формула.

С точки зрения базы данных плотные вектора тяжелой формулы это исходные вектора, потому что другие ей не показывают. Поэтому давайте договоримся исходными называть плотные вектора, а сжатыми — вектора ещё меньшей размерности. Хотя формально они оба сжатые, а вторые по сути дважды сжатые. Но этот нюанс относится к предподготовке данных и нас не касается.

Верхнеуровневое описание поиска

Рассмотрим пока как задача решается в рамках одного сегмента.

Наш запрос состоит из:

Вектора для ANN поиска.
Размер конечного топа.
Размер предварительного топа для легкой формулы.
Фильтраций по обратному индексу
Фильтрации по прямому индексу условно трёх типов: высокоселективные, среднеселективные, низкоселективные.

В общем случае для легкой формулы нам необходимо подобрать такое преобразование, которое выдает вектора такие, что скалярное умножение сжатых векторов максимально близко к скалярному умножению полных. В общем случае это может быть сложное нелинейное преобразование зависящее от данных. Возможно даже у каждого neighbors set будет свое преобразование. Выбор центрального элемента тоже можно обобщить, т.к. вообще нас интересует не просто элемент наиболее похожий на данные, а тот который даёт наименьшую погрешность при скалярном умножении. Поэтому лучше назвать его не центральным, а типичным.

Получая запрос, мы сперва достаем его ANN часть и идём с ней в AnnIndex, который по размеру топа и паттерну относительно быстро набирает нам проскорированный список Neighbors Set. Затем мы скармливаем этот список планировщику, который возвращает план обхода.

Данные внутри neighbors set обходятся последовательно. Сперва мы фильтруем документы через обратный индекс и список живых документов. Затем применяем высокоселективные фильтры через прямой индекс, потом скорируем выжившие документы легкой формулой и набираем предварительный топ с запасом, отсекая документы по минимальному предварительному скору. Потом применяем среднеселективные фильтры.

Затем предварительные топы сливаются и в конце мы проходим ещё раз и набираем финальный топ тяжелой формулой.

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

Ещё один несомненный плюс — это очевидная возможность распараллелить обход отдельных neighbors set.

Эффективное поддержание топа

Самый наивный способ поддержать топ — это воспользоваться бинарной кучей: кладём в кучу элемент и, если куча слишком большая, вынимаем минимальный. Это требует $O(n + \log k)$ операций, где размер топа (и размер кучи), а — число кандидатов.

Но есть способ лучше. Для этого надо знать про алгоритм поиска k-ой порядковой статистики, по английски просто selection или partial sorting (это немного разные, но близкие задачи). Так же известен под именем nth_element в C++ и в Rust. Этот алгоритм, в частности, позволяет за линейное время найти медиану в случайном массиве и, более того, передвинуть все элементы меньше медианы в одну сторону, а больше медианы — в другую.

Допустим, мы ищем топ размера k, заводим массив размера и кладем туда кандидатов. Когда он заполнится делаем selection медианы за O(k) . Таким образом первые элементов массива это топ, а нижние — мусор. Далее мы начинаем добавлять элементы в нижнюю половину массива, причём только те, которые лучше медианы, и когда место заканчивается — повторяем selection. У такого алгоритма асимптотика O(n) , т.к. вставка амортизировано O(1) , поскольку обычно мы просто кладем элемент в массив и делаем selection, который стоит O(k) , не чаще одного раза на вставок.

Рекомендую статью про современные эффективные selection алгоритмы и также реализацию на C++ от того же автора.

Эффективное объединение постинг листов

Если мы хотим поддержать булеву алгебру для фильтрующих запросов, нам надо придумать как сделать операции NOT, AND и OR для постинг листов.

В общем случае, постинг лист представляет из себя итератор по документам, которые под него подходят. Мы обговорим API в следующей части, но, забегая вперед, скажу, что итератор должен уметь говорить, где он стоит сейчас и какой следующий DocId, а также прыгать на минимальный DocId не меньший чем заданный.

Операция NOT делается тривиально: надо просто получить следующий элемент оригинального итератора, запомнить его и далее возвращать все элементы до него (не включительно) и так же поступать далее.

Операцию OR по множеству итераторов можно реализовать так: кладем все итераторы в бинарную мин-кучу по DocId на котором он сейчас стоит. Когда нас просят сдвинуться на какой-то DocId мы достаем из кучи минимальный итератор и, если он меньше чем искомый, то сдвигаем его, а иначе останавливаемся. Повторяем это действие до остановки.

Операция AND делается так: предварительно сортируем итераторы по селективности, берем самый селективный и сдвигаем его, потом пытаемся сдвинуть следующий и далее, пока либо все они не сдвинуться на нужный DocId, либо пока какой-то итератор не перескочит его дальше. Во втором случае начинаем всё с начала.

Для production ready системы стоит ещё реализовать WEAK AND: как AND, но необходима истинность лишь некоторого числа условий. Но мы это опустим.

Формат хранения данных

В рамках дизайна нет смысла слишком подробно описывать формат хранения данных, потому что это очевидная точка бесконечной вариативности, однако будет полезно обрисовать основные темы, чтобы было ясно куда копать тому, кто решит написать своё хранилище.

Помимо метаинформации об индексе, который можно хранить хоть в JSON, мы храним следующую информацию о каждом сегменте:

Битсет живых объектов.
ANN индекс.
Обратные индексы.
Поколоночно все остальные поля-кортежи с разреженным индексом для быстрого доступа.

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

Потом идёт название кодека, с помощью которого он был сохранен, это позволит легко переключать версии кодека без необходимости предварительно конвертировать старые файлы в новый формат.

Потом контрольная сумма. Она обязательно нужна для надежности. Причём нам достаточно одной контрольной суммы на файл, т.к. мы его весь загружаем в память, но если бы мы разрабатывали формат данных для постоянного чтения с диска, то было бы разумнее хешировать отдельные чанки файла, чтобы можно было проверять корректность только нужных частей. Обратите внимание, что даже контрольная сумма идёт после имени кодека на случай, если мы захотим сменить алгоритм в будущем.

А в конце файла идет небольшая запись, необходимая для выравнивания. Даже если файл лежит в памяти, эффективнее считывать за раз целый регистр из нескольких байт, чем делать это побайтово. Однако можно случайно напороться на конец файла и сделать некорректное чтение. Что бы чтения оставались корректными, а код простым и эффективным, можно просто записать в конец немного данных заведомо больше любого регистра.

Битсет живых объектов

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

Можно динамически перестраивать структуру при удалении.
Можно ставить tombstone: делать дырку в файле, для пропуска её при чтении.
Можно завести отдельный файлик с битсетом живых объектов.

Первый вариант требует updatable структур данных и нам совсем не подходит. Отдельно хранящиеся битсеты лучше чем tombstone-ы, т.к. можно обновлять только их, оставляя весь остальной файл сегмента полностью иммутабельным. К тому же мы положим туда статистику общего числа живых и удаленных объектов, которая поможет лучше выбирать политику слияния старых сегментов. А также в разбивке по neighbors set, чтобы помочь планировщику запросов, выбирать neighbors set для обхода.

Формат хранения статистики живых документов

Формат векторов и ANN индекса

Для векторов мы должны максимизировать локальность данных во время вычислений, обеспечить быстрый доступ к neighbors set и быстрое последовательное сканирование. Поэтому поступим так:

Сперва идет индекс neighbors set-ов, который для каждого neighbors set хранит типичный вектор, функцию сжатия вектора для легкой формулы, указатель на начало сжатых векторов этого neighbors set и указатель на начало полных векторов.
После идут все сжатые вектора всех neighbors set подряд.
Затем — все полные вектора всех neighbors set подряд.

Сжатые представления можно было бы вычислять на лету, но это повторяющееся тяжёлое детерминированное вычисление, которое лучше заранее «закешировать». Сжатые и полные представления хранятся отдельно для максимизации локальности. Это и с диском, и с кешем, и в распределенной системе будет хорошо работать.

Формат хранения векторов и ANN индекса к ним

Формат хранения обратных индексов

Любой обратный индекс представляет из себя predecessor и набор posting-list’-ов, на которые тот ссылается.

Predecessor — это штука, которая может вернуть наибольший элемент, не превышающий заданный, и наименьший элемент, не меньший чем заданный. Простейшее решение — это дерево поиска, но т.к. наши сегменты иммутабельные, то нас интересует static predecessor. Самое просто решение — это сортированный по ключу массив пар + бинарный поиск, но есть варианты поинтереснее. Стандартный вариант с которого стоит начать: для строк использовать FST, а для чисел — сортированный массив с разреженной индексацией. Массив чисел стоит сжать следующим образом: в начало положить минимум и наибольший общий делитель, потом из сортированных чисел вычесть минимум и поделить на наибольший общий делитель, а то что от них останется уложить с помощью bit packing.

Таким образом, predecessor позволяет делать точечные и range-запросы. Понятно, что для наших целей в результате мы хотим получить не только сами ключи, но и ассоциированные с ними значения, а именно оффсет, начала постинг листа в файле.

Постинг листы хранят сортированными, чтобы быстро находить пересечения, и сжатыми, т.к. они могут быть огромными. Самый известный алгоритм сжатия строго возрастающих числовых значений это PFOR. Им сжимаются блоки DocId по 128 штук, и на них делаются скиплисты. Вообще, размер блока должен быть кратен 64, чтобы во время bit-packing стадии независимо от получившейся битности весь блок целиком занимал целое число 64-битных кусков. А 128 это просто эмпирически подобранное оптимальное значение.

Формат хранения остальных колонок

Для остальных полей за основу мы возьмем формат хранения колоночных данных Clickhouse, но с учетом наших neighbors set.

Каждая колонка группируется в блоки, и один блок не может содержать документы из разных neighbors set. Строится разреженный индекс на начало каждого блока (очевидно, среди них есть так же и индекс на начало каждого neighbors set). Блоки сжимаются независимо друг от друга. Чтобы увеличить локальность данных при последовательном сканировании, мы будем располагать блоки одной колонки не подряд, что было бы логично, а чередованно. Получатся такие чанки колонок: сперва все первые блоки всех колонок, потом все вторые и т.д. Таким образом, при сканировании документов последовательно у нас будет меньше случайных блужданий. Если я верно понял, Clickhouse формирует такие чанки динамически во время чтения данных с диска, однако т.к. у нас основной юзкейс — это mmap всех данных в память, то будет лучше сразу класть данные порезанными на чанки.

Колоночно-чанковый формат хранения документов

Запись

Нам осталось ответить на вопрос как мы будем писать данные, чтобы потом их так лихо читать. Формирование новых сегментов достаточно очевидно: надо набирать в памяти большой батч документов и затем сохранять его в нужном формате. Но есть два нюанса: как удалять документы и как мерджить сегменты.

Удаления и обновления

Чтобы отличить апдейт старого документа от вставки нового, нам надо уметь их как-то различать, даже если какие-то поля у них различаются. Т.е. нам нужен уникальный индекс. Условимся его задавать в схеме.

В общем случае, обновление — это удаление и вставка, поэтому нам достаточно решить задачу удаления, чтобы всё сошлось.

Вообще, мы могли бы поступать так: после сохранения нового сегмента брать все ключи уникального индекса, потом искать их нашим механизмом в остальных сегментах и обновлять списки живых документов там, где что-то найдётся.

Однако, очевидно будет эффективнее держать в памяти хеш-таблицу, отображающую уникальный ключ документа в имя сегмента и его DocId, чтобы делать эту операцию быстрее. Она будет занимать совсем немного места по сравнению со всем индексом целиком.

Merge сегментов

В простейшем виде политика мерджа сегментов должна стараться удерживать их общее число ниже определенного порога, а также, в первую очередь мерджить маленькие сегменты и те сегменты, в которых много удалённых документов.

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

Вообще, для обычных LSM-деревьев можно поступать эффективнее, т.к. есть первичный ключ по которому сегменты отсортированы и объединение не требует пересортировки данных. Однако в нашем случае расположение объектов зависит от значений векторов и того, как ведёт себя кластеризация, и хотя наверное можно придумать такие алгоритмы, которые позволят дешево мерджить два ANN индекса (скажем пытаться найти похожие neighbors set, объединять их и пересчитывать типичный вектор), в рамках этой статьи мы останемся с наивным решением.

Что будет дальше

Хотя мы уже описали алгоритмы, структуры данных и форматы хранения, которые собираемся использовать, этого пока не достаточно, чтобы сесть за реализацию. Мы ещё должны аккуратно разрезать систему на слои и описать API взаимодействия между ними. На этом этапе проектирования нам как раз и пригодится типизация.

Особая благодарность:

ссылка на оригинал статьи https://habr.com/ru/articles/832016/