Как хранить массивы NumPy в PostgreSQL: tobytes vs msgpack-numpy

Сравниваем способы хранения массивов NumPy в PostgreSQL: быстрый tobytes с shape/dtype и удобный msgpack-numpy. Примеры на psycopg и рекомендации по выбору.

Сохранять массивы NumPy в PostgreSQL вроде бы просто, пока не упираешься в метаданные: сырые байты не содержат ни формы (shape), ни типа (dtype). Распространённый обходной путь — сериализовать с помощью NumPy.tobytes и складывать shape и dtype в отдельных столбцах. Это работает, но трение ощущается, когда делаешь это снова и снова в разных таблицах и местах вызова.

Базовый вариант: хранение массивов как сырых байтов

Подход ниже использует psycopg, numpy.tobytes и белый список целевых таблиц. Он выполняет вставку или обновление, а при чтении восстанавливает массивы, комбинируя сохранённые байты с формой и типом.

import psycopg
from psycopg import sql
import os
import orjson
import numpy as np
pg_conn_opts = {
    'dbname': os.getenv("PG_DATABASE"),
    'user': os.getenv("PG_USERNAME"),
    'password': os.getenv("PG_PASSWORD"),
    'host': os.getenv("PG_HOST"),
    'connect_timeout': os.getenv("PG_CONNECT_TIMEOUT")
}
def persist_ndarray(arr, arr_key, target_table):
    if target_table in ['table_a', 'table_b', 'table_c']:
        with psycopg.connect(**pg_conn_opts) as cxn:
            with cxn.cursor() as cur:
                try:
                    cur.execute(sql.SQL(
                        """
                        INSERT INTO {tbl}(array_id, np_array_bytes, np_array_shape, np_array_dtype)
                        VALUES (%s, %s, %s, %s)
                        ON CONFLICT (array_id) DO UPDATE SET
                        np_array_bytes = EXCLUDED.np_array_bytes,
                        np_array_shape = EXCLUDED.np_array_shape,
                        np_array_dtype = EXCLUDED.np_array_dtype
                        """
                    ).format(tbl=sql.Identifier(target_table)),
                    [arr_key, arr.tobytes(), orjson.dumps(arr.shape), str(arr.dtype)])
                    cxn.commit()
                except:
                    print("Error while saving NumPy array to database.")
                finally:
                    cxn.close()
    else:
        print("You're attempting to write arrays to a non-approved table.")
def fetch_ndarray(arr_key, target_table):
    if target_table in ['table_a', 'table_b', 'table_c']:
        with psycopg.connect(**pg_conn_opts) as cxn:
            with cxn.cursor() as cur:
                cur.execute(sql.SQL(
                    """
                    SELECT np_array_bytes, np_array_shape, np_array_dtype
                    FROM {tbl}
                    WHERE array_id = %s
                    """
                ).format(tbl=sql.Identifier(target_table)), [arr_key])
                row = cur.fetchone()
                if row:
                    blob, shp, dt = row
                    return np.frombuffer(blob, dtype=dt).reshape(orjson.loads(shp))

Почему это тяжелее, чем должно быть

tobytes быстрый, но сырые байты ничего не знают о массиве, из которого пришли. Поэтому приходится отдельно отслеживать shape и dtype, протаскивать их через SQL и вручную собирать при чтении. Это надёжно, но многословно, и каждый вызов повторяет одну и ту же хореографию. Если нужно хранить массивы в нескольких таблицах, накладные расходы растут с каждым вставкой и выборкой.

Простой альтернативный вариант: msgpack-numpy

msgpack-numpy умеет сериализовать массив NumPy в компактный бинарный блок, в котором уже есть shape и dtype. Значит, можно сохранять одно бинарное значение и декодировать его обратно в исходный массив без жонглирования дополнительными столбцами. Для ориентира:

Он немного медленнее, чем A.tobytes(), но компенсирует это удобством — вместе с данными сохраняются shape и dtype. К тому же формат получается меньше и быстрее, чем нативный формат NumPy np.save().

Есть и практическое наблюдение по тестам на крупных массивах:

tobytes примерно в 2 раза быстрее при упаковке массивов с миллионом чисел и намного быстрее (в 1000–2000 раз) при распаковке. Однако оба варианта достаточно быстры, так что это важно лишь для очень больших массивов.

Иными словами, если важны удобство и минимальные накладные расходы, msgpack-numpy отлично подходит. Если же нужна абсолютная скорость на очень больших массивах, сырой путь через tobytes остаётся привлекательным.

Обновлённый код: упаковка и распаковка с msgpack-numpy

Фрагмент ниже показывает, как сохранять и загружать с помощью msgpack-numpy, не меняя модель взаимодействия с базой. Байтам отводится столбец np_array_bytes; shape и dtype по-прежнему записываются как раньше, хотя в упакованной нагрузке они уже присутствуют.

import psycopg
from psycopg import sql
import os
import orjson
import numpy as np
import msgpack
import msgpack_numpy as mpn
pg_conn_opts = {
    'dbname': os.getenv("PG_DATABASE"),
    'user': os.getenv("PG_USERNAME"),
    'password': os.getenv("PG_PASSWORD"),
    'host': os.getenv("PG_HOST"),
    'connect_timeout': os.getenv("PG_CONNECT_TIMEOUT")
}
def persist_ndarray_msgpack(arr, arr_key, target_table):
    if target_table in ['table_a', 'table_b', 'table_c']:
        with psycopg.connect(**pg_conn_opts) as cxn:
            with cxn.cursor() as cur:
                try:
                    packed = msgpack.packb(arr, default=mpn.encode)
                    cur.execute(sql.SQL(
                        """
                        INSERT INTO {tbl}(array_id, np_array_bytes, np_array_shape, np_array_dtype)
                        VALUES (%s, %s, %s, %s)
                        ON CONFLICT (array_id) DO UPDATE SET
                        np_array_bytes = EXCLUDED.np_array_bytes,
                        np_array_shape = EXCLUDED.np_array_shape,
                        np_array_dtype = EXCLUDED.np_array_dtype
                        """
                    ).format(tbl=sql.Identifier(target_table)),
                    [arr_key, packed, orjson.dumps(arr.shape), str(arr.dtype)])
                    cxn.commit()
                except:
                    print("Failed to persist NumPy array via msgpack.")
                finally:
                    cxn.close()
    else:
        print("Refusing to write to unexpected table.")
def fetch_ndarray_msgpack(arr_key, target_table):
    if target_table in ['table_a', 'table_b', 'table_c']:
        with psycopg.connect(**pg_conn_opts) as cxn:
            with cxn.cursor() as cur:
                cur.execute(sql.SQL(
                    """
                    SELECT np_array_bytes, np_array_shape, np_array_dtype
                    FROM {tbl}
                    WHERE array_id = %s
                    """
                ).format(tbl=sql.Identifier(target_table)), [arr_key])
                row = cur.fetchone()
                if row:
                    packed, shp, dt = row
                    return msgpack.unpackb(packed, object_hook=mpn.decode)

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

Почему это важно

Выбор способа сериализации определяет объём «клеевого» кода, который придётся поддерживать. Сырые байты переносят эту сложность в приложение через дополнительные метаданные. msgpack-numpy концентрирует её на этапе сериализации. Компромисс — в нюансах производительности: tobytes может быть быстрее, особенно при декодировании очень больших массивов, тогда как msgpack-numpy делает ставку на удобство и при этом меньше и быстрее, чем np.save.

Выводы

Если приоритет — чистый и легко поддерживаемый конвейер через несколько таблиц PostgreSQL, msgpack-numpy — практичный вариант, который несёт с собой shape и dtype. Если критична пропускная способность и нужна максимальная скорость кодирования/декодирования на очень больших массивах, можно остаться на tobytes плюс явные shape и dtype. В любом случае ограничивайте доступ к таблицам списком допустимых и соблюдайте единообразие в работе с dtype, чтобы читать ровно то, что сохраняете.