LRU cache for decoded chunks

zarr/core.py

@@ -103,7 +103,8 @@ class Array(object):
     """
 
     def __init__(self, store, path=None, read_only=False, chunk_store=None,
-                 synchronizer=None, cache_metadata=True, cache_attrs=True):
+                 synchronizer=None, cache_metadata=True, cache_attrs=True,
+                 chunk_cache=None):
         # N.B., expect at this point store is fully initialized with all
         # configuration metadata fully specified and normalized
 
@@ -118,6 +119,7 @@ def __init__(self, store, path=None, read_only=False, chunk_store=None,
         self._synchronizer = synchronizer
         self._cache_metadata = cache_metadata
         self._is_view = False
+        self._chunk_cache = chunk_cache
 
         # initialize metadata
         self._load_metadata()
@@ -1562,8 +1564,21 @@ def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection,
         ckey = self._chunk_key(chunk_coords)
 
         try:
+
+            cdata = None
+            chunk_was_cached = False
+
+            # first try getting from cache (if one has been provided)
+            if self._chunk_cache is not None:
+                try:
+                    cdata = self._chunk_cache[ckey]
+                    chunk_was_cached = True
+                except KeyError:
+                    pass
+
             # obtain compressed data for chunk
-            cdata = self.chunk_store[ckey]
+            if not chunk_was_cached:
+                cdata = self.chunk_store[ckey]
 
         except KeyError:
             # chunk not initialized
@@ -1593,19 +1608,30 @@ def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection,
                     # contiguous, so we can decompress directly from the chunk
                     # into the destination array
 
-                    if self._compressor:
+                    if chunk_was_cached:
+                        np.copyto(dest, cdata)
+                    elif self._compressor:
                         self._compressor.decode(cdata, dest)
+                        if self._chunk_cache is not None:
+                            self._chunk_cache[ckey] = np.copy(dest)
                     else:
                         if isinstance(cdata, np.ndarray):
                             chunk = cdata.view(self._dtype)
                         else:
                             chunk = np.frombuffer(cdata, dtype=self._dtype)
                         chunk = chunk.reshape(self._chunks, order=self._order)
                         np.copyto(dest, chunk)
+                        if self._chunk_cache is not None:
+                            self._chunk_cache[ckey] = np.copy(chunk)
                     return
 
             # decode chunk
-            chunk = self._decode_chunk(cdata)
+            if not chunk_was_cached:
+                chunk = self._decode_chunk(cdata)
+                if self._chunk_cache is not None:
+                    self._chunk_cache[ckey] = np.copy(chunk)
+            else:
+                chunk = cdata
 
             # select data from chunk
             if fields:
@@ -1714,6 +1740,8 @@ def _chunk_setitem_nosync(self, chunk_coords, chunk_selection, value, fields=Non
             else:
                 chunk[chunk_selection] = value
 
+        if self._chunk_cache is not None:
+            self._chunk_cache[ckey] = np.copy(chunk)
         # encode chunk
         cdata = self._encode_chunk(chunk)
 

zarr/storage.py

@@ -1883,3 +1883,121 @@ def __delitem__(self, key):
         with self._mutex:
             self._invalidate_keys()
             self._invalidate_value(key)
+
+
+class ChunkCache(MutableMapping):
+    """Storage class that implements a least-recently-used (LRU) cache for array chunks.
+    Intended primarily for use with stores that can be slow to access, e.g., remote stores that
+    require network communication to store and retrieve data.
+
+    Parameters
+    ----------
+    max_size : int
+        The maximum size that the cache may grow to, in number of bytes. Provide `None`
+        if you would like the cache to have unlimited size.
+
+    """
+
+    def __init__(self, max_size):
+        self._max_size = max_size
+        self._current_size = 0
+        self._values_cache = OrderedDict()
+        self._mutex = Lock()
+        self.hits = self.misses = 0
+
+    def __getstate__(self):
+        return (self._max_size, self._current_size,
+                self._values_cache, self.hits,
+                self.misses)
+
+    def __setstate__(self, state):
+        (self._store, self._max_size, self._current_size,
+         self._values_cache, self.hits,
+         self.misses) = state
+        self._mutex = Lock()
+
+    def __len__(self):
+        return len(self._keys())
+
+    def __iter__(self):
+        return self.keys()
+
+    def __contains__(self, key):
+        with self._mutex:
+            return key in self._keys()
+
+    def clear(self):
+        self.invalidate()
+
+    def keys(self):
+        with self._mutex:
+            return iter(self._keys())
+
+    def _keys(self):
+        return self._values_cache.keys()
+
+    def _pop_value(self):
+        # remove the first value from the cache, as this will be the least recently
+        # used value
+        _, v = self._values_cache.popitem(last=False)
+        return v
+
+    def _accommodate_value(self, value_size):
+        if self._max_size is None:
+            return
+        # ensure there is enough space in the cache for a new value
+        while self._current_size + value_size > self._max_size:
+            v = self._pop_value()
+            self._current_size -= buffer_size(v)
+
+    def _cache_value(self, key, value):
+        # cache a value
+        value_size = buffer_size(value)
+        # check size of the value against max size, as if the value itself exceeds max
+        # size then we are never going to cache it
+        if self._max_size is None or value_size <= self._max_size:
+            self._accommodate_value(value_size)
+            self._values_cache[key] = value
+            self._current_size += value_size
+
+    def invalidate(self):
+        """Completely clear the cache."""
+        with self._mutex:
+            self._values_cache.clear()
+
+    def invalidate_values(self):
+        """Clear the values cache."""
+        with self._mutex:
+            self._values_cache.clear()
+
+    def _invalidate_value(self, key):
+        if key in self._values_cache:
+            value = self._values_cache.pop(key)
+            self._current_size -= buffer_size(value)
+
+    def __getitem__(self, key):
+        try:
+            # try to obtain the value from the cache
+            with self._mutex:
+                value = self._values_cache[key]
+                # cache hit if no KeyError is raised
+                self.hits += 1
+                # treat the end as most recently used
+                OrderedDict_move_to_end(self._values_cache, key)
+
+        except KeyError:
+            # cache miss
+            with self._mutex:
+                self.misses += 1
+            raise KeyError
+
+        return value
+
+    def __setitem__(self, key, value):
+        with self._mutex:
+            self._invalidate_value(key)
+            self._cache_value(key, value)
+
+    def __delitem__(self, key):
+        with self._mutex:
+            self._invalidate_value(key)

zarr/tests/test_core.py

@@ -16,7 +16,7 @@
 
 from zarr.storage import (DirectoryStore, init_array, init_group, NestedDirectoryStore,
                           DBMStore, LMDBStore, atexit_rmtree, atexit_rmglob,
-                          LRUStoreCache)
+                          LRUStoreCache, ChunkCache)
 from zarr.core import Array
 from zarr.errors import PermissionError
 from zarr.compat import PY2, text_type, binary_type
@@ -1698,3 +1698,48 @@ def create_array(read_only=False, **kwargs):
         init_array(store, **kwargs)
         return Array(store, read_only=read_only, cache_metadata=cache_metadata,
                      cache_attrs=cache_attrs)
+
+
+class TestArrayWithChunkCache(TestArray):
+
+    @staticmethod
+    def create_array(read_only=False, **kwargs):
+        store = dict()
+        kwargs.setdefault('compressor', Zlib(level=1))
+        cache_metadata = kwargs.pop('cache_metadata', True)
+        cache_attrs = kwargs.pop('cache_attrs', True)
+        init_array(store, **kwargs)
+        return Array(store, read_only=read_only, cache_metadata=cache_metadata,
+                     cache_attrs=cache_attrs, chunk_cache=ChunkCache(max_size=None))
+
+    @staticmethod
+    def create_array_with_cache(read_only=False, **kwargs):
+        store = dict()
+        kwargs.setdefault('compressor', Zlib(level=1))
+        cache_metadata = kwargs.pop('cache_metadata', True)
+        cache_attrs = kwargs.pop('cache_attrs', True)
+        init_array(store, **kwargs)
+        cache = ChunkCache(max_size=None)
+        return Array(store, read_only=read_only, cache_metadata=cache_metadata,
+                     cache_attrs=cache_attrs, chunk_cache=cache), cache
+
+    def test_hit_miss(self):
+        a = np.arange(100).reshape((10, 10))
+        z, cache = self.create_array_with_cache(shape=a.shape, chunks=(10,1), dtype=a.dtype)
+
+        # test write cache
+        z[:] = a
+        assert cache.misses == 0 and cache.hits == 0
+        _ = z[:]
+        assert cache.misses == 0 and cache.hits == 10
+
+        cache.clear()
+        cache.misses = 0
+        cache.hits = 0
+
+        # test read cache
+        assert cache.misses == 0 and cache.hits == 0
+        _ = z[:]
+        assert cache.misses == 10 and cache.hits == 0
+        _ = z[:]
+        assert cache.misses == 10 and cache.hits == 10