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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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
 

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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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
 

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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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
 

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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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
 

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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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
 

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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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
 

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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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
 

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

@@ -851,10 +851,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
  ----------
@@ -138,6 +147,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
@@ -154,6 +164,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()
@@ -795,19 +806,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:
@@ -1590,6 +1615,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
 
@@ -1649,6 +1680,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 = []
@@ -1671,6 +1703,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,
@@ -1680,6 +1724,7 @@ def _process_chunk(
  out_is_ndarray,
  fields,
  out_selection,
+ ckey,
  partial_read_decode=False,
  ):
  """Take binary data from storage and fill output array"""
@@ -1743,16 +1788,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):
@@ -1785,22 +1826,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):
@@ -1844,6 +1901,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:
@@ -1949,7 +2007,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 + self._dimension_separator.join(map(str, chunk_coords))

zarr/creation.py

@@ -21,7 +21,8 @@ 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, dimension_separator=None, **kwargs):
+ object_codec=None, dimension_separator=None, chunk_cache=None,
+ **kwargs):
  """Create an array.
 
  Parameters
@@ -53,6 +54,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
@@ -142,7 +152,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
 
@@ -400,6 +411,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.
@@ -456,6 +468,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
  -------
@@ -545,7 +566,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