Как я улучшил векторный поиск в YDB

TL;DR. В распределённой СУБД YDB ^[1] (читается вай‑ди‑би) векторный поиск по kmeans‑tree индексу раскрывался оптимизатором в цепочку из нескольких стадий StreamLookup. Это работало, но порождало большие планы запросов и существенно затрудняло оптимизации самого поиска. Я заменил эту цепочку одним специализированным read‑актором TKqpVectorSearchActor, который берёт всю логику обхода индекса под свой контроль, а не размазывает её по независимым стадиям.

Результат на стенде: 226.8 → 290.9 запросов/с (+28%), p50 42 → 33 мс, p99 66 → 53 мс при неизменном recall, равным 0.883. PR: ydb‑platform/ydb#42140 ^[2].

Часть 0. Зачем векторный поиск в OLTP‑СУБД

Пара слов для тех, кто пришёл не из мира векторных баз.

VSS (Vector Similarity Search) — это задача «найти K объектов, чьё векторное представление ближе всего к заданному вектору». В технической литературе такие векторные представления объектов называют эмбеддингами. Близость же меряется по косинусу или скалярному произведению, также используют евклидово расстояние.

Наивное решение — полный перебор: посчитать расстояние до каждого вектора и взять K ближайших. Для миллиона векторов, каждый из которых представлен массивом из 1024 чисел с плавающей точностью, это порядка 4 ГБ арифметики на один запрос. Такой подход полностью рабочий, однако работает он непозволительно долго. Поэтому строят ANN‑индексы (Approximate Nearest Neighbours), которые за счёт уменьшения точности позволяют получить выигрыш в скорости. По сути можно выделить два основных семейства:

1. Графовые (HNSW — Hierarchical Navigable Small World, и родственники): строим многослойный граф соседства, ищем жадным спуском. Отличный recall/latency, но структура плохо переживает шардирование и частые записи.

2. Кластерные / IVF‑подобные (Inverted File): разбиваем пространство на кластеры, при поиске выбираем несколько ближайших кластеров и перебираем только их содержимое.

У кластерных / IVF‑подобных индексов хуже recall при том же объёме работы, однако структура тривиально раскладывается на строки обычной таблицы, а значит — шардируется, реплицируется и переживает перезапуски ровно так же, как обычные данные.

YDB — распределённая OLTP‑СУБД с построчным хранением, шардированием по диапазонам первичного ключа и распределёнными транзакциями. Для неё второй вариант естественнее: индекс — это просто ещё несколько таблиц. Поэтому в YDB реализован kmeans‑tree — иерархический IVF: дерево кластеров, полученное рекурсивным применением k‑means.

Ценность истории в том, что векторный поиск живёт не в отдельном сервисе, а внутри той же базы, что и остальные данные. Это значит: один язык запросов, одни транзакции, один снапшот. Можно писать SELECT ... VIEW index ORDER BY Knn::CosineSimilarity(...) LIMIT 10 и джойнить это с чем угодно, не думая о синхронизации двух хранилищ.

Часть 1. Как устроен kmeans‑tree индекс в YDB

Индекс раскладывается в две (для префиксного варианта — три) служебные таблицы:

indexImplLevelTable   // — узлы дерева (центроиды)
  PK: (__ydb_parent, __ydb_id)
  columns: __ydb_centroid

indexImplPostingTable // — листья: принадлежность строк кластерам
  PK: (__ydb_parent, <PK основной таблицы>)
  columns: [опционально — покрывающие колонки]

indexImplPrefixTable  // — только для префиксных индексов:
                      // отображение значений префикса в корневой кластер группы

__ydb_parent — идентификатор родительского кластера, __ydb_id — идентификатор самого кластера, __ydb_centroid — вектор той же размерности, что и эмбеддинги строк, вычисленный как среднее арифметическое всех векторов, попавших в кластер на очередной итерации k‑means. Корень дерева — искусственный кластер с id = 0.

Поиск ближайших соседей выглядит следующим образом:

Ширину обхода задаёт параметр levelTop: сколько ближайших кластеров использовать на следующем уровене. Это классическая ручка recall/latency. Также есть overlap_clusters — строка может лежать сразу в нескольких кластерах, что повышает recall ценой размера posting‑таблицы (и требует дедупликации PK на выходе).

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

Часть 2. Как это работало раньше

До моей работы весь этот обход выражался средствами логического оптимизатора KQP (Kikimr Query Processor). Функция DoRewriteTopSortOverKMeansTree брала выражение вида «TopSort над чтением по вектору» и переписывала его в цепочку узлов TKqlStreamLookupTable.

Выглядит это примерно так (сильно сокращённый вариант):

auto read = Build<TKqlStreamLookupTable>(ctx, pos)
    .Table(levelTable)
    .LookupKeys(lookupKeys) // буквально: WHERE __ydb_parent = 0
    .Columns(levelColumns)
    .Settings(settingsNode) // VectorTopColumn/Target/Limit
    .Done().Ptr();

VectorReadLevel(indexDesc, ctx, pos, levelLambda, top, levelTable,
                levelColumns, settings.VectorTopLimit, settingsNode, read);
// ^ добавляет ЕЩЁ ОДИН StreamLookup на каждый уровень дерева

settings.VectorTopColumn = indexDesc.KeyColumns.back();
settings.VectorTopLimit = top.Count().Ptr();
settings.VectorTopDistinct = withOverlap;
VectorReadMain(ctx, pos, postingTable, ..., mainTable, ..., settings, read);
// ^ добавляет StreamLookup на posting + StreamLookup на main

VectorTopMain(ctx, top, read); // финальный TopSort по расстоянию

Каждый StreamLookup — это отдельная стадия физического плана. Стадии в YDB соединяются каналами и исполняются compute‑акторами, потенциально на разных узлах кластера. Таким образом для двухуровневого индекса план запроса получится следующий план запроса:

Формально — работает. Практически — три проблемы, ровно те, что перечислены в issue #30408 ^[3]:

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

2. Оптимизации некуда вставлять. Хочется кэшировать верхние уровни дерева в RAM (они неизменяемы между перестроениями индекса), хочется пробовать разные стратегии обхода — best‑first вместо строго послойного, адаптивную ширину, ранний выход по границе расстояния. Всё это — логика, у которой должно быть состояние на протяжении всего поиска. В цепочке независимых стадий такого места нет: каждая стадия видит только свой вход и выход.

3. Конвейеризация невозможна. Скан posting‑таблицы и точечные чтения основной — это две разные стадии, соединённые каналом. Пока posting‑скан не отдаст свою порцию строк дальше, чтения основной таблицы не начнутся. А ведь как только у нас есть первая сотня PK‑кандидатов, за ними уже можно идти — параллельно с продолжающимся сканом.

Часть 3. Идея: один актор вместо цепочки стадий

YDB построена на акторной модели. Compute actor исполняет одну стадию плана; данные в стадию могут приходить не только по каналу от другой стадии, но и из источника (source) или через input transform — актора, который реализует интерфейс IDqComputeActorAsyncInput и подкладывает compute‑актору строки.

Ровно сюда и напрашивается векторный поиск: это один актор, который получает на вход целевой вектор, а на выход отдаёт готовые top‑K строк, отсортированные по расстоянию. Всё, что между, — его личное дело.

Логический оптимизатор при этом сильно упрощается. Вместо построения цепочки — один узел:

TExprNode::TPtr read = Build<TKqlReadTableVectorIndex>(ctx, pos)
    .Table(match.Table())
    .Index(ctx.NewAtom(pos, indexDesc.Name))
    .Columns<TCoAtomList>().Add(columns).Build()
    .TopK(top.Count())
    .TargetVector(TExprBase{targetVector})
    .Done().Ptr();

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

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

// Ensure the embedding column is among the read columns so the actor can rank rows
// even if it is not part of the final projection.
if (!hasEmbedding) {
    columns.push_back(Build<TCoAtom>(ctx, pos).Value(embeddingColumn).Done());
}

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

Часть 4. Устройство актора

TKqpVectorSearchActor — это конечный автомат из четырёх состояний:

Внутри каждой фазы актор запускает вложенные чтения (inner reads). Ключевое решение: он не изобретает работу с даташардами заново, а переиспользует существующий CreateKqpReadActor — тот же самый read actor, что обслуживает обычные ReadTableRanges в планах. Актор создаёт его вложенным, подкладывает настройки чтения и вычитывает результат.

using TKqpVectorInnerReadFactory = std::function<std::pair<IDqComputeActorAsyncInput*, IActor*>(
    const NKikimrTxDataShard::TKqpReadRangesSourceSettings* settings,
    TIntrusivePtr<NActors::TProtoArenaHolder> arena,
    const NActors::TActorId& parentId,
    TVector<TSerializedCellVec>&& keyPoints)>;

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

Фаза Level: один читатель на раунд, а не на родителя

Все родители раунда идут в одно чтение, в виде списка точечных диапазонов по __ydb_parent. Read actor сам раскидает эти диапазоны по шардам level‑таблицы и опросит шарды параллельно.

// Read all parents of this round in a single inner read: it fans the
// per-parent ranges out across the level table's shards in parallel
// instead of doing one sequential round-trip per parent.

Ранжирование детей идёт инкрементально, по мере прихода строк, в ограниченную по размеру кучу:

void PushLevelCandidate(TClusterId id, double distance) {
    auto cmp = [](const auto& a, const auto& b) { return a.second < b.second; };
    if (LevelCandidates.size() < LevelTop || distance < LevelCandidates.front().second) {
        PushBoundedMaxHeap(LevelCandidates, LevelTop, std::make_pair(id, distance), cmp);
    }
}

Max‑heap с худшим элементом на вершине: как только куча заполнена, кандидат, который не ближе текущего худшего, отбрасывается сразу. Структура ранжирования никогда не растёт больше levelTop, сколько бы шардов ни сыпало строками. То же самое — для финального top‑K по основной таблице.

Фаза Posting: конвейер вместо барьера

Здесь живёт основной архитектурный выигрыш. Скан posting‑таблицы и точечные чтения основной перекрываются во времени:

Каждый цикл вычитывания, который оставил в буфере хоть один новый PK, немедленно отправляет его в чтение основной таблицы:

// Non-covered posting rows stream in across many drain cycles; every cycle that
// leaves keys buffered dispatches them straight into a main read that overlaps the
// still-running posting read -- there is no accumulation threshold, so a drain cycle
// with a single new key still launches a read. Because every cycle's keys go out
// immediately, nothing is left stranded and no explicit final-tail flush is needed.

Порога накопления сознательно нет. Он был в промежуточных версиях — и создавал проблему «хвоста»: последние несколько ключей зависали в буфере, требуя явного финального сброса. Без порога инвариант простой: буфер пуст на границе каждого цикла.

Из‑за этого перекрытия, кстати, пришлось ввести отдельный тег EReadKind для вложенных чтений:

// The posting and main reads of a non-covered search overlap (a main read is
// launched once enough candidate PKs have accumulated, while the posting read is
// still running), so a single EPhase no longer identifies a read's table -- the tag does.
enum class EReadKind { Level, Posting, Main };

Классический признак того, что конвейеризация состоялась: фаза перестала однозначно определять, что именно сейчас читается.

Покрывающие и частично покрывающие индексы

Если posting‑таблица содержит все нужные колонки (PostingCovers), чтение основной таблицы пропускается целиком — поиск завершается на скане posting. Тут возникла отдельная сложность: чтобы собрать строку результата, актору нужны и список выходных колонок, и отдельно — колонки первичного ключа (по ним строка идентифицируется). Но PK часто входит в список выходных колонок — например, если пользователь и так запросил id, а id это и есть PK. Если склеить оба списка наивно, один и тот же id колонки окажется в запросе дважды, а такое чтение даташард отвергает. Пришлось один раз в конструкторе заранее построить план раскладки: для каждой PK‑колонки решить, переиспользовать ли её позицию среди выходных колонок или добавить отдельной:

// Output columns occupy positions 0..N-1 (distinct, matching VectorColumnIndex and
// the result row layout); each PK column reuses its output position if it is also
// an output column, else is appended (recorded in CoveredExtraPkIndices).

Более интересен частично покрывающий случай: posting‑таблица содержит эмбеддинг, но не все выходные колонки (типичный пример — префиксный индекс, у которого в posting нет колонки префикса). Основную таблицу читать придётся, но ранжировать можно уже на posting‑скане — и тогда в даташард можно протолкнуть top‑K: каждый шард вернёт не все свои строки, а только TopK ближайших. Экономия на сети и на числе последующих точечных чтений получается заметная.

Кэш уровней

Level‑таблица неизменяема между перестроениями индекса. Значит, её строки можно кэшировать в процессе, ключ — (path id level-таблицы, id родительского кластера):

// The KMeans level table is immutable, so its rows can be cached across
// queries keyed by (level table path, parent cluster id).
UseLevelCache = LevelsCache && LevelsCache->MaxBytes() > 0;

При старте раунда родители‑попадания обслуживаются из кэша и вообще не попадают в чтение; промахи накапливаются в CachingLevelBatches и после завершения чтения кладутся в кэш.

Часть 5. Результаты

Стенд

Конфигурация: один узел, 2× Xeon E5-2650v4, 512 ГБ RAM, 6× SATA SSD. Датасет — Wikipedia english_simple (Cohere/wikipedia-22-12-simple‑embeddings ^[4]), 485 859 строк. Индекс: 2 уровня × 100 кластеров, overlap_clusters = 3. Поиск: LIMIT 20, levelTop = 10, 100 целевых векторов. Векторная impl‑таблица не партиционирована.

Последняя строка важнее всех остальных: алгоритмически поиск не изменился. Тот же обход, та же ширина, та же выборка кластеров, та же точность. Изменилась только механика исполнения. Это чистое ускорение, а не размен recall на latency.

Локальный A/B

Отдельно гонял штатный скрипт сравнения производительности векторного поиска на своей машине: сборка release, 100 секунд на прогон, 10 итераций, медиана.

Разброс между двумя измерениями (+28% против +9.7%) — не противоречие, а ожидаемое следствие того, что реальный прирост зависит от параметр ов индекса и поиска, настроек партиционирования таблицы и общей конфигурации стенда.

Часть 6. Что дальше

Главное, ради чего всё затевалось, — не 28%, а то, что теперь есть место для оптимизаций поиска. Появившиеся возможности:

1. Стратегии обхода дерева. Сейчас обход строго послойный: взяли levelTop лучших на уровне, спустились. Best‑first обход по глобальной очереди с приоритетом даёт при том же бюджете чтений лучший recall. Это чисто внутренняя логика актора — переписывается локально.

2. Адаптивная ширина и ранний выход. Если после первого листового кластера top‑K уже заполнен, а расстояние до его худшего кандидата достаточно мало, дальние кластеры можно не смотреть — маловероятно, что там найдётся что‑то ближе. Нужен доступ к текущему состоянию ранжирования во время обхода — он теперь есть.

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