Add support for serializing LazyUDF with DSL functions

Author: FrancescAlted

doc/reference/batch_store.rst

@@ -47,15 +47,18 @@ Structured reference-style Blosc2 objects currently supported:
 
 - ``C2Array``
 - ``LazyExpr``
+- ``LazyUDF`` backed by ``@blosc2.dsl_kernel``
 
-``LazyExpr`` values preserve reference semantics and are serialized as a recipe
-plus durable operand references. Supported operands are:
+``LazyExpr`` values and supported ``LazyUDF`` values preserve reference
+semantics and are serialized as a recipe plus durable operand references.
+Supported operands are:
 
 - persistent local Blosc2 operands reopenable from ``urlpath``
 - remote ``C2Array`` operands
 - ``DictStore`` members reopenable from ``(.b2d|.b2z, key)``
 
-Purely in-memory operands are intentionally not supported.
+Purely in-memory operands are intentionally not supported. Plain Python
+``LazyUDF`` callables are not serialized by msgpack.
 
 Quick example
 -------------

examples/vlstore-lazyudf.py

@@ -0,0 +1,47 @@
+#######################################################################
+# 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
+
+
+def show(label, value):
+    print(f"{label}: {value}")
+
+
+@blosc2.dsl_kernel
+def kernel_add_twice(x, y):
+    return x + y * 2
+
+
+urlpath = "example_vlstore_lazyudf.b2frame"
+a_path = "example_vlstore_lazyudf_a.b2nd"
+b_path = "example_vlstore_lazyudf_b.b2nd"
+blosc2.remove_urlpath(urlpath)
+blosc2.remove_urlpath(a_path)
+blosc2.remove_urlpath(b_path)
+
+a = blosc2.asarray(np.arange(5, dtype=np.float32), urlpath=a_path, mode="w")
+b = blosc2.asarray(np.arange(5, dtype=np.float32) * 2, urlpath=b_path, mode="w")
+lazy_udf = blosc2.lazyudf(kernel_add_twice, (a, b), dtype=a.dtype, shape=a.shape)
+
+vla = blosc2.VLArray(urlpath=urlpath, mode="w", contiguous=True)
+vla.append({"kind": "lazyudf", "value": lazy_udf})
+
+restored = vla[0]["value"]
+show("Stored type", type(vla[0]["value"]).__name__)
+show("Computed values", restored[:])
+
+reopened = blosc2.open(urlpath, mode="r")
+restored_reopened = reopened[0]["value"]
+show("Reopened type", type(restored_reopened).__name__)
+show("Reopened values", restored_reopened[:])
+
+blosc2.remove_urlpath(urlpath)
+blosc2.remove_urlpath(a_path)
+blosc2.remove_urlpath(b_path)

src/blosc2/_msgpack_utils.py

@@ -7,9 +7,17 @@
 
 from __future__ import annotations
 
+import builtins
+import inspect
+import linecache
+import textwrap
+from dataclasses import asdict
+
+import numpy as np
 from msgpack import ExtType, packb, unpackb
 
 from blosc2 import blosc2_ext
+from blosc2.dsl_kernel import DSLKernel
 
 # Msgpack extension type codes are application-defined.  Reserve code 42 in
 # python-blosc2 for values serialized as Blosc2 CFrames via ``to_cframe()`` and
@@ -21,6 +29,7 @@
 # stable ``kind`` and ``version`` envelope.
 _BLOSC2_STRUCTURED_EXT_CODE = 43
 _BLOSC2_STRUCTURED_VERSION = 1
+_BLOSC2_DSL_VERSION = 1
 
 
 def _encode_operand_reference(obj):
@@ -74,6 +83,43 @@ def _encode_structured_reference(obj):
             "operands": {key: _encode_operand_reference(value) for key, value in operands.items()},
         }
         return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
+    if isinstance(obj, blosc2.LazyUDF):
+        if not isinstance(obj.func, DSLKernel):
+            raise TypeError("Structured Blosc2 msgpack payload only supports LazyUDF backed by DSLKernel")
+        udf_func = obj.func.func
+        udf_name = getattr(udf_func, "__name__", obj.func.__name__)
+        try:
+            udf_source = textwrap.dedent(inspect.getsource(udf_func)).lstrip()
+        except Exception:
+            udf_source = obj.func.dsl_source
+        if udf_source is None:
+            raise ValueError("Structured LazyUDF msgpack payload requires recoverable DSL kernel source")
+        kwargs = {}
+        for key, value in obj.kwargs.items():
+            if key in {"dtype", "shape"}:
+                continue
+            if isinstance(value, blosc2.CParams | blosc2.DParams):
+                kwargs[key] = asdict(value)
+            else:
+                kwargs[key] = value
+        # Keep both source forms:
+        # - udf_source recreates the executable Python function object
+        # - dsl_source preserves the DSLKernel's normalized DSL metadata so the
+        #   reconstructed function can keep its DSL identity and fast-path hints
+        payload = {
+            "kind": "lazyudf",
+            "version": _BLOSC2_STRUCTURED_VERSION,
+            "function_kind": "dsl",
+            "dsl_version": _BLOSC2_DSL_VERSION,
+            "name": udf_name,
+            "udf_source": udf_source,
+            "dsl_source": obj.func.dsl_source,
+            "dtype": np.dtype(obj.dtype).str,
+            "shape": list(obj.shape),
+            "operands": {f"o{i}": _encode_operand_reference(value) for i, value in enumerate(obj.inputs)},
+            "kwargs": kwargs,
+        }
+        return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
     return None
 
 
@@ -113,8 +159,6 @@ def _decode_operand_reference(payload):
 
 
 def _decode_structured_reference(data):
-    import blosc2
-
     payload = unpackb(data)
     if not isinstance(payload, dict):
         raise TypeError("Structured Blosc2 msgpack payload must decode to a mapping")
@@ -127,29 +171,80 @@ def _decode_structured_reference(data):
     if kind == "c2array":
         return _decode_operand_reference(payload)
     if kind == "lazyexpr":
-        expression = payload.get("expression")
-        if not isinstance(expression, str):
-            raise TypeError("Structured LazyExpr msgpack payload requires a string 'expression'")
-        operands_payload = payload.get("operands")
-        if not isinstance(operands_payload, dict):
-            raise TypeError("Structured LazyExpr msgpack payload requires a mapping 'operands'")
-        operands = {key: _decode_operand_reference(value) for key, value in operands_payload.items()}
-        return blosc2.lazyexpr(expression, operands=operands)
+        return _decode_structured_lazyexpr(payload)
+    if kind == "lazyudf":
+        return _decode_structured_lazyudf(payload)
     raise ValueError(f"Unsupported structured Blosc2 msgpack payload kind: {kind!r}")
 
 
+def _decode_structured_lazyexpr(payload):
+    import blosc2
+
+    expression = payload.get("expression")
+    if not isinstance(expression, str):
+        raise TypeError("Structured LazyExpr msgpack payload requires a string 'expression'")
+    operands_payload = payload.get("operands")
+    if not isinstance(operands_payload, dict):
+        raise TypeError("Structured LazyExpr msgpack payload requires a mapping 'operands'")
+    operands = {key: _decode_operand_reference(value) for key, value in operands_payload.items()}
+    return blosc2.lazyexpr(expression, operands=operands)
+
+
+def _decode_structured_lazyudf(payload):
+    import blosc2
+
+    function_kind = payload.get("function_kind")
+    if function_kind != "dsl":
+        raise ValueError(f"Unsupported structured LazyUDF function kind: {function_kind!r}")
+    dsl_version = payload.get("dsl_version")
+    if dsl_version != _BLOSC2_DSL_VERSION:
+        raise ValueError(f"Unsupported structured LazyUDF DSL version: {dsl_version!r}")
+    udf_source = payload.get("udf_source")
+    if not isinstance(udf_source, str):
+        raise TypeError("Structured LazyUDF msgpack payload requires a string 'udf_source'")
+    name = payload.get("name")
+    if not isinstance(name, str):
+        raise TypeError("Structured LazyUDF msgpack payload requires a string 'name'")
+    dtype = payload.get("dtype")
+    if not isinstance(dtype, str):
+        raise TypeError("Structured LazyUDF msgpack payload requires a string 'dtype'")
+    shape_payload = payload.get("shape")
+    if not isinstance(shape_payload, list):
+        raise TypeError("Structured LazyUDF msgpack payload requires a list 'shape'")
+    operands_payload = payload.get("operands")
+    if not isinstance(operands_payload, dict):
+        raise TypeError("Structured LazyUDF msgpack payload requires a mapping 'operands'")
+    kwargs = payload.get("kwargs", {})
+    if not isinstance(kwargs, dict):
+        raise TypeError("Structured LazyUDF msgpack payload requires a mapping 'kwargs'")
+
+    local_ns = {}
+    filename = f"<{name}>"
+    safe_globals = {
+        "__builtins__": {k: v for k, v in builtins.__dict__.items() if k != "__import__"},
+        "np": np,
+        "blosc2": blosc2,
+    }
+    linecache.cache[filename] = (len(udf_source), None, udf_source.splitlines(True), filename)
+    exec(compile(udf_source, filename, "exec"), safe_globals, local_ns)
+    func = local_ns[name]
+    if not isinstance(func, DSLKernel):
+        func = DSLKernel(func)
+    dsl_source = payload.get("dsl_source")
+    if dsl_source is not None and func.dsl_source is None:
+        func.dsl_source = dsl_source
+
+    operands = tuple(
+        _decode_operand_reference(operands_payload[f"o{n}"]) for n in range(len(operands_payload))
+    )
+    return blosc2.lazyudf(func, operands, dtype=np.dtype(dtype), shape=tuple(shape_payload), **kwargs)
+
+
 def _encode_msgpack_ext(obj):
     import blosc2
 
     if isinstance(
-        obj,
-        (
-            blosc2.NDArray,
-            blosc2.SChunk,
-            blosc2.VLArray,
-            blosc2.BatchStore,
-            blosc2.EmbedStore,
-        ),
+        obj, blosc2.NDArray | blosc2.SChunk | blosc2.VLArray | blosc2.BatchStore | blosc2.EmbedStore
     ):
         return ExtType(_BLOSC2_EXT_CODE, obj.to_cframe())
     structured = _encode_structured_reference(obj)

src/blosc2/batch_store.py

@@ -173,11 +173,14 @@ class BatchStore:
         :class:`blosc2.EmbedStore`, which are serialized transparently via
         :meth:`to_cframe` / :func:`blosc2.from_cframe`. Msgpack also supports
         structured Blosc2 reference objects, currently
-        :class:`blosc2.C2Array` and :class:`blosc2.LazyExpr`. Lazy expressions
-        preserve reference semantics, so only persistent local operands,
-        :class:`blosc2.C2Array` operands, and :class:`blosc2.DictStore`
-        members are supported; purely in-memory operands are rejected.
-        ``"arrow"`` is optional and requires ``pyarrow``.
+        :class:`blosc2.C2Array`, :class:`blosc2.LazyExpr`, and
+        :class:`blosc2.LazyUDF` backed by :func:`blosc2.dsl_kernel`. These lazy
+        objects preserve reference semantics, so only persistent local
+        operands, :class:`blosc2.C2Array` operands, and
+        :class:`blosc2.DictStore` members are supported; purely in-memory
+        operands are rejected. Plain Python :class:`blosc2.LazyUDF` callables
+        are not serialized by msgpack. ``"arrow"`` is optional and requires
+        ``pyarrow``.
     _from_schunk : blosc2.SChunk, optional
         Internal hook used when reopening an already-tagged BatchStore.
     **kwargs

src/blosc2/vlarray.py

@@ -38,11 +38,13 @@ class VLArray:
     transparently via :meth:`to_cframe` / :func:`blosc2.from_cframe`.
 
     Msgpack also supports structured Blosc2 reference objects. Currently this
-    includes :class:`blosc2.C2Array` and :class:`blosc2.LazyExpr`. Lazy
-    expressions are serialized as recipes plus durable operand references, so
-    only persistent local operands, :class:`blosc2.C2Array` operands, and
-    :class:`blosc2.DictStore` members are supported. Purely in-memory operands
-    are intentionally rejected.
+    includes :class:`blosc2.C2Array`, :class:`blosc2.LazyExpr`, and
+    :class:`blosc2.LazyUDF` backed by :func:`blosc2.dsl_kernel`. Lazy
+    expressions and supported lazy UDFs are serialized as recipes plus durable
+    operand references, so only persistent local operands,
+    :class:`blosc2.C2Array` operands, and :class:`blosc2.DictStore` members are
+    supported. Purely in-memory operands are intentionally rejected. Plain
+    Python :class:`blosc2.LazyUDF` callables are not serialized by msgpack.
     """
 
     @staticmethod

tests/test_batch_store.py

@@ -12,6 +12,17 @@
 import blosc2.c2array as blosc2_c2array
 from blosc2._msgpack_utils import msgpack_packb, msgpack_unpackb
 
+
+@blosc2.dsl_kernel
+def _kernel_add_twice(x, y):
+    return x + y * 2
+
+
+def _python_udf_add(inputs_tuple, output, offset):
+    x, y = inputs_tuple
+    output[:] = x + y
+
+
 BATCHES = [
     [b"bytes\x00payload", "plain text", 42],
     [{"nested": [1, 2]}, None, {"tail": True}],
@@ -69,6 +80,20 @@ def _make_in_memory_lazyexpr():
     return blosc2.lazyexpr("a + b", operands={"a": a, "b": b})
 
 
+def _make_persistent_lazyudf(tmp_path):
+    a = blosc2.asarray(np.arange(5, dtype=np.float32), urlpath=tmp_path / "a_udf.b2nd", mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.float32) * 2, urlpath=tmp_path / "b_udf.b2nd", mode="w")
+    udf = blosc2.lazyudf(_kernel_add_twice, (a, b), dtype=a.dtype, shape=a.shape)
+    expected = a[:] + b[:] * 2
+    return udf, expected
+
+
+def _make_persistent_python_lazyudf(tmp_path):
+    a = blosc2.asarray(np.arange(5, dtype=np.float32), urlpath=tmp_path / "a_pyudf.b2nd", mode="w")
+    b = blosc2.asarray(np.arange(5, dtype=np.float32) * 2, urlpath=tmp_path / "b_pyudf.b2nd", mode="w")
+    return blosc2.lazyudf(_python_udf_add, (a, b), dtype=a.dtype, shape=a.shape)
+
+
 @pytest.mark.parametrize(
     ("contiguous", "urlpath"),
     [
@@ -296,6 +321,33 @@ def test_batchstore_msgpack_supports_lazyexpr(tmp_path):
     np.testing.assert_array_equal(restored[:], expected)
 
 
+def test_msgpack_supports_lazyudf_dslkernel(tmp_path):
+    udf, expected = _make_persistent_lazyudf(tmp_path)
+
+    restored = msgpack_unpackb(msgpack_packb({"udf": udf}))["udf"]
+
+    assert isinstance(restored, blosc2.LazyUDF)
+    np.testing.assert_allclose(restored[:], expected)
+
+
+def test_batchstore_msgpack_supports_lazyudf_dslkernel(tmp_path):
+    udf, expected = _make_persistent_lazyudf(tmp_path)
+
+    barray = blosc2.BatchStore(items_per_block=2)
+    barray.append([udf])
+    restored = barray[0][0]
+
+    assert isinstance(restored, blosc2.LazyUDF)
+    np.testing.assert_allclose(restored[:], expected)
+
+
+def test_msgpack_rejects_plain_python_lazyudf(tmp_path):
+    udf = _make_persistent_python_lazyudf(tmp_path)
+
+    with pytest.raises(TypeError, match="DSLKernel"):
+        msgpack_packb({"udf": udf})
+
+
 def test_msgpack_rejects_lazyexpr_with_in_memory_operands():
     expr = _make_in_memory_lazyexpr()