Getitems: support meta_array

zarr/_storage/store.py

@@ -2,7 +2,9 @@
 import os
 from collections.abc import MutableMapping
 from string import ascii_letters, digits
-from typing import Any, List, Mapping, Optional, Union
+from typing import Any, Sequence, List, Mapping, Optional, Union
+
+from numcodecs.ndarray_like import NDArrayLike
 
 from zarr.meta import Metadata2, Metadata3
 from zarr.util import normalize_storage_path
@@ -129,6 +131,32 @@ def _ensure_store(store: Any):
             f"wrap it in Zarr.storage.KVStore. Got {store}"
         )
 
+    def getitems(self, keys: Sequence[str], meta_array: NDArrayLike) -> Mapping[str, Any]:
+        """Retrieve data from multiple keys.
+
+        Parameters
+        ----------
+        keys : Iterable[str]
+            The keys to retrieve
+        meta_array : array-like
+            An array instance to use for determining the output type. For now, this is
+            only a hint and can be ignore by the implementation, in which case the type
+            of the output is the same as calling __getitem__() for each key in keys.
+
+        Returns
+        -------
+        Mapping
+            A collection mapping the input keys to their results.
+
+        Developer Notes
+        ---------------
+        This default implementation use __getitem__() to read each key sequential and
+        ignores the meta_array argument. Overwrite this method to implement concurrent
+        reads of multiple keys and/or to utilize the meta_array argument.
+        """
+
+        return {k: self[k] for k in keys if k in self}
+
 
 class Store(BaseStore):
     """Abstract store class used by implementations following the Zarr v2 spec.

zarr/core.py

@@ -1257,20 +1257,13 @@ def _get_selection(self, indexer, out=None, fields=None):
         else:
             check_array_shape('out', out, out_shape)
 
-        # iterate over chunks
-        if not hasattr(self.chunk_store, "getitems") or \
-           any(map(lambda x: x == 0, self.shape)):
-            # sequentially get one key at a time from storage
-            for chunk_coords, chunk_selection, out_selection in indexer:
-
-                # load chunk selection into output array
-                self._chunk_getitem(chunk_coords, chunk_selection, out, out_selection,
-                                    drop_axes=indexer.drop_axes, fields=fields)
-        else:
-            # allow storage to get multiple items at once
+        if math.prod(out_shape) > 0:
+            # get chunks
             lchunk_coords, lchunk_selection, lout_selection = zip(*indexer)
-            self._chunk_getitems(lchunk_coords, lchunk_selection, out, lout_selection,
-                                 drop_axes=indexer.drop_axes, fields=fields)
+            self._chunk_getitems(
+                lchunk_coords, lchunk_selection, out, lout_selection,
+                drop_axes=indexer.drop_axes, fields=fields
+            )
 
         if out.shape:
             return out
@@ -1930,76 +1923,44 @@ def _process_chunk(
         # store selected data in output
         out[out_selection] = tmp
 
-    def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection,
-                       drop_axes=None, fields=None):
-        """Obtain part or whole of a chunk.
+    def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection,
+                        drop_axes=None, fields=None):
+        """Obtain part or whole of chunks.
 
         Parameters
         ----------
-        chunk_coords : tuple of ints
-            Indices of the chunk.
-        chunk_selection : selection
-            Location of region within the chunk to extract.
+        chunk_coords : list of tuple of ints
+            Indices of the chunks.
+        chunk_selection : list of selections
+            Location of region within the chunks to extract.
         out : ndarray
             Array to store result in.
-        out_selection : selection
-            Location of region within output array to store results in.
+        out_selection : list of selections
+            Location of regions within output array to store results in.
         drop_axes : tuple of ints
             Axes to squeeze out of the chunk.
         fields
             TODO
-
         """
-        out_is_ndarray = True
-        try:
-            out = ensure_ndarray_like(out)
-        except TypeError:
-            out_is_ndarray = False
-
-        assert len(chunk_coords) == len(self._cdata_shape)
-
-        # obtain key for chunk
-        ckey = self._chunk_key(chunk_coords)
 
-        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
-
-        else:
-            self._process_chunk(out, cdata, chunk_selection, drop_axes,
-                                out_is_ndarray, fields, out_selection)
-
-    def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection,
-                        drop_axes=None, fields=None):
-        """As _chunk_getitem, but for lists of chunks
-
-        This gets called where the storage supports ``getitems``, so that
-        it can decide how to fetch the keys, allowing concurrency.
-        """
         out_is_ndarray = True
         try:
             out = ensure_ndarray_like(out)
         except TypeError:  # pragma: no cover
             out_is_ndarray = False
 
+        # Keys to retrieve
         ckeys = [self._chunk_key(ch) for ch in lchunk_coords]
+
+        partial_read_decode = False
+        # Check if we can do a partial read
         if (
             self._partial_decompress
             and self._compressor
             and self._compressor.codec_id == "blosc"
             and hasattr(self._compressor, "decode_partial")
             and not fields
             and self.dtype != object
-            and hasattr(self.chunk_store, "getitems")
         ):
             partial_read_decode = True
             cdatas = {
@@ -2008,8 +1969,8 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection,
                 if ckey in self.chunk_store
             }
         else:
-            partial_read_decode = False
-            cdatas = self.chunk_store.getitems(ckeys, on_error="omit")
+            cdatas = self.chunk_store.getitems(ckeys, meta_array=self._meta_array)
+
         for ckey, chunk_select, out_select in zip(ckeys, lchunk_selection, lout_selection):
             if ckey in cdatas:
                 self._process_chunk(

zarr/storage.py

@@ -31,7 +31,7 @@
 from os import scandir
 from pickle import PicklingError
 from threading import Lock, RLock
-from typing import Optional, Union, List, Tuple, Dict, Any
+from typing import Sequence, Mapping, Optional, Union, List, Tuple, Dict, Any
 import uuid
 import time
 
@@ -41,6 +41,7 @@
     ensure_text,
     ensure_contiguous_ndarray_like
 )
+from numcodecs.ndarray_like import NDArrayLike
 from numcodecs.registry import codec_registry
 
 from zarr.errors import (
@@ -1363,7 +1364,10 @@ def _normalize_key(self, key):
 
         return key.lower() if self.normalize_keys else key
 
-    def getitems(self, keys, **kwargs):
+    def getitems(
+        self, keys: Sequence[str], meta_array: NDArrayLike
+    ) -> Mapping[str, Any]:
+
         keys_transformed = [self._normalize_key(key) for key in keys]
         results = self.map.getitems(keys_transformed, on_error="omit")
         # The function calling this method may not recognize the transformed keys