Support for serializing LazyArray as well

Author: FrancescAlted

src/blosc2/_msgpack_utils.py

@@ -23,24 +23,54 @@
 _BLOSC2_STRUCTURED_VERSION = 1
 
 
-def _encode_structured_reference(obj):
+def _encode_operand_reference(obj):
     import blosc2
 
     if isinstance(obj, blosc2.C2Array):
-        payload = {
+        return {
             "kind": "c2array",
             "version": _BLOSC2_STRUCTURED_VERSION,
             "path": obj.path,
             "urlbase": obj.urlbase,
         }
+    if isinstance(obj, blosc2.Proxy):
+        obj = obj._cache
+    if hasattr(obj, "schunk"):
+        urlpath = obj.schunk.urlpath
+        if urlpath is None:
+            raise ValueError(
+                "Structured Blosc2 msgpack payload requires operands to be stored on disk/network"
+            )
+        return {
+            "kind": "urlpath",
+            "version": _BLOSC2_STRUCTURED_VERSION,
+            "urlpath": urlpath,
+        }
+    raise TypeError("Structured Blosc2 msgpack payload requires NDArray, C2Array, or Proxy operands")
+
+
+def _encode_structured_reference(obj):
+    import blosc2
+
+    if isinstance(obj, blosc2.C2Array):
+        payload = _encode_operand_reference(obj)
+        return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
+    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 = {
+            "kind": "lazyexpr",
+            "version": _BLOSC2_STRUCTURED_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))
     return None
 
 
-def _decode_structured_reference(data):
+def _decode_operand_reference(payload):
     import blosc2
 
-    payload = unpackb(data)
     if not isinstance(payload, dict):
         raise TypeError("Structured Blosc2 msgpack payload must decode to a mapping")
 
@@ -57,6 +87,37 @@ def _decode_structured_reference(data):
         if urlbase is not None and not isinstance(urlbase, str):
             raise TypeError("Structured C2Array msgpack payload requires 'urlbase' to be a string or None")
         return blosc2.C2Array(path, urlbase=urlbase)
+    if kind == "urlpath":
+        urlpath = payload.get("urlpath")
+        if not isinstance(urlpath, str):
+            raise TypeError("Structured urlpath msgpack payload requires a string 'urlpath'")
+        return blosc2.open(urlpath, mode="r")
+    raise ValueError(f"Unsupported structured Blosc2 msgpack payload operand kind: {kind!r}")
+
+
+def _decode_structured_reference(data):
+    import blosc2
+
+    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}")
+
+    kind = payload.get("kind")
+    if kind == "c2array":
+        return _decode_operand_reference(payload)
+    if kind == "lazyexpr":
+        expression = payload.get("expression")
+        if not isinstance(expression, str):
+            raise TypeError("Structured LazyExpr msgpack 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'")
+        operands = {key: _decode_operand_reference(value) for key, value in operands_payload.items()}
+        return blosc2.lazyexpr(expression, operands=operands)
     raise ValueError(f"Unsupported structured Blosc2 msgpack payload kind: {kind!r}")
 
 

tests/test_batch_store.py

@@ -55,6 +55,20 @@ def fake_info(path_, urlbase_, params=None, headers=None, model=None, auth_token
     return blosc2.C2Array(path, urlbase=urlbase)
 
 
+def _make_persistent_lazyexpr(tmp_path):
+    a = blosc2.asarray(np.arange(5, dtype=np.int64), urlpath=tmp_path / "a.b2nd", mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.int64) * 2, urlpath=tmp_path / "b.b2nd", mode="w")
+    expr = blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
+    expected = np.arange(5, dtype=np.int64) * 3
+    return expr, expected
+
+
+def _make_in_memory_lazyexpr():
+    a = blosc2.asarray(np.arange(5, dtype=np.int64))
+    b = blosc2.asarray(np.arange(5, dtype=np.int64) * 2)
+    return blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
+
+
 @pytest.mark.parametrize(
     ("contiguous", "urlpath"),
     [
@@ -260,6 +274,96 @@ def test_batchstore_msgpack_supports_c2array(monkeypatch):
     assert restored.auth_token is None
 
 
+def test_msgpack_supports_lazyexpr(tmp_path):
+    expr, expected = _make_persistent_lazyexpr(tmp_path)
+
+    payload = msgpack_packb({"expr": expr})
+    restored = msgpack_unpackb(payload)["expr"]
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], expected)
+
+
+def test_batchstore_msgpack_supports_lazyexpr(tmp_path):
+    expr, expected = _make_persistent_lazyexpr(tmp_path)
+
+    barray = blosc2.BatchStore(items_per_block=2)
+    barray.append([expr])
+
+    restored = barray[0][0]
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], expected)
+
+
+def test_msgpack_rejects_lazyexpr_with_in_memory_operands():
+    expr = _make_in_memory_lazyexpr()
+
+    with pytest.raises(ValueError, match="stored on disk/network"):
+        msgpack_packb({"expr": expr})
+
+
+def test_batchstore_msgpack_rejects_lazyexpr_with_in_memory_operands():
+    expr = _make_in_memory_lazyexpr()
+
+    barray = blosc2.BatchStore(items_per_block=2)
+    with pytest.raises(ValueError, match="stored on disk/network"):
+        barray.append([expr])
+
+
+@pytest.mark.network
+def test_msgpack_supports_lazyexpr_with_c2array_operand(cat2_context, tmp_path):
+    path = "@public/expr/ds-1-2-linspace-float64-b2-(5,)d.b2nd"
+    a = blosc2.C2Array(path)
+    a_values = np.asarray(a[:])
+    b = blosc2.asarray(a_values * 2, urlpath=tmp_path / "b.b2nd", mode="w")
+    expr = blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
+
+    restored = msgpack_unpackb(msgpack_packb({"expr": expr}))["expr"]
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_allclose(restored[:], a_values + b[:])
+
+
+@pytest.mark.network
+def test_batchstore_msgpack_supports_lazyexpr_with_c2array_operand(cat2_context, tmp_path):
+    path = "@public/expr/ds-1-2-linspace-float64-b2-(5,)d.b2nd"
+    a = blosc2.C2Array(path)
+    a_values = np.asarray(a[:])
+    b = blosc2.asarray(a_values * 2, urlpath=tmp_path / "b.b2nd", mode="w")
+    expr = blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
+
+    barray = blosc2.BatchStore(items_per_block=2)
+    barray.append([expr])
+    restored = barray[0][0]
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_allclose(restored[:], a_values + b[:])
+
+
+@pytest.mark.xfail(
+    strict=True, reason="Structured LazyExpr urlpath operands do not preserve .b2z member offsets"
+)
+def test_msgpack_lazyexpr_with_b2z_operands(tmp_path):
+    store_path = tmp_path / "operands.b2z"
+    ext_a = tmp_path / "a.b2nd"
+    ext_b = tmp_path / "b.b2nd"
+    expected = np.arange(5, dtype=np.int64) * 3
+
+    a = blosc2.asarray(np.arange(5, dtype=np.int64), urlpath=str(ext_a), mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.int64) * 2, urlpath=str(ext_b), mode="w")
+    with blosc2.DictStore(str(store_path), mode="w", threshold=None) as dstore:
+        dstore["/a"] = a
+        dstore["/b"] = b
+
+    with blosc2.DictStore(str(store_path), mode="r") as dstore:
+        expr = blosc2.lazyexpr("a + b", operands={"a": dstore["/a"], "b": dstore["/b"]})
+        restored = msgpack_unpackb(msgpack_packb({"expr": expr}))["expr"]
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], expected)
+
+
 @pytest.mark.network
 def test_msgpack_roundtrip_c2array_network(cat2_context):
     path = "@public/expr/ds-1-2-linspace-float64-b2-(5,)d.b2nd"

tests/test_vlarray.py

@@ -54,6 +54,20 @@ def fake_info(path_, urlbase_, params=None, headers=None, model=None, auth_token
     return blosc2.C2Array(path, urlbase=urlbase)
 
 
+def _make_persistent_lazyexpr(tmp_path):
+    a = blosc2.asarray(np.arange(5, dtype=np.int64), urlpath=tmp_path / "a.b2nd", mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.int64) * 2, urlpath=tmp_path / "b.b2nd", mode="w")
+    expr = blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
+    expected = np.arange(5, dtype=np.int64) * 3
+    return expr, expected
+
+
+def _make_in_memory_lazyexpr():
+    a = blosc2.asarray(np.arange(5, dtype=np.int64))
+    b = blosc2.asarray(np.arange(5, dtype=np.int64) * 2)
+    return blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
+
+
 @pytest.mark.parametrize(
     ("contiguous", "urlpath"),
     [
@@ -192,6 +206,26 @@ def test_vlarray_msgpack_supports_c2array(monkeypatch):
     assert restored.auth_token is None
 
 
+def test_vlarray_msgpack_supports_lazyexpr(tmp_path):
+    expr, expected = _make_persistent_lazyexpr(tmp_path)
+
+    vlarray = blosc2.VLArray()
+    vlarray.append(expr)
+
+    restored = vlarray[0]
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], expected)
+
+
+def test_vlarray_msgpack_rejects_lazyexpr_with_in_memory_operands():
+    expr = _make_in_memory_lazyexpr()
+
+    vlarray = blosc2.VLArray()
+    with pytest.raises(ValueError, match="stored on disk/network"):
+        vlarray.append(expr)
+
+
 @pytest.mark.network
 def test_vlarray_msgpack_roundtrip_c2array_network(cat2_context):
     path = "@public/expr/ds-1-2-linspace-float64-b2-(5,)d.b2nd"