Add support for multiple zstd frames in decompression

Author: mkitti

numcodecs/tests/test_zstd.py

@@ -90,3 +90,16 @@ def test_native_functions():
     assert Zstd.default_level() == 3
     assert Zstd.min_level() == -131072
     assert Zstd.max_level() == 22
+
+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

@@ -42,10 +42,11 @@ cdef extern from "zstd.h":
                            const void* src,
                            size_t compressedSize) 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
@@ -186,7 +187,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_UNKNOWN or dest_size == ZSTD_CONTENTSIZE_ERROR:
             raise RuntimeError('Zstd decompression error: invalid input data')
 
@@ -226,6 +231,33 @@ def decompress(source, dest=None):
     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.