add optional caching in open_zarr

ecmwf · May 2, 2024 · 7961f16 · 7961f16
1 parent 4426fd7
commit 7961f16
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/anemoi/datasets/data/stores.py b/src/anemoi/datasets/data/stores.py
@@ -102,7 +102,7 @@ def __contains__(self, key):
         return key in self.store
 
 
-def open_zarr(path, dont_fail=False):
+def open_zarr(path, dont_fail=False, cache=None):
     try:
         store = path
 
@@ -117,6 +117,9 @@ def open_zarr(path, dont_fail=False):
                 store = zarr.storage.DirectoryStore(store)
             store = DebugStore(store)
 
+        if cache is not None:
+            store = zarr.LRUStoreCache(store, max_size=cache)
+
         return zarr.convenience.open(store, "r")
     except zarr.errors.PathNotFoundError:
         if not dont_fail: