Fix a performance regression (and address #1532)

neo/rawio/biocamrawio.py

@@ -61,6 +61,14 @@ class BiocamRawIO(BaseRawIO):
     extensions = ["h5", "brw"]
     rawmode = "one-file"
 
+    # Ceiling on the temporary all-channel buffer used when only a subset of channels is requested.
+    # The dense layouts read a contiguous slice, so their cost scales with bytes and a small buffer
+    # stays in cache. The event-based sparse decoder re-reads both tables of contents and re-walks
+    # the event blob on every call, so its cost is per call rather than per byte and it wants a
+    # buffer large enough that it almost never splits.
+    max_read_bytes_dense = 4 * 1024 * 1024
+    max_read_bytes_sparse = 512 * 1024 * 1024
+
     def __init__(self, filename="", fill_gaps_strategy="zeros"):
         BaseRawIO.__init__(self)
         self.filename = filename
@@ -136,12 +144,8 @@ def _get_signal_t_start(self, block_index, seg_index, stream_index):
             raise ValueError("`stream_index must be 0")
         return self._segment_t_start(block_index, seg_index)
 
-    def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes):
-        if i_start is None:
-            i_start = 0
-        if i_stop is None:
-            i_stop = self._num_frames
-
+    def _read_block(self, i_start, i_stop):
+        """Read frames [i_start, i_stop) for all channels, as (n_frames, n_channels)."""
         # read functions are different based on the version of biocam
         if self._read_function is readHDF5t_brw4_sparse:
             if self._fill_gaps_strategy is None:
@@ -163,34 +167,49 @@ def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, strea
         else:
             data = self._read_function(self._filehandle, i_start, i_stop, self._num_channels)
 
-        # older style data returns array of (n_samples, n_channels), should be a view
-        # but if memory issues come up we should doublecheck out how the file is being stored
-        if data.ndim > 1:
-            if channel_indexes is None:
-                channel_indexes = slice(None)
-            sig_chunk = data[:, channel_indexes]
-
-        # newer style data returns an initial flat array (n_samples * n_channels)
-        # we iterate through channels rather than slicing
-        # Due to the fact that Neo and SpikeInterface tend to prefer slices we need to add
-        # some careful checks around slicing of None in the case we need to iterate through
-        # channels. First check if None. Then check if slice and only if slice check that it is slice(None)
+        # Newer style data is read as a flat (n_samples * n_channels) array. The h5py read is
+        # C-contiguous, so this reshape is a free view: it copies nothing and allocates nothing.
+        # Older style data is already (n_samples, n_channels).
+        if data.ndim == 1:
+            data = data.reshape(i_stop - i_start, self._num_channels)
+        return data
+
+    def _get_analogsignal_chunk(self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes):
+        if i_start is None:
+            i_start = 0
+        if i_stop is None:
+            i_stop = self._num_frames
+        n_frames = i_stop - i_start
+        if channel_indexes is None:
+            channel_indexes = slice(None)
+        # Number of selected channels, for every form `channel_indexes` can take: a slice, an
+        # integer array or a boolean mask.
+        n_ch = np.arange(self._num_channels)[channel_indexes].size
+        dtype = self.header["signal_channels"]["dtype"][0]
+        itemsize = np.dtype(dtype).itemsize
+        read_bytes = n_frames * self._num_channels * itemsize
+        out_bytes = n_frames * n_ch * itemsize
+        max_read_bytes = (
+            self.max_read_bytes_sparse if self._read_function is readHDF5t_brw4_sparse else self.max_read_bytes_dense
+        )
+
+        # Reading a subset of channels from a very large recording.
+        # The BioCam hdf5 layout is such that all channels are chunked together, so they cannot be
+        # read independently. In the event that only a subset of channels of a very large recording
+        # is needed, we cannot slice a single channel. Instead, we must loop and read all channels,
+        # slicing as we go. We will do this chunking if the read is large and the output to be
+        # returned is smaller.
+        read_by_chunks = read_bytes > max_read_bytes and out_bytes < read_bytes // 2
+        if read_by_chunks:
+            # Read in chunks, copying desired channels into a preallocated array.
+            out = np.empty((n_frames, n_ch), dtype=dtype)
+            max_frames = max(1, max_read_bytes // (itemsize * self._num_channels))
+            for start in range(i_start, i_stop, max_frames):
+                stop = min(start + max_frames, i_stop)
+                out[start - i_start : stop - i_start] = self._read_block(start, stop)[:, channel_indexes]
         else:
-            if channel_indexes is None:
-                channel_indexes = [ch for ch in range(self._num_channels)]
-            elif isinstance(channel_indexes, slice):
-                start = channel_indexes.start or 0
-                stop = channel_indexes.stop or self._num_channels
-                step = channel_indexes.step or 1
-                channel_indexes = [ch for ch in range(start, stop, step)]
-
-            sig_chunk = np.zeros((i_stop - i_start, len(channel_indexes)), dtype=data.dtype)
-            # iterate through channels to prevent loading all channels into memory which can cause
-            # memory exhaustion. See https://github.com/SpikeInterface/spikeinterface/issues/3303
-            for index, channel_index in enumerate(channel_indexes):
-                sig_chunk[:, index] = data[channel_index :: self._num_channels]
-
-        return sig_chunk
+            out = self._read_block(i_start, i_stop)[:, channel_indexes]
+        return out
 
 
 def open_biocam_file_header(filename) -> dict:

neo/test/rawiotest/test_biocamrawio.py

@@ -24,8 +24,76 @@ class TestBiocamRawIO(
     entities_to_test = [
         "biocam/biocam_hw3.0_fw1.6.brw",
         "biocam/biocam_hw3.0_fw1.7.0.12_raw.brw",
+        "biocam/raw_shape_variants/version_100.brw",
+        "biocam/raw_shape_variants/version_101.brw",
     ]
 
+    # The two files in `raw_shape_variants` hold the same recording, 64 channels over 1000 samples,
+    # written in the two shapes that the BRW v3.x specification defines for `3BData/Raw`: a T x W
+    # matrix for version 100 and a flat interleaved array for version 101. The value stored at a
+    # given position is `sample * 64 + channel`, so no two positions share a value and a misread
+    # returns a number that says where it actually came from.
+    shape_variants = [
+        "biocam/raw_shape_variants/version_100.brw",
+        "biocam/raw_shape_variants/version_101.brw",
+    ]
+    variant_num_channels = 64
+    variant_num_frames = 1000
+
+    # Some of the ways that neo and SpikeInterface are known to ask for channels.
+    channel_queries = [None, slice(None), slice(2, 6), [0, 1, 2], [3, 0, 2], [5], np.array([1, 4])]
+    # None means the whole segment. The rest include a non-zero start, a single frame, and the
+    # empty range at the end of the recording.
+    frame_ranges = [(None, None), (0, 5), (7, 13), (999, 1000), (1000, 1000)]
+
+    def test_get_analogsignal_chunk_values(self):
+        """Values, shape and dtype are correct for every way of selecting channels."""
+        expected_all = np.arange(self.variant_num_frames * self.variant_num_channels, dtype="uint16")
+        expected_all = expected_all.reshape(self.variant_num_frames, self.variant_num_channels)
+
+        for entity in self.shape_variants:
+            reader = BiocamRawIO(filename=self.get_local_path(entity))
+            reader.parse_header()
+            for channel_indexes in self.channel_queries:
+                for i_start, i_stop in self.frame_ranges:
+                    with self.subTest(entity=entity, channel_indexes=channel_indexes, i_start=i_start, i_stop=i_stop):
+                        sig = reader.get_analogsignal_chunk(
+                            block_index=0,
+                            seg_index=0,
+                            i_start=i_start,
+                            i_stop=i_stop,
+                            channel_indexes=channel_indexes,
+                        )
+                        start = 0 if i_start is None else i_start
+                        stop = self.variant_num_frames if i_stop is None else i_stop
+                        columns = slice(None) if channel_indexes is None else channel_indexes
+                        expected = expected_all[start:stop][:, columns]
+                        assert sig.dtype == expected.dtype
+                        np.testing.assert_array_equal(sig, expected)
+
+    def test_chunked_read_matches_direct_read(self):
+        """Shrinking the read budget must not change the data that comes back.
+
+        A request for a few channels over a long stretch of a large recording is served by
+        reading all the channels in blocks and copying out the requested ones, the BioCAM
+        layout interleaving the channels so that they cannot be read independently.
+        """
+        for entity in self.entities_to_test:
+            reader = BiocamRawIO(filename=self.get_local_path(entity))
+            reader.parse_header()
+            for channel_indexes in self.channel_queries:
+                with self.subTest(entity=entity, channel_indexes=channel_indexes):
+                    kwargs = dict(block_index=0, seg_index=0, i_start=0, i_stop=100, channel_indexes=channel_indexes)
+                    direct = reader.get_analogsignal_chunk(**kwargs)
+                    # A budget of ten frames splits the hundred frames requested above into ten
+                    # blocks that the reader has to stitch back together.
+                    reader.max_read_bytes_dense = 10 * reader._num_channels * np.dtype("uint16").itemsize
+                    chunked = reader.get_analogsignal_chunk(**kwargs)
+                    del reader.max_read_bytes_dense
+
+                    assert chunked.dtype == direct.dtype
+                    np.testing.assert_array_equal(chunked, direct)
+
 
 def test_biocamrawio_gain(tmp_path):
     """Test that BiocamRawIO correctly reads the gain from a Biocam HDF5 file.