File size: 2,192 Bytes
bc84119 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """Chunk-parallel wrapper around the context-mixing codec.
The input is split into fixed-size chunks, each compressed independently in a worker
process. Output layout is deterministic regardless of worker count:
b'M4' | uvarint(chunk_count) | uvarint(blob_len)*count | blobs...
Chunk size is a fixed constant, so compressed bytes are identical on any machine.
"""
import multiprocessing as mp
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
CHUNK = 1 << 20 # 1 MiB
MAGIC = b'M4'
def _uvarint(value: int) -> bytes:
out = bytearray()
while True:
b = value & 0x7F
value >>= 7
if value:
out.append(b | 0x80)
else:
out.append(b)
return bytes(out)
def _read_uvarint(buf: bytes, pos: int) -> tuple[int, int]:
shift = 0
value = 0
while True:
b = buf[pos]
pos += 1
value |= (b & 0x7F) << shift
if not b & 0x80:
return value, pos
shift += 7
def _compress_chunk(chunk: bytes) -> bytes:
from codec import compress
return compress(chunk)
def _decompress_chunk(blob: bytes) -> bytes:
from codec import decompress
return decompress(blob)
def _pool_map(fn, items):
if len(items) <= 1:
return [fn(item) for item in items]
workers = min(len(items), os.cpu_count() or 1)
with mp.get_context('fork').Pool(workers) as pool:
return pool.map(fn, items, chunksize=1)
def compress(data: bytes) -> bytes:
chunks = [data[i:i + CHUNK] for i in range(0, len(data), CHUNK)] or [b'']
blobs = _pool_map(_compress_chunk, chunks)
head = bytearray(MAGIC)
head += _uvarint(len(blobs))
for blob in blobs:
head += _uvarint(len(blob))
return bytes(head) + b''.join(blobs)
def decompress(blob: bytes) -> bytes:
assert blob[:2] == MAGIC, 'bad magic'
count, pos = _read_uvarint(blob, 2)
lengths = []
for _ in range(count):
n, pos = _read_uvarint(blob, pos)
lengths.append(n)
blobs = []
for n in lengths:
blobs.append(blob[pos:pos + n])
pos += n
return b''.join(_pool_map(_decompress_chunk, blobs))
|