diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py index 4b4377b159..bd03f0b55e 100644 --- a/src/spikeinterface/core/sortinganalyzer.py +++ b/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): """ diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index 023de78dbe..776a33dfb0 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/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." ) diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index d2a604e38b..4f478e425c 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/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