[WIP] Attempt to continue LRU cache for decoded chunks

zarr/core.py

@@ -179,6 +179,9 @@ def __init__(
         cache_attrs=True,
         partial_decompress=False,
         chunk_cache=None,
+        write_empty_chunks=True,
+        zarr_version=None,
+        meta_array=None,
     ):
         # N.B., expect at this point store is fully initialized with all
         # configuration metadata fully specified and normalized
@@ -207,6 +210,17 @@ def __init__(
         self._is_view = False
         self._partial_decompress = partial_decompress
         self._chunk_cache = chunk_cache
+        self._write_empty_chunks = write_empty_chunks
+        if meta_array is not None:
+            self._meta_array = np.empty_like(meta_array, shape=())
+        else:
+            self._meta_array = np.empty(())
+        self._version = zarr_version
+        if self._version == 3:
+            self._data_key_prefix = 'data/root/' + self._key_prefix
+            self._data_path = 'data/root/' + self._path
+            self._hierarchy_metadata = _get_hierarchy_metadata(store=self._store)
+            self._metadata_key_suffix = self._hierarchy_metadata['metadata_key_suffix']
 
         # initialize metadata
         self._load_metadata()
@@ -955,7 +969,7 @@ def _get_basic_selection_zd(self, selection, out=None, fields=None):
 
             except KeyError:
                 # chunk not initialized
-                chunk = np.zeros((), dtype=self._dtype)
+                chunk = np.zeros_like(self._meta_array, shape=(), dtype=self._dtype)
                 if self._fill_value is not None:
                     chunk.fill(self._fill_value)
 
@@ -1763,6 +1777,10 @@ def _set_basic_selection_zd(self, selection, value, fields=None):
 
         # remove chunk if write_empty_chunks is false and it only contains the fill value
         if (not self.write_empty_chunks) and all_equal(self.fill_value, chunk):
+            # invalidate value in cache
+            if self._chunk_cache is not None:
+                if ckey in self._chunk_cache:
+                    del self._chunk_cache[ckey]
             try:
                 del self.chunk_store[ckey]
                 return
@@ -2140,6 +2158,9 @@ def _chunk_setitem_nosync(self, chunk_coords, chunk_selection, value, fields=Non
         # attempt to delete chunk if it only contains the fill value
         if (not self.write_empty_chunks) and all_equal(self.fill_value, cdata):
             self._chunk_delitem(ckey)
+            if self._chunk_cache is not None:
+                #TODO christian
+                self._chunk_cache
         else:
             self.chunk_store[ckey] = self._encode_chunk(cdata)
 
@@ -2203,17 +2224,14 @@ def _process_for_setitem(self, ckey, chunk_selection, value, fields=None):
             else:
                 chunk[chunk_selection] = value
 
-        # encode 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)
+                chunk = self._decode_chunk(self._encode_chunk(chunk))
             self._chunk_cache[ckey] = np.copy(chunk)
 
-        return cdata
+        return chunk
 
     def _chunk_key(self, chunk_coords):
         if self._version == 3:
@@ -2556,6 +2574,13 @@ def _resize_nosync(self, *args):
                 except KeyError:
                     # chunk not initialized
                     pass
+                if self._chunk_cache is not None:
+                    try:
+                        del self._chunk_cache[key]
+                    except KeyError:
+                        # chunk not cached
+                        pass
+
             old_cdata_shape_working_list[idx_cdata] = min(val_old_cdata, val_new_cdata)
 
     def append(self, data, axis=0):

zarr/creation.py

@@ -20,8 +20,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, chunk_cache=None,
-           **kwargs):
+           object_codec=None, dimension_separator=None, chunk_cache=None, write_empty_chunks=True,
+           *, zarr_version=None, meta_array=None, **kwargs):
     """Create an array.
 
     Parameters
@@ -183,7 +183,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,
-              chunk_cache=chunk_cache)
+              chunk_cache=chunk_cache, write_empty_chunks=write_empty_chunks, 
+              meta_array=meta_array)
 
     return z
 
@@ -421,6 +422,10 @@ def open_array(
     storage_options=None,
     partial_decompress=False,
     chunk_cache=None,
+    write_empty_chunks=True,
+    *,
+    zarr_version=None,
+    dimension_separator=None,
     **kwargs
 ):
     """Open an array using file-mode-like semantics.
@@ -494,6 +499,17 @@ def open_array(
         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.
+        .. versionadded:: 2.11
+
+    zarr_version : {None, 2, 3}, optional
+        The zarr protocol version of the array to be opened. If None, it will
+        be inferred from ``store`` or ``chunk_store`` if they are provided,
+        otherwise defaulting to 2.
+    dimension_separator : {None, '.', '/'}, optional
+        Can be used to specify whether the array is in a flat ('.') or nested
+        ('/') format. If None, the appropriate value will be read from `store`
+        when present. Otherwise, defaults to '.' when ``zarr_version == 2``
+        and `/` otherwise.
 
     Returns
     -------
@@ -603,7 +619,8 @@ 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_cache=chunk_cache)
+              chunk_store=chunk_store, chunk_cache=chunk_cache, 
+              write_empty_chunks=write_empty_chunks)
 
     return z
 

zarr/hierarchy.py

@@ -121,7 +121,17 @@ class Group(MutableMapping):
     """
 
     def __init__(self, store, path=None, read_only=False, chunk_store=None,
-                 cache_attrs=True, synchronizer=None, chunk_cache=None):
+                 cache_attrs=True, synchronizer=None, chunk_cache=None, zarr_version=None, *,
+                 meta_array=None):
+        store: BaseStore = _normalize_store_arg(store, zarr_version=zarr_version)
+        if zarr_version is None:
+            zarr_version = getattr(store, '_store_version', DEFAULT_ZARR_VERSION)
+
+        if zarr_version != 2:
+            assert_zarr_v3_api_available()
+
+        if chunk_store is not None:
+            chunk_store: BaseStore = _normalize_store_arg(chunk_store, zarr_version=zarr_version)
         self._store = store
         self._chunk_store = chunk_store
         self._chunk_cache = chunk_cache
@@ -427,11 +437,23 @@ def __getitem__(self, item):
         if contains_array(self._store, path):
             return Array(self._store, read_only=self._read_only, path=path,
                          chunk_store=self._chunk_store, chunk_cache=self._chunk_cache,
-                         synchronizer=self._synchronizer, cache_attrs=self.attrs.cache)
-        elif contains_group(self._store, path):
+                         synchronizer=self._synchronizer, cache_attrs=self.attrs.cache,
+                         zarr_version=self._version, meta_array=self._meta_array)
+        elif contains_group(self._store, path, explicit_only=True):
             return Group(self._store, read_only=self._read_only, path=path,
-                         chunk_store=self._chunk_store, cache_attrs=self.attrs.cache,
-                         synchronizer=self._synchronizer, chunk_cache=self._chunk_cache)
+                         chunk_store=self._chunk_store, chunk_cache=self._chunk_cache, cache_attrs=self.attrs.cache,
+                         synchronizer=self._synchronizer, zarr_version=self._version,
+                         meta_array=self._meta_array)
+        elif self._version == 3:
+            implicit_group = meta_root + path + '/'
+            # non-empty folder in the metadata path implies an implicit group
+            if self._store.list_prefix(implicit_group):
+                return Group(self._store, read_only=self._read_only, path=path,
+                             chunk_store=self._chunk_store, cache_attrs=self.attrs.cache,
+                             synchronizer=self._synchronizer, zarr_version=self._version,
+                             meta_array=self._meta_array)
+            else:
+                raise KeyError(item)
         else:
             raise KeyError(item)
 
@@ -516,14 +538,39 @@ def groups(self):
         foo <class 'zarr.hierarchy.Group'>
 
         """
-        for key in sorted(listdir(self._store, self._path)):
-            path = self._key_prefix + key
-            if contains_group(self._store, path):
-                yield key, Group(self._store, path=path, read_only=self._read_only,
-                                 chunk_store=self._chunk_store,
-                                 chunk_cache=self._chunk_cache,
-                                 cache_attrs=self.attrs.cache,
-                                 synchronizer=self._synchronizer)
+        if self._version == 2:
+            for key in sorted(listdir(self._store, self._path)):
+                path = self._key_prefix + key
+                if contains_group(self._store, path, explicit_only=False):
+                    yield key, Group(
+                        self._store,
+                        path=path,
+                        read_only=self._read_only,
+                        chunk_store=self._chunk_store,
+                        chunk_cache=self._chunk_cache,
+                        cache_attrs=self.attrs.cache,
+                        synchronizer=self._synchronizer,
+                        zarr_version=self._version)
+
+        else:
+            dir_name = meta_root + self._path
+            group_sfx = '.group' + self._metadata_key_suffix
+            for key in sorted(listdir(self._store, dir_name)):
+                if key.endswith(group_sfx):
+                    key = key[:-len(group_sfx)]
+                path = self._key_prefix + key
+                if path.endswith(".array" + self._metadata_key_suffix):
+                    # skip array keys
+                    continue
+                if contains_group(self._store, path, explicit_only=False):
+                    yield key, Group(
+                        self._store,
+                        path=path,
+                        read_only=self._read_only,
+                        chunk_store=self._chunk_store,
+                        cache_attrs=self.attrs.cache,
+                        synchronizer=self._synchronizer,
+                        zarr_version=self._version)
 
     def array_keys(self, recurse=False):
         """Return an iterator over member names for arrays only.
@@ -816,7 +863,7 @@ def _create_group_nosync(self, name, overwrite=False):
 
         return Group(self._store, path=path, read_only=self._read_only,
                      chunk_store=self._chunk_store, cache_attrs=self.attrs.cache,
-                     synchronizer=self._synchronizer, chunk_cache=self._chunk_cache)
+                     synchronizer=self._synchronizer, chunk_cache=self._chunk_cache, zarr_version=self._version)
 
     def create_groups(self, *names, **kwargs):
         """Convenience method to create multiple groups in a single call."""
@@ -860,7 +907,8 @@ def _require_group_nosync(self, name, overwrite=False):
 
         return Group(self._store, path=path, read_only=self._read_only,
                      chunk_store=self._chunk_store, cache_attrs=self.attrs.cache,
-                     synchronizer=self._synchronizer, chunk_cache=self._chunk_cache)
+                     synchronizer=self._synchronizer, chunk_cache=self._chunk_cache, 
+                     zarr_version=self._version)
 
     def require_groups(self, *names):
         """Convenience method to require multiple groups in a single call."""
@@ -989,7 +1037,7 @@ def _require_dataset_nosync(self, name, shape, dtype=None, exact=False,
             a = Array(self._store, path=path, read_only=self._read_only,
                       chunk_store=self._chunk_store, synchronizer=synchronizer,
                       cache_metadata=cache_metadata, cache_attrs=cache_attrs,
-                      chunk_cache=self._chunk_cache)
+                      chunk_cache=self._chunk_cache, meta_array=self._meta_array)
             shape = normalize_shape(shape)
             if shape != a.shape:
                 raise TypeError('shape do not match existing array; expected {}, got {}'
@@ -1178,8 +1226,7 @@ def _normalize_store_arg(store, *, storage_options=None, mode="r",
 
 
 def group(store=None, overwrite=False, chunk_store=None,
-          cache_attrs=True, synchronizer=None, path=None,
-          chunk_cache=None):
+          cache_attrs=True, synchronizer=None, path=None, chunk_cache=None, *, zarr_version=None):
     """Create a group.
 
     Parameters
@@ -1244,11 +1291,11 @@ def group(store=None, overwrite=False, chunk_store=None,
 
     return Group(store, read_only=False, chunk_store=chunk_store,
                  cache_attrs=cache_attrs, synchronizer=synchronizer, path=path,
-                 chunk_cache=chunk_cache)
+                 zarr_version=zarr_version)
 
 
 def open_group(store=None, mode='a', cache_attrs=True, synchronizer=None, path=None,
-               chunk_store=None, storage_options=None, chunk_cache=None):
+               chunk_store=None, storage_options=None, chunk_cache=None, *, zarr_version=None, meta_array=None):
     """Open a group using file-mode-like semantics.
 
     Parameters
@@ -1351,4 +1398,4 @@ def open_group(store=None, mode='a', cache_attrs=True, synchronizer=None, path=N
 
     return Group(store, read_only=read_only, cache_attrs=cache_attrs,
                  synchronizer=synchronizer, path=path, chunk_store=chunk_store,
-                 chunk_cache=chunk_cache)
+                 chunk_cache=chunk_cache, zarr_version=zarr_version, meta_array=meta_array)

zarr/storage.py

@@ -2348,7 +2348,7 @@ def __delitem__(self, key):
         raise NotImplementedError  # pragma: no cover
 
 
-class LRUStoreCache(LRUMappingCache):
+class LRUStoreCache(Store, LRUMappingCache):
     """Storage class that implements a least-recently-used (LRU) cache layer over
     some other store. Intended primarily for use with stores that can be slow to
     access, e.g., remote stores that require network communication to store and
@@ -2385,12 +2385,17 @@ class LRUStoreCache(LRUMappingCache):
 
     """
 
-    def __init__(self, store, max_size):
+    def __init__(self, store: StoreLike, max_size: int):
         super(LRUStoreCache, self).__init__(max_size)
-        self._store = store
+        self._store: BaseStore = BaseStore._ensure_store(store)
+        self._max_size = max_size
+        self._current_size = 0
         self._keys_cache = None
         self._contains_cache = None
-        self._listdir_cache = dict()
+        self._listdir_cache: Dict[Path, Any] = dict()
+        self._values_cache: Dict[Path, Any] = OrderedDict()
+        self._mutex = Lock()
+        self.hits = self.misses = 0
 
     def __getstate__(self):
         return (self._store, self._max_size, self._current_size, self._keys_cache,
@@ -2430,6 +2435,30 @@ def listdir(self, path: Path = None):
     def getsize(self, path=None) -> int:
         return getsize(self._store, path=path)
 
+    def _pop_value(self):
+        # remove the first value from the cache, as this will be the least recently
+        # used value
+        _, v = self._values_cache.popitem(last=False)
+        return v
+
+    def _accommodate_value(self, value_size):
+        if self._max_size is None:
+            return
+        # ensure there is enough space in the cache for a new value
+        while self._current_size + value_size > self._max_size:
+            v = self._pop_value()
+            self._current_size -= buffer_size(v)
+
+    def _cache_value(self, key: Path, value):
+        # cache a value
+        value_size = buffer_size(value)
+        # check size of the value against max size, as if the value itself exceeds max
+        # size then we are never going to cache it
+        if self._max_size is None or value_size <= self._max_size:
+            self._accommodate_value(value_size)
+            self._values_cache[key] = value
+            self._current_size += value_size
+
     def invalidate(self):
         """Completely clear the cache."""
         with self._mutex:
@@ -2986,6 +3015,9 @@ def invalidate(self):
         with self._mutex:
             self._values_cache.clear()
 
+    def close(self):
+        pass
+
     def __getitem__(self, key):
         try:
             # try to obtain the value from the cache