Comparación de almacenes de pares clave-valor embebidos para Python

ACTUALIZADO: Se agregó RocksDict

Necesito almacenar una gran cantidad de archivos de texto pequeños (~3 mil millones y crece en 70 millones cada día) con tamaños de 100B a unos pocos kB en mi proyecto de tghub. Solo hay 2 requisitos:

• acceso rápido por id (cada archivo tiene una clave única)
• almacenarlos lo más compactamente posible, idealmente con compresión

En realidad, puedo hacer una estructura jerárquica y simplemente almacenarlos en el sistema de archivos (también puedo usar ZFS para la compresión), sin embargo, temo desperdiciar demasiado espacio porque el tamaño promedio del archivo es de solo 1Kb.

Soluciones como Cassandra, Hbase son exageradas para mí. No necesito sus características para nada. Redis no es adecuado porque mantiene todos los datos en memoria. Probemos soluciones embebidas:

1. Sqlite (sería demasiado lento a priori, por la naturaleza de RDBMS)
2. Sqlitedict (sería demasiado lento a priori, porque es un envoltorio para sqlite)
3. Pysos
4. LevelDB
5. Shelve
6. Diskcache
7. Lmdb
8. RocksDict

Comparación de características

Nombre	Seguros a nivel de hilos	Seguros a nivel de procesos	Serialización
pysos	No	No	Personalizado
LevelDB	Sí	No	Ninguno
Shelve	No	No	Pickle
Diskcache	Sí	Sí	Personalizable
Lmdb	Sí	Sí	Ninguno
RocksDict	Sí	Sí	Personalizable

• Lmdb soporta lecturas concurrentes, pero la operación de escritura es de un solo hilo
• RocksDict soporta lecturas concurrentes a través de un índice secundario
• RocksDict soporta rocksdb y speedb (se considera una versión mejorada de rocksdb)

Preparación

Script que genera 1 000 000 archivos 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, 2000)),
        '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, indent=4)

Genera archivos con este esquema:

{
  "text": "cadena aleatoria con longitud de 0 a 2000",
  "created_at": 1727290164
}

Necesitamos un generador para leer los archivos preparados:

def json_file_reader(directory):
    for filename in os.listdir(directory):
        file_path = os.path.join(directory, filename)
        if os.path.isfile(file_path) and filename.endswith('.json'):
            with open(file_path, 'r') as json_file:
                yield os.path.splitext(filename)[0], json_file.read()

Además, sería bueno comparar los resultados con un generador ordenado (asc):

def json_file_reader_sorted(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()

Instalar librerías de Python:

pip install pysos
pip install diskcache
pip install plyvel-ci # para leveldb
pip install lmdb
pip install speedict # RocksDict

Scripts de prueba

Pysos

import pysos

pysos_dir = f'{workspace_dir}/pysos'
db = pysos.Dict(pysos_dir)
for id, data in json_file_reader(output_dir):
    db[id] = data

Shelve

import shelve

shelve_dir = f'{workspace_dir}/shelve'
with shelve.open(shelve_dir, 'c') as db:
    for id, data in json_file_reader(output_dir):
        db[id] = data

Diskcache

import diskcache as dc

diskcache_dir = f'{workspace_dir}/diskcache'
cache = dc.Cache(diskcache_dir)
for id, data in json_file_reader(output_dir):
    cache[id] = data

LevelDB

import plyvel

leveldb_dir = f'{workspace_dir}/leveldb'
with plyvel.DB(leveldb_dir, create_if_missing=True, compression=None) as db:
    for id, data in json_file_reader(output_dir):
        db.put(int(id).to_bytes(4, 'big'), data.encode())

Leveldb con compresión

import plyvel

leveldb_snappy_dir = f'{workspace_dir}/leveldb_snappy'
with plyvel.DB(leveldb_snappy_dir, create_if_missing=True, compression='snappy') as db:
    for id, data in json_file_reader(output_dir):
        db.put(int(id).to_bytes(4, 'big'), data.encode())

Lmdb

import lmdb

lmdb_dir = f'{workspace_dir}/lmdb'
# reservemos 100Gb
with lmdb.open(lmdb_dir, 10 ** 11) as env:
    with env.begin(write=True) as txn:
        for id, data in json_file_reader(output_dir):
            txn.put(int(id).to_bytes(4, 'big'), data.encode())

RocksDict

from speedict import Rdict

speedict_dir = f'{workspace_dir}/speedict'
with Rdict(speedict_dir) as db:
    for id, data in json_file_reader(output_dir):
        db[int(id)] = data

versión comprimida:

from rocksdict import Rdict, Options, DBCompressionType

def db_options():
    opt = Options()
    opt.set_compression_type(DBCompressionType.zstd())
    return opt

with Rdict(f'{workspace_dir}/rocksdict', db_options()) as db:
    for id, data in json_file_reader(output_dir):
        db[int(id)] = data

Para usar speedb solo necesitamos cambiar la importación de rocksdict a speedict

Resultados

Revisé el tamaño de cada conjunto de datos con el comando de terminal du -sh $dataset

nombre	espacio ocupado	tiempo de ejecución
Archivos sin procesar	3.8G	4m 25s
Un archivo de texto	1.0G	-
Archivo de texto comprimido	820Mb	-
Pysos	1.1G	4m 37s
Shelve	-	-
Diskcache	1.0Gb	7m 29s
LevelDB	1.0Gb	5m 2s
LevelDB(snappy)	1.0Gb	5m 16s
Lmdb	1.1Gb	4m 9s
Lmdb (ordenado)	1.5Gb	1m 27s
RocksDict (rocksdb)	1.0Gb	4m 26s
RocksDict (rocksdb, ordenado)	1.0Gb	1m 31s
RocksDict (rocksdb, ordenado, comprimido)	854Mb	1m 31s
RocksDict (speedb)	1.0Gb	4m 14s
RocksDict (speedb, ordenado, comprimido)	854Mb	1m 39s

• Desafortunadamente shelve falló después de 18s con HASH: Out of overflow pages. Increase page size.
• LevelDB tiene el mismo tamaño con/sin compresión pero el tiempo de ejecución es diferente.
• Se espera que Lmdb sea más grande que LevelDB. Lmdb usa B+tree (las actualizaciones ocupan más espacio), otros usan LSM-tree.
• Para la compresión, usé zstd

Ajuste adicional de Lmdb

Formato binario

Probemos Cap'n Proto, parece prometedor.

1. Instalar paquete del sistema, en mi caso (os x): brew install cproto
2. Instalar paquete de python: pip install pycapnp

Ahora necesitamos un esquema:

archivo: msg.capnp

@0xd9e822aa834af2fe;

struct Msg {
  createdAt @0 :Int64;
  text @1 :Text;
}

Ahora podemos importarlo y usarlo en nuestra aplicación:

import msg_capnp as schema

lmdb_capnp_dir = f'{workspace_dir}/lmdb_capnp'
# reservemos 100Gb
with lmdb.open(lmdb_capnp_dir, 10 ** 11) as env:
    with env.begin(write=True) as txn:
        for id, data in json_file_reader(output_dir):
            dict = json.loads(data)
            msg = schema.Msg.new_message(createdAt=int(dict['created_at']), text=dict['text'])
            txn.put(int(id).to_bytes(4, 'big'), msg.to_bytes())

Desafortunadamente, nuestra base de datos tiene el mismo tamaño ~1.5Gb. Es un poco extraño... Estaba seguro de que el tamaño sería mucho menor.

Compactar

Intentemos compactar la base de datos (análogo a VACUUM)

import lmdb

# Abrir el entorno original
original_path = '/tmp/data/lmdb'
with lmdb.open(original_path, readonly=True, lock=False) as env:
    compacted_path = '/tmp/data/lmdb_compact'
    env.copy(compacted_path, compact=True)

Toma 9s y produce el mismo tamaño de base de datos: 1.5Gb

Compresión

Por defecto, Lmdb no soporta compresión, sin embargo, podemos intentar usar zstd.

tar -cf - lmdb | zstd -o lmdb.tar.zst
/*stdin*\            : 61.72%   (  1.54 GiB =>    971 MiB, lmdb.tar.zst)

Ahora es mejor, usar zfs con zstd puede ahorrar algo de espacio en el futuro. El tamaño es casi el mismo si comprimimos texto sin procesar. P.D. Si tratamos de comprimir rocksdb con zstd obtenemos 836 MiB, es incluso mejor que la compresión interna.

Resumen

En mi opinión, lmdb es el ganador. A pesar de que no he proporcionado resultados detallados sobre el rendimiento de lectura, en mis pruebas rápidas, esta cosa es realmente rápida. RocksDb puede ser una solución alternativa.