Быстрая обработка CSV с помощью ОКМД (SIMD)

Недавно автор узнал об инструменте csvquote, который кодирует проблемные символы CSV так, чтобы утилиты unix правильно их обрабатывали. Он меняет кодировку в конце конвейера, восстанавливая исходный ввод. Оригинальная реализация работает с кавычками CSV простым, примитивным методом. Но на современном «железе» есть подход лучше, проще и в 3 раза быстрее.

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

Есть ещё один подход, сочетающий принцип ОКМД (одиночный поток команд, множественный поток данных, также называется SIMD) с приёмами битового жонглирования, который на порядок увеличивает скорость обработки. Моя реализация — csvquote включает оба подхода.

Введение

Записи в данных CSV разделяются символами переноса строки, а поля — запятыми. Поля могут заключаться в кавычки:

aaa,bbb,ccc
xxx,"yyy",zzz

Поля, содержащие символ переноса строки (U+000A), кавычки (U+0022) или запятую (U+002C), должны заключаться в кавычки, иначе они будут восприниматься как элементы форматирования CSV. Закавыченные кавычки превращаются в пару кавычек. Например, вот две записи с двумя полями в каждой:

"George Herman ""Babe"" Ruth","1919–1921, 1923, 1926"
"Frankenstein;
or, The Modern Prometheus",Mary Shelley

Инструмент, не поддерживающий применяемое в CSV разделение данных посредством запятых и символов переноса строки (например, awk), обработает эти записи некорректно.

Но csvquote переводит закавыченные символы переноса строки в разделители записей (U+001E), а запятые — в разделители элементов данных (U+001F). Эти управляющие символы редко появляются в обычных текстовых данных и могут быть легко обработаны в тексте в кодировке UTF-8 без декодирования или кодирования. Записи превращаются в:

"George Herman ""Babe"" Ruth","1919–1921\x1f 1923\x1f 1926"
"Frankenstein;\x1eor\x1f The Modern Prometheus",Mary Shelley

Управляющие символы здесь — \x1e и \x1f .

Данные имеют ровно ту же длину, ведь это прямая замена байт на байт. Кавычки остаются нетронутыми. Задача — спарсить кавычки и отследить, окажутся ли два этих символа внутри пар кавычек или нет.

Совершенствование конечного автомата

Оригинальная реализация csvquote обходит входные данные побайтово и находится в одном из трёх состояний:

1. Вне кавычек (исходное состояние).

2. Внутри кавычек.

3. В «сбежавшей» кавычке (первой " из "").

В моей реализации я использую конечный автомат вместе со switch:

// Return the next state given an input character.
int next(int state, int c)
{
    switch (state) {
    case 1: return c == '"' ? 2 : 1;
    case 2: return c == '"' ? 3 : 2;
    case 3: return c == '"' ? 2 : 1;
    }
}

В реальной программе условий потенциальной замены больше. Убивающих производительность ветвлений ужасно много.

Но в данном контексте мы имеем дело с поиском «снаружи-внутри», а не проверкой CSV, поэтому состояние в «сбежавшей» кавычке лишнее. Надо лишь сопоставить пары кавычек. «Сбежавшую» кавычку можно считать завершающей закавыченную область и сразу начинающей новую. А значит, в этой простой конфигурации есть только два первых состояния:

int next(int state, int c)
{
    switch (state) {
    case 1: return c == '"' ? 2 : 1;
    case 2: return c == '"' ? 1 : 2;
    }
}

Текст может обрабатываться в виде байтов, поэтому возможно только 256 вариантов входных данных. Имея 2 состояния и 256 входных данных, этот конечный автомат благодаря механизму замены может быть реализован без ветвлений с 512-байтовой таблицей. Вот инициализация таблицы:

unsigned char table[2][256];

void init(void)
{
    for (int i = 0; i < 256; i++) {
        table[0][i] = i;
        table[1][i] = i;
    }
    table[1]['\n'] = 0x1e;
    table[1][',']  = 0x1f;
}

В первом состоянии символы отображаются сами на себя. Во втором — на свои замены. Вот энкодер и декодер целиком:

void encode(unsigned char *buf, size_t len)
{
    int state = 0;
    for (size_t i = 0; i < len; i++) {
        state ^= (buf[i] == '"');
        buf[i] = table[state][buf[i]];
    }
}

Декодеру, строго говоря, не нужно обрабатывать кавычки. По бенчмарку (в моей реализации это csvdump) он на ноутбуке обрабатывает данные со скоростью ~1 ГБ/с — в 3 раза быстрее, чем в оригинальной реализации. И можно ускориться!

ОКМД и дополнение двух подходов

В любой хорошей реализации ОКМД используется маскирование. Находим кавычки, вычисляем маску по закавыченным областям, и другую маску — для совпадений с заменами, объединяем эти маски, затем используем полученную маску для смешивания входных данных с заменами. Примерно так:

quotes    = find_quoted_regions(input)
linefeeds = input == '\n'
commas    = input == ','
output    = blend(input, '\n', quotes & linefeeds)
output    = blend(output, ',', quotes & commas)

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

Вот что мы сделаем. Загрузим 32 байта в ОКМД-регистр (например, AVX2) и вычислим 32-битную маску, где каждый бит соответствует одному байту. Если в байте кавычка, соответствующий бит установлен в 1:

"George Herman ""Babe"" Ruth","1
10000000000000011000011000001010

Последний бит соответствует началу закавыченной области. Для маски желательно установить все биты, следующие за этим битом. Сделаем это, вычитая 1:

"George Herman ""Babe"" Ruth","1
10000000000000011000011000001001

Используя метод Кернигана, удалим этот бит из исходного ввода, выполнив операцию «И» для всех битов:

"George Herman ""Babe"" Ruth","1
10000000000000011000011000001000

Остался новый младший бит. Повторив операцию, создадим слои масок — по одному для каждой кавычки ввода:

10000000000000011000011000001001
10000000000000011000011000000111
10000000000000011000010111111111
10000000000000011000001111111111
10000000000000010111111111111111
10000000000000001111111111111111
01111111111111111111111111111111

Видели исключающее «ИЛИ» в конечном автомате для переключения между состояниями в коде выше? Выполнив операцию исключающего «ИЛИ» для всех их вместе, мы включаем и выключаем кавычки, создавая закавыченные области:

"George Herman ""Babe"" Ruth","1
01111111111111100111100111110001

Но крайне важно, чтобы открывающая кавычка была включена в эту маску (скоро объясню почему). Выполнив операцию исключающего «ИЛИ» для предварительно вычитаемого значения с маской при вычислении маски, включаем и выключаем оставшиеся кавычки так, чтобы открывающие кавычки были включены. Вот соответствующая функция:

uint32_t find_quoted_regions(uint32_t x)
{
    uint32_t r = 0;
    while (x) {
        r ^= x;
        r ^= x - 1;
        x &= x - 1;
    }
    return r;
}

Она даёт именно то, что нужно:

"George Herman ""Babe"" Ruth","1
11111111111111101111101111110011

Важно, чтобы открывающая кавычка была включена, ведь в области, начинающейся в последнем байте, будет этот последний набор битов. Используем этот последний бит, чтобы определить, начинается ли следующий фрагмент в закавыченном состоянии. Если это так, то нужно только выполнить логическую операцию «НЕ» для всего результата, чтобы изменить закавыченные области.

Как добавить знак к 1 в набор всех битов или ничего не сделать для нуля? Выполнить логическую операцию «НЕ»:

uint32_t carry  = -(prev & 1);
uint32_t quotes = find_quoted_regions(input) ^ carry;
// ...
prev = quotes;

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

Однако допущена серьёзная ошибка. Я использую _mm256_movemask_epi8, а он помещает первый байт в младший бит! Вот как это выглядит:

1","htuR ""ebaB"" namreH egroeG"
01010000011000011000000000000001

Эффективного способа перевернуть биты нет, поэтому нужно просто поработать в другом направлении. Чтобы перевернуть биты влево от установленного бита, выполняем логическую операцию «НЕ»:

00000000000000000000000010000000 = +0x00000080
11111111111111111111111110000000 = -0x00000080

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

uint32_t find_quoted_regions(uint32_t x)
{
    uint32_t r = x;
    while (x) {
        r ^= -x ^ x;
        x &= x - 1;
    }
    return r;
}

Результат:

1","htuR ""ebaB"" namreH egroeG"
11001111110111110111111111111111

Теперь перенос зависит от старшего бита, а не от младшего:

uint32_t carry = -(prev >> 31);

Маска, менющая направление

Следующая проблема: по непонятным причинам AVX2 не включает инверсию _mm256_movemask_epi8. Чтобы преобразовать битовую маску обратно в байтовую, надо кое-что перетасовать. К счастью, я не первый сталкиваюсь с этой проблемой, поэтому не пришлось решать её с нуля.

Сначала заполняем 32-байтовый регистр повторяющимися копиями 32-битной маски:

abcdabcdabcdabcdabcdabcdabcdabcd

Перетасуем байты так, чтобы первые 8 байтов регистра имели одну и ту же копию первого байта битовой маски и т. д.:

aaaaaaaabbbbbbbbccccccccdddddddd

В байте 0 нас волнует только бит 0, в байте 1 — бит 1... в байте N — только бит N%8. Заранее вычислим маску, чтобы изолировать каждый из этих битов и сделать из битовой маски правильную побайтовую маску. К счастью, всё не так плохо, как я думал: вместо одной инструкции получилось 4.

Результаты

В моём бенчмарке встречаются случайно встречающиеся закавыченные поля, версия ОКМД обрабатывает данные со скоростью ~4 ГиБ/с — в 10 раз быстрее, чем в оригинальной реализации. Код не профилировался, но надеюсь, что ошибочные предсказания ветвлений в цикле битовой маски — это основная трудность, препятствующая гипотетическому ускорению в 32 раза.

Также в моей версии дополнительно отклоняются входные данные, содержащие два специальных управляющих символа, так как кодировка будет необратимой. По мере возможности это реализовано в ОКМД, но замедляет обработку примерно на 10%.

Что дальше: PCLMULQDQ

Джефф Лэнгдейл и другие любезно указали на PCLMULQDQ, который может вычислять маски кавычек с использованием умножения без переноса (ещё ссылка) полностью в ОКМД и без цикла. Я ещё не разобрался, как именно его применять, но он должен быть намного быстрее.

Продолжить работу с данными и прокачать ваши навыки вы сможете на наших курсах.

• Профессия Data Scientist (24 месяца)

• Курс «Алгоритмы и структуры данных» (5 месяцев)

Узнайте подробности акции.