Add vlmeta support to LazyArray objects

Author: FrancescAlted

doc/reference/lazyarray.rst

@@ -14,6 +14,11 @@ You can get an object following the LazyArray API with any of the following ways
 
 The LazyArray object is a thin wrapper around the expression or user-defined function that allows for lazy computation. This means that the expression is not computed until the ``compute`` or ``__getitem__`` methods are called. The ``compute`` method will return a new NDArray object with the result of the expression evaluation. The ``__getitem__`` method will return an NumPy object instead.
 
+LazyArray objects also support user metadata via :attr:`LazyArray.vlmeta`. For
+in-memory objects, this metadata lives on the Python object itself. For
+persisted LazyArrays reopened from disk, metadata is synchronized with the
+underlying carrier and survives reopening.
+
 See the `LazyExpr`_ and `LazyUDF`_ sections for more information.
 
 .. currentmodule:: blosc2
@@ -33,6 +38,10 @@ See the `LazyExpr`_ and `LazyUDF`_ sections for more information.
     ---------------
     .. automethod:: __getitem__
 
+    Attributes
+    ----------
+    .. autoattribute:: vlmeta
+
 .. _LazyExpr:
 
 LazyExpr
@@ -51,7 +60,7 @@ LazyUDF
 
 For getting a LazyUDF object (which is LazyArray-compliant) from a user-defined Python function, you can use the lazyudf constructor below. See  `a tutorial on how this works <../getting_started/tutorials/03.lazyarray-udf.html>`_.
 
-This object follows the `LazyArray`_ API for computation, although storage is not supported yet.
+This object follows the `LazyArray`_ API for computation and storage.
 
 .. autofunction:: lazyudf
 

src/blosc2/b2objects.py

@@ -23,6 +23,7 @@
 _B2OBJECT_META_KEY = "b2o"
 _B2OBJECT_VERSION = 1
 _B2OBJECT_DSL_VERSION = 1
+_B2OBJECT_USER_VLMETA_KEY = "_b2o_user_vlmeta"
 
 
 def make_b2object_carrier(
@@ -44,6 +45,17 @@ def write_b2object_payload(array, payload: dict[str, Any]) -> None:
     array.schunk.vlmeta[_B2OBJECT_META_KEY] = payload
 
 
+def write_b2object_user_vlmeta(array, user_vlmeta: dict[str, Any]) -> None:
+    array.schunk.vlmeta[_B2OBJECT_USER_VLMETA_KEY] = user_vlmeta
+
+
+def read_b2object_user_vlmeta(obj) -> dict[str, Any]:
+    schunk = getattr(obj, "schunk", obj)
+    if _B2OBJECT_USER_VLMETA_KEY not in schunk.vlmeta:
+        return {}
+    return schunk.vlmeta[_B2OBJECT_USER_VLMETA_KEY]
+
+
 def encode_operand_reference(obj):
     return blosc2.Ref.from_object(obj).to_dict()
 
@@ -232,4 +244,5 @@ def open_b2object(obj):
     if isinstance(opened, blosc2.LazyExpr | blosc2.LazyUDF):
         opened.array = obj
         opened.schunk = schunk
+        opened._set_user_vlmeta(read_b2object_user_vlmeta(obj), sync=False)
     return opened

src/blosc2/lazyexpr.py

@@ -24,6 +24,7 @@
 import textwrap
 import threading
 from abc import ABC, abstractmethod, abstractproperty
+from collections.abc import MutableMapping
 from dataclasses import asdict
 from enum import Enum
 from pathlib import Path
@@ -42,7 +43,13 @@
 
 import blosc2
 
-from .b2objects import encode_b2object_payload, make_b2object_carrier, write_b2object_payload
+from .b2objects import (
+    encode_b2object_payload,
+    make_b2object_carrier,
+    read_b2object_user_vlmeta,
+    write_b2object_payload,
+    write_b2object_user_vlmeta,
+)
 from .dsl_kernel import DSLKernel, DSLSyntaxError, DSLValidator, specialize_miniexpr_inputs
 
 if blosc2._HAS_NUMBA:
@@ -325,7 +332,64 @@ class LazyArrayEnum(Enum):
     UDF = 1
 
 
+class LazyArrayVLMeta(MutableMapping):
+    """User metadata attached to a LazyArray."""
+
+    def __init__(self, lazyarr: LazyArray):
+        self.lazyarr = lazyarr
+
+    def __getitem__(self, key):
+        return self.lazyarr._get_user_vlmeta()[key]
+
+    def __setitem__(self, key, value):
+        data = self.lazyarr._get_user_vlmeta()
+        data[key] = value
+        self.lazyarr._sync_user_vlmeta()
+
+    def __delitem__(self, key):
+        data = self.lazyarr._get_user_vlmeta()
+        del data[key]
+        self.lazyarr._sync_user_vlmeta()
+
+    def __iter__(self):
+        return iter(self.lazyarr._get_user_vlmeta())
+
+    def __len__(self):
+        return len(self.lazyarr._get_user_vlmeta())
+
+    def getall(self):
+        return self.lazyarr._get_user_vlmeta().copy()
+
+    def __repr__(self):
+        return repr(self.getall())
+
+    def __str__(self):
+        return str(self.getall())
+
+
 class LazyArray(ABC, blosc2.Operand):
+    def _get_user_vlmeta(self) -> dict[str, Any]:
+        if not hasattr(self, "_vlmeta_user"):
+            self._vlmeta_user = {}
+        return self._vlmeta_user
+
+    def _set_user_vlmeta(self, metadata: dict[str, Any], *, sync: bool = True) -> None:
+        self._vlmeta_user = dict(metadata)
+        if sync:
+            self._sync_user_vlmeta()
+
+    def _sync_user_vlmeta(self) -> None:
+        array = getattr(self, "array", None)
+        if array is not None:
+            write_b2object_user_vlmeta(array, self._get_user_vlmeta())
+
+    @property
+    def vlmeta(self) -> LazyArrayVLMeta:
+        """User variable-length metadata for this LazyArray."""
+        if not hasattr(self, "_vlmeta_proxy"):
+            self._vlmeta_proxy = LazyArrayVLMeta(self)
+        return self._vlmeta_proxy
+
     @abstractmethod
     def indices(self, order: str | list[str] | None = None) -> blosc2.LazyArray:
         """
@@ -494,6 +558,9 @@ def save(self, **kwargs: Any) -> None:
           section for more info).
         * This is currently only supported for :ref:`LazyExpr` and :ref:`LazyUDF`
           (including kernels decorated with :func:`blosc2.dsl_kernel`).
+        * User metadata can be attached via :attr:`vlmeta`. For in-memory LazyArrays
+          this stays in memory; for persisted LazyArrays it is serialized and restored
+          on reopen.
 
         Examples
         --------
@@ -3764,6 +3831,7 @@ def _to_b2object_carrier(self, **kwargs):
             **kwargs,
         )
         write_b2object_payload(array, payload)
+        write_b2object_user_vlmeta(array, self._get_user_vlmeta())
         return array
 
     @classmethod
@@ -4125,6 +4193,7 @@ def save(self, urlpath=None, **kwargs):
             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
+            write_b2object_user_vlmeta(array, self._get_user_vlmeta())
 
     def to_cframe(self) -> bytes:
         return self._to_b2object_carrier().to_cframe()
@@ -4160,6 +4229,7 @@ def _to_b2object_carrier(self, **kwargs):
             **kwargs,
         )
         write_b2object_payload(array, payload)
+        write_b2object_user_vlmeta(array, self._get_user_vlmeta())
         return array
 
 
@@ -4518,6 +4588,7 @@ def open_lazyarray(array):
     new_expr.array = array
     # We want to expose schunk too, so that .info() can be used on the LazyArray
     new_expr.schunk = array.schunk
+    new_expr._set_user_vlmeta(read_b2object_user_vlmeta(array), sync=False)
     return new_expr
 
 

tests/ndarray/test_lazyexpr.py

@@ -1768,6 +1768,29 @@ def test_save_proxy_operands_reopen_default_mode(tmp_path):
     np.testing.assert_array_equal(restored[:], np.arange(10, dtype=np.int64) * 2)
 
 
+def test_lazyexpr_vlmeta_in_memory_and_persisted(tmp_path):
+    a = blosc2.asarray(np.arange(5, dtype=np.int64), urlpath=str(tmp_path / "a.b2nd"), mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.int64), urlpath=str(tmp_path / "b.b2nd"), mode="w")
+    expr = a + b
+
+    expr.vlmeta["name"] = "sum"
+    expr.vlmeta["config"] = {"scale": 1}
+    assert expr.vlmeta["name"] == "sum"
+    assert expr.vlmeta["config"] == {"scale": 1}
+
+    expr_path = tmp_path / "expr_vlmeta.b2nd"
+    expr.save(str(expr_path))
+    restored = blosc2.open(str(expr_path))
+
+    assert restored.vlmeta["name"] == "sum"
+    assert restored.vlmeta["config"] == {"scale": 1}
+
+    restored.vlmeta["note"] = "persisted"
+    reopened = blosc2.open(str(expr_path))
+    assert reopened.vlmeta["note"] == "persisted"
+    np.testing.assert_array_equal(reopened[:], np.arange(5, dtype=np.int64) * 2)
+
+
 # Test the chaining of multiple lazy expressions
 def test_chain_expressions():
     N = 1_000

tests/ndarray/test_lazyudf.py

@@ -501,6 +501,23 @@ def test_save_ludf():
     blosc2.remove_urlpath(urlpath)
 
 
+def test_lazyudf_vlmeta_roundtrip(tmp_path):
+    a_path = tmp_path / "a.b2nd"
+    expr_path = tmp_path / "lazyudf_vlmeta.b2nd"
+    array = blosc2.asarray(np.arange(5, dtype=np.int64), urlpath=str(a_path), mode="w")
+    expr = blosc2.lazyudf(udf1p, (array,), np.float64)
+
+    expr.vlmeta["name"] = "increment"
+    expr.vlmeta["attrs"] = {"version": 1}
+    expr.save(urlpath=str(expr_path))
+
+    restored = blosc2.open(str(expr_path))
+
+    assert isinstance(restored, blosc2.LazyUDF)
+    assert restored.vlmeta["name"] == "increment"
+    assert restored.vlmeta["attrs"] == {"version": 1}
+
+
 # Test get_chunk method
 def test_get_chunk():
     a = blosc2.linspace(0, 100, 100, shape=(10, 10), chunks=(3, 4), blocks=(2, 3))