Merge pull request #2884 from alejoe91/fix-zarr-properties

Protect zarr properties of type object
SpikeInterface · May 21, 2024 · ce1deaf · ce1deaf
2 parents 2e7c79c + c1f2404
commit ce1deaf
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from pathlib import Path
 import numpy as np
 import zarr
@@ -308,13 +309,14 @@ def get_default_zarr_compressor(clevel: int = 5):
     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]
-):
+def add_properties_and_annotations(zarr_group: zarr.hierarchy.Group, recording_or_sorting: 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)
+        if values.dtype.kind == "O":
+            warnings.warn(f"Property {key} not saved because it is a python Object type")
+            continue
         prop_group.create_dataset(name=key, data=values, compressor=None)
 
     # save annotations

diff --git a/src/spikeinterface/extractors/phykilosortextractors.py b/src/spikeinterface/extractors/phykilosortextractors.py
@@ -165,14 +165,18 @@ def __init__(
                 self.set_property(key="group", values=cluster_info[prop_name])
             elif prop_name == "cluster_id":
                 self.set_property(key="original_cluster_id", values=cluster_info[prop_name])
-            elif prop_name != "group":
-                self.set_property(key=prop_name, values=cluster_info[prop_name])
             elif prop_name == "group":
                 # rename group property to 'quality'
                 self.set_property(key="quality", values=cluster_info[prop_name])
             else:
                 if load_all_cluster_properties:
-                    self.set_property(key=prop_name, values=cluster_info[prop_name])
+                    # pandas loads strings as objects
+                    if cluster_info[prop_name].values.dtype.kind == "O":
+                        prop_dtype = type(cluster_info[prop_name].values[0])
+                        values_ = cluster_info[prop_name].values.astype(prop_dtype)
+                    else:
+                        values_ = cluster_info[prop_name].values
+                    self.set_property(key=prop_name, values=values_)
 
         self.annotate(phy_folder=str(phy_folder.resolve()))