Add b2o-backed persistence for DSL LazyUDF objects

Author: FrancescAlted

src/blosc2/b2objects.py

@@ -49,9 +49,14 @@ def encode_operand_reference(obj):
 
 
 def decode_operand_reference(payload, *, base_path=None):
+    if (
+        payload.get("kind") in {"urlpath", "dictstore_key"}
+        and base_path is not None
+        and not pathlib.Path(payload["urlpath"]).is_absolute()
+    ):
+        payload = dict(payload)
+        payload["urlpath"] = (base_path / payload["urlpath"]).as_posix()
     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()
 
 
@@ -123,23 +128,28 @@ def decode_structured_lazyexpr(payload, *, carrier_path=None):
     operands_payload = payload.get("operands")
     if not isinstance(operands_payload, dict):
         raise TypeError("Structured LazyExpr payload requires a mapping 'operands'")
+    operands, missing_ops = decode_operand_mapping(operands_payload, base_path=carrier_path)
+    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_operand_mapping(operands_payload, *, base_path=None):
     operands = {}
     missing_ops = {}
     for key, value in operands_payload.items():
         try:
-            operands[key] = decode_operand_reference(value, base_path=carrier_path)
+            operands[key] = decode_operand_reference(value, base_path=base_path)
         except FileNotFoundError:
             ref = blosc2.Ref.from_dict(value)
-            if ref.kind == "urlpath":
+            if ref.kind in {"urlpath", "dictstore_key"}:
                 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)
+    return operands, missing_ops
 
 
 def decode_structured_lazyudf(payload, *, carrier_path=None):
@@ -180,12 +190,16 @@ def decode_structured_lazyudf(payload, *, carrier_path=None):
     func = local_ns[name]
     if not isinstance(func, DSLKernel):
         func = DSLKernel(func)
-
-    operands = tuple(
-        decode_operand_reference(operands_payload[f"o{n}"], base_path=carrier_path)
-        for n in range(len(operands_payload))
+    ordered_operands_payload = {f"o{n}": operands_payload[f"o{n}"] for n in range(len(operands_payload))}
+    operands, missing_ops = decode_operand_mapping(ordered_operands_payload, base_path=carrier_path)
+    if missing_ops:
+        exc = blosc2.exceptions.MissingOperands(name, missing_ops)
+        exc.expr = name
+        exc.missing_ops = missing_ops
+        raise exc
+    return blosc2.lazyudf(
+        func, tuple(operands.values()), dtype=np.dtype(dtype), shape=tuple(shape_payload), **kwargs
     )
-    return blosc2.lazyudf(func, operands, dtype=np.dtype(dtype), shape=tuple(shape_payload), **kwargs)
 
 
 def read_b2object_marker(obj) -> dict[str, Any] | None:

src/blosc2/lazyexpr.py

@@ -3718,10 +3718,7 @@ def save(self, urlpath=None, **kwargs):
         self._to_b2object_carrier(**kwargs)
 
     def to_cframe(self) -> bytes:
-        try:
-            return self._to_b2object_carrier().to_cframe()
-        except (TypeError, ValueError):
-            return self.compute().to_cframe()
+        return self._to_b2object_carrier().to_cframe()
 
     def _to_b2object_carrier(self, **kwargs):
         expression = self.expression_tosave if hasattr(self, "expression_tosave") else self.expression
@@ -4079,50 +4076,91 @@ def save(self, urlpath=None, **kwargs):
         if urlpath is None:
             raise ValueError("To save a LazyArray you must provide an urlpath")
 
-        meta = kwargs.get("meta", {})
-        meta["LazyArray"] = LazyArrayEnum.UDF.value
         kwargs["urlpath"] = urlpath
-        kwargs["meta"] = meta
         kwargs["mode"] = "w"  # always overwrite the file in urlpath
+        try:
+            self._to_b2object_carrier(**kwargs)
+        except (TypeError, ValueError):
+            meta = kwargs.get("meta", {})
+            meta["LazyArray"] = LazyArrayEnum.UDF.value
+            kwargs["meta"] = meta
 
-        # Create an empty array; useful for providing the shape and dtype of the outcome
-        array = blosc2.empty(shape=self.shape, dtype=self.dtype, **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 = {}
-        operands_ = self.inputs_dict
-        for i, (_key, value) in enumerate(operands_.items()):
-            pos_key = f"o{i}"  # always use positional keys for consistent loading
-            if isinstance(value, blosc2.C2Array):
-                operands[pos_key] = {
-                    "path": str(value.path),
-                    "urlbase": value.urlbase,
-                }
+            # Save the expression and operands in the metadata
+            operands = {}
+            operands_ = self.inputs_dict
+            for i, (_key, value) in enumerate(operands_.items()):
+                pos_key = f"o{i}"  # always use positional keys for consistent loading
+                if isinstance(value, blosc2.C2Array):
+                    operands[pos_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"
+                    ) from None
+                if value.schunk.urlpath is None:
+                    raise ValueError(
+                        "To save a LazyArray, all operands must be stored on disk/network"
+                    ) from None
+                operands[pos_key] = value.schunk.urlpath
+            udf_func = self.func.func if isinstance(self.func, DSLKernel) else self.func
+            udf_name = getattr(udf_func, "__name__", self.func.__name__)
+            try:
+                udf_source = textwrap.dedent(inspect.getsource(udf_func)).lstrip()
+            except Exception:
+                udf_source = None
+            meta = {
+                "UDF": udf_source,
+                "operands": operands,
+                "name": udf_name,
+            }
+            if isinstance(self.func, DSLKernel) and self.func.dsl_source is not None:
+                meta["dsl_source"] = self.func.dsl_source
+            array.schunk.vlmeta["_LazyArray"] = meta
+
+    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("Persistent Blosc2 object payload is not supported for this LazyUDF")
+
+        carrier_urlpath = kwargs.get("urlpath")
+        carrier_parent = Path(carrier_urlpath).parent if carrier_urlpath is not None else None
+        for operand_payload in payload["operands"].values():
+            if operand_payload["kind"] not in {"urlpath", "dictstore_key"}:
                 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[pos_key] = value.schunk.urlpath
-        udf_func = self.func.func if isinstance(self.func, DSLKernel) else self.func
-        udf_name = getattr(udf_func, "__name__", self.func.__name__)
-        try:
-            udf_source = textwrap.dedent(inspect.getsource(udf_func)).lstrip()
-        except Exception:
-            udf_source = None
-        meta = {
-            "UDF": udf_source,
-            "operands": operands,
-            "name": udf_name,
-        }
-        if isinstance(self.func, DSLKernel) and self.func.dsl_source is not None:
-            meta["dsl_source"] = self.func.dsl_source
-        array.schunk.vlmeta["_LazyArray"] = meta
+            operand_urlpath = Path(operand_payload["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()
+            operand_payload["urlpath"] = ref_urlpath
+
+        array = make_b2object_carrier(
+            "lazyudf",
+            self.shape,
+            self.dtype,
+            chunks=self.chunks,
+            blocks=self.blocks,
+            **kwargs,
+        )
+        write_b2object_payload(array, payload)
+        return array
 
 
 def _numpy_eval_expr(expression, operands, prefer_blosc=False):

tests/data/legacy_lazyudf_v1/a.b2nd

@@ -975,3 +975,33 @@ def test_dsl_save_input_names_match(tmp_path):
     assert isinstance(reloaded.func, DSLKernel)
     assert reloaded.func.input_names == ["x", "y"]
     assert list(reloaded.inputs_dict.keys()) == ["x", "y"]
+
+
+def test_dsl_save_dictstore_operands(tmp_path):
+    shape = (10,)
+    store_path = tmp_path / "ops.b2z"
+    ext_a = tmp_path / "a.b2nd"
+    ext_b = tmp_path / "b.b2nd"
+    expr_path = tmp_path / "lazy.b2nd"
+
+    a = blosc2.asarray(np.arange(shape[0], dtype=np.float64), urlpath=str(ext_a), mode="w")
+    b = blosc2.asarray(np.arange(shape[0], dtype=np.float64) * 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"]
+        lazy = blosc2.lazyudf(kernel_save_simple, (a, b), dtype=np.float64)
+        lazy.save(urlpath=str(expr_path))
+
+    carrier = blosc2.open(str(expr_path), mode="r").array
+    assert carrier.schunk.vlmeta["b2o"]["operands"] == {
+        "o0": {"kind": "dictstore_key", "version": 1, "urlpath": "ops.b2z", "key": "/a"},
+        "o1": {"kind": "dictstore_key", "version": 1, "urlpath": "ops.b2z", "key": "/b"},
+    }
+
+    reloaded = blosc2.open(str(expr_path), mode="r")
+    expected = (np.arange(shape[0], dtype=np.float64) * 3) ** 2
+    np.testing.assert_allclose(reloaded.compute()[...], expected, rtol=1e-5, atol=1e-6)

tests/data/legacy_lazyudf_v1/b.b2nd

@@ -1874,12 +1874,8 @@ def test_to_cframe():
     dtype = "float64"
     a = blosc2.linspace(0, 1, N * N, dtype=dtype, shape=(N, N))
     expr = a**3 + blosc2.sin(a**2)
-    cframe = expr.to_cframe()
-    assert len(cframe) > 0
-    arr = blosc2.ndarray_from_cframe(cframe)
-    assert arr.shape == expr.shape
-    assert arr.dtype == expr.dtype
-    assert np.allclose(arr[:], expr[:])
+    with pytest.raises(ValueError, match="stored on disk/network"):
+        expr.to_cframe()
 
 
 # Test for the bug where multiplying two complex lazy expressions would fail with:

tests/data/legacy_lazyudf_v1/expr.b2nd

@@ -15,6 +15,11 @@
 import blosc2.c2array as blosc2_c2array
 
 
+@blosc2.dsl_kernel
+def kernel_add_square(x, y):
+    return x * x + y * y + 2 * x * y
+
+
 def _make_c2array(
     monkeypatch,
     path="@public/examples/ds-1d.b2nd",
@@ -128,3 +133,52 @@ def test_legacy_lazyexpr_open_backward_compat():
 
     assert isinstance(restored, blosc2.LazyExpr)
     np.testing.assert_array_equal(restored[:], np.arange(5, dtype=np.int64) * 3)
+
+
+def test_legacy_lazyudf_open_backward_compat():
+    fixture = Path(__file__).parent / "data" / "legacy_lazyudf_v1" / "expr.b2nd"
+
+    restored = blosc2.open(fixture, mode="r")
+
+    assert isinstance(restored, blosc2.LazyUDF)
+    np.testing.assert_allclose(restored.compute()[:], (np.arange(5, dtype=np.float64) * 3) ** 2)
+
+
+def test_lazyudf_from_cframe_roundtrip(tmp_path):
+    a = blosc2.asarray(np.arange(5, dtype=np.float64), urlpath=tmp_path / "a.b2nd", mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.float64) * 2, urlpath=tmp_path / "b.b2nd", mode="w")
+    expr = blosc2.lazyudf(kernel_add_square, (a, b), dtype=np.float64)
+    carrier = blosc2.ndarray_from_cframe(expr.to_cframe())
+
+    assert carrier.schunk.meta["b2o"] == {"kind": "lazyudf", "version": 1}
+    payload = carrier.schunk.vlmeta["b2o"]
+    assert payload["kind"] == "lazyudf"
+    assert payload["version"] == 1
+    assert payload["function_kind"] == "dsl"
+    assert payload["dsl_version"] == 1
+    assert payload["name"] == "kernel_add_square"
+    assert "kernel_add_square" in payload["udf_source"]
+    assert payload["dtype"] == np.dtype(np.float64).str
+    assert payload["shape"] == [5]
+    assert payload["operands"] == {
+        "o0": {"kind": "urlpath", "version": 1, "urlpath": str(tmp_path / "a.b2nd")},
+        "o1": {"kind": "urlpath", "version": 1, "urlpath": str(tmp_path / "b.b2nd")},
+    }
+
+    restored = blosc2.from_cframe(expr.to_cframe())
+
+    assert isinstance(restored, blosc2.LazyUDF)
+    np.testing.assert_allclose(restored[:], (np.arange(5, dtype=np.float64) * 3) ** 2)
+
+
+def test_lazyudf_open_roundtrip(tmp_path):
+    a = blosc2.asarray(np.arange(5, dtype=np.float64), urlpath=tmp_path / "a.b2nd", mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.float64) * 2, urlpath=tmp_path / "b.b2nd", mode="w")
+    expr = blosc2.lazyudf(kernel_add_square, (a, b), dtype=np.float64)
+    urlpath = tmp_path / "expr.b2nd"
+
+    expr.save(urlpath)
+    restored = blosc2.open(urlpath, mode="r")
+
+    assert isinstance(restored, blosc2.LazyUDF)
+    np.testing.assert_allclose(restored[:], (np.arange(5, dtype=np.float64) * 3) ** 2)