Расширяемый Postgres

На прошедшем PGConf.Russia был доклад про расширение MobilityDB, а Андрей Бородин предложил идею расширять методы индексов под задачу.

Продолжу тему с расширением Postgres под решаемую задачу на примере расширения сделанного в рамках HighLoad Cup 2018, код доступен на GithHub. На хабре уже есть статья от blackmaster. Расширение добавляет два типа с поддержкой btree и GIN индексов.

Начальная схема:

CREATE TABLE accounts
(
  id           INTEGER PRIMARY KEY,
  email        VARCHAR(100) UNIQUE,
  fname        SMALLINT,
  sname        VARCHAR(50),
  phone        VARCHAR(16) UNIQUE,
  birth        TIMESTAMP NOT NULL,
  country      SMALLINT,
  city         SMALLINT,
  email_domain SMALLINT,
  joined       TIMESTAMP NOT NULL,
  status       status,
  sex          sex       NOT NULL,
  interests    SMALLINT[],
  premium      tsrange,
  likes_ids    INTEGER[],
  likes_tss    INTEGER[],
  CHECK ( joined >= '01.01.2011'::timestamp ),
  CHECK ( joined <= '01.01.2018'::timestamp )
);

CREATE INDEX IF NOT EXISTS interests_ids_index ON accounts USING GIN(interests);

Количество записей 1 300 000, а размеры интересующих отношений:

Конечная схема:

CREATE UNLOGGED TABLE accounts
(
  id           INTEGER PRIMARY KEY,
  email        VARCHAR(100) UNIQUE,
  fname        SMALLINT,
  sname        VARCHAR(50),
  phone        phone UNIQUE,
  birth        TIMESTAMP NOT NULL,
  country      SMALLINT,
  city         SMALLINT,
  email_domain SMALLINT,
  joined       TIMESTAMP NOT NULL,
  status       status,
  sex          sex       NOT NULL,
  interests    bit128,
  premium      tsrange,
  likes_ids    INTEGER[],
  likes_tss    INTEGER[],
  CHECK ( joined >= '01.01.2011'::timestamp ),
  CHECK ( joined <= '01.01.2018'::timestamp )
);

Размеры:

Были добавлены два типа: phone и bit128.

phone

Номер телефона имеет вид 8(929)5481819, восьмерку можно отбросить. Код оператора укладывается в 10 бит, номер 24 бит, получается нужно 5 байт. Не очень удобное число из-за выравнивания данных в памяти. На вопрос как лучше укладывать данные в 5, 6 или 8 байт, ответ может дать только тесты.

Если следовать примеру из документации, то нужно немного. Определить тип:

class Phone : public PAllocNew<Phone>
{
public:

    bool fromCString(char* str) noexcept;
    char* toCString() const noexcept;

    int code() const noexcept;

    bool operator==(const Phone& b) const noexcept;
    bool operator<(const Phone& b) const noexcept;
    bool operator<=(const Phone& b) const noexcept;
    bool operator>(const Phone& b) const noexcept;
    bool operator>=(const Phone& b) const noexcept;

private:

    uint64_t m_data = 0;
};

#define GETARG_PHONE(x) reinterpret_cast<Phone*>(PG_GETARG_POINTER(x))

PAllocNew вспомогательный класс перегружающий new:

template<typename T>
class PAllocNew
{
public:

    static void* operator new(std::size_t count, const std::nothrow_t& tag)
    {
        return palloc(count);
    }
};

Память, выделенная через palloc, освобождается при завершении транзакции. Добавить функции ввода/вывода:

Datum phone_in(PG_FUNCTION_ARGS)
{
    char* str = PG_GETARG_CSTRING(0);

    auto v = new(std::nothrow) Phone;
    if (!v->fromCString(str))
    {
        ereport(ERROR,
            (errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
                errmsg(
                    "invalid input syntax for phone: \"%s\"",
                    str
                )));
    }

    PG_RETURN_POINTER(v);
}

Datum phone_out(PG_FUNCTION_ARGS)
{
    const auto phone = GETARG_PHONE(0);

    PG_RETURN_CSTRING(phone->toCString());
}

И регистрация типа:

CREATE TYPE phone;

CREATE OR REPLACE FUNCTION phone_in ( cstring )
RETURNS phone AS
'libhighloadcup'
LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

CREATE OR REPLACE FUNCTION phone_out ( phone )
RETURNS cstring AS
'libhighloadcup'
LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

CREATE TYPE phone (
   INTERNALLENGTH = 8,
   INPUT = phone_in,
   OUTPUT = phone_out,
   ALIGNMENT = int4,
   COLLATE TRUE
);

Т.к. новый тип имеет размер не более 8 байт, то можно передавать не по ссылке, а по значению. В этом случае в CREATE TYPE нужно указать флаг PASSEDBYVALUE. Или использовать LIKE:

CREATE TYPE phone (
   INPUT = phone_in,
   OUTPUT = phone_out,
   LIKE = int8,
   COLLATE TRUE
);

Тогда INTERNALLENGTH, ALIGNMENT и PASSEDBYVALUE унаследуются от int8.

btree индекс

Поддежка уникальности поля осуществляется созданием уникального индекса B-дерева. Делается это через классы операторов и стратегий.

CREATE OPERATOR = (
   PROCEDURE = phone_equal,
   LEFTARG = phone,
   RIGHTARG = phone,
   COMMUTATOR = =,
   NEGATOR = !=
);

CREATE OPERATOR < (
  PROCEDURE = phone_lt,
  LEFTARG = phone,
  RIGHTARG = phone,
  NEGATOR = >
);

CREATE OPERATOR <= (
  PROCEDURE = phone_le,
  LEFTARG = phone,
  RIGHTARG = phone
);

CREATE OPERATOR >= (
  PROCEDURE = phone_ge,
  LEFTARG = phone,
  RIGHTARG = phone
);

CREATE OPERATOR > (
  PROCEDURE = phone_gt,
  LEFTARG = phone,
  RIGHTARG = phone,
  NEGATOR = <
);

CREATE OPERATOR CLASS btree_phone_ops
  DEFAULT FOR TYPE phone USING btree AS
   OPERATOR        1        <,
   OPERATOR        2        <=,
   OPERATOR        3        =,
   OPERATOR        4        >=,
   OPERATOR        5        >,
   FUNCTION        1       phone_equal_internal ( phone, phone );

Операторы возвращают bool, а phone_equal_internal int:

Datum phone_equal_internal(PG_FUNCTION_ARGS)
{
    const auto a = GETARG_PHONE(0);
    const auto b = GETARG_PHONE(1);
    if (*a < *b)
    {
        PG_RETURN_INT32(-1);
    }

    if (*a > *b)
    {
        PG_RETURN_INT32(1);
    }

    PG_RETURN_INT32(0);
}

Все это дало небольшое уменьшение данных:

Количество аккаунтов с номерами 533 289, что составляет 41%. По крайней мере нет сравнения строк, при работе с индексом.

bit128

Захотелось иметь битовое поле с операциями пересечения(&&), вхождения(<@) и с поддержкой GIN. Достаточно было 96 бит, но пошел по простому пути и взял uint128.

class BitSet128: public PAllocNew<BitSet128>
{
public:

    bool haveIntersection(const BitSet128& b) const noexcept;
    bool contains(const BitSet128& b)  const noexcept;

    bool operator==(const BitSet128& b) const noexcept;
    bool operator<(const BitSet128& b) const noexcept;
    bool operator>(const BitSet128& b) const noexcept;

    bool fromCString(char* str) noexcept;
    char* toCString() const noexcept;

    std::vector<uint8_t> keys() const;

private:

    uint128 m_bits = 0;
};

#define GETARG_BIT128(x) reinterpret_cast<BitSet128*>(PG_GETARG_POINTER(x))

Регистрация типа и операции:

CREATE TYPE bit128 (
   INTERNALLENGTH = 16,
   INPUT = bit128_in,
   OUTPUT = bit128_out,
   ALIGNMENT = int4
);

CREATE OPERATOR = (
   PROCEDURE = bit128_equal,
   LEFTARG = bit128,
   RIGHTARG = bit128,
   COMMUTATOR = =,
   NEGATOR = !=
);

CREATE OPERATOR && (
   PROCEDURE = bit128_intersection,
   LEFTARG = bit128,
   RIGHTARG = bit128
);

CREATE OPERATOR @> (
   PROCEDURE = bit128_containts,
   LEFTARG = bit128,
   RIGHTARG = bit128
);

Generalized Inverted Index/GIN

Про GIN в Postgres хорошо написал Егор Рогов erogov в статье Индексы в PostgreSQL — 7, применяя к задаче полнотекстового поиска. Документация по расширяемости GIN говорит о том, что достаточно реализовать 4 функции:

Извлечение массива ключей из индексируемого объекта:

Datum *extractValue(Datum itemValue, int32 *nkeys, bool **nullFlags)

Извлекает ключи для запрашиваемого значения:

Datum *extractQuery(Datum query, int32 *nkeys, StrategyNumber n, bool **pmatch, Pointer **extra_data, bool **nullFlags, int32 *searchMode)

Например в запросе:

SELECT * FROM test WHERE a && ARRAY[1, 2, 3]

extractQuery применяется к массиву ARRAY[1, 2, 3], а ключами могут быть значения 1, 2, 3.
Проверяет удовлетворяет ли значение запросу:

bool consistent(bool check[], StrategyNumber n, Datum query, int32 nkeys, Pointer extra_data[], bool *recheck, Datum queryKeys[], bool nullFlags[])

Так как извлеченные ключи хранятся в дереве поиска, то нужна функция сравнения:

int compare(Datum a, Datum b)

API может показаться избыточным, но предоставляет все, что может пригодиться. Перейдем к реализации. Извлечение ключей:

Datum gin_extract_value_bit128(PG_FUNCTION_ARGS)
{
    auto bitset = GETARG_BIT128(0);
    const auto bits = bitset->keys();

    int32* nentries = (int32*) PG_GETARG_POINTER(1);
    *nentries = bits.size();
    Datum* entries = NULL;

    entries = (Datum*) palloc(sizeof(Datum) * bits.size());
    for (int i = 0; i < bits.size(); ++i)
    {
        entries[i] = DatumGetInt16(bits[i]);
    }

    PG_RETURN_POINTER(entries);
}

В выходной параметр nentries записываем количество ключей и возвращаем массив ключей entries. Сравнение ключей:

Datum bit128_cmp(PG_FUNCTION_ARGS)
{
    const auto a = PG_GETARG_INT16(0);
    const auto b = PG_GETARG_INT16(1);

    if (a < b)
    {
        PG_RETURN_INT32(-1);
    }

    if (a > b)
    {
        PG_RETURN_INT32(1);
    }

    PG_RETURN_INT32(0);
}

Этих двух функций достаточно для создания индекса. Остальные функции используются при поиске по индексу. Извлечение ключей из запроса:

Datum gin_extract_query_bit128(PG_FUNCTION_ARGS)
{
    const auto a = GETARG_BIT128(0);
    int32* nentries = (int32*) PG_GETARG_POINTER(1);
    StrategyNumber strategy = PG_GETARG_UINT16(2);
    int32* searchMode = (int32*) PG_GETARG_POINTER(6);

    Datum* res = nullptr;

    const auto keys = a->keys();

    *nentries = keys.size();
    res = (Datum*) palloc(sizeof(Datum) * keys.size());
    for (int i = 0; i < keys.size(); ++i)
    {
        res[i] = DatumGetInt16(keys[i]);
    }

    switch (strategy)
    {
        case RTOverlapStrategyNumber:  //  &&
            if (*nentries > 0)
            {
                *searchMode = GIN_SEARCH_MODE_DEFAULT;
            }
            else
            {
                *searchMode = GIN_SEARCH_MODE_ALL;
            }
            break;
        case RTContainsStrategyNumber: //  @>
            if (*nentries > 0)
            {
                *searchMode = GIN_SEARCH_MODE_DEFAULT;
            }
            else
            {
                *searchMode = GIN_SEARCH_MODE_ALL;
            }
            break;
    }

    PG_RETURN_POINTER(res);
}

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

Datum gin_bit128_consistent(PG_FUNCTION_ARGS)
{
    bool* check = (bool*) PG_GETARG_POINTER(0);

    StrategyNumber strategy = PG_GETARG_UINT16(1);
    int32 nkeys = PG_GETARG_INT32(3);
    bool* recheck = (bool*) PG_GETARG_POINTER(5);
    bool res = true;
    switch (strategy)
    {
        case RTOverlapStrategyNumber:  //  &&
        {
            for (int i = 0; i < nkeys; ++i)
            {
                if (check[i])
                {
                    res = true;
                }
            };
            *recheck = false;
        }
            break;
        case RTContainsStrategyNumber: //  @>
        {
            for (int i = 0; i < nkeys; ++i)
            {
                if (check[i])
                {
                    res = true;
                }
            };
            *recheck = true;
        }
            break;
    }

    PG_RETURN_BOOL(res);
}

Возвращает true, если индексированный объект удовлетворяет оператору запроса с номером стратегии. Массив check имеет длину nkeys, что соотвествует числу ключей возвращенной gin_extract_query_bit128. Проверка check[i] = true означает, что i-ый ключ из gin_extract_query_bit128 присутствует в индексированном объекте. Этого достаточно, для пересечения, но т.к. в GIN мы не работаем с самими значениями, то для стратегии вхождения выставляем recheck в true. Это приведет к вызову соотвествующего оператора, для проверки результата. Класс оператора:

CREATE OR REPLACE FUNCTION gin_extract_value_bit128 ( bit128, internal, internal)
RETURNS internal AS
'libhighloadcup'
LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

CREATE OR REPLACE FUNCTION gin_extract_query_bit128(bit128, internal, int2, internal, internal)
RETURNS internal AS
'libhighloadcup'
LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

CREATE OR REPLACE FUNCTION gin_bit128_consistent(internal, int2, anyelement, int4, internal, internal)
RETURNS bool AS
'libhighloadcup'
LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

CREATE OPERATOR CLASS bit128_ops
DEFAULT FOR TYPE bit128 USING gin
AS
    OPERATOR        3       &&,
    OPERATOR        6       =,
    OPERATOR        7       @>,
    FUNCTION        1       bit128_cmp (int2, int2 ),
    FUNCTION        2       gin_extract_value_bit128(bit128, internal, internal),
    FUNCTION        3       gin_extract_query_bit128(bit128, internal, int2, internal, internal),
    FUNCTION        4       gin_bit128_consistent(internal, int2, anyelement, int4, internal, internal),
STORAGE  int2;

Добавился STORAGE, он указывает какой тип данных истользуется для ключей. Какой результат стал по диску:

Интересный резльтат, потому что есть нюансы:

• физическое хранение базы данных;
• распределение данных.

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

Распределение по интересам имеет следующий вид:

Честно не знаю как релизован SMALLINT[], но предположим что 4 байта на размер и 2 байта на значение. Тогда в 16 байт можно уместить массив из 6 элементов. И 98% элементов укладывались бы в 16 байт и разницы в 16 MB. Вроде бы тип bit128 избыточен, но и стандартный тип избыточен на этом наборе данных.