fixe, and allow partial reads for uncompressed v3 arrays

Author: jstriebel

zarr/_storage/store.py

@@ -272,13 +272,17 @@ def get_partial_values(self, key_ranges):
         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."""
+        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():
-            value = self[key]
+            try:
+                value = self[key]
+            except KeyError:
+                continue
             for i, (range_from, range_length) in indexed_ranges:
                 if range_length is None:
                     results[i] = value[range_from:]

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):
 

zarr/_storage/v3_storage_transformers.py

@@ -35,7 +35,7 @@ def get_chunk_slice(self, chunk: Tuple[int, ...]) -> Optional[slice]:
         if (chunk_start, chunk_len) == (MAX_UINT_64, MAX_UINT_64):
             return None
         else:
-            return slice(chunk_start, chunk_start + chunk_len)
+            return slice(int(chunk_start), int(chunk_start + chunk_len))
 
     def set_chunk_slice(
         self, chunk: Tuple[int, ...], chunk_slice: Optional[slice]
@@ -126,10 +126,13 @@ def _key_to_shard(self, chunk_key: str) -> Tuple[str, Tuple[int, ...]]:
 
     def _get_index_from_store(self, shard_key: str) -> _ShardIndex:
         # At the end of each shard 2*64bit per chunk for offset and length define the index:
+        index_bytes = self.inner_store.get_partial_values(
+            [(shard_key, (-16 * self._num_chunks_per_shard, None))]
+        )[0]
+        if index_bytes is None:
+            raise KeyError(shard_key)
         return _ShardIndex.from_bytes(
-            self.inner_store.get_partial_values(
-                [(shard_key, (-16 * self._num_chunks_per_shard, None))]
-            )[0],
+            index_bytes,
             self,
         )
 
@@ -156,7 +159,11 @@ def __getitem__(self, key):
         if _is_data_key(key):
             if self.supports_efficient_get_partial_values():
                 # Use the partial implementation, which fetches the index seperately
-                return self.get_partial_values([(key, (0, None))])[0]
+                value = self.get_partial_values([(key, (0, None))])[0]
+                if value is None:
+                    raise KeyError(key)
+                else:
+                    return value
             shard_key, chunk_subkey = self._key_to_shard(key)
             try:
                 full_shard_value = self.inner_store[shard_key]
@@ -254,7 +261,7 @@ def __delitem__(self, key):
                 del self.inner_store[shard_key]
             else:
                 index_bytes = index.to_bytes()
-                self.set_partial_values([(shard_key, -len(index_bytes), index_bytes)])
+                self.inner_store.set_partial_values([(shard_key, -len(index_bytes), index_bytes)])
         else:
             del self.inner_store[key]
 
@@ -281,22 +288,35 @@ def get_partial_values(self, key_ranges):
         if self.supports_efficient_get_partial_values():
             transformed_key_ranges = []
             cached_indices = {}
-            for key, range_ in key_ranges:
+            none_indices = []
+            for i, (key, range_) in enumerate(key_ranges):
                 if _is_data_key(key):
                     shard_key, chunk_subkey = self._key_to_shard(key)
                     try:
                         index = cached_indices[shard_key]
                     except KeyError:
-                        index = self._get_index_from_store(shard_key)
+                        try:
+                            index = self._get_index_from_store(shard_key)
+                        except KeyError:
+                            none_indices.append(i)
+                            continue
                         cached_indices[shard_key] = index
                     chunk_slice = index.get_chunk_slice(chunk_subkey)
+                    if chunk_slice is None:
+                        none_indices.append(i)
+                        continue
                     range_start, range_length = range_
+                    if range_length is None:
+                        range_length = chunk_slice.stop - chunk_slice.start
                     transformed_key_ranges.append(
-                        (shard_key, (range_start + chunk_slice.satrt, range_length))
+                        (shard_key, (range_start + chunk_slice.start, range_length))
                     )
                 else:
                     transformed_key_ranges.append((key, range_))
-            return self.inner_store.get_partial_values(transformed_key_ranges)
+            values = self.inner_store.get_partial_values(transformed_key_ranges)
+            for i in none_indices:
+                values.insert(i, None)
+            return values
         else:
             return StoreV3.get_partial_values(self, key_ranges)
 

zarr/core.py

@@ -51,6 +51,7 @@
     normalize_shape,
     normalize_storage_path,
     PartialReadBuffer,
+    UncompressedPartialReadBufferV3,
 )
 
 
@@ -180,6 +181,7 @@ def __init__(
 
         self._store = store
         self._chunk_store = chunk_store
+        self._transformed_chunk_store = None
         self._path = normalize_storage_path(path)
         if self._path:
             self._key_prefix = self._path + '/'
@@ -282,17 +284,13 @@ def _load_metadata_nosync(self):
 
             if self._version == 3:
                 storage_transformers = meta.get('storage_transformers', [])
-                transformed_store = self._store
-                for storage_transformer in storage_transformers:
-                    transformed_store = storage_transformer._copy_for_array(self, transformed_store)
-                self._transformed_store = transformed_store
-                if self._chunk_store is not None:
-                    transformed_chunk_store = self._chunk_store
+                if storage_transformers:
+                    transformed_store = self._chunk_store or self._store
                     for storage_transformer in storage_transformers:
-                        transformed_chunk_store = (
-                            storage_transformer._copy_for_array(self, transformed_chunk_store)
+                        transformed_store = storage_transformer._copy_for_array(
+                            self, transformed_store
                         )
-                    self._transformed_chunk_store = transformed_chunk_store
+                    self._transformed_chunk_store = transformed_store
 
     def _refresh_metadata(self):
         if not self._cache_metadata:
@@ -336,7 +334,7 @@ def _flush_metadata_nosync(self):
     @property
     def store(self):
         """A MutableMapping providing the underlying storage for the array."""
-        return self._transformed_store
+        return self._store
 
     @property
     def path(self):
@@ -373,10 +371,12 @@ def read_only(self, value):
     @property
     def chunk_store(self):
         """A MutableMapping providing the underlying storage for array chunks."""
-        if self._chunk_store is None:
-            return self._store
-        else:
+        if self._transformed_chunk_store is not None:
             return self._transformed_chunk_store
+        elif self._chunk_store is not None:
+            return self._chunk_store
+        else:
+            return self._store
 
     @property
     def shape(self):
@@ -1252,8 +1252,12 @@ def _get_selection(self, indexer, out=None, fields=None):
             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)):
+        if (
+            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:
 
@@ -1790,7 +1794,7 @@ def _set_selection(self, indexer, value, fields=None):
             check_array_shape('value', value, sel_shape)
 
         # iterate over chunks in range
-        if not hasattr(self.store, "setitems") or self._synchronizer is not None \
+        if not hasattr(self.chunk_store, "setitems") or self._synchronizer is not None \
            or any(map(lambda x: x == 0, self.shape)):
             # iterative approach
             for chunk_coords, chunk_selection, out_selection in indexer:
@@ -1868,10 +1872,12 @@ def _process_chunk(
                 # contiguous, so we can decompress directly from the chunk
                 # into the destination array
                 if self._compressor:
-                    if isinstance(cdata, PartialReadBuffer):
+                    if isinstance(cdata, (PartialReadBuffer, UncompressedPartialReadBufferV3)):
                         cdata = cdata.read_full()
                     self._compressor.decode(cdata, dest)
                 else:
+                    if isinstance(cdata, UncompressedPartialReadBufferV3):
+                        cdata = cdata.read_full()
                     chunk = ensure_ndarray(cdata).view(self._dtype)
                     chunk = chunk.reshape(self._chunks, order=self._order)
                     np.copyto(dest, chunk)
@@ -1893,13 +1899,21 @@ def _process_chunk(
                         else dim
                         for i, dim in enumerate(self.chunks)
                     ]
-                    cdata.read_part(start, nitems)
-                    chunk_partial = self._decode_chunk(
-                        cdata.buff,
-                        start=start,
-                        nitems=nitems,
-                        expected_shape=expected_shape,
-                    )
+                    if isinstance(cdata, UncompressedPartialReadBufferV3):
+                        chunk_partial = self._decode_chunk(
+                            cdata.read_part(start, nitems),
+                            start=start,
+                            nitems=nitems,
+                            expected_shape=expected_shape,
+                        )
+                    else:
+                        cdata.read_part(start, nitems)
+                        chunk_partial = self._decode_chunk(
+                            cdata.buff,
+                            start=start,
+                            nitems=nitems,
+                            expected_shape=expected_shape,
+                        )
                     tmp[partial_out_selection] = chunk_partial
                 out[out_selection] = tmp[chunk_selection]
                 return
@@ -1994,9 +2008,29 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection,
                 for ckey in ckeys
                 if ckey in self.chunk_store
             }
+        elif (
+            self._partial_decompress
+            and not self._compressor
+            and not fields
+            and self.dtype != object
+            and hasattr(self.chunk_store, "get_partial_values")
+            and self.chunk_store.supports_efficient_get_partial_values()
+        ):
+            partial_read_decode = True
+            cdatas = {
+                ckey: UncompressedPartialReadBufferV3(
+                    ckey, self.chunk_store, itemsize=self.itemsize
+                )
+                for ckey in ckeys
+                if ckey in self.chunk_store
+            }
         else:
             partial_read_decode = False
-            cdatas = self.chunk_store.getitems(ckeys, on_error="omit")
+            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:
+                cdatas = self.chunk_store.getitems(ckeys, on_error="omit")
         for ckey, chunk_select, out_select in zip(ckeys, lchunk_selection, lout_selection):
             if ckey in cdatas:
                 self._process_chunk(

zarr/tests/test_core.py

@@ -3287,6 +3287,54 @@ def expected(self):
         ]
 
 
+@pytest.mark.skipif(have_fsspec is False, reason="needs fsspec")
+@pytest.mark.skipif(not v3_api_available, reason="V3 is disabled")
+class TestArrayWithFSStoreV3PartialReadUncompressedSharded(
+    TestArrayWithPathV3, TestArrayWithFSStorePartialRead
+):
+
+    @staticmethod
+    def create_array(array_path='arr1', read_only=False, **kwargs):
+        path = mkdtemp()
+        atexit.register(shutil.rmtree, path)
+        store = FSStoreV3(path)
+        cache_metadata = kwargs.pop("cache_metadata", True)
+        cache_attrs = kwargs.pop("cache_attrs", True)
+        write_empty_chunks = kwargs.pop('write_empty_chunks', True)
+        kwargs.setdefault('compressor', None)
+        num_dims = 1 if isinstance(kwargs["shape"], int) else len(kwargs["shape"])
+        sharding_transformer = ShardingStorageTransformer(
+            "indexed", chunks_per_shard=(2, ) * num_dims
+        )
+        init_array(store, path=array_path, storage_transformers=[sharding_transformer], **kwargs)
+        return Array(
+            store,
+            path=array_path,
+            read_only=read_only,
+            cache_metadata=cache_metadata,
+            cache_attrs=cache_attrs,
+            partial_decompress=True,
+            write_empty_chunks=write_empty_chunks,
+        )
+
+    def test_nbytes_stored(self):
+        z = self.create_array(shape=1000, chunks=100)
+        expect_nbytes_stored = sum(buffer_size(v) for k, v in z._store.items() if k != 'zarr.json')
+        assert expect_nbytes_stored == z.nbytes_stored
+        z[:] = 42
+        expect_nbytes_stored = sum(buffer_size(v) for k, v in z._store.items() if k != 'zarr.json')
+        assert expect_nbytes_stored == z.nbytes_stored
+
+    def expected(self):
+        return [
+            "90109fc2a4e17efbcb447003ea1c08828b91f71e",
+            "2b73519f7260dba3ddce0d2b70041888856fec6b",
+            "bca5798be2ed71d444f3045b05432d937682b7dd",
+            "9ff1084501e28520e577662a6e3073f1116c76a2",
+            "882a97cad42417f90f111d0cb916a21579650467",
+        ]
+
+
 @pytest.mark.skipif(have_fsspec is False, reason="needs fsspec")
 @pytest.mark.skipif(not v3_api_available, reason="V3 is disabled")
 class TestArrayWithFSStoreV3Nested(TestArrayWithPathV3, TestArrayWithFSStoreNested):
@@ -3398,9 +3446,6 @@ def create_array(array_path='arr1', read_only=False, **kwargs):
                      cache_metadata=cache_metadata,
                      cache_attrs=cache_attrs, write_empty_chunks=write_empty_chunks)
 
-    # def test_nbytes_stored(self):
-    #     pass  # not implemented
-
     def test_nbytes_stored(self):
         z = self.create_array(shape=1000, chunks=100)
         expect_nbytes_stored = sum(buffer_size(v) for k, v in z._store.items() if k != 'zarr.json')
@@ -3410,9 +3455,19 @@ def test_nbytes_stored(self):
         assert expect_nbytes_stored == z.nbytes_stored
 
         # mess with store
-        z.chunk_store[data_root + z._key_prefix + 'foo'] = list(range(10))
+        z.store[data_root + z._key_prefix + 'foo'] = list(range(10))
         assert -1 == z.nbytes_stored
 
+    def test_keys_inner_store(self):
+        z = self.create_array(shape=1000, chunks=100)
+        assert z.chunk_store.keys() == z._store.keys()
+        meta_keys = set(z.store.keys())
+        z[:] = 42
+        assert len(z.chunk_store.keys() - meta_keys) == 10
+        # inner store should have half the data keys,
+        # since chunks_per_shard is 2:
+        assert len(z._store.keys() - meta_keys) == 5
+
     def expected(self):
         return [
             "b46294e25b1d816055e7937780265c0d8d5d6c47",