New iter_objects for iterating over objects in batch store

Author: FrancescAlted

bench/batch_store.py

@@ -125,6 +125,13 @@ def main() -> None:
     else:
         read_store = blosc2.BatchStore(urlpath=URLPATH, mode="r", contiguous=True, blocksize_max=BLOCKSIZE_MAX)
     samples, timings_ns = measure_random_reads(read_store)
+    t0 = time.perf_counter()
+    checksum = 0
+    nobjects = 0
+    for obj in read_store.iter_objects():
+        checksum += obj["blue"]
+        nobjects += 1
+    iter_time_s = time.perf_counter() - t0
 
     print()
     print("BatchStore benchmark")
@@ -138,6 +145,9 @@ def main() -> None:
     print(f"  mean: {statistics.fmean(timings_ns) / 1_000:.2f} us")
     print(f"  max:  {max(timings_ns) / 1_000:.2f} us")
     print(f"  min:  {min(timings_ns) / 1_000:.2f} us")
+    print(f"Object iteration via iter_objects(): {iter_time_s:.3f} s")
+    print(f"  per object: {iter_time_s * 1_000_000 / nobjects:.2f} us")
+    print(f"  checksum: {checksum}")
     print("Sample reads:")
     for timing_ns, batch_index, item_index, value in samples[:5]:
         print(f"  {timing_ns / 1_000:.2f} us -> read_store[{batch_index}][{item_index}] = {value}")

src/blosc2/batch_store.py

@@ -418,10 +418,17 @@ def __delitem__(self, index: int | slice) -> None:
     def __len__(self) -> int:
         return self.schunk.nchunks
 
-    def __iter__(self) -> Iterator[Batch]:
+    def iter_batches(self) -> Iterator[Batch]:
         for i in range(len(self)):
             yield self[i]
 
+    def iter_objects(self) -> Iterator[Any]:
+        for batch in self.iter_batches():
+            yield from batch
+
+    def __iter__(self) -> Iterator[Batch]:
+        yield from self.iter_batches()
+
     @property
     def meta(self):
         return self.schunk.meta
@@ -474,7 +481,7 @@ def info(self) -> InfoReporter:
     @property
     def info_items(self) -> list:
         """A list of tuples with summary information about this BatchStore."""
-        batch_sizes = [len(batch) for batch in self]
+        batch_sizes = [len(batch) for batch in self.iter_batches()]
         if batch_sizes:
             batch_stats = (
                 f"mean={statistics.fmean(batch_sizes):.2f}, max={max(batch_sizes)}, min={min(batch_sizes)}"
@@ -515,7 +522,7 @@ def copy(self, **kwargs: Any) -> BatchStore:
         if "storage" not in kwargs and len(self.vlmeta) > 0:
             for key, value in self.vlmeta.getall().items():
                 out.vlmeta[key] = value
-        out.extend(self)
+        out.extend(self.iter_batches())
         return out
 
     def __enter__(self) -> BatchStore:

tests/test_batch_store.py

@@ -48,9 +48,9 @@ def test_batchstore_roundtrip(contiguous, urlpath):
     assert len(barray) == len(BATCHES)
     assert barray.blocksize_max is not None
     assert 1 <= barray.blocksize_max <= len(BATCHES[0])
-    assert [batch[:] for batch in barray] == BATCHES
+    assert [batch[:] for batch in barray.iter_batches()] == BATCHES
     assert barray.append([1, 2]) == len(BATCHES) + 1
-    assert [batch[:] for batch in barray][-1] == [1, 2]
+    assert [batch[:] for batch in barray.iter_batches()][-1] == [1, 2]
 
     batch0 = barray[0]
     assert isinstance(batch0, blosc2.Batch)
@@ -78,13 +78,13 @@ def test_batchstore_roundtrip(contiguous, urlpath):
     del expected[2]
     del barray[-2]
     del expected[-2]
-    assert [batch[:] for batch in barray] == expected
+    assert [batch[:] for batch in barray.iter_batches()] == expected
 
     if urlpath is not None:
         reopened = blosc2.open(urlpath, mode="r")
         assert isinstance(reopened, blosc2.BatchStore)
         assert reopened.blocksize_max is None
-        assert [batch[:] for batch in reopened] == expected
+        assert [batch[:] for batch in reopened.iter_batches()] == expected
         with pytest.raises(ValueError):
             reopened.append(["nope"])
         with pytest.raises(ValueError):
@@ -105,12 +105,12 @@ def test_batchstore_roundtrip(contiguous, urlpath):
         reopened_rw = blosc2.open(urlpath, mode="a")
         reopened_rw[0] = ["changed", "batch", 0]
         expected[0] = ["changed", "batch", 0]
-        assert [batch[:] for batch in reopened_rw] == expected
+        assert [batch[:] for batch in reopened_rw.iter_batches()] == expected
 
         if contiguous:
             reopened_mmap = blosc2.open(urlpath, mode="r", mmap_mode="r")
             assert isinstance(reopened_mmap, blosc2.BatchStore)
-            assert [batch[:] for batch in reopened_mmap] == expected
+            assert [batch[:] for batch in reopened_mmap.iter_batches()] == expected
 
     blosc2.remove_urlpath(urlpath)
 
@@ -126,11 +126,11 @@ def test_batchstore_from_cframe():
 
     restored = blosc2.from_cframe(barray.to_cframe())
     assert isinstance(restored, blosc2.BatchStore)
-    assert [batch[:] for batch in restored] == expected
+    assert [batch[:] for batch in restored.iter_batches()] == expected
 
     restored2 = blosc2.from_cframe(barray.to_cframe())
     assert isinstance(restored2, blosc2.BatchStore)
-    assert [batch[:] for batch in restored2] == expected
+    assert [batch[:] for batch in restored2.iter_batches()] == expected
 
 
 def test_batchstore_info():
@@ -172,7 +172,7 @@ def test_batchstore_explicit_blocksize_max():
     assert barray.blocksize_max == 2
     barray.append([1, 2, 3])
     barray.append([4])
-    assert [batch[:] for batch in barray] == [[1, 2, 3], [4]]
+    assert [batch[:] for batch in barray.iter_batches()] == [[1, 2, 3], [4]]
 
 
 def test_batchstore_get_vlblock_and_scalar_access():
@@ -226,6 +226,15 @@ def wrapped_get_vlblock(nchunk, nblock):
         barray.schunk.get_vlblock = original_get_vlblock
 
 
+def test_batchstore_iter_objects():
+    barray = blosc2.BatchStore(blocksize_max=2)
+    batches = [[1, 2, 3], [4], [5, 6]]
+    barray.extend(batches)
+
+    assert [batch[:] for batch in barray] == batches
+    assert list(barray.iter_objects()) == [1, 2, 3, 4, 5, 6]
+
+
 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
@@ -287,7 +296,7 @@ def test_batchstore_constructor_kwargs():
     barray.extend(BATCHES)
 
     reopened = blosc2.BatchStore(urlpath=urlpath, mode="r", contiguous=True, mmap_mode="r")
-    assert [batch[:] for batch in reopened] == BATCHES
+    assert [batch[:] for batch in reopened.iter_batches()] == BATCHES
 
     blosc2.remove_urlpath(urlpath)
 
@@ -306,21 +315,21 @@ def test_batchstore_list_like_ops(contiguous, urlpath):
 
     barray = blosc2.BatchStore(storage=_storage(contiguous, urlpath))
     barray.extend([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
-    assert [batch[:] for batch in barray] == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+    assert [batch[:] for batch in barray.iter_batches()] == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
     assert barray.pop() == [7, 8, 9]
     assert barray.pop(0) == [1, 2, 3]
-    assert [batch[:] for batch in barray] == [[4, 5, 6]]
+    assert [batch[:] for batch in barray.iter_batches()] == [[4, 5, 6]]
 
     barray.clear()
     assert len(barray) == 0
-    assert [batch[:] for batch in barray] == []
+    assert [batch[:] for batch in barray.iter_batches()] == []
 
     barray.extend([["a", "b", "c"], ["d", "e", "f"]])
-    assert [batch[:] for batch in barray] == [["a", "b", "c"], ["d", "e", "f"]]
+    assert [batch[:] for batch in barray.iter_batches()] == [["a", "b", "c"], ["d", "e", "f"]]
 
     if urlpath is not None:
         reopened = blosc2.open(urlpath, mode="r")
-        assert [batch[:] for batch in reopened] == [["a", "b", "c"], ["d", "e", "f"]]
+        assert [batch[:] for batch in reopened.iter_batches()] == [["a", "b", "c"], ["d", "e", "f"]]
 
     blosc2.remove_urlpath(urlpath)
 
@@ -346,15 +355,15 @@ def test_batchstore_slices(contiguous, urlpath):
 
     barray[2:5] = [["a", "b", "c"], ["d", "e", "f"], ["g", "h", "i"]]
     expected[2:5] = [["a", "b", "c"], ["d", "e", "f"], ["g", "h", "i"]]
-    assert [batch[:] for batch in barray] == expected
+    assert [batch[:] for batch in barray.iter_batches()] == expected
 
     barray[1:6:2] = [[100, 101, 102], [103, 104, 105], [106, 107, 108]]
     expected[1:6:2] = [[100, 101, 102], [103, 104, 105], [106, 107, 108]]
-    assert [batch[:] for batch in barray] == expected
+    assert [batch[:] for batch in barray.iter_batches()] == expected
 
     del barray[::3]
     del expected[::3]
-    assert [batch[:] for batch in barray] == expected
+    assert [batch[:] for batch in barray.iter_batches()] == expected
 
     if urlpath is not None:
         reopened = blosc2.open(urlpath, mode="r")
@@ -392,14 +401,14 @@ def test_batchstore_copy():
     copied = original.copy(
         urlpath=copy_path, contiguous=False, cparams={"codec": blosc2.Codec.LZ4, "clevel": 5}
     )
-    assert [batch[:] for batch in copied] == [batch[:] for batch in original]
+    assert [batch[:] for batch in copied.iter_batches()] == [batch[:] for batch in original.iter_batches()]
     assert copied.urlpath == copy_path
     assert copied.schunk.contiguous is False
     assert copied.cparams.codec == blosc2.Codec.LZ4
     assert copied.cparams.clevel == 5
 
     inmem = original.copy()
-    assert [batch[:] for batch in inmem] == [batch[:] for batch in original]
+    assert [batch[:] for batch in inmem.iter_batches()] == [batch[:] for batch in original.iter_batches()]
     assert inmem.urlpath is None
 
     with pytest.raises(ValueError, match="meta should not be passed to copy"):
@@ -434,7 +443,7 @@ def test_batchstore_multithreaded_inner_vl(contiguous, nthreads):
     )
     barray.extend(batches)
 
-    assert [batch[:] for batch in barray] == batches
+    assert [batch[:] for batch in barray.iter_batches()] == batches
     assert [barray[i][:] for i in range(len(barray))] == batches
 
 
@@ -453,7 +462,7 @@ def test_batchstore_validation_errors():
         blosc2.BatchStore().pop()
     barray.extend([[1, 2, 3]])
     assert barray.append([2, 3]) == 2
-    assert [batch[:] for batch in barray] == [[1, 2, 3], [2, 3]]
+    assert [batch[:] for batch in barray.iter_batches()] == [[1, 2, 3], [2, 3]]
     with pytest.raises(NotImplementedError):
         barray.pop(slice(0, 1))
 
@@ -466,7 +475,7 @@ def test_batchstore_in_embed_store():
     estore["/batch"] = barray
     restored = estore["/batch"]
     assert isinstance(restored, blosc2.BatchStore)
-    assert [batch[:] for batch in restored] == BATCHES
+    assert [batch[:] for batch in restored.iter_batches()] == BATCHES
 
 
 def test_batchstore_in_dict_store():
@@ -481,6 +490,6 @@ def test_batchstore_in_dict_store():
     with blosc2.DictStore(path, mode="r") as dstore:
         restored = dstore["/batch"]
         assert isinstance(restored, blosc2.BatchStore)
-        assert [batch[:] for batch in restored] == BATCHES
+        assert [batch[:] for batch in restored.iter_batches()] == BATCHES
 
     blosc2.remove_urlpath(path)