From 79516e2ecc23a77431f034752864f6f0dbc4794a Mon Sep 17 00:00:00 2001 From: Parker Loflin Date: Fri, 6 Nov 2026 17:30:20 -0500 Subject: [PATCH] Compress large IGzipFile writes in 4 MiB pieces to save memory IGzipFile passed each write to a single compress() call, which grows one bytes object until it holds all of the compressed data. Writing 100 MiB in one call took 100 MiB of extra memory for random data and 29 MiB for text (peak RSS on Linux, Python 3.11 and 3.14). Writes larger than 4 MiB are now compressed in 4 MiB pieces of a memoryview, and each compressed piece is written out before the next one is compressed. The same writes now take at most 6.8 MiB of extra memory for random data and 3.0 MiB for text. On Linux the write of random data is also about 15% faster, because fewer new pages are touched; the write of text takes the same time. Writes of 4 MiB or less are unchanged. The views are released in with statements, because a bytearray cannot be resized while a view of it exists, and PyPy frees views only when it collects garbage. Data that is not contiguous still goes to compress() whole, so it raises BufferError as before. The compressed bytes of large writes can differ slightly, because ISA-L sees the data in pieces. They decompress to the same data. --- CHANGELOG.rst | 6 +++ src/isal/igzip.py | 16 ++++++- tests/test_igzip.py | 104 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 125 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 481b7c5f..89ea2f44 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,6 +9,12 @@ Changelog version 1.10.0-dev ------------------ ++ Lower the memory use of large writes to ``igzip.open`` and ``IGzipFile`` + files. Writes larger than 4 MiB are now compressed in 4 MiB pieces, so the + compressed data of the whole write is no longer held in memory at once. + Writing 100 MiB in one call took 29 MiB of extra memory for text and + 100 MiB for data that does not compress, and now takes less than 8 MiB. + The compressed bytes of such writes can differ slightly. + Fix the gzip header that ``igzip_threaded.open`` writes when ``threads`` is not 0. It had the XFL and OS bytes swapped, so the ``file`` command reported the output as coming from VM/CMS or a FAT file system. diff --git a/src/isal/igzip.py b/src/isal/igzip.py index b18af11c..dbb98128 100644 --- a/src/isal/igzip.py +++ b/src/isal/igzip.py @@ -59,6 +59,8 @@ # _write_raw in large chunks, which compress much faster. _GZIP_BUFFERS_WRITES = hasattr(gzip.GzipFile, "_write_raw") +_WRITE_CHUNK_SIZE = 4 * 1024 * 1024 + # The open method was copied from the CPython source with minor adjustments. def open(filename, mode="rb", compresslevel=_COMPRESS_LEVEL_TRADEOFF, @@ -227,7 +229,19 @@ def _write_raw(self, data): data = memoryview(data) length = data.nbytes - if length > 0: + # Non-contiguous data goes to compress(), which raises BufferError. + if length > _WRITE_CHUNK_SIZE and (not isinstance(data, memoryview) + or data.c_contiguous): + # Compress in pieces to bound memory. Release the views at once so + # a bytearray can be resized afterwards, also on PyPy. + with memoryview(data) as view, view.cast("B") as byte_view: + for start in range(0, length, _WRITE_CHUNK_SIZE): + with byte_view[start:start + _WRITE_CHUNK_SIZE] as chunk: + self.fileobj.write(self.compress.compress(chunk)) + self.crc = isal_zlib.crc32(chunk, self.crc) + self.size += length + self.offset += length + elif length > 0: self.fileobj.write(self.compress.compress(data)) self.size += length self.crc = isal_zlib.crc32(data, self.crc) diff --git a/tests/test_igzip.py b/tests/test_igzip.py index 457f13f6..5ce51814 100644 --- a/tests/test_igzip.py +++ b/tests/test_igzip.py @@ -9,10 +9,12 @@ from CPython. Uses pytest which is easier to work with. Meant to complement the gzip module compliance tests. It should improve coverage as well.""" +import array import gzip import io import itertools import os +import random import re import struct import subprocess @@ -27,6 +29,11 @@ import pytest +try: + import tracemalloc +except ImportError: # PyPy + tracemalloc = None # type: ignore + DATA = b'This is a simple test with igzip' COMPRESSED_DATA = gzip.compress(DATA) TEST_FILE = str((Path(__file__).parent / "data" / "test.fastq.gz")) @@ -771,6 +778,103 @@ def test_small_writes_tell_and_flush(): assert gzip.decompress(output.getvalue()) == b"abcdefgh" +def large_write_data(size): + # Half incompressible, half compressible. + rng = random.Random(size) + half = size // 2 + repeated = rng.randbytes(1000) * (size // 1000 + 1) + return rng.randbytes(half) + repeated[:size - half] + + +@pytest.mark.parametrize("size", [ + igzip._WRITE_CHUNK_SIZE - 1, + igzip._WRITE_CHUNK_SIZE, + igzip._WRITE_CHUNK_SIZE + 1, + 3 * igzip._WRITE_CHUNK_SIZE, + 3 * igzip._WRITE_CHUNK_SIZE + 12345, +]) +def test_large_write(size): + data = large_write_data(size) + output = io.BytesIO() + with igzip.IGzipFile(fileobj=output, mode="wb") as gzip_file: + assert gzip_file.write(b"head") == 4 + assert gzip_file.write(data) == size + assert gzip_file.write(b"tail") == 4 + assert gzip_file.tell() == size + 8 + assert gzip.decompress(output.getvalue()) == b"head" + data + b"tail" + + +@pytest.mark.parametrize("make_buffer", [ + bytearray, + lambda data: memoryview(b"x" + data)[1:], + lambda data: array.array("I", data), + lambda data: memoryview(data).cast("B", (len(data) // 1024, 1024)), +], ids=["bytearray", "memoryview", "array", "2d-memoryview"]) +def test_large_write_buffer_types(make_buffer): + data = large_write_data(2 * igzip._WRITE_CHUNK_SIZE + 1024) + output = io.BytesIO() + with igzip.IGzipFile(fileobj=output, mode="wb") as gzip_file: + assert gzip_file.write(make_buffer(data)) == len(data) + assert gzip_file.tell() == len(data) + assert gzip.decompress(output.getvalue()) == data + + +def test_large_write_releases_bytearray(): + data = bytearray(large_write_data(igzip._WRITE_CHUNK_SIZE + 1)) + with igzip.IGzipFile(fileobj=io.BytesIO(), mode="wb") as gzip_file: + gzip_file.write(data) + # Raises BufferError if a view of the bytearray still exists + data.clear() + + +def test_large_write_error_releases_bytearray(): + class SmallWritesOnly(io.BytesIO): + def write(self, data): + if len(data) > 1024 * 1024: + raise OSError("write too large") + return super().write(data) + + # Incompressible, so the first piece is too large to write. + data = bytearray(random.Random(0).randbytes(igzip._WRITE_CHUNK_SIZE + 1)) + with igzip.IGzipFile(fileobj=SmallWritesOnly(), mode="wb") as gzip_file: + with pytest.raises(OSError) as error: + gzip_file.write(data) + # The traceback in error must not keep a view of the bytearray + data.clear() + error.match("write too large") + + +@pytest.mark.parametrize("size", [100, igzip._WRITE_CHUNK_SIZE + 1]) +def test_write_not_contiguous(size): + output = io.BytesIO() + with igzip.IGzipFile(fileobj=output, mode="wb") as gzip_file: + with pytest.raises(BufferError): + gzip_file.write(memoryview(bytes(2 * size))[::2]) + gzip_file.write(b"abc") + assert gzip.decompress(output.getvalue()) == b"abc" + + +@pytest.mark.skipif(tracemalloc is None, reason="PyPy has no tracemalloc") +def test_large_write_compresses_in_pieces(tmp_path): + # Incompressible, so one compress() call would need 24 MiB. + data = random.Random(0).randbytes(6 * igzip._WRITE_CHUNK_SIZE) + was_tracing = tracemalloc.is_tracing() + if not was_tracing: + tracemalloc.start() + try: + with igzip.open(tmp_path / "test.gz", "wb") as gzip_file: + tracemalloc.reset_peak() + before = tracemalloc.get_traced_memory()[0] + gzip_file.write(data) + peak = tracemalloc.get_traced_memory()[1] - before + finally: + if not was_tracing: + tracemalloc.stop() + assert peak < 3 * igzip._WRITE_CHUNK_SIZE + with igzip.open(tmp_path / "test.gz", "rb") as gzip_file: + assert gzip_file.read() == data + + @pytest.mark.skipif(sys.platform != "linux", reason="bash required") def test_no_broken_pipe(): subprocess.run(