Add b2o carrier support for persisted C2Array objects

Author: FrancescAlted

src/blosc2/__init__.py

@@ -538,6 +538,7 @@ def _raise(exc):
 from .batch_store import Batch, BatchStore
 from .vlarray import VLArray, vlarray_from_cframe
 from .ref import Ref
+from .b2objects import _open_b2object
 
 from .c2array import c2context, C2Array, URLPath
 

src/blosc2/_msgpack_utils.py src/blosc2/b2objects.pysrc/blosc2/_msgpack_utils.py renamed to src/blosc2/b2objects.py

@@ -12,50 +12,57 @@
 import linecache
 import textwrap
 from dataclasses import asdict
+from typing import Any
 
 import numpy as np
-from msgpack import ExtType, packb, unpackb
 
-from blosc2 import blosc2_ext
+import blosc2
 from blosc2.dsl_kernel import DSLKernel
-from blosc2.ref import Ref
 
-# Msgpack extension type codes are application-defined.  Reserve code 42 in
-# python-blosc2 for values serialized as Blosc2 CFrames via ``to_cframe()`` and
-# reconstructed with ``blosc2.from_cframe()``.  Keep this stable for backward
-# compatibility with persisted msgpack payloads produced by this package.
-_BLOSC2_EXT_CODE = 42
-# Reserve code 43 for structured Blosc2 reference objects that are not naturally
-# serialized as CFrames.  The payload is a msgpack-encoded mapping with a
-# stable ``kind`` and ``version`` envelope.
-_BLOSC2_STRUCTURED_EXT_CODE = 43
-_BLOSC2_STRUCTURED_VERSION = 1
-_BLOSC2_DSL_VERSION = 1
+_B2OBJECT_META_KEY = "b2o"
+_B2OBJECT_VERSION = 1
+_B2OBJECT_DSL_VERSION = 1
+
+
+def _make_b2object_carrier(
+    kind: str,
+    shape,
+    dtype,
+    *,
+    chunks=None,
+    blocks=None,
+    **kwargs,
+):
+    meta = dict(kwargs.pop("meta", {}))
+    meta[_B2OBJECT_META_KEY] = {"kind": kind, "version": _B2OBJECT_VERSION}
+    kwargs["meta"] = meta
+    return blosc2.empty(shape=shape, dtype=dtype, chunks=chunks, blocks=blocks, **kwargs)
+
+
+def _write_b2object_payload(array, payload: dict[str, Any]) -> None:
+    array.schunk.vlmeta[_B2OBJECT_META_KEY] = payload
 
 
 def _encode_operand_reference(obj):
-    return Ref.from_object(obj).to_dict()
+    return blosc2.Ref.from_object(obj).to_dict()
+
 
+def _decode_operand_reference(payload):
+    return blosc2.Ref.from_dict(payload).open()
 
-def _encode_structured_reference(obj):
-    import blosc2
 
-    if isinstance(obj, blosc2.Ref):
-        payload = {"kind": "ref", "version": _BLOSC2_STRUCTURED_VERSION, "ref": obj.to_dict()}
-        return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
+def _encode_b2object_payload(obj) -> dict[str, Any] | None:
     if isinstance(obj, blosc2.C2Array):
-        payload = _encode_operand_reference(obj)
-        return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
+        return blosc2.Ref.c2array_ref(obj.path, obj.urlbase).to_dict()
     if isinstance(obj, blosc2.LazyExpr):
         expression = obj.expression_tosave if hasattr(obj, "expression_tosave") else obj.expression
         operands = obj.operands_tosave if hasattr(obj, "operands_tosave") else obj.operands
-        payload = {
+        return {
             "kind": "lazyexpr",
-            "version": _BLOSC2_STRUCTURED_VERSION,
+            "version": _B2OBJECT_VERSION,
             "expression": expression,
             "operands": {key: _encode_operand_reference(value) for key, value in operands.items()},
         }
-        return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
     if isinstance(obj, blosc2.LazyUDF):
         if not isinstance(obj.func, DSLKernel):
             raise TypeError("Structured Blosc2 msgpack payload only supports LazyUDF backed by DSLKernel")
@@ -75,15 +82,11 @@ def _encode_structured_reference(obj):
                 kwargs[key] = asdict(value)
             else:
                 kwargs[key] = value
-        # Keep both source forms:
-        # - udf_source recreates the executable Python function object
-        # - dsl_source preserves the DSLKernel's normalized DSL metadata so the
-        #   reconstructed function can keep its DSL identity and fast-path hints
-        payload = {
+        return {
             "kind": "lazyudf",
-            "version": _BLOSC2_STRUCTURED_VERSION,
+            "version": _B2OBJECT_VERSION,
             "function_kind": "dsl",
-            "dsl_version": _BLOSC2_DSL_VERSION,
+            "dsl_version": _B2OBJECT_DSL_VERSION,
             "name": udf_name,
             "udf_source": udf_source,
             "dsl_source": obj.func.dsl_source,
@@ -92,76 +95,60 @@ def _encode_structured_reference(obj):
             "operands": {f"o{i}": _encode_operand_reference(value) for i, value in enumerate(obj.inputs)},
             "kwargs": kwargs,
         }
-        return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
     return None
 
 
-def _decode_operand_reference(payload):
-    return Ref.from_dict(payload).open()
-
-
-def _decode_structured_reference(data):
-    payload = unpackb(data)
-    if not isinstance(payload, dict):
-        raise TypeError("Structured Blosc2 msgpack payload must decode to a mapping")
-
-    version = payload.get("version")
-    if version != _BLOSC2_STRUCTURED_VERSION:
-        raise ValueError(f"Unsupported structured Blosc2 msgpack payload version: {version!r}")
-
+def _decode_b2object_payload(payload: dict[str, Any]):
     kind = payload.get("kind")
-    if kind == "ref":
-        ref_payload = payload.get("ref")
-        return Ref.from_dict(ref_payload)
+    version = payload.get("version")
+    if version != _B2OBJECT_VERSION:
+        raise ValueError(f"Unsupported persisted Blosc2 object version: {version!r}")
     if kind == "c2array":
-        return _decode_operand_reference(payload)
+        ref = blosc2.Ref.from_dict(payload)
+        return ref.open()
     if kind == "lazyexpr":
         return _decode_structured_lazyexpr(payload)
     if kind == "lazyudf":
         return _decode_structured_lazyudf(payload)
-    raise ValueError(f"Unsupported structured Blosc2 msgpack payload kind: {kind!r}")
+    raise ValueError(f"Unsupported persisted Blosc2 object kind: {kind!r}")
 
 
 def _decode_structured_lazyexpr(payload):
-    import blosc2
-
     expression = payload.get("expression")
     if not isinstance(expression, str):
-        raise TypeError("Structured LazyExpr msgpack payload requires a string 'expression'")
+        raise TypeError("Structured LazyExpr payload requires a string 'expression'")
     operands_payload = payload.get("operands")
     if not isinstance(operands_payload, dict):
-        raise TypeError("Structured LazyExpr msgpack payload requires a mapping 'operands'")
+        raise TypeError("Structured LazyExpr payload requires a mapping 'operands'")
     operands = {key: _decode_operand_reference(value) for key, value in operands_payload.items()}
     return blosc2.lazyexpr(expression, operands=operands)
 
 
 def _decode_structured_lazyudf(payload):
-    import blosc2
-
     function_kind = payload.get("function_kind")
     if function_kind != "dsl":
         raise ValueError(f"Unsupported structured LazyUDF function kind: {function_kind!r}")
     dsl_version = payload.get("dsl_version")
-    if dsl_version != _BLOSC2_DSL_VERSION:
+    if dsl_version != _B2OBJECT_DSL_VERSION:
         raise ValueError(f"Unsupported structured LazyUDF DSL version: {dsl_version!r}")
     udf_source = payload.get("udf_source")
     if not isinstance(udf_source, str):
-        raise TypeError("Structured LazyUDF msgpack payload requires a string 'udf_source'")
+        raise TypeError("Structured LazyUDF payload requires a string 'udf_source'")
     name = payload.get("name")
     if not isinstance(name, str):
-        raise TypeError("Structured LazyUDF msgpack payload requires a string 'name'")
+        raise TypeError("Structured LazyUDF payload requires a string 'name'")
     dtype = payload.get("dtype")
     if not isinstance(dtype, str):
-        raise TypeError("Structured LazyUDF msgpack payload requires a string 'dtype'")
+        raise TypeError("Structured LazyUDF payload requires a string 'dtype'")
     shape_payload = payload.get("shape")
     if not isinstance(shape_payload, list):
-        raise TypeError("Structured LazyUDF msgpack payload requires a list 'shape'")
+        raise TypeError("Structured LazyUDF payload requires a list 'shape'")
     operands_payload = payload.get("operands")
     if not isinstance(operands_payload, dict):
-        raise TypeError("Structured LazyUDF msgpack payload requires a mapping 'operands'")
+        raise TypeError("Structured LazyUDF payload requires a mapping 'operands'")
     kwargs = payload.get("kwargs", {})
     if not isinstance(kwargs, dict):
-        raise TypeError("Structured LazyUDF msgpack payload requires a mapping 'kwargs'")
+        raise TypeError("Structured LazyUDF payload requires a mapping 'kwargs'")
 
     local_ns = {}
     filename = f"<{name}>"
@@ -185,38 +172,26 @@ def _decode_structured_lazyudf(payload):
     return blosc2.lazyudf(func, operands, dtype=np.dtype(dtype), shape=tuple(shape_payload), **kwargs)
 
 
-def _encode_msgpack_ext(obj):
-    import blosc2
-
-    if isinstance(
-        obj, blosc2.NDArray | blosc2.SChunk | blosc2.VLArray | blosc2.BatchStore | blosc2.EmbedStore
-    ):
-        return ExtType(_BLOSC2_EXT_CODE, obj.to_cframe())
-    structured = _encode_structured_reference(obj)
-    if structured is not None:
-        return structured
-    return blosc2_ext.encode_tuple(obj)
-
-
-def msgpack_packb(value):
-    return packb(value, default=_encode_msgpack_ext, strict_types=True, use_bin_type=True)
-
-
-def decode_tuple_list_hook(obj):
-    if obj and isinstance(obj[0], str) and obj[0] == "__tuple__":
-        return tuple(obj[1:])
-    return obj
+def _read_b2object_marker(obj) -> dict[str, Any] | None:
+    schunk = getattr(obj, "schunk", obj)
+    if _B2OBJECT_META_KEY not in schunk.meta:
+        return None
+    return schunk.meta[_B2OBJECT_META_KEY]
 
 
-def _decode_msgpack_ext(code, data):
-    import blosc2
+def _read_b2object_payload(obj) -> dict[str, Any]:
+    schunk = getattr(obj, "schunk", obj)
+    return schunk.vlmeta[_B2OBJECT_META_KEY]
 
-    if code == _BLOSC2_EXT_CODE:
-        return blosc2.from_cframe(data, copy=True)
-    if code == _BLOSC2_STRUCTURED_EXT_CODE:
-        return _decode_structured_reference(data)
-    return ExtType(code, data)
 
+def _open_b2object(obj):
+    marker = _read_b2object_marker(obj)
+    if marker is None:
+        return None
 
-def msgpack_unpackb(payload):
-    return unpackb(payload, list_hook=decode_tuple_list_hook, ext_hook=_decode_msgpack_ext)
+    payload = _read_b2object_payload(obj)
+    if marker.get("version") != _B2OBJECT_VERSION:
+        raise ValueError(f"Unsupported persisted Blosc2 object version: {marker.get('version')!r}")
+    if marker.get("kind") != payload.get("kind"):
+        raise ValueError("Persisted Blosc2 object marker/payload kind mismatch")
+    return _decode_b2object_payload(payload)

src/blosc2/batch_store.py

@@ -18,8 +18,8 @@
 import numpy as np
 
 import blosc2
-from blosc2._msgpack_utils import msgpack_packb, msgpack_unpackb
 from blosc2.info import InfoReporter, format_nbytes_info
+from blosc2.msgpack_utils import msgpack_packb, msgpack_unpackb
 
 _BATCHSTORE_META = {"version": 1, "serializer": "msgpack", "items_per_block": None, "arrow_schema": None}
 _SUPPORTED_SERIALIZERS = {"msgpack", "arrow"}

src/blosc2/c2array.py

@@ -18,6 +18,7 @@
 import requests
 
 import blosc2
+from blosc2.b2objects import _encode_b2object_payload, _make_b2object_carrier, _write_b2object_payload
 from blosc2.info import InfoReporter, format_nbytes_info
 
 _subscriber_data = {
@@ -238,6 +239,37 @@ def __init__(self, path: str, /, urlbase: str | None = None, auth_token: str | N
         cparams.pop("filters, meta", None)
         self._cparams = blosc2.CParams(**cparams)
 
+    def _to_b2object_payload(self) -> dict:
+        payload = _encode_b2object_payload(self)
+        if payload is None:
+            raise TypeError("Unsupported persisted Blosc2 object")
+        return payload
+
+    def _to_b2object_carrier(self, **kwargs):
+        array = _make_b2object_carrier(
+            "c2array",
+            self.shape,
+            self.dtype,
+            chunks=self.chunks,
+            blocks=self.blocks,
+            cparams=self.cparams,
+            **kwargs,
+        )
+        _write_b2object_payload(array, self._to_b2object_payload())
+        return array
+
+    def to_cframe(self) -> bytes:
+        """Serialize the remote array reference as a CFrame-backed Blosc2 object."""
+        return self._to_b2object_carrier().to_cframe()
+
+    def save(self, urlpath: str, contiguous: bool = True, **kwargs) -> None:
+        """Persist the remote array reference using a CFrame-backed carrier."""
+        blosc2.blosc2_ext.check_access_mode(urlpath, "w")
+        kwargs["urlpath"] = urlpath
+        kwargs["contiguous"] = contiguous
+        kwargs["mode"] = "w"
+        self._to_b2object_carrier(**kwargs)
+
     def __getitem__(self, slice_: int | slice | Sequence[slice]) -> np.ndarray:
         """
         Get a slice of the array (returning NumPy array).

src/blosc2/core.py

@@ -1918,7 +1918,9 @@ def ndarray_from_cframe(cframe: bytes | str, copy: bool = False) -> blosc2.NDArr
 
 def from_cframe(
     cframe: bytes | str, copy: bool = True
-) -> blosc2.EmbedStore | blosc2.NDArray | blosc2.SChunk | blosc2.BatchStore | blosc2.VLArray:
+) -> (
+    blosc2.EmbedStore | blosc2.NDArray | blosc2.SChunk | blosc2.BatchStore | blosc2.VLArray | blosc2.C2Array
+):
     """Create a :ref:`EmbedStore <EmbedStore>`, :ref:`NDArray <NDArray>`, :ref:`SChunk <SChunk>`,
     :ref:`BatchStore <BatchStore>` or :ref:`VLArray <VLArray>` instance
     from a contiguous frame buffer.
@@ -1956,6 +1958,8 @@ def from_cframe(
         return blosc2.BatchStore(_from_schunk=schunk_from_cframe(cframe, copy=copy))
     if "vlarray" in schunk.meta:
         return blosc2.vlarray_from_cframe(cframe, copy=copy)
+    if "b2o" in schunk.meta:
+        return blosc2._open_b2object(ndarray_from_cframe(cframe, copy=copy))
     if "b2nd" in schunk.meta:
         return ndarray_from_cframe(cframe, copy=copy)
     return schunk_from_cframe(cframe, copy=copy)