Properly extract chunk keys for arrays with a single chunk (#523)

* Properly extract chunk keys for arrays with a single chunk

* Update test to specify chunk_key_encoding

* Lint

* Raise NotImplementedError on scalar arrays

Author: maxrjones

virtualizarr/manifests/store.py

@@ -108,36 +108,31 @@ def get_zarr_metadata(manifest_group: ManifestGroup, key: str) -> Buffer:
         return dict_to_buffer(metadata, prototype=default_buffer_prototype())
 
 
-def parse_manifest_index(
-    key: str, chunk_key_encoding: str = "."
-) -> tuple[str, tuple[int, ...]]:
+def parse_manifest_index(key: str, chunk_key_encoding: str = ".") -> tuple[int, ...]:
     """
     Splits `key` provided to a zarr store into the variable indicated
     by the first part and the chunk index from the 3rd through last parts,
     which can be used to index into the ndarrays containing paths, offsets,
     and lengths in ManifestArrays.
 
-    Currently only works for 1d+ arrays with a tree depth of one from the
-    root Zarr group.
-
     Parameters
     ----------
     key : str
     chunk_key_encoding : str
 
     Returns
     -------
-    ManifestIndex
+    tuple containing chunk indexes
     """
-    parts = key.split("/")
-    var = parts[0]
-    # Assume "c" is the second part
-    # TODO: Handle scalar array case with "c" holds the data
-    if chunk_key_encoding == "/":
-        indexes = tuple(int(ind) for ind in parts[2:])
-    else:
-        indexes = tuple(int(ind) for ind in parts[2].split(chunk_key_encoding))
-    return var, indexes
+    if key.endswith("c"):
+        # Scalar arrays hold the data in the "c" key
+        raise NotImplementedError(
+            "Scalar arrays are not yet supported by ManifestStore"
+        )
+    parts = key.split(
+        "c/"
+    )  # TODO: Open an issue upstream about the Zarr spec indicating this should be f"c{chunk_key_encoding}" rather than always "c/"
+    return tuple(int(ind) for ind in parts[1].split(chunk_key_encoding))
 
 
 def find_matching_store(stores: StoreDict, request_key: str) -> StoreRequest:
@@ -263,13 +258,16 @@ async def get(
 
         if key.endswith("zarr.json"):
             return get_zarr_metadata(self._group, key)
-        var, chunk_key = parse_manifest_index(key)
+        var = key.split("/")[0]
         marr = self._group.arrays[var]
         manifest = marr.manifest
 
-        path = manifest._paths[*chunk_key]
-        offset = manifest._offsets[*chunk_key]
-        length = manifest._lengths[*chunk_key]
+        chunk_indexes = parse_manifest_index(
+            key, marr.metadata.chunk_key_encoding.separator
+        )
+        path = manifest._paths[*chunk_indexes]
+        offset = manifest._offsets[*chunk_indexes]
+        length = manifest._lengths[*chunk_indexes]
         # Get the  configured object store instance that matches the path
         store_request = find_matching_store(stores=self._stores, request_key=path)
         # Transform the input byte range to account for the chunk location in the file

virtualizarr/manifests/utils.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from zarr import Array
+from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike
 from zarr.core.metadata.v3 import ArrayV3Metadata
 
 from virtualizarr.codecs import convert_to_codec_pipeline, get_codecs
@@ -14,9 +15,11 @@ def create_v3_array_metadata(
     shape: tuple[int, ...],
     data_type: np.dtype,
     chunk_shape: tuple[int, ...],
+    chunk_key_encoding: ChunkKeyEncodingLike = {"name": "default"},
     fill_value: Any = None,
     codecs: Optional[list[Dict[str, Any]]] = None,
     attributes: Optional[Dict[str, Any]] = None,
+    dimension_names: Optional[tuple[str, ...]] = None,
 ) -> ArrayV3Metadata:
     """
     Create an ArrayV3Metadata instance with standard configuration.
@@ -30,12 +33,16 @@ def create_v3_array_metadata(
         The numpy dtype of the array
     chunk_shape : tuple[int, ...]
         The shape of each chunk
+    chunk_key_encoding : ChunkKeyEncodingLike
+        The mapping from chunk grid cell coordinates to keys.
     fill_value : Any, optional
         The fill value for the array
     codecs : list[Dict[str, Any]], optional
         List of codec configurations
     attributes : Dict[str, Any], optional
         Additional attributes for the array
+    dimension_names : tuple[str], optional
+        Names of the dimensions
 
     Returns
     -------
@@ -49,14 +56,14 @@ def create_v3_array_metadata(
             "name": "regular",
             "configuration": {"chunk_shape": chunk_shape},
         },
-        chunk_key_encoding={"name": "default"},
+        chunk_key_encoding=chunk_key_encoding,
         fill_value=fill_value,
         codecs=convert_to_codec_pipeline(
             codecs=codecs or [],
             dtype=data_type,
         ),
         attributes=attributes or {},
-        dimension_names=None,
+        dimension_names=dimension_names,
         storage_transformers=None,
     )
 

virtualizarr/tests/test_manifests/test_store.py

@@ -2,8 +2,9 @@
 
 import json
 import pickle
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
+import numpy as np
 import pytest
 from zarr.abc.store import (
     OffsetByteRequest,
@@ -19,14 +20,15 @@
     ManifestGroup,
     ManifestStore,
 )
+from virtualizarr.manifests.utils import create_v3_array_metadata
 from virtualizarr.tests import requires_minio, requires_obstore
 
 if TYPE_CHECKING:
     from obstore.store import ObjectStore
 
 
 def _generate_manifest_store(
-    store: ObjectStore, *, prefix: str, filepath: str, array_v3_metadata: Callable
+    store: ObjectStore, *, prefix: str, filepath: str
 ) -> ManifestStore:
     """
     Generate a ManifestStore for testing.
@@ -66,9 +68,15 @@ def _generate_manifest_store(
         "1.1": {"path": f"{prefix}/{filepath}", "offset": 12, "length": 4},
     }
     manifest = ChunkManifest(entries=chunk_dict)
-    chunks = (1, 4)
-    shape = (2, 8)
-    array_metadata = array_v3_metadata(shape=shape, chunks=chunks)
+    codecs = [{"configuration": {"endian": "little"}, "name": "bytes"}]
+    array_metadata = create_v3_array_metadata(
+        shape=(4, 4),
+        chunk_shape=(2, 2),
+        data_type=np.dtype("int32"),
+        codecs=codecs,
+        chunk_key_encoding={"name": "default", "separator": "."},
+        fill_value=0,
+    )
     manifest_array = ManifestArray(metadata=array_metadata, chunkmanifest=manifest)
     manifest_group = ManifestGroup(
         arrays={"foo": manifest_array, "bar": manifest_array},
@@ -78,7 +86,7 @@ def _generate_manifest_store(
 
 
 @pytest.fixture()
-def local_store(tmpdir, array_v3_metadata):
+def local_store(tmpdir):
     import obstore as obs
 
     store = obs.store.LocalStore()
@@ -88,12 +96,11 @@ def local_store(tmpdir, array_v3_metadata):
         store=store,
         prefix=prefix,
         filepath=filepath,
-        array_v3_metadata=array_v3_metadata,
     )
 
 
 @pytest.fixture()
-def s3_store(minio_bucket, array_v3_metadata):
+def s3_store(minio_bucket):
     import obstore as obs
 
     store = obs.store.S3Store(
@@ -110,7 +117,6 @@ def s3_store(minio_bucket, array_v3_metadata):
         store=store,
         prefix=prefix,
         filepath=filepath,
-        array_v3_metadata=array_v3_metadata,
     )
 
 
@@ -164,7 +170,7 @@ async def test_get_metadata(self, manifest_store, request):
             "foo/zarr.json", prototype=default_buffer_prototype()
         )
         metadata = json.loads(observed.to_bytes())
-        assert metadata["chunk_grid"]["configuration"]["chunk_shape"] == [1, 4]
+        assert metadata["chunk_grid"]["configuration"]["chunk_shape"] == [2, 2]
         assert metadata["node_type"] == "array"
         assert metadata["zarr_format"] == 3