Sharding storage transformer for v3

docs/release.rst

@@ -35,6 +35,18 @@ Bug fixes
 * Fix bug in LRUEStoreCache in which the current size wasn't reset on invalidation.
   By :user:`BGCMHou <BGCMHou>` and :user:`Josh Moore <joshmoore>` :issue:`1076`, :issue:`1077`.
 
+Enhancements
+~~~~~~~~~~~~
+
+* **Improve Zarr V3 support, adding partial store read/write, storage transformers and sharding.**
+  Add two features of the [v3 spec](https://zarr-specs.readthedocs.io/en/latest/core/v3.0.html):
+  * `get_partial_values` and `set_partial_values`
+  * efficient `get_partial_values` implementation for `FSStoreV3`
+  * storage transformers interface
+  * sharding storage transformer
+  By :user:`Jonathan Striebel <jstriebel>`; :issue:`1096`, :issue:`1111`.
+
+
 Documentation
 ~~~~~~~~~~~~~
 

zarr/_storage/store.py

@@ -1,6 +1,8 @@
 import abc
 import os
+from collections import defaultdict
 from collections.abc import MutableMapping
+from copy import copy
 from string import ascii_letters, digits
 from typing import Any, List, Mapping, Optional, Union
 
@@ -254,6 +256,74 @@ def __setitem__(self, key, value):
     def __getitem__(self, key):
         """Get a value."""
 
+    @abc.abstractmethod
+    def rmdir(self, path=None):
+        """Remove a data path and all its subkeys and related metadata.
+        Expects a path without the data or meta root prefix."""
+
+    def supports_efficient_get_partial_values(self):
+        return False
+
+    def get_partial_values(self, key_ranges):
+        """Get multiple partial values.
+        key_ranges can be an iterable of key, range pairs,
+        where a range specifies two integers range_start and range_length
+        as a tuple, (range_start, range_length).
+        range_length may be None to indicate to read until the end.
+        range_start may be negative to start reading range_start bytes
+        from the end of the file.
+        A key may occur multiple times with different ranges.
+        Inserts None for missing keys into the returned list."""
+        results = [None] * len(key_ranges)
+        indexed_ranges_by_key = defaultdict(list)
+        for i, (key, range_) in enumerate(key_ranges):
+            indexed_ranges_by_key[key].append((i, range_))
+        for key, indexed_ranges in indexed_ranges_by_key.items():
+            try:
+                value = self[key]
+            except KeyError:  # pragma: no cover
+                continue
+            for i, (range_from, range_length) in indexed_ranges:
+                if range_length is None:
+                    results[i] = value[range_from:]
+                else:
+                    results[i] = value[range_from:range_from + range_length]
+        return results
+
+    def supports_efficient_set_partial_values(self):
+        return False
+
+    def set_partial_values(self, key_start_values):
+        """Set multiple partial values.
+        key_start_values can be an iterable of key, start and value triplets
+        as tuples, (key, start, value), where start defines the offset in bytes.
+        A key may occur multiple times with different starts and non-overlapping values.
+        Also, start may only be beyond the current value if other values fill the gap.
+        start may be negative to start writing start bytes from the current
+        end of the file, ending the file with the new value."""
+        unique_keys = set(next(zip(*key_start_values)))
+        values = {}
+        for key in unique_keys:
+            old_value = self.get(key)
+            values[key] = None if old_value is None else bytearray(old_value)
+        for key, start, value in key_start_values:
+            if values[key] is None:
+                assert start == 0
+                values[key] = value
+            else:
+                if start > len(values[key]):  # pragma: no cover
+                    raise ValueError(
+                        f"Cannot set value at start {start}, "
+                        + f"since it is beyond the data at key {key}, "
+                        + f"having length {len(values[key])}."
+                    )
+                if start < 0:
+                    values[key][start:] = value
+                else:
+                    values[key][start:start + len(value)] = value
+        for key, value in values.items():
+            self[key] = value
+
     def clear(self):
         """Remove all items from store."""
         self.erase_prefix("/")
@@ -303,6 +373,150 @@ def _ensure_store(store):
         )
 
 
+class StorageTransformer(MutableMapping, abc.ABC):
+    """Base class for storage transformers. The methods simply pass on the data as-is
+    and should be overwritten by sub-classes."""
+
+    _store_version = 3
+    _metadata_class = Metadata3
+
+    def __init__(self, _type) -> None:
+        if _type not in self.valid_types:  # pragma: no cover
+            raise ValueError(
+                f"Storage transformer cannot be initialized with type {_type}, "
+                + f"must be one of {list(self.valid_types)}."
+            )
+        self.type = _type
+        self._inner_store = None
+
+    def _copy_for_array(self, array, inner_store):
+        transformer_copy = copy(self)
+        transformer_copy._inner_store = inner_store
+        return transformer_copy
+
+    @abc.abstractproperty
+    def extension_uri(self):
+        pass  # pragma: no cover
+
+    @abc.abstractproperty
+    def valid_types(self):
+        pass  # pragma: no cover
+
+    def get_config(self):
+        """Return a dictionary holding configuration parameters for this
+        storage transformer. All values must be compatible with JSON encoding."""
+        # Override in sub-class if need special encoding of config values.
+        # By default, assume all non-private members are configuration
+        # parameters except for type .
+        return {
+            k: v for k, v in self.__dict__.items()
+            if not k.startswith('_') and k != "type"
+        }
+
+    @classmethod
+    def from_config(cls, _type, config):
+        """Instantiate storage transformer from a configuration object."""
+        # override in sub-class if need special decoding of config values
+
+        # by default, assume constructor accepts configuration parameters as
+        # keyword arguments without any special decoding
+        return cls(_type, **config)
+
+    @property
+    def inner_store(self) -> Union["StorageTransformer", StoreV3]:
+        assert self._inner_store is not None, (
+            "inner_store is not initialized, first get a copy via _copy_for_array."
+        )
+        return self._inner_store
+
+    # The following implementations are usually fine to keep as-is:
+
+    def __eq__(self, other):
+        return (
+            type(self) == type(other) and
+            self._inner_store == other._inner_store and
+            self.get_config() == other.get_config()
+        )
+
+    def erase(self, key):
+        self.__delitem__(key)
+
+    def list(self):
+        return list(self.keys())
+
+    def list_dir(self, prefix):
+        return StoreV3.list_dir(self, prefix)
+
+    def is_readable(self):
+        return self.inner_store.is_readable()
+
+    def is_writeable(self):
+        return self.inner_store.is_writeable()
+
+    def is_listable(self):
+        return self.inner_store.is_listable()
+
+    def is_erasable(self):
+        return self.inner_store.is_erasable()
+
+    def clear(self):
+        return self.inner_store.clear()
+
+    def __enter__(self):
+        return self.inner_store.__enter__()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        return self.inner_store.__exit__(exc_type, exc_value, traceback)
+
+    def close(self) -> None:
+        return self.inner_store.close()
+
+    # The following implementations might need to be re-implemented
+    # by subclasses implementing storage transformers:
+
+    def rename(self, src_path: str, dst_path: str) -> None:
+        return self.inner_store.rename(src_path, dst_path)
+
+    def list_prefix(self, prefix):
+        return self.inner_store.list_prefix(prefix)
+
+    def erase_prefix(self, prefix):
+        return self.inner_store.erase_prefix(prefix)
+
+    def rmdir(self, path=None):
+        return self.inner_store.rmdir(path)
+
+    def __contains__(self, key):
+        return self.inner_store.__contains__(key)
+
+    def __setitem__(self, key, value):
+        return self.inner_store.__setitem__(key, value)
+
+    def __getitem__(self, key):
+        return self.inner_store.__getitem__(key)
+
+    def __delitem__(self, key):
+        return self.inner_store.__delitem__(key)
+
+    def __iter__(self):
+        return self.inner_store.__iter__()
+
+    def __len__(self):
+        return self.inner_store.__len__()
+
+    def supports_efficient_get_partial_values(self):
+        return self.inner_store.supports_efficient_get_partial_values()
+
+    def get_partial_values(self, key_ranges):
+        return self.inner_store.get_partial_values(key_ranges)
+
+    def supports_efficient_set_partial_values(self):
+        return self.inner_store.supports_efficient_set_partial_values()
+
+    def set_partial_values(self, key_start_values):
+        return self.inner_store.set_partial_values(key_start_values)
+
+
 # allow MutableMapping for backwards compatibility
 StoreLike = Union[BaseStore, MutableMapping]
 

zarr/_storage/v3.py

@@ -182,6 +182,37 @@ def rmdir(self, path=None):
             if self.fs.isdir(store_path):
                 self.fs.rm(store_path, recursive=True)
 
+    def supports_efficient_get_partial_values(self):
+        return True
+
+    def get_partial_values(self, key_ranges):
+        """Get multiple partial values.
+        key_ranges can be an iterable of key, range pairs,
+        where a range specifies two integers range_start and range_length
+        as a tuple, (range_start, range_length).
+        range_length may be None to indicate to read until the end.
+        range_start may be negative to start reading range_start bytes
+        from the end of the file.
+        A key may occur multiple times with different ranges.
+        Inserts None for missing keys into the returned list."""
+        results = []
+        for key, (range_start, range_length) in key_ranges:
+            key = self._normalize_key(key)
+            path = self.dir_path(key)
+            try:
+                if range_start < 0:
+                    if range_length is None:
+                        result = self.fs.tail(path, size=-range_start)
+                    else:
+                        size = self.fs.size(path)
+                        result = self.fs.read_block(path, size + range_start, range_length)
+                else:
+                    result = self.fs.read_block(path, range_start, range_length)
+            except self.map.missing_exceptions:
+                result = None
+            results.append(result)
+        return results
+
 
 class MemoryStoreV3(MemoryStore, StoreV3):