Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 17 additions & 33 deletions virtualizarr/manifests/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
from zarr.core.common import BytesLike

from virtualizarr.manifests.group import ManifestGroup
from virtualizarr.vendor.zarr.core.metadata import dict_to_buffer

if TYPE_CHECKING:
from obstore.store import (
Expand Down Expand Up @@ -56,36 +55,6 @@ def get_store_prefix(url: str) -> str:
return "" if scheme in {"", "file"} else f"{scheme}://{netloc}"


def get_zarr_metadata(manifest_group: ManifestGroup, key: str) -> Buffer:
"""
Generate the expected Zarr V3 metadata from a virtual dataset.

Group metadata is returned for all Datasets and Array metadata
is returned for all DataArrays.

Combines the ManifestArray metadata with the attrs from the DataArray
and adds `dimension_names` for all arrays if not already provided.

Parameters
----------
manifest_group : ManifestGroup
key : str

Returns
-------
Buffer
"""
# If requesting the root metadata, return the standard group metadata with additional dataset specific attributes

if key == "zarr.json":
metadata = manifest_group.metadata.to_dict()
return dict_to_buffer(metadata, prototype=default_buffer_prototype())
else:
var, _ = key.split("/")
metadata = manifest_group.arrays[var].metadata.to_dict()
return dict_to_buffer(metadata, prototype=default_buffer_prototype())


def parse_manifest_index(key: str, chunk_key_encoding: str = ".") -> tuple[int, ...]:
"""
Splits `key` provided to a zarr store into the variable indicated
Expand Down Expand Up @@ -243,8 +212,19 @@ async def get(
byte_range: ByteRequest | None = None,
) -> Buffer | None:
# docstring inherited
if key.endswith("zarr.json"):
return get_zarr_metadata(self._group, key)

if key == "zarr.json":
# Return group metadata
return self._group.metadata.to_buffer_dict(
prototype=default_buffer_prototype()
)["zarr.json"]
elif key.endswith("zarr.json"):
# Return array metadata
# TODO: Handle nested groups
var, _ = key.split("/")
return self._group.arrays[var].metadata.to_buffer_dict(
prototype=default_buffer_prototype()
)["zarr.json"]
var = key.split("/")[0]
marr = self._group.arrays[var]
manifest = marr.manifest
Expand All @@ -258,22 +238,26 @@ async def get(
return None
offset = manifest._offsets[chunk_indexes]
length = manifest._lengths[chunk_indexes]

# Get the configured object store instance that matches the path
store = self._store_registry.get_store(path)
if not store:
raise ValueError(
f"Could not find a store to use for {path} in the store registry"
)

# Truncate path to match Obstore expectations
key = urlparse(path).path
if hasattr(store, "prefix") and store.prefix:
# strip the prefix from key
key = key.removeprefix(str(store.prefix))

# Transform the input byte range to account for the chunk location in the file
chunk_end_exclusive = offset + length
byte_range = _transform_byte_range(
byte_range, chunk_start=offset, chunk_end_exclusive=chunk_end_exclusive
)

# Actually get the bytes
try:
bytes = await store.get_range_async(
Expand Down
5 changes: 1 addition & 4 deletions virtualizarr/tests/test_parsers/test_fits.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,6 @@

@requires_kerchunk
@requires_network
@pytest.mark.xfail(
reason="Big endian not yet supported by zarr-python 3.0"
) # https://github.com/zarr-developers/zarr-python/issues/2324
def test_open_hubble_data():
# data from https://registry.opendata.aws/hst/
file_url = "s3://stpubdata/hst/public/f05i/f05i0201m/f05i0201m_a1f.fits"
Expand All @@ -28,4 +25,4 @@ def test_open_hubble_data():
assert list(vds.variables) == ["PRIMARY"]
var = vds["PRIMARY"].variable
assert var.sizes == {"y": 17, "x": 589}
assert var.dtype == ">i4"
assert var.dtype == "int32"
32 changes: 0 additions & 32 deletions virtualizarr/vendor/zarr/core/metadata.py

This file was deleted.