diff --git a/.github/workflows/zarr-metadata.yml b/.github/workflows/zarr-metadata.yml index df7d96cc1c..b5f56dd508 100644 --- a/.github/workflows/zarr-metadata.yml +++ b/.github/workflows/zarr-metadata.yml @@ -1,5 +1,8 @@ name: zarr-metadata +# Job steps delegate to packages/zarr-metadata/justfile, the single source of +# truth for this package's verbs; CI owns only the python matrix and caching. + on: push: branches: [main] @@ -39,12 +42,14 @@ jobs: uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 with: enable-cache: true + - name: Install just + uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 - name: Set up Python ${{ matrix.python-version }} run: uv python install ${{ matrix.python-version }} - name: Sync test dependency group run: uv sync --group test --python ${{ matrix.python-version }} - name: Run pytest - run: uv run --group test pytest tests + run: just test ruff: name: ruff @@ -59,8 +64,10 @@ jobs: persist-credentials: false - name: Install uv uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - name: Install just + uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 - name: Run ruff - run: uvx ruff check . + run: just lint pyright: name: pyright @@ -77,19 +84,35 @@ jobs: uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 with: enable-cache: true - - name: Set up Python - run: uv python install 3.11 - - name: Sync test dependency group - run: uv sync --group test --python 3.11 + - name: Install just + uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 - name: Run pyright - # Pinned to the last version that types PEP 661 sentinels in class - # attributes correctly; 1.1.405+ regressed (microsoft/pyright#11115). - # Unpin when the fix lands. - run: uv run --group test --with 'pyright==1.1.404' pyright src + # The pyright version and interpreter pins live in the justfile. + run: just typecheck + + docs: + name: docs + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-metadata + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + with: + enable-cache: true + - name: Install just + uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 + - name: Build docs + run: just docs-check zarr-metadata-complete: name: zarr-metadata complete - needs: [test, ruff, pyright] + needs: [test, ruff, pyright, docs] if: always() runs-on: ubuntu-latest steps: diff --git a/changes/4201.bugfix.md b/changes/4201.bugfix.md new file mode 100644 index 0000000000..d837a8a9e2 --- /dev/null +++ b/changes/4201.bugfix.md @@ -0,0 +1 @@ +Fixed the opt-in `FusedCodecPipeline` for serializers that advertise the partial-decode/encode mixins with only the documented async partial methods: the partial dispatch previously asserted on the private `_decode_partial_sync`/`_encode_partial_sync` hooks (an `AssertionError`, or an `AttributeError` mid-IO under `python -O`); such codecs now take the full-chunk sync path. diff --git a/changes/4202.bugfix.md b/changes/4202.bugfix.md new file mode 100644 index 0000000000..6130fc5b33 --- /dev/null +++ b/changes/4202.bugfix.md @@ -0,0 +1,10 @@ +Fixed `FusedCodecPipeline` (the opt-in synchronous pipeline) silently skipping +array-array/bytes-bytes codecs placed outside a sharding serializer on its +partial-decode/partial-encode fast paths. With an outer compressor (e.g. +`compressors=[GzipCodec()]` around a `ShardingCodec` serializer), the fused +pipeline wrote non-conforming stored bytes that `BatchedCodecPipeline` (and any +other conforming reader) could not read, and could fail to read data that +`BatchedCodecPipeline` had written. With an outer array-array codec (e.g. +`TransposeCodec`), it silently returned wrong data in both directions with no +error. Only the opt-in `FusedCodecPipeline` was affected; the default +`BatchedCodecPipeline` was never impacted. diff --git a/mkdocs.yml b/mkdocs.yml index 46bfc1764c..87aaf23430 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -92,6 +92,7 @@ nav: - ' zarr.testing.utils': api/zarr/testing/utils.md - ' zarr.zeros': api/zarr/functions/zeros.md - ' zarr.zeros_like': api/zarr/functions/zeros_like.md + - 'zarr-metadata ↪': https://zarr-metadata.readthedocs.io/ - release-notes.md - contributing.md hooks: diff --git a/packages/zarr-metadata/.readthedocs.yaml b/packages/zarr-metadata/.readthedocs.yaml new file mode 100644 index 0000000000..b89846f570 --- /dev/null +++ b/packages/zarr-metadata/.readthedocs.yaml @@ -0,0 +1,20 @@ +# Read the Docs configuration for the zarr-metadata docs site, separate from +# the zarr-python site configured by the repo-root .readthedocs.yaml. The RTD +# project for zarr-metadata must set its configuration-file path to +# packages/zarr-metadata/.readthedocs.yaml. +version: 2 + +build: + os: ubuntu-22.04 + tools: + python: "3.12" + jobs: + install: + - pip install --upgrade pip + - pip install ./packages/zarr-metadata --group packages/zarr-metadata/pyproject.toml:docs + build: + html: + - mkdocs build --strict -f packages/zarr-metadata/mkdocs.yml --site-dir $READTHEDOCS_OUTPUT/html + +mkdocs: + configuration: packages/zarr-metadata/mkdocs.yml diff --git a/packages/zarr-metadata/README.md b/packages/zarr-metadata/README.md index 69b80d7332..6b6b172aec 100644 --- a/packages/zarr-metadata/README.md +++ b/packages/zarr-metadata/README.md @@ -2,6 +2,8 @@ Python types, models, and validators for Zarr v2 and v3 metadata. +Documentation: + ## What this is Two layers and an optional integration: @@ -82,6 +84,23 @@ store I/O. The models begin and end at the metadata documents themselves — `from_key_value` / `to_key_value` map documents to store keys and bytes, and everything past that belongs to consumer libraries. +## Developing + +Package-scoped development commands live in the [`justfile`](./justfile) +(requires [just](https://github.com/casey/just)): + +``` +just test # run the test suite (extra args go to pytest) +just lint # ruff, same invocation as CI +just typecheck # pyright, pinned to the version CI uses +just docs-check # strict build of the docs site +just check # all of the above +just docs-serve # serve the docs site locally +``` + +Run them from this directory, or from anywhere in the repository as +`just packages/zarr-metadata/`. + ## Releasing The package version is derived from git tags by `hatch-vcs`. Tags must diff --git a/packages/zarr-metadata/docs/api/index.md b/packages/zarr-metadata/docs/api/index.md new file mode 100644 index 0000000000..2aa39ab161 --- /dev/null +++ b/packages/zarr-metadata/docs/api/index.md @@ -0,0 +1,31 @@ +--- +title: API reference +--- + +# API reference + +The package is organized to mirror the structure of the Zarr specifications: + +- [`zarr_metadata.model`](model.md) — frozen-dataclass document models, + structural validators, loc-aware parsers, and the `UNSET` sentinel +- [`zarr_metadata.pydantic`](pydantic.md) — optional Pydantic field types + over the models +- [`zarr_metadata.v2`](v2.md) — `TypedDict` shapes for Zarr v2 documents + (`.zarray`, `.zgroup`, `.zattrs`, `.zmetadata`) +- [`zarr_metadata.v3`](v3/index.md) — `TypedDict` shapes for Zarr v3 + documents, with subpackages for [chunk grids](v3/chunk_grid.md), + [chunk key encodings](v3/chunk_key_encoding.md), [codecs](v3/codec.md), + and [data types](v3/data_type.md) + +Every public name is also re-exported at the top level, so +`from zarr_metadata import ZarrV3ArrayMetadataJSON` and +`from zarr_metadata.v3.array import ZarrV3ArrayMetadataJSON` are equivalent. + +## Common types + +A few cross-cutting aliases are exported only from the top-level +`zarr_metadata` namespace: + +::: zarr_metadata.JSONValue + +::: zarr_metadata.ZarrV3NamedConfigJSON diff --git a/packages/zarr-metadata/docs/api/model.md b/packages/zarr-metadata/docs/api/model.md new file mode 100644 index 0000000000..c82ba98f2d --- /dev/null +++ b/packages/zarr-metadata/docs/api/model.md @@ -0,0 +1,5 @@ +--- +title: model +--- + +::: zarr_metadata.model diff --git a/packages/zarr-metadata/docs/api/pydantic.md b/packages/zarr-metadata/docs/api/pydantic.md new file mode 100644 index 0000000000..edecb416a7 --- /dev/null +++ b/packages/zarr-metadata/docs/api/pydantic.md @@ -0,0 +1,5 @@ +--- +title: pydantic +--- + +::: zarr_metadata.pydantic diff --git a/packages/zarr-metadata/docs/api/v2.md b/packages/zarr-metadata/docs/api/v2.md new file mode 100644 index 0000000000..2fe5b6ec56 --- /dev/null +++ b/packages/zarr-metadata/docs/api/v2.md @@ -0,0 +1,17 @@ +--- +title: v2 +--- + +::: zarr_metadata.v2 + options: + members: false + +::: zarr_metadata.v2.array + +::: zarr_metadata.v2.group + +::: zarr_metadata.v2.attributes + +::: zarr_metadata.v2.codec + +::: zarr_metadata.v2.consolidated diff --git a/packages/zarr-metadata/docs/api/v3/chunk_grid.md b/packages/zarr-metadata/docs/api/v3/chunk_grid.md new file mode 100644 index 0000000000..724b1c9d8d --- /dev/null +++ b/packages/zarr-metadata/docs/api/v3/chunk_grid.md @@ -0,0 +1,11 @@ +--- +title: chunk_grid +--- + +::: zarr_metadata.v3.chunk_grid + options: + members: false + +::: zarr_metadata.v3.chunk_grid.regular + +::: zarr_metadata.v3.chunk_grid.rectilinear diff --git a/packages/zarr-metadata/docs/api/v3/chunk_key_encoding.md b/packages/zarr-metadata/docs/api/v3/chunk_key_encoding.md new file mode 100644 index 0000000000..bb063deb25 --- /dev/null +++ b/packages/zarr-metadata/docs/api/v3/chunk_key_encoding.md @@ -0,0 +1,11 @@ +--- +title: chunk_key_encoding +--- + +::: zarr_metadata.v3.chunk_key_encoding + options: + members: false + +::: zarr_metadata.v3.chunk_key_encoding.default + +::: zarr_metadata.v3.chunk_key_encoding.v2 diff --git a/packages/zarr-metadata/docs/api/v3/codec.md b/packages/zarr-metadata/docs/api/v3/codec.md new file mode 100644 index 0000000000..cb96d2c7d5 --- /dev/null +++ b/packages/zarr-metadata/docs/api/v3/codec.md @@ -0,0 +1,25 @@ +--- +title: codec +--- + +::: zarr_metadata.v3.codec + options: + members: false + +::: zarr_metadata.v3.codec.blosc + +::: zarr_metadata.v3.codec.bytes + +::: zarr_metadata.v3.codec.cast_value + +::: zarr_metadata.v3.codec.crc32c + +::: zarr_metadata.v3.codec.gzip + +::: zarr_metadata.v3.codec.scale_offset + +::: zarr_metadata.v3.codec.sharding_indexed + +::: zarr_metadata.v3.codec.transpose + +::: zarr_metadata.v3.codec.zstd diff --git a/packages/zarr-metadata/docs/api/v3/data_type.md b/packages/zarr-metadata/docs/api/v3/data_type.md new file mode 100644 index 0000000000..f482c33201 --- /dev/null +++ b/packages/zarr-metadata/docs/api/v3/data_type.md @@ -0,0 +1,45 @@ +--- +title: data_type +--- + +::: zarr_metadata.v3.data_type + options: + members: false + +::: zarr_metadata.v3.data_type.bool + +::: zarr_metadata.v3.data_type.int8 + +::: zarr_metadata.v3.data_type.int16 + +::: zarr_metadata.v3.data_type.int32 + +::: zarr_metadata.v3.data_type.int64 + +::: zarr_metadata.v3.data_type.uint8 + +::: zarr_metadata.v3.data_type.uint16 + +::: zarr_metadata.v3.data_type.uint32 + +::: zarr_metadata.v3.data_type.uint64 + +::: zarr_metadata.v3.data_type.float16 + +::: zarr_metadata.v3.data_type.float32 + +::: zarr_metadata.v3.data_type.float64 + +::: zarr_metadata.v3.data_type.complex64 + +::: zarr_metadata.v3.data_type.complex128 + +::: zarr_metadata.v3.data_type.raw + +::: zarr_metadata.v3.data_type.bytes + +::: zarr_metadata.v3.data_type.string + +::: zarr_metadata.v3.data_type.numpy_datetime64 + +::: zarr_metadata.v3.data_type.numpy_timedelta64 diff --git a/packages/zarr-metadata/docs/api/v3/index.md b/packages/zarr-metadata/docs/api/v3/index.md new file mode 100644 index 0000000000..f20267d372 --- /dev/null +++ b/packages/zarr-metadata/docs/api/v3/index.md @@ -0,0 +1,15 @@ +--- +title: v3 +--- + +::: zarr_metadata.v3 + options: + members: false + +::: zarr_metadata.v3.ZarrV3MetadataFieldJSON + +::: zarr_metadata.v3.array + +::: zarr_metadata.v3.group + +::: zarr_metadata.v3.consolidated diff --git a/packages/zarr-metadata/docs/index.md b/packages/zarr-metadata/docs/index.md new file mode 100644 index 0000000000..2004f2dc54 --- /dev/null +++ b/packages/zarr-metadata/docs/index.md @@ -0,0 +1,102 @@ +--- +title: zarr-metadata +--- + +# zarr-metadata + +Basic tools for modelling Zarr metadata, with minimal dependencies. + +`zarr-metadata` is developed in the +[zarr-python repository](https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-metadata) +and released independently of `zarr` itself. Install it with: + +``` +pip install zarr-metadata +``` + +## Who needs this + +This library might be useful to you if your software interacts with Zarr metadata documents. + +## What this is + +This library is *not* a full Zarr implementation. Instead, it's a collection of data structures and routines that +closely model the content of the Zarr specifications, such as: + +- **Typed JSON shapes** ([`zarr_metadata.v2`](api/v2.md) and + [`zarr_metadata.v3`](api/v3/index.md)): `TypedDict` definitions and + `Literal` aliases for the JSON documents specified by the + [Zarr v2](https://zarr-specs.readthedocs.io/en/latest/v2/v2.0.html) and + [Zarr v3](https://zarr-specs.readthedocs.io/en/latest/v3/core/index.html) + specifications, plus types for + [zarr-extensions](https://github.com/zarr-developers/zarr-extensions/) and a + few widely-used-but-unspecified entities (e.g. consolidated metadata). +- **Document models** ([`zarr_metadata.model`](api/model.md)): canonical + frozen-dataclass models of whole metadata documents, with structural + validators, loc-aware parsers, and store-key (de)serialization. A document + produced by `to_json` shares no mutable state with the model that produced + it. +- **Optional Pydantic integration** ([`zarr_metadata.pydantic`](api/pydantic.md), + requires Pydantic 2.13 or newer): each model as a Pydantic field type that + validates raw documents through the same strict parser. + +## What this is for + +The public `TypedDict` definitions describe the static JSON shape of Zarr +metadata. For strict, loc-aware validation of JSON loaded from disk, use the +model parser: + +```python +import json +from zarr_metadata.model import ZarrV3ArrayMetadata + +with open("zarr.json", "rb") as f: + raw = json.load(f) + +metadata = ZarrV3ArrayMetadata.from_json(raw) +``` + +The optional Pydantic integration delegates raw input to the same strict +parser and returns the same normalized model class: + +```python +from pydantic import TypeAdapter +import zarr_metadata.pydantic as zmp + +metadata = TypeAdapter(zmp.ZarrV3ArrayMetadata).validate_python(raw) +encoded = metadata.to_key_value()["zarr.json"] +``` + +A bare `TypeAdapter` over a public document `TypedDict` is a coercive shape +adapter, not a Zarr conformance validator; it may coerce values or discard +members that the strict model parser rejects. + +## Validation boundary + +The model validators enforce the declared document structure and a small set +of context-free consistency rules, including fixed format literals, finite +JSON numbers, non-negative dimensions, non-empty v3 codec pipelines, and one +`dimension_names` entry per array dimension. They do not interpret extension +names or configurations, resolve codec pipelines, or decide whether a data +type, chunk grid, codec, or storage transformer is supported. Those decisions +belong to consumer implementations. + +## Scope + +At minimum, this library supports what Zarr-Python needs: the complete +Zarr v2 and v3 specs, consolidated metadata, and a subset of the metadata +defined in `zarr-extensions`. We are generally open to contributions that +add types, models, or structural validation for Zarr metadata with a +published spec. + +Runtime array behavior is out of scope: nothing here encodes or decodes +chunks, resolves codec or data type names to implementations, or performs +store I/O. The models begin and end at the metadata documents themselves — +`from_key_value` / `to_key_value` map documents to store keys and bytes, +and everything past that belongs to consumer libraries. + +## Reference + +- [API reference](api/index.md) +- [Changelog](https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-metadata/CHANGELOG.md) +- [License (MIT)](https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-metadata/LICENSE.txt) diff --git a/packages/zarr-metadata/justfile b/packages/zarr-metadata/justfile new file mode 100644 index 0000000000..0f1861ed7d --- /dev/null +++ b/packages/zarr-metadata/justfile @@ -0,0 +1,58 @@ +# Development verbs for the zarr-metadata package. Recipes run with this +# directory as the working directory regardless of where `just` is invoked. + +# List available recipes +default: + @just --list + +# Run the test suite; extra args are passed to pytest +test *args: + uv run --group test pytest tests {{ args }} + +# Lint with the same invocation CI uses +lint: + uvx ruff check . + +# Pinned to the last pyright that types PEP 661 sentinels in class attributes +# correctly; 1.1.405+ regressed (microsoft/pyright#11115). Unpin when fixed. +pyright_version := "1.1.404" + +# CI runs pyright on python 3.11; the pinned pyright predates 3.14, whose +# stdlib it cannot parse, so pin the interpreter to match CI. +# Type-check the package sources +typecheck: + uv run --python 3.11 --group test --with 'pyright=={{ pyright_version }}' pyright src + +# Run everything CI runs for this package +check: lint typecheck test docs-check + +# Preview the changelog that the next release would generate +changelog-draft: + uvx towncrier build --draft --version Unreleased + +# Build this package's documentation site, warnings as errors +docs-check: + env DISABLE_MKDOCS_2_WARNING=true uv run --group docs mkdocs build --strict + +# With no argument, uses port 8000 if free, otherwise an ephemeral free port; +# an explicitly requested port is used as-is so a conflict fails loudly. +# Serve this package's documentation site +docs-serve port="": + #!/usr/bin/env bash + set -euo pipefail + port="{{ port }}" + if [ -z "$port" ]; then + port=$(uv run --group docs python -c ' + import socket + s = socket.socket() + try: + s.bind(("127.0.0.1", 8000)) + except OSError: + s.close() + s = socket.socket() + s.bind(("127.0.0.1", 0)) + print(s.getsockname()[1]) + s.close() + ') + fi + exec env DISABLE_MKDOCS_2_WARNING=true uv run --group docs mkdocs serve -a "localhost:$port" diff --git a/packages/zarr-metadata/mkdocs.yml b/packages/zarr-metadata/mkdocs.yml new file mode 100644 index 0000000000..fc54597230 --- /dev/null +++ b/packages/zarr-metadata/mkdocs.yml @@ -0,0 +1,106 @@ +site_name: zarr-metadata +# The package lives in the zarr-python monorepo; point the header source +# widget at the package directory rather than the repository root. +repo_name: zarr-python/packages/zarr-metadata +repo_url: https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-metadata +# Absolute because mkdocs would otherwise append this to repo_url's subpath. +edit_uri: https://github.com/zarr-developers/zarr-python/edit/main/packages/zarr-metadata/docs/ +site_description: Spec-defined metadata types, models, and validators for Zarr v2 and v3. +site_author: Davis Bennett +site_url: !ENV [READTHEDOCS_CANONICAL_URL, 'https://zarr-metadata.readthedocs.io/'] +docs_dir: docs +use_directory_urls: true + +nav: + - index.md + - API Reference: + - api/index.md + - ' zarr_metadata.model': api/model.md + - ' zarr_metadata.pydantic': api/pydantic.md + - ' zarr_metadata.v2': api/v2.md + - ' zarr_metadata.v3': + - api/v3/index.md + - ' zarr_metadata.v3.chunk_grid': api/v3/chunk_grid.md + - ' zarr_metadata.v3.chunk_key_encoding': api/v3/chunk_key_encoding.md + - ' zarr_metadata.v3.codec': api/v3/codec.md + - ' zarr_metadata.v3.data_type': api/v3/data_type.md + - Changelog: https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-metadata/CHANGELOG.md + +watch: + - src + +theme: + language: en + name: material + + palette: + # Light mode + - media: "(prefers-color-scheme: light)" + scheme: default + toggle: + icon: material/brightness-7 + name: Switch to dark mode + + # Dark mode + - media: "(prefers-color-scheme: dark)" + scheme: slate + toggle: + icon: material/brightness-4 + name: Switch to light mode + + font: + text: Roboto + code: Roboto Mono + + features: + - content.code.annotate + - content.code.copy + - navigation.indexes + - navigation.instant + - navigation.tracking + - search.suggest + - search.share + +plugins: + - autorefs + - search + - mkdocstrings: + enable_inventory: true + handlers: + python: + paths: [src] + options: + allow_inspection: true + docstring_section_style: list + docstring_style: numpy + inherited_members: true + line_length: 60 + separate_signature: true + show_root_heading: true + show_signature_annotations: true + show_source: true + show_symbol_type_toc: true + signature_crossrefs: true + show_if_no_docstring: true + extensions: + - griffe_inherited_docstrings + + inventories: + - https://docs.python.org/3/objects.inv + - https://zarr.readthedocs.io/en/stable/objects.inv + +markdown_extensions: + - admonition + - attr_list + - def_list + - footnotes + - md_in_html + - pymdownx.details + - pymdownx.superfences + - toc: + permalink: true + - pymdownx.highlight: + anchor_linenums: true + line_spans: __span + pygments_lang_class: true + - pymdownx.inlinehilite diff --git a/packages/zarr-metadata/pyproject.toml b/packages/zarr-metadata/pyproject.toml index edc4b696a6..6e97d26409 100644 --- a/packages/zarr-metadata/pyproject.toml +++ b/packages/zarr-metadata/pyproject.toml @@ -44,10 +44,21 @@ Homepage = "https://github.com/zarr-developers/zarr-python" Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-metadata" Issues = "https://github.com/zarr-developers/zarr-python/issues" Changelog = "https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-metadata/CHANGELOG.md" -Documentation = "https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-metadata/README.md" +Documentation = "https://zarr-metadata.readthedocs.io/" [dependency-groups] test = ["pytest", "pydantic>=2.13", "jsonschema"] +docs = [ + # Pins match the zarr-python docs environment in the repo-root + # pyproject.toml so the two sites render with the same toolchain. + "mkdocs-material==9.7.6", + "mkdocs==1.6.1", + "mkdocstrings==1.0.4", + "mkdocstrings-python==2.0.5", + "griffe-inherited-docstrings==1.1.3", + # mkdocstrings uses ruff to format rendered signatures + "ruff==0.15.20", +] [tool.hatch.version] source = "vcs" diff --git a/src/zarr/core/codec_pipeline.py b/src/zarr/core/codec_pipeline.py index 4b8831bc7b..ca760ece59 100644 --- a/src/zarr/core/codec_pipeline.py +++ b/src/zarr/core/codec_pipeline.py @@ -143,10 +143,10 @@ def pipeline_supports_partial_decode( selection non-contiguous, a BB codec can rewrite the bytes), making partial decode infeasible. - NOTE: the two pipelines currently pass different ``require_no_aa_bb`` values - (Batched: True; Fused: False). That divergence is intentional-for-now and - tracked separately; this function centralizes the predicate without changing - either pipeline's behavior. + Both pipelines pass `require_no_aa_bb=True`: an outer AA/BB codec (e.g. a + compressor wrapping a sharding serializer) must see every byte of the + chunk, so a partial branch that only re-decodes/re-encodes the inner + sharding codec would silently bypass it. """ if require_no_aa_bb and (len(array_array_codecs) + len(bytes_bytes_codecs)) != 0: return False @@ -162,8 +162,7 @@ def pipeline_supports_partial_encode( ) -> bool: """Whether a codec pipeline can encode a partial selection without a full rewrite. - Mirror of ``pipeline_supports_partial_decode`` for encoding. See its note re: - the per-pipeline ``require_no_aa_bb`` divergence. + Mirror of `pipeline_supports_partial_decode` for encoding. """ if require_no_aa_bb and (len(array_array_codecs) + len(bytes_bytes_codecs)) != 0: return False @@ -934,14 +933,11 @@ def __iter__(self) -> Iterator[Codec]: @property def supports_partial_decode(self) -> bool: - # NOTE: unlike BatchedCodecPipeline this does NOT require the AA/BB codec - # lists to be empty (require_no_aa_bb=False). That divergence is tracked - # separately; see pipeline_supports_partial_decode. return pipeline_supports_partial_decode( self.array_bytes_codec, array_array_codecs=self.array_array_codecs, bytes_bytes_codecs=self.bytes_bytes_codecs, - require_no_aa_bb=False, + require_no_aa_bb=True, ) @property @@ -950,7 +946,7 @@ def supports_partial_encode(self) -> bool: self.array_bytes_codec, array_array_codecs=self.array_array_codecs, bytes_bytes_codecs=self.bytes_bytes_codecs, - require_no_aa_bb=False, + require_no_aa_bb=True, ) def validate( @@ -1039,10 +1035,14 @@ def read_sync( # Partial-decode fast path: the AB codec owns IO (read only the # byte ranges needed for the requested selection). Same condition - # and dispatch as BatchedCodecPipeline.read_batch. - if self.supports_partial_decode: - codec = self.array_bytes_codec - assert hasattr(codec, "_decode_partial_sync") + # and dispatch as BatchedCodecPipeline.read_batch, plus a gate on the + # sync partial method: the public partial-decode contract + # (`ArrayBytesCodecPartialDecodeMixin`) only requires the async + # `_decode_partial_single`, so a codec may support partial decode + # without `_decode_partial_sync` — such codecs take the full-chunk + # path below instead. + codec = self.array_bytes_codec + if self.supports_partial_decode and hasattr(codec, "_decode_partial_sync"): def _read_one( item: tuple[Any, ArraySpec, SelectorTuple, SelectorTuple, bool], @@ -1111,10 +1111,14 @@ def write_sync( # Partial-encode path: the AB codec owns IO (read, merge, encode, # write). Same condition and calling convention as - # BatchedCodecPipeline.write_batch. - if self.supports_partial_encode: - codec = self.array_bytes_codec - assert hasattr(codec, "_encode_partial_sync") + # BatchedCodecPipeline.write_batch, plus a gate on the sync partial + # method: the public partial-encode contract + # (`ArrayBytesCodecPartialEncodeMixin`) only requires the async + # `_encode_partial_single`, so a codec may support partial encode + # without `_encode_partial_sync` — such codecs take the full-chunk + # path below instead. + codec = self.array_bytes_codec + if self.supports_partial_encode and hasattr(codec, "_encode_partial_sync"): scalar = len(value.shape) == 0 def _write_one( diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 02b4026fd9..fd86936853 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -2,21 +2,32 @@ from __future__ import annotations -from typing import Any +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING, Any import numpy as np import pytest import zarr -from zarr.abc.codec import BytesBytesCodec +from zarr.abc.codec import ( + ArrayBytesCodec, + ArrayBytesCodecPartialDecodeMixin, + ArrayBytesCodecPartialEncodeMixin, + BytesBytesCodec, +) from zarr.codecs.bytes import BytesCodec from zarr.codecs.gzip import GzipCodec from zarr.codecs.transpose import TransposeCodec from zarr.codecs.zstd import ZstdCodec from zarr.core.codec_pipeline import FusedCodecPipeline from zarr.core.config import config as zarr_config +from zarr.registry import register_codec from zarr.storage import MemoryStore, StorePath +if TYPE_CHECKING: + from zarr.core.array_spec import ArraySpec + from zarr.core.buffer import Buffer, NDBuffer + @pytest.mark.parametrize( "codecs", @@ -261,7 +272,7 @@ def test_chunk_transform_uses_runtime_prototype() -> None: """ from zarr.abc.codec import BytesBytesCodec from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import Buffer, BufferPrototype, default_buffer_prototype + from zarr.core.buffer import BufferPrototype, default_buffer_prototype from zarr.core.chunk_utils import ChunkTransform from zarr.core.dtype import get_data_type_from_native_dtype @@ -831,3 +842,127 @@ def test_async_decode_encode_passes_through_none_chunks() -> None: assert decoded[1] is None assert decoded[0] is not None np.testing.assert_array_equal(decoded[0].as_numpy_array(), data) + + +# --------------------------------------------------------------------------- +# Graceful fallback for partial-mixin codecs without private sync-partial hooks +# +# The public partial-decode/encode contract (`ArrayBytesCodecPartialDecodeMixin` +# / `ArrayBytesCodecPartialEncodeMixin`) only requires the async +# `_decode_partial_single` / `_encode_partial_single`. The fused pipeline must +# route such codecs through its full-chunk sync path instead of asserting on +# the private `_decode_partial_sync` / `_encode_partial_sync` hooks. The double +# below is a minimal conforming implementer of that contract; it guards the +# public extension API, so it must not grow the private sync-partial methods. +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class PartialMixinCodec( + ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin +): + """Serializer with sync whole-chunk methods plus ONLY async partial methods. + + This is the pre-fused public contract for partial-capable codecs: the + mixins' `_decode_partial_single` / `_encode_partial_single`. It must not + implement `_decode_partial_sync` / `_encode_partial_sync`. + """ + + inner: BytesCodec = field(default_factory=BytesCodec) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> PartialMixinCodec: + return cls() + + def to_dict(self) -> dict[str, Any]: + return {"name": "test-partial-mixin"} + + def evolve_from_array_spec(self, array_spec: ArraySpec) -> PartialMixinCodec: + return replace(self, inner=self.inner.evolve_from_array_spec(array_spec)) + + def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int: + return self.inner.compute_encoded_size(input_byte_length, chunk_spec) + + def _decode_sync(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> NDBuffer: + return self.inner._decode_sync(chunk_bytes, chunk_spec) + + def _encode_sync(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> Buffer | None: + return self.inner._encode_sync(chunk_array, chunk_spec) + + async def _decode_single(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> NDBuffer: + return self._decode_sync(chunk_bytes, chunk_spec) + + async def _encode_single(self, chunk_array: NDBuffer, chunk_spec: ArraySpec) -> Buffer | None: + return self._encode_sync(chunk_array, chunk_spec) + + async def _decode_partial_single( + self, byte_getter: Any, selection: Any, chunk_spec: ArraySpec + ) -> NDBuffer | None: + chunk_bytes = await byte_getter.get(prototype=chunk_spec.prototype) + if chunk_bytes is None: + return None + return self._decode_sync(chunk_bytes, chunk_spec)[selection] + + async def _encode_partial_single( + self, byte_setter: Any, chunk_array: NDBuffer, selection: Any, chunk_spec: ArraySpec + ) -> None: + existing = await byte_setter.get(prototype=chunk_spec.prototype) + if existing is None: + full = chunk_spec.prototype.nd_buffer.create( + shape=chunk_spec.shape, + dtype=chunk_spec.dtype.to_native_dtype(), + fill_value=chunk_spec.fill_value, + ) + else: + full = self._decode_sync(existing, chunk_spec) + full[selection] = chunk_array + encoded = self._encode_sync(full, chunk_spec) + assert encoded is not None + await byte_setter.set(encoded) + + +register_codec("test-partial-mixin", PartialMixinCodec) + +_FUSED = {"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"} +_BATCHED = {"codec_pipeline.path": "zarr.core.codec_pipeline.BatchedCodecPipeline"} + + +@pytest.mark.filterwarnings("ignore::zarr.errors.UnstableSpecificationWarning") +@pytest.mark.parametrize("dtype", ["uint8", "float64"]) +def test_partial_mixin_codec_async_partial_only_round_trip(dtype: str) -> None: + """A serializer advertising the partial mixins with only async partial + methods must round-trip under the fused pipeline: full write, full read, + partial read, partial write, plus cross-pipeline parity with + BatchedCodecPipeline.""" + data = np.arange(64, dtype=dtype).reshape(8, 8) + + with zarr_config.set(_FUSED): + store = MemoryStore() + arr = zarr.create_array( + store, + shape=(8, 8), + chunks=(4, 4), + dtype=dtype, + serializer=PartialMixinCodec(), + compressors=None, + filters=None, + fill_value=0, + ) + + pipeline = arr._async_array.codec_pipeline + assert isinstance(pipeline, FusedCodecPipeline) + assert pipeline.supports_partial_decode + assert pipeline.supports_partial_encode + assert pipeline.sync_transform is not None + + arr[:] = data + np.testing.assert_array_equal(arr[:], data) + np.testing.assert_array_equal(arr[1:5, 2:7], data[1:5, 2:7]) + + expected = data.copy() + expected[2:6, 1:3] = 7 + arr[2:6, 1:3] = expected[2:6, 1:3] + np.testing.assert_array_equal(arr[:], expected) + + with zarr_config.set(_BATCHED): + np.testing.assert_array_equal(zarr.open_array(store, mode="r")[:], expected) diff --git a/tests/test_pipeline_parity.py b/tests/test_pipeline_parity.py index 717f0f48f1..94d95c4c24 100644 --- a/tests/test_pipeline_parity.py +++ b/tests/test_pipeline_parity.py @@ -33,6 +33,8 @@ from __future__ import annotations +import warnings +from contextlib import contextmanager from typing import TYPE_CHECKING, Any import numpy as np @@ -48,7 +50,9 @@ ShardingCodec, SubchunkWriteOrder, ) +from zarr.codecs.transpose import TransposeCodec from zarr.core.config import config as zarr_config +from zarr.errors import ZarrUserWarning from zarr.storage import MemoryStore if TYPE_CHECKING: @@ -107,11 +111,15 @@ def _store_snapshot(store: MemoryStore) -> dict[str, bytes]: ("2d-unsharded", {"shape": (20, 20), "chunks": (5, 5), "shards": None}), ("2d-sharded", {"shape": (20, 20), "chunks": (5, 5), "shards": (10, 10)}), # Nested sharding: outer chunk (10,10) sharded into inner chunks (5,5). - # Restricted to bytes-only codec because combining an outer ShardingCodec - # with a compressor (gzip) triggers a ZarrUserWarning and results in a - # checksum mismatch inside the inner shard index — a known limitation, not - # a pipeline-parity bug. The bytes-only path still exercises the full - # two-level shard encoding/decoding in both pipelines. + # Restricted to the codec configs that don't set their own `serializer` + # (bytes-only, gzip): this layout supplies an explicit nested-ShardingCodec + # `serializer`, and a codec config that also sets `serializer` (e.g. + # bytes-big-endian) would silently clobber it via dict merge, dropping + # sharding from the test entirely rather than exercising it. The gzip + # config applies as an outer bytes-bytes codec around the outer + # ShardingCodec -- this is the regression coverage for the fused pipeline + # applying outer AA/BB codecs around sharding (see + # `pipeline_supports_partial_decode`/`pipeline_supports_partial_encode`). ( "2d-nested-sharded", { @@ -122,9 +130,7 @@ def _store_snapshot(store: MemoryStore) -> dict[str, bytes]: chunk_shape=(10, 10), codecs=[ShardingCodec(chunk_shape=(5, 5))], ), - # Only run with the bytes-only codec config; gzip is incompatible - # with nested sharding (see comment above). - "_codec_ids": {"bytes-only"}, + "_codec_ids": {"bytes-only", "gzip"}, }, ), ] @@ -226,6 +232,23 @@ def _matrix() -> Iterator[Any]: # --------------------------------------------------------------------------- +@contextmanager +def _ignore_sharding_combo_warning() -> Iterator[None]: + """Suppress the "combining sharding_indexed disables partial reads" warning. + + Only the nested-sharded-plus-outer-codec matrix cell emits this; scoping the + ignore filter to just its message/category (rather than blanket-disabling + warnings) keeps every other warning in the run promoted to an error as usual. + """ + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"Combining a `sharding_indexed` codec.*", + category=ZarrUserWarning, + ) + yield + + def _write_under_pipeline( pipeline_path: str, codec_kwargs: CodecConfig, @@ -244,12 +267,13 @@ def _write_under_pipeline( create_kwargs = {"dtype": "float64", **array_layout, **codec_kwargs} store = MemoryStore() with zarr_config.set({"codec_pipeline.path": pipeline_path}): - arr = zarr.create_array( - store=store, - fill_value=0, - config={"write_empty_chunks": write_empty_chunks}, - **create_kwargs, - ) + with _ignore_sharding_combo_warning(): + arr = zarr.create_array( + store=store, + fill_value=0, + config={"write_empty_chunks": write_empty_chunks}, + **create_kwargs, + ) for sel, val in sequence: arr[sel] = val contents = arr[...] @@ -259,7 +283,8 @@ def _write_under_pipeline( def _read_under_pipeline(pipeline_path: str, store: MemoryStore) -> Any: """Re-open an existing store under the chosen pipeline and read it whole.""" with zarr_config.set({"codec_pipeline.path": pipeline_path}): - arr = zarr.open_array(store=store, mode="r") + with _ignore_sharding_combo_warning(): + arr = zarr.open_array(store=store, mode="r") return arr[...] @@ -418,3 +443,81 @@ def run(pipeline_path: str) -> tuple[dict[str, bytes], Any]: f"(index_location={index_location!r}) — byte-range write fast path likely assumed " f"the wrong physical chunk order" ) + + +# --------------------------------------------------------------------------- +# Outer array-array / bytes-bytes codecs around a sharding serializer +# --------------------------------------------------------------------------- +# +# Regression coverage for FusedCodecPipeline.supports_partial_decode/encode: +# it used to allow AA/BB codecs outside the sharding codec, so its partial +# branches called ShardingCodec._decode_partial_sync/_encode_partial_sync +# directly on the raw stored value, skipping any outer filter/compressor. +# That corrupted on-disk bytes for an outer bytes-bytes codec (unreadable by +# the other pipeline) and silently produced wrong data for an outer +# array-array codec. Both configs below force the partial branches: a +# region write and a region read are included alongside the full ones. + +_OUTER_AA_BB_CONFIGS: list[tuple[str, CodecConfig]] = [ + ( + "outer-gzip-around-sharding", + { + "serializer": ShardingCodec(chunk_shape=(2, 2)), + "compressors": [GzipCodec(level=1)], + }, + ), + ( + "outer-transpose-around-sharding", + { + "filters": [TransposeCodec(order=(1, 0))], + "serializer": ShardingCodec(chunk_shape=(2, 2)), + "compressors": None, + }, + ), +] + + +@pytest.mark.parametrize(("config_id", "codec_kwargs"), _OUTER_AA_BB_CONFIGS) +@pytest.mark.parametrize( + ("writer", "reader"), + [(_BATCHED, _FUSED), (_FUSED, _BATCHED)], + ids=["batched-write-fused-read", "fused-write-batched-read"], +) +def test_pipeline_parity_outer_aa_bb_codecs( + config_id: str, + codec_kwargs: CodecConfig, + writer: str, + reader: str, +) -> None: + """Data written under one pipeline with outer AA/BB codecs must read back + correctly under the other, including through a partial write and a + partial read. + """ + shape = (8, 8) + data = (np.arange(int(np.prod(shape))).reshape(shape) + 1).astype("uint16") + store = MemoryStore() + + with zarr_config.set({"codec_pipeline.path": writer}): + with _ignore_sharding_combo_warning(): + arr = zarr.create_array( + store=store, + shape=shape, + chunks=(4, 4), + dtype=data.dtype, + fill_value=0, + **codec_kwargs, + ) + arr[...] = data + arr[2:5, 1:3] = 99 # region write -- exercises the partial-encode branch + + expected = data.copy() + expected[2:5, 1:3] = 99 + + with zarr_config.set({"codec_pipeline.path": reader}): + with _ignore_sharding_combo_warning(): + arr2 = zarr.open_array(store=store, mode="r") + full = arr2[...] + partial = arr2[1:3, 2:7] # region read -- exercises the partial-decode branch + + np.testing.assert_array_equal(full, expected) + np.testing.assert_array_equal(partial, expected[1:3, 2:7])