From 3a8799c8612e27577ed5dd0eca89984408ec1fe0 Mon Sep 17 00:00:00 2001 From: Graham Findlay Date: Thu, 8 Oct 2026 14:58:17 -0500 Subject: [PATCH 1/2] Override search_cached_spikes_sorted for identity selections In #4783 an identity selection of a lazy Zarr array shares the parent's lazy vector, but after #4809 UnitSelectionSorting inherited the BaseSorting's search_cached_spike_sorted, which would see a non-ndarray and raise an error. The fix here is to just override the method and delegate the search to the parent for identity selections. Since the identity selection literally shares the parent's vector and segment slices, the parent's search is guaranteed to give the right positions. --- .../core/tests/test_unitsselectionsorting.py | 5 +++++ src/spikeinterface/core/unitsselectionsorting.py | 10 ++++++++++ 2 files changed, 15 insertions(+) diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py index 0b9cf6c813..c240d25d88 100644 --- a/src/spikeinterface/core/tests/test_unitsselectionsorting.py +++ b/src/spikeinterface/core/tests/test_unitsselectionsorting.py @@ -262,6 +262,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"])) diff --git a/src/spikeinterface/core/unitsselectionsorting.py b/src/spikeinterface/core/unitsselectionsorting.py index 31fbe2098f..33fec6fa94 100644 --- a/src/spikeinterface/core/unitsselectionsorting.py +++ b/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): From 1e660706451dcea177fba370f8c1017167fa39fd Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 8 Oct 2026 16:33:37 +0200 Subject: [PATCH 2/2] test: add missing require zarr write (cherry picked from commit edf5334c9d37458d9cff2178ef8e11f3b417a6be) --- src/spikeinterface/core/tests/test_unitsselectionsorting.py | 1 + src/spikeinterface/core/tests/test_zarrextractors.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/spikeinterface/core/tests/test_unitsselectionsorting.py b/src/spikeinterface/core/tests/test_unitsselectionsorting.py index c240d25d88..309563df62 100644 --- a/src/spikeinterface/core/tests/test_unitsselectionsorting.py +++ b/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 diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index e1b315e592..2b388839d1 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/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"