Better support for VLArrays in dict and tree stores

FrancescAlted · FrancescAlted · commit 2ba2532addd2 · 2026-03-14T19:14:16.000+01:00
diff --git a/src/blosc2/dict_store.py b/src/blosc2/dict_store.py
@@ -5,19 +5,23 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #######################################################################
 
+from __future__ import annotations
+
 import os
 import shutil
 import tempfile
 import zipfile
-from collections.abc import Iterator, Set
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 import blosc2
 from blosc2.c2array import C2Array
 from blosc2.embed_store import EmbedStore
-from blosc2.schunk import SChunk
+from blosc2.schunk import SChunk, _process_opened_object
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Set
 
 
 class DictStore:
@@ -244,20 +248,36 @@ def estore(self) -> EmbedStore:
         """Access the underlying EmbedStore."""
         return self._estore
 
-    def __setitem__(self, key: str, value: blosc2.Array | SChunk) -> None:
+    @staticmethod
+    def _value_nbytes(value: blosc2.Array | SChunk | blosc2.VLArray) -> int:
+        if isinstance(value, blosc2.VLArray):
+            return value.schunk.nbytes
+        return value.nbytes
+
+    @staticmethod
+    def _is_external_value(value: blosc2.Array | SChunk | blosc2.VLArray) -> bool:
+        return isinstance(value, (blosc2.NDArray, SChunk, blosc2.VLArray)) and bool(
+            getattr(value, "urlpath", None)
+        )
+
+    @staticmethod
+    def _external_ext(value: blosc2.Array | SChunk | blosc2.VLArray) -> str:
+        if isinstance(value, blosc2.NDArray):
+            return ".b2nd"
+        return ".b2f"
+
+    def __setitem__(self, key: str, value: blosc2.Array | SChunk | blosc2.VLArray) -> None:
         """Add a node to the DictStore."""
         if isinstance(value, np.ndarray):
             value = blosc2.asarray(value, cparams=self.cparams, dparams=self.dparams)
         # C2Array should always go to embed store; let estore handle it directly
         if isinstance(value, C2Array):
             self._estore[key] = value
             return
-        exceeds_threshold = self.threshold is not None and value.nbytes >= self.threshold
-        # Consider both NDArray and SChunk external files (have urlpath)
-        external_file = isinstance(value, (blosc2.NDArray, SChunk)) and getattr(value, "urlpath", None)
+        exceeds_threshold = self.threshold is not None and self._value_nbytes(value) >= self.threshold
+        external_file = self._is_external_value(value)
         if exceeds_threshold or (external_file and self.threshold is None):
-            # Choose extension based on type
-            ext = ".b2f" if isinstance(value, SChunk) else ".b2nd"
+            ext = self._external_ext(value)
             # Convert key to a proper file path within the tree directory
             rel_key = key.lstrip("/")
             dest_path = os.path.join(self.working_dir, rel_key + ext)
@@ -272,7 +292,7 @@ def __setitem__(self, key: str, value: blosc2.Array | SChunk) -> None:
                 if hasattr(value, "save"):
                     value.save(urlpath=dest_path)
                 else:
-                    # An SChunk does not have a save() method
+                    # SChunk and VLArray can both be persisted via their cframe.
                     with open(dest_path, "wb") as f:
                         f.write(value.to_cframe())
             else:
@@ -290,20 +310,21 @@ def __setitem__(self, key: str, value: blosc2.Array | SChunk) -> None:
                 value = blosc2.from_cframe(value.to_cframe())
             self._estore[key] = value
 
-    def __getitem__(self, key: str) -> blosc2.NDArray | SChunk | C2Array:
+    def __getitem__(self, key: str) -> blosc2.NDArray | SChunk | blosc2.VLArray | C2Array:
         """Retrieve a node from the DictStore."""
         # Check map_tree first
         if key in self.map_tree:
             filepath = self.map_tree[key]
             if filepath in self.offsets:
                 offset = self.offsets[filepath]["offset"]
-                return blosc2.blosc2_ext.open(
+                opened = blosc2.blosc2_ext.open(
                     self.b2z_path,
                     mode="r",
                     offset=offset,
                     mmap_mode=self.mmap_mode,
                     dparams=self.dparams,
                 )
+                return _process_opened_object(opened)
             else:
                 urlpath = os.path.join(self.working_dir, filepath)
                 if os.path.exists(urlpath):
@@ -319,7 +340,7 @@ def __getitem__(self, key: str) -> blosc2.NDArray | SChunk | C2Array:
         # Fall back to EmbedStore
         return self._estore[key]
 
-    def get(self, key: str, default: Any = None) -> blosc2.NDArray | SChunk | C2Array | Any:
+    def get(self, key: str, default: Any = None) -> blosc2.NDArray | SChunk | blosc2.VLArray | C2Array | Any:
         """Retrieve a node, or default if not found."""
         try:
             return self[key]
diff --git a/src/blosc2/embed_store.py b/src/blosc2/embed_store.py
@@ -5,15 +5,20 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #######################################################################
 
+from __future__ import annotations
+
 import copy
-from collections.abc import Iterator, KeysView
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 import blosc2
 from blosc2.c2array import C2Array
-from blosc2.schunk import SChunk
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, KeysView
+
+    from blosc2.schunk import SChunk
 
 PROFILE = False  # Set to True to enable PROFILE prints in EmbedStore
 
@@ -168,7 +173,7 @@ def _ensure_capacity(self, needed_bytes: int) -> None:
             new_size = max(required_size, int(self._store.shape[0] * 1.5))
             self._store.resize((new_size,))
 
-    def __setitem__(self, key: str, value: blosc2.Array | SChunk) -> None:
+    def __setitem__(self, key: str, value: blosc2.Array | SChunk | blosc2.VLArray) -> None:
         """Add a node to the embed store."""
         if self.mode == "r":
             raise ValueError("Cannot set items in read-only mode.")
@@ -191,7 +196,7 @@ def __setitem__(self, key: str, value: blosc2.Array | SChunk) -> None:
             self._embed_map[key] = {"offset": offset, "length": data_len}
         self._save_metadata()
 
-    def __getitem__(self, key: str) -> blosc2.NDArray | SChunk:
+    def __getitem__(self, key: str) -> blosc2.NDArray | SChunk | blosc2.VLArray:
         """Retrieve a node from the embed store."""
         if key not in self._embed_map:
             raise KeyError(f"Key '{key}' not found in the embed store.")
@@ -207,7 +212,7 @@ def __getitem__(self, key: str) -> blosc2.NDArray | SChunk:
         # Use from_cframe so we can deserialize either an NDArray or an SChunk
         return blosc2.from_cframe(serialized_data, copy=True)
 
-    def get(self, key: str, default: Any = None) -> blosc2.NDArray | SChunk | Any:
+    def get(self, key: str, default: Any = None) -> blosc2.NDArray | SChunk | blosc2.VLArray | Any:
         """Retrieve a node, or default if not found."""
         return self[key] if key in self._embed_map else default
 
@@ -234,12 +239,12 @@ def keys(self) -> KeysView[str]:
         """Return all keys."""
         return self._embed_map.keys()
 
-    def values(self) -> Iterator[blosc2.NDArray | SChunk]:
+    def values(self) -> Iterator[blosc2.NDArray | SChunk | blosc2.VLArray]:
         """Iterate over all values."""
         for key in self._embed_map:
             yield self[key]
 
-    def items(self) -> Iterator[tuple[str, blosc2.NDArray | SChunk]]:
+    def items(self) -> Iterator[tuple[str, blosc2.NDArray | SChunk | blosc2.VLArray]]:
         """Iterate over (key, value) pairs."""
         for key in self._embed_map:
             yield key, self[key]
diff --git a/src/blosc2/tree_store.py b/src/blosc2/tree_store.py
@@ -5,6 +5,8 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #######################################################################
 
+from __future__ import annotations
+
 import contextlib
 import os
 from collections.abc import Iterator, MutableMapping
@@ -14,11 +16,11 @@
 
 import blosc2
 from blosc2.dict_store import DictStore
-from blosc2.schunk import SChunk
 
 if TYPE_CHECKING:
     from blosc2.c2array import C2Array
     from blosc2.ndarray import NDArray
+    from blosc2.schunk import SChunk
 
 
 class vlmetaProxy(MutableMapping):
@@ -29,7 +31,7 @@ class vlmetaProxy(MutableMapping):
     - Delegates iteration and length to the underlying vlmeta object.
     """
 
-    def __init__(self, tstore: "TreeStore", inner_vlmeta):
+    def __init__(self, tstore: TreeStore, inner_vlmeta):
         self._tstore = tstore
         self._inner = inner_vlmeta
 
@@ -224,7 +226,7 @@ def _validate_key(self, key: str) -> str:
 
         return key
 
-    def __setitem__(self, key: str, value: blosc2.Array | SChunk) -> None:
+    def __setitem__(self, key: str, value: blosc2.Array | SChunk | blosc2.VLArray) -> None:
         """Add a node with hierarchical key validation.
 
         Parameters
@@ -266,7 +268,7 @@ def __setitem__(self, key: str, value: blosc2.Array | SChunk) -> None:
         full_key = self._translate_key_to_full(key)
         super().__setitem__(full_key, value)
 
-    def __getitem__(self, key: str) -> "NDArray | C2Array | SChunk | TreeStore":
+    def __getitem__(self, key: str) -> NDArray | C2Array | SChunk | blosc2.VLArray | TreeStore:
         """Retrieve a node or subtree view.
 
         If the key points to a subtree (intermediate path with children),
@@ -280,7 +282,7 @@ def __getitem__(self, key: str) -> "NDArray | C2Array | SChunk | TreeStore":
 
         Returns
         -------
-        out : blosc2.NDArray or blosc2.C2Array or blosc2.SChunk or TreeStore
+        out : blosc2.NDArray or blosc2.C2Array or blosc2.SChunk or blosc2.VLArray or TreeStore
             The stored array/chunk if key is a leaf node, or a TreeStore subtree view
             if key is an intermediate path with children.
 
@@ -416,7 +418,7 @@ def __iter__(self) -> Iterator[str]:
         """Iterate over keys, excluding vlmeta keys."""
         return iter(self.keys())
 
-    def items(self) -> Iterator[tuple[str, "NDArray | C2Array | SChunk | TreeStore"]]:
+    def items(self) -> Iterator[tuple[str, NDArray | C2Array | SChunk | TreeStore]]:
         """Return key-value pairs in the current subtree view."""
         for key in self.keys():
             yield key, self[key]
@@ -575,7 +577,7 @@ def walk(self, path: str = "/", topdown: bool = True) -> Iterator[tuple[str, lis
             # Yield current level after children (post-order)
             yield path, children_dirs, leaf_nodes
 
-    def get_subtree(self, path: str) -> "TreeStore":
+    def get_subtree(self, path: str) -> TreeStore:
         """Create a subtree view with the specified path as root.
 
         Parameters
diff --git a/tests/test_dict_store.py b/tests/test_dict_store.py
@@ -223,6 +223,49 @@ def test_external_schunk_file_and_reopen():
         os.remove(path)
 
 
+def test_store_and_retrieve_vlarray_in_dict(tmp_path):
+    path = tmp_path / "test_dstore_vlarray_embed.b2z"
+    values = [{"name": "alpha", "count": 1}, None, ("tuple", 2), [1, "two", b"three"]]
+
+    vlarray = blosc2.VLArray()
+    vlarray.extend(values)
+
+    with DictStore(str(path), mode="w") as dstore:
+        dstore["/vlarray"] = vlarray
+        value = dstore["/vlarray"]
+        assert isinstance(value, blosc2.VLArray)
+        assert list(value) == values
+
+    with DictStore(str(path), mode="r") as dstore_read:
+        value = dstore_read["/vlarray"]
+        assert isinstance(value, blosc2.VLArray)
+        assert list(value) == values
+
+
+def test_external_vlarray_file_and_reopen(tmp_path):
+    ext_path = tmp_path / "ext_vlarray.b2frame"
+    path = tmp_path / "test_dstore_vlarray_external.b2z"
+    values = ["alpha", {"nested": True}, None, (1, 2, 3)]
+
+    vlarray = blosc2.VLArray(urlpath=str(ext_path), mode="w", contiguous=True)
+    vlarray.extend(values)
+    vlarray.vlmeta["description"] = "External VLArray"
+
+    with DictStore(str(path), mode="w", threshold=None) as dstore:
+        dstore["/dir1/vlarray_ext"] = vlarray
+        assert "/dir1/vlarray_ext" in dstore.map_tree
+        assert dstore.map_tree["/dir1/vlarray_ext"].endswith(".b2f")
+
+    with zipfile.ZipFile(path, "r") as zf:
+        assert "dir1/vlarray_ext.b2f" in zf.namelist()
+
+    with DictStore(str(path), mode="r") as dstore_read:
+        value = dstore_read["/dir1/vlarray_ext"]
+        assert isinstance(value, blosc2.VLArray)
+        assert list(value) == values
+        assert value.vlmeta["description"] == "External VLArray"
+
+
 def _digest_value(value):
     """Return a bytes digest of a stored value."""
     if isinstance(value, blosc2.SChunk):
diff --git a/tests/test_tree_store.py b/tests/test_tree_store.py
@@ -604,6 +604,56 @@ def test_schunk_support():
     os.remove("test_schunk.b2z")
 
 
+def test_vlarray_support():
+    """Test that TreeStore supports embedded VLArray objects."""
+    values = [{"name": "alpha", "count": 1}, None, ("tuple", 2), [1, "two", b"three"]]
+    with TreeStore("test_vlarray.b2z", mode="w") as tstore:
+        vlarray = blosc2.VLArray()
+        vlarray.extend(values)
+        tstore["/data/vlarray1"] = vlarray
+
+        retrieved = tstore["/data/vlarray1"]
+        assert isinstance(retrieved, blosc2.VLArray)
+        assert list(retrieved) == values
+
+        data_subtree = tstore["/data"]
+        assert isinstance(data_subtree, TreeStore)
+        assert set(data_subtree.keys()) == {"/vlarray1"}
+
+    with TreeStore("test_vlarray.b2z", mode="r") as tstore:
+        retrieved = tstore["/data/vlarray1"]
+        assert isinstance(retrieved, blosc2.VLArray)
+        assert list(retrieved) == values
+
+    os.remove("test_vlarray.b2z")
+
+
+def test_external_vlarray_support():
+    """Test that TreeStore supports external VLArray objects."""
+    ext_path = "ext_vlarray.b2frame"
+    values = ["alpha", {"nested": True}, None, (1, 2, 3)]
+    if os.path.exists(ext_path):
+        os.remove(ext_path)
+
+    vlarray = blosc2.VLArray(urlpath=ext_path, mode="w", contiguous=True)
+    vlarray.extend(values)
+    vlarray.vlmeta["description"] = "External VLArray for TreeStore"
+
+    with TreeStore("test_vlarray_external.b2z", mode="w", threshold=None) as tstore:
+        tstore["/data/vlarray_ext"] = vlarray
+        assert "/data/vlarray_ext" in tstore
+
+    with TreeStore("test_vlarray_external.b2z", mode="r") as tstore:
+        retrieved = tstore["/data/vlarray_ext"]
+        assert isinstance(retrieved, blosc2.VLArray)
+        assert list(retrieved) == values
+        assert retrieved.vlmeta["description"] == "External VLArray for TreeStore"
+
+    if os.path.exists(ext_path):
+        os.remove(ext_path)
+    os.remove("test_vlarray_external.b2z")
+
+
 def test_walk_topdown_argument_ordering():
     """Ensure walk supports topdown argument mimicking os.walk order semantics."""
     with TreeStore("test_walk_topdown.b2z", mode="w") as tstore:
