Merge pull request #45 from cta-observatory/fix-zstd

Fix reading of zstd compressed files (no backseeking in iter_blocks)

Author: maxnoe

corsikaio/io.py

@@ -61,32 +61,41 @@ def read_buffer_size(path):
 
 def iter_blocks(f, thinning=False):
     is_fortran_file = True
-    if thinning == False:
+    if not thinning:
         block_size = BLOCK_SIZE_BYTES
         buffer_size = DEFAULT_BUFFER_SIZE
     else:
         block_size = BLOCK_SIZE_BYTES_THIN
         buffer_size = DEFAULT_BUFFER_SIZE_THIN
 
-
     data = f.read(4)
-    f.seek(0)
+    first = True
     if data == b'RUNH':
         is_fortran_file = False
 
     while True:
         # for the fortran-chunked output, we need to read the record size
         if is_fortran_file:
-            data = f.read(RECORD_MARKER.size)
+            if first is True:
+                data = data + f.read(RECORD_MARKER.size - len(data))
+            else:
+                data = f.read(RECORD_MARKER.size)
+
             if len(data) == 0:
                 return
 
             if len(data) < RECORD_MARKER.size:
                 raise IOError("Read less bytes than expected, file seems to be truncated")
 
             buffer_size, = RECORD_MARKER.unpack(data)
+            data = b""
+
+        if first is True:
+            data = data + f.read(buffer_size - len(data))
+            first = False
+        else:
+            data = f.read(buffer_size)
 
-        data = f.read(buffer_size)
         if is_fortran_file:
             if len(data) < buffer_size:
                 raise IOError("Read less bytes than expected, file seems to be truncated")

setup.cfg

@@ -25,7 +25,8 @@ zstd =
 	zstandard
 tests =
 	pytest
-    scipy
+  scipy
+  zstandard
 all =
 	%(zstd)s
 	%(tests)s

tests/test_file.py

@@ -1,6 +1,12 @@
+from contextlib import ExitStack
+import gzip
+from pathlib import Path
+
 import pytest
 import numpy as np
 
+from zstandard import ZstdCompressor
+
 from corsikaio.constants import BLOCK_SIZE_BYTES
 from corsikaio.io import RECORD_MARKER
 
@@ -167,3 +173,40 @@ def test_longitudinal_parameters():
             parameters = event.end["longitudinal_fit_parameters"]
             np.testing.assert_array_equal(parameters != 0, True)
         assert n_events == 5
+
+
+@pytest.mark.parametrize (
+    "test_path",
+    [
+        'tests/resources/mmcs65',
+        'tests/resources/corsika74100',
+    ]
+)
+@pytest.mark.parametrize( "compression", ["gz", "zst"])
+def test_compressed(test_path, compression, tmp_path):
+    from corsikaio import CorsikaCherenkovFile
+
+    test_path = Path(test_path)
+    compressed = tmp_path / f"{test_path.name}.{compression}"
+
+    ctx = ExitStack()
+
+    with ctx:
+        infile = ctx.enter_context(test_path.open("rb"))
+        outfile = ctx.enter_context(compressed.open("wb"))
+
+        if compression == "gz":
+            outstream = ctx.enter_context(gzip.GzipFile(fileobj=outfile, mode="wb"))
+        elif compression == "zst":
+            compressor = ZstdCompressor(level=10)
+            outstream = ctx.enter_context(compressor.stream_writer(outfile))
+        else:
+            raise ValueError(f"Unknown compression: {compression}")
+
+        for chunk in iter(lambda : infile.read(102400), b""):
+            outstream.write(chunk)
+
+    with CorsikaCherenkovFile(compressed) as cf, CorsikaCherenkovFile(test_path) as f:
+        for event in f:
+            compressed_event = next(cf)
+            assert event.header["event_number"] == compressed_event.header["event_number"]