Политика обновления через reinforcement learning

Для одного запроса запускаются четыре RL прохода. На каждом проходе Qwen генерирует кандидата в план запроса и отправляет его в Postgres для сравнения с собственным планом. Каждому проходу присваивается скалярное вознаграждение, которое распространяется в обратном направлении для обновления весов модели.

Веса направлены в сторону более быстрых планов.

Насколько хороши оптимизаторы запросов на самом деле?

Leis и коллеги задали этот вопрос в 2015 году. Через 10 лет они вернулись к этому вопросу снова.

Несмотря на огромный объём исследований за десятилетие с момента первоначального анализа, оптимизаторы запросов остаются далеки от идеала.

Сначала это было удивительно. База данных Postgres должна знать всё о данных в своих таблицах, верно? Как это может быть сложно?

Как выясняется: невероятно сложно. На самом деле одна из критических операций в оптимизаторе — упорядочение присоединений (join ordering) — известна как NP-трудная.

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

Далее описывается эксперимент, проведённый для ответа на вопрос: может ли небольшая модель с открытыми весами получить дополнительную подготовку через supervised fine-tuning (SFT) и agentic reinforcement learning (RL) для создания планов запросов Postgres, которые превосходят встроенные планы?

Ответ утвердительный. Ключевые результаты:

  • Достигнут 44,7% прирост скорости отклика на 113 запросах с большим количеством присоединений, начиная с модели, которая не могла создать план для 99 из них
  • Построен стенд измерения Postgres, который минимизирует помехи кэша страниц Linux между параллельными контейнерами
  • Разработан пользовательский вариант GRPO для оценки прохождений RL в среде с естественной неточностью
  • Распределение RL между двумя машинами: vLLM и тренер на арендованном узле 2x H100 и четыре контейнера Postgres на рабочем столе
  • Запуск off-policy дистилляции на половине тысячи траекторий агента GPT-6 Astra

Начнём с самого начала.

Внутри оптимизатора запросов

Рассмотрим следующий срез набора данных IMDb:

-- An IMDb title (movie, series, episode, etc.) [~1M rows]
title (
  id              integer PRIMARY KEY,
  title           text,
  production_year integer,
  kind_id         integer -- FK -> kind_type
)

-- Movie <> company junction table [~2M rows]
movie_companies (
  id              integer PRIMARY KEY,
  movie_id        integer, -- FK -> title.id
  company_id      integer, -- FK -> company_name.id
  company_type_id integer, -- FK -> company_type.id
  note            text
)

-- A company's name, origin, etc. [~100k rows]
company_name (
  id           integer PRIMARY KEY,
  name         text,
  country_code text     -- '[us]', '[jp]', ...
)

-- Lookup table of company roles for a title [4 rows]
company_type (
  id   integer PRIMARY KEY,
  kind text -- 'production companies', 'distributors', ...
)

-- Lookup table for what a title _is_ [7 rows]
kind_type (
  id   integer PRIMARY KEY,
  kind text -- 'movie', 'tv series', 'episode', ...
)

Допустим, мне нужно ответить на вопрос: «Какие японские компании выпустили больше всего фильмов в 2000-х годах?» Можно написать следующий запрос:

SELECT cn.name,
       COUNT(*) AS titles
FROM   title AS t,
       movie_companies AS mc,
       company_name AS cn
WHERE  t.id = mc.movie_id
  AND  mc.company_id = cn.id
  AND  cn.country_code = '[jp]'
  AND  t.production_year BETWEEN 2000 AND 2009
GROUP  BY cn.name
ORDER  BY titles DESC
LIMIT  10;

Запрос выводит 10 японских компаний с количеством фильмов, связанных с ними в период 2000–2009, отсортированных от наибольшего к наименьшему.

Но как Postgres получил эти результаты?

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

Чтобы это проиллюстрировать, представим тот же запрос без фильтра японской компании и диапазона дат:

SELECT cn.name,
       COUNT(*) AS titles
FROM   title AS t,
       movie_companies AS mc,
       company_name AS cn
WHERE  t.id = mc.movie_id
  AND  mc.company_id = cn.id
GROUP  BY cn.name
ORDER  BY titles DESC
LIMIT  10;

mc может присоединиться к cn только через mc.company_id = cn.id, а t может присоединиться к mc только через t.id = mc.movie_id.

Эти ограничения создают два действительных дерева присоединений:

Первое дерево (cn ⋈ mc) ⋈ t выполняет нижнее присоединение первым; результат становится входом для корневого присоединения.

Мощность таблицы или результата запроса — это количество строк, которые она содержит. Предположим, что соответствующие таблицы имеют следующие мощности:

  1. cn = 100k
  2. mc = 2m
  3. t = 1m

Учитывая наши присоединения, получаем следующие мощности:

(cn ⋈ mc) = 2m, затем ⋈ t = 2m

(t ⋈ mc) = 2m, затем ⋈ cn = 2m

Независимо от порядка, в который эти три таблицы присоединяются, одни и те же 2m строк всегда передаются во второе присоединение.

Теперь добавим обратно наши избирательные предикаты:

  1. cn' = 5k (предположим, что 5% из 100k компаний — японские)
  2. mc = 2m (не изменяется)
  3. t' = 200k (предположим, что 20% из 1m названий были созданы в 2000-х)

(cn' ⋈ mc) ≈ 100k, затем ⋈ t' ≈ 20k

(t' ⋈ mc) ≈ 400k, затем ⋈ cn' ≈ 20k

Первое упорядочение присоединений фильтрует 2m записей movie_companies до 5% среза компаний, которые японские. Предполагая равномерное распределение (обсудим позже, почему мы это предполагаем), это присоединение даёт примерно 100k строк. Присоединение результата с отфильтрованной таблицей title сохраняет только 20% из этих строк из 2000-х годов.

Второе упорядочение присоединений фильтрует 2m записей movie_companies до 20% среза названий, созданных в 2000-х. То же предположение о равномерности, поэтому первое присоединение даёт 400k строк, что означает, что во второе присоединение передаётся 400k строк.

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

К сожалению, на этом не заканчивается.

Комбинаторный взрыв

Каждое присоединение может использовать любое из:

  1. Hash join
  2. Merge join
  3. Nested-loop join

Учитывая коммутативность, существует 4 различные ориентации внешнего/внутреннего присоединения, что приводит к 8 возможным комбинациям.

Кроме того, каждую таблицу можно сканировать разными способами. Рассмотрим только четыре типа сканирования:

  1. Sequential
  2. Index
  3. Index-only
  4. Bitmap

2 дерева присоединений × 2² ориентации × 3² алгоритма × 4³ сканирования = 4,608

Существует 4608 различных способов выполнения этого запроса!

Чтобы усугубить ситуацию, каждое присоединение комбинаторно взрывает пространство поиска. С двумя таблицами — 96 вариантов. С тремя таблицами — 4,608. С четырьмя — 442,368. С пятью — уже 33 миллиона. С двенадцатью таблицами — более 72 триллионов.

Оценка вместо подсчёта

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

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

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

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

Заключение

Языковые модели могут быть обучены через supervised fine-tuning и reinforcement learning для создания оптимальных планов запросов SQL, превосходящих встроенные оптимизаторы баз данных. Подход показывает обещание для других областей, где требуется оптимизация комбинаторных поисков.