Каков наиболее эффективный способ упаковки различных массивов numpy в двоичный файл?

по умолчанию вам следует использовать np.savez, он быстрый и может использовать сжатие. но если вы по какой-то причине этого не хотите (например, вы хотите прочитать файл с другого языка, который не может импортировать функции Python), вам не следует сериализовать массивы по строкам. вместо этого сгладьте каждый массив и запишите их один за другим следующим образом:

def arr_write(stream, arr):
    shape = np.array([len(arr.shape), *arr.shape], dtype='<i4')
    if arr.dtype == np.dtype('u1'): dtype = b' u1'
    elif arr.dtype == np.dtype('<f4'): dtype = b'<f4'
    else: raise ValueError(f'unknown dtype: {arr.dtype}')
    stream.write(dtype)
    stream.write(shape.tobytes())
    stream.write(arr.ravel().tobytes())

arrays = [a, b, c, d]
with open(file, 'wb') as f:
    f.write(len(arrays).to_bytes(4, 'little')) # optional 
    for arr in arrays: arr_write(f, arr)
    

таким образом, вся основная информация о массивах может быть получена из файла.

Это легко реализовать с помощью структурированных массивов.

def structured_save(path, a, b, c, d):
    arr = np.empty(
        len(a),
        dtype=np.dtype(
            [
                ("a", a.dtype, a.shape[1:]),
                ("b", b.dtype, b.shape[1:]),
                ("c", c.dtype, c.shape[1:]),
                ("d", d.dtype, d.shape[1:]),
            ]
        ),
    )
    arr["a"] = a
    arr["b"] = b
    arr["c"] = c
    arr["d"] = d
    arr.tofile(path)  # faster than f.write(arr.tobytes())

Это просто, но arr является копией a, b, c, d и требует дополнительных 160 МБ (32 байта x 5 миллионов строк) памяти. При таком большом использовании памяти нельзя игнорировать влияние на скорость выполнения.

Чтобы решить эту проблему, мы можем использовать фрагментированную запись.

def chunked_structured_save(path, a, b, c, d, chunk_size=100_000):
    arr = np.empty(
        chunk_size,
        dtype=np.dtype(
            [
                ("a", a.dtype, a.shape[1:]),
                ("b", b.dtype, b.shape[1:]),
                ("c", c.dtype, c.shape[1:]),
                ("d", d.dtype, d.shape[1:]),
            ]
        ),
    )

    with open(path, "wb") as f:
        for i in range(0, len(a), chunk_size):
            n_elem = len(a[i: i + chunk_size])
            arr["a"][:n_elem] = a[i: i + chunk_size]
            arr["b"][:n_elem] = b[i: i + chunk_size]
            arr["c"][:n_elem] = c[i: i + chunk_size]
            arr["d"][:n_elem] = d[i: i + chunk_size]
            arr[:n_elem].tofile(f)

Он немного сложнее, но превосходит как по использованию памяти, так и по скорости выполнения.

Тест:

import time
from io import BytesIO
from pathlib import Path

import numpy as np


def baseline(path, a, b, c, d):
    buffer = BytesIO()
    for row in zip(a, b, c, d):
        buffer.write(row[0].tobytes())
        buffer.write(row[1].tobytes())
        buffer.write(row[2].tobytes())
        buffer.write(row[3].tobytes())

    with open(path, "wb") as f:
        f.write(buffer.getvalue())


def structured_save(path, a, b, c, d):
    arr = np.empty(
        len(a),
        dtype=np.dtype(
            [
                ("a", a.dtype, a.shape[1:]),
                ("b", b.dtype, b.shape[1:]),
                ("c", c.dtype, c.shape[1:]),
                ("d", d.dtype, d.shape[1:]),
            ]
        ),
    )
    arr["a"] = a
    arr["b"] = b
    arr["c"] = c
    arr["d"] = d
    arr.tofile(path)


def chunked_structured_save(path, a, b, c, d, chunk_size=100_000):
    arr = np.empty(
        chunk_size,
        dtype=np.dtype(
            [
                ("a", a.dtype, a.shape[1:]),
                ("b", b.dtype, b.shape[1:]),
                ("c", c.dtype, c.shape[1:]),
                ("d", d.dtype, d.shape[1:]),
            ]
        ),
    )

    with open(path, "wb") as f:
        for i in range(0, len(a), chunk_size):
            n_elem = len(a[i: i + chunk_size])
            arr["a"][:n_elem] = a[i: i + chunk_size]
            arr["b"][:n_elem] = b[i: i + chunk_size]
            arr["c"][:n_elem] = c[i: i + chunk_size]
            arr["d"][:n_elem] = d[i: i + chunk_size]
            arr[:n_elem].tofile(f)


def main():
    n = 5_000_000
    a = np.arange((n * 3)).reshape(n, 3).astype(np.float32)
    b = np.arange((n * 3)).reshape(n, 3).astype(np.float32)
    c = np.arange((n * 4)).reshape(n, 4).astype(np.uint8)
    d = np.arange((n * 4)).reshape(n, 4).astype(np.uint8)

    candidates = [
        baseline,
        structured_save,
        chunked_structured_save,
    ]
    name_len = max(len(f.__name__) for f in candidates)

    path = Path("temp.bin")
    baseline(path, a, b, c, d)
    expected = path.read_bytes()

    start = time.perf_counter()
    np.savez(path, a=a, b=b, c=c, d=d)
    print(f"{'np.savez (reference)':{name_len}}: {time.perf_counter() - start}")

    for f in candidates:
        started = time.perf_counter()
        f(path, a, b, c, d)
        elapsed = time.perf_counter() - started
        print(f"{f.__name__:{name_len}}: {elapsed}")
        assert path.read_bytes() == expected, f"{f.__name__} failed"


if __name__ == "__main__":
    main()

Результат:

np.savez (reference)   : 0.8915687190001336
baseline               : 5.309623991999615
structured_save        : 0.2205286160005926
chunked_structured_save: 0.18220391599970753

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

Вот результат на моей другой машине:

np.savez (reference)   : 0.1376536000170745
baseline               : 3.5804199000122026
structured_save        : 0.4771533999883104
chunked_structured_save: 0.13709589999052696

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

Следуя ответу @ken, я изменил код, чтобы сделать его более универсальным для использования kwargs:

def chunked_structured_save(save_path, chunk_size=100000, **kwargs):
    arr = np.empty(
        chunk_size,
        dtype=np.dtype([(k, v.dtype, v.shape[1:]) for k, v in kwargs.items()]),
    )

    n_rows = len(list(kwargs.values())[0])
    with open(save_path, "wb") as f:
        for i in range(0, n_rows, chunk_size):
            n_elem = n_rows - i if (i + chunk_size) > n_rows else chunk_size
            for k, v in kwargs.items():
                arr[k][:n_elem] = v[i: i + chunk_size]

            arr[:n_elem].tofile(f)


chunked_structured_save(path, a=a, b=b, c=c, d=d)