Revert "Revert "Use zstandard to decompress chunks to save memory (#977)" (…" (#981)

This reverts commit 07a1168.

Author: dachengx

.github/workflows/pytest.yml

@@ -42,6 +42,8 @@ jobs:
       - name: Install requirements for Python 3.11
         if: matrix.python-version == '3.11'
         run: pip install git+https://github.com/XENONnT/base_environment.git --force-reinstall
+      - name: Install strax
+        run: pip install .
       - name: Start MongoDB
         uses: supercharge/mongodb-github-action@1.11.0
         with:

pyproject.toml

@@ -43,6 +43,7 @@ scipy = "*"
 tqdm = ">=4.46.0"
 zarr = "<3.0.0"
 zstd = "*"
+zstandard = "*"
 sphinx = { version = "*", optional = true }
 sphinx_rtd_theme = { version = "*", optional = true }
 nbsphinx = { version = "*", optional = true }

strax/io.py

@@ -7,6 +7,7 @@
 import numpy as np
 import blosc
 import zstd
+import zstandard
 import lz4.frame as lz4
 from ast import literal_eval
 
@@ -40,9 +41,11 @@ def _bz2_decompress(f, buffer_size=DECOMPRESS_BUFFER_SIZE):
 _zstd_compress = lambda data: zstd.compress(data, 3, 1)
 
 
-def _zstd_decompress(f):
-    data = f.read()
-    data = zstd.decompress(data)
+def _zstd_decompress(f, chunk_size=64 * 1024 * 1024):
+    decompressor = zstandard.ZstdDecompressor().decompressobj()
+    data = bytearray()  # Efficient mutable storage
+    for d in iter(lambda: f.read(chunk_size), b""):
+        data.extend(decompressor.decompress(d))
     return data
 
 
@@ -181,5 +184,10 @@ def load_chunk(chunk_info):
         x = load_chunk(chunk_info)
         x = strax.apply_selection(x, **kwargs)
         results.append(x)
-    results = np.hstack(results)
+
+    # No need to hstack if only one chunk is loaded
+    if len(results) == 1:
+        results = results[0]
+    else:
+        results = np.hstack(results)
     return results if len(results) else np.empty(0, dtype)