Поиск «токсичных» SQL-запросов

Мы, студенты из МИФИ, Даниил и Александр, пришли на стажировку в Сбербанк в департамент SberData, который занимается развитием внутренней корпоративной аналитической платформы (КАП).Это современная платформа с удобными инструментами созданная для закрытия полного спектра потребностей Сбера в работе с данными, таких как хранение, интеграция, разнообразная аналитика, отчетность, моделирование и контроль качества данных. Все эти направления было бы трудно развивать без отдельного R&D подразделения, в составе которого мы и работаем. Сегодня мы хотим поделиться нашим исследованием в области проектирования алгоритмов в выявлении «токсичных» SQL‑запросов с помощью машинного обучения. Почему же запросы называются именно «токсичные»? Они затрачивают на своё выполнение слишком большое количество ресурсов, а именно времени. На самом деле не только время, но для упрощения мы будем считать только время, так как это ключевой параметр.

Статья посвящена исследованию существующих подходов и их апробации на открытых данных. В качестве общедоступных данных были выбраны данные из таких бенчмарков, как TPC‑H и BIRD. Помимо этого, в статье рассматриваются некоторые трудности, с которыми мы столкнулись при работе над задачей, например, генерация данных и SQL‑запросов, а также миграция между диалектами SQL. В конце статьи мы опишем оригинальный подход, к которому по итогу пришли. В следующей статье мы расскажем о применении полученного опыта для реальной промышленной системы.

Наше решение задачи можно разделить на несколько пунктов:

• Обзор существующих методов

• Описание выбранных моделей и метрик

• Тестирование моделей на TPC‑H

  • Модуль генерации SQL‑запросов

• Тестирование моделей на BIRD

  • Модуль генерации данных в БД

  • Миграция с SQLite на PostgreSQL

• Выводы

Обзор существующих методов

Рассмотрим три различных подхода:

1.1 Одним из таких являются графовые нейронные сети — архитектура нейронный сети позволяет работать с данными, представленными в виде матрицы смежности, за счет графовой свертки. В данной работе (Query Performance Prediction for Concurrent Queries using Graph Embedding) план выполнения запроса представляется в виде графа, вершинами которого являются операторы запроса и их параметры, ребра же отражают связи между этими операторами. Этот метод дает возможность рассматривать сразу несколько запросов, что позволяет учитывать конкуренцию за ресурсы при одновременном выполнении нескольких запросов. Метод обеспечивает высокую точность предсказания, но требует больших объемов данных для эффективного обучения. Пример построения нейронной сети для набора операторов представлен на рисунке ниже.

1.2 Затем мы обратили внимание на идею (Plan‑Structured Deep Neural Network Models for Query Performance Prediction) рассмотрения каждого оператора запроса, как отдельного элемента. Предлагается использовать модульную архитектуру глубокой нейронной сети, где каждый оператор плана запроса отображается отдельным модулем. Эти модули объединяются в единую сеть, которая соответствует структуре плана запроса и позволяет точно прогнозировать задержки на каждом этапе выполнения запроса. Такой подход позволяет гибко и точно моделировать сложные SQL‑запросы, учитывая их структуру и ресурсоемкость. Процесс работы модели (полный цикл) показан на рисунке ниже.

1.3 Эти методы используются для предсказания времени выполнения SQL‑запросов на основе анализа их планов выполнения, однако стоит учитывать, что в реальной системе не всегда есть доступ к плану выполнения запроса. По этой причине на фоне остальных подходов выделяется метод, основанный на векторизации текста запроса, используя такие методы, как TF‑IDF и Bag of Words (BoW). Например, в исследовании используются методы векторизации текста в сочетании с различными моделями, для прогнозирования времени выполнения запросов и их ресурсов. Архитектура модели, представленная в статье, показана на рисунке ниже.

Таким образом опишем плюсы и минусы каждого метода:

Описание выбранных моделей и метрик

Ниже мы будем применять понятие плана выполнения запроса — детализированное описание последовательности операций, которые СУБД (система управления базами данных) выполняет для получения результата SQL‑запроса. Для того чтобы не вмешиваться в промышленные системы, мы будем использовать текст запроса. Но из плюсов нам будет доступно только два параметра из плана запроса — cost и row. Cost (оценка стоимости) указывает на оценку стоимости выполнения каждой операции в плане запроса. Row показывает количество строк, которое, по оценке СУБД, будет прочитано.

Для выбранных нами признаков опишем используемые алгоритмы векторизации и модели регрессии.

Алгоритмы векторизации:

• BoW (Bag of Words): преобразует текст в числовые признаки, подсчитывая частоту каждого слова в документе;

• TF‑IDF (Term Frequency‑Inverse Document Frequency): преобразует текст в числовые признаки, учитывая частоту слова в документе и обратную частоту слов в корпусе текстов;

• BERT (Bidirectional Encoder Representations from Transformers): используется для создания векторных представлений текстов. Модель предобучена (bert‑base‑uncased) на больших данных и способна улавливать контекстные зависимости между словами в предложении. Тексты запросов были преобразованы в токены с помощью BERT Tokenizer, а затем токены использовались для извлечения признаков с помощью BERT Model.

Модели регрессии:

• Lasso — регуляризованный алгоритм линейной регрессии, использующий штраф L1;

• Ridge — алгоритм, аналогичный Lasso, но использующий штраф L2;

• Random forest — случайный лес, используемый для задачи регрессии;

• Neural network — классическая полносвязная многослойная нейронная сеть, реализованная на Pytorch;

• SVR — метод опорных векторов для задачи регрессии.

Метрики:

• MAPE — среднее абсолютное отклонение (погрешность предсказания в долях / процентах):

  

• Accuracy — для задачи регрессии, если предсказанное значение совпадает с истинным с погрешностью %, то считается, как верное предсказание. Далее используется окно в 25 % погрешности.

  

В этих формулах— это ответ модели, а — это истинное значение.

Тестирование моделей на TPC-H

Сразу получить доступ к данным у нас не вышло, и первым шагом решили потренироваться на «кошечках». Первоначально для обучения моделей и предварительной оценки был выбран бенчмарк TPC‑H, который в основном используется для сравнения производительности аналитических систем и хранилищ данных. Решили взять для теста именно его, так как в нем уже есть более пяти связанных по ключу таблиц и генерировать данные для теста нет необходимости. Во время тестирования моделей столкнулись с тем, что SQL‑запросов, содержащихся в выбранном бенчмарке, не хватает, поэтому был разработан модуль для их генерации при помощи языковых моделей, который будет описан ниже. Для проверки работы моделей сгенерировали данные — среднее выполнение SQL‑запросов до десятков секунд. При тестировании были получены результаты, представленные в таблице ниже.

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

Модуль генерации SQL-запросов

Из‑за недостатка SQL‑запросов решили написать свой небольшой модуль по генерации, в который можно найти по ссылке. Принцип работы основан на большой языковой модели — GigaChat. Для наших задач он отлично справился с генерацией дополнительных SQL запросов. Вот алгоритм действия:

1. Идет обращение в GigaChat с указанием на генерацию некоторого числа SQL‑запросов;

2. Из полученного текста извлекаются SQL‑запросы;

3. Они валидируются на БД малого объема для проверки корректности синтаксиса.

При разработке были выявлены следующие особенности:

• лучше указывать ограничения в генерируемых запросах за раз, тогда их разнообразие будет выше (подбирается эмпирическим методом);

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

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

Тестирование моделей на BIRD

С целью получения набора данных с изначально большим числом запросов и таблиц на следующем этапе был выбран бенчмарк BIRD, предназначенный для решения задачи text2sql. Он включает в себя множество баз данных, к каждой из которых есть определенное число запросов. Таблица с топ-7 БД по количеству запросов представлена ниже.

Наш окончательный выбор пал на базу public_review_platform из‑за большого количества запросов. При этом количество таблиц в ней небольшое — это позволит легко мигрировать на postgres с изначального sqlite, на котором сделан весь бенчмарк BIRD. Подробнее про миграцию будет ниже.

Также на этом шаге мы решили добавить дополнительную предобработку запросов, чтобы повысить точность, а именно:

• Форматирование (приведение к единому регистру и замена множественных пробелов);

• Удаление комментариев (как однострочных, так и многострочных);

• Замена алиасов на полные названия таблиц;

  • Было:

    SELECT u.user_id, u.user_average_stars, r.review_stars FROM users u JOIN reviews r ON u.user_id = r.user_id WHERE r.review_length = ‘Short’ AND r.review_votes_useful = ‘Medium’;

  • Стало:

    select users.user_id, users.user_average_stars, reviews.review_stars from users join reviews on users.user_id = reviews.user_id where reviews.review_length = ‘short’ and reviews.review_votes_useful = ‘medium’;

• Добавление дополнительных признаков из плана запроса (примерная стоимость запроса и количество обрабатываемых строк).

Более того, к базе данных были сгенерированы дополнительные запросы с помощью модуля генерации запросов, описанного выше. Тем самым число запросов увеличилось до 600, и, помимо этого, размер базы данных был расширен с 100 Мб до 15 Гб суммарно. Затем на полученных данных было проведено тестирование, результаты которого представлены ниже.

TF‑IDF

BoW 

BERT

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

При рассмотрении различных методов векторизации текста было выявлено, что мешок слов в среднем оказался хуже, поскольку метод просто считает количество вхождений слов, не учитывая их взаимосвязи и контекст. В противоположность этому, TF‑IDF показал высокие результаты, так как учитывает частоту встречаемости каждого слова в документах, что предотвращает переоценку часто встречающихся терминов (например, операторов SQL) и позволяет выделять редкие, но значимые слова, такие как названия таблиц или столбцов. BERT также показал высокие результаты, но дальнейшее тестирование проводилось на векторизаторе TF‑IDF из‑за его высокой скорости работы. Тем не менее, итоговую модель можно запустить на BERT, так как он способен учитывать семантику слов, что повышает точность модели.

При анализе зависимости точности моделей от максимального количества признаков в векторизаторе TF‑IDF было замечено, что с увеличением количества признаков точность моделей возрастает. Это связано с тем, что больший размер вектора позволяет учитывать большее количество столбцов таблиц и ключевых слов запросов. Однако избыточный размер вектора перестает приносить дополнительные преимущества, и точность моделей стабилизируется. Поэтому размерность вектора следует подбирать, исходя из среднего и максимального количества таблиц и столбцов в системе. Для нашего случая оптимальным значением максимального размера вектора оказалось 4096, но при необходимости более быстрой работы можно использовать и размер, равный 1024. Выбор размерности объясняется размером «словаря», который используется в запросах, в частности, с количеством уникальных названий таблиц и столбцов. Забегая вперед, скажем, что на реальных данных необходим вектор большего размера, поскольку мощность словаря больше.

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

Модуль генерации данных в БД

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

1. Ввод схемы БД для генерации данных;

2. Поддержка внешних ключей для корректной генерации (происходит наивная проверка среди уже сгенерированных данных);

3. Для каждого столбца можно задать значения для генерации (например, список городов);

4. Запись в файл или построчная запись в PostgreSQL.

Миграция с SQLite на PostgreSQL

Бенчмарк BIRD в качестве диалекта использует SQLite, типы данных которого отличаются от PostgreSQL. Просто использовать SQLite нежелательно, и мы мигрировали в PostgreSQL. Поэтому нам необходимо было выполнить следующие шаги для миграции (их можно выполнить при помощи ручного анализа и регулярных выражений для замены):

1. Перевести описание таблиц в корректный формат;

1. Изменить все типы данных в соответствии с содержанием (например, некоторые колонки «текст» перевести во «время»);

2. Перевести табличные файлы описания в.sql файлы;

2. Внести изменения данных в запросы (например, «11AM» перевести в «11:00:00»).

Выводы

Для PostgreSQL получилось достичь результатов, с которыми уже можно продолжать работать на реальных СУБД (на бенчмарке BIRD: ~30% отклонения / 75–80% точности). Лучшими векторизаторами оказались TF‑IDF и BERT, но последний более ресурсозатратный. А лучшую точность показали регрессоры градиентный бустинг (XGBoost) и случайный лес (Random Forest) за счет их ансамблевой структуры.

Основные преимущества модели:

• Базируется на тексте запроса, который несложно получить в системе;

• Хорошая интерпретируемость.

Недостатки и возможные нюансы при дальнейшей работе:

• При большом количестве таблиц вектор будет слишком большим (вероятно, необходимо уменьшить признаковое пространство).

Авторы: Усков Даниил(@Dusik207) и Диков Александр(@DikovAlexandr) под кураторством Максима Радионова(@tr1ggers), Дмитриева Игоря (@indmitriev) и Владимира Зворыгина(@Wocher), участники профессионального сообщества Сбера DWH/BigData.