LRU cache for decoded chunks

docs/api/storage.rst

@@ -39,6 +39,10 @@ Storage (``zarr.storage``)
 
 .. autoclass:: ConsolidatedMetadataStore
 
+.. autoclass:: LRUChunkCache
+
+    .. automethod:: invalidate
+
 .. autofunction:: init_array
 .. autofunction:: init_group
 .. autofunction:: contains_array

docs/tutorial.rst

@@ -849,10 +849,34 @@ store. E.g.::
     b'Hello from the cloud!'
     0.0009490990014455747
 
-If you are still experiencing poor performance with distributed/cloud storage,
-please raise an issue on the GitHub issue tracker with any profiling data you
-can provide, as there may be opportunities to optimise further either within
-Zarr or within the mapping interface to the storage.
+The above :class:`zarr.storage.LRUStoreCache` wraps any Zarr storage class, and stores
+encoded chunks. So every time cache is accessed, the chunk has to be decoded. For cases
+where decoding is computationally expensive, Zarr also provides a
+:class:`zarr.storage.LRUChunkCache` which can store decoded chunks, e.g.::
+
+    >>> import zarr
+    >>> from numcodecs import LZMA
+    >>> import numpy as np
+    >>> store = zarr.DictStore()
+    >>> z = zarr.array(np.random.randn(1000000).reshape(1000,1000), chunks=(100,100),
+    ...                store=store, compressor=LZMA())
+    >>> from timeit import timeit
+    >>> # data access without cache
+    ... timeit('z[:]', number=1, globals=globals())  # doctest: +SKIP
+    0.6703157789888792
+    >>> z_with_cache = zarr.Array(store=store, chunk_cache=zarr.LRUChunkCache(max_size=None))
+    >>> # first data access about the same as without cache
+    ... timeit('z_with_cache[:]', number=1, globals=globals())  # doctest: +SKIP
+    0.681269913999131
+    >>> # second time accesses the decoded chunks in the cache
+    ... timeit('z_with_cache[:]', number=1, globals=globals())  # doctest: +SKIP
+    0.007617925992235541
+
+
+If you are still experiencing poor performance with distributed/cloud storage, please
+raise an issue on the GitHub issue tracker with any profiling data you can provide, as
+there may be opportunities to optimise further either within Zarr or within the mapping
+interface to the storage.
 
 IO with ``fsspec``
 ~~~~~~~~~~~~~~~~~~

zarr/__init__.py

@@ -13,7 +13,7 @@
 from zarr.storage import (ABSStore, DBMStore, DictStore, DirectoryStore,
                           LMDBStore, LRUStoreCache, MemoryStore, MongoDBStore,
                           NestedDirectoryStore, RedisStore, SQLiteStore,
-                          TempStore, ZipStore)
+                          TempStore, ZipStore, LRUChunkCache)
 from zarr.sync import ProcessSynchronizer, ThreadSynchronizer
 from zarr.version import version as __version__
 

zarr/core.py

@@ -76,6 +76,15 @@ class Array:
         read and decompressed when possible.
 
         .. versionadded:: 2.7
+    chunk_cache: MutableMapping, optional
+        Mapping to store decoded chunks for caching. Can be used in repeated
+        chunk access scenarios when decoding of data is computationally
+        expensive.
+        NOTE: When using the write cache feature with object arrays(i.e.
+        when dtype of array is 'object' and when writing to the array with
+        chunk_cache provided) could result in a slight slowdown as some
+        dtypes, like VLenArray, have to go through the encode-decode phase
+        before having the correct dtype.
 
     Attributes
     ----------
@@ -137,6 +146,7 @@ def __init__(
         cache_metadata=True,
         cache_attrs=True,
         partial_decompress=False,
+        chunk_cache=None,
     ):
         # N.B., expect at this point store is fully initialized with all
         # configuration metadata fully specified and normalized
@@ -153,6 +163,7 @@ def __init__(
         self._cache_metadata = cache_metadata
         self._is_view = False
         self._partial_decompress = partial_decompress
+        self._chunk_cache = chunk_cache
 
         # initialize metadata
         self._load_metadata()
@@ -793,19 +804,33 @@ def _get_basic_selection_zd(self, selection, out=None, fields=None):
         if selection not in ((), (Ellipsis,)):
             err_too_many_indices(selection, ())
 
-        try:
-            # obtain encoded data for chunk
-            ckey = self._chunk_key((0,))
-            cdata = self.chunk_store[ckey]
+        # obtain key for chunk
+        ckey = self._chunk_key((0,))
 
-        except KeyError:
-            # chunk not initialized
-            chunk = np.zeros((), dtype=self._dtype)
-            if self._fill_value is not None:
-                chunk.fill(self._fill_value)
+        # setup variable to hold decoded chunk
+        chunk = None
 
-        else:
-            chunk = self._decode_chunk(cdata)
+        # check for cached chunk
+        if self._chunk_cache is not None:
+            chunk = self._chunk_cache.get(ckey)
+
+        if chunk is None:
+            try:
+                # obtain encoded data for chunk
+                cdata = self.chunk_store[ckey]
+
+            except KeyError:
+                # chunk not initialized
+                chunk = np.zeros((), dtype=self._dtype)
+                if self._fill_value is not None:
+                    chunk.fill(self._fill_value)
+
+            else:
+                chunk = self._decode_chunk(cdata)
+
+            # cache decoded chunk
+            if self._chunk_cache is not None:
+                self._chunk_cache[ckey] = chunk
 
         # handle fields
         if fields:
@@ -1588,6 +1613,12 @@ def _set_basic_selection_zd(self, selection, value, fields=None):
         cdata = self._encode_chunk(chunk)
         self.chunk_store[ckey] = cdata
 
+        if self._chunk_cache is not None:
+            # ensure cached chunk has been round tripped through encode-decode if dtype=object
+            if self.dtype == object:
+                chunk = self._decode_chunk(cdata)
+            self._chunk_cache[ckey] = chunk
+
     def _set_basic_selection_nd(self, selection, value, fields=None):
         # implementation of __setitem__ for array with at least one dimension
 
@@ -1647,6 +1678,7 @@ def _set_selection(self, indexer, value, fields=None):
 
                 # put data
                 self._chunk_setitem(chunk_coords, chunk_selection, chunk_value, fields=fields)
+
         else:
             lchunk_coords, lchunk_selection, lout_selection = zip(*indexer)
             chunk_values = []
@@ -1669,6 +1701,18 @@ def _set_selection(self, indexer, value, fields=None):
             self._chunk_setitems(lchunk_coords, lchunk_selection, chunk_values,
                                  fields=fields)
 
+    def _select_and_set_out(self, fields, chunk, chunk_selection, drop_axes,
+                            out, out_selection):
+        # select data from chunk
+        if fields:
+            chunk = chunk[fields]
+        tmp = chunk[chunk_selection]
+        if drop_axes:
+            tmp = np.squeeze(tmp, axis=drop_axes)
+
+        # store selected data in output
+        out[out_selection] = tmp
+
     def _process_chunk(
         self,
         out,
@@ -1678,6 +1722,7 @@ def _process_chunk(
         out_is_ndarray,
         fields,
         out_selection,
+        ckey,
         partial_read_decode=False,
     ):
         """Take binary data from storage and fill output array"""
@@ -1741,16 +1786,12 @@ def _process_chunk(
         except ArrayIndexError:
             cdata = cdata.read_full()
         chunk = self._decode_chunk(cdata)
+        if self._chunk_cache is not None:
+            # cache the decoded chunk
+            self._chunk_cache[ckey] = chunk
 
-        # select data from chunk
-        if fields:
-            chunk = chunk[fields]
-        tmp = chunk[chunk_selection]
-        if drop_axes:
-            tmp = np.squeeze(tmp, axis=drop_axes)
-
-        # store selected data in output
-        out[out_selection] = tmp
+        self._select_and_set_out(fields, chunk, chunk_selection, drop_axes,
+                                 out, out_selection)
 
     def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection,
                        drop_axes=None, fields=None):
@@ -1783,22 +1824,38 @@ def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection,
         # obtain key for chunk
         ckey = self._chunk_key(chunk_coords)
 
-        try:
-            # obtain compressed data for chunk
-            cdata = self.chunk_store[ckey]
+        # setup variable to hold decoded chunk
+        chunk = None
 
-        except KeyError:
-            # chunk not initialized
-            if self._fill_value is not None:
-                if fields:
-                    fill_value = self._fill_value[fields]
-                else:
-                    fill_value = self._fill_value
-                out[out_selection] = fill_value
+        # check for cached chunk
+        if self._chunk_cache is not None:
+            try:
+                chunk = self._chunk_cache[ckey]
+                self._select_and_set_out(fields, chunk, chunk_selection,
+                                         drop_axes, out, out_selection)
+            except KeyError:
+                pass
 
-        else:
-            self._process_chunk(out, cdata, chunk_selection, drop_axes,
-                                out_is_ndarray, fields, out_selection)
+        if chunk is None:
+
+            try:
+                # obtain compressed data for chunk
+                cdata = self.chunk_store[ckey]
+
+            except KeyError:
+                # chunk not initialized
+                if self._fill_value is not None:
+                    if fields:
+                        fill_value = self._fill_value[fields]
+                    else:
+                        fill_value = self._fill_value
+                    out[out_selection] = fill_value
+                return
+
+            else:
+                self._process_chunk(out, cdata, chunk_selection, drop_axes,
+                                    out_is_ndarray, fields, out_selection,
+                                    ckey)
 
     def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection,
                         drop_axes=None, fields=None):
@@ -1842,6 +1899,7 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection,
                     out_is_ndarray,
                     fields,
                     out_select,
+                    ckey,
                     partial_read_decode=partial_read_decode,
                 )
             else:
@@ -1947,7 +2005,16 @@ def _process_for_setitem(self, ckey, chunk_selection, value, fields=None):
                 chunk[chunk_selection] = value
 
         # encode chunk
-        return self._encode_chunk(chunk)
+        cdata = self._encode_chunk(chunk)
+
+        # cache the chunk
+        if self._chunk_cache is not None:
+            # ensure cached chunk has been round tripped through encode-decode if dtype=object
+            if self.dtype == object:
+                chunk = self._decode_chunk(cdata)
+            self._chunk_cache[ckey] = np.copy(chunk)
+
+        return cdata
 
     def _chunk_key(self, chunk_coords):
         return self._key_prefix + '.'.join(map(str, chunk_coords))

zarr/creation.py

@@ -19,7 +19,7 @@ def create(shape, chunks=True, dtype=None, compressor='default',
            fill_value=0, order='C', store=None, synchronizer=None,
            overwrite=False, path=None, chunk_store=None, filters=None,
            cache_metadata=True, cache_attrs=True, read_only=False,
-           object_codec=None, **kwargs):
+           object_codec=None, chunk_cache=None, **kwargs):
     """Create an array.
 
     Parameters
@@ -51,6 +51,15 @@ def create(shape, chunks=True, dtype=None, compressor='default',
     chunk_store : MutableMapping, optional
         Separate storage for chunks. If not provided, `store` will be used
         for storage of both chunks and metadata.
+    chunk_cache: MutableMapping, optional
+        Mapping to store decoded chunks for caching. Can be used in repeated
+        chunk access scenarios when decoding of data is computationally
+        expensive.
+        NOTE: When using the write cache feature with object arrays(i.e.
+        when dtype of array is 'object' and when writing to the array with
+        chunk_cache provided) could result in a slight slowdown as some
+        dtypes, like VLenArray, have to go through the encode-decode phase
+        before having the correct dtype.
     filters : sequence of Codecs, optional
         Sequence of filters to use to encode chunk data prior to compression.
     cache_metadata : bool, optional
@@ -124,7 +133,8 @@ def create(shape, chunks=True, dtype=None, compressor='default',
 
     # instantiate array
     z = Array(store, path=path, chunk_store=chunk_store, synchronizer=synchronizer,
-              cache_metadata=cache_metadata, cache_attrs=cache_attrs, read_only=read_only)
+              cache_metadata=cache_metadata, cache_attrs=cache_attrs, read_only=read_only,
+              chunk_cache=chunk_cache)
 
     return z
 
@@ -380,6 +390,7 @@ def open_array(
     chunk_store=None,
     storage_options=None,
     partial_decompress=False,
+    chunk_cache=None,
     **kwargs
 ):
     """Open an array using file-mode-like semantics.
@@ -436,6 +447,15 @@ def open_array(
         read and decompressed when possible.
 
         .. versionadded:: 2.7
+    chunk_cache: MutableMapping, optional
+        Mapping to store decoded chunks for caching. Can be used in repeated
+        chunk access scenarios when decoding of data is computationally
+        expensive.
+        NOTE: When using the write cache feature with object arrays(i.e.
+        when dtype of array is 'object' and when writing to the array with
+        chunk_cache provided) could result in a slight slowdown as some
+        dtypes, like VLenArray, have to go through the encode-decode phase
+        before having the correct dtype.
 
     Returns
     -------
@@ -525,7 +545,7 @@ def open_array(
     # instantiate array
     z = Array(store, read_only=read_only, synchronizer=synchronizer,
               cache_metadata=cache_metadata, cache_attrs=cache_attrs, path=path,
-              chunk_store=chunk_store)
+              chunk_store=chunk_store, chunk_cache=chunk_cache)
 
     return z