Add support for multiple zstd frames in decompression

Author: mkitti

numcodecs/tests/test_zstd.py

@@ -87,7 +87,6 @@ def test_native_functions():
     assert Zstd.min_level() == -131072
     assert Zstd.max_level() == 22
 
-
 def test_streaming_decompression():
     # Test input frames with unknown frame content size
     codec = Zstd()
@@ -156,3 +155,16 @@ def zstd_cli_available() -> bool:
     return not subprocess.run(
         ["zstd", "-V"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
     ).returncode
+
+def test_multi_frame():
+    codec = Zstd()
+
+    hello_world = codec.encode(b"Hello world!")
+    assert codec.decode(hello_world) == b"Hello world!"
+    assert codec.decode(hello_world*2) == b"Hello world!Hello world!"
+
+    hola = codec.encode(b"Hola ")
+    mundo = codec.encode(b"Mundo!")
+    assert codec.decode(hola) == b"Hola "
+    assert codec.decode(mundo) == b"Mundo!"
+    assert codec.decode(hola+mundo) == b"Hola Mundo!"

numcodecs/zstd.pyx

@@ -68,10 +68,12 @@ cdef extern from "zstd.h":
     size_t ZSTD_freeDStream(ZSTD_DStream* zds) nogil
     size_t ZSTD_initDStream(ZSTD_DStream* zds) nogil
 
-    cdef long ZSTD_CONTENTSIZE_UNKNOWN
-    cdef long ZSTD_CONTENTSIZE_ERROR
+    cdef unsigned long long ZSTD_CONTENTSIZE_UNKNOWN
+    cdef unsigned long long ZSTD_CONTENTSIZE_ERROR
+
     unsigned long long ZSTD_getFrameContentSize(const void* src,
                                                 size_t srcSize) nogil
+    size_t ZSTD_findFrameCompressedSize(const void* src, size_t srcSize) nogil
 
     int ZSTD_minCLevel() nogil
     int ZSTD_maxCLevel() nogil
@@ -216,7 +218,11 @@ def decompress(source, dest=None):
     try:
 
         # determine uncompressed size
-        dest_size = ZSTD_getFrameContentSize(source_ptr, source_size)
+        try:
+            dest_size = findTotalContentSize(source_ptr, source_size)
+        except RuntimeError:
+            raise RuntimeError('Zstd decompression error: invalid input data')
+
         if dest_size == 0 or dest_size == ZSTD_CONTENTSIZE_ERROR:
             raise RuntimeError('Zstd decompression error: invalid input data')
 
@@ -353,6 +359,33 @@ cdef stream_decompress(const Py_buffer* source_pb):
 
     return dest
 
+cdef findTotalContentSize(const void* source_ptr, size_t source_size):
+    cdef:
+        unsigned long long frame_content_size = 0
+        unsigned long long total_content_size = 0
+        size_t frame_compressed_size = 0
+        size_t offset = 0
+
+    while offset < source_size:
+        frame_compressed_size = ZSTD_findFrameCompressedSize(source_ptr + offset, source_size - offset);
+
+        if ZSTD_isError(frame_compressed_size):
+            error = ZSTD_getErrorName(frame_compressed_size)
+            raise RuntimeError('Could not set determine zstd frame size: %s' % error)
+
+        frame_content_size = ZSTD_getFrameContentSize(source_ptr + offset, frame_compressed_size);
+
+        if frame_content_size == ZSTD_CONTENTSIZE_ERROR:
+            return ZSTD_CONTENTSIZE_ERROR
+
+        if frame_content_size == ZSTD_CONTENTSIZE_UNKNOWN:
+            return ZSTD_CONTENTSIZE_UNKNOWN
+
+        total_content_size += frame_content_size
+        offset += frame_compressed_size
+
+    return total_content_size
+
 class Zstd(Codec):
     """Codec providing compression using Zstandard.