Repository: SpikeInterface/spikeinterface

Issue:
Add zarr backend for sorting object
Currently, both `Recording` and `WaveformExtractor` have a zarr backend, but not the `Sorting` object

Discussion:

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 78c32e590..603ab627f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,9 +21,10 @@ classifiers = [
 
 dependencies = [
     "numpy",
-    "neo>=0.12.0",
     "threadpoolctl",
     "tqdm",
+    "zarr",
+    "neo>=0.12.0",
     "probeinterface>=0.2.19",
 ]
 
@@ -83,7 +84,6 @@ streaming_extractors = [
 ]
 
 full = [
-    "zarr",
     "h5py",
     "pandas",
     "xarray",
@@ -112,16 +112,14 @@ qualitymetrics = [
 
 test_core = [
     "pytest",
-    "zarr",
     "psutil",
 ]
 
 test = [
     "pytest",
+    "pytest-dependency",
     "pytest-cov",
 
-    # zarr is needed for testing
-    "zarr",
     "xarray",
     "huggingface_hub",
 
@@ -180,6 +178,7 @@ dev = [
 [tool.pytest.ini_options]
 markers = [
     "core",
+    "generation",
     "extractors",
     "preprocessing",
     "postprocessing",
diff --git a/src/spikeinterface/core/__init__.py b/src/spikeinterface/core/__init__.py
index 9f91c8759..f49088733 100644
--- a/src/spikeinterface/core/__init__.py
+++ b/src/spikeinterface/core/__init__.py
@@ -1,6 +1,6 @@
 from .base import load_extractor  # , load_extractor_from_dict, load_extractor_from_json, load_extractor_from_pickle
 from .baserecording import BaseRecording, BaseRecordingSegment
-from .basesorting import BaseSorting, BaseSortingSegment
+from .basesorting import BaseSorting, BaseSortingSegment, SpikeVectorSortingSegment
 from .baseevent import BaseEvent, BaseEventSegment
 from .basesnippets import BaseSnippets, BaseSnippetsSegment
 from .baserecordingsnippets import BaseRecordingSnippets
@@ -9,7 +9,7 @@ from .baserecordingsnippets import BaseRecordingSnippets
 from .binaryrecordingextractor import BinaryRecordingExtractor, read_binary
 from .npzsortingextractor import NpzSortingExtractor, read_npz_sorting
 from .numpyextractors import NumpyRecording, NumpySorting, SharedMemorySorting, NumpyEvent, NumpySnippets
-from .zarrrecordingextractor import ZarrRecordingExtractor, read_zarr, get_default_zarr_compressor
+from .zarrextractors import ZarrRecordingExtractor, ZarrSortingExtractor, read_zarr, get_default_zarr_compressor
 from .binaryfolder import BinaryFolderRecording, read_binary_folder
 from .sortingfolder import NumpyFolderSorting, NpzFolderSorting, read_numpy_sorting_folder, read_npz_folder
 from .npysnippetsextractor import NpySnippetsExtractor, read_npy_snippets
diff --git a/src/spikeinterface/core/base.py b/src/spikeinterface/core/base.py
index ed2bb91df..5e8b8f9f1 100644
--- a/src/spikeinterface/core/base.py
+++ b/src/spikeinterface/core/base.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 from pathlib import Path
-import re
+import shutil
 from typing import Any, Iterable, List, Optional, Sequence, Union
 import importlib
 import warnings
@@ -761,7 +761,7 @@ class BaseExtractor:
             file = None
 
             if folder.suffix == ".zarr":
-                from .zarrrecordingextractor import read_zarr
+                from .zarrextractors import read_zarr
 
                 extractor = read_zarr(folder)
             else:
@@ -942,52 +942,41 @@ class BaseExtractor:
         self,
         name=None,
         folder=None,
+        overwrite=False,
         storage_options=None,
         channel_chunk_size=None,
         verbose=True,
-        zarr_path=None,
         **save_kwargs,
     ):
         """
         Save extractor to zarr.
 
-        The save consist of:
-            * extracting traces by calling get_trace() method in chunks
-            * saving data into a zarr file
-            * dumping the original extractor for provenance in attributes
-
         Parameters
         ----------
         name: str or None, default: None
             Name of the subfolder in get_global_tmp_folder()
-            If "name" is given, "folder" must be None.
+            If "name" is given, "folder" must be None
         folder: str, Path, or None, default: None
             The folder used to save the zarr output. If the folder does not have a ".zarr" suffix,
-            it will be automatically appended.
+            it will be automatically appended
+        overwrite: bool, default: False
+            If True, the folder is removed if it already exists
         storage_options: dict or None, default: None
             Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
             For cloud storage locations, this should not be None (in case of default values, use an empty dict)
         channel_chunk_size: int or None, default: None
-            Channels per chunk
+            Channels per chunk (only for BaseRecording)
         verbose: bool, default: True
             If True, the output is verbose
-        zarr_path: str, Path, or None, default: None
-            (Deprecated) Name of the zarr folder (.zarr)
-        **save_kwargs: Keyword arguments for saving.
+        **save_kwargs: Keyword arguments for saving to zarr
 
         Returns
         -------
-        cached: ZarrRecordingExtractor
+        cached: ZarrExtractor
             Saved copy of the extractor.
         """
         import zarr
-        from .zarrrecordingextractor import read_zarr
-
-        if zarr_path is not None:
-            warnings.warn(
-                "The 'zarr_path' argument is deprecated. " "Use 'folder' instead", DeprecationWarning, stacklevel=2
-            )
-            folder = zarr_path
+        from .zarrextractors import read_zarr
 
         if folder is None:
             cache_folder = get_global_tmp_folder()
@@ -1006,38 +995,18 @@ class BaseExtractor:
                 folder = Path(folder)
                 if folder.suffix != ".zarr":
                     folder = folder.parent / f"{folder.stem}.zarr"
+                if folder.is_dir() and overwrite:
+                    shutil.rmtree(folder)
                 zarr_path = folder
-                zarr_path_init = str(zarr_path)
             else:
                 zarr_path = folder
-                zarr_path_init = zarr_path
 
         if isinstance(zarr_path, Path):
             assert not zarr_path.exists(), f"Path {zarr_path} already exists, choose another name"
-
-        zarr_root = zarr.open(zarr_path_init, mode="w", storage_options=storage_options)
-
-        if self.check_if_json_serializable():
-            zarr_root.attrs["provenance"] = check_json(self.to_dict())
-        else:
-            zarr_root.attrs["provenance"] = None
-
-        # save data (done the subclass)
-        save_kwargs["zarr_root"] = zarr_root
         save_kwargs["zarr_path"] = zarr_path
         save_kwargs["storage_options"] = storage_options
         save_kwargs["channel_chunk_size"] = channel_chunk_size
         cached = self._save(verbose=verbose, **save_kwargs)
-
-        # save properties
-        prop_group = zarr_root.create_group("properties")
-        for key in self.get_property_keys():
-            values = self.get_property(key)
-            prop_group.create_dataset(name=key, data=values, compressor=None)
-
-        # save annotations
-        zarr_root.attrs["annotations"] = check_json(self._annotations)
-
         cached = read_zarr(zarr_path)
 
         return cached
diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py
index c112447c3..802578426 100644
--- a/src/spikeinterface/core/baserecording.py
+++ b/src/spikeinterface/core/baserecording.py
@@ -9,12 +9,10 @@ from probeinterface import Probe, ProbeGroup, read_probeinterface, select_axes,
 from .base import BaseSegment
 from .baserecordingsnippets import BaseRecordingSnippets
 from .core_tools import (
-    check_json,
     convert_bytes_to_str,
     convert_seconds_to_str,
     write_binary_recording,
     write_memory_recording,
-    write_traces_to_zarr,
 )
 from .job_tools import split_job_kwargs
 
@@ -497,51 +495,12 @@ class BaseRecording(BaseRecordingSnippets):
             )
 
         elif format == "zarr":
-            from .zarrrecordingextractor import get_default_zarr_compressor, ZarrRecordingExtractor
+            from .zarrextractors import ZarrRecordingExtractor
 
-            zarr_kwargs = kwargs.copy()
-
-            zarr_root = zarr_kwargs["zarr_root"]
-            zarr_root.attrs["sampling_frequency"] = float(self.get_sampling_frequency())
-            zarr_root.attrs["num_segments"] = int(self.get_num_segments())
-            zarr_root.create_dataset(name="channel_ids", data=self.get_channel_ids(), compressor=None)
-
-            zarr_kwargs["dataset_paths"] = [f"traces_seg{i}" for i in range(self.get_num_segments())]
-            zarr_kwargs["dtype"] = kwargs.get("dtype", None) or self.get_dtype()
-
-            if "compressor" not in zarr_kwargs:
-                zarr_kwargs["compressor"] = compressor = get_default_zarr_compressor()
-                print(
-                    f"Using default zarr compressor: {compressor}. To use a different compressor, use the "
-                    f"'compressor' argument"
-                )
-
-            write_traces_to_zarr(self, **zarr_kwargs, **job_kwargs)
-
-            # save probe
-            if self.get_property("contact_vector") is not None:
-                probegroup = self.get_probegroup()
-                zarr_root.attrs["probe"] = check_json(probegroup.to_dict(array_as_list=True))
-
-            # save time vector if any
-            t_starts = np.zeros(self.get_num_segments(), dtype="float64") * np.nan
-            for segment_index, rs in enumerate(self._recording_segments):
-                d = rs.get_times_kwargs()
-                time_vector = d["time_vector"]
-                if time_vector is not None:
-                    _ = zarr_root.create_dataset(
-                        name=f"times_seg{segment_index}",
-                        data=time_vector,
-                        filters=zarr_kwargs.get("filters", None),
-                        compressor=zarr_kwargs["compressor"],
-                    )
-                elif d["t_start"] is not None:
-                    t_starts[segment_index] = d["t_start"]
-
-            if np.any(~np.isnan(t_starts)):
-                zarr_root.create_dataset(name="t_starts", data=t_starts, compressor=None)
-
-            cached = ZarrRecordingExtractor(zarr_kwargs["zarr_path"], zarr_kwargs["storage_options"])
+            zarr_path = kwargs.pop("zarr_path")
+            storage_options = kwargs.pop("storage_options")
+            ZarrRecordingExtractor.write_recording(self, zarr_path, storage_options, **kwargs)
+            cached = ZarrRecordingExtractor(zarr_path, storage_options)
 
         elif format == "nwb":
             # TODO implement a format based on zarr
diff --git a/src/spikeinterface/core/basesorting.py b/src/spikeinterface/core/basesorting.py
index dda825c59..465807cc9 100644
--- a/src/spikeinterface/core/basesorting.py
+++ b/src/spikeinterface/core/basesorting.py
@@ -240,6 +240,18 @@ class BaseSorting(BaseExtractor):
                 warnings.warn("The registered recording will not be persistent on disk, but only available in memory")
                 cached.register_recording(self._recording)
 
+        elif format == "zarr":
+            from .zarrextractors import ZarrSortingExtractor
+
+            zarr_path = save_kwargs.pop("zarr_path")
+            storage_options = save_kwargs.pop("storage_options")
+            ZarrSortingExtractor.write_sorting(self, zarr_path, storage_options, **save_kwargs)
+            cached = ZarrSortingExtractor(zarr_path, storage_options)
+
+            if self.has_recording():
+                warnings.warn("The registered recording will not be persistent on disk, but only available in memory")
+                cached.register_recording(self._recording)
+
         elif format == "npz_folder":
             from .sortingfolder import NpzFolderSorting
 
@@ -496,7 +508,9 @@ class BaseSorting(BaseExtractor):
                 for unit_id in unit_ids:
                     self.get_unit_spike_train(unit_id, segment_index=segment_index, use_cache=True)
 
-    def to_spike_vector(self, concatenated=True, extremum_channel_inds=None, use_cache=True):
+    def to_spike_vector(
+        self, concatenated=True, extremum_channel_inds=None, use_cache=True
+    ) -> np.ndarray | list[np.ndarray]:
         """
         Construct a unique structured numpy vector concatenating all spikes
         with several fields: sample_index, unit_index, segment_index.
@@ -678,3 +692,35 @@ class BaseSortingSegment(BaseSegment):
         """
         # must be implemented in subclass
         raise NotImplementedError
+
+
+class SpikeVectorSortingSegment(BaseSortingSegment):
+    """
+    A sorting segment that stores spike times as a spike vector.
+    """
+
+    def __init__(self, spikes, segment_index, unit_ids):
+        BaseSortingSegment.__init__(self)
+        self.spikes = spikes
+        self.segment_index = segment_index
+        self.unit_ids = list(unit_ids)
+        self.spikes_in_seg = None
+
+    def get_unit_spike_train(self, unit_id, start_frame, end_frame):
+        if self.spikes_in_seg is None:
+            # the slicing of segment is done only once the first time
+            # this fasten the constructor a lot
+            s0, s1 = np.searchsorted(self.spikes["segment_index"], [self.segment_index, self.segment_index + 1])
+            self.spikes_in_seg = self.spikes[s0:s1]
+
+        start = 0 if start_frame is None else np.searchsorted(self.spikes_in_seg["sample_index"], start_frame)
+        end = (
+            len(self.spikes_in_seg)
+            if end_frame is None
+            else np.searchsorted(self.spikes_in_seg["sample_index"], end_frame)
+        )
+
+        unit_index = self.unit_ids.index(unit_id)
+        times = self.spikes_in_seg[start:end][self.spikes_in_seg[start:end]["unit_index"] == unit_index]["sample_index"]
+
+        return times
diff --git a/src/spikeinterface/core/core_tools.py b/src/spikeinterface/core/core_tools.py
index 5892c62b6..5d30d9448 100644
--- a/src/spikeinterface/core/core_tools.py
+++ b/src/spikeinterface/core/core_tools.py
@@ -5,9 +5,7 @@ import sys
 import datetime
 import json
 from copy import deepcopy
-import gc
 import mmap
-import inspect
 
 import numpy as np
 from tqdm import tqdm
@@ -623,147 +621,6 @@ def write_to_h5_dataset_format(
     return save_path
 
 
-def write_traces_to_zarr(
-    recording,
-    zarr_root,
-    zarr_path,
-    storage_options,
-    dataset_paths,
-    channel_chunk_size=None,
-    dtype=None,
-    compressor=None,
-    filters=None,
-    verbose=False,
-    auto_cast_uint=True,
-    **job_kwargs,
-):
-    """
-    Save the trace of a recording extractor in several zarr format.
-
-
-    Parameters
-    ----------
-    recording: RecordingExtractor
-        The recording extractor object to be saved in .dat format
-    zarr_root: zarr.Group
-        The zarr root
-    zarr_path: str or Path
-        The path to the zarr file
-    storage_options: dict or None
-        Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
-    dataset_paths: list
-        List of paths to traces datasets in the zarr group
-    channel_chunk_size: int or None, default: None (chunking in time only)
-        Channels per chunk
-    dtype: dtype, default: None
-        Type of the saved data
-    compressor: zarr compressor or None, default: None
-        Zarr compressor
-    filters: list, default: None
-        List of zarr filters
-    verbose: bool, default: False
-        If True, output is verbose (when chunks are used)
-    auto_cast_uint: bool, default: True
-        If True, unsigned integers are automatically cast to int if the specified dtype is signed
-    {}
-    """
-    assert dataset_paths is not None, "Provide 'file_path'"
-
-    if not isinstance(dataset_paths, list):
-        dataset_paths = [dataset_paths]
-    assert len(dataset_paths) == recording.get_num_segments()
-
-    if dtype is None:
-        dtype = recording.get_dtype()
-    if auto_cast_uint:
-        cast_unsigned = determine_cast_unsigned(recording, dtype)
-    else:
-        cast_unsigned = False
-
-    job_kwargs = fix_job_kwargs(job_kwargs)
-    chunk_size = ensure_chunk_size(recording, **job_kwargs)
-
-    # create zarr datasets files
-    for segment_index in range(recording.get_num_segments()):
-        num_frames = recording.get_num_samples(segment_index)
-        num_channels = recording.get_num_channels()
-        dset_name = dataset_paths[segment_index]
-        shape = (num_frames, num_channels)
-        _ = zarr_root.create_dataset(
-            name=dset_name,
-            shape=shape,
-            chunks=(chunk_size, channel_chunk_size),
-            dtype=dtype,
-            filters=filters,
-            compressor=compressor,
-        )
-        # synchronizer=zarr.ThreadSynchronizer())
-
-    # use executor (loop or workers)
-    func = _write_zarr_chunk
-    init_func = _init_zarr_worker
-    init_args = (recording, zarr_path, storage_options, dataset_paths, dtype, cast_unsigned)
-    executor = ChunkRecordingExecutor(
-        recording, func, init_func, init_args, verbose=verbose, job_name="write_zarr_recording", **job_kwargs
-    )
-    executor.run()
-
-
-# used by write_zarr_recording + ChunkRecordingExecutor
-def _init_zarr_worker(recording, zarr_path, storage_options, dataset_paths, dtype, cast_unsigned):
-    import zarr
-
-    # create a local dict per worker
-    worker_ctx = {}
-    if isinstance(recording, dict):
-        from spikeinterface.core import load_extractor
-
-        worker_ctx["recording"] = load_extractor(recording)
-    else:
-        worker_ctx["recording"] = recording
-
-    # reload root and datasets
-    if storage_options is None:
-        if isinstance(zarr_path, str):
-            zarr_path_init = zarr_path
-            zarr_path = Path(zarr_path)
-        else:
-            zarr_path_init = str(zarr_path)
-    else:
-        zarr_path_init = zarr_path
-
-    root = zarr.open(zarr_path_init, mode="r+", storage_options=storage_options)
-    zarr_datasets = []
-    for dset_name in dataset_paths:
-        z = root[dset_name]
-        zarr_datasets.append(z)
-    worker_ctx["zarr_datasets"] = zarr_datasets
-    worker_ctx["dtype"] = np.dtype(dtype)
-    worker_ctx["cast_unsigned"] = cast_unsigned
-
-    return worker_ctx
-
-
-# used by write_zarr_recording + ChunkRecordingExecutor
-def _write_zarr_chunk(segment_index, start_frame, end_frame, worker_ctx):
-    # recover variables of the worker
-    recording = worker_ctx["recording"]
-    dtype = worker_ctx["dtype"]
-    zarr_dataset = worker_ctx["zarr_datasets"][segment_index]
-    cast_unsigned = worker_ctx["cast_unsigned"]
-
-    # apply function
-    traces = recording.get_traces(
-        start_frame=start_frame, end_frame=end_frame, segment_index=segment_index, cast_unsigned=cast_unsigned
-    )
-    traces = traces.astype(dtype, copy=False)
-    zarr_dataset[start_frame:end_frame, :] = traces
-
-    # fix memory leak by forcing garbage collection
-    del traces
-    gc.collect()
-
-
 def determine_cast_unsigned(recording, dtype):
     recording_dtype = np.dtype(recording.get_dtype())
 
diff --git a/src/spikeinterface/core/numpyextractors.py b/src/spikeinterface/core/numpyextractors.py
index 1a736206c..88db04251 100644
--- a/src/spikeinterface/core/numpyextractors.py
+++ b/src/spikeinterface/core/numpyextractors.py
@@ -5,7 +5,7 @@ from spikeinterface.core import (
     BaseRecording,
     BaseSorting,
     BaseRecordingSegment,
-    BaseSortingSegment,
+    SpikeVectorSortingSegment,
     BaseEvent,
     BaseEventSegment,
     BaseSnippets,
@@ -140,7 +140,7 @@ class NumpySorting(BaseSorting):
             nseg = spikes[-1]["segment_index"] + 1
 
         for segment_index in range(nseg):
-            self.add_sorting_segment(NumpySortingSegment(spikes, segment_index, unit_ids))
+            self.add_sorting_segment(SpikeVectorSortingSegment(spikes, segment_index, unit_ids))
 
         # important trick : the cache is already spikes vector
         self._cached_spike_vector = spikes
@@ -332,34 +332,6 @@ class NumpySorting(BaseSorting):
         return sorting
 
 
-class NumpySortingSegment(BaseSortingSegment):
-    def __init__(self, spikes, segment_index, unit_ids):
-        BaseSortingSegment.__init__(self)
-        self.spikes = spikes
-        self.segment_index = segment_index
-        self.unit_ids = list(unit_ids)
-        self.spikes_in_seg = None
-
-    def get_unit_spike_train(self, unit_id, start_frame, end_frame):
-        if self.spikes_in_seg is None:
-            # the slicing of segment is done only once the first time
-            # this fasten the constructor a lot
-            s0, s1 = np.searchsorted(self.spikes["segment_index"], [self.segment_index, self.segment_index + 1])
-            self.spikes_in_seg = self.spikes[s0:s1]
-
-        start = 0 if start_frame is None else np.searchsorted(self.spikes_in_seg["sample_index"], start_frame)
-        end = (
-            len(self.spikes_in_seg)
-            if end_frame is None
-            else np.searchsorted(self.spikes_in_seg["sample_index"], end_frame)
-        )
-
-        unit_index = self.unit_ids.index(unit_id)
-        times = self.spikes_in_seg[start:end][self.spikes_in_seg[start:end]["unit_index"] == unit_index]["sample_index"]
-
-        return times
-
-
 class SharedMemorySorting(BaseSorting):
     def __init__(self, shm_name, shape, sampling_frequency, unit_ids, dtype=minimum_spike_dtype, main_shm_owner=True):
         assert len(shape) == 1
@@ -376,7 +348,7 @@ class SharedMemorySorting(BaseSorting):
 
         nseg = self.shm_spikes[-1]["segment_index"] + 1
         for segment_index in range(nseg):
-            self.add_sorting_segment(NumpySortingSegment(self.shm_spikes, segment_index, unit_ids))
+            self.add_sorting_segment(SpikeVectorSortingSegment(self.shm_spikes, segment_index, unit_ids))
 
         # important trick : the cache is already spikes vector
         self._cached_spike_vector = self.shm_spikes
diff --git a/src/spikeinterface/core/sortingfolder.py b/src/spikeinterface/core/sortingfolder.py
index 40614d468..08d38c451 100644
--- a/src/spikeinterface/core/sortingfolder.py
+++ b/src/spikeinterface/core/sortingfolder.py
@@ -3,10 +3,9 @@ import json
 
 import numpy as np
 
-from .basesorting import BaseSorting, BaseSortingSegment
+from .basesorting import BaseSorting, SpikeVectorSortingSegment
 from .npzsortingextractor import NpzSortingExtractor
 from .core_tools import define_function_from_class, make_paths_absolute
-from .numpyextractors import NumpySortingSegment
 
 
 class NumpyFolderSorting(BaseSorting):
@@ -41,7 +40,7 @@ class NumpyFolderSorting(BaseSorting):
         self.spikes = np.load(folder_path / "spikes.npy", mmap_mode="r")
 
         for segment_index in range(num_segments):
-            self.add_sorting_segment(NumpySortingSegment(self.spikes, segment_index, unit_ids))
+            self.add_sorting_segment(SpikeVectorSortingSegment(self.spikes, segment_index, unit_ids))
 
         # important trick : the cache is already spikes vector
         self._cached_spike_vector = self.spikes
diff --git a/src/spikeinterface/core/waveform_extractor.py b/src/spikeinterface/core/waveform_extractor.py
index fb6670907..cbf1ab9b4 100644
--- a/src/spikeinterface/core/waveform_extractor.py
+++ b/src/spikeinterface/core/waveform_extractor.py
@@ -969,7 +969,7 @@ class WaveformExtractor:
                     np.save(waveform_folder / f"sampled_index_{unit_id}.npy", sampled_indices)
         elif format == "zarr":
             import zarr
-            from .zarrrecordingextractor import get_default_zarr_compressor
+            from .zarrextractors import get_default_zarr_compressor
 
             if folder.suffix != ".zarr":
                 folder = folder.parent / f"{folder.stem}.zarr"
@@ -2044,7 +2044,7 @@ class BaseWaveformExtractorExtension:
                     except:
                         raise Exception(f"Could not save {ext_data_name} as extension data")
         elif self.format == "zarr":
-            from .zarrrecordingextractor import get_default_zarr_compressor
+            from .zarrextractors import get_default_zarr_compressor
             import pandas as pd
             import numcodecs
 
diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
new file mode 100644
index 000000000..c9540ffe4
--- /dev/null
+++ b/src/spikeinterface/core/zarrextractors.py
@@ -0,0 +1,534 @@
+from __future__ import annotations
+
+from pathlib import Path
+import numpy as np
+import zarr
+
+from probeinterface import ProbeGroup
+
+from .baserecording import BaseRecording, BaseRecordingSegment
+from .basesorting import BaseSorting, SpikeVectorSortingSegment, minimum_spike_dtype
+from .core_tools import define_function_from_class, check_json
+from .job_tools import split_job_kwargs
+
+
+class ZarrRecordingExtractor(BaseRecording):
+    """
+    RecordingExtractor for a zarr format
+
+    Parameters
+    ----------
+    folder_path: str or Path
+        Path to the zarr root folder
+    storage_options: dict or None
+        Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
+
+    Returns
+    -------
+    recording: ZarrRecordingExtractor
+        The recording Extractor
+    """
+
+    extractor_name = "ZarrRecording"
+    installed = True
+    mode = "folder"
+    installation_mesg = ""
+    name = "zarr"
+
+    def __init__(self, folder_path: Path | str, storage_options: dict | None = None):
+        assert self.installed, self.installation_mesg
+
+        folder_path, folder_path_kwarg = resolve_zarr_path(folder_path)
+
+        self._root = zarr.open(str(folder_path), mode="r", storage_options=storage_options)
+
+        sampling_frequency = self._root.attrs.get("sampling_frequency", None)
+        num_segments = self._root.attrs.get("num_segments", None)
+        assert "channel_ids" in self._root.keys(), "'channel_ids' dataset not found!"
+        channel_ids = self._root["channel_ids"][:]
+
+        assert sampling_frequency is not None, "'sampling_frequency' attiribute not found!"
+        assert num_segments is not None, "'num_segments' attiribute not found!"
+
+        channel_ids = np.array(channel_ids)
+
+        dtype = self._root["traces_seg0"].dtype
+
+        BaseRecording.__init__(self, sampling_frequency, channel_ids, dtype)
+
+        dtype = np.dtype(dtype)
+        t_starts = self._root.get("t_starts", None)
+
+        total_nbytes = 0
+        total_nbytes_stored = 0
+        cr_by_segment = {}
+        for segment_index in range(num_segments):
+            trace_name = f"traces_seg{segment_index}"
+            assert (
+                len(channel_ids) == self._root[trace_name].shape[1]
+            ), f"Segment {segment_index} has the wrong number of channels!"
+
+            time_kwargs = {}
+            time_vector = self._root.get(f"times_seg{segment_index}", None)
+            if time_vector is not None:
+                time_kwargs["time_vector"] = time_vector
+            else:
+                if t_starts is None:
+                    t_start = None
+                else:
+                    t_start = t_starts[segment_index]
+                    if np.isnan(t_start):
+                        t_start = None
+                time_kwargs["t_start"] = t_start
+                time_kwargs["sampling_frequency"] = sampling_frequency
+
+            rec_segment = ZarrRecordingSegment(self._root, trace_name, **time_kwargs)
+
+            nbytes_segment = self._root[trace_name].nbytes
+            nbytes_stored_segment = self._root[trace_name].nbytes_stored
+            cr_by_segment[segment_index] = nbytes_segment / nbytes_stored_segment
+
+            total_nbytes += nbytes_segment
+            total_nbytes_stored += nbytes_stored_segment
+            self.add_recording_segment(rec_segment)
+
+        # load probe
+        probe_dict = self._root.attrs.get("probe", None)
+        if probe_dict is not None:
+            probegroup = ProbeGroup.from_dict(probe_dict)
+            self.set_probegroup(probegroup, in_place=True)
+
+        # load properties
+        if "properties" in self._root:
+            prop_group = self._root["properties"]
+            for key in prop_group.keys():
+                values = self._root["properties"][key]
+                self.set_property(key, values)
+
+        # load annotations
+        annotations = self._root.attrs.get("annotations", None)
+        if annotations is not None:
+            self.annotate(**annotations)
+        # annotate compression ratios
+        cr = total_nbytes / total_nbytes_stored
+        self.annotate(compression_ratio=cr, compression_ratio_segments=cr_by_segment)
+
+        self._kwargs = {"folder_path": folder_path_kwarg, "storage_options": storage_options}
+
+    @staticmethod
+    def write_recording(
+        recording: BaseRecording, folder_path: str | Path, storage_options: dict | None = None, **kwargs
+    ):
+        zarr_root = zarr.open(str(folder_path), mode="w", storage_options=storage_options)
+        add_recording_to_zarr_group(recording, zarr_root, **kwargs)
+
+
+class ZarrRecordingSegment(BaseRecordingSegment):
+    def __init__(self, root, dataset_name, **time_kwargs):
+        BaseRecordingSegment.__init__(self, **time_kwargs)
+        self._timeseries = root[dataset_name]
+
+    def get_num_samples(self) -> int:
+        """Returns the number of samples in this signal block
+
+        Returns:
+            SampleIndex: Number of samples in the signal block
+        """
+        return self._timeseries.shape[0]
+
+    def get_traces(
+        self,
+        start_frame: int | None = None,
+        end_frame: int | None = None,
+        channel_indices: list[int | str] | None = None,
+    ) -> np.ndarray:
+        traces = self._timeseries[start_frame:end_frame]
+        if channel_indices is not None:
+            traces = traces[:, channel_indices]
+        return traces
+
+
+class ZarrSortingExtractor(BaseSorting):
+    """
+    SortingExtractor for a zarr format
+
+    Parameters
+    ----------
+    root_path: str or Path
+        Path to the zarr root file
+    storage_options: dict or None
+        Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
+
+    Returns
+    -------
+    sorting: ZarrSortingExtractor
+        The sorting Extractor
+    """
+
+    extractor_name = "ZarrSorting"
+    installed = True
+    mode = "folder"
+    installation_mesg = ""
+    name = "zarr"
+
+    def __init__(self, folder_path: Path | str, storage_options: dict | None = None):
+        assert self.installed, self.installation_mesg
+
+        folder_path, folder_path_kwarg = resolve_zarr_path(folder_path)
+
+        self._root = zarr.open(str(folder_path), mode="r", storage_options=storage_options)
+
+        sampling_frequency = self._root.attrs.get("sampling_frequency", None)
+        num_segments = self._root.attrs.get("num_segments", None)
+        assert "unit_ids" in self._root.keys(), "'unit_ids' dataset not found!"
+        unit_ids = self._root["unit_ids"][:]
+
+        assert sampling_frequency is not None, "'sampling_frequency' attiribute not found!"
+        assert num_segments is not None, "'num_segments' attiribute not found!"
+
+        unit_ids = np.array(unit_ids)
+        assert "spikes" in self._root.keys(), "'spikes' dataset not found!"
+        spikes_group = self._root["spikes"]
+        segment_slices_list = spikes_group["segment_slices"][:]
+
+        BaseSorting.__init__(self, sampling_frequency, unit_ids)
+
+        spikes = np.zeros(len(spikes_group["sample_index"]), dtype=minimum_spike_dtype)
+        spikes["sample_index"] = spikes_group["sample_index"][:]
+        spikes["unit_index"] = spikes_group["unit_index"][:]
+        for i, (start, end) in enumerate(segment_slices_list):
+            spikes["segment_index"][start:end] = i
+
+        for segment_index in range(num_segments):
+            soring_segment = SpikeVectorSortingSegment(spikes, segment_index, unit_ids)
+            self.add_sorting_segment(soring_segment)
+
+        # load properties
+        if "properties" in self._root:
+            prop_group = self._root["properties"]
+            for key in prop_group.keys():
+                values = self._root["properties"][key]
+                self.set_property(key, values)
+
+        # load annotations
+        annotations = self._root.attrs.get("annotations", None)
+        if annotations is not None:
+            self.annotate(**annotations)
+
+        self._kwargs = {"root_path": folder_path_kwarg, "storage_options": storage_options}
+
+    @staticmethod
+    def write_sorting(sorting: BaseSorting, folder_path: str | Path, storage_options: dict | None = None, **kwargs):
+        """
+        Write a sorting extractor to zarr format.
+        """
+        zarr_root = zarr.open(str(folder_path), mode="w", storage_options=storage_options)
+        add_sorting_to_zarr_group(sorting, zarr_root, **kwargs)
+
+
+read_zarr_recording = define_function_from_class(source_class=ZarrRecordingExtractor, name="read_zarr_recording")
+read_zarr_sorting = define_function_from_class(source_class=ZarrSortingExtractor, name="read_zarr_sorting")
+
+
+def read_zarr(
+    folder_path: str | Path, storage_options: dict | None = None
+) -> ZarrRecordingExtractor | ZarrSortingExtractor:
+    """
+    Read recording or sorting from a zarr format
+
+    Parameters
+    ----------
+    folder_path: str or Path
+        Path to the zarr root file
+    storage_options: dict or None
+        Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
+
+    Returns
+    -------
+    extractor: ZarrExtractor
+        The loaded extractor
+    """
+    root = zarr.open(str(folder_path), mode="r", storage_options=storage_options)
+    if "channel_ids" in root.keys():
+        return read_zarr_recording(folder_path, storage_options=storage_options)
+    elif "unit_ids" in root.keys():
+        return read_zarr_sorting(folder_path, storage_options=storage_options)
+    else:
+        raise ValueError("Cannot find 'channel_ids' or 'unit_ids' in zarr root. Not a valid SpikeInterface zarr format")
+
+
+### UTILITY FUNCTIONS ###
+def resolve_zarr_path(folder_path: str | Path):
+    """
+    Resolve a path to a zarr folder.
+
+    Parameters
+    ----------
+    """
+    if str(folder_path).startswith("s3:") or str(folder_path).startswith("gcs:"):
+        # cloud location, no need to resolve
+        return folder_path, folder_path
+    else:
+        folder_path = Path(folder_path)
+        folder_path_kwarg = str(Path(folder_path).resolve())
+        return folder_path, folder_path_kwarg
+
+
+def get_default_zarr_compressor(clevel: int = 5):
+    """
+    Return default Zarr compressor object for good preformance in int16
+    electrophysiology data.
+
+    cname: zstd (zstandard)
+    clevel: 5
+    shuffle: BITSHUFFLE
+
+    Parameters
+    ----------
+    clevel : int, default: 5
+        Compression level (higher -> more compressed).
+        Minimum 1, maximum 9. By default 5
+
+    Returns
+    -------
+    Blosc.compressor
+        The compressor object that can be used with the save to zarr function
+    """
+    assert ZarrRecordingExtractor.installed, ZarrRecordingExtractor.installation_mesg
+    from numcodecs import Blosc
+
+    return Blosc(cname="zstd", clevel=clevel, shuffle=Blosc.BITSHUFFLE)
+
+
+def add_properties_and_annotations(
+    zarr_group: zarr.hierarchy.Group, recording_or_sorting: Union[BaseRecording, BaseSorting]
+):
+    # save properties
+    prop_group = zarr_group.create_group("properties")
+    for key in recording_or_sorting.get_property_keys():
+        values = recording_or_sorting.get_property(key)
+        prop_group.create_dataset(name=key, data=values, compressor=None)
+
+    # save annotations
+    zarr_group.attrs["annotations"] = check_json(recording_or_sorting._annotations)
+
+
+def add_sorting_to_zarr_group(sorting: BaseSorting, zarr_group: zarr.hierarchy.Group, **kwargs):
+    """
+    Add a sorting extractor to a zarr group.
+
+    Parameters
+    ----------
+    sorting: BaseSorting
+        The sorting extractor object to be added to the zarr group
+    zarr_group: zarr.hierarchy.Group
+        The zarr group
+    kwargs: dict
+        Other arguments passed to the zarr compressor
+    """
+    from numcodecs import Delta
+
+    num_segments = sorting.get_num_segments()
+    zarr_group.attrs["sampling_frequency"] = float(sorting.sampling_frequency)
+    zarr_group.attrs["num_segments"] = int(num_segments)
+    zarr_group.create_dataset(name="unit_ids", data=sorting.unit_ids, compressor=None)
+
+    if "compressor" not in kwargs:
+        compressor = get_default_zarr_compressor()
+
+    # save sub fields
+    spikes_group = zarr_group.create_group(name="spikes")
+    spikes = sorting.to_spike_vector()
+    for field in spikes.dtype.fields:
+        if field != "segment_index":
+            spikes_group.create_dataset(
+                name=field,
+                data=spikes[field],
+                compressor=compressor,
+                filters=[Delta(dtype=spikes[field].dtype)],
+            )
+        else:
+            segment_slices = []
+            for segment_index in range(num_segments):
+                i0, i1 = np.searchsorted(spikes["segment_index"], [segment_index, segment_index + 1])
+                segment_slices.append([i0, i1])
+            spikes_group.create_dataset(name="segment_slices", data=segment_slices, compressor=None)
+
+    add_properties_and_annotations(zarr_group, sorting)
+
+
+# Recording
+def add_recording_to_zarr_group(recording: BaseRecording, zarr_group: zarr.hierarchy.Group, **kwargs):
+    zarr_kwargs, job_kwargs = split_job_kwargs(kwargs)
+
+    if recording.check_if_json_serializable():
+        zarr_group.attrs["provenance"] = check_json(recording.to_dict(recursive=True))
+    else:
+        zarr_group.attrs["provenance"] = None
+
+    # save data (done the subclass)
+    zarr_group.attrs["sampling_frequency"] = float(recording.get_sampling_frequency())
+    zarr_group.attrs["num_segments"] = int(recording.get_num_segments())
+    zarr_group.create_dataset(name="channel_ids", data=recording.get_channel_ids(), compressor=None)
+    dataset_paths = [f"traces_seg{i}" for i in range(recording.get_num_segments())]
+
+    zarr_kwargs["dtype"] = kwargs.get("dtype", None) or recording.get_dtype()
+    if "compressor" not in zarr_group:
+        zarr_kwargs["compressor"] = get_default_zarr_compressor()
+
+    add_traces_to_zarr(
+        recording=recording,
+        zarr_group=zarr_group,
+        dataset_paths=dataset_paths,
+        **zarr_kwargs,
+        **job_kwargs,
+    )
+
+    # # save probe
+    if recording.get_property("contact_vector") is not None:
+        probegroup = recording.get_probegroup()
+        zarr_group.attrs["probe"] = check_json(probegroup.to_dict(array_as_list=True))
+
+    # save time vector if any
+    t_starts = np.zeros(recording.get_num_segments(), dtype="float64") * np.nan
+    for segment_index, rs in enumerate(recording._recording_segments):
+        d = rs.get_times_kwargs()
+        time_vector = d["time_vector"]
+        if time_vector is not None:
+            _ = zarr_group.create_dataset(
+                name=f"times_seg{segment_index}",
+                data=time_vector,
+                filters=zarr_kwargs.get("filters", None),
+                compressor=zarr_kwargs["compressor"],
+            )
+        elif d["t_start"] is not None:
+            t_starts[segment_index] = d["t_start"]
+
+    if np.any(~np.isnan(t_starts)):
+        zarr_group.create_dataset(name="t_starts", data=t_starts, compressor=None)
+
+    add_properties_and_annotations(zarr_group, recording)
+
+
+def add_traces_to_zarr(
+    recording,
+    zarr_group,
+    dataset_paths,
+    channel_chunk_size=None,
+    dtype=None,
+    compressor=None,
+    filters=None,
+    verbose=False,
+    auto_cast_uint=True,
+    **job_kwargs,
+):
+    """
+    Save the trace of a recording extractor in several zarr format.
+
+    Parameters
+    ----------
+    recording: RecordingExtractor
+        The recording extractor object to be saved in .dat format
+    zarr_group: zarr.Group
+        The zarr group to add traces to
+    dataset_paths: list
+        List of paths to traces datasets in the zarr group
+    channel_chunk_size: int or None, default: None (chunking in time only)
+        Channels per chunk
+    dtype: dtype, default: None
+        Type of the saved data
+    compressor: zarr compressor or None, default: None
+        Zarr compressor
+    filters: list, default: None
+        List of zarr filters
+    verbose: bool, default: False
+        If True, output is verbose (when chunks are used)
+    auto_cast_uint: bool, default: True
+        If True, unsigned integers are automatically cast to int if the specified dtype is signed
+    {}
+    """
+    from .job_tools import (
+        ensure_chunk_size,
+        fix_job_kwargs,
+        ChunkRecordingExecutor,
+    )
+    from .core_tools import determine_cast_unsigned
+
+    assert dataset_paths is not None, "Provide 'file_path'"
+
+    if not isinstance(dataset_paths, list):
+        dataset_paths = [dataset_paths]
+    assert len(dataset_paths) == recording.get_num_segments()
+
+    if dtype is None:
+        dtype = recording.get_dtype()
+    if auto_cast_uint:
+        cast_unsigned = determine_cast_unsigned(recording, dtype)
+    else:
+        cast_unsigned = False
+
+    job_kwargs = fix_job_kwargs(job_kwargs)
+    chunk_size = ensure_chunk_size(recording, **job_kwargs)
+
+    # create zarr datasets files
+    zarr_datasets = []
+    for segment_index in range(recording.get_num_segments()):
+        num_frames = recording.get_num_samples(segment_index)
+        num_channels = recording.get_num_channels()
+        dset_name = dataset_paths[segment_index]
+        shape = (num_frames, num_channels)
+        dset = zarr_group.create_dataset(
+            name=dset_name,
+            shape=shape,
+            chunks=(chunk_size, channel_chunk_size),
+            dtype=dtype,
+            filters=filters,
+            compressor=compressor,
+        )
+        zarr_datasets.append(dset)
+        # synchronizer=zarr.ThreadSynchronizer())
+
+    # use executor (loop or workers)
+    func = _write_zarr_chunk
+    init_func = _init_zarr_worker
+    init_args = (recording, zarr_datasets, dtype, cast_unsigned)
+    executor = ChunkRecordingExecutor(
+        recording, func, init_func, init_args, verbose=verbose, job_name="write_zarr_recording", **job_kwargs
+    )
+    executor.run()
+
+
+# used by write_zarr_recording + ChunkRecordingExecutor
+def _init_zarr_worker(recording, zarr_datasets, dtype, cast_unsigned):
+    import zarr
+
+    # create a local dict per worker
+    worker_ctx = {}
+    worker_ctx["recording"] = recording
+    worker_ctx["zarr_datasets"] = zarr_datasets
+    worker_ctx["dtype"] = np.dtype(dtype)
+    worker_ctx["cast_unsigned"] = cast_unsigned
+
+    return worker_ctx
+
+
+# used by write_zarr_recording + ChunkRecordingExecutor
+def _write_zarr_chunk(segment_index, start_frame, end_frame, worker_ctx):
+    import gc
+
+    # recover variables of the worker
+    recording = worker_ctx["recording"]
+    dtype = worker_ctx["dtype"]
+    zarr_dataset = worker_ctx["zarr_datasets"][segment_index]
+    cast_unsigned = worker_ctx["cast_unsigned"]
+
+    # apply function
+    traces = recording.get_traces(
+        start_frame=start_frame, end_frame=end_frame, segment_index=segment_index, cast_unsigned=cast_unsigned
+    )
+    traces = traces.astype(dtype)
+    zarr_dataset[start_frame:end_frame, :] = traces
+
+    # fix memory leak by forcing garbage collection
+    del traces
+    gc.collect()
diff --git a/src/spikeinterface/core/zarrrecordingextractor.py b/src/spikeinterface/core/zarrrecordingextractor.py
deleted file mode 100644
index 6c15044ad..000000000
--- a/src/spikeinterface/core/zarrrecordingextractor.py
+++ /dev/null
@@ -1,184 +0,0 @@
-from typing import List, Union
-
-from pathlib import Path
-from probeinterface import ProbeGroup
-
-import numpy as np
-
-from .baserecording import BaseRecording, BaseRecordingSegment
-from .core_tools import define_function_from_class
-
-try:
-    import zarr
-
-    HAVE_ZARR = True
-except ImportError:
-    HAVE_ZARR = False
-
-
-class ZarrRecordingExtractor(BaseRecording):
-    """
-    RecordingExtractor for a zarr format
-
-    Parameters
-    ----------
-    root_path: str or Path
-        Path to the zarr root file
-    storage_options: dict or None
-        Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
-
-    Returns
-    -------
-    recording: ZarrRecordingExtractor
-        The recording Extractor
-    """
-
-    extractor_name = "ZarrRecording"
-    installed = HAVE_ZARR  # check at class level if installed or not
-    mode = "file"
-    # error message when not installed
-    installation_mesg = "To use the ZarrRecordingExtractor install zarr: \n\n pip install zarr\n\n"
-    name = "zarr"
-
-    def __init__(self, root_path: Union[Path, str], storage_options=None):
-        assert self.installed, self.installation_mesg
-
-        if storage_options is None:
-            if isinstance(root_path, str):
-                root_path_init = root_path
-                root_path = Path(root_path)
-            else:
-                root_path_init = str(root_path)
-            root_path_kwarg = str(Path(root_path).absolute())
-        else:
-            root_path_init = root_path
-            root_path_kwarg = root_path_init
-
-        self._root = zarr.open(root_path_init, mode="r", storage_options=storage_options)
-
-        sampling_frequency = self._root.attrs.get("sampling_frequency", None)
-        num_segments = self._root.attrs.get("num_segments", None)
-        assert "channel_ids" in self._root.keys(), "'channel_ids' dataset not found!"
-        channel_ids = self._root["channel_ids"][:]
-
-        assert sampling_frequency is not None, "'sampling_frequency' attiribute not found!"
-        assert num_segments is not None, "'num_segments' attiribute not found!"
-
-        channel_ids = np.array(channel_ids)
-
-        dtype = self._root["traces_seg0"].dtype
-
-        BaseRecording.__init__(self, sampling_frequency, channel_ids, dtype)
-
-        dtype = np.dtype(dtype)
-        t_starts = self._root.get("t_starts", None)
-
-        total_nbytes = 0
-        total_nbytes_stored = 0
-        cr_by_segment = {}
-        for segment_index in range(num_segments):
-            trace_name = f"traces_seg{segment_index}"
-            assert (
-                len(channel_ids) == self._root[trace_name].shape[1]
-            ), f"Segment {segment_index} has the wrong number of channels!"
-
-            time_kwargs = {}
-            time_vector = self._root.get(f"times_seg{segment_index}", None)
-            if time_vector is not None:
-                time_kwargs["time_vector"] = time_vector
-            else:
-                if t_starts is None:
-                    t_start = None
-                else:
-                    t_start = t_starts[segment_index]
-                    if np.isnan(t_start):
-                        t_start = None
-                time_kwargs["t_start"] = t_start
-                time_kwargs["sampling_frequency"] = sampling_frequency
-
-            rec_segment = ZarrRecordingSegment(self._root, trace_name, **time_kwargs)
-
-            nbytes_segment = self._root[trace_name].nbytes
-            nbytes_stored_segment = self._root[trace_name].nbytes_stored
-            cr_by_segment[segment_index] = nbytes_segment / nbytes_stored_segment
-
-            total_nbytes += nbytes_segment
-            total_nbytes_stored += nbytes_stored_segment
-            self.add_recording_segment(rec_segment)
-
-        # load probe
-        probe_dict = self._root.attrs.get("probe", None)
-        if probe_dict is not None:
-            probegroup = ProbeGroup.from_dict(probe_dict)
-            self.set_probegroup(probegroup, in_place=True)
-
-        # load properties
-        if "properties" in self._root:
-            prop_group = self._root["properties"]
-            for key in prop_group.keys():
-                values = self._root["properties"][key]
-                self.set_property(key, values)
-
-        # load annotations
-        annotations = self._root.attrs.get("annotations", None)
-        if annotations is not None:
-            self.annotate(**annotations)
-        # annotate compression ratios
-        cr = total_nbytes / total_nbytes_stored
-        self.annotate(compression_ratio=cr, compression_ratio_segments=cr_by_segment)
-
-        self._kwargs = {"root_path": root_path_kwarg, "storage_options": storage_options}
-
-
-class ZarrRecordingSegment(BaseRecordingSegment):
-    def __init__(self, root, dataset_name, **time_kwargs):
-        BaseRecordingSegment.__init__(self, **time_kwargs)
-        self._timeseries = root[dataset_name]
-
-    def get_num_samples(self) -> int:
-        """Returns the number of samples in this signal block
-
-        Returns:
-            SampleIndex: Number of samples in the signal block
-        """
-        return self._timeseries.shape[0]
-
-    def get_traces(
-        self,
-        start_frame: Union[int, None] = None,
-        end_frame: Union[int, None] = None,
-        channel_indices: Union[List, None] = None,
-    ) -> np.ndarray:
-        traces = self._timeseries[start_frame:end_frame]
-        if channel_indices is not None:
-            traces = traces[:, channel_indices]
-        return traces
-
-
-read_zarr = define_function_from_class(source_class=ZarrRecordingExtractor, name="read_zarr")
-
-
-def get_default_zarr_compressor(clevel=5):
-    """
-    Return default Zarr compressor object for good preformance in int16
-    electrophysiology data.
-
-    cname: zstd (zstandard)
-    clevel: 5
-    shuffle: BITSHUFFLE
-
-    Parameters
-    ----------
-    clevel : int, default: 5
-        Compression level (higher -> more compressed).
-        Minimum 1, maximum 9. By default 5
-
-    Returns
-    -------
-    Blosc.compressor
-        The compressor object that can be used with the save to zarr function
-    """
-    assert ZarrRecordingExtractor.installed, ZarrRecordingExtractor.installation_mesg
-    from numcodecs import Blosc
-
-    return Blosc(cname="zstd", clevel=clevel, shuffle=Blosc.BITSHUFFLE)
diff --git a/src/spikeinterface/extractors/extractorlist.py b/src/spikeinterface/extractors/extractorlist.py
index f8198c3d1..1ea8bb8ec 100644
--- a/src/spikeinterface/extractors/extractorlist.py
+++ b/src/spikeinterface/extractors/extractorlist.py
@@ -11,6 +11,7 @@ from spikeinterface.core import (
     NumpySorting,
     NpySnippetsExtractor,
     ZarrRecordingExtractor,
+    ZarrSortingExtractor,
     read_binary,
     read_zarr,
     read_npz_sorting,
@@ -79,8 +80,9 @@ recording_extractor_full_list += neo_recording_extractors_list
 
 sorting_extractor_full_list = [
     NpzSortingExtractor,
-    # natively implemented in spikeinterface.extractors
+    ZarrSortingExtractor,
     NumpySorting,
+    # natively implemented in spikeinterface.extractors
     MdaSortingExtractor,
     SHYBRIDSortingExtractor,
     ALFSortingExtractor,


Test patch:
diff --git a/src/spikeinterface/core/testing.py b/src/spikeinterface/core/testing.py
index 1a13974b5..c6843c64c 100644
--- a/src/spikeinterface/core/testing.py
+++ b/src/spikeinterface/core/testing.py
@@ -78,6 +78,8 @@ def check_sortings_equal(
 ) -> None:
     assert SX1.get_num_segments() == SX2.get_num_segments()
 
+    max_spike_index = SX1.to_spike_vector()["sample_index"].max()
+
     # TODO for later  use to_spike_vector() to do this without looping
     for segment_idx in range(SX1.get_num_segments()):
         # get_unit_ids
@@ -88,6 +90,20 @@ def check_sortings_equal(
             train1 = np.sort(SX1.get_unit_spike_train(id, segment_index=segment_idx))
             train2 = np.sort(SX2.get_unit_spike_train(id, segment_index=segment_idx))
             assert np.array_equal(train1, train2)
+            train1 = np.sort(SX1.get_unit_spike_train(id, segment_index=segment_idx, start_frame=30))
+            train2 = np.sort(SX2.get_unit_spike_train(id, segment_index=segment_idx, start_frame=30))
+            assert np.array_equal(train1, train2)
+            # test that slicing works correctly
+            train1 = np.sort(SX1.get_unit_spike_train(id, segment_index=segment_idx, end_frame=max_spike_index - 30))
+            train2 = np.sort(SX2.get_unit_spike_train(id, segment_index=segment_idx, end_frame=max_spike_index - 30))
+            assert np.array_equal(train1, train2)
+            train1 = np.sort(
+                SX1.get_unit_spike_train(id, segment_index=segment_idx, start_frame=30, end_frame=max_spike_index - 30)
+            )
+            train2 = np.sort(
+                SX2.get_unit_spike_train(id, segment_index=segment_idx, start_frame=30, end_frame=max_spike_index - 30)
+            )
+            assert np.array_equal(train1, train2)
 
     if check_annotations:
         check_extractor_annotations_equal(SX1, SX2)
diff --git a/src/spikeinterface/core/tests/test_basesorting.py b/src/spikeinterface/core/tests/test_basesorting.py
index 9e974387f..b0ac99089 100644
--- a/src/spikeinterface/core/tests/test_basesorting.py
+++ b/src/spikeinterface/core/tests/test_basesorting.py
@@ -130,6 +130,17 @@ def test_BaseSorting():
     del sorting6
     del sorting5
 
+    # test save to zarr
+    # compressor = get_default_zarr_compressor()
+    sorting_zarr = sorting.save(format="zarr", folder=cache_folder / "sorting")
+    sorting_zarr_loaded = load_extractor(cache_folder / "sorting.zarr")
+    # annotations is False because Zarr adds compression ratios
+    check_sortings_equal(sorting, sorting_zarr, check_annotations=False, check_properties=True)
+    check_sortings_equal(sorting_zarr, sorting_zarr_loaded, check_annotations=False, check_properties=True)
+    for annotation_name in sorting.get_annotation_keys():
+        assert sorting.get_annotation(annotation_name) == sorting_zarr.get_annotation(annotation_name)
+        assert sorting.get_annotation(annotation_name) == sorting_zarr_loaded.get_annotation(annotation_name)
+
 
 def test_npy_sorting():
     sfreq = 10
diff --git a/src/spikeinterface/generation/tests/test_mock.py b/src/spikeinterface/generation/tests/test_mock.py
new file mode 100644
index 000000000..37c6bde47
--- /dev/null
+++ b/src/spikeinterface/generation/tests/test_mock.py
@@ -0,0 +1,3 @@
+def test_mock():
+    # TODO: Add test logic here
+    pass