diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 7f185535df..0f7c2edb2b 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -13,11 +13,12 @@ from zarr.core.array import ( DEFAULT_FILL_VALUE, Array, + ArrayMetadataProbe, AsyncArray, CompressorLike, + _probe_array_metadata, create_array, from_array, - get_array_metadata, ) from zarr.core.array_spec import ArrayConfigLike, parse_array_config from zarr.core.buffer import NDArrayLike @@ -35,6 +36,7 @@ AsyncGroup, ConsolidatedMetadata, GroupMetadata, + _PreFetchedGroupMetadata, create_hierarchy, ) from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata @@ -108,6 +110,24 @@ def _infer_overwrite(mode: AccessModeLiteral) -> bool: return mode in _OVERWRITE_MODES +def _array_probe_to_group_prefetch( + probe: ArrayMetadataProbe, +) -> _PreFetchedGroupMetadata | None: + """Translate array-probe buffers into group-open pre-fetched metadata. + + The ``zarr.json`` and ``.zattrs`` documents fetched while probing for an array + are the same documents ``AsyncGroup.open`` would otherwise fetch, so they can + be reused to avoid duplicate requests. Returns None if nothing reusable was + fetched (e.g. an explicit ``zarr_format`` probe, which does not overlap). + """ + pre_fetched: _PreFetchedGroupMetadata = {} + if "zarr_json" in probe: + pre_fetched["zarr_json"] = probe["zarr_json"] + if "zattrs" in probe: + pre_fetched["zattrs"] = probe["zattrs"] + return pre_fetched or None + + def _get_shape_chunks(a: ArrayLike | Any) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]: """Helper function to get the shape and chunks from an array-like object""" shape = None @@ -358,20 +378,33 @@ async def open( # TODO: the mode check below seems wrong! if "shape" not in kwargs and mode in {"a", "r", "r+", "w"}: - try: - metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) + # Probe for an array once. The probe also returns the zarr.json/.zattrs + # buffers it fetched so the group open below can reuse them instead of + # re-fetching (only populated for zarr_format=None). + result = await _probe_array_metadata(store_path, zarr_format=zarr_format) + if result.metadata is not None: # TODO: remove this cast when we fix typing for array metadata dicts - _metadata_dict = cast("ArrayMetadataDict", metadata_dict) - # for v2, the above would already have raised an exception if not an array + _metadata_dict = cast("ArrayMetadataDict", result.metadata) zarr_format = _metadata_dict["zarr_format"] is_v3_array = zarr_format == 3 and _metadata_dict.get("node_type") == "array" if is_v3_array or zarr_format == 2: return AsyncArray( store_path=store_path, metadata=_metadata_dict, config=kwargs.get("config") ) - except (AssertionError, FileNotFoundError, NodeTypeValidationError): - pass - return await open_group(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs) + elif not isinstance(result.error, (FileNotFoundError, NodeTypeValidationError)): + # Only "not an array" errors trigger the group fallback; anything else + # (e.g. an invalid zarr_format) is propagated, as it was before. + assert result.error is not None + raise result.error + # Fall back to opening a group, reusing the buffers already fetched during + # the array probe (zarr.json, .zattrs) so they are not fetched again. + return await open_group( + store=store_path, + zarr_format=zarr_format, + mode=mode, + _pre_fetched_metadata=_array_probe_to_group_prefetch(result.probe), + **kwargs, + ) try: return await open_array(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs) @@ -756,6 +789,7 @@ async def open_group( meta_array: Any | None = None, # not used attributes: dict[str, JSON] | None = None, use_consolidated: bool | str | None = None, + _pre_fetched_metadata: _PreFetchedGroupMetadata | None = None, ) -> AsyncGroup: """Open a group using file-mode-like semantics. @@ -806,6 +840,10 @@ async def open_group( Zarr format 2 allowed configuring the key storing the consolidated metadata (``.zmetadata`` by default). Specify the custom key as ``use_consolidated`` to load consolidated metadata from a non-default key. + _pre_fetched_metadata : _PreFetchedGroupMetadata or None, default None + Private. Already-fetched metadata buffers that ``AsyncGroup.open`` may + reuse instead of re-fetching. Used by ``zarr.open`` to avoid duplicate + fetches when falling back from an array probe to a group open. Returns ------- @@ -829,7 +867,10 @@ async def open_group( try: if mode in _READ_MODES: return await AsyncGroup.open( - store_path, zarr_format=zarr_format, use_consolidated=use_consolidated + store_path, + zarr_format=zarr_format, + use_consolidated=use_consolidated, + _pre_fetched_metadata=_pre_fetched_metadata, ) except (KeyError, FileNotFoundError): pass diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 8386427b3f..29d1302164 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -40,6 +40,7 @@ ZarrFormat, ) from zarr.core.dtype import ZDTypeLike + from zarr.core.group import _PreFetchedGroupMetadata from zarr.storage import StoreLike from zarr.types import AnyArray @@ -473,6 +474,7 @@ def open_group( meta_array: Any | None = None, # not used in async api attributes: dict[str, JSON] | None = None, use_consolidated: bool | str | None = None, + _pre_fetched_metadata: _PreFetchedGroupMetadata | None = None, ) -> Group: """Open a group using file-mode-like semantics. @@ -523,6 +525,10 @@ def open_group( Zarr format 2 allowed configuring the key storing the consolidated metadata (``.zmetadata`` by default). Specify the custom key as ``use_consolidated`` to load consolidated metadata from a non-default key. + _pre_fetched_metadata : _PreFetchedGroupMetadata or None, default None + Private. Already-fetched metadata buffers that ``AsyncGroup.open`` may + reuse instead of re-fetching. Used by ``zarr.open`` to avoid duplicate + fetches when falling back from an array probe to a group open. Returns ------- @@ -543,6 +549,7 @@ def open_group( meta_array=meta_array, attributes=attributes, use_consolidated=use_consolidated, + _pre_fetched_metadata=_pre_fetched_metadata, ) ) ) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index d15c70064b..484496a1f2 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -10,6 +10,7 @@ TYPE_CHECKING, Any, Literal, + NotRequired, TypedDict, cast, overload, @@ -133,6 +134,7 @@ ArrayNotFoundError, ChunkNotFoundError, MetadataValidationError, + NodeTypeValidationError, ZarrDeprecationWarning, ZarrUserWarning, ) @@ -154,6 +156,7 @@ from zarr.abc.codec import CodecPipeline from zarr.abc.store import Store from zarr.codecs.sharding import ShardingCodecIndexLocation + from zarr.core.buffer import Buffer from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -234,9 +237,48 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None raise TypeError # pragma: no cover -async def get_array_metadata( +class ArrayMetadataProbe(TypedDict): + """Buffers fetched while probing for array metadata. + + Private. Carried on ``ArrayProbeResult`` so callers (e.g. ``zarr.open``) can + reuse the ``zarr.json`` and ``.zattrs`` documents when falling back to a group + open, instead of fetching them a second time. A key is present only if that + document was actually fetched (a ``None`` value means fetched-but-absent). Only + the ``zarr_format=None`` probe overlaps with a group probe, so only that case + populates these keys. + """ + + zarr_json: NotRequired[Buffer | None] + zattrs: NotRequired[Buffer | None] + + +@dataclass(frozen=True, kw_only=True) +class ArrayProbeResult: + """Outcome of probing a path for array metadata. + + Private. Either ``metadata`` is set (the path holds an array) or ``error`` is + set (it does not, or the request was invalid). ``probe`` carries any reusable + buffers fetched during a ``zarr_format=None`` probe so a caller can hand them + to ``AsyncGroup.open`` without re-fetching. + """ + + metadata: dict[str, JSON] | None = None + error: Exception | None = None + probe: ArrayMetadataProbe = field(default_factory=ArrayMetadataProbe) + + +async def _probe_array_metadata( store_path: StorePath, zarr_format: ZarrFormat | None = 3 -) -> dict[str, JSON]: +) -> ArrayProbeResult: + """Fetch and parse array metadata without raising on "not an array". + + Returns an ``ArrayProbeResult`` whose ``metadata`` is set on success and whose + ``error`` is set (but not raised) when the path does not hold an array or the + request is invalid. ``probe`` always carries the reusable buffers fetched + during a ``zarr_format=None`` probe, even when ``error`` is set, which is + exactly when ``zarr.open`` needs to fall back to a group open. + """ + probe: ArrayMetadataProbe = {} if zarr_format == 2: zarray_bytes, zattrs_bytes = await gather( (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype), @@ -247,7 +289,7 @@ async def get_array_metadata( "A Zarr V2 array metadata document was not found in store " f"{store_path.store!r} at path {store_path.path!r}." ) - raise ArrayNotFoundError(msg) + return ArrayProbeResult(error=ArrayNotFoundError(msg), probe=probe) elif zarr_format == 3: zarr_json_bytes = await (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype) if zarr_json_bytes is None: @@ -255,13 +297,16 @@ async def get_array_metadata( "A Zarr V3 array metadata document was not found in store " f"{store_path.store!r} at path {store_path.path!r}." ) - raise ArrayNotFoundError(msg) + return ArrayProbeResult(error=ArrayNotFoundError(msg), probe=probe) elif zarr_format is None: zarr_json_bytes, zarray_bytes, zattrs_bytes = await gather( (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype), (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype), (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype), ) + # Record the buffers that overlap with a group probe so the caller may + # reuse them. These are recorded even when array detection fails below. + probe = {"zarr_json": zarr_json_bytes, "zattrs": zattrs_bytes} if zarr_json_bytes is not None and zarray_bytes is not None: # warn and favor v3 msg = f"Both zarr.json (Zarr format 3) and .zarray (Zarr format 2) metadata objects exist at {store_path}. Zarr v3 will be used." @@ -271,7 +316,7 @@ async def get_array_metadata( f"Neither Zarr V3 nor Zarr V2 array metadata documents " f"were found in store {store_path.store!r} at path {store_path.path!r}." ) - raise ArrayNotFoundError(msg) + return ArrayProbeResult(error=ArrayNotFoundError(msg), probe=probe) # set zarr_format based on which keys were found if zarr_json_bytes is not None: zarr_format = 3 @@ -279,7 +324,7 @@ async def get_array_metadata( zarr_format = 2 else: msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." # type: ignore[unreachable] - raise MetadataValidationError(msg) + return ArrayProbeResult(error=MetadataValidationError(msg), probe=probe) metadata_dict: dict[str, JSON] if zarr_format == 2: @@ -293,9 +338,22 @@ async def get_array_metadata( assert zarr_json_bytes is not None metadata_dict = buffer_to_json_object(zarr_json_bytes) - parse_node_type_array(metadata_dict.get("node_type")) + try: + parse_node_type_array(metadata_dict.get("node_type")) + except NodeTypeValidationError as exc: + return ArrayProbeResult(error=exc, probe=probe) + + return ArrayProbeResult(metadata=metadata_dict, probe=probe) - return metadata_dict + +async def get_array_metadata( + store_path: StorePath, zarr_format: ZarrFormat | None = 3 +) -> dict[str, JSON]: + result = await _probe_array_metadata(store_path, zarr_format=zarr_format) + if result.error is not None: + raise result.error + assert result.metadata is not None + return result.metadata @dataclass(frozen=True) diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 52eaa3e144..c18be46990 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -8,7 +8,7 @@ from collections import defaultdict from dataclasses import asdict, dataclass, field, fields, replace from itertools import accumulate -from typing import TYPE_CHECKING, Literal, assert_never, cast, overload +from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict, assert_never, cast, overload import numpy as np import numpy.typing as npt @@ -83,6 +83,21 @@ logger = logging.getLogger("zarr.group") +class _PreFetchedGroupMetadata(TypedDict): + """Already-fetched metadata buffers that ``AsyncGroup.open`` may reuse. + + This is a private optimization used by ``zarr.open`` to avoid re-fetching + the ``zarr.json`` and ``.zattrs`` documents it already read while probing + for an array. Only the format-detection documents are cached here; the + consolidated-metadata document is still fetched by ``AsyncGroup.open`` + because the key depends on ``use_consolidated``. + """ + + zarr_json: NotRequired[Buffer | None] + zgroup: NotRequired[Buffer | None] + zattrs: NotRequired[Buffer | None] + + def parse_zarr_format(data: Any) -> ZarrFormat: """Parse the zarr_format field from metadata.""" if data in (2, 3): @@ -490,6 +505,7 @@ async def open( store: StoreLike, zarr_format: ZarrFormat | None = 3, use_consolidated: bool | str | None = None, + _pre_fetched_metadata: _PreFetchedGroupMetadata | None = None, ) -> AsyncGroup: """Open a new AsyncGroup @@ -514,6 +530,11 @@ async def open( Zarr format 2 allowed configuring the key storing the consolidated metadata (``.zmetadata`` by default). Specify the custom key as ``use_consolidated`` to load consolidated metadata from a non-default key. + _pre_fetched_metadata : _PreFetchedGroupMetadata or None, default None + Private. Already-fetched metadata buffers (``zarr.json``, ``.zgroup``, + ``.zattrs``) that this method may reuse instead of re-fetching. Only + consulted when ``zarr_format`` is None. Used by ``zarr.open`` to avoid + duplicate fetches when falling back from an array probe to a group open. """ store_path = await make_store_path(store) if not store_path.store.supports_consolidated_metadata: @@ -554,16 +575,38 @@ async def open( if zarr_json_bytes is None: raise FileNotFoundError(store_path) elif zarr_format is None: + # Reuse any buffers that a caller (e.g. ``zarr.open``) already + # fetched while probing for an array, only fetching what is missing. + pre_fetched = _pre_fetched_metadata or {} + keys = (ZARR_JSON, ZGROUP_JSON, ZATTRS_JSON, str(consolidated_key)) + cached = ( + pre_fetched.get("zarr_json"), + pre_fetched.get("zgroup"), + pre_fetched.get("zattrs"), + None, # consolidated metadata is never pre-fetched + ) + present = ( + "zarr_json" in pre_fetched, + "zgroup" in pre_fetched, + "zattrs" in pre_fetched, + False, + ) + results = await asyncio.gather( + *( + (store_path / key).get() + for key, is_present in zip(keys, present, strict=True) + if not is_present + ) + ) + results_iter = iter(results) ( zarr_json_bytes, zgroup_bytes, zattrs_bytes, maybe_consolidated_metadata_bytes, - ) = await asyncio.gather( - (store_path / ZARR_JSON).get(), - (store_path / ZGROUP_JSON).get(), - (store_path / ZATTRS_JSON).get(), - (store_path / str(consolidated_key)).get(), + ) = tuple( + value if is_present else next(results_iter) + for value, is_present in zip(cached, present, strict=True) ) if zarr_json_bytes is not None and zgroup_bytes is not None: # warn and favor v3 diff --git a/tests/test_api.py b/tests/test_api.py index 788519969d..1b1e7793e2 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,5 +1,6 @@ from __future__ import annotations +import collections import inspect import re from typing import TYPE_CHECKING, Any @@ -13,8 +14,10 @@ if TYPE_CHECKING: from collections.abc import Callable from pathlib import Path + from typing import Self - from zarr.abc.store import Store + from zarr.abc.store import ByteRequest + from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.common import JSON, MemoryOrder, ZarrFormat from zarr.types import AnyArray @@ -30,6 +33,7 @@ import zarr.api.synchronous import zarr.core.group from zarr import Array, Group +from zarr.abc.store import Store from zarr.api.synchronous import ( create, create_array, @@ -51,6 +55,7 @@ ) from zarr.storage import MemoryStore from zarr.storage._utils import normalize_path +from zarr.storage._wrapper import WrapperStore from zarr.testing.utils import gpu_test @@ -1313,6 +1318,115 @@ async def test_open_falls_back_to_open_group_async(zarr_format: ZarrFormat) -> N assert group.attrs == {"key": "value"} +class _CountingStore(WrapperStore[Store]): + """A wrapper store that records the key of every ``get`` request. + + Used to assert that ``zarr.open`` does not fetch the same metadata key twice + when it falls back from an array probe to a group open. + """ + + get_counts: collections.Counter[str] + + def __init__(self, store: Store) -> None: + super().__init__(store) + self.get_counts = collections.Counter() + + def _with_store(self, store: Store) -> Self: + new = type(self)(store) + new.get_counts = self.get_counts + return new + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + self.get_counts[key] += 1 + return await self._store.get(key, prototype, byte_range) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_open_group_no_duplicate_fetches(zarr_format: ZarrFormat) -> None: + """``zarr.open`` on a group path must not fetch any metadata key twice. + + Regression test for the array-then-group double probe: ``open`` used to call + ``get_array_metadata`` (fetching zarr.json + .zattrs) and then re-fetch those + same keys via ``open_group``. + """ + store = _CountingStore(MemoryStore()) + await zarr.api.asynchronous.open_group( + store, attributes={"key": "value"}, zarr_format=zarr_format + ) + + store.get_counts.clear() + group = await zarr.api.asynchronous.open(store=store) + assert isinstance(group, zarr.core.group.AsyncGroup) + assert group.metadata.zarr_format == zarr_format + assert group.attrs == {"key": "value"} + + duplicates = {key: count for key, count in store.get_counts.items() if count > 1} + assert duplicates == {}, f"duplicate metadata fetches: {duplicates}" + # Each relevant key is fetched at most once: zarr.json, .zarray, .zattrs, + # .zgroup, and the consolidated key .zmetadata. Before the refactor, open() + # fetched 7 keys (zarr.json and .zattrs twice). Now it is 5 with no + # duplicates -- the two overlapping keys are no longer re-fetched. + assert sum(store.get_counts.values()) <= 5 + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_open_array_via_open_unchanged(zarr_format: ZarrFormat) -> None: + """``zarr.open`` on an array path still returns the array (no group fallback).""" + store = MemoryStore() + await zarr.api.asynchronous.create_array( + store, shape=(10,), dtype="uint8", zarr_format=zarr_format, attributes={"k": "v"} + ) + arr = await zarr.api.asynchronous.open(store=store) + assert isinstance(arr, AsyncArray) + assert arr.metadata.zarr_format == zarr_format + assert arr.attrs == {"k": "v"} + + +async def test_open_both_formats_present_warns() -> None: + """The both-formats warning still fires through ``zarr.open``'s probe.""" + store = MemoryStore() + zarr.create_array(store, shape=(10,), dtype="uint8", zarr_format=3) + zarr.create_array(store, shape=(10,), dtype="uint8", zarr_format=2) + msg = ( + f"Both zarr.json (Zarr format 3) and .zarray (Zarr format 2) metadata " + f"objects exist at {store}. Zarr v3 will be used." + ) + with pytest.warns(ZarrUserWarning, match=re.escape(msg)): + result = await zarr.api.asynchronous.open(store=store) + assert isinstance(result, AsyncArray) + assert result.metadata.zarr_format == 3 + + +async def test_open_missing_node_raises() -> None: + """``zarr.open`` on an empty store raises the same error as before.""" + store = MemoryStore() + with pytest.raises(FileNotFoundError): + await zarr.api.asynchronous.open(store=store, mode="r") + + +async def test_open_consolidated_via_open() -> None: + """``use_consolidated`` still works when passed through ``zarr.open``.""" + store = MemoryStore() + root = await zarr.api.asynchronous.open_group(store, zarr_format=2) + await root.create_array(name="a", shape=(3,), dtype="uint8") + await zarr.api.asynchronous.consolidate_metadata(store, zarr_format=2) + + group = await zarr.api.asynchronous.open(store=store, use_consolidated=True) + assert isinstance(group, zarr.core.group.AsyncGroup) + assert group.metadata.consolidated_metadata is not None + assert "a" in group.metadata.consolidated_metadata.metadata + + # use_consolidated=False must discard consolidated metadata via open(). + group_no_cons = await zarr.api.asynchronous.open(store=store, use_consolidated=False) + assert isinstance(group_no_cons, zarr.core.group.AsyncGroup) + assert group_no_cons.metadata.consolidated_metadata is None + + @pytest.mark.parametrize("mode", ["r", "r+", "w", "a"]) def test_open_modes_creates_group(tmp_path: Path, mode: str) -> None: # https://github.com/zarr-developers/zarr-python/issues/2490