Preserve legacy LazyExpr persistence behavior and support .b2z operands

Author: FrancescAlted

src/blosc2/b2objects.py

@@ -10,6 +10,7 @@
 import builtins
 import inspect
 import linecache
+import pathlib
 import textwrap
 from dataclasses import asdict
 from typing import Any
@@ -47,8 +48,11 @@ def encode_operand_reference(obj):
     return blosc2.Ref.from_object(obj).to_dict()
 
 
-def decode_operand_reference(payload):
-    return blosc2.Ref.from_dict(payload).open()
+def decode_operand_reference(payload, *, base_path=None):
+    ref = blosc2.Ref.from_dict(payload)
+    if ref.kind == "urlpath" and base_path is not None and not pathlib.Path(ref.urlpath).is_absolute():
+        return blosc2.open(base_path / ref.urlpath, mode="r")
+    return ref.open()
 
 
 def encode_b2object_payload(obj) -> dict[str, Any] | None:
@@ -98,7 +102,7 @@ def encode_b2object_payload(obj) -> dict[str, Any] | None:
     return None
 
 
-def decode_b2object_payload(payload: dict[str, Any]):
+def decode_b2object_payload(payload: dict[str, Any], *, carrier_path=None):
     kind = payload.get("kind")
     version = payload.get("version")
     if version != _B2OBJECT_VERSION:
@@ -107,24 +111,39 @@ def decode_b2object_payload(payload: dict[str, Any]):
         ref = blosc2.Ref.from_dict(payload)
         return ref.open()
     if kind == "lazyexpr":
-        return decode_structured_lazyexpr(payload)
+        return decode_structured_lazyexpr(payload, carrier_path=carrier_path)
     if kind == "lazyudf":
-        return decode_structured_lazyudf(payload)
+        return decode_structured_lazyudf(payload, carrier_path=carrier_path)
     raise ValueError(f"Unsupported persisted Blosc2 object kind: {kind!r}")
 
 
-def decode_structured_lazyexpr(payload):
+def decode_structured_lazyexpr(payload, *, carrier_path=None):
     expression = payload.get("expression")
     if not isinstance(expression, str):
         raise TypeError("Structured LazyExpr payload requires a string 'expression'")
     operands_payload = payload.get("operands")
     if not isinstance(operands_payload, dict):
         raise TypeError("Structured LazyExpr payload requires a mapping 'operands'")
-    operands = {key: decode_operand_reference(value) for key, value in operands_payload.items()}
+    operands = {}
+    missing_ops = {}
+    for key, value in operands_payload.items():
+        try:
+            operands[key] = decode_operand_reference(value, base_path=carrier_path)
+        except FileNotFoundError:
+            ref = blosc2.Ref.from_dict(value)
+            if ref.kind == "urlpath":
+                missing_ops[key] = pathlib.Path(ref.urlpath)
+            else:
+                raise
+    if missing_ops:
+        exc = blosc2.exceptions.MissingOperands(expression, missing_ops)
+        exc.expr = expression
+        exc.missing_ops = missing_ops
+        raise exc
     return blosc2.lazyexpr(expression, operands=operands)
 
 
-def decode_structured_lazyudf(payload):
+def decode_structured_lazyudf(payload, *, carrier_path=None):
     function_kind = payload.get("function_kind")
     if function_kind != "dsl":
         raise ValueError(f"Unsupported structured LazyUDF function kind: {function_kind!r}")
@@ -167,7 +186,8 @@ def decode_structured_lazyudf(payload):
         func.dsl_source = dsl_source
 
     operands = tuple(
-        decode_operand_reference(operands_payload[f"o{n}"]) for n in range(len(operands_payload))
+        decode_operand_reference(operands_payload[f"o{n}"], base_path=carrier_path)
+        for n in range(len(operands_payload))
     )
     return blosc2.lazyudf(func, operands, dtype=np.dtype(dtype), shape=tuple(shape_payload), **kwargs)
 
@@ -194,4 +214,12 @@ def open_b2object(obj):
         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)
+    carrier_path = None
+    schunk = getattr(obj, "schunk", obj)
+    if getattr(schunk, "urlpath", None) is not None:
+        carrier_path = pathlib.Path(schunk.urlpath).parent
+    opened = decode_b2object_payload(payload, carrier_path=carrier_path)
+    if isinstance(opened, blosc2.LazyExpr | blosc2.LazyUDF):
+        opened.array = obj
+        opened.schunk = schunk
+    return opened

src/blosc2/lazyexpr.py

@@ -3718,12 +3718,46 @@ def save(self, urlpath=None, **kwargs):
         self._to_b2object_carrier(**kwargs)
 
     def to_cframe(self) -> bytes:
-        return self._to_b2object_carrier().to_cframe()
+        try:
+            return self._to_b2object_carrier().to_cframe()
+        except (TypeError, ValueError):
+            return self.compute().to_cframe()
 
     def _to_b2object_carrier(self, **kwargs):
-        payload = encode_b2object_payload(self)
-        if payload is None:
-            raise TypeError("Unsupported persisted Blosc2 object")
+        expression = self.expression_tosave if hasattr(self, "expression_tosave") else self.expression
+        operands_ = self.operands_tosave if hasattr(self, "operands_tosave") else self.operands
+        validate_expr(expression)
+
+        payload = {"kind": "lazyexpr", "version": 1, "expression": expression, "operands": {}}
+        carrier_urlpath = kwargs.get("urlpath")
+        carrier_parent = Path(carrier_urlpath).parent if carrier_urlpath is not None else None
+        for key, value in operands_.items():
+            if isinstance(value, blosc2.C2Array):
+                payload["operands"][key] = encode_b2object_payload(value)
+                continue
+            if isinstance(value, blosc2.Proxy):
+                value = value._cache
+            ref = getattr(value, "_blosc2_ref", None)
+            if isinstance(ref, blosc2.Ref):
+                payload["operands"][key] = ref.to_dict()
+                continue
+            if not hasattr(value, "schunk"):
+                raise ValueError(
+                    "To save a LazyArray, all operands must be blosc2.NDArray or blosc2.C2Array objects"
+                )
+            if value.schunk.urlpath is None:
+                raise ValueError("To save a LazyArray, all operands must be stored on disk/network")
+            operand_urlpath = Path(value.schunk.urlpath)
+            if carrier_parent is not None and not operand_urlpath.is_absolute():
+                ref_urlpath = operand_urlpath.as_posix()
+            elif carrier_parent is not None:
+                try:
+                    ref_urlpath = operand_urlpath.relative_to(carrier_parent).as_posix()
+                except ValueError:
+                    ref_urlpath = operand_urlpath.as_posix()
+            else:
+                ref_urlpath = operand_urlpath.as_posix()
+            payload["operands"][key] = {"kind": "urlpath", "version": 1, "urlpath": ref_urlpath}
         array = make_b2object_carrier(
             "lazyexpr",
             self.shape,

tests/ndarray/test_lazyexpr.py

@@ -1721,6 +1721,37 @@ def test_missing_operator():
     blosc2.remove_urlpath("expr.b2nd")
 
 
+def test_save_dictstore_operands(tmp_path):
+    store_path = tmp_path / "operands.b2z"
+    ext_a = tmp_path / "a.b2nd"
+    ext_b = tmp_path / "b.b2nd"
+    expr_path = tmp_path / "expr.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:
+        a = dstore["/a"]
+        b = dstore["/b"]
+        expr = blosc2.lazyexpr("a + b")
+        expr.save(expr_path)
+
+    carrier = blosc2.open(expr_path, mode="r").array
+    assert carrier.schunk.vlmeta["b2o"]["operands"] == {
+        "a": {"kind": "dictstore_key", "version": 1, "urlpath": str(store_path), "key": "/a"},
+        "b": {"kind": "dictstore_key", "version": 1, "urlpath": str(store_path), "key": "/b"},
+    }
+
+    restored = blosc2.open(expr_path)
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], expected)
+
+
 # Test the chaining of multiple lazy expressions
 def test_chain_expressions():
     N = 1_000