Override search_cached_spikes_sorted for identity selections

src/spikeinterface/core/tests/test_unitsselectionsorting.py

@@ -248,6 +248,7 @@ def test_non_identity_selection_does_not_share(unit_ids):
     assert len(parent._cached_lexsorted_spike_vector) == 0
 
 
+@pytest.mark.requires_zarr_write
 def test_identity_selection_keeps_lazy_zarr_vector(tmp_path):
     """A lazy parent spike vector should stay lazy through an identity selection."""
     from spikeinterface.core.zarrextractors import ZarrSortingExtractor, ZarrSpikeVector
@@ -262,6 +263,11 @@ def test_identity_selection_keeps_lazy_zarr_vector(tmp_path):
     assert renamed._get_spike_vector_segment_slices() is lazy_parent._get_spike_vector_segment_slices()
 
     eager_parent = ZarrSortingExtractor(folder)
+    frames = [-1, 0, 5, 7, 10_000]
+    np.testing.assert_array_equal(
+        renamed.search_cached_spikes_sorted(frames, segment_index=0),
+        eager_parent.search_cached_spikes_sorted(frames, segment_index=0),
+    )
     subset = lazy_parent.select_units(["29", "22"])
     assert isinstance(subset.to_spike_vector(), np.ndarray)
     assert np.array_equal(subset.to_spike_vector(), _mask_and_remap(eager_parent, ["29", "22"]))

src/spikeinterface/core/tests/test_zarrextractors.py

@@ -102,6 +102,7 @@ def test_ZarrSampleIndexSearch(tmp_path):
         np.testing.assert_array_equal(search.searchsorted([5], 10, 10), [0])
 
 
+@pytest.mark.requires_zarr_write
 def test_ZarrSortingExtractor_lazy_search(tmp_path):
     sorting = generate_sorting(num_units=10, durations=[5.0, 3.0, 4.0], firing_rates=40.0, seed=0)
     folder = tmp_path / "sorting.zarr"

src/spikeinterface/core/unitsselectionsorting.py

@@ -71,6 +71,16 @@ def _compute_and_cache_spike_vector(self) -> None:
         )
         self._cached_spike_vector = spike_vector
 
+    def search_cached_spikes_sorted(
+        self,
+        indices: list[int],
+        segment_index: int | None = None,
+    ):
+        if self._is_identity_selection:
+            # same spikes in the same positions, and the parent knows how to search a lazy vector
+            return self._parent_sorting.search_cached_spikes_sorted(indices=indices, segment_index=segment_index)
+        return super().search_cached_spikes_sorted(indices=indices, segment_index=segment_index)
+
 
 class UnitsSelectionSortingSegment(BaseSortingSegment):
     def __init__(self, parent_segment, ids_conversion):