From 8da13336dfa60c2b4412097bc5e58488aebbb66c Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sun, 31 May 2026 19:28:07 +0200 Subject: [PATCH 1/2] chore(deps): bump the actions group across 1 directory with 8 updates (#176) Bumps the actions group with 8 updates in the / directory: | Package | From | To | | --- | --- | --- | | [prefix-dev/setup-pixi](https://github.com/prefix-dev/setup-pixi) | `0.9.5` | `0.9.6` | | [codecov/codecov-action](https://github.com/codecov/codecov-action) | `6.0.0` | `6.0.1` | | [github/issue-metrics](https://github.com/github/issue-metrics) | `4.2.2` | `4.2.7` | | [j178/prek-action](https://github.com/j178/prek-action) | `2.0.3` | `2.0.4` | | [actions/upload-artifact](https://github.com/actions/upload-artifact) | `7.0.0` | `7.0.1` | | [actions/download-artifact](https://github.com/actions/download-artifact) | `7.0.0` | `8.0.1` | | [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) | `1.13.0` | `1.14.0` | | [zizmorcore/zizmor-action](https://github.com/zizmorcore/zizmor-action) | `0.5.3` | `0.5.6` | Updates `prefix-dev/setup-pixi` from 0.9.5 to 0.9.6 - [Release notes](https://github.com/prefix-dev/setup-pixi/releases) - [Commits](https://github.com/prefix-dev/setup-pixi/compare/1b2de7f3351f171c8b4dfeb558c639cb58ed4ec0...5185adfbffb4bd703da3010310260805d89ebb11) Updates `codecov/codecov-action` from 6.0.0 to 6.0.1 - [Release notes](https://github.com/codecov/codecov-action/releases) - [Changelog](https://github.com/codecov/codecov-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/codecov/codecov-action/compare/57e3a136b779b570ffcdbf80b3bdc90e7fab3de2...e79a6962e0d4c0c17b229090214935d2e33f8354) Updates `github/issue-metrics` from 4.2.2 to 4.2.7 - [Release notes](https://github.com/github/issue-metrics/releases) - [Commits](https://github.com/github/issue-metrics/compare/c9e9838147fd355dace335ba787f01b6641a400a...1e38d5e62363e14db8019ed7d106b9855bdba6cc) Updates `j178/prek-action` from 2.0.3 to 2.0.4 - [Release notes](https://github.com/j178/prek-action/releases) - [Commits](https://github.com/j178/prek-action/compare/6ad80277337ad479fe43bd70701c3f7f8aa74db3...bdca6f102f98e2b4c7029491a53dfd366469e33d) Updates `actions/upload-artifact` from 7.0.0 to 7.0.1 - [Release notes](https://github.com/actions/upload-artifact/releases) - [Commits](https://github.com/actions/upload-artifact/compare/v7...043fb46d1a93c77aae656e7c1c64a875d1fc6a0a) Updates `actions/download-artifact` from 7.0.0 to 8.0.1 - [Release notes](https://github.com/actions/download-artifact/releases) - [Commits](https://github.com/actions/download-artifact/compare/v7...3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c) Updates `pypa/gh-action-pypi-publish` from 1.13.0 to 1.14.0 - [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases) - [Commits](https://github.com/pypa/gh-action-pypi-publish/compare/v1.13.0...cef221092ed1bacb1cc03d23a2d87d1d172e277b) Updates `zizmorcore/zizmor-action` from 0.5.3 to 0.5.6 - [Release notes](https://github.com/zizmorcore/zizmor-action/releases) - [Commits](https://github.com/zizmorcore/zizmor-action/compare/b1d7e1fb5de872772f31590499237e7cce841e8e...5f14fd08f7cf1cb1609c1e344975f152c7ee938d) --- updated-dependencies: - dependency-name: prefix-dev/setup-pixi dependency-version: 0.9.6 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: codecov/codecov-action dependency-version: 6.0.1 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: github/issue-metrics dependency-version: 4.2.7 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: j178/prek-action dependency-version: 2.0.4 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: actions/upload-artifact dependency-version: 7.0.1 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: actions/download-artifact dependency-version: 8.0.1 dependency-type: direct:production update-type: version-update:semver-major dependency-group: actions - dependency-name: pypa/gh-action-pypi-publish dependency-version: 1.14.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: actions - dependency-name: zizmorcore/zizmor-action dependency-version: 0.5.6 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/downstream.yml | 2 +- .github/workflows/gpu_test.yml | 2 +- .github/workflows/hypothesis.yaml | 2 +- .github/workflows/issue-metrics.yml | 2 +- .github/workflows/lint.yml | 2 +- .github/workflows/test.yml | 4 ++-- .github/workflows/zarr-metadata-release.yml | 12 ++++++------ .github/workflows/zizmor.yml | 2 +- 8 files changed, 14 insertions(+), 14 deletions(-) diff --git a/.github/workflows/downstream.yml b/.github/workflows/downstream.yml index 74026233c4..3eb6898895 100644 --- a/.github/workflows/downstream.yml +++ b/.github/workflows/downstream.yml @@ -34,7 +34,7 @@ jobs: persist-credentials: false - name: Set up pixi - uses: prefix-dev/setup-pixi@1b2de7f3351f171c8b4dfeb558c639cb58ed4ec0 # v0.9.5 + uses: prefix-dev/setup-pixi@5185adfbffb4bd703da3010310260805d89ebb11 # v0.9.6 with: manifest-path: xarray/pixi.toml diff --git a/.github/workflows/gpu_test.yml b/.github/workflows/gpu_test.yml index 403441b306..333769cb9e 100644 --- a/.github/workflows/gpu_test.yml +++ b/.github/workflows/gpu_test.yml @@ -76,7 +76,7 @@ jobs: hatch env run --env "$HATCH_ENV" run-coverage - name: Upload coverage - uses: codecov/codecov-action@57e3a136b779b570ffcdbf80b3bdc90e7fab3de2 # v6.0.0 + uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 with: token: ${{ secrets.CODECOV_TOKEN }} flags: gpu diff --git a/.github/workflows/hypothesis.yaml b/.github/workflows/hypothesis.yaml index 4f9467be7d..a456b2aa0a 100644 --- a/.github/workflows/hypothesis.yaml +++ b/.github/workflows/hypothesis.yaml @@ -93,7 +93,7 @@ jobs: key: cache-hypothesis-${{ runner.os }}-${{ github.run_id }} - name: Upload coverage - uses: codecov/codecov-action@57e3a136b779b570ffcdbf80b3bdc90e7fab3de2 # v6.0.0 + uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 with: token: ${{ secrets.CODECOV_TOKEN }} flags: tests diff --git a/.github/workflows/issue-metrics.yml b/.github/workflows/issue-metrics.yml index 14fba5b9ec..510849ef3e 100644 --- a/.github/workflows/issue-metrics.yml +++ b/.github/workflows/issue-metrics.yml @@ -33,7 +33,7 @@ jobs: echo "last_month=$first_day..$last_day" >> "$GITHUB_ENV" - name: Run issue-metrics tool - uses: github/issue-metrics@c9e9838147fd355dace335ba787f01b6641a400a # v4.2.2 + uses: github/issue-metrics@1e38d5e62363e14db8019ed7d106b9855bdba6cc # v4.2.7 env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} SEARCH_QUERY: 'repo:zarr-developers/zarr-python is:issue created:${{ env.last_month }} -reason:"not planned"' diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 768e660ec2..fec211b4dd 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -30,4 +30,4 @@ jobs: uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 with: enable-cache: true - - uses: j178/prek-action@6ad80277337ad479fe43bd70701c3f7f8aa74db3 # v2.0.3 + - uses: j178/prek-action@bdca6f102f98e2b4c7029491a53dfd366469e33d # v2.0.4 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 03143d3e5b..62e571856b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -78,7 +78,7 @@ jobs: hatch env run --env "$HATCH_ENV" run-coverage - name: Upload coverage if: ${{ matrix.dependency-set == 'optional' && matrix.os == 'ubuntu-latest' }} - uses: codecov/codecov-action@57e3a136b779b570ffcdbf80b3bdc90e7fab3de2 # v6.0.0 + uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 with: token: ${{ secrets.CODECOV_TOKEN }} flags: tests @@ -125,7 +125,7 @@ jobs: run: | hatch env run --env "$HATCH_ENV" run-coverage - name: Upload coverage - uses: codecov/codecov-action@57e3a136b779b570ffcdbf80b3bdc90e7fab3de2 # v6.0.0 + uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 with: token: ${{ secrets.CODECOV_TOKEN }} flags: tests diff --git a/.github/workflows/zarr-metadata-release.yml b/.github/workflows/zarr-metadata-release.yml index 809d502f16..9639fcfdd3 100644 --- a/.github/workflows/zarr-metadata-release.yml +++ b/.github/workflows/zarr-metadata-release.yml @@ -35,7 +35,7 @@ jobs: - name: Build run: hatch build - - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0 + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: zarr-metadata-dist path: packages/zarr-metadata/dist @@ -45,7 +45,7 @@ jobs: needs: [build] runs-on: ubuntu-latest steps: - - uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7.0.0 + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: zarr-metadata-dist path: dist @@ -76,7 +76,7 @@ jobs: id-token: write # required for OIDC trusted publishing attestations: write # required for artifact attestations steps: - - uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7.0.0 + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: zarr-metadata-dist path: dist @@ -87,7 +87,7 @@ jobs: subject-path: dist/* - name: Publish package to PyPI - uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 upload_testpypi: name: Upload to TestPyPI @@ -101,7 +101,7 @@ jobs: id-token: write attestations: write steps: - - uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7.0.0 + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: zarr-metadata-dist path: dist @@ -112,6 +112,6 @@ jobs: subject-path: dist/* - name: Publish package to TestPyPI - uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # v1.13.0 + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 with: repository-url: https://test.pypi.org/legacy/ diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index da19f22421..7ac4fe5d0e 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -32,4 +32,4 @@ jobs: persist-credentials: false - name: Run zizmor - uses: zizmorcore/zizmor-action@b1d7e1fb5de872772f31590499237e7cce841e8e # v0.5.3 + uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6 From b1616e8abca8f48de40b76c0e74bf18e935e9149 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 12 Jun 2026 14:33:21 +0200 Subject: [PATCH 2/2] Avoid duplicate metadata fetches when zarr.open() falls back from array to group `zarr.open()` with `zarr_format=None` probed for array metadata (3 GETs: zarr.json, .zarray, .zattrs) and, on a group path, then re-probed via `open_group()` (4 GETs: zarr.json, .zgroup, .zattrs, .zmetadata), fetching zarr.json and .zattrs twice. This refactors the array probe (`_probe_array_metadata`) to return the buffers it fetched, and threads them into the group-open path via a private `_pre_fetched_metadata` parameter on `open_group`/`AsyncGroup.open`. The group open now reuses the already-fetched zarr.json and .zattrs documents instead of re-fetching them. v2/v3 precedence, the both-formats warning, consolidated-metadata handling, and all exception types/messages are preserved unchanged. For a v3 group via `open()`: 7 GETs (2 duplicated) -> 5 GETs (0 duplicated). Co-Authored-By: Claude Fable 5 --- src/zarr/api/asynchronous.py | 59 +++++++++++++++--- src/zarr/api/synchronous.py | 7 +++ src/zarr/core/array.py | 74 +++++++++++++++++++--- src/zarr/core/group.py | 55 +++++++++++++++-- tests/test_api.py | 116 ++++++++++++++++++++++++++++++++++- 5 files changed, 287 insertions(+), 24 deletions(-) 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