Como armazenar arquivos de texto pequenos de forma eficiente

Preparação

Vamos criar um conjunto de dados de teste com muitos arquivos de texto.

import os
import json
import random
import string
from datetime import datetime

workspace_dir = '/tmp/data'
output_dir = f'{workspace_dir}/jsons'
os.makedirs(output_dir, exist_ok=True)


def generate_random_string(length):
    return ''.join(random.choices(string.ascii_letters + string.digits + string.punctuation + ' ', k=length))


for i in range(1000000):
    data = {
        'text': generate_random_string(random.randint(0, 1000)),
        'created_at': int(datetime.now().timestamp())
    }
    filename = os.path.join(output_dir, f'{i}.json')
    with open(filename, 'w') as json_file:
        json.dump(data, json_file, separators=(',', ':'))

Por exemplo, vamos gerar 1.000.000 de arquivos json com a estrutura (para eficiência, eles não têm espaços nos arquivos, o que é alcançado com o parâmetro separators=(',', ':') no método json.dump):

{
  "text": "string aleatória",
  "created_at": 1727625937
}

Vamos medir o tamanho real:

du -sk /tmp/data/jsons
4000000 jsons

São 4Gb, no entanto, o tamanho deveria ser <= 1Gb (1000000 * 1Kb). Isso ocorre porque uso um sistema de arquivos com tamanho de registro de 4Kb (uso OS X, mas no Linux a situação será a mesma), então tenho 1000000 * 4Kb = 4Gb.

Criando arquivo de texto mesclado

Precisamos de um gerador primeiro para iterar os arquivos:

import os
import json


def json_file_reader(directory):
    json_files = [filename for filename in os.listdir(directory) if filename.endswith('.json')]
    sorted_files = sorted(json_files, key=lambda x: int(os.path.splitext(x)[0]))
    for filename in sorted_files:
        file_path = os.path.join(directory, filename)
        if os.path.isfile(file_path):
            with open(file_path, 'r', encoding='utf-8') as json_file:
                yield os.path.splitext(filename)[0], json_file.read()

Além disso, eu os ordeno por nome de arquivo, mas isso não é necessário.

Script mesclado, precisamos adicionar id ao dicionário json também, caso contrário seria impossível entender o arquivo fonte:

with open(f'{workspace_dir}/merged.json', 'w', encoding='utf-8') as infile:
    for id, data in json_file_reader(output_dir):
        dict = json.loads(data)
        dict['id'] = int(id)
        infile.write(f'{json.dumps(dict, ensure_ascii=False, separators=(',', ':'))}\n')

du -k /tmp/data/merged.json 
557060	merged.json

Agora é apenas 557Mb! Mais de 7x de diferença! Não desperdiçamos clusters mais. O tamanho médio do arquivo é de 557b que combina perfeitamente com nossa função de conteúdo random.randint(0, 1000).

Formato binário

Agora vamos tentar usar uma estrutura binária otimizada.

Primeiro precisamos serializar nossa estrutura:

import struct


def pack(data):
    text_bytes = data['text'].encode('utf-8')
    format_string = f'iiH{len(text_bytes)}s'
    return struct.pack(format_string, data['id'], data['created_at'], len(text_bytes), text_bytes)


def unpack(data):
    offset = 10
    i, ts, l = struct.unpack('iiH', data[:offset])
    text = struct.unpack(f'{l}s', data[offset:offset + l])[0].decode()
    return {
        'id': i,
        'created_at': ts,
        'text': text
    }


# test
packed = pack({"id": 1, "created_at": 1727625937, "text": "Hey!"})
print(f"{packed} -> {unpack(packed)}")

Agora crie um arquivo binário:

# %%
with open(f'{workspace_dir}/merged.struct.bin', 'wb') as infile:
    for id, data in json_file_reader(output_dir):
        dict = json.loads(data)
        dict['id'] = int(id)
        infile.write(pack(dict))

du -k /tmp/data/merged.struct.bin 
507908	merged.struct.bin

Reduzimos ainda mais o arquivo para ~ 508M.

Messagepack

Às vezes, não é possível ter um esquema predefinido, especialmente se quisermos armazenar json arbitrário. Vamos tentar o MessagePack.

Instale primeiro: pip install msgpack.

Agora gere o arquivo binário mesclado:

import msgpack

with open(f'{workspace_dir}/merged.msgpack.bin', 'wb') as infile:
    for id, data in json_file_reader(output_dir):
        dict = json.loads(data)
        dict['id'] = int(id)
        infile.write(msgpack.packb(dict))

du -k /tmp/data/merged.msgpack.bin 
524292	merged.msgpack.bin

O tamanho é um pouco maior do que com o protocolo personalizado ~524M, mas acho que é um preço razoável por não ter esquema. O Messagepack possui uma ferramenta muito útil para acesso aleatório:

with open(f'{workspace_dir}/merged.msgpack.bin', 'rb') as file:
    file.seek(0)

    unpacker = msgpack.Unpacker(file)
    obj = unpacker.unpack()
    print(obj)
    print(unpacker.tell())  # offset atual

Compressão

Vamos ainda mais longe e tentemos comprimir esses arquivos com lz4 e zstd:

# dir
-> time tar -cf - jsons | lz4 - jsons.tar.lz4
Compressed 3360880640 bytes into 628357651 bytes ==> 18.70%                    
tar -cf - jsons  19.93s user 299.75s system 61% cpu 8:38.46 total
lz4 - jsons.tar.lz4  2.20s user 1.41s system 0% cpu 8:38.46 total
-> time tar -cf - jsons | zstd -o jsons.tar.zst

-> time lz4 merged.json
Compressed 558739918 bytes into 522708179 bytes ==> 93.55%                     
lz4 merged.json  0.47s user 0.35s system 152% cpu 0.534 total

Resultados

nome	tamanho bruto	diminuído	tamanho lz4	tempo lz4	tamanho zstd	tempo zstd
Arquivos no diretório	4G	1x	628M	8:38	464M	8:55
Arquivo de texto mesclado	557M	7.18x	523M	0.53s	409M	1.62s
Arquivo bin mesclado (struct)	508M	7.87x	510M	0.46s	410M	1.47s
Arquivo bin mesclado (msgpack)	524M	7.63x	514M	0.36s	411M	1.6s