Store pickled extension data in zarr SortingAnalyzers as byte arrays instead of object arrays

src/spikeinterface/core/sortinganalyzer.py

@@ -40,8 +40,8 @@
 from .numpyextractors import NumpySorting
 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 .zarrextractors import get_default_zarr_compressor, ZarrSortingExtractor, super_zarr_open
+from .zarr_tools import write_json_array, write_pickle_array, read_encoded_array
 from .node_pipeline import run_node_pipeline
 
 # Typing hints
@@ -3370,10 +3370,12 @@ def load_data(self, lazy=False):
             for ext_data_name in extension_group.keys():
                 ext_data_ = extension_group[ext_data_name]
                 if "encoding" in ext_data_.attrs:
+                    if lazy and ext_data_.attrs["encoding"] == "pickle":
+                        continue
                     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]
+                    ext_data = read_encoded_array(ext_data_)
                 elif "dataframe" in ext_data_.attrs:
                     import pandas as pd
 
@@ -3384,9 +3386,10 @@ def load_data(self, lazy=False):
                             ext_data.loc[:, col] = ext_data_[col][:]
                     ext_data = ext_data.convert_dtypes()
                 elif "object" in ext_data_.attrs:
+                    # legacy pickle object array (spikeinterface<=0.105)
                     if lazy:
                         continue
-                    ext_data = ext_data_[0]
+                    ext_data = read_encoded_array(ext_data_)
                 else:
                     ext_data = ext_data_ if lazy else np.array(ext_data_[:])
                 self.set_data(ext_data_name, ext_data)
@@ -3577,10 +3580,11 @@ def _save_data(self):
                 else:
                     # any object
                     try:
-                        _write_object_array(extension_group, ext_data_name, ext_data, codec="pickle")
+                        write_pickle_array(
+                            extension_group, ext_data_name, ext_data, compressor=saving_options["compressor"]
+                        )
                     except:
                         raise Exception(f"Could not save {ext_data_name} as extension data")
-                    extension_group[ext_data_name].attrs["object"] = True
 
     def _reset_extension_folder(self):
         """

src/spikeinterface/core/zarr_tools.py

@@ -1,4 +1,5 @@
 import json
+import pickle
 
 import numpy as np
 
@@ -35,10 +36,40 @@ def write_json_array(group, name: str, data, compressor=None):
     return array
 
 
+def write_pickle_array(group, name: str, data, compressor=None):
+    """
+    Save a python object as its pickle bytes in a uint8 array.
+
+    This replaces the length-1 object arrays with a pickle codec, which zarr-python 3 cannot write.
+    It is used for extension data that is not json-serializable, such as the PCA models of the
+    "principal_components" extension. 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 : Any
+        The picklable object to save.
+    compressor : numcodecs.Codec | None, default: None
+        The compressor of the array.
+
+    Returns
+    -------
+    array : zarr.Array
+        The created array.
+    """
+    payload = np.frombuffer(pickle.dumps(data), dtype="uint8")
+    array = group.create_dataset(name=name, data=payload, compressor=compressor)
+    array.attrs["encoding"] = "pickle"
+    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.
+    Read data saved by `write_json_array` or `write_pickle_array`, or a legacy length-1 object array
+    (json or pickle codec) saved by spikeinterface<=0.105.
 
     Parameters
     ----------
@@ -56,6 +87,8 @@ def read_encoded_array(array):
         return array[0]
     if encoding == "json":
         return json.loads(np.asarray(array[:]).tobytes().decode("utf-8"))
+    if encoding == "pickle":
+        return pickle.loads(np.asarray(array[:]).tobytes())
     raise ValueError(
         f"Unknown encoding {encoding!r} for {array.name}. It was probably saved by a newer version of spikeinterface."
     )

src/spikeinterface/core/zarrextractors.py

@@ -744,51 +744,6 @@ def create_zarr_path_for_write(folder_path: str | Path, overwrite: bool = False)
     return folder_path
 
 
-def _write_object_array(
-    group,
-    name: str,
-    data,
-    codec: str = "json",
-    overwrite: bool = True,
-):
-    """
-    Write a length-1 object-dtype array holding a Python dict/list/object.
-
-    Centralizes the v2/v3 codec-placement difference for object blobs: under zarr-v2
-    the object codec goes in ``object_codec=``; under zarr-v3 it goes in ``filters=``
-    (wrapped via ``numcodecs.zarr3.*``). The helper picks the right path automatically.
-
-    Parameters
-    ----------
-    group : zarr.Group
-        The zarr group to write into.
-    name : str
-        Name of the array inside ``group``.
-    data : Any
-        The Python object to store. Wrapped into ``np.array([data], dtype=object)``.
-    codec : {"json", "pickle"}, default: "json"
-        Which object codec to use.
-    overwrite : bool, default: True
-        Whether to overwrite an existing array with the same name.
-    """
-    import numcodecs
-
-    if codec == "json":
-        codec_instance = numcodecs.JSON()
-    elif codec == "pickle":
-        codec_instance = numcodecs.Pickle()
-    else:
-        raise ValueError(f"codec must be 'json' or 'pickle', got {codec!r}")
-
-    arr = np.array([data], dtype=object)
-    return group.create_dataset(
-        name=name,
-        data=arr,
-        object_codec=codec_instance,
-        overwrite=overwrite,
-    )
-
-
 def get_default_zarr_compressor(clevel: int = 5):
     """
     Return default Zarr compressor object for good preformance in int16