Parallel decompression for detector data

.readthedocs.yml

@@ -1,7 +1,7 @@
 version: 2  # Required
 
 build:
-  os: ubuntu-20.04
+  os: ubuntu-24.04
   tools:
     python: "3.12"
 

extra_data/components.py

@@ -1,6 +1,8 @@
 """Interfaces to data from specific instruments
 """
 import logging
+import os
+
 import math
 import re
 from collections.abc import Iterable
@@ -13,6 +15,7 @@
 from .exceptions import SourceNameError
 from .reader import DataCollection, by_id, by_index
 from .read_machinery import DataChunk, roi_shape, split_trains
+from .utils import default_num_threads
 from .writer import FileWriter
 from .write_cxi import XtdfCXIWriter, JUNGFRAUCXIWriter
 
@@ -1216,7 +1219,56 @@ def _read_chunk(self, chunk: DataChunk, mod_out, roi):
                 axis=0, out=mod_out[tgt_pulse_sel]
             )
 
-    def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, module_gaps=False):
+    def _read_parallel_decompress(self, out, module_gaps, threads=16):
+        from .compression import multi_dataset_decompressor, parallel_decompress_chunks
+
+        modno_to_keydata_no_virtual = {}
+        all_datasets = []
+        for (m, vkd) in self.modno_to_keydata.items():
+            modno_to_keydata_no_virtual[m] = kd = vkd._without_virtual_overview()
+            all_datasets.extend([f.file[kd.hdf5_data_path] for f in kd.files])
+
+        if any(d.chunks != (1,) + d.shape[1:] for d in all_datasets):
+            return False  # Chunking not as we expect
+
+        decomp_proto = multi_dataset_decompressor(all_datasets)
+        if decomp_proto is None:
+            return False  # No suitable fast decompression path
+
+        load_tasks = []
+        for i, (modno, kd) in enumerate(sorted(modno_to_keydata_no_virtual.items())):
+            mod_ix = (modno - self.det._modnos_start_at) if module_gaps else i
+            # 'chunk' in the lines below means a range of consecutive indices
+            # in one HDF5 dataset, as elsewhere in EXtra-data.
+            # We use this to build a list of HDF5 chunks (1 frame per chunk)
+            # to be loaded & decompressed. Sorry about that.
+            for chunk in kd._data_chunks:
+                dset = chunk.dataset
+
+                for tgt_slice, chunk_slice in self.det._split_align_chunk(
+                        chunk, self.det.train_ids_perframe,
+                ):
+                    inc_pulses_chunk = self._sel_frames[tgt_slice]
+                    if inc_pulses_chunk.sum() == 0:  # No data from this chunk selected
+                        continue
+
+                    dataset_ixs = np.nonzero(inc_pulses_chunk)[0] + chunk_slice.start
+
+                    # Where does this data go in the target array?
+                    tgt_start_ix = self._sel_frames[:tgt_slice.start].sum()
+
+                    # Each task is a h5py.h5d.DatasetID, coordinates & array destination
+                    load_tasks.extend([
+                        (dset.id, (ds_ix, 0, 0), out[mod_ix, tgt_start_ix + i])
+                        for i, ds_ix in enumerate(dataset_ixs)]
+                    )
+
+        parallel_decompress_chunks(load_tasks, decomp_proto, threads=threads)
+
+        return True
+
+    def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None,
+                module_gaps=False, decompress_threads=None):
         """Get an array of per-pulse data (image.*) for xtdf detector"""
         out_shape = self.buffer_shape(module_gaps=module_gaps, roi=roi)
 
@@ -1226,6 +1278,14 @@ def ndarray(self, *, fill_value=None, out=None, roi=(), astype=None, module_gaps
         elif out.shape != out_shape:
             raise ValueError(f'requires output array of shape {out_shape}')
 
+        if roi == () and astype is None:
+            if decompress_threads is None:
+                decompress_threads = default_num_threads(fixed_limit=16)
+
+            if decompress_threads > 1:
+                if self._read_parallel_decompress(out, module_gaps, decompress_threads):
+                    return out
+
         reading_view = out.view()
         if self._extraneous_dim:
             reading_view.shape = out.shape[:2] + (1,) + out.shape[2:]
@@ -1252,8 +1312,13 @@ def _wrap_xarray(self, arr, subtrain_index='pulseId'):
         })
 
     def xarray(self, *, pulses=None, fill_value=None, roi=(), astype=None,
-               subtrain_index='pulseId', unstack_pulses=False):
-        arr = self.ndarray(fill_value=fill_value, roi=roi, astype=astype)
+               subtrain_index='pulseId', unstack_pulses=False, decompress_threads=None):
+        arr = self.ndarray(
+            fill_value=fill_value,
+            roi=roi,
+            astype=astype,
+            decompress_threads=decompress_threads,
+        )
         out = self._wrap_xarray(arr, subtrain_index)
 
         if unstack_pulses:

extra_data/compression.py

@@ -0,0 +1,122 @@
+import threading
+from copy import copy
+from multiprocessing.pool import ThreadPool
+
+import h5py
+import numpy as np
+from zlib_into import decompress_into
+
+
+def filter_ids(dset: h5py.Dataset):
+    dcpl = dset.id.get_create_plist()
+    return [dcpl.get_filter(i)[0] for i in range(dcpl.get_nfilters())]
+
+
+class DeflateDecompressor:
+    def __init__(self, deflate_filter_idx=0):
+        self.deflate_filter_bit = 1 << deflate_filter_idx
+
+    @classmethod
+    def for_dataset(cls, dset: h5py.Dataset):
+        filters = filter_ids(dset)
+        if filters == [h5py.h5z.FILTER_DEFLATE]:
+            return cls()
+        if dset.dtype.itemsize == 1 and filters == [
+            h5py.h5z.FILTER_SHUFFLE,
+            h5py.h5z.FILTER_DEFLATE,
+        ]:
+            # The shuffle filter doesn't change single byte values, so we can
+            # skip it.
+            return cls(deflate_filter_idx=1)
+
+        return None
+
+    def clone(self):
+        return copy(self)
+
+    def apply_filters(self, data, filter_mask, out):
+        if filter_mask & self.deflate_filter_bit:
+            # The deflate filter is skipped, so just copy the data
+            memoryview(out)[:] = data
+        else:
+            decompress_into(data, out)
+
+
+class ShuffleDeflateDecompressor:
+    def __init__(self, chunk_shape, dtype):
+        self.chunk_shape = chunk_shape
+        self.dtype = dtype
+        chunk_nbytes = dtype.itemsize
+        for l in chunk_shape:
+            chunk_nbytes *= l
+        # This will hold the decompressed data before shuffling
+        self.chunk_buf = np.zeros(chunk_nbytes, dtype=np.uint8)
+        self.shuffled_view = (     # E.g. for int32 data with chunks (10, 5):
+                self.chunk_buf                  # (200,) uint8
+                .reshape((dtype.itemsize, -1))  # (4, 50)
+                .transpose()                    # (50, 4)
+        )
+        # Check this is still a view on the buffered data
+        assert self.shuffled_view.base is self.chunk_buf
+
+    @classmethod
+    def for_dataset(cls, dset: h5py.Dataset):
+        if filter_ids(dset) == [h5py.h5z.FILTER_SHUFFLE, h5py.h5z.FILTER_DEFLATE]:
+            return cls(dset.chunks, dset.dtype)
+
+        return None
+
+    def clone(self):
+        return type(self)(self.chunk_shape, self.dtype)
+
+    def apply_filters(self, data, filter_mask, out):
+        if filter_mask & 2:
+            # The deflate filter is skipped
+            memoryview(self.chunk_buf)[:] = data
+        else:
+            decompress_into(data, self.chunk_buf)
+
+        if filter_mask & 1:
+            # The shuffle filter is skipped
+            memoryview(out)[:] = self.chunk_buf
+        else:
+            # Numpy does the shuffling by copying data between views
+            out.reshape((-1, 1)).view(np.uint8)[:] = self.shuffled_view
+
+
+def dataset_decompressor(dset):
+    for cls in [DeflateDecompressor, ShuffleDeflateDecompressor]:
+        if (inst := cls.for_dataset(dset)) is not None:
+            return inst
+
+    return None
+
+
+def multi_dataset_decompressor(dsets):
+    if not dsets:
+        return None
+
+    chunk = dsets[0].chunks
+    dtype = dsets[0]
+    filters = filter_ids(dsets[0])
+    for d in dsets[1:]:
+        if d.chunks != chunk or d.dtype != dtype or filter_ids(d) != filters:
+            return None  # Datasets are not consistent
+
+    return dataset_decompressor(dsets[0])
+
+
+def parallel_decompress_chunks(tasks, decompressor_proto, threads=16):
+    tlocal = threading.local()
+
+    def load_one(dset_id, coord, dest):
+        try:
+            decomp = tlocal.decompressor
+        except AttributeError:
+            tlocal.decompressor = decomp = decompressor_proto.clone()
+
+        filter_mask, compdata = dset_id.read_direct_chunk(coord)
+        decomp.apply_filters(compdata, filter_mask, dest)
+
+    with ThreadPool(threads) as pool:
+        pool.starmap(load_one, tasks)

extra_data/keydata.py

@@ -165,6 +165,21 @@ def source_file_paths(self):
         from pathlib import Path
         return [Path(p) for p in paths]
 
+    def _without_virtual_overview(self):
+        if not self.files[0].file[self.hdf5_data_path].is_virtual:
+            # We're already looking at regular source files
+            return self
+
+        return KeyData(
+            self.source, self.key,
+            train_ids=self.train_ids,
+            files=[FileAccess(p) for p in self.source_file_paths],
+            section=self.section,
+            dtype=self.dtype,
+            eshape=self.entry_shape,
+            inc_suspect_trains=self.inc_suspect_trains,
+        )
+
     def _find_attributes(self, dset):
         """Find Karabo attributes belonging to a dataset."""
         attrs = dict(dset.attrs)

extra_data/utils.py

@@ -23,6 +23,16 @@ def available_cpu_cores():
         return min(ncpu, 8)
 
 
+def default_num_threads(fixed_limit=16):
+    # Default to 16, OMP_NUM_THREADS, or available CPU cores (picking lowest)
+    threads_limits = ([fixed_limit, available_cpu_cores()])
+    try:
+        threads_limits.append(int(os.environ['OMP_NUM_THREADS']))
+    except (KeyError, ValueError):
+        pass
+    return min(threads_limits)
+
+
 def progress_bar(done, total, suffix=" "):
     line = f"Progress: {done}/{total}{suffix}[{{}}]"
     length = min(get_terminal_size().columns - len(line), 50)

setup.py

@@ -58,6 +58,7 @@ def find_version(*parts):
           'pandas',
           'xarray',
           'pyyaml',
+          'zlib_into',
       ],
       extras_require={
           'bridge': [