From b614ef19e1b03293340b75a21117f69d4b210acc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 15:58:08 +0300 Subject: [PATCH] switch compression from gzip to zstd, use stdlib for python >= 3.14 --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/store/compression.py | 48 ++++++++++++++-- haiku_rag_slim/pyproject.toml | 1 + tests/store/test_compression.py | 56 ++++++++++++++++++- uv.lock | 2 + 5 files changed, 102 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b91c85f..e46244d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ - **Chunk expansion performance**: Fetch only chunks in the needed order range during context expansion instead of all chunks for a document - **Embedding batching**: Batch embedding calls in groups of 512 to avoid request size limits and timeouts with large documents - **DoclingDocument validation**: Strip page images before validation on the read path — pages are only needed for visualize_chunk and account for ~99% of the JSON size +- **Compression**: Switch from gzip to zstd for docling document storage (Python 3.14 stdlib, zstandard package for older versions) ## [0.36.3] - 2026-04-01 diff --git a/haiku_rag_slim/haiku/rag/store/compression.py b/haiku_rag_slim/haiku/rag/store/compression.py index 881efb89..1e848920 100644 --- a/haiku_rag_slim/haiku/rag/store/compression.py +++ b/haiku_rag_slim/haiku/rag/store/compression.py @@ -1,11 +1,49 @@ -import gzip +import json + +try: + from compression.zstd import ( # ty: ignore[unresolved-import] + compress as _zstd_compress, # type: ignore[import-not-found] + ) + from compression.zstd import ( # ty: ignore[unresolved-import] + decompress as _zstd_decompress, # type: ignore[import-not-found] + ) +except ImportError: + from zstandard import ZstdCompressor, ZstdDecompressor + + _zstd_compressor = ZstdCompressor() + _zstd_decompressor = ZstdDecompressor() + + def _zstd_compress(data: bytes) -> bytes: + return _zstd_compressor.compress(data) + + def _zstd_decompress(data: bytes) -> bytes: + return _zstd_decompressor.decompress(data, max_output_size=len(data) * 20) def compress_json(json_str: str) -> bytes: - """Compress a JSON string with gzip.""" - return gzip.compress(json_str.encode("utf-8")) + """Compress a JSON string with zstd.""" + return _zstd_compress(json_str.encode("utf-8")) def decompress_json(data: bytes) -> str: - """Decompress gzip-compressed data to a JSON string.""" - return gzip.decompress(data).decode("utf-8") + """Decompress zstd-compressed data to a JSON string.""" + return _zstd_decompress(data).decode("utf-8") + + +def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]: + """Split a DoclingDocument JSON into structure and pages, compress both with zstd. + + Returns: + Tuple of (structure_bytes, pages_bytes). pages_bytes is None if the + document has no page images. + """ + data = json.loads(json_str) + pages = data.pop("pages", None) + + structure_bytes = _zstd_compress(json.dumps(data).encode("utf-8")) + + pages_bytes = None + if pages: + pages_bytes = _zstd_compress(json.dumps(pages).encode("utf-8")) + + return structure_bytes, pages_bytes diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 1949fd14..834dfd99 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "rich>=14.3.3", "typer>=0.21.0,<0.22.0", "watchfiles>=1.1.1", + "zstandard>=0.23.0; python_version<'3.14'", ] [project.optional-dependencies] diff --git a/tests/store/test_compression.py b/tests/store/test_compression.py index 068e756b..9532b312 100644 --- a/tests/store/test_compression.py +++ b/tests/store/test_compression.py @@ -1,4 +1,10 @@ -from haiku.rag.store.compression import compress_json, decompress_json +import json + +from haiku.rag.store.compression import ( + compress_docling_split, + compress_json, + decompress_json, +) class TestJsonCompression: @@ -13,3 +19,51 @@ class TestJsonCompression: compressed = compress_json(json_str) decompressed = decompress_json(compressed) assert decompressed == json_str + + def test_compress_json_produces_zstd(self): + compressed = compress_json('{"test": true}') + assert compressed[:4] == b"\x28\xb5\x2f\xfd" + + +class TestDoclingCompressionSplit: + def test_split_with_pages(self): + data = { + "name": "test_doc", + "texts": [{"text": "hello"}], + "pages": {"1": {"image": "base64data"}, "2": {"image": "more"}}, + } + json_str = json.dumps(data) + structure_bytes, pages_bytes = compress_docling_split(json_str) + + assert structure_bytes is not None + assert pages_bytes is not None + + # Structure should not contain pages + structure = json.loads(decompress_json(structure_bytes)) + assert "pages" not in structure + assert structure["name"] == "test_doc" + assert structure["texts"] == [{"text": "hello"}] + + # Pages should contain only pages + pages = json.loads(decompress_json(pages_bytes)) + assert "1" in pages + assert "2" in pages + + def test_split_without_pages(self): + data = {"name": "test_doc", "texts": []} + json_str = json.dumps(data) + structure_bytes, pages_bytes = compress_docling_split(json_str) + + assert structure_bytes is not None + assert pages_bytes is None + + structure = json.loads(decompress_json(structure_bytes)) + assert structure["name"] == "test_doc" + + def test_split_with_empty_pages(self): + data = {"name": "test_doc", "texts": [], "pages": {}} + json_str = json.dumps(data) + structure_bytes, pages_bytes = compress_docling_split(json_str) + + assert structure_bytes is not None + assert pages_bytes is None diff --git a/uv.lock b/uv.lock index 3067cd66..1da254a7 100644 --- a/uv.lock +++ b/uv.lock @@ -1519,6 +1519,7 @@ dependencies = [ { name = "rich" }, { name = "typer" }, { name = "watchfiles" }, + { name = "zstandard", marker = "python_full_version < '3.14'" }, ] [package.optional-dependencies] @@ -1599,6 +1600,7 @@ requires-dist = [ { name = "typer", specifier = ">=0.21.0,<0.22.0" }, { name = "watchfiles", specifier = ">=1.1.1" }, { name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a11" }, + { name = "zstandard", marker = "python_full_version < '3.14'", specifier = ">=0.23.0" }, ] provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jina", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]