From feeccf0322042de5a55953ae0b6fde75c78ce7d6 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Tue, 29 Sep 2026 11:10:40 +0200
Subject: [PATCH 01/12] perf: speed up and reduce RAM for SpikeRetriever

---
 src/spikeinterface/core/basesorting.py   | 15 ++++--
 src/spikeinterface/core/node_pipeline.py | 69 +++++++++++++++---------
 2 files changed, 54 insertions(+), 30 deletions(-)

diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py
index 1e136f2dff..d6ab8b7601 100644
--- a/src/spikeinterface/core/basesorting.py
+++ b/src/spikeinterface/core/basesorting.py
@@ -1076,11 +1076,16 @@ def _get_spike_vector_segment_slices(self):
         if self._cached_spike_vector_segment_slices is None:
             # compute the, this is needed when spikevector is loaded from format and not computed
             num_seg = self.get_num_segments()
-            slices = np.searchsorted(self._cached_spike_vector["segment_index"], np.arange(num_seg + 1))
-            self._cached_spike_vector_segment_slices = np.zeros((num_seg, 2), dtype="int64")
-            for seg_index in range(num_seg):
-                self._cached_spike_vector_segment_slices[seg_index, 0] = slices[seg_index]
-                self._cached_spike_vector_segment_slices[seg_index, 1] = slices[seg_index + 1]
+            if num_seg == 1:
+                self._cached_spike_vector_segment_slices = np.array(
+                    [[0, self._cached_spike_vector.size]], dtype="int64"
+                )
+            else:
+                slices = np.searchsorted(self._cached_spike_vector["segment_index"], np.arange(num_seg + 1))
+                self._cached_spike_vector_segment_slices = np.zeros((num_seg, 2), dtype="int64")
+                for seg_index in range(num_seg):
+                    self._cached_spike_vector_segment_slices[seg_index, 0] = slices[seg_index]
+                    self._cached_spike_vector_segment_slices[seg_index, 1] = slices[seg_index + 1]
         return self._cached_spike_vector_segment_slices
 
     def to_reordered_spike_vector(
diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index 6ecf66918d..f09bd036f2 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -110,10 +110,13 @@ def __init__(self, recording, peaks):
         self.peaks = peaks
 
         # precompute segment slice
-        self.segment_slices = []
-        for segment_index in range(recording.get_num_segments()):
-            i0, i1 = np.searchsorted(peaks["segment_index"], [segment_index, segment_index + 1])
-            self.segment_slices.append(slice(i0, i1))
+        if recording.get_num_segments() > 1:
+            self.segment_slices = []
+            for segment_index in range(recording.get_num_segments()):
+                i0, i1 = np.searchsorted(peaks["segment_index"], [segment_index, segment_index + 1])
+                self.segment_slices.append(slice(i0, i1))
+        else:
+            self.segment_slices = None
 
     def get_margin(self):
         return 0
@@ -122,15 +125,22 @@ def get_dtype(self):
         return base_peak_dtype
 
     def get_peak_slice(self, segment_index, start_frame, end_frame, max_margin):
-        sl = self.segment_slices[segment_index]
-        peaks_in_segment = self.peaks[sl]
+        if self.segment_slices is not None:
+            sl = self.segment_slices[segment_index]
+            peaks_in_segment = self.peaks[sl]
+        else:
+            peaks_in_segment = self.peaks
         i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
         return i0, i1
 
     def compute(self, traces, start_frame, end_frame, segment_index, max_margin, peak_slice):
         # get local peaks
-        sl = self.segment_slices[segment_index]
-        peaks_in_segment = self.peaks[sl]
+        if self.segment_slices is not None:
+            sl = self.segment_slices[segment_index]
+            peaks_in_segment = self.peaks[sl]
+        else:
+            peaks_in_segment = self.peaks
+
         # i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
         i0, i1 = peak_slice
         local_peaks = peaks_in_segment[i0:i1]
@@ -204,19 +214,20 @@ def __init__(
 
         main_channel_ids = sorting.get_property("main_channel_id")
         assert main_channel_ids is not None, "SpikeRetriever needs the sorting to have `main_channel_id`s."
-        main_channel_indices = recording.ids_to_indices(main_channel_ids)
-        self.peaks = sorting_to_peaks(sorting, main_channel_indices, self._dtype)
+        self.main_channel_indices = recording.ids_to_indices(main_channel_ids)
+        self.spike_vector, segment_slices = sorting.to_spike_vector(return_slices=True)
+        self.spike_sample_indices = np.asarray(self.spike_vector["sample_index"])
 
         if not channel_from_template:
             channel_distance = get_channel_distances(recording)
             self.neighbours_mask = channel_distance <= radius_um
             self.peak_sign = peak_sign
 
-        # precompute segment slice
-        self.segment_slices = []
-        for segment_index in range(recording.get_num_segments()):
-            i0, i1 = np.searchsorted(self.peaks["segment_index"], [segment_index, segment_index + 1])
-            self.segment_slices.append(slice(i0, i1))
+        # For mono-segment, we avoid an extra slice, otherwise make them a tuple for slicing
+        if sorting.get_num_segments() == 1:
+            self.segment_slices = None
+        else:
+            self.segment_slices = [slice(int(s0), int(s1)) for s0, s1 in segment_slices]
 
     def get_margin(self):
         return 0
@@ -225,26 +236,34 @@ def get_dtype(self):
         return self._dtype
 
     def get_peak_slice(self, segment_index, start_frame, end_frame, max_margin):
-        sl = self.segment_slices[segment_index]
-        peaks_in_segment = self.peaks[sl]
+        if self.segment_slices is not None:
+            sl = self.segment_slices[segment_index]
+            sample_indices_in_segment = self.spike_sample_indices[sl]
+        else:
+            sample_indices_in_segment = self.spike_sample_indices
         if self.include_spikes_in_margin:
-            i0, i1 = np.searchsorted(
-                peaks_in_segment["sample_index"], [start_frame - max_margin, end_frame + max_margin]
-            )
+            i0, i1 = np.searchsorted(sample_indices_in_segment, [start_frame - max_margin, end_frame + max_margin])
         else:
-            i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
+            i0, i1 = np.searchsorted(sample_indices_in_segment, [start_frame, end_frame])
         return i0, i1
 
     def compute(self, traces, start_frame, end_frame, segment_index, max_margin, peak_slice):
         # get local peaks
-        sl = self.segment_slices[segment_index]
-        peaks_in_segment = self.peaks[sl]
         i0, i1 = peak_slice
+        if self.segment_slices is not None:
+            s0 = self.segment_slices[segment_index].start
+            spikes = self.spike_vector[s0 + i0 : s0 + i1]
+        else:
+            spikes = self.spike_vector[i0:i1]
 
-        local_peaks = peaks_in_segment[i0:i1]
+        local_peaks = np.zeros(spikes.size, dtype=self._dtype)
+        local_peaks["sample_index"] = spikes["sample_index"]
+        local_peaks["channel_index"] = self.main_channel_indices[spikes["unit_index"]]
+        local_peaks["amplitude"] = 0.0
+        local_peaks["segment_index"] = spikes["segment_index"]
+        local_peaks["unit_index"] = spikes["unit_index"]
 
         # make sample index local to traces
-        local_peaks = local_peaks.copy()
         local_peaks["sample_index"] -= start_frame - max_margin
 
         # handle flag for margin

From 1871f5bd7aabf720928054ca4f4154094d2bad0c Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Tue, 29 Sep 2026 12:38:22 +0200
Subject: [PATCH 02/12] fix: complete node pipeline kwargs and peak locations
 as float32 instead of 64

---
 src/spikeinterface/core/node_pipeline.py      | 35 ++++++++++++++++++-
 .../postprocessing/unit_locations.py          |  8 ++---
 .../preprocessing/detect_artifacts.py         |  9 +++++
 .../waveforms/features_from_peaks.py          |  6 ++--
 4 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index f09bd036f2..0faa82f332 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -51,7 +51,11 @@ def __init__(
             parents = [parents]
         self.parents = parents
 
-        self._kwargs = dict()
+        self._kwargs = dict(
+            time_series=time_series,
+            return_output=return_output,
+            parents=parents,
+        )
 
     def get_margin(self):
         # can optionally be overwritten
@@ -118,6 +122,8 @@ def __init__(self, recording, peaks):
         else:
             self.segment_slices = None
 
+        self._kwargs.update(dict(peaks=peaks))
+
     def get_margin(self):
         return 0
 
@@ -229,6 +235,17 @@ def __init__(
         else:
             self.segment_slices = [slice(int(s0), int(s1)) for s0, s1 in segment_slices]
 
+        self._kwargs.update(
+            dict(
+                sorting=sorting,
+                channel_from_template=channel_from_template,
+                extremum_channel_inds=extremum_channel_inds,
+                radius_um=radius_um,
+                peak_sign=peak_sign,
+                include_spikes_in_margin=include_spikes_in_margin,
+            )
+        )
+
     def get_margin(self):
         return 0
 
@@ -366,6 +383,15 @@ def __init__(
             self.nafter = ms_to_samples(ms_after, sampling_frequency)
         self.neighbours_mask = None
 
+        self._kwargs.update(
+            dict(
+                ms_before=ms_before,
+                ms_after=ms_after,
+                nbefore=nbefore,
+                nafter=nafter,
+            )
+        )
+
 
 class ExtractDenseWaveforms(WaveformsNode):
     def __init__(
@@ -489,6 +515,13 @@ def __init__(
             self.neighbours_mask = self.channel_distance <= radius_um
         self.max_num_chans = np.max(np.sum(self.neighbours_mask, axis=1))
 
+        self._kwargs.update(
+            dict(
+                radius_um=radius_um,
+                sparsity_mask=sparsity_mask,
+            )
+        )
+
     def get_margin(self):
         return max(self.nbefore, self.nafter)
 
diff --git a/src/spikeinterface/postprocessing/unit_locations.py b/src/spikeinterface/postprocessing/unit_locations.py
index 4d913ab4f9..7cd94d7ee9 100644
--- a/src/spikeinterface/postprocessing/unit_locations.py
+++ b/src/spikeinterface/postprocessing/unit_locations.py
@@ -7,10 +7,10 @@
 
 # this dict is for peak location
 dtype_localize_by_method = {
-    "center_of_mass": [("x", "float64"), ("y", "float64")],
-    "grid_convolution": [("x", "float64"), ("y", "float64"), ("z", "float64")],
-    "peak_channel": [("x", "float64"), ("y", "float64")],
-    "monopolar_triangulation": [("x", "float64"), ("y", "float64"), ("z", "float64"), ("alpha", "float64")],
+    "center_of_mass": [("x", "float32"), ("y", "float32")],
+    "grid_convolution": [("x", "float32"), ("y", "float32"), ("z", "float32")],
+    "peak_channel": [("x", "float32"), ("y", "float32")],
+    "monopolar_triangulation": [("x", "float32"), ("y", "float32"), ("z", "float32"), ("alpha", "float32")],
 }
 
 possible_localization_methods = list(dtype_localize_by_method.keys())
diff --git a/src/spikeinterface/preprocessing/detect_artifacts.py b/src/spikeinterface/preprocessing/detect_artifacts.py
index c29199ae3b..ab19c49edc 100644
--- a/src/spikeinterface/preprocessing/detect_artifacts.py
+++ b/src/spikeinterface/preprocessing/detect_artifacts.py
@@ -126,6 +126,15 @@ def __init__(
         else:
             self.diff_threshold_unscaled = None
 
+        self._kwargs.update(
+            dict(
+                saturation_threshold_uV=saturation_threshold_uV,
+                diff_threshold_uV=diff_threshold_uV,
+                proportion=proportion,
+                signed=signed,
+            )
+        )
+
     def get_margin(self) -> int:
         """Return the number of margin samples required on each side of a chunk."""
         return 0
diff --git a/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py b/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py
index 1695f2cc59..0c7ffe7447 100644
--- a/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py
+++ b/src/spikeinterface/sortingcomponents/waveforms/features_from_peaks.py
@@ -91,8 +91,8 @@ def __init__(
 
         self.all_channels = all_channels
         self.peak_sign = peak_sign
-        self._kwargs.update(dict(all_channels=all_channels, peak_sign=peak_sign))
         self._dtype = recording.get_dtype()
+        self._kwargs.update(dict(all_channels=all_channels, peak_sign=peak_sign))
 
     def get_dtype(self):
         return self._dtype
@@ -125,8 +125,8 @@ def __init__(
         self.channel_distance = get_channel_distances(recording)
         self.neighbours_mask = self.channel_distance <= radius_um
         self.all_channels = all_channels
-        self._kwargs.update(dict(radius_um=radius_um, all_channels=all_channels))
         self._dtype = recording.get_dtype()
+        self._kwargs.update(dict(radius_um=radius_um, all_channels=all_channels))
 
     def get_dtype(self):
         return self._dtype
@@ -168,6 +168,7 @@ def __init__(
         self.radius_um = radius_um
         self.sparse = sparse
         self.noise_threshold = noise_threshold
+        self._dtype = recording.get_dtype()
         self._kwargs.update(
             dict(
                 projections=projections,
@@ -177,7 +178,6 @@ def __init__(
                 feature=feature,
             )
         )
-        self._dtype = recording.get_dtype()
 
     def get_dtype(self):
         return self._dtype

From 1f1bdd2c64bc9deb8eeac0920ffc753ecec9f330 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Tue, 29 Sep 2026 12:58:16 +0200
Subject: [PATCH 03/12] fix: amplitude_scalings float32 dtype

---
 src/spikeinterface/postprocessing/amplitude_scalings.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/spikeinterface/postprocessing/amplitude_scalings.py b/src/spikeinterface/postprocessing/amplitude_scalings.py
index d28d2ab1c0..69a1420638 100644
--- a/src/spikeinterface/postprocessing/amplitude_scalings.py
+++ b/src/spikeinterface/postprocessing/amplitude_scalings.py
@@ -178,13 +178,12 @@ def __init__(
         PipelineNode.__init__(self, recording, parents=parents, return_output=return_output)
         self.return_in_uV = return_in_uV
         if return_in_uV and recording.has_scaleable_traces():
-            self._dtype = np.float32
             self._gains = recording.get_channel_gains()
             self._offsets = recording.get_channel_offsets()
         else:
-            self._dtype = recording.get_dtype()
             self._gains = None
             self._offsets = None
+        self._dtype = np.float32
         spike_retriever = find_parent_of_type(parents, SpikeRetriever)
         assert isinstance(
             spike_retriever, SpikeRetriever
@@ -268,7 +267,7 @@ def compute(self, traces, peaks):
             collisions = {}
 
         # compute the scaling for each spike
-        scalings = np.zeros(len(local_spikes), dtype=float)
+        scalings = np.zeros(len(local_spikes), dtype=self._dtype)
         spike_collision_mask = np.zeros(len(local_spikes), dtype=bool)
 
         for spike_index, spike in enumerate(local_spikes):

From ea80b9538d8113c9fd23c750288e0dd9f4b7d7ac Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Tue, 29 Sep 2026 18:40:33 +0200
Subject: [PATCH 04/12] fix: add peaks property to SpikeRetriever

---
 src/spikeinterface/core/node_pipeline.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index 0faa82f332..5da8678c03 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -223,6 +223,7 @@ def __init__(
         self.main_channel_indices = recording.ids_to_indices(main_channel_ids)
         self.spike_vector, segment_slices = sorting.to_spike_vector(return_slices=True)
         self.spike_sample_indices = np.asarray(self.spike_vector["sample_index"])
+        self._peaks = None
 
         if not channel_from_template:
             channel_distance = get_channel_distances(recording)
@@ -246,6 +247,13 @@ def __init__(
             )
         )
 
+    @property
+    def peaks(self):
+        if self._peaks is not None:
+            return self._peaks
+        self._peaks = sorting_to_peaks(self.sorting)
+        return self._peaks
+
     def get_margin(self):
         return 0
 

From 90043c501cecc3377b38f155e2a0e3265e0dc8aa Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Tue, 29 Sep 2026 18:58:54 +0200
Subject: [PATCH 05/12] fix: add sorting attribute to SpikeRetriever

---
 src/spikeinterface/core/node_pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index 5da8678c03..dbc8840a86 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -211,7 +211,6 @@ def __init__(
                 category=FutureWarning,
                 stacklevel=2,
             )
-
         self._dtype = spike_peak_dtype
 
         self.include_spikes_in_margin = include_spikes_in_margin
@@ -223,6 +222,7 @@ def __init__(
         self.main_channel_indices = recording.ids_to_indices(main_channel_ids)
         self.spike_vector, segment_slices = sorting.to_spike_vector(return_slices=True)
         self.spike_sample_indices = np.asarray(self.spike_vector["sample_index"])
+        self.sorting = sorting
         self._peaks = None
 
         if not channel_from_template:

From 4e8e9f9e83047663f6ef17408af1d3ebd078ae1e Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Thu, 1 Oct 2026 11:55:09 +0200
Subject: [PATCH 06/12] fix: sorting_to_peaks need main_channel_indices

---
 src/spikeinterface/core/node_pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index dbc8840a86..dec8d6b617 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -251,7 +251,7 @@ def __init__(
     def peaks(self):
         if self._peaks is not None:
             return self._peaks
-        self._peaks = sorting_to_peaks(self.sorting)
+        self._peaks = sorting_to_peaks(self.sorting, self.main_channel_indices)
         return self._peaks
 
     def get_margin(self):

From fb34821a3f7ac8d663001bda75d599a30ce9c48f Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Mon, 5 Oct 2026 12:58:09 +0200
Subject: [PATCH 07/12] wip: search_cached_spikes_sorted

---
 src/spikeinterface/core/basesorting.py    | 13 +++++++++++++
 src/spikeinterface/core/node_pipeline.py  | 21 ++++++---------------
 src/spikeinterface/core/zarrextractors.py | 16 +++++++++++++++-
 3 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py
index d6ab8b7601..4660de6e14 100644
--- a/src/spikeinterface/core/basesorting.py
+++ b/src/spikeinterface/core/basesorting.py
@@ -149,6 +149,19 @@ def get_total_duration(self) -> float:
         ), "This methods requires an associated recording. Call self.register_recording() first."
         return self._recording.get_total_duration()
 
+    def search_cached_spikes_sorted(
+        self,
+        indices: list[int],
+        segment_index: int | None = None,
+    ):
+        if self._cached_spike_vector is None:
+            self._compute_and_cache_spike_vector()
+        if segment_index is None and self.get_num_segments() == 1:
+            return np.searchsorted(self._cached_spike_vector["sample_index"], indices)
+        else:
+            sl = self._cached_spike_vector_segment_slices[segment_index]
+            return np.searchsorted(self._cached_spike_vector[sl[0] : sl[1]]["sample_index"], indices)
+
     def get_unit_spike_train(
         self,
         unit_id: str | int,
diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index dec8d6b617..87fee796e3 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -221,7 +221,6 @@ def __init__(
         assert main_channel_ids is not None, "SpikeRetriever needs the sorting to have `main_channel_id`s."
         self.main_channel_indices = recording.ids_to_indices(main_channel_ids)
         self.spike_vector, segment_slices = sorting.to_spike_vector(return_slices=True)
-        self.spike_sample_indices = np.asarray(self.spike_vector["sample_index"])
         self.sorting = sorting
         self._peaks = None
 
@@ -230,12 +229,6 @@ def __init__(
             self.neighbours_mask = channel_distance <= radius_um
             self.peak_sign = peak_sign
 
-        # For mono-segment, we avoid an extra slice, otherwise make them a tuple for slicing
-        if sorting.get_num_segments() == 1:
-            self.segment_slices = None
-        else:
-            self.segment_slices = [slice(int(s0), int(s1)) for s0, s1 in segment_slices]
-
         self._kwargs.update(
             dict(
                 sorting=sorting,
@@ -261,16 +254,14 @@ def get_dtype(self):
         return self._dtype
 
     def get_peak_slice(self, segment_index, start_frame, end_frame, max_margin):
-        if self.segment_slices is not None:
-            sl = self.segment_slices[segment_index]
-            sample_indices_in_segment = self.spike_sample_indices[sl]
-        else:
-            sample_indices_in_segment = self.spike_sample_indices
         if self.include_spikes_in_margin:
-            i0, i1 = np.searchsorted(sample_indices_in_segment, [start_frame - max_margin, end_frame + max_margin])
+            indices = [start_frame - max_margin, end_frame + max_margin]
         else:
-            i0, i1 = np.searchsorted(sample_indices_in_segment, [start_frame, end_frame])
-        return i0, i1
+            indices = [start_frame, end_frame]
+        return self.sorting.search_cached_spikes_sorted(
+            indices=indices,
+            segment_index=segment_index,
+        )
 
     def compute(self, traces, start_frame, end_frame, segment_index, max_margin, peak_slice):
         # get local peaks
diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
index 78c10df8d4..14d4a24fc1 100644
--- a/src/spikeinterface/core/zarrextractors.py
+++ b/src/spikeinterface/core/zarrextractors.py
@@ -524,7 +524,7 @@ def __init__(
             # we do not need to lexsort at init (very high cost) because there already sorted by frame before to be saved.
             # In version 0.104.X this was fully lexsorted, but we don't need it anymore because it's only important in the context of SpikeVectorBased extensions in the SortingAnalyzer, which stores its own copy of the Sorting object. This makes the extension data and the spike vector always matching their order.
             # spikes = spikes[np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"]))]
-
+        self._lazy_spike_vector = lazy_spike_vector
         self._cached_spike_vector = spikes
         # pre-populate segment slices so _get_spike_vector_segment_slices() never
         # needs to materialise the full segment_index array
@@ -553,6 +553,20 @@ def __init__(
             "lazy_spike_vector": lazy_spike_vector,
         }
 
+    def search_cached_spikes_sorted(
+        self,
+        indices: list[int],
+        segment_index: int | None = None,
+    ):
+        if not self._lazy_spike_vector:
+            return super().search_cached_spikes_sorted(
+                indices=indices,
+                segment_index=segment_index,
+            )
+        else:
+            # TODO: implement search function minimizing zarr chunks reads to reduce RAM usage.
+            pass
+
     @staticmethod
     def write_sorting(
         sorting: BaseSorting,

From bee887dccdb8eafe656e17dbbfbae908558d3edc Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Mon, 5 Oct 2026 15:18:08 +0200
Subject: [PATCH 08/12] feat: implement search_cached_spikes_sorted and lazy
 zarr searchsorted

---
 src/spikeinterface/core/basesorting.py        | 35 ++++++-
 src/spikeinterface/core/node_pipeline.py      | 10 ++
 src/spikeinterface/core/sortinganalyzer.py    | 19 +---
 .../core/tests/test_zarrextractors.py         | 59 ++++++++++++
 src/spikeinterface/core/zarrextractors.py     | 93 ++++++++++++++++++-
 5 files changed, 194 insertions(+), 22 deletions(-)

diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py
index e74c507789..fdd15e08e6 100644
--- a/src/spikeinterface/core/basesorting.py
+++ b/src/spikeinterface/core/basesorting.py
@@ -154,13 +154,38 @@ def search_cached_spikes_sorted(
         indices: list[int],
         segment_index: int | None = None,
     ):
+        """
+        Search sample indices (frames) in the cached spike vector of one segment.
+
+        Equivalent to `np.searchsorted(segment_sample_index, indices, side="left")`.
+        Sortings with a lazy spike vector override this to avoid materialising it.
+
+        Parameters
+        ----------
+        indices : list[int]
+            The sample indices (frames) to search.
+        segment_index : int | None, default: None
+            The segment to search. Can be None for mono-segment sortings.
+
+        Returns
+        -------
+        positions : np.ndarray
+            The insertion positions, relative to the start of the segment in the spike vector.
+        """
         if self._cached_spike_vector is None:
             self._compute_and_cache_spike_vector()
-        if segment_index is None and self.get_num_segments() == 1:
-            return np.searchsorted(self._cached_spike_vector["sample_index"], indices)
-        else:
-            sl = self._cached_spike_vector_segment_slices[segment_index]
-            return np.searchsorted(self._cached_spike_vector[sl[0] : sl[1]]["sample_index"], indices)
+        spikes = self._cached_spike_vector
+        if not isinstance(spikes, np.ndarray):  # np.memmap is an ndarray
+            # np.searchsorted would materialise a lazy vector on every call
+            raise TypeError(
+                f"{type(self).__name__} holds a lazy spike vector ({type(spikes).__name__}) "
+                "and must override search_cached_spikes_sorted()"
+            )
+        if segment_index is None:
+            assert self.get_num_segments() == 1, "segment_index is required for multi-segment sortings"
+            segment_index = 0
+        start, stop = self._get_spike_vector_segment_slices()[segment_index]
+        return np.searchsorted(spikes["sample_index"][start:stop], indices)
 
     def get_unit_spike_train(
         self,
diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index 87fee796e3..7978ada949 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -224,6 +224,16 @@ def __init__(
         self.sorting = sorting
         self._peaks = None
 
+        # get_peak_slice() returns positions relative to the segment start
+        if sorting.get_num_segments() == 1:
+            self.segment_slices = None
+        else:
+            self.segment_slices = [slice(int(s0), int(s1)) for s0, s1 in segment_slices]
+
+        # build any lazy search index now, in the parent, so forked workers inherit it
+        # instead of each building its own
+        sorting.search_cached_spikes_sorted([0], segment_index=0)
+
         if not channel_from_template:
             channel_distance = get_channel_distances(recording)
             self.neighbours_mask = channel_distance <= radius_um
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index 1c6604516b..42cfd41792 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -1171,22 +1171,13 @@ def load_from_zarr(
         settings = cls._handle_backward_compatibility_settings_pre_init(settings)
 
         # Load sorting (in memory or lazy)
+        zarr_sorting = ZarrSortingExtractor(
+            folder, zarr_group="sorting", storage_options=storage_options, lazy_spike_vector=lazy
+        )
         if lazy:
-            copy_spike_vector = False
-            lazy_spike_vector = True
+            sorting = zarr_sorting
         else:
-            copy_spike_vector = True
-            lazy_spike_vector = False
-        sorting = NumpySorting.from_sorting(
-            ZarrSortingExtractor(
-                folder,
-                zarr_group="sorting",
-                storage_options=storage_options,
-                lazy_spike_vector=lazy_spike_vector,
-            ),
-            with_metadata=True,
-            copy_spike_vector=copy_spike_vector,
-        )
+            sorting = NumpySorting.from_sorting(zarr_sorting, with_metadata=True, copy_spike_vector=True)
 
         # Load recording (if available)
         if recording is None:
diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py
index 6278541698..8c3a08da0f 100644
--- a/src/spikeinterface/core/tests/test_zarrextractors.py
+++ b/src/spikeinterface/core/tests/test_zarrextractors.py
@@ -1,6 +1,7 @@
 import pytest
 from pathlib import Path
 
+import numpy as np
 import zarr
 
 from spikeinterface.core import (
@@ -10,6 +11,7 @@
 )
 from spikeinterface.core.zarrextractors import (
     ZarrRecordingExtractor,
+    ZarrSampleIndexSearch,
     ZarrSortingExtractor,
     add_sorting_to_zarr_group,
     get_default_zarr_compressor,
@@ -78,7 +80,64 @@ def test_ZarrSortingExtractor(tmp_path):
     sorting = load(sorting.to_dict())
 
 
+def test_ZarrSampleIndexSearch(tmp_path):
+    rng = np.random.default_rng(0)
+    # two "segments", each sorted, with long runs of equal values so that runs cross
+    # the (tiny) chunk boundaries
+    segments = [np.sort(rng.integers(0, 40, size=101)), np.sort(rng.integers(0, 25, size=58))]
+    sample_index = np.concatenate(segments)
+    bounds = np.cumsum([0] + [len(s) for s in segments])
+    z = zarr.open(tmp_path / "sample_index.zarr", mode="w", shape=sample_index.shape, chunks=(7,), dtype="int64")
+    z[:] = sample_index
+
+    values = np.arange(-3, 45)
+    for chunk_firsts in (sample_index[::7], None):
+        search = ZarrSampleIndexSearch(z, chunk_firsts)
+        for start, stop in zip(bounds[:-1], bounds[1:]):
+            expected = np.searchsorted(sample_index[start:stop], values, side="left")
+            np.testing.assert_array_equal(search.searchsorted(values, start, stop), expected)
+        # empty range
+        np.testing.assert_array_equal(search.searchsorted([5], 10, 10), [0])
+
+
+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"
+    ZarrSortingExtractor.write_sorting(sorting, folder)
+    # re-store sample_index in small chunks, so that the search crosses chunk boundaries
+    # and segments start in the middle of a chunk
+    spikes_group = zarr.open(folder, mode="a")["spikes"]
+    sample_index = spikes_group["sample_index"][:]
+    del spikes_group["sample_index"], spikes_group["sample_index_chunk_firsts"]
+    spikes_group.create_dataset("sample_index", data=sample_index, chunks=(97,))
+    spikes_group.create_dataset("sample_index_chunk_firsts", data=sample_index[::97], compressor=None)
+    assert spikes_group["sample_index"].nchunks > 3
+
+    in_ram = ZarrSortingExtractor(folder)
+    lazy = ZarrSortingExtractor(folder, lazy_spike_vector=True)
+    assert type(lazy.to_spike_vector()).__name__ == "ZarrSpikeVector"
+
+    rng = np.random.default_rng(1)
+    num_samples = int(sorting.to_spike_vector()["sample_index"].max()) + 100
+    for segment_index in range(sorting.get_num_segments()):
+        frames = np.sort(rng.integers(-10, num_samples, size=200))
+        expected = in_ram.search_cached_spikes_sorted(frames, segment_index=segment_index)
+        np.testing.assert_array_equal(lazy.search_cached_spikes_sorted(frames, segment_index=segment_index), expected)
+
+    # stores written before the chunk index existed rebuild it from the data
+    del zarr.open(folder, mode="a")["spikes/sample_index_chunk_firsts"]
+    old = ZarrSortingExtractor(folder, lazy_spike_vector=True)
+    frames = np.arange(-5, num_samples, 37)
+    for segment_index in range(sorting.get_num_segments()):
+        np.testing.assert_array_equal(
+            old.search_cached_spikes_sorted(frames, segment_index=segment_index),
+            in_ram.search_cached_spikes_sorted(frames, segment_index=segment_index),
+        )
+
+
 if __name__ == "__main__":
     tmp_path = Path("tmp")
     test_zarr_compression_options(tmp_path)
     test_ZarrSortingExtractor(tmp_path)
+    test_ZarrSampleIndexSearch(tmp_path)
+    test_ZarrSortingExtractor_lazy_search(tmp_path)
diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
index f83efe8d61..8a5a9b2d8f 100644
--- a/src/spikeinterface/core/zarrextractors.py
+++ b/src/spikeinterface/core/zarrextractors.py
@@ -457,6 +457,74 @@ def copy(self):
         return np.copy(np.asarray(self))
 
 
+class ZarrSampleIndexSearch:
+    """
+    `np.searchsorted` on a zarr `sample_index` array without materialising it.
+
+    The array is sorted within each segment and stored in fixed-size zarr chunks. The
+    search keeps the first value of every chunk (8 B per chunk), picks the one chunk that
+    can hold the answer, and decodes only that chunk. The last decoded chunk is cached, so
+    consecutive searches in the same region decode nothing.
+
+    Parameters
+    ----------
+    sample_index : zarr.Array
+        The 1D `sample_index` array of a spike vector.
+    chunk_firsts : np.ndarray | None, default: None
+        `sample_index[::chunk_length]`, as written by `add_sorting_to_zarr_group`.
+        If None (stores written before this index existed), it is read from the array,
+        which decodes every chunk once.
+    """
+
+    def __init__(self, sample_index, chunk_firsts=None):
+        self._sample_index = sample_index
+        self._num_spikes = sample_index.shape[0]
+        self._chunk_length = sample_index.chunks[0]
+        if chunk_firsts is None:
+            starts = np.arange(0, self._num_spikes, self._chunk_length)
+            chunk_firsts = sample_index.get_orthogonal_selection(starts) if starts.size else []
+        self._chunk_firsts = np.asarray(chunk_firsts, dtype="int64")
+        self._cached_chunk = (-1, None)  # (chunk index, decoded values)
+
+    def _get_chunk(self, chunk_index):
+        # read the (index, values) pair once into a local and replace it as a whole: threads
+        # sharing this object (pool_engine="thread") then always see a consistent pair, and a
+        # race can only cost an extra decode
+        cached = self._cached_chunk
+        if cached[0] != chunk_index:
+            start = chunk_index * self._chunk_length
+            stop = min(start + self._chunk_length, self._num_spikes)
+            cached = (chunk_index, self._sample_index[start:stop])
+            self._cached_chunk = cached
+        return cached[1]
+
+    def searchsorted(self, values, start, stop):
+        """
+        Equivalent to `np.searchsorted(sample_index[start:stop], values, side="left")`.
+
+        Returns positions relative to `start`. `start:stop` must be a range over which
+        `sample_index` is sorted, i.e. one segment.
+        """
+        values = np.atleast_1d(np.asarray(values))
+        out = np.zeros(values.size, dtype="int64")
+        if stop <= start:
+            return out
+        first_chunk = start // self._chunk_length
+        last_chunk = (stop - 1) // self._chunk_length
+        # first values of the chunks after `first_chunk`, all of which lie inside start:stop
+        later_firsts = self._chunk_firsts[first_chunk + 1 : last_chunk + 1]
+        for i, value in enumerate(values):
+            # the answer lies in the last chunk whose first value is < value: side="left" keeps
+            # runs of equal values that cross a chunk boundary in the earlier chunk
+            chunk_index = first_chunk + int(np.searchsorted(later_firsts, value, side="left"))
+            chunk_start = chunk_index * self._chunk_length
+            lo = max(chunk_start, start)
+            hi = min(chunk_start + self._chunk_length, stop)
+            block = self._get_chunk(chunk_index)[lo - chunk_start : hi - chunk_start]
+            out[i] = lo + int(np.searchsorted(block, value, side="left")) - start
+        return out
+
+
 class ZarrSortingExtractor(BaseSorting):
     """
     SortingExtractor for a zarr format
@@ -528,6 +596,8 @@ def __init__(
             # In version 0.104.X this was fully lexsorted, but we don't need it anymore because it's only important in the context of SpikeVectorBased extensions in the SortingAnalyzer, which stores its own copy of the Sorting object. This makes the extension data and the spike vector always matching their order.
             # spikes = spikes[np.lexsort((spikes["unit_index"], spikes["sample_index"], spikes["segment_index"]))]
         self._lazy_spike_vector = lazy_spike_vector
+        self._spikes_group = spikes_group
+        self._sample_index_search = None
         self._cached_spike_vector = spikes
         # pre-populate segment slices so _get_spike_vector_segment_slices() never
         # needs to materialise the full segment_index array
@@ -566,9 +636,17 @@ def search_cached_spikes_sorted(
                 indices=indices,
                 segment_index=segment_index,
             )
-        else:
-            # TODO: implement search function minimizing zarr chunks reads to reduce RAM usage.
-            pass
+        if segment_index is None:
+            assert self.get_num_segments() == 1, "segment_index is required for multi-segment sortings"
+            segment_index = 0
+        if self._sample_index_search is None:
+            # the chunk index is written with the sorting; older stores rebuild it from the data
+            chunk_firsts = None
+            if "sample_index_chunk_firsts" in self._spikes_group:
+                chunk_firsts = self._spikes_group["sample_index_chunk_firsts"][:]
+            self._sample_index_search = ZarrSampleIndexSearch(self._spikes_group["sample_index"], chunk_firsts)
+        start, stop = self._cached_spike_vector_segment_slices[segment_index]
+        return self._sample_index_search.searchsorted(indices, int(start), int(stop))
 
     @staticmethod
     def write_sorting(
@@ -799,6 +877,15 @@ def add_sorting_to_zarr_group(
                 segment_slices.append([i0, i1])
             spikes_group.create_dataset(name="segment_slices", data=segment_slices, compressor=None)
 
+    # first sample_index of every zarr chunk: lets a lazy reader search sample_index
+    # one chunk at a time (see ZarrSampleIndexSearch) instead of materialising it
+    chunk_length = spikes_group["sample_index"].chunks[0]
+    spikes_group.create_dataset(
+        name="sample_index_chunk_firsts",
+        data=np.asarray(spikes["sample_index"][::chunk_length], dtype="int64"),
+        compressor=None,
+    )
+
     add_properties_and_annotations(zarr_group, sorting)
 
 

From c787e254267f27437b17b4df6bc7eefd0ad810f5 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Wed, 7 Oct 2026 12:35:21 +0200
Subject: [PATCH 09/12] fix: remove direct access to retriever peaks

---
 src/spikeinterface/core/node_pipeline.py      | 17 ++++-------
 .../postprocessing/spike_locations.py         |  5 ++++
 .../tests/test_peak_detection.py              |  2 +-
 .../peak_localization/main.py                 | 17 ++++++++++-
 .../sortingcomponents/peak_selection.py       | 21 +++++++++++---
 src/spikeinterface/sortingcomponents/tools.py | 28 +++++++++++++++----
 6 files changed, 67 insertions(+), 23 deletions(-)

diff --git a/src/spikeinterface/core/node_pipeline.py b/src/spikeinterface/core/node_pipeline.py
index 7978ada949..0287010e73 100644
--- a/src/spikeinterface/core/node_pipeline.py
+++ b/src/spikeinterface/core/node_pipeline.py
@@ -111,7 +111,7 @@ class PeakRetriever(PeakSource):
     def __init__(self, recording, peaks):
         PipelineNode.__init__(self, recording, return_output=False)
 
-        self.peaks = peaks
+        self._peaks = peaks
 
         # precompute segment slice
         if recording.get_num_segments() > 1:
@@ -133,9 +133,9 @@ def get_dtype(self):
     def get_peak_slice(self, segment_index, start_frame, end_frame, max_margin):
         if self.segment_slices is not None:
             sl = self.segment_slices[segment_index]
-            peaks_in_segment = self.peaks[sl]
+            peaks_in_segment = self._peaks[sl]
         else:
-            peaks_in_segment = self.peaks
+            peaks_in_segment = self._peaks
         i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
         return i0, i1
 
@@ -143,9 +143,9 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, pea
         # get local peaks
         if self.segment_slices is not None:
             sl = self.segment_slices[segment_index]
-            peaks_in_segment = self.peaks[sl]
+            peaks_in_segment = self._peaks[sl]
         else:
-            peaks_in_segment = self.peaks
+            peaks_in_segment = self._peaks
 
         # i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
         i0, i1 = peak_slice
@@ -250,13 +250,6 @@ def __init__(
             )
         )
 
-    @property
-    def peaks(self):
-        if self._peaks is not None:
-            return self._peaks
-        self._peaks = sorting_to_peaks(self.sorting, self.main_channel_indices)
-        return self._peaks
-
     def get_margin(self):
         return 0
 
diff --git a/src/spikeinterface/postprocessing/spike_locations.py b/src/spikeinterface/postprocessing/spike_locations.py
index 32334e6b3e..d864a0d73e 100644
--- a/src/spikeinterface/postprocessing/spike_locations.py
+++ b/src/spikeinterface/postprocessing/spike_locations.py
@@ -69,19 +69,24 @@ def _get_pipeline_nodes(self):
 
         recording = self.sorting_analyzer.recording
         sorting = self.sorting_analyzer.sorting
+        main_channel_indices = self.sorting_analyzer.get_main_channels(outputs="index", with_dict=False)
 
         retriever_kwargs = {
             "channel_from_template": True,
             **self.params["spike_retriever_kwargs"],
         }
         retriever = SpikeRetriever(sorting, recording, **retriever_kwargs)
+        spikes = sorting.to_spike_vector()
+        channel_index_array = main_channel_indices[spikes["unit_index"]]
         nodes = get_localization_pipeline_nodes(
             recording,
             retriever,
+            peaks_or_spikes=spikes,
             method=self.params["method"],
             method_kwargs=self.params["method_kwargs"],
             ms_before=self.params["ms_before"],
             ms_after=self.params["ms_after"],
+            channel_index_array=channel_index_array,
             seed=self.params.get("seed"),
         )
         return nodes
diff --git a/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py b/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py
index 25b9d9515a..a0b0a40521 100644
--- a/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py
+++ b/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py
@@ -352,7 +352,7 @@ def test_detect_peaks_matched_filtering(recording, job_kwargs):
     ms_before = 1.0
     ms_after = 1.0
     prototype, _, _ = get_prototype_and_waveforms_from_peaks(
-        recording, peaks=peaks_le, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs
+        recording, peaks_or_spikes=peaks_le, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs
     )
 
     peaks_mf_neg = detect_peaks(
diff --git a/src/spikeinterface/sortingcomponents/peak_localization/main.py b/src/spikeinterface/sortingcomponents/peak_localization/main.py
index c788f8b7fa..ffbdb5973d 100644
--- a/src/spikeinterface/sortingcomponents/peak_localization/main.py
+++ b/src/spikeinterface/sortingcomponents/peak_localization/main.py
@@ -20,10 +20,12 @@
 def get_localization_pipeline_nodes(
     recording,
     peak_source,
+    peaks_or_spikes=None,
     method="center_of_mass",
     method_kwargs=None,
     ms_before=0.5,
     ms_after=0.5,
+    channel_index_array=None,
     job_kwargs=None,
     seed=None,
 ):
@@ -44,16 +46,28 @@ def get_localization_pipeline_nodes(
 
     if method == "grid_convolution" and "prototype" not in method_kwargs:
         assert isinstance(peak_source, (PeakRetriever, SpikeRetriever))
+        assert peaks_or_spikes is not None, "peaks_or_spikes must be provided for grid_convolution method."
         # extract prototypes silently
 
         from ..tools import get_prototype_and_waveforms_from_peaks
 
+        if isinstance(peak_source, SpikeRetriever):
+            assert channel_index_array is not None, "For SpikeRetriever, channel_index_array must be provided."
+
         job_kwargs = fix_job_kwargs(job_kwargs)
         job_kwargs["progress_bar"] = False
 
+        # How do you get peaks
+
         method_kwargs = method_kwargs.copy()
         method_kwargs["prototype"], _, _ = get_prototype_and_waveforms_from_peaks(
-            recording, peaks=peak_source.peaks, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs, seed=seed
+            recording,
+            peaks=peaks_or_spikes,
+            ms_before=ms_before,
+            ms_after=ms_after,
+            channel_index_array=channel_index_array,
+            job_kwargs=job_kwargs,
+            seed=seed,
         )
 
     localization_nodes = method_class(recording, parents=[peak_source, extract_dense_waveforms], **method_kwargs)
@@ -141,6 +155,7 @@ def localize_peaks(
     pipeline_nodes = get_localization_pipeline_nodes(
         recording,
         peak_source,
+        peaks_or_spikes=peaks,
         method=method,
         method_kwargs=method_kwargs,
         ms_before=ms_before,
diff --git a/src/spikeinterface/sortingcomponents/peak_selection.py b/src/spikeinterface/sortingcomponents/peak_selection.py
index ddec61c1ae..018e627217 100644
--- a/src/spikeinterface/sortingcomponents/peak_selection.py
+++ b/src/spikeinterface/sortingcomponents/peak_selection.py
@@ -108,6 +108,11 @@ def select_peak_indices(peaks, method, seed, **method_kwargs):
 
     rng = np.random.default_rng(seed=seed)
 
+    # raise error if "channel_index" is not in peaks dtype and select_channel
+    channel_index_array = method_kwargs.get("channel_index_array", None)
+    if channel_index_array is None and "channel_index" in peaks.dtype.names:
+        channel_index_array = peaks["channel_index"]
+
     if method == "uniform":
         params = {"select_per_channel": False, "n_peaks": None}
 
@@ -116,9 +121,11 @@ def select_peak_indices(peaks, method, seed, **method_kwargs):
         assert params["n_peaks"] is not None, "n_peaks should be defined!"
 
         if params["select_per_channel"]:
+            if channel_index_array is None:
+                raise ValueError("Peaks do not have 'channel_index' field and channel_index_array is not provided.")
             ## This method will randomly select max_peaks_per_channel peaks per channels
-            for channel in np.unique(peaks["channel_index"]):
-                peaks_indices = np.where(peaks["channel_index"] == channel)[0]
+            for channel in np.unique(channel_index_array):
+                peaks_indices = np.where(channel_index_array == channel)[0]
                 max_peaks = min(peaks_indices.size, params["n_peaks"])
                 selected_indices += [rng.choice(peaks_indices, size=max_peaks, replace=False)]
         else:
@@ -144,8 +151,10 @@ def select_peak_indices(peaks, method, seed, **method_kwargs):
             assert params["noise_levels"] is not None, "Noise levels should be provided"
 
             if params["select_per_channel"]:
-                for channel in np.unique(peaks["channel_index"]):
-                    peaks_indices = np.where(peaks["channel_index"] == channel)[0]
+                if channel_index_array is None:
+                    raise ValueError("Peaks do not have 'channel_index' field and channel_index_array is not provided.")
+                for channel in np.unique(channel_index_array):
+                    peaks_indices = np.where(channel_index_array == channel)[0]
                     if params["n_peaks"] > peaks_indices.size:
                         selected_indices += [peaks_indices]
                     else:
@@ -170,6 +179,10 @@ def select_peak_indices(peaks, method, seed, **method_kwargs):
                 if params["n_peaks"] > peaks.size:
                     selected_indices += [np.arange(peaks.size)]
                 else:
+                    if channel_index_array is None:
+                        raise ValueError(
+                            "Peaks do not have 'channel_index' field and channel_index_array is not provided."
+                        )
                     snrs = peaks["amplitude"] / params["noise_levels"][peaks["channel_index"]]
                     preprocessing = QuantileTransformer(output_distribution="uniform", n_quantiles=min(100, len(snrs)))
                     snrs = preprocessing.fit_transform(snrs[:, np.newaxis])
diff --git a/src/spikeinterface/sortingcomponents/tools.py b/src/spikeinterface/sortingcomponents/tools.py
index b445513a6c..6a69cea0fc 100644
--- a/src/spikeinterface/sortingcomponents/tools.py
+++ b/src/spikeinterface/sortingcomponents/tools.py
@@ -41,7 +41,9 @@ def make_multi_method_doc(methods, indent="    "):
     return doc
 
 
-def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job_name=None, job_kwargs=None):
+def extract_waveform_at_max_channel(
+    rec, peaks, ms_before=0.5, ms_after=1.5, channel_index_array=None, job_name=None, job_kwargs=None
+):
     """
     Helper function to extract waveforms at the max channel from a peak list
 
@@ -57,8 +59,14 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job
         peaks.size, dtype=[("sample_index", "int64"), ("unit_index", "int64"), ("segment_index", "int64")]
     )
     spikes["sample_index"] = peaks["sample_index"]
-    spikes["unit_index"] = peaks["channel_index"]
     spikes["segment_index"] = peaks["segment_index"]
+    if channel_index_array is not None:
+        assert len(channel_index_array) == peaks.size, "Length of channel_index_array must match number of peaks."
+        spikes["unit_index"] = channel_index_array
+    else:
+        if "channel_index" not in peaks.dtype.names:
+            raise ValueError("The peaks array must have a 'channel_index' field.")
+        spikes["unit_index"] = peaks["channel_index"]
 
     nbefore = ms_to_samples(ms_before, rec.sampling_frequency)
     nafter = ms_to_samples(ms_after, rec.sampling_frequency)
@@ -83,7 +91,7 @@ def extract_waveform_at_max_channel(rec, peaks, ms_before=0.5, ms_after=1.5, job
 
 
 def get_prototype_and_waveforms_from_peaks(
-    recording, peaks, n_peaks=5000, ms_before=0.5, ms_after=0.5, seed=None, job_kwargs=None
+    recording, peaks, n_peaks=5000, ms_before=0.5, ms_after=0.5, channel_index_array=None, seed=None, job_kwargs=None
 ):
     """
     Function to extract a prototype waveform from peaks.
@@ -93,13 +101,16 @@ def get_prototype_and_waveforms_from_peaks(
     recording : Recording
         The recording object containing the data.
     peaks : numpy.array, optional
-        Array of peaks, if None, peaks will be detected, by default None.
+        Array of peaks or spikes, if None, peaks will be detected, by default None.
+        If spikes, the channel_index_array must be provided to indicate the channel of each spike.
     n_peaks : int, optional
         Number of peaks to consider, by default 5000.
     ms_before : float, optional
         Time in milliseconds before the peak to extract the waveform, by default 0.5.
     ms_after : float, optional
         Time in milliseconds after the peak to extract the waveform, by default 0.5.
+    channel_index_array : numpy.array | None, optional
+        Array of channel indices corresponding to each peak, by default None.
     seed : int or None, optional
         Seed for random number generator, by default None.
     job_kwargs : dict
@@ -122,7 +133,13 @@ def get_prototype_and_waveforms_from_peaks(
     nafter = ms_to_samples(ms_after, recording.sampling_frequency)
 
     few_peaks = select_peaks(
-        peaks, recording=recording, method="uniform", n_peaks=n_peaks, margin=(nbefore, nafter), seed=seed
+        peaks,
+        recording=recording,
+        method="uniform",
+        n_peaks=n_peaks,
+        margin=(nbefore, nafter),
+        channel_index_array=channel_index_array,
+        seed=seed,
     )
     waveforms = extract_waveform_at_max_channel(
         recording,
@@ -131,6 +148,7 @@ def get_prototype_and_waveforms_from_peaks(
         ms_after=ms_after,
         job_kwargs=job_kwargs,
         job_name="waveform prototype",
+        channel_index_array=channel_index_array,
     )
 
     with np.errstate(divide="ignore", invalid="ignore"):

From c58a25f2630a895b946f6348ac7d059e8c1c4b4b Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Wed, 7 Oct 2026 13:13:02 +0200
Subject: [PATCH 10/12] fix: tests

---
 .../peak_detection/tests/test_peak_detection.py                 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py b/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py
index a0b0a40521..25b9d9515a 100644
--- a/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py
+++ b/src/spikeinterface/sortingcomponents/peak_detection/tests/test_peak_detection.py
@@ -352,7 +352,7 @@ def test_detect_peaks_matched_filtering(recording, job_kwargs):
     ms_before = 1.0
     ms_after = 1.0
     prototype, _, _ = get_prototype_and_waveforms_from_peaks(
-        recording, peaks_or_spikes=peaks_le, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs
+        recording, peaks=peaks_le, ms_before=ms_before, ms_after=ms_after, job_kwargs=job_kwargs
     )
 
     peaks_mf_neg = detect_peaks(

From 7da203327210b2efb1115c7945c44552236a2e6b Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Wed, 7 Oct 2026 13:27:28 +0200
Subject: [PATCH 11/12] perf: perform search by unique chunks and remove cache

---
 src/spikeinterface/core/zarrextractors.py | 31 ++++++++---------------
 1 file changed, 10 insertions(+), 21 deletions(-)

diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
index 8a5a9b2d8f..04dec55b46 100644
--- a/src/spikeinterface/core/zarrextractors.py
+++ b/src/spikeinterface/core/zarrextractors.py
@@ -463,8 +463,7 @@ class ZarrSampleIndexSearch:
 
     The array is sorted within each segment and stored in fixed-size zarr chunks. The
     search keeps the first value of every chunk (8 B per chunk), picks the one chunk that
-    can hold the answer, and decodes only that chunk. The last decoded chunk is cached, so
-    consecutive searches in the same region decode nothing.
+    can hold the answer, and decodes only that chunk — once per unique chunk per call.
 
     Parameters
     ----------
@@ -484,19 +483,6 @@ def __init__(self, sample_index, chunk_firsts=None):
             starts = np.arange(0, self._num_spikes, self._chunk_length)
             chunk_firsts = sample_index.get_orthogonal_selection(starts) if starts.size else []
         self._chunk_firsts = np.asarray(chunk_firsts, dtype="int64")
-        self._cached_chunk = (-1, None)  # (chunk index, decoded values)
-
-    def _get_chunk(self, chunk_index):
-        # read the (index, values) pair once into a local and replace it as a whole: threads
-        # sharing this object (pool_engine="thread") then always see a consistent pair, and a
-        # race can only cost an extra decode
-        cached = self._cached_chunk
-        if cached[0] != chunk_index:
-            start = chunk_index * self._chunk_length
-            stop = min(start + self._chunk_length, self._num_spikes)
-            cached = (chunk_index, self._sample_index[start:stop])
-            self._cached_chunk = cached
-        return cached[1]
 
     def searchsorted(self, values, start, stop):
         """
@@ -513,15 +499,18 @@ def searchsorted(self, values, start, stop):
         last_chunk = (stop - 1) // self._chunk_length
         # first values of the chunks after `first_chunk`, all of which lie inside start:stop
         later_firsts = self._chunk_firsts[first_chunk + 1 : last_chunk + 1]
-        for i, value in enumerate(values):
-            # the answer lies in the last chunk whose first value is < value: side="left" keeps
-            # runs of equal values that cross a chunk boundary in the earlier chunk
-            chunk_index = first_chunk + int(np.searchsorted(later_firsts, value, side="left"))
+        # assign every value to its chunk in one vectorised call: side="left" keeps runs of
+        # equal values that cross a chunk boundary in the earlier chunk
+        chunk_indices = first_chunk + np.searchsorted(later_firsts, values, side="left")
+        # decode each touched chunk exactly once and answer all values that fall in it
+        unique_chunks, inverse = np.unique(chunk_indices, return_inverse=True)
+        for k, chunk_index in enumerate(unique_chunks):
+            mask = inverse == k
             chunk_start = chunk_index * self._chunk_length
             lo = max(chunk_start, start)
             hi = min(chunk_start + self._chunk_length, stop)
-            block = self._get_chunk(chunk_index)[lo - chunk_start : hi - chunk_start]
-            out[i] = lo + int(np.searchsorted(block, value, side="left")) - start
+            block = self._sample_index[lo:hi]
+            out[mask] = lo + np.searchsorted(block, values[mask], side="left") - start
         return out
 
 

From f284b29804d9b599f6f03e9cf29f0feae34b6792 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Wed, 7 Oct 2026 13:28:47 +0200
Subject: [PATCH 12/12] Apply suggestion from @alejoe91

---
 src/spikeinterface/core/tests/test_zarrextractors.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py
index 8c3a08da0f..f7449811e5 100644
--- a/src/spikeinterface/core/tests/test_zarrextractors.py
+++ b/src/spikeinterface/core/tests/test_zarrextractors.py
@@ -117,7 +117,7 @@ def test_ZarrSortingExtractor_lazy_search(tmp_path):
     lazy = ZarrSortingExtractor(folder, lazy_spike_vector=True)
     assert type(lazy.to_spike_vector()).__name__ == "ZarrSpikeVector"
 
-    rng = np.random.default_rng(1)
+    rng = np.random.default_rng(2308)
     num_samples = int(sorting.to_spike_vector()["sample_index"].max()) + 100
     for segment_index in range(sorting.get_num_segments()):
         frames = np.sort(rng.integers(-10, num_samples, size=200))