feat(extension_types): support SpikeInterface BaseSorting objects (ITL-468)

src/orcapod/contexts/data/v0.1.json

@@ -50,6 +50,11 @@
                             "_config": {},
                             "_optional": true
                         },
+                        {
+                            "_class": "orcapod.extension_types.spikeinterface_types.LogicalSISorting",
+                            "_config": {},
+                            "_optional": true
+                        },
                         {
                             "_class": "orcapod.extension_types.pandas_type.LogicalPandasDataFrame",
                             "_config": {}
@@ -108,6 +113,7 @@
                 [{"_type": "pyarrow.Table"},            {"_class": "orcapod.hashing.semantic_hashing.builtin_handlers.ArrowTableHandler",   "_config": {}}],
                 [{"_type": "pyarrow.RecordBatch"},      {"_class": "orcapod.hashing.semantic_hashing.builtin_handlers.ArrowTableHandler",   "_config": {}}],
                 [{"_type": "spikeinterface.core.BaseRecording", "_optional": true}, {"_class": "orcapod.extension_types.spikeinterface_types.SIRecordingHandler", "_config": {}, "_optional": true}],
+                [{"_type": "spikeinterface.core.BaseSorting", "_optional": true}, {"_class": "orcapod.extension_types.spikeinterface_types.SISortingHandler", "_config": {}, "_optional": true}],
                 [{"_type": "numpy.ndarray"},            {"_class": "orcapod.hashing.semantic_hashing.builtin_handlers.NumpyArrayHandler",        "_config": {}}],
                 [{"_type": "pandas.DataFrame"},         {"_class": "orcapod.hashing.semantic_hashing.builtin_handlers.PandasDataFrameHandler",  "_config": {}}],
                 [{"_type": "pandas.Series"},            {"_class": "orcapod.hashing.semantic_hashing.builtin_handlers.PandasSeriesHandler",      "_config": {}}]
@@ -145,7 +151,8 @@
             "Added orcapod.Directory content-identified type with recursive Merkle tree hashing, LogicalDirectory Arrow extension, ignore filter support, and DirectoryHandler (ITL-451)",
             "Added numpy.ndarray as a native value type via LogicalNumpyArray (large_binary/.npy) and NumpyArrayHandler; object-dtype arrays are rejected eagerly (ITL-460)",
             "Added spikeinterface.BaseRecording as a native value type via LogicalSIRecording (large_string/JSON) and SIRecordingHandler (SHA-256 of JSON bytes); auto-registered when spikeinterface is installed via _optional entries in v0.1.json; added _optional flag support to parse_objectspec for optional-extras types (ITL-459)",
-            "Added pandas.DataFrame and pandas.Series as native value types via LogicalPandasDataFrame and LogicalPandasSeries (Arrow IPC / large_binary), with index preservation and PandasDataFrameHandler / PandasSeriesHandler for content hashing using StarfixArrowHasher (PLT-1869)"
+            "Added pandas.DataFrame and pandas.Series as native value types via LogicalPandasDataFrame and LogicalPandasSeries (Arrow IPC / large_binary), with index preservation and PandasDataFrameHandler / PandasSeriesHandler for content hashing using StarfixArrowHasher (PLT-1869)",
+            "Added spikeinterface.BaseSorting as a native value type via LogicalSISorting (large_string/JSON) and SISortingHandler (SHA-256 of JSON bytes); auto-registered when spikeinterface is installed via _optional entries in v0.1.json (ITL-468)"
         ]
     }
 }

src/orcapod/extension_types/__init__.py

@@ -30,9 +30,13 @@
 from .numpy_type import LogicalNumpyArray  # ITL-460
 from .pandas_type import LogicalPandasDataFrame, LogicalPandasSeries  # PLT-1869
 
-# ITL-459 — SpikeInterface support (optional; requires pip install orcapod[spikeinterface])
+# ITL-459, ITL-468 — SpikeInterface support (optional; requires pip install orcapod[spikeinterface])
 try:
-    from .spikeinterface_types import LogicalSIRecording, register_spikeinterface_types
+    from .spikeinterface_types import (
+        LogicalSIRecording,
+        LogicalSISorting,
+        register_spikeinterface_types,
+    )
     _SI_AVAILABLE = True
 except ImportError:
     _SI_AVAILABLE = False
@@ -64,8 +68,8 @@
     "LogicalDirectory",
     # ITL-460
     "LogicalNumpyArray",
-    # ITL-459 (conditional — only present when spikeinterface is installed)
-    *( ["LogicalSIRecording", "register_spikeinterface_types"] if _SI_AVAILABLE else [] ),
+    # ITL-459, ITL-468 (conditional — only present when spikeinterface is installed)
+    *( ["LogicalSIRecording", "LogicalSISorting", "register_spikeinterface_types"] if _SI_AVAILABLE else [] ),
     # PLT-1869
     "LogicalPandasDataFrame",
     "LogicalPandasSeries",

src/orcapod/extension_types/spikeinterface_types.py

@@ -1,16 +1,20 @@
-"""SpikeInterface LogicalType and handler for orcapod (ITL-459).
+"""SpikeInterface LogicalTypes and handlers for orcapod (ITL-459, ITL-468).
 
-`LogicalSIRecording` maps `spikeinterface.core.BaseRecording` ↔ Arrow
-`large_string` using SpikeInterface's own `to_dict(recursive=True,
-include_annotations=True, include_properties=False)` JSON dump (encoded
-via `SIJsonEncoder`) as the storage envelope. `SIRecordingHandler` hashes
+``LogicalSIRecording`` maps ``spikeinterface.core.BaseRecording`` ↔ Arrow
+``large_string`` using SpikeInterface's own ``to_dict(recursive=True,
+include_annotations=True, include_properties=False)`` JSON dump (encoded
+via ``SIJsonEncoder``) as the storage envelope. ``SIRecordingHandler`` hashes
 the same JSON bytes via SHA-256 for content identity.
 
-This module requires the optional `spikeinterface` extras group:
-`pip install orcapod[spikeinterface]`
+``LogicalSISorting`` maps ``spikeinterface.core.BaseSorting`` ↔ Arrow
+``large_string`` using the same serialization approach. ``SISortingHandler``
+hashes the JSON bytes via SHA-256.
+
+This module requires the optional ``spikeinterface`` extras group:
+``pip install orcapod[spikeinterface]``
 
 Register SI types into the default orcapod context before using them in
-pods: call `register_spikeinterface_types()` once at startup.
+pods: call ``register_spikeinterface_types()`` once at startup.
 """
 
 from __future__ import annotations
@@ -32,7 +36,7 @@
     from orcapod.protocols.hashing_protocols import SemanticHasherProtocol
 
 try:
-    from spikeinterface.core import BaseRecording
+    from spikeinterface.core import BaseRecording, BaseSorting
 except ImportError as _exc:
     raise ImportError(
         "spikeinterface is not installed. "
@@ -168,6 +172,134 @@ def storage_to_python(
         return si_load(si_dict)
 
 
+class LogicalSISorting(BaseLogicalType):
+    """Logical type for ``spikeinterface.core.BaseSorting``.
+
+    Stores ``BaseSorting`` instances as Arrow ``large_string`` columns
+    tagged with extension name ``"spikeinterface.sorting"``. The stored
+    value is SpikeInterface's own ``to_dict(recursive=True,
+    include_annotations=True, include_properties=False)`` output, encoded
+    via ``SIJsonEncoder``. Loading reconstructs the sorting via
+    ``spikeinterface.core.load(dict)``.
+
+    Only sortings whose ``check_serializability("json")`` returns ``True``
+    are accepted. File-backed sortings (zarr, numpy_folder, npz_folder,
+    sorter folder) qualify. In-memory ``NumpySorting`` objects do not and
+    raise ``ValueError`` with clear save instructions.
+
+    Example:
+        >>> import tempfile, numpy as np
+        >>> import spikeinterface.core as si
+        >>> from orcapod.extension_types.spikeinterface_types import LogicalSISorting
+        >>> lt = LogicalSISorting()
+        >>> with tempfile.TemporaryDirectory() as tmp:
+        ...     sorting = si.NumpySorting.from_unit_dict(
+        ...         {0: np.array([0, 100, 200])}, sampling_frequency=30000
+        ...     )
+        ...     saved = sorting.save_to_folder(tmp + "/sorting")
+        ...     storage = lt.python_to_storage(saved)
+        ...     recovered = lt.storage_to_python(storage)
+        ...     saved.get_unit_ids().tolist() == recovered.get_unit_ids().tolist()
+        True
+    """
+
+    _arrow_ext_class = make_arrow_extension_type("spikeinterface.sorting", pa.large_string())
+    _arrow_ext: pa.ExtensionType | None = None
+    _polars_ext_class = make_polars_extension_type("spikeinterface.sorting", pa.large_string())
+    _polars_ext: pl.BaseExtension | None = None
+
+    logical_type_name: str = "spikeinterface.sorting"
+    python_type: type = BaseSorting
+
+    def get_arrow_extension_type(self) -> pa.ExtensionType:
+        """Return the cached Arrow extension type for ``BaseSorting``.
+
+        Returns:
+            A ``pa.ExtensionType`` with extension name
+            ``"spikeinterface.sorting"`` and storage type ``pa.large_string()``.
+        """
+        if LogicalSISorting._arrow_ext is None:
+            LogicalSISorting._arrow_ext = LogicalSISorting._arrow_ext_class()
+        return LogicalSISorting._arrow_ext
+
+    def get_polars_extension_type(self) -> pl.BaseExtension:
+        """Return the cached Polars extension type for ``BaseSorting``.
+
+        Returns:
+            A ``pl.BaseExtension`` registered under ``"spikeinterface.sorting"``.
+        """
+        if LogicalSISorting._polars_ext is None:
+            LogicalSISorting._polars_ext = LogicalSISorting._polars_ext_class()
+        return LogicalSISorting._polars_ext
+
+    def python_to_storage(
+        self, value: Any, converter: TypeConverterProtocol | None = None
+    ) -> str:
+        """Serialise a ``BaseSorting`` to its JSON storage representation.
+
+        Args:
+            value: A ``BaseSorting`` instance whose
+                ``check_serializability("json")`` returns ``True``.
+            converter: Ignored. Present for protocol conformance.
+
+        Returns:
+            A JSON string produced by ``sorting.to_dict(recursive=True,
+            include_annotations=True, include_properties=False)`` encoded
+            via ``SIJsonEncoder``.
+
+        Raises:
+            ValueError: If the sorting is not JSON-serialisable (e.g. an
+                in-memory ``NumpySorting``).
+        """
+        if not value.check_serializability("json"):
+            raise ValueError(
+                "This BaseSorting is not JSON-serializable and cannot be stored "
+                "by orcapod. This typically means it holds data in memory (e.g. "
+                "NumpySorting). Sortings built on top of file-backed data "
+                "(zarr, numpy_folder, npz_folder, etc.) are fine and do not need "
+                "to be materialized first. If your sorting is in-memory, call "
+                "sorting.save_to_zarr(path) or sorting.save_to_folder(path) "
+                "first, then pass the returned extractor to the pod."
+            )
+        from spikeinterface.core.core_tools import SIJsonEncoder
+        return json.dumps(
+            value.to_dict(
+                include_annotations=True,
+                include_properties=False,
+                recursive=True,
+            ),
+            cls=SIJsonEncoder,
+        )
+
+    def storage_to_python(
+        self, storage_value: Any, converter: TypeConverterProtocol | None = None
+    ) -> BaseSorting:
+        """Reconstruct a ``BaseSorting`` from its JSON storage string.
+
+        Args:
+            storage_value: A JSON string as stored in Arrow.
+            converter: Ignored. Present for protocol conformance.
+
+        Returns:
+            A ``BaseSorting`` instance reconstructed via
+            ``spikeinterface.core.load``.
+
+        Raises:
+            ValueError: If ``storage_value`` is not valid JSON.
+            FileNotFoundError: If the backing zarr/folder no longer exists
+                (raised by SpikeInterface, propagated as-is).
+        """
+        from spikeinterface.core import load as si_load
+        try:
+            si_dict = json.loads(storage_value)
+        except (json.JSONDecodeError, TypeError) as exc:
+            raise ValueError(
+                f"LogicalSISorting: cannot deserialise storage value "
+                f"{storage_value!r}; expected a JSON string."
+            ) from exc
+        return si_load(si_dict)
+
+
 class SIRecordingHandler:
     """Semantic hash handler for `spikeinterface.core.BaseRecording`.
 
@@ -221,22 +353,77 @@ def handle(self, obj: Any, hasher: SemanticHasherProtocol | None) -> ContentHash
         )
 
 
+class SISortingHandler:
+    """Semantic hash handler for ``spikeinterface.core.BaseSorting``.
+
+    Computes a SHA-256 ``ContentHash`` of the JSON bytes produced by
+    ``sorting.to_dict(recursive=True, include_annotations=True,
+    include_properties=False)`` encoded via ``SIJsonEncoder``. This is
+    identical to the bytes that ``LogicalSISorting`` stores in Arrow, so
+    hash input and storage representation are always consistent.
+
+    The ``hasher`` argument is accepted for protocol conformance but not used —
+    hashing is done directly via ``hashlib.sha256`` to avoid overhead.
+    """
+
+    def handle(self, obj: Any, hasher: SemanticHasherProtocol | None) -> ContentHash:
+        """Return a SHA-256 ``ContentHash`` of the sorting's JSON dump.
+
+        Args:
+            obj: A ``BaseSorting`` instance.
+            hasher: Accepted for protocol conformance; not used.
+
+        Returns:
+            A ``ContentHash`` with ``method="sha256"`` and digest equal to the
+            SHA-256 of the JSON bytes from ``to_dict(recursive=True,
+            include_annotations=True, include_properties=False)`` encoded
+            via ``SIJsonEncoder``.
+
+        Raises:
+            TypeError: If ``obj`` is not a ``BaseSorting``.
+            ValueError: If the sorting is not JSON-serialisable (in-memory).
+        """
+        if not isinstance(obj, BaseSorting):
+            raise TypeError(
+                f"SISortingHandler: expected BaseSorting, got {type(obj)!r}"
+            )
+        if not obj.check_serializability("json"):
+            raise ValueError(
+                "Cannot hash an in-memory BaseSorting "
+                "(check_serializability('json') is False). "
+                "Save it to disk first with save_to_zarr() or save_to_folder()."
+            )
+        # TODO(ITL-468): phase 2 — also hash backing source directory contents
+        from spikeinterface.core.core_tools import SIJsonEncoder
+        json_bytes = json.dumps(
+            obj.to_dict(include_annotations=True, include_properties=False, recursive=True),
+            cls=SIJsonEncoder,
+        ).encode()
+        logger.debug("SISortingHandler: hashing %d JSON bytes", len(json_bytes))
+        return ContentHash(
+            method="sha256",
+            digest=hashlib.sha256(json_bytes).digest(),
+        )
+
+
 def register_spikeinterface_types(context: Any = None) -> None:
-    """Register SpikeInterface LogicalTypes into an orcapod `DataContext`.
+    """Register SpikeInterface LogicalTypes into an orcapod ``DataContext``.
+
+    Registers both ``LogicalSIRecording`` / ``SIRecordingHandler`` (ITL-459)
+    and ``LogicalSISorting`` / ``SISortingHandler`` (ITL-468).
 
     For the default context this is called automatically at startup (the
-    default `v0.1.json` context config lists `LogicalSIRecording` and
-    `SIRecordingHandler` with `"_optional": true`, so they are wired in
-    whenever `spikeinterface` is installed). Call this function explicitly
-    only when you are working with a custom `DataContext` that was not
-    constructed from the default config.
-
-    If `context` is `None`, the default context (from
-    `orcapod.contexts.get_default_context()`) is used. The function is
+    default ``v0.1.json`` context config lists all four with ``"_optional": true``,
+    so they are wired in whenever ``spikeinterface`` is installed). Call this
+    function explicitly only when working with a custom ``DataContext`` that was
+    not constructed from the default config.
+
+    If ``context`` is ``None``, the default context (from
+    ``orcapod.contexts.get_default_context()``) is used. The function is
     idempotent — calling it more than once on the same context is safe.
 
     Args:
-        context: A `DataContext` instance, or `None` to use the default.
+        context: A ``DataContext`` instance, or ``None`` to use the default.
 
     Example:
         >>> from orcapod.extension_types.spikeinterface_types import register_spikeinterface_types
@@ -246,23 +433,40 @@ def register_spikeinterface_types(context: Any = None) -> None:
         from orcapod.contexts import get_default_context
         context = get_default_context()
 
-    lt = LogicalSIRecording()
+    # --- Recording ---
+    lt_recording = LogicalSIRecording()
     try:
-        context.type_converter.register_logical_type(lt)
+        context.type_converter.register_logical_type(lt_recording)
     except ValueError as exc:
         # A different LogicalSIRecording instance is already registered (e.g.
-        # auto-registered from v0.1.json at context creation time).  That is
-        # fine — both instances are equivalent.  Any other ValueError propagates.
+        # auto-registered from v0.1.json at context creation time). That is
+        # fine — both instances are equivalent. Any other ValueError propagates.
         if "already bound to" not in str(exc):
             raise
         logger.debug(
             "register_spikeinterface_types: LogicalSIRecording already registered, skipping"
         )
     else:
+        logger.debug("register_spikeinterface_types: registered LogicalSIRecording")
+
+    # Handler registration silently replaces an existing entry, so always safe.
+    context.semantic_hasher.type_handler_registry.register(BaseRecording, SIRecordingHandler())
+
+    # --- Sorting ---
+    lt_sorting = LogicalSISorting()
+    try:
+        context.type_converter.register_logical_type(lt_sorting)
+    except ValueError as exc:
+        # A different LogicalSISorting instance is already registered (e.g.
+        # auto-registered from v0.1.json at context creation time). That is
+        # fine — both instances are equivalent. Any other ValueError propagates.
+        if "already bound to" not in str(exc):
+            raise
         logger.debug(
-            "register_spikeinterface_types: registered LogicalSIRecording and SIRecordingHandler"
+            "register_spikeinterface_types: LogicalSISorting already registered, skipping"
         )
+    else:
+        logger.debug("register_spikeinterface_types: registered LogicalSISorting")
 
-    # SIRecordingHandler registration silently replaces an existing entry, so
-    # this call is always safe regardless of prior registration state.
-    context.semantic_hasher.type_handler_registry.register(BaseRecording, SIRecordingHandler())
+    # Handler registration silently replaces an existing entry, so always safe.
+    context.semantic_hasher.type_handler_registry.register(BaseSorting, SISortingHandler())