Avoid eager DictStore rehydration and new b2o bundle demo

Author: FrancescAlted

examples/embeded-expr-udf-b2z.py

@@ -0,0 +1,76 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org>
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#######################################################################
+
+import numpy as np
+
+import blosc2
+from blosc2 import where
+
+
+def show(label, value):
+    print(f"{label}: {value}")
+
+
+@blosc2.dsl_kernel
+def masked_energy(a, b, mask):
+    return where(mask > 0, a * a + 2 * b, 0.0)
+
+
+bundle_path = "example_b2o_bundle.b2z"
+blosc2.remove_urlpath(bundle_path)
+
+# Build a portable bundle with ordinary arrays plus persisted lazy recipes.
+with blosc2.DictStore(bundle_path, mode="w", threshold=1) as store:
+    store["/data/a"] = np.linspace(0.0, 1.0, 10, dtype=np.float32)
+    store["/data/b"] = np.linspace(1.0, 2.0, 10, dtype=np.float32)
+    store["/data/mask"] = np.asarray([0, 1, 1, 0, 1, 0, 1, 1, 0, 1], dtype=np.int8)
+
+    # Reopen the array members through the store so operand refs point back to
+    # the .b2z container via dictstore_key payloads.
+    a = store["/data/a"]
+    b = store["/data/b"]
+    mask = store["/data/mask"]
+
+    expr = blosc2.lazyexpr("(a - b) / (a + b + 1e-6)", operands={"a": a, "b": b})
+    udf = blosc2.lazyudf(masked_energy, (a, b, mask), dtype=np.float32, shape=a.shape)
+
+    # DictStore currently stores array-like external leaves, so persist the
+    # logical lazy objects through their b2o NDArray carriers.
+    store["/recipes/expr"] = blosc2.ndarray_from_cframe(expr.to_cframe())
+    store["/recipes/udf"] = blosc2.ndarray_from_cframe(udf.to_cframe())
+
+show("Bundle created", bundle_path)
+
+# Reopen the bundle read-only.  The persisted LazyExpr and LazyUDF can be
+# evaluated directly without re-saving, rebuilding, or re-deploying the .b2z.
+with blosc2.open(bundle_path, mode="r") as store:
+    show("Read-only keys", sorted(store.keys()))
+
+    expr = store["/recipes/expr"]
+    udf = store["/recipes/udf"]
+
+    expr_result = expr.compute()
+    udf_result = udf.compute()
+
+    show("Reopened expr type", type(expr).__name__)
+    show("Reopened udf type", type(udf).__name__)
+    show("Expr operand refs", expr.array.schunk.vlmeta["b2o"]["operands"])
+    show("UDF operand refs", udf.array.schunk.vlmeta["b2o"]["operands"])
+    show("Expr values", np.round(expr_result[:], 4))
+    show("UDF values", udf_result[:])
+
+    expected_expr = (store["/data/a"][:] - store["/data/b"][:]) / (
+        store["/data/a"][:] + store["/data/b"][:] + 1e-6
+    )
+    expected_udf = np.where(
+        store["/data/mask"][:] > 0,
+        store["/data/a"][:] ** 2 + 2 * store["/data/b"][:],
+        0.0,
+    ).astype(np.float32)
+    np.testing.assert_allclose(expr_result[:], expected_expr, rtol=1e-6, atol=1e-6)
+    np.testing.assert_allclose(udf_result[:], expected_udf, rtol=1e-6, atol=1e-6)
+    show("Read-only evaluation", "ok")

src/blosc2/dict_store.py

@@ -227,30 +227,39 @@ def _opened_external_kind(
         rel_path: str,
     ) -> str | None:
         """Return the supported external leaf kind for an already opened object."""
-        processed = process_opened_object(opened)
-        if isinstance(processed, blosc2.BatchStore):
-            kind = "batchstore"
-        elif isinstance(processed, blosc2.VLArray):
-            kind = "vlarray"
-        elif isinstance(processed, blosc2.NDArray):
+        meta = getattr(opened, "schunk", opened).meta
+        if "b2o" in meta and isinstance(opened, blosc2.NDArray):
+            # Keep b2o carriers as NDArray external leaves during discovery.
+            # Rehydrating them here can recurse when a lazy recipe points back
+            # into the same DictStore via dictstore_key refs.
             kind = "ndarray"
-        elif isinstance(processed, SChunk):
-            kind = "schunk"
+            processed_name = type(opened).__name__
         else:
-            warnings.warn(
-                f"Ignoring unsupported Blosc2 object at '{rel_path}' during DictStore discovery: "
-                f"{type(processed).__name__}",
-                UserWarning,
-                stacklevel=2,
-            )
-            return None
+            processed = process_opened_object(opened)
+            processed_name = type(processed).__name__
+            if isinstance(processed, blosc2.BatchStore):
+                kind = "batchstore"
+            elif isinstance(processed, blosc2.VLArray):
+                kind = "vlarray"
+            elif isinstance(processed, blosc2.NDArray):
+                kind = "ndarray"
+            elif isinstance(processed, SChunk):
+                kind = "schunk"
+            else:
+                warnings.warn(
+                    f"Ignoring unsupported Blosc2 object at '{rel_path}' during DictStore discovery: "
+                    f"{processed_name}",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                return None
 
         expected_ext = cls._expected_ext_from_kind(kind)
         found_ext = os.path.splitext(rel_path)[1]
         if found_ext != expected_ext:
             warnings.warn(
                 f"External leaf '{rel_path}' uses extension '{found_ext}' but metadata resolves to "
-                f"{type(processed).__name__}; expected '{expected_ext}'.",
+                f"{processed_name}; expected '{expected_ext}'.",
                 UserWarning,
                 stacklevel=2,
             )

tests/test_b2objects.py

@@ -182,3 +182,28 @@ def test_lazyudf_open_roundtrip(tmp_path):
 
     assert isinstance(restored, blosc2.LazyUDF)
     np.testing.assert_allclose(restored[:], (np.arange(5, dtype=np.float64) * 3) ** 2)
+
+
+def test_b2z_bundle_with_lazy_recipes_opens_read_only(tmp_path):
+    bundle_path = tmp_path / "bundle.b2z"
+
+    with blosc2.DictStore(str(bundle_path), mode="w", threshold=1) as store:
+        store["/data/a"] = np.arange(5, dtype=np.float64)
+        store["/data/b"] = np.arange(5, dtype=np.float64) * 2
+
+        a = store["/data/a"]
+        b = store["/data/b"]
+        expr = blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
+        udf = blosc2.lazyudf(kernel_add_square, (a, b), dtype=np.float64, shape=a.shape)
+
+        store["/recipes/expr"] = blosc2.ndarray_from_cframe(expr.to_cframe())
+        store["/recipes/udf"] = blosc2.ndarray_from_cframe(udf.to_cframe())
+
+    with blosc2.open(str(bundle_path), mode="r") as store:
+        restored_expr = store["/recipes/expr"]
+        restored_udf = store["/recipes/udf"]
+
+        assert isinstance(restored_expr, blosc2.LazyExpr)
+        assert isinstance(restored_udf, blosc2.LazyUDF)
+        np.testing.assert_allclose(restored_expr.compute()[:], np.arange(5, dtype=np.float64) * 3)
+        np.testing.assert_allclose(restored_udf.compute()[:], (np.arange(5, dtype=np.float64) * 3) ** 2)