Initial support for Blosc2 cframes serialized with msgpack

Author: FrancescAlted

src/blosc2/_msgpack_utils.py

@@ -7,20 +7,51 @@
 
 from __future__ import annotations
 
-from msgpack import packb, unpackb
+from msgpack import ExtType, packb, unpackb
 
 from blosc2 import blosc2_ext
 
+# 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
+
+
+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())
+    return blosc2_ext.encode_tuple(obj)
+
 
 def msgpack_packb(value):
-    return packb(value, default=blosc2_ext.encode_tuple, strict_types=True, use_bin_type=True)
+    return packb(value, default=_encode_msgpack_ext, strict_types=True, use_bin_type=True)
 
 
 def decode_tuple_list_hook(obj):
-    if obj and obj[0] == "__tuple__":
+    if obj and isinstance(obj[0], str) and obj[0] == "__tuple__":
         return tuple(obj[1:])
     return obj
 
 
+def _decode_msgpack_ext(code, data):
+    import blosc2
+
+    if code == _BLOSC2_EXT_CODE:
+        return blosc2.from_cframe(data, copy=True)
+    return ExtType(code, data)
+
+
 def msgpack_unpackb(payload):
-    return unpackb(payload, list_hook=decode_tuple_list_hook)
+    return unpackb(payload, list_hook=decode_tuple_list_hook, ext_hook=_decode_msgpack_ext)

src/blosc2/batch_store.py

@@ -166,8 +166,12 @@ class BatchStore:
         If not provided, a value is inferred from the first batch.
     serializer : {"msgpack", "arrow"}, optional
         Serializer used for batch payloads. ``"msgpack"`` is the default and is
-        the general-purpose choice for Python items. ``"arrow"`` is optional and
-        requires ``pyarrow``.
+        the general-purpose choice for Python items, including nested Blosc2
+        containers such as :class:`blosc2.NDArray`, :class:`blosc2.SChunk`,
+        :class:`blosc2.VLArray`, :class:`blosc2.BatchStore`, and
+        :class:`blosc2.EmbedStore`, which are serialized transparently via
+        :meth:`to_cframe` / :func:`blosc2.from_cframe`. ``"arrow"`` is optional
+        and requires ``pyarrow``.
     _from_schunk : blosc2.SChunk, optional
         Internal hook used when reopening an already-tagged BatchStore.
     **kwargs

src/blosc2/vlarray.py

@@ -29,7 +29,14 @@ def _check_serialized_size(buffer: bytes) -> None:
 
 
 class VLArray:
-    """A variable-length array backed by an :class:`blosc2.SChunk`."""
+    """A variable-length array backed by an :class:`blosc2.SChunk`.
+
+    Entries are serialized with msgpack before compression. Standard Python
+    objects are supported, and Blosc2 containers such as
+    :class:`blosc2.NDArray`, :class:`blosc2.SChunk`, :class:`blosc2.VLArray`,
+    :class:`blosc2.BatchStore`, and :class:`blosc2.EmbedStore` are serialized
+    transparently via :meth:`to_cframe` / :func:`blosc2.from_cframe`.
+    """
 
     @staticmethod
     def _set_typesize_one(cparams: blosc2.CParams | dict | None) -> blosc2.CParams | dict:

tests/test_batch_store.py

@@ -5,6 +5,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #######################################################################
 
+import numpy as np
 import pytest
 
 import blosc2
@@ -27,6 +28,24 @@ def _storage(contiguous, urlpath, mode="w"):
     return blosc2.Storage(contiguous=contiguous, urlpath=urlpath, mode=mode)
 
 
+def _make_nested_blosc2_objects():
+    ndarray = blosc2.arange(6, dtype=np.int32)
+
+    schunk = blosc2.SChunk(chunksize=16)
+    schunk.append_data(np.arange(4, dtype=np.int32))
+
+    nested_vlarray = blosc2.VLArray()
+    nested_vlarray.extend(["alpha", {"beta": 2}])
+
+    nested_batchstore = blosc2.BatchStore(max_blocksize=2)
+    nested_batchstore.extend([[1, 2], ["x", {"y": 3}]])
+
+    estore = blosc2.EmbedStore()
+    estore["/node"] = blosc2.arange(3, dtype=np.int32)
+
+    return ndarray, schunk, nested_vlarray, nested_batchstore, estore
+
+
 @pytest.mark.parametrize(
     ("contiguous", "urlpath"),
     [
@@ -181,6 +200,31 @@ def test_batchstore_from_cframe():
     assert [batch[:] for batch in restored2] == expected
 
 
+def test_batchstore_msgpack_supports_blosc2_objects():
+    ndarray, schunk, nested_vlarray, nested_batchstore, estore = _make_nested_blosc2_objects()
+
+    barray = blosc2.BatchStore(max_blocksize=2)
+    barray.append([ndarray, schunk, nested_vlarray, nested_batchstore, estore])
+
+    restored = barray[0][:]
+
+    assert isinstance(restored[0], blosc2.NDArray)
+    assert np.array_equal(restored[0][:], ndarray[:])
+
+    assert isinstance(restored[1], blosc2.SChunk)
+    assert restored[1].decompress_chunk(0) == schunk.decompress_chunk(0)
+
+    assert isinstance(restored[2], blosc2.VLArray)
+    assert list(restored[2]) == list(nested_vlarray)
+
+    assert isinstance(restored[3], blosc2.BatchStore)
+    assert [batch[:] for batch in restored[3]] == [batch[:] for batch in nested_batchstore]
+
+    assert isinstance(restored[4], blosc2.EmbedStore)
+    assert list(restored[4].keys()) == ["/node"]
+    assert np.array_equal(restored[4]["/node"][:], estore["/node"][:])
+
+
 def test_batchstore_info():
     barray = blosc2.BatchStore()
     barray.extend(BATCHES)

tests/test_vlarray.py

@@ -5,6 +5,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #######################################################################
 
+import numpy as np
 import pytest
 
 import blosc2
@@ -26,6 +27,24 @@ def _storage(contiguous, urlpath, mode="w"):
     return blosc2.Storage(contiguous=contiguous, urlpath=urlpath, mode=mode)
 
 
+def _make_nested_blosc2_objects():
+    ndarray = blosc2.arange(6, dtype=np.int32)
+
+    schunk = blosc2.SChunk(chunksize=16)
+    schunk.append_data(np.arange(4, dtype=np.int32))
+
+    nested_vlarray = blosc2.VLArray()
+    nested_vlarray.extend(["alpha", {"beta": 2}])
+
+    nested_batchstore = blosc2.BatchStore(max_blocksize=2)
+    nested_batchstore.extend([[1, 2], ["x", {"y": 3}]])
+
+    estore = blosc2.EmbedStore()
+    estore["/node"] = blosc2.arange(3, dtype=np.int32)
+
+    return ndarray, schunk, nested_vlarray, nested_batchstore, estore
+
+
 @pytest.mark.parametrize(
     ("contiguous", "urlpath"),
     [
@@ -117,6 +136,39 @@ def test_vlarray_from_cframe():
     assert list(restored2) == expected
 
 
+def test_vlarray_msgpack_supports_blosc2_objects():
+    ndarray, schunk, nested_vlarray, nested_batchstore, estore = _make_nested_blosc2_objects()
+
+    vlarray = blosc2.VLArray()
+    vlarray.append(
+        {
+            "ndarray": ndarray,
+            "schunk": schunk,
+            "vlarray": nested_vlarray,
+            "batchstore": nested_batchstore,
+            "estore": estore,
+        }
+    )
+
+    restored = vlarray[0]
+
+    assert isinstance(restored["ndarray"], blosc2.NDArray)
+    assert np.array_equal(restored["ndarray"][:], ndarray[:])
+
+    assert isinstance(restored["schunk"], blosc2.SChunk)
+    assert restored["schunk"].decompress_chunk(0) == schunk.decompress_chunk(0)
+
+    assert isinstance(restored["vlarray"], blosc2.VLArray)
+    assert list(restored["vlarray"]) == list(nested_vlarray)
+
+    assert isinstance(restored["batchstore"], blosc2.BatchStore)
+    assert [batch[:] for batch in restored["batchstore"]] == [batch[:] for batch in nested_batchstore]
+
+    assert isinstance(restored["estore"], blosc2.EmbedStore)
+    assert list(restored["estore"].keys()) == ["/node"]
+    assert np.array_equal(restored["estore"]["/node"][:], estore["/node"][:])
+
+
 def test_vlarray_info():
     vlarray = blosc2.VLArray()
     vlarray.extend(VALUES)