diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index b702237445..4b4377b159 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/src/spikeinterface/core/sortinganalyzer.py @@ -41,6 +41,7 @@ from .sparsity import ChannelSparsity, estimate_sparsity from .sortingfolder import NumpyFolderSorting from .zarrextractors import get_default_zarr_compressor, ZarrSortingExtractor, super_zarr_open, _write_object_array +from .zarr_tools import write_json_array, read_encoded_array from .node_pipeline import run_node_pipeline # Typing hints @@ -1076,17 +1077,14 @@ def create_zarr( if recording is not None: rec_dict = recording.to_dict(relative_to=relative_to, recursive=True) if recording.check_serializability("json"): - _write_object_array(zarr_root, "recording", check_json(rec_dict), codec="json") - elif recording.check_serializability("pickle"): - try: - _write_object_array(zarr_root, "recording", rec_dict, codec="pickle") - except: - warnings.warn( - "Failed to serialize recording with Pickle Codec! " - "The recording link will be lost for future load" - ) + compressor = saving_options.get("compressor", get_default_zarr_compressor()) + write_json_array(zarr_root, "recording", check_json(rec_dict), compressor=compressor) else: - warnings.warn("The Recording is not serializable! The recording link will be lost for future load") + warnings.warn( + "The recording is not JSON-serializable, so it is not saved in the zarr SortingAnalyzer, which " + "will load without a recording. Use `set_temporary_recording()` after loading to attach it " + "again, or save the SortingAnalyzer in 'binary_folder' format, which keeps it." + ) else: assert rec_attributes is not None, "recording or rec_attributes must be provided" warnings.warn("Recording not provided, instantiating SortingAnalyzer in recordingless mode.") @@ -1187,7 +1185,7 @@ def load_from_zarr( if recording is None: rec_field = zarr_root.get("recording") if rec_field is not None: - rec_dict = rec_field[0] + rec_dict = read_encoded_array(rec_field) try: recording = load(rec_dict, base_folder=folder) except: @@ -3371,7 +3369,10 @@ def load_data(self, lazy=False): extension_group = self._get_zarr_extension_group(mode="r") for ext_data_name in extension_group.keys(): ext_data_ = extension_group[ext_data_name] - if "dict" in ext_data_.attrs: + if "encoding" in ext_data_.attrs: + ext_data = read_encoded_array(ext_data_) + elif "dict" in ext_data_.attrs: + # legacy json object array (spikeinterface<=0.105) ext_data = ext_data_[0] elif "dataframe" in ext_data_.attrs: import pandas as pd @@ -3552,9 +3553,9 @@ def _save_data(self): if ext_data_name in extension_group: del extension_group[ext_data_name] if isinstance(ext_data, (dict, list)): - ext_data_ = check_json(ext_data) - _write_object_array(extension_group, ext_data_name, ext_data_, codec="json") - extension_group[ext_data_name].attrs["dict"] = True + write_json_array( + extension_group, ext_data_name, check_json(ext_data), compressor=saving_options["compressor"] + ) elif isinstance(ext_data, np.ndarray): # only save the array if the dataset does not already exist, since it is created directly # by the run_node_pipeline() function in the case of nodepipeline extensions diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py new file mode 100644 index 0000000000..023de78dbe --- /dev/null +++ b/src/spikeinterface/core/zarr_tools.py @@ -0,0 +1,61 @@ +import json + +import numpy as np + + +def write_json_array(group, name: str, data, compressor=None): + """ + Save json-serializable data as the utf-8 bytes of its json text in a uint8 array. + + This replaces the length-1 object arrays with a json codec, which zarr-python 3 cannot write. + The data is stored as array data rather than as an attribute, so it is only read when needed: + attributes are read every time the group is opened, which is slow for large data such as a + recording dict. The array is marked with an "encoding" attribute, so that other encodings can be + added later. See `read_encoded_array` for the reverse operation. + + Parameters + ---------- + group : zarr.Group + The zarr group in which the array is created. + name : str + The name of the array. + data : dict | list + The json-serializable data to save. + compressor : numcodecs.Codec | None, default: None + The compressor of the array. + + Returns + ------- + array : zarr.Array + The created array. + """ + payload = np.frombuffer(json.dumps(data).encode("utf-8"), dtype="uint8") + array = group.create_dataset(name=name, data=payload, compressor=compressor) + array.attrs["encoding"] = "json" + return array + + +def read_encoded_array(array): + """ + Read data saved by `write_json_array`, or a legacy length-1 object array (json or pickle codec) + saved by spikeinterface<=0.105. + + Parameters + ---------- + array : zarr.Array + The array to read. + + Returns + ------- + data : Any + The saved data. + """ + encoding = array.attrs.get("encoding", None) + if encoding is None: + # legacy length-1 object array + return array[0] + if encoding == "json": + return json.loads(np.asarray(array[:]).tobytes().decode("utf-8")) + raise ValueError( + f"Unknown encoding {encoding!r} for {array.name}. It was probably saved by a newer version of spikeinterface." + ) diff --git a/src/spikeinterface/preprocessing/pipeline.py b/src/spikeinterface/preprocessing/pipeline.py index 0fe320681c..b75cae49f8 100644 --- a/src/spikeinterface/preprocessing/pipeline.py +++ b/src/spikeinterface/preprocessing/pipeline.py @@ -4,6 +4,7 @@ from spikeinterface.core import BaseRecording from spikeinterface.core.core_tools import is_dict_extractor, is_path_remote from spikeinterface.core.zarrextractors import super_zarr_open +from spikeinterface.core.zarr_tools import read_encoded_array from spikeinterface.preprocessing.preprocessing_classes import preprocessor_dict, _all_preprocesser_dict pp_names_to_functions = {preprocessor.__name__: preprocessor for preprocessor in preprocessor_dict.values()} @@ -265,7 +266,7 @@ def get_preprocessing_list_from_analyzer(analyzer_folder, format="auto", backend rec_field = zarr_root.get("recording") if rec_field is not None: - recording_dict = rec_field[0] + recording_dict = read_encoded_array(rec_field) else: recording_dict = {}