Store JSON data in zarr SortingAnalyzers as byte arrays instead of object arrays

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

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."
+    )

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 = {}