switch compression from gzip to zstd, use stdlib for python >= 3.14

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
 

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."""
+    """Compress a JSON string with zstd."""
-    return gzip.compress(json_str.encode("utf-8"))
+    return _zstd_compress(json_str.encode("utf-8"))
 
 
 def decompress_json(data: bytes) -> str:
-    """Decompress gzip-compressed data to a JSON string."""
+    """Decompress zstd-compressed data to a JSON string."""
-    return gzip.decompress(data).decode("utf-8")
+    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

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]

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

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"]