Add b2o carrier support for persisted LazyExpr objects

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -42,6 +42,7 @@
 
 import blosc2
 
+from .b2objects import _encode_b2object_payload, _make_b2object_carrier, _write_b2object_payload
 from .dsl_kernel import DSLKernel, DSLSyntaxError, _DSLValidator, specialize_miniexpr_inputs
 
 if blosc2._HAS_NUMBA:
@@ -570,7 +571,7 @@ def convert_inputs(inputs):
         return []
     inputs_ = []
     for obj in inputs:
-        if not isinstance(obj, (np.ndarray, blosc2.Operand)) and not np.isscalar(obj):
+        if not isinstance(obj, np.ndarray | blosc2.Operand) and not np.isscalar(obj):
             try:
                 obj = blosc2.SimpleProxy(obj)
             except Exception:
@@ -2855,7 +2856,7 @@ def result_type(
     # Follow NumPy rules for scalar-array operations
     # Create small arrays with the same dtypes and let NumPy's type promotion determine the result type
     arrs = [
-        (np.array(value).dtype if isinstance(value, (str, bytes)) else value)
+        (np.array(value).dtype if isinstance(value, str | bytes) else value)
         if (np.isscalar(value) or not hasattr(value, "dtype"))
         else np.array([0], dtype=_convert_dtype(value.dtype))
         for value in arrays_and_dtypes
@@ -2902,7 +2903,7 @@ def __init__(self, new_op):  # noqa: C901
         # Check that operands are proper Operands, LazyArray or scalars; if not, convert to NDArray objects
         value1 = (
             blosc2.SimpleProxy(value1)
-            if not (isinstance(value1, (blosc2.Operand, np.ndarray)) or np.isscalar(value1))
+            if not (isinstance(value1, blosc2.Operand | np.ndarray) or np.isscalar(value1))
             else value1
         )
         # Reset values represented as np.int64 etc. to be set as Python natives
@@ -2926,7 +2927,7 @@ def __init__(self, new_op):  # noqa: C901
             return
         value2 = (
             blosc2.SimpleProxy(value2)
-            if not (isinstance(value2, (blosc2.Operand, np.ndarray)) or np.isscalar(value2))
+            if not (isinstance(value2, blosc2.Operand | np.ndarray) or np.isscalar(value2))
             else value2
         )
         # Reset values represented as np.int64 etc. to be set as Python natives
@@ -3711,44 +3712,27 @@ def save(self, urlpath=None, **kwargs):
         if urlpath is None:
             raise ValueError("To save a LazyArray you must provide an urlpath")
 
-        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 expression
-        validate_expr(expression)
-
-        meta = kwargs.get("meta", {})
-        meta["LazyArray"] = LazyArrayEnum.Expr.value
         kwargs["urlpath"] = urlpath
-        kwargs["meta"] = meta
         kwargs["mode"] = "w"  # always overwrite the file in urlpath
+        self._to_b2object_carrier(**kwargs)
 
-        # Create an empty array; useful for providing the shape and dtype of the outcome
-        array = blosc2.empty(shape=self.shape, dtype=self.dtype, **kwargs)
-
-        # Save the expression and operands in the metadata
-        operands = {}
-        for key, value in operands_.items():
-            if isinstance(value, blosc2.C2Array):
-                operands[key] = {
-                    "path": str(value.path),
-                    "urlbase": value.urlbase,
-                }
-                continue
-            if isinstance(value, blosc2.Proxy):
-                # Take the required info from the Proxy._cache container
-                value = value._cache
-            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")
-            operands[key] = value.schunk.urlpath
-        array.schunk.vlmeta["_LazyArray"] = {
-            "expression": expression,
-            "UDF": None,
-            "operands": operands,
-        }
+    def to_cframe(self) -> bytes:
+        return self._to_b2object_carrier().to_cframe()
+
+    def _to_b2object_carrier(self, **kwargs):
+        payload = _encode_b2object_payload(self)
+        if payload is None:
+            raise TypeError("Unsupported persisted Blosc2 object")
+        array = _make_b2object_carrier(
+            "lazyexpr",
+            self.shape,
+            self.dtype,
+            chunks=self.chunks,
+            blocks=self.blocks,
+            **kwargs,
+        )
+        _write_b2object_payload(array, payload)
+        return array
 
     @classmethod
     def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=None):
@@ -3764,7 +3748,7 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
             _operands = operands | local_vars
             # Check that operands are proper Operands, LazyArray or scalars; if not, convert to NDArray objects
             for op, val in _operands.items():
-                if not (isinstance(val, (blosc2.Operand, np.ndarray)) or np.isscalar(val)):
+                if not (isinstance(val, blosc2.Operand | np.ndarray) or np.isscalar(val)):
                     _operands[op] = blosc2.SimpleProxy(val)
             # for scalars just return value (internally converts to () if necessary)
             opshapes = {k: v if not hasattr(v, "shape") else v.shape for k, v in _operands.items()}

tests/data/legacy_lazyexpr_v1/a.b2nd

@@ -7,6 +7,8 @@
 
 from __future__ import annotations
 
+from pathlib import Path
+
 import numpy as np
 
 import blosc2
@@ -69,7 +71,7 @@ def test_c2array_from_cframe_roundtrip(monkeypatch):
 
 def test_c2array_open_roundtrip(tmp_path, monkeypatch):
     original = _make_c2array(monkeypatch, shape=(8,), chunks=(4,), blocks=(2,))
-    urlpath = tmp_path / "remote-array.b2o"
+    urlpath = tmp_path / "remote-array.b2nd"
 
     original.save(urlpath)
     restored = blosc2.open(urlpath, mode="r")
@@ -81,3 +83,48 @@ def test_c2array_open_roundtrip(tmp_path, monkeypatch):
     assert restored.shape == original.shape
     assert restored.chunks == original.chunks
     assert restored.blocks == original.blocks
+
+
+def test_lazyexpr_from_cframe_roundtrip(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})
+    carrier = blosc2.ndarray_from_cframe(expr.to_cframe())
+
+    assert carrier.schunk.meta["b2o"] == {"kind": "lazyexpr", "version": 1}
+    assert carrier.schunk.vlmeta["b2o"] == {
+        "kind": "lazyexpr",
+        "version": 1,
+        "expression": "a + b",
+        "operands": {
+            "a": {"kind": "urlpath", "version": 1, "urlpath": str(tmp_path / "a.b2nd")},
+            "b": {"kind": "urlpath", "version": 1, "urlpath": str(tmp_path / "b.b2nd")},
+        },
+    }
+
+    restored = blosc2.from_cframe(expr.to_cframe())
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], np.arange(5, dtype=np.int64) * 3)
+
+
+def test_lazyexpr_open_roundtrip(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})
+    urlpath = tmp_path / "expr.b2nd"
+
+    expr.save(urlpath)
+    restored = blosc2.open(urlpath, mode="r")
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], np.arange(5, dtype=np.int64) * 3)
+
+
+def test_legacy_lazyexpr_open_backward_compat():
+    fixture = Path(__file__).parent / "data" / "legacy_lazyexpr_v1" / "expr.b2nd"
+
+    restored = blosc2.open(fixture, mode="r")
+
+    assert isinstance(restored, blosc2.LazyExpr)
+    np.testing.assert_array_equal(restored[:], np.arange(5, dtype=np.int64) * 3)