Repository: zarr-developers/zarr-python

Issue:
tifffile.ZarrTiffStore wrapped in zarr.LRUStoreCache fails to read any chunks
### Zarr version

2.16.0

### Numcodecs version

0.11.0

### Python Version

3.11.4

### Operating System

Windows

### Installation

py -m pip install -U zarr

### Description

Since zarr v2.15.0, the `tifffile.ZarrTiffStore` fails to read any chunks when wrapping the store in `zarr.LRUStoreCache`. The chunks returned by `LRUStoreCache` are always zeroed.

I don't understand yet what change in v2.15 causes the failure. It may be related to `ZarrTiffStore` not returning all chunk keys from the `keys` method because there may be many million keys and it would take significant resources to parse the TIFF file(s) for all keys upfront. Instead, the store relies on the `__contains__` method.

The following patch fixes the issue for me: 
```patch
diff --git a/zarr/storage.py b/zarr/storage.py
index 4f7b9905..c3260638 100644
--- a/zarr/storage.py
+++ b/zarr/storage.py
@@ -2436,7 +2436,7 @@ class LRUStoreCache(Store):
         with self._mutex:
             if self._contains_cache is None:
                 self._contains_cache = set(self._keys())
-            return key in self._contains_cache
+            return key in self._contains_cache or key in self._store

     def clear(self):
         self._store.clear()
```

### Steps to reproduce

```Python
import numpy
import zarr
import tifffile

data = numpy.random.rand(8, 8)
tifffile.imwrite('test.tif', data)
store = tifffile.imread('test.tif', aszarr=True)

try:
    # without cache
    zarray = zarr.open(store, mode='r')
    numpy.testing.assert_array_equal(zarray, data)
    # with cache
    cache = zarr.LRUStoreCache(store, max_size=2**10)
    zarray = zarr.open(cache, mode='r')
    numpy.testing.assert_array_equal(zarray, data)
finally:
    store.close()
```

### Additional output

```Python traceback
Traceback (most recent call last):
  File "D:\Dev\Python\tifffile\test_zarr_lrucache.py", line 16, in <module>
    numpy.testing.assert_array_equal(zarray, data)
  File "X:\Python311\Lib\site-packages\numpy\testing\_private\utils.py", line 920, in assert_array_equal
    assert_array_compare(operator.__eq__, x, y, err_msg=err_msg,
  File "X:\Python311\Lib\contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "X:\Python311\Lib\site-packages\numpy\testing\_private\utils.py", line 797, in assert_array_compare
    raise AssertionError(msg)
AssertionError:
Arrays are not equal

Mismatched elements: 64 / 64 (100%)
Max absolute difference: 0.99263945
Max relative difference: 1.
 x: array([[0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0.],...
 y: array([[0.992639, 0.580277, 0.454157, 0.055261, 0.664422, 0.82762 ,
        0.144393, 0.240842],
       [0.241279, 0.190615, 0.37115 , 0.078737, 0.706201, 0.593853,...
```

Comments:

Patch:
diff --git a/docs/release.rst b/docs/release.rst
index cf1400d3..188edd62 100644
--- a/docs/release.rst
+++ b/docs/release.rst
@@ -27,6 +27,9 @@ Maintenance
 * Add ``docs`` requirements to ``pyproject.toml``
   By :user:`John A. Kirkham <jakirkham>` :issue:`1494`.
 
+* Fixed caching issue in ``LRUStoreCache``.
+  By :user:`Mads R. B. Kristensen <madsbk>` :issue:`1499`.
+
 .. _release_2.16.0:
 
 2.16.0
diff --git a/zarr/_storage/v3.py b/zarr/_storage/v3.py
index 1a50265c..00dc085d 100644
--- a/zarr/_storage/v3.py
+++ b/zarr/_storage/v3.py
@@ -509,7 +509,7 @@ class LRUStoreCacheV3(RmdirV3, LRUStoreCache, StoreV3):
         self._max_size = max_size
         self._current_size = 0
         self._keys_cache = None
-        self._contains_cache = None
+        self._contains_cache = {}
         self._listdir_cache: Dict[Path, Any] = dict()
         self._values_cache: Dict[Path, Any] = OrderedDict()
         self._mutex = Lock()
diff --git a/zarr/storage.py b/zarr/storage.py
index 4f7b9905..b36f804e 100644
--- a/zarr/storage.py
+++ b/zarr/storage.py
@@ -2393,7 +2393,7 @@ class LRUStoreCache(Store):
         self._max_size = max_size
         self._current_size = 0
         self._keys_cache = None
-        self._contains_cache = None
+        self._contains_cache: Dict[Any, Any] = {}
         self._listdir_cache: Dict[Path, Any] = dict()
         self._values_cache: Dict[Path, Any] = OrderedDict()
         self._mutex = Lock()
@@ -2434,9 +2434,9 @@ class LRUStoreCache(Store):
 
     def __contains__(self, key):
         with self._mutex:
-            if self._contains_cache is None:
-                self._contains_cache = set(self._keys())
-            return key in self._contains_cache
+            if key not in self._contains_cache:
+                self._contains_cache[key] = key in self._store
+            return self._contains_cache[key]
 
     def clear(self):
         self._store.clear()
@@ -2506,7 +2506,7 @@ class LRUStoreCache(Store):
 
     def _invalidate_keys(self):
         self._keys_cache = None
-        self._contains_cache = None
+        self._contains_cache.clear()
         self._listdir_cache.clear()
 
     def _invalidate_value(self, key):


Test patch:
diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py
index 95570004..ca6a6c1a 100644
--- a/zarr/tests/test_storage.py
+++ b/zarr/tests/test_storage.py
@@ -2196,7 +2196,10 @@ class TestLRUStoreCache(StoreTests):
         assert keys == sorted(cache.keys())
         assert 1 == store.counter["keys"]
         assert foo_key in cache
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 1 == store.counter["__contains__", foo_key]
+        # the next check for `foo_key` is cached
+        assert foo_key in cache
+        assert 1 == store.counter["__contains__", foo_key]
         assert keys == sorted(cache)
         assert 0 == store.counter["__iter__"]
         assert 1 == store.counter["keys"]
@@ -2215,23 +2218,23 @@ class TestLRUStoreCache(StoreTests):
         keys = sorted(cache.keys())
         assert keys == [bar_key, baz_key, foo_key]
         assert 3 == store.counter["keys"]
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 1 == store.counter["__contains__", foo_key]
         assert 0 == store.counter["__iter__"]
         cache.invalidate_keys()
         keys = sorted(cache)
         assert keys == [bar_key, baz_key, foo_key]
         assert 4 == store.counter["keys"]
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 1 == store.counter["__contains__", foo_key]
         assert 0 == store.counter["__iter__"]
         cache.invalidate_keys()
         assert foo_key in cache
-        assert 5 == store.counter["keys"]
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 4 == store.counter["keys"]
+        assert 2 == store.counter["__contains__", foo_key]
         assert 0 == store.counter["__iter__"]
 
         # check these would get counted if called directly
         assert foo_key in store
-        assert 1 == store.counter["__contains__", foo_key]
+        assert 3 == store.counter["__contains__", foo_key]
         assert keys == sorted(store)
         assert 1 == store.counter["__iter__"]