New cache for the last block read

Author: FrancescAlted

src/blosc2/batch_store.py

@@ -34,6 +34,8 @@ def __init__(self, parent: BatchStore, nbatch: int, lazybatch: bytes) -> None:
         self._nbatch = nbatch
         self._lazybatch = lazybatch
         self._items: list[Any] | None = None
+        self._cached_block_index: int | None = None
+        self._cached_block: list[Any] | None = None
         self._nbytes, self._cbytes, self._nblocks = blosc2.get_cbuffer_sizes(lazybatch)
 
     def _normalize_index(self, index: int) -> int:
@@ -51,6 +53,14 @@ def _decode_items(self) -> list[Any]:
             self._items = [item for block in blocks for item in block]
         return self._items
 
+    def _get_block(self, block_index: int) -> list[Any]:
+        if self._cached_block_index == block_index and self._cached_block is not None:
+            return self._cached_block
+        block = msgpack_unpackb(self._parent.schunk.get_vlblock(self._nbatch, block_index))
+        self._cached_block_index = block_index
+        self._cached_block = block
+        return block
+
     def __getitem__(self, index: int | slice) -> Any | list[Any]:
         if isinstance(index, slice):
             items = self._decode_items()
@@ -64,7 +74,7 @@ def __getitem__(self, index: int | slice) -> Any | list[Any]:
             block_index, item_index = divmod(index, blocksize_max)
             if block_index >= self._nblocks:
                 raise IndexError("Batch index out of range")
-            block = msgpack_unpackb(self._parent.schunk.get_vlblock(self._nbatch, block_index))
+            block = self._get_block(block_index)
             try:
                 return block[item_index]
             except IndexError as exc:

tests/test_batch_store.py

@@ -202,6 +202,30 @@ def test_batchstore_get_vlblock_and_scalar_access():
     blosc2.remove_urlpath(urlpath)
 
 
+def test_batchstore_scalar_reads_cache_vlblocks():
+    barray = blosc2.BatchStore(blocksize_max=2)
+    barray.append([0, 1, 2, 3, 4])
+
+    batch = barray[0]
+    original_get_vlblock = barray.schunk.get_vlblock
+    calls = []
+
+    def wrapped_get_vlblock(nchunk, nblock):
+        calls.append((nchunk, nblock))
+        return original_get_vlblock(nchunk, nblock)
+
+    barray.schunk.get_vlblock = wrapped_get_vlblock
+    try:
+        assert batch[0] == 0
+        assert batch[1] == 1
+        assert batch[0] == 0
+        assert batch[2] == 2
+        assert batch[3] == 3
+        assert calls == [(0, 0), (0, 1)]
+    finally:
+        barray.schunk.get_vlblock = original_get_vlblock
+
+
 def test_batchstore_respects_explicit_use_dict_and_non_zstd():
     barray = blosc2.BatchStore(cparams={"codec": blosc2.Codec.LZ4, "clevel": 5})
     assert barray.cparams.codec == blosc2.Codec.LZ4