From c268c6379895ca74ef1132b3d3163a1b41ef8312 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 12 Sep 2022 11:39:19 +0200 Subject: [PATCH 01/19] Use _chunk_getitems() always --- zarr/core.py | 118 +++++++++++++++++++-------------------------------- 1 file changed, 44 insertions(+), 74 deletions(-) diff --git a/zarr/core.py b/zarr/core.py index e5b2045160..5caf7ccd06 100644 --- a/zarr/core.py +++ b/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,86 +1923,63 @@ 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] - 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 = { - ckey: PartialReadBuffer(ckey, self.chunk_store) - for ckey in ckeys - if ckey in self.chunk_store - } + + # Check if we should retrieve multiple keys at a time + use_getitems = ( + hasattr(self.chunk_store, "getitems") + and math.prod(self.shape) > 0 + ) + + partial_read_decode = False + if use_getitems: + # 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 + ): + partial_read_decode = True + cdatas = { + ckey: PartialReadBuffer(ckey, self.chunk_store) + for ckey in ckeys + if ckey in self.chunk_store + } + else: + cdatas = self.chunk_store.getitems(ckeys, on_error="omit") else: - partial_read_decode = False - cdatas = self.chunk_store.getitems(ckeys, on_error="omit") + cdatas = {k: self.chunk_store[k] for k in ckeys if k in self.chunk_store} + for ckey, chunk_select, out_select in zip(ckeys, lchunk_selection, lout_selection): if ckey in cdatas: self._process_chunk( From c7023f984a917c94879392b32913ffa6628dd5b8 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 12 Sep 2022 14:52:30 +0200 Subject: [PATCH 02/19] Implement getitems() always --- zarr/_storage/store.py | 41 ++++++++++++++++++++++++++++++++++++++- zarr/core.py | 44 ++++++++++++++++++------------------------ 2 files changed, 59 insertions(+), 26 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 9e265cf383..6b39c52bda 100644 --- a/zarr/_storage/store.py +++ b/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, Iterable, 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,43 @@ def _ensure_store(store: Any): f"wrap it in Zarr.storage.KVStore. Got {store}" ) + def getitems( + self, keys: Iterable[str], meta_array: NDArrayLike, *, on_error: str = "omit" + ) -> 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. + on_error : str, optional + The policy on how to handle exceptions when retrieving keys. For now, the + only supported policy is "omit", which means that failing keys are omitted + from the returned result. + + 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. + """ + + # Please overwrite `getitems` to support non-default values of `on_error` + if on_error != "omit": + raise ValueError(f"{self.__class__} doesn't support on_error='{on_error}'") + + # Please overwrite `getitems` to support non-default values of `meta_array` + 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. diff --git a/zarr/core.py b/zarr/core.py index 5caf7ccd06..2c6223cb97 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -1949,36 +1949,30 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, except TypeError: # pragma: no cover out_is_ndarray = False + if math.prod(self.shape) <= 0: + return + # Keys to retrieve ckeys = [self._chunk_key(ch) for ch in lchunk_coords] - # Check if we should retrieve multiple keys at a time - use_getitems = ( - hasattr(self.chunk_store, "getitems") - and math.prod(self.shape) > 0 - ) - partial_read_decode = False - if use_getitems: - # 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 - ): - partial_read_decode = True - cdatas = { - ckey: PartialReadBuffer(ckey, self.chunk_store) - for ckey in ckeys - if ckey in self.chunk_store - } - else: - cdatas = self.chunk_store.getitems(ckeys, on_error="omit") + # 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 + ): + partial_read_decode = True + cdatas = { + ckey: PartialReadBuffer(ckey, self.chunk_store) + for ckey in ckeys + if ckey in self.chunk_store + } else: - cdatas = {k: self.chunk_store[k] for k in ckeys if k in self.chunk_store} + 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: From 89fa599f60f477be4f163c4466e9420de9c7851b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 12 Sep 2022 15:39:29 +0200 Subject: [PATCH 03/19] FSStore.getitems(): accept meta_array and on_error --- zarr/storage.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/zarr/storage.py b/zarr/storage.py index f5459990ba..571680cc49 100644 --- a/zarr/storage.py +++ b/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 Iterable, 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,9 +1364,15 @@ def _normalize_key(self, key): return key.lower() if self.normalize_keys else key - def getitems(self, keys, **kwargs): + def getitems( + self, keys: Iterable[str], meta_array: NDArrayLike, *, on_error: str = "omit" + ) -> Mapping[str, Any]: + + if on_error != "omit": + raise ValueError(f"{self.__class__} doesn't support on_error='{on_error}'") + keys_transformed = [self._normalize_key(key) for key in keys] - results = self.map.getitems(keys_transformed, on_error="omit") + results = self.map.getitems(keys_transformed, on_error=on_error) # The function calling this method may not recognize the transformed keys # So we send the values returned by self.map.getitems back into the original key space. return {keys[keys_transformed.index(rk)]: rv for rk, rv in results.items()} From f05ee3ae739a5e1ccfd5303c008dce3c7c9f6919 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 12 Sep 2022 15:44:14 +0200 Subject: [PATCH 04/19] getitems(): handle on_error="omit" --- zarr/_storage/store.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 6b39c52bda..3ff5e59555 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -165,8 +165,13 @@ def getitems( if on_error != "omit": raise ValueError(f"{self.__class__} doesn't support on_error='{on_error}'") - # Please overwrite `getitems` to support non-default values of `meta_array` - return {k: self[k] for k in keys if k in self} + ret = {} + for k in keys: + try: + ret[k] = self[k] + except Exception: + pass # Omit keys that fails + return ret class Store(BaseStore): From 0eed377be58e46480d1b6a12453510247962996f Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 13 Sep 2022 08:20:56 +0200 Subject: [PATCH 05/19] Removed the `on_error argument` --- zarr/_storage/store.py | 20 ++------------------ zarr/storage.py | 7 ++----- 2 files changed, 4 insertions(+), 23 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 3ff5e59555..c9bd272e2c 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -131,9 +131,7 @@ def _ensure_store(store: Any): f"wrap it in Zarr.storage.KVStore. Got {store}" ) - def getitems( - self, keys: Iterable[str], meta_array: NDArrayLike, *, on_error: str = "omit" - ) -> Mapping[str, Any]: + def getitems(self, keys: Iterable[str], meta_array: NDArrayLike) -> Mapping[str, Any]: """Retrieve data from multiple keys. Parameters @@ -144,10 +142,6 @@ def getitems( 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. - on_error : str, optional - The policy on how to handle exceptions when retrieving keys. For now, the - only supported policy is "omit", which means that failing keys are omitted - from the returned result. Returns ------- @@ -161,17 +155,7 @@ def getitems( reads of multiple keys and/or to utilize the meta_array argument. """ - # Please overwrite `getitems` to support non-default values of `on_error` - if on_error != "omit": - raise ValueError(f"{self.__class__} doesn't support on_error='{on_error}'") - - ret = {} - for k in keys: - try: - ret[k] = self[k] - except Exception: - pass # Omit keys that fails - return ret + return {k: self[k] for k in keys if k in self} class Store(BaseStore): diff --git a/zarr/storage.py b/zarr/storage.py index 571680cc49..21976f72b5 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -1365,14 +1365,11 @@ def _normalize_key(self, key): return key.lower() if self.normalize_keys else key def getitems( - self, keys: Iterable[str], meta_array: NDArrayLike, *, on_error: str = "omit" + self, keys: Iterable[str], meta_array: NDArrayLike ) -> Mapping[str, Any]: - if on_error != "omit": - raise ValueError(f"{self.__class__} doesn't support on_error='{on_error}'") - keys_transformed = [self._normalize_key(key) for key in keys] - results = self.map.getitems(keys_transformed, on_error=on_error) + results = self.map.getitems(keys_transformed, on_error="omit") # The function calling this method may not recognize the transformed keys # So we send the values returned by self.map.getitems back into the original key space. return {keys[keys_transformed.index(rk)]: rv for rk, rv in results.items()} From e578051430792f0572920dc52fd373e0b26f92f1 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 13 Sep 2022 11:18:41 +0200 Subject: [PATCH 06/19] remove redundant check --- zarr/core.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/zarr/core.py b/zarr/core.py index 2c6223cb97..3e55914ac3 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -1949,9 +1949,6 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, except TypeError: # pragma: no cover out_is_ndarray = False - if math.prod(self.shape) <= 0: - return - # Keys to retrieve ckeys = [self._chunk_key(ch) for ch in lchunk_coords] From 03c97a8a8c9f703e9790994a2a4d4e1aa6697b8c Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 13 Sep 2022 11:49:36 +0200 Subject: [PATCH 07/19] getitems(): use Sequence instead of Iterable --- zarr/_storage/store.py | 4 ++-- zarr/storage.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index c9bd272e2c..84431dd41f 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -2,7 +2,7 @@ import os from collections.abc import MutableMapping from string import ascii_letters, digits -from typing import Any, Iterable, List, Mapping, Optional, Union +from typing import Any, Sequence, List, Mapping, Optional, Union from numcodecs.ndarray_like import NDArrayLike @@ -131,7 +131,7 @@ def _ensure_store(store: Any): f"wrap it in Zarr.storage.KVStore. Got {store}" ) - def getitems(self, keys: Iterable[str], meta_array: NDArrayLike) -> Mapping[str, Any]: + def getitems(self, keys: Sequence[str], meta_array: NDArrayLike) -> Mapping[str, Any]: """Retrieve data from multiple keys. Parameters diff --git a/zarr/storage.py b/zarr/storage.py index 21976f72b5..d5810ab038 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -31,7 +31,7 @@ from os import scandir from pickle import PicklingError from threading import Lock, RLock -from typing import Iterable, Mapping, Optional, Union, List, Tuple, Dict, Any +from typing import Sequence, Mapping, Optional, Union, List, Tuple, Dict, Any import uuid import time @@ -1365,7 +1365,7 @@ def _normalize_key(self, key): return key.lower() if self.normalize_keys else key def getitems( - self, keys: Iterable[str], meta_array: NDArrayLike + self, keys: Sequence[str], meta_array: NDArrayLike ) -> Mapping[str, Any]: keys_transformed = [self._normalize_key(key) for key in keys] From 1d2b6ea2ff956601f85334a3bcb451b6dc66c9fc Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 7 Oct 2022 15:27:04 +0200 Subject: [PATCH 08/19] Typo Co-authored-by: Josh Moore --- zarr/_storage/store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 84431dd41f..081d65a49f 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -150,7 +150,7 @@ def getitems(self, keys: Sequence[str], meta_array: NDArrayLike) -> Mapping[str, Developer Notes --------------- - This default implementation use __getitem__() to read each key sequential and + This default implementation uses __getitem__() to read each key sequentially and ignores the meta_array argument. Overwrite this method to implement concurrent reads of multiple keys and/or to utilize the meta_array argument. """ From 05be1d4ddc853b3bfa10ace7d939a8ac48d2bec6 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 12 Oct 2022 09:51:05 +0200 Subject: [PATCH 09/19] Introduce a contexts argument --- zarr/_storage/store.py | 22 +++++++++++----------- zarr/core.py | 5 ++++- zarr/storage.py | 3 +-- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 081d65a49f..f49aaf397c 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -4,8 +4,6 @@ from string import ascii_letters, digits 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 @@ -131,28 +129,30 @@ 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]: + def getitems( + self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} + ) -> 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. + contexts: Mapping[str, Mapping] + A mapping of keys to their context. Each context is a mapping of store + specific information. E.g. a context could be a dict telling the store + the preferred output array type: `{"meta_array": cupy.empty(())}` Returns ------- Mapping A collection mapping the input keys to their results. - Developer Notes - --------------- + Notes + ----- This default implementation uses __getitem__() to read each key sequentially and - ignores the meta_array argument. Overwrite this method to implement concurrent - reads of multiple keys and/or to utilize the meta_array argument. + ignores contexts. Overwrite this method to implement concurrent reads of multiple + keys and/or to utilize the contexts. """ return {k: self[k] for k in keys if k in self} diff --git a/zarr/core.py b/zarr/core.py index 3e55914ac3..1765169bc8 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -1969,7 +1969,10 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, if ckey in self.chunk_store } else: - cdatas = self.chunk_store.getitems(ckeys, meta_array=self._meta_array) + contexts = {} + if not isinstance(self._meta_array, np.ndarray): + contexts = {k: {"meta_array": self._meta_array} for k in ckeys} + cdatas = self.chunk_store.getitems(ckeys, contexts=contexts) for ckey, chunk_select, out_select in zip(ckeys, lchunk_selection, lout_selection): if ckey in cdatas: diff --git a/zarr/storage.py b/zarr/storage.py index d5810ab038..8d5f5c0b64 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -41,7 +41,6 @@ ensure_text, ensure_contiguous_ndarray_like ) -from numcodecs.ndarray_like import NDArrayLike from numcodecs.registry import codec_registry from zarr.errors import ( @@ -1365,7 +1364,7 @@ def _normalize_key(self, key): return key.lower() if self.normalize_keys else key def getitems( - self, keys: Sequence[str], meta_array: NDArrayLike + self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} ) -> Mapping[str, Any]: keys_transformed = [self._normalize_key(key) for key in keys] From 5513d6f0d1243c436d7336b6744a90934d10ed13 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 12 Oct 2022 10:20:44 +0200 Subject: [PATCH 10/19] CountingDict: impl. getitems() --- zarr/tests/util.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/zarr/tests/util.py b/zarr/tests/util.py index faa2f35d25..2b5649c8b7 100644 --- a/zarr/tests/util.py +++ b/zarr/tests/util.py @@ -1,6 +1,7 @@ import collections import os import tempfile +from typing import Any, Mapping, Sequence from zarr.storage import Store from zarr._storage.v3 import StoreV3 @@ -42,6 +43,13 @@ def __delitem__(self, key): self.counter['__delitem__', key] += 1 del self.wrapped[key] + def getitems( + self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} + ) -> Mapping[str, Any]: + for key in keys: + self.counter['__getitem__', key] += 1 + return {k: self.wrapped[k] for k in keys if k in self.wrapped} + class CountingDictV3(CountingDict, StoreV3): pass From 40ecba0c89a619022deccfda9707ca2cc7dc18b3 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 12 Oct 2022 11:04:52 +0200 Subject: [PATCH 11/19] added test_getitems() --- zarr/tests/test_storage.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py index 39d4b5988d..e8cacf3f2a 100644 --- a/zarr/tests/test_storage.py +++ b/zarr/tests/test_storage.py @@ -9,6 +9,7 @@ import tempfile from contextlib import contextmanager from pickle import PicklingError +from typing import Any, Mapping, Sequence from zipfile import ZipFile import numpy as np @@ -2572,3 +2573,29 @@ def test_meta_prefix_6853(): fixtures = group(store=DirectoryStore(str(fixture))) assert list(fixtures.arrays()) + + +def test_getitems_contexts(): + + class MyStore(CountingDict): + def __init__(self): + super().__init__() + self.last_contexts = None + + def getitems( + self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} + ) -> Mapping[str, Any]: + self.last_contexts = contexts + return {k: self.wrapped[k] for k in keys if k in self.wrapped} + + store = MyStore() + z = zarr.create(shape=(10,), store=store) + + # By default, not contexts are given to the store's getitems() + z[0] + assert len(store.last_contexts) == 0 + + # Setting a non-default meta_array, will create contexts for the store's getitems() + z._meta_array = "my_meta_array" + z[0] + assert store.last_contexts == {'0': {'meta_array': 'my_meta_array'}} From 81549f568bbfa4d27e369033c705a5eb18fbe639 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 4 Nov 2022 10:22:49 +0100 Subject: [PATCH 12/19] Introduce Context --- zarr/_storage/store.py | 5 +++-- zarr/context.py | 19 +++++++++++++++++++ zarr/storage.py | 3 ++- zarr/tests/test_storage.py | 3 ++- zarr/tests/util.py | 3 ++- 5 files changed, 28 insertions(+), 5 deletions(-) create mode 100644 zarr/context.py diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index f49aaf397c..95c90fb807 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -6,6 +6,7 @@ from zarr.meta import Metadata2, Metadata3 from zarr.util import normalize_storage_path +from zarr.context import Context # v2 store keys array_meta_key = '.zarray' @@ -130,7 +131,7 @@ def _ensure_store(store: Any): ) def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} + self, keys: Sequence[str], contexts: Mapping[str, Context] = {} ) -> Mapping[str, Any]: """Retrieve data from multiple keys. @@ -138,7 +139,7 @@ def getitems( ---------- keys : Iterable[str] The keys to retrieve - contexts: Mapping[str, Mapping] + contexts: Mapping[str, Context] A mapping of keys to their context. Each context is a mapping of store specific information. E.g. a context could be a dict telling the store the preferred output array type: `{"meta_array": cupy.empty(())}` diff --git a/zarr/context.py b/zarr/context.py new file mode 100644 index 0000000000..09d7f4ef0a --- /dev/null +++ b/zarr/context.py @@ -0,0 +1,19 @@ + +from typing import TypedDict + +from numcodecs.compat import NDArrayLike + + +class Context(TypedDict, total=False): + """ A context for component specific information + + All keys are optional. Any component reading the context must provide + a default value for any key not in the context. + + Items + ----- + meta_array : array-like, optional + An array-like instance to use for determining the preferred output + array type. + """ + meta_array: NDArrayLike diff --git a/zarr/storage.py b/zarr/storage.py index 5dbafb6f7c..2e30ebac91 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -42,6 +42,7 @@ ensure_contiguous_ndarray_like ) from numcodecs.registry import codec_registry +from zarr.context import Context from zarr.errors import ( MetadataError, @@ -1362,7 +1363,7 @@ def _normalize_key(self, key): return key.lower() if self.normalize_keys else key def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} + self, keys: Sequence[str], contexts: Mapping[str, Context] = {} ) -> Mapping[str, Any]: keys_transformed = [self._normalize_key(key) for key in keys] diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py index e8cacf3f2a..ae87f5dda3 100644 --- a/zarr/tests/test_storage.py +++ b/zarr/tests/test_storage.py @@ -21,6 +21,7 @@ import zarr from zarr._storage.store import _get_hierarchy_metadata from zarr.codecs import BZ2, AsType, Blosc, Zlib +from zarr.context import Context from zarr.convenience import consolidate_metadata from zarr.errors import ContainsArrayError, ContainsGroupError, MetadataError from zarr.hierarchy import group @@ -2583,7 +2584,7 @@ def __init__(self): self.last_contexts = None def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} + self, keys: Sequence[str], contexts: Mapping[str, Context] = {} ) -> Mapping[str, Any]: self.last_contexts = contexts return {k: self.wrapped[k] for k in keys if k in self.wrapped} diff --git a/zarr/tests/util.py b/zarr/tests/util.py index 2b5649c8b7..0c6c6c01a0 100644 --- a/zarr/tests/util.py +++ b/zarr/tests/util.py @@ -2,6 +2,7 @@ import os import tempfile from typing import Any, Mapping, Sequence +from zarr.context import Context from zarr.storage import Store from zarr._storage.v3 import StoreV3 @@ -44,7 +45,7 @@ def __delitem__(self, key): del self.wrapped[key] def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Mapping] = {} + self, keys: Sequence[str], contexts: Mapping[str, Context] = {} ) -> Mapping[str, Any]: for key in keys: self.counter['__getitem__', key] += 1 From 2bfe68a6fa848480570f32a3f6ae694128b49a85 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 4 Nov 2022 10:36:27 +0100 Subject: [PATCH 13/19] doc --- zarr/context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zarr/context.py b/zarr/context.py index 09d7f4ef0a..83fbaafa9b 100644 --- a/zarr/context.py +++ b/zarr/context.py @@ -8,7 +8,7 @@ class Context(TypedDict, total=False): """ A context for component specific information All keys are optional. Any component reading the context must provide - a default value for any key not in the context. + a default implementation in the case a key cannot be found. Items ----- From 02fc80d05a2d781778a63988c18d50839117088d Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 13 Mar 2023 13:25:57 +0100 Subject: [PATCH 14/19] support the new get_partial_values() method --- zarr/_storage/store.py | 16 +++++++++++++--- zarr/core.py | 28 +++++----------------------- zarr/storage.py | 2 +- zarr/tests/test_storage.py | 8 ++------ zarr/tests/util.py | 2 +- 5 files changed, 22 insertions(+), 34 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index e56f712cba..279b8899be 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -133,7 +133,7 @@ def _ensure_store(store: Any): ) def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Context] = {} + self, keys: Sequence[str], contexts: Mapping[str, Context] ) -> Mapping[str, Any]: """Retrieve data from multiple keys. @@ -157,8 +157,15 @@ def getitems( ignores contexts. Overwrite this method to implement concurrent reads of multiple keys and/or to utilize the contexts. """ - - return {k: self[k] for k in keys if k in self} + keys_in_self = [k for k in keys if k in self] # Ignoring keys not in `self` + if hasattr(self, "get_partial_values"): + # Optimization, if `get_partial_values` is available, we use it to retrieve + # all the values in a single call. + values = self.get_partial_values([(k, (0, None)) for k in keys_in_self]) + else: + # Otherwise, we just retrieve each key sequentially. + values = [self[k] for k in keys_in_self] + return dict(zip(keys_in_self, values)) class Store(BaseStore): @@ -545,6 +552,9 @@ def __len__(self): def supports_efficient_get_partial_values(self): return self.inner_store.supports_efficient_get_partial_values + def getitems(self, keys, contexts): + return self.inner_store.getitems(keys, contexts) + def get_partial_values(self, key_ranges): return self.inner_store.get_partial_values(key_ranges) diff --git a/zarr/core.py b/zarr/core.py index 199d2fb58f..37afb2c3cb 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -1275,27 +1275,14 @@ def _get_selection(self, indexer, out=None, fields=None): check_array_shape('out', out, out_shape) # iterate over chunks - if ( - math.prod(out_shape) > 0 and - not hasattr(self.chunk_store, "getitems") and not ( - hasattr(self.chunk_store, "get_partial_values") and - self.chunk_store.supports_efficient_get_partial_values - ) - ) 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: + if math.prod(out_shape) > 0: # allow storage to get multiple items at once 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 ) - if out.shape: return out else: @@ -2029,15 +2016,10 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, } else: partial_read_decode = False - if not hasattr(self.chunk_store, "getitems"): - values = self.chunk_store.get_partial_values([(ckey, (0, None)) for ckey in ckeys]) - cdatas = {key: value for key, value in zip(ckeys, values) if value is not None} - else: - contexts = {} - if not isinstance(self._meta_array, np.ndarray): - contexts = {k: {"meta_array": self._meta_array} for k in ckeys} - cdatas = self.chunk_store.getitems(ckeys, contexts=contexts, on_error="omit") - + contexts = {} + if not isinstance(self._meta_array, np.ndarray): + contexts = {k: {"meta_array": self._meta_array} for k in ckeys} + cdatas = self.chunk_store.getitems(ckeys, contexts) for ckey, chunk_select, out_select in zip(ckeys, lchunk_selection, lout_selection): if ckey in cdatas: diff --git a/zarr/storage.py b/zarr/storage.py index 0892d4a93b..2138d55c4e 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -1382,7 +1382,7 @@ def _normalize_key(self, key): return key.lower() if self.normalize_keys else key def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Context] = {} + self, keys: Sequence[str], contexts: Mapping[str, Context] ) -> Mapping[str, Any]: keys_transformed = [self._normalize_key(key) for key in keys] diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py index c6736c1109..eac7454ddb 100644 --- a/zarr/tests/test_storage.py +++ b/zarr/tests/test_storage.py @@ -9,7 +9,6 @@ import tempfile from contextlib import contextmanager from pickle import PicklingError -from typing import Any, Mapping, Sequence from zipfile import ZipFile import numpy as np @@ -21,7 +20,6 @@ import zarr from zarr._storage.store import _get_hierarchy_metadata from zarr.codecs import BZ2, AsType, Blosc, Zlib -from zarr.context import Context from zarr.convenience import consolidate_metadata from zarr.errors import ContainsArrayError, ContainsGroupError, MetadataError from zarr.hierarchy import group @@ -2595,11 +2593,9 @@ def __init__(self): super().__init__() self.last_contexts = None - def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Context] = {} - ) -> Mapping[str, Any]: + def getitems(self, keys, contexts): self.last_contexts = contexts - return {k: self.wrapped[k] for k in keys if k in self.wrapped} + return super().getitems(keys, contexts) store = MyStore() z = zarr.create(shape=(10,), store=store) diff --git a/zarr/tests/util.py b/zarr/tests/util.py index 0c6c6c01a0..0db00c9fb8 100644 --- a/zarr/tests/util.py +++ b/zarr/tests/util.py @@ -45,7 +45,7 @@ def __delitem__(self, key): del self.wrapped[key] def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Context] = {} + self, keys: Sequence[str], contexts: Mapping[str, Context] ) -> Mapping[str, Any]: for key in keys: self.counter['__getitem__', key] += 1 From d0afcdec2764f983d82aa965ccfc96395b9745f8 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 13 Mar 2023 17:00:41 +0100 Subject: [PATCH 15/19] Resolve conflict with get_partial_values() --- zarr/_storage/store.py | 13 +------------ zarr/core.py | 6 +++++- zarr/tests/test_storage_v3.py | 2 ++ 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 279b8899be..94c3999924 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -157,15 +157,7 @@ def getitems( ignores contexts. Overwrite this method to implement concurrent reads of multiple keys and/or to utilize the contexts. """ - keys_in_self = [k for k in keys if k in self] # Ignoring keys not in `self` - if hasattr(self, "get_partial_values"): - # Optimization, if `get_partial_values` is available, we use it to retrieve - # all the values in a single call. - values = self.get_partial_values([(k, (0, None)) for k in keys_in_self]) - else: - # Otherwise, we just retrieve each key sequentially. - values = [self[k] for k in keys_in_self] - return dict(zip(keys_in_self, values)) + return {k: self[k] for k in keys if k in self} class Store(BaseStore): @@ -552,9 +544,6 @@ def __len__(self): def supports_efficient_get_partial_values(self): return self.inner_store.supports_efficient_get_partial_values - def getitems(self, keys, contexts): - return self.inner_store.getitems(keys, contexts) - def get_partial_values(self, key_ranges): return self.inner_store.get_partial_values(key_ranges) diff --git a/zarr/core.py b/zarr/core.py index 37afb2c3cb..84b8cd4d6c 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -1982,7 +1982,6 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, # 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 @@ -1991,6 +1990,7 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, and hasattr(self._compressor, "decode_partial") and not fields and self.dtype != object + and hasattr(self.chunk_store, "getitems") ): partial_read_decode = True cdatas = { @@ -2014,6 +2014,10 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, for ckey in ckeys if ckey in self.chunk_store } + elif hasattr(self.chunk_store, "get_partial_values"): + partial_read_decode = False + values = self.chunk_store.get_partial_values([(ckey, (0, None)) for ckey in ckeys]) + cdatas = {key: value for key, value in zip(ckeys, values) if value is not None} else: partial_read_decode = False contexts = {} diff --git a/zarr/tests/test_storage_v3.py b/zarr/tests/test_storage_v3.py index cc031f0db4..418f7d506b 100644 --- a/zarr/tests/test_storage_v3.py +++ b/zarr/tests/test_storage_v3.py @@ -666,6 +666,8 @@ def _get_public_and_dunder_methods(some_class): def test_storage_transformer_interface(): store_v3_methods = _get_public_and_dunder_methods(StoreV3) store_v3_methods.discard("__init__") + # Note, getitems() isn't mandatory when get_partial_values() is available + store_v3_methods.discard("getitems") storage_transformer_methods = _get_public_and_dunder_methods(StorageTransformer) storage_transformer_methods.discard("__init__") storage_transformer_methods.discard("get_config") From d9838ef4ca4207dcd7556586dd9fd29850b81285 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 11 Apr 2023 11:06:41 +0200 Subject: [PATCH 16/19] make contexts keyword-only --- zarr/_storage/store.py | 2 +- zarr/core.py | 2 +- zarr/storage.py | 2 +- zarr/tests/test_storage.py | 4 ++-- zarr/tests/util.py | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/zarr/_storage/store.py b/zarr/_storage/store.py index 94c3999924..0594dc22de 100644 --- a/zarr/_storage/store.py +++ b/zarr/_storage/store.py @@ -133,7 +133,7 @@ def _ensure_store(store: Any): ) def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Context] + self, keys: Sequence[str], *, contexts: Mapping[str, Context] ) -> Mapping[str, Any]: """Retrieve data from multiple keys. diff --git a/zarr/core.py b/zarr/core.py index 84b8cd4d6c..0d8d474b98 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -2023,7 +2023,7 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, contexts = {} if not isinstance(self._meta_array, np.ndarray): contexts = {k: {"meta_array": self._meta_array} for k in ckeys} - cdatas = self.chunk_store.getitems(ckeys, contexts) + cdatas = self.chunk_store.getitems(ckeys, contexts=contexts) for ckey, chunk_select, out_select in zip(ckeys, lchunk_selection, lout_selection): if ckey in cdatas: diff --git a/zarr/storage.py b/zarr/storage.py index 2138d55c4e..e6c3f62faf 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -1382,7 +1382,7 @@ def _normalize_key(self, key): return key.lower() if self.normalize_keys else key def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Context] + self, keys: Sequence[str], *, contexts: Mapping[str, Context] ) -> Mapping[str, Any]: keys_transformed = [self._normalize_key(key) for key in keys] diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py index eac7454ddb..d7d425ef5b 100644 --- a/zarr/tests/test_storage.py +++ b/zarr/tests/test_storage.py @@ -2593,9 +2593,9 @@ def __init__(self): super().__init__() self.last_contexts = None - def getitems(self, keys, contexts): + def getitems(self, keys, *, contexts): self.last_contexts = contexts - return super().getitems(keys, contexts) + return super().getitems(keys, contexts=contexts) store = MyStore() z = zarr.create(shape=(10,), store=store) diff --git a/zarr/tests/util.py b/zarr/tests/util.py index 0db00c9fb8..19ac8c0bfa 100644 --- a/zarr/tests/util.py +++ b/zarr/tests/util.py @@ -45,7 +45,7 @@ def __delitem__(self, key): del self.wrapped[key] def getitems( - self, keys: Sequence[str], contexts: Mapping[str, Context] + self, keys: Sequence[str], *, contexts: Mapping[str, Context] ) -> Mapping[str, Any]: for key in keys: self.counter['__getitem__', key] += 1 From c3ee95fc9f7200f608d220224a7b0aae1650977a Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 11 Apr 2023 13:29:25 +0200 Subject: [PATCH 17/19] Introduce ConstantMap --- zarr/core.py | 4 ++- zarr/tests/test_storage.py | 13 ++++++++-- zarr/util.py | 52 ++++++++++++++++++++++++++++++++++++-- 3 files changed, 64 insertions(+), 5 deletions(-) diff --git a/zarr/core.py b/zarr/core.py index 0d8d474b98..5537733b4b 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -13,6 +13,7 @@ from zarr._storage.store import _prefix_to_attrs_key, assert_zarr_v3_api_available from zarr.attrs import Attributes from zarr.codecs import AsType, get_codec +from zarr.context import Context from zarr.errors import ArrayNotFoundError, ReadOnlyError, ArrayIndexError from zarr.indexing import ( BasicIndexer, @@ -41,6 +42,7 @@ normalize_store_arg, ) from zarr.util import ( + ConstantMap, all_equal, InfoReporter, check_array_shape, @@ -2022,7 +2024,7 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, partial_read_decode = False contexts = {} if not isinstance(self._meta_array, np.ndarray): - contexts = {k: {"meta_array": self._meta_array} for k in ckeys} + contexts = ConstantMap(ckeys, constant=Context(meta_array=self._meta_array)) cdatas = self.chunk_store.getitems(ckeys, contexts=contexts) for ckey, chunk_select, out_select in zip(ckeys, lchunk_selection, lout_selection): diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py index d7d425ef5b..f157e2a3d2 100644 --- a/zarr/tests/test_storage.py +++ b/zarr/tests/test_storage.py @@ -20,6 +20,7 @@ import zarr from zarr._storage.store import _get_hierarchy_metadata from zarr.codecs import BZ2, AsType, Blosc, Zlib +from zarr.context import Context from zarr.convenience import consolidate_metadata from zarr.errors import ContainsArrayError, ContainsGroupError, MetadataError from zarr.hierarchy import group @@ -37,7 +38,7 @@ from zarr.storage import FSStore, rename, listdir from zarr._storage.v3 import KVStoreV3 from zarr.tests.util import CountingDict, have_fsspec, skip_test_env_var, abs_container, mktemp -from zarr.util import json_dumps +from zarr.util import ConstantMap, json_dumps @contextmanager @@ -2598,7 +2599,7 @@ def getitems(self, keys, *, contexts): return super().getitems(keys, contexts=contexts) store = MyStore() - z = zarr.create(shape=(10,), store=store) + z = zarr.create(shape=(10,), chunks=1, store=store) # By default, not contexts are given to the store's getitems() z[0] @@ -2608,3 +2609,11 @@ def getitems(self, keys, *, contexts): z._meta_array = "my_meta_array" z[0] assert store.last_contexts == {'0': {'meta_array': 'my_meta_array'}} + assert isinstance(store.last_contexts, ConstantMap) + # Accseeing different chunks should trigger different key request + z[1] + assert store.last_contexts == {'1': {'meta_array': 'my_meta_array'}} + assert isinstance(store.last_contexts, ConstantMap) + z[2:4] + assert store.last_contexts == ConstantMap(['2', '3'], Context({'meta_array': 'my_meta_array'})) + assert isinstance(store.last_contexts, ConstantMap) diff --git a/zarr/util.py b/zarr/util.py index be5f174aab..ef905994e5 100644 --- a/zarr/util.py +++ b/zarr/util.py @@ -1,3 +1,4 @@ +import collections.abc import inspect import json import math @@ -5,12 +6,21 @@ from textwrap import TextWrapper import mmap import time -from typing import Any, Callable, Dict, Optional, Tuple, Union +from typing import ( + Any, + Callable, + Dict, + Iterator, + Optional, + Tuple, + TypeVar, + Union, + Iterable +) import numpy as np from asciitree import BoxStyle, LeftAligned from asciitree.traversal import Traversal -from collections.abc import Iterable from numcodecs.compat import ( ensure_text, ensure_ndarray_like, @@ -21,6 +31,9 @@ from numcodecs.registry import codec_registry from numcodecs.blosc import cbuffer_sizes, cbuffer_metainfo +KeyType = TypeVar('KeyType') +ValueType = TypeVar('ValueType') + def flatten(arg: Iterable) -> Iterable: for element in arg: @@ -745,3 +758,38 @@ def ensure_contiguous_ndarray_or_bytes(buf) -> Union[NDArrayLike, bytes]: except TypeError: # An error is raised if `buf` couldn't be zero-copy converted return ensure_bytes(buf) + + +class ConstantMap(collections.abc.Mapping[KeyType, ValueType]): + """A read-only map that maps all keys to the same constant value + + Useful if you want to call `getitems()` with the same context for all keys. + + Parameters + ---------- + keys + The keys of the map. Will be copied to a frozenset if it isn't already. + constant + The constant that all keys are mapping to. + """ + + def __init__(self, keys: Iterable[KeyType], constant: ValueType) -> None: + self._keys = keys if isinstance(keys, frozenset) else frozenset(keys) + self._constant = constant + + def __getitem__(self, key: KeyType) -> ValueType: + if key not in self._keys: + raise KeyError(repr(key)) + return self._constant + + def __iter__(self) -> Iterator[KeyType]: + return iter(self._keys) + + def __len__(self) -> int: + return len(self._keys) + + def __contains__(self, key: object) -> bool: + return key in self._keys + + def __repr__(self) -> str: + return repr({k: v for k, v in self.items()}) From ec5f39606521f4f9728e08aa067a50f9eb1adb4b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 11 Apr 2023 13:42:32 +0200 Subject: [PATCH 18/19] use typing.Mapping --- zarr/util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zarr/util.py b/zarr/util.py index ef905994e5..68a238fbe4 100644 --- a/zarr/util.py +++ b/zarr/util.py @@ -1,4 +1,3 @@ -import collections.abc import inspect import json import math @@ -11,6 +10,7 @@ Callable, Dict, Iterator, + Mapping, Optional, Tuple, TypeVar, @@ -760,7 +760,7 @@ def ensure_contiguous_ndarray_or_bytes(buf) -> Union[NDArrayLike, bytes]: return ensure_bytes(buf) -class ConstantMap(collections.abc.Mapping[KeyType, ValueType]): +class ConstantMap(Mapping[KeyType, ValueType]): """A read-only map that maps all keys to the same constant value Useful if you want to call `getitems()` with the same context for all keys. From a1d3520f33528df09ebb58bdb5659b7ccb2ca589 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 11 Apr 2023 14:14:49 +0200 Subject: [PATCH 19/19] test_constant_map --- zarr/tests/test_util.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/zarr/tests/test_util.py b/zarr/tests/test_util.py index e9e1786abe..0a717b8f28 100644 --- a/zarr/tests/test_util.py +++ b/zarr/tests/test_util.py @@ -5,7 +5,7 @@ import pytest from zarr.core import Array -from zarr.util import (all_equal, flatten, guess_chunks, human_readable_size, +from zarr.util import (ConstantMap, all_equal, flatten, guess_chunks, human_readable_size, info_html_report, info_text_report, is_total_slice, json_dumps, normalize_chunks, normalize_dimension_separator, @@ -248,3 +248,16 @@ def test_json_dumps_numpy_dtype(): # Check that we raise the error of the superclass for unsupported object with pytest.raises(TypeError): json_dumps(Array) + + +def test_constant_map(): + val = object() + m = ConstantMap(keys=[1, 2], constant=val) + assert len(m) == 2 + assert m[1] is val + assert m[2] is val + assert 1 in m + assert 0 not in m + with pytest.raises(KeyError): + m[0] + assert repr(m) == repr({1: val, 2: val})