Merge branch 'main' into handle-zarr-3.1.0

Author: d-v-b

docs/compression/zstd.rst

@@ -7,6 +7,9 @@ Zstd
     .. autoattribute:: codec_id
     .. automethod:: encode
     .. automethod:: decode
+        .. note::
+            If the compressed data does not contain the decompressed size, streaming
+            decompression will be used.
     .. automethod:: get_config
     .. automethod:: from_config
 

docs/release.rst

@@ -95,6 +95,8 @@ Maintenance
 
 Improvements
 ~~~~~~~~~~~~
+* Add streaming decompression for ZSTD (:issue:`699`)
+  By :user:`Mark Kittisopikul <mkitti>`.
 * Raise a custom `UnknownCodecError` when trying to retrieve an unavailable codec.
   By :user:`Cas Wognum <cwognum>`.
 

numcodecs/abc.py

@@ -29,14 +29,13 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Optional
 
 
 class Codec(ABC):
     """Codec abstract base class."""
 
     # override in sub-class
-    codec_id: Optional[str] = None
+    codec_id: str | None = None
     """Codec identifier."""
 
     @abstractmethod

numcodecs/checksum32.py

@@ -3,7 +3,7 @@
 import zlib
 from contextlib import suppress
 from types import ModuleType
-from typing import Literal, Optional
+from typing import Literal
 
 import numpy as np
 from typing_extensions import Buffer
@@ -12,7 +12,7 @@
 from .compat import ensure_contiguous_ndarray, ndarray_copy
 from .jenkins import jenkins_lookup3
 
-_crc32c: Optional[ModuleType] = None
+_crc32c: ModuleType | None = None
 with suppress(ImportError):
     import crc32c as _crc32c  # type: ignore[no-redef, unused-ignore]
 

numcodecs/lzma.py

@@ -1,7 +1,6 @@
 from types import ModuleType
-from typing import Optional
 
-_lzma: Optional[ModuleType] = None
+_lzma: ModuleType | None = None
 try:
     import lzma as _lzma
 except ImportError:  # pragma: no cover

numcodecs/tests/test_pyzstd.py

@@ -0,0 +1,77 @@
+# Check Zstd against pyzstd package
+
+import numpy as np
+import pytest
+import pyzstd
+
+from numcodecs.zstd import Zstd
+
+test_data = [
+    b"Hello World!",
+    np.arange(113).tobytes(),
+    np.arange(10, 15).tobytes(),
+    np.random.randint(3, 50, size=(53,), dtype=np.uint16).tobytes(),
+]
+
+
+@pytest.mark.parametrize("input", test_data)
+def test_pyzstd_simple(input):
+    """
+    Test if Zstd.[decode, encode] can perform the inverse operation to
+    pyzstd.[compress, decompress] in the simple case.
+    """
+    z = Zstd()
+    assert z.decode(pyzstd.compress(input)) == input
+    assert pyzstd.decompress(z.encode(input)) == input
+
+
+@pytest.mark.xfail
+@pytest.mark.parametrize("input", test_data)
+def test_pyzstd_simple_multiple_frames_decode(input):
+    """
+    Test decompression of two concatenated frames of known sizes
+
+    numcodecs.zstd.Zstd currently fails because it only assesses the size of the
+    first frame. Rather, it should keep iterating through all the frames until
+    the end of the input buffer.
+    """
+    z = Zstd()
+    assert pyzstd.decompress(pyzstd.compress(input) * 2) == input * 2
+    assert z.decode(pyzstd.compress(input) * 2) == input * 2
+
+
+@pytest.mark.parametrize("input", test_data)
+def test_pyzstd_simple_multiple_frames_encode(input):
+    """
+    Test if pyzstd can decompress two concatenated frames from Zstd.encode
+    """
+    z = Zstd()
+    assert pyzstd.decompress(z.encode(input) * 2) == input * 2
+
+
+@pytest.mark.parametrize("input", test_data)
+def test_pyzstd_streaming(input):
+    """
+    Test if Zstd can decode a single frame and concatenated frames in streaming
+    mode where the decompressed size is not recorded in the frame header.
+    """
+    pyzstd_c = pyzstd.ZstdCompressor()
+    pyzstd_d = pyzstd.ZstdDecompressor()
+    pyzstd_e = pyzstd.EndlessZstdDecompressor()
+    z = Zstd()
+
+    d_bytes = input
+    pyzstd_c.compress(d_bytes)
+    c_bytes = pyzstd_c.flush()
+    assert z.decode(c_bytes) == d_bytes
+    assert pyzstd_d.decompress(z.encode(d_bytes)) == d_bytes
+
+    # Test multiple streaming frames
+    assert z.decode(c_bytes * 2) == pyzstd_e.decompress(c_bytes * 2)
+    assert z.decode(c_bytes * 3) == pyzstd_e.decompress(c_bytes * 3)
+    assert z.decode(c_bytes * 4) == pyzstd_e.decompress(c_bytes * 4)
+    assert z.decode(c_bytes * 5) == pyzstd_e.decompress(c_bytes * 5)
+    assert z.decode(c_bytes * 7) == pyzstd_e.decompress(c_bytes * 7)
+    assert z.decode(c_bytes * 11) == pyzstd_e.decompress(c_bytes * 11)
+    assert z.decode(c_bytes * 13) == pyzstd_e.decompress(c_bytes * 13)
+    assert z.decode(c_bytes * 99) == pyzstd_e.decompress(c_bytes * 99)

numcodecs/tests/test_zarr3.py

@@ -122,8 +122,8 @@ def test_generic_filter(
             ],
         )
 
-        a[:, :] = data.copy()
-        a = zarr.open_array(store / "generic", mode="r")
+    a[:, :] = data.copy()
+    a = zarr.open_array(store / "generic", mode="r")
     np.testing.assert_array_equal(data, a[:, :])
 
 
@@ -140,8 +140,8 @@ def test_generic_filter_bitround(store: StorePath):
             filters=[numcodecs.zarr3.BitRound(keepbits=3)],
         )
 
-        a[:, :] = data.copy()
-        a = zarr.open_array(store / "generic_bitround", mode="r")
+    a[:, :] = data.copy()
+    a = zarr.open_array(store / "generic_bitround", mode="r")
     assert np.allclose(data, a[:, :], atol=0.1)
 
 
@@ -158,8 +158,8 @@ def test_generic_filter_quantize(store: StorePath):
             filters=[numcodecs.zarr3.Quantize(digits=3)],
         )
 
-        a[:, :] = data.copy()
-        a = zarr.open_array(store / "generic_quantize", mode="r")
+    a[:, :] = data.copy()
+    a = zarr.open_array(store / "generic_quantize", mode="r")
     assert np.allclose(data, a[:, :], atol=0.001)
 
 
@@ -177,8 +177,8 @@ def test_generic_filter_packbits(store: StorePath):
             filters=[numcodecs.zarr3.PackBits()],
         )
 
-        a[:, :] = data.copy()
-        a = zarr.open_array(store / "generic_packbits", mode="r")
+    a[:, :] = data.copy()
+    a = zarr.open_array(store / "generic_packbits", mode="r")
     np.testing.assert_array_equal(data, a[:, :])
 
     with pytest.raises(ValueError, match=".*requires bool dtype.*"):
@@ -217,8 +217,8 @@ def test_generic_checksum(
             compressors=[codec_class()],
         )
 
-        a[:, :] = data.copy()
-        a = zarr.open_array(store / "generic_checksum", mode="r")
+    a[:, :] = data.copy()
+    a = zarr.open_array(store / "generic_checksum", mode="r")
     np.testing.assert_array_equal(data, a[:, :])
 
 
@@ -267,8 +267,8 @@ def test_delta_astype(store: StorePath):
             ],
         )
 
-        a[:, :] = data.copy()
-        a = zarr.open_array(store / "generic", mode="r")
+    a[:, :] = data.copy()
+    a = zarr.open_array(store / "generic", mode="r")
     np.testing.assert_array_equal(data, a[:, :])
 
 

numcodecs/tests/test_zstd.py

@@ -1,14 +1,9 @@
 import itertools
+import subprocess
 
 import numpy as np
 import pytest
 
-try:
-    from numcodecs.zstd import Zstd
-except ImportError:  # pragma: no cover
-    pytest.skip("numcodecs.zstd not available", allow_module_level=True)
-
-
 from numcodecs.tests.common import (
     check_backwards_compatibility,
     check_config,
@@ -17,6 +12,7 @@
     check_err_encode_object_buffer,
     check_repr,
 )
+from numcodecs.zstd import Zstd
 
 codecs = [
     Zstd(),
@@ -90,3 +86,73 @@ def test_native_functions():
     assert Zstd.default_level() == 3
     assert Zstd.min_level() == -131072
     assert Zstd.max_level() == 22
+
+
+def test_streaming_decompression():
+    # Test input frames with unknown frame content size
+    codec = Zstd()
+
+    # If the zstd command line interface is available, check the bytes
+    cli = zstd_cli_available()
+    if cli:
+        view_zstd_streaming_bytes()
+
+    # Encode bytes directly that were the result of streaming compression
+    bytes_val = b'(\xb5/\xfd\x00Xa\x00\x00Hello World!'
+    dec = codec.decode(bytes_val)
+    dec_expected = b'Hello World!'
+    assert dec == dec_expected
+    if cli:
+        assert bytes_val == generate_zstd_streaming_bytes(dec_expected)
+        assert dec_expected == generate_zstd_streaming_bytes(bytes_val, decompress=True)
+
+    # Two consecutive frames given as input
+    bytes2 = bytes(bytearray(bytes_val * 2))
+    dec2 = codec.decode(bytes2)
+    dec2_expected = b'Hello World!Hello World!'
+    assert dec2 == dec2_expected
+    if cli:
+        assert dec2_expected == generate_zstd_streaming_bytes(bytes2, decompress=True)
+
+    # Single long frame that decompresses to a large output
+    bytes3 = b'(\xb5/\xfd\x00X$\x02\x00\xa4\x03ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz\x01\x00:\xfc\xdfs\x05\x05L\x00\x00\x08s\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08k\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08c\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08[\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08S\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08K\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08C\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08u\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08m\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08e\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08]\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08U\x01\x00\xfc\xff9\x10\x02L\x00\x00\x08M\x01\x00\xfc\xff9\x10\x02M\x00\x00\x08E\x01\x00\xfc\x7f\x1d\x08\x01'
+    dec3 = codec.decode(bytes3)
+    dec3_expected = b'ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz' * 1024 * 32
+    assert dec3 == dec3_expected
+    if cli:
+        assert bytes3 == generate_zstd_streaming_bytes(dec3_expected)
+        assert dec3_expected == generate_zstd_streaming_bytes(bytes3, decompress=True)
+
+    # Garbage input results in an error
+    bytes4 = bytes(bytearray([0, 0, 0, 0, 0, 0, 0, 0]))
+    with pytest.raises(RuntimeError, match='Zstd decompression error: invalid input data'):
+        codec.decode(bytes4)
+
+
+def generate_zstd_streaming_bytes(input: bytes, *, decompress: bool = False) -> bytes:
+    """
+    Use the zstd command line interface to compress or decompress bytes in streaming mode.
+    """
+    if decompress:
+        args = ["-d"]
+    else:
+        args = []
+
+    p = subprocess.run(["zstd", "--no-check", *args], input=input, capture_output=True)
+    return p.stdout
+
+
+def view_zstd_streaming_bytes():
+    bytes_val = generate_zstd_streaming_bytes(b"Hello world!")
+    print(f"    bytes_val = {bytes_val}")
+
+    bytes3 = generate_zstd_streaming_bytes(
+        b"ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz" * 1024 * 32
+    )
+    print(f"    bytes3 = {bytes3}")
+
+
+def zstd_cli_available() -> bool:
+    return not subprocess.run(
+        ["zstd", "-V"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+    ).returncode

numcodecs/zfpy.py

@@ -2,9 +2,8 @@
 from contextlib import suppress
 from importlib.metadata import PackageNotFoundError, version
 from types import ModuleType
-from typing import Optional
 
-_zfpy: Optional[ModuleType] = None
+_zfpy: ModuleType | None = None
 
 _zfpy_version: tuple = ()
 with suppress(PackageNotFoundError):