diff --git a/changes/4277.feature.md b/changes/4277.feature.md new file mode 100644 index 0000000000..f5c247496d --- /dev/null +++ b/changes/4277.feature.md @@ -0,0 +1,28 @@ +`zarr.registry.get_codec_class` now raises `zarr.errors.UnknownCodecError` instead of `KeyError` +when no implementation is registered for a codec, and `zarr.core.config.BadConfigError` instead of +`KeyError` when the implementation named in `config["codecs"][name]` is not registered. +`zarr.registry.get_numcodec` raises `UnknownCodecError` instead of the `ValueError` numcodecs +raises for an unregistered Zarr format 2 codec id (`numcodecs.errors.UnknownCodecError` on +numcodecs 0.15.1 and later). All of these are subclasses of `ValueError`, so `except ValueError` +is unaffected, but `except KeyError` and `except numcodecs.errors.UnknownCodecError` are. + +These errors now name Python packages known to provide the codec, so that a user who cannot read +an array learns what to install: + +``` +An implementation for codec 'wavpack' is not available. Register one explicitly using the codec +registry (see ...), or install a Python package that registers a codec implementation with +numcodecs. Known packages supporting this codec: wavpack-numcodecs. +``` + +The tables covering this live in `src/zarr/registry.py`, one per Zarr format, and include the +codecs `numcodecs` gates behind its own optional dependencies (`zfpy`, `pcodec`, `crc32c`, +`msgpack2`). Codec authors can add their published package to them. + +A codec whose `from_dict` raises `KeyError` on a malformed configuration now surfaces as +`zarr.errors.MetadataValidationError` naming the codec and the missing key. Previously it was +reported as `UnknownCodecError: Unknown codec: ''`, presenting a configuration +key as though it were a codec name, and on the `zarr.open` path a bare `KeyError` could be +swallowed by the array-then-group fallback and reported as an unrelated group error. + +`zarr.errors.UnknownCodecError` is now exported from `zarr.errors`. diff --git a/docs/user-guide/consolidated_metadata.md b/docs/user-guide/consolidated_metadata.md index 9cb4d87c89..b237a7e986 100644 --- a/docs/user-guide/consolidated_metadata.md +++ b/docs/user-guide/consolidated_metadata.md @@ -11,6 +11,10 @@ entire hierarchy, especially when the metadata is being served over a network. Consolidated metadata essentially stores all the metadata for a hierarchy in the metadata of the root Group. +This page describes how to use consolidated metadata from Python. For a precise +description of what is written to the store, intended for other implementations, +see [Consolidated metadata format](consolidated_metadata_format.md). + ## Usage If consolidated metadata is present in a Zarr Group's metadata then it is used @@ -84,9 +88,9 @@ print(output.getvalue()) !!! info "Added in version 3.1.1" The keys in the consolidated metadata are sorted prior to writing. Keys are sorted in ascending order by path depth, where a path is defined as a sequence - of strings joined by `"/"`. For keys with the same path length, lexicographic - order is used to break the tie. This behavior ensures deterministic metadata - output for a given group. + of strings joined by `"/"`. For keys with the same depth, the tie is broken by + comparing the paths after Unicode NFKC normalization and case-folding. This + behavior ensures deterministic metadata output for a given group. ### Controlling the use of consolidated metadata diff --git a/docs/user-guide/consolidated_metadata_format.md b/docs/user-guide/consolidated_metadata_format.md new file mode 100644 index 0000000000..9189523887 --- /dev/null +++ b/docs/user-guide/consolidated_metadata_format.md @@ -0,0 +1,506 @@ +# Consolidated metadata format + +This page is a specification-style description of the consolidated metadata +that Zarr-Python reads and writes. It is intended for people implementing +consolidated metadata in another library or language who want to interoperate +with Zarr-Python, and for people who need to know exactly what Zarr-Python +puts on disk. For an introduction to *using* consolidated metadata from +Python, see [Consolidated metadata](consolidated_metadata.md). + + +The key words "MUST", "MUST NOT", "SHOULD", "SHOULD NOT", and "MAY" in this +document are to be interpreted as described in +[RFC 2119](https://www.rfc-editor.org/rfc/rfc2119). + +## Motivation + +Text in this section is quoted from +[zarr-specs#309](https://github.com/zarr-developers/zarr-specs/pull/309) +(Tom Augspurger, 2024), licensed +[CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/): + +> Consolidated metadata can help reduce the time needed to load the metadata +> for an entire hierarchy, especially when the metadata is being served over a +> network. Without consolidated metadata, opening an entire hierarchy over the +> network requires an HTTP request per node. Consolidated metadata enables +> loading the metadata for every node in a hierarchy with a single HTTP +> request. + +and, from the proposal's description: + +> This PR adds a new optional field to the core metadata for consolidating all +> the metadata of all child nodes under a single object. The motivation is +> similar to consolidated metadata in Zarr V2: without consolidated metadata, +> the time to load metadata for an entire hierarchy scales linearly with the +> number of nodes. This can be costly, especially for large hierarchies served +> HTTP from a remote storage system (like Blob Storage). + +## Concepts + +A hierarchy is **consolidated at** a group (the *consolidating group*). The +consolidating group's own metadata document gains a copy of the metadata of +every node below it, at every depth. Nothing about the child nodes' own +metadata documents changes. + +Two representations of the same information appear in this document: + +Flat +: A single mapping from *path* to *metadata document*, where path is the + full path of a node relative to the consolidating group. This is the + persisted (on-disk) form for both Zarr formats. + +Nested +: Each group holds only its *immediate* children, and child groups hold + their own children recursively. This is Zarr-Python's in-memory form + (`GroupMetadata.consolidated_metadata.metadata`). It is never written to a + store and is mentioned here only because it leaks into the on-disk form in + one place (the [empty child marker](#child-groups-carry-an-empty-marker)). + +Converting between the two is lossless. `ConsolidatedMetadata.flattened_metadata` +goes nested → flat; `ConsolidatedMetadata._flat_to_nested` goes flat → nested. + +## Paths + +A **path** is the name of a node relative to the consolidating group: + +* Segments are joined with `/`. +* There is no leading or trailing `/`. +* The consolidating group itself is **not** included (it would have the empty + path). + +Given the hierarchy below, where capital letters are groups and lowercase +letters are arrays: + +```text +A/ + x + B/ + y + C/ +``` + +consolidating at `A` produces the paths `B`, `B/C`, `B/y`, `x`; consolidating +at `B` produces `C`, `y`; consolidating at `C` produces no paths (an empty +mapping, which is still written). + +!!! note "Difference from the zarr-specs#309 text" + The worked example in zarr-specs#309 reads: + + > If we consolidate the metadata at the Group ``A``, the consolidated + > metadata would have the keys ``"A", "A/B", "A/B/C", "A/B/C/x", ...``. + > + > If we consolidate the metadata at the Group ``B``, the consolidated + > metadata would have the keys ``"C", "C/x", "C/y"``. + + The second sentence matches Zarr-Python; the first does not, since it + includes the consolidating group's own name as a prefix. Zarr-Python + always uses paths relative to the consolidating group and never includes + the consolidating group itself. This is the behaviour the rest of the + proposal describes ("the path of the node relative to the node at which + the metadata is being consolidated"), so we read the first sentence as a + typo. + +## Zarr format 3 + +### Location + +Consolidated metadata is stored **inline** in the consolidating group's +`zarr.json`, under a top-level key named `consolidated_metadata`. No other +file is written. + +### Schema + +From zarr-specs#309 (Tom Augspurger, CC-BY-4.0): + +> `consolidated_metadata` +> +> An object consolidating all the Array and Group metadata of members below +> the root node in a hierarchy. +> +> | Field | Type | Description | +> |-------------------|---------------------------|-------------| +> | `metadata` | `Map` | A mapping from node path to Group or Array `Metadata` object. | +> | `kind` | const `'inline'` | The string literal `'inline'`. Reserved for future use. | +> | `must_understand` | const `False` | The boolean literal `False`. Indicates that the field is not required to load the Zarr hierarchy. | +> +> Note that *all* children Arrays and Groups should be included in +> consolidated metadata, not just the nodes immediately below the root Group. +> Children nested inside other groups should be included too as a flat list +> of nodes. The keys of `metadata` should be the path of the node relative to +> the node at which the metadata is being consolidated (i.e. the `Group` +> where this `consolidated_metadata` object is stored). +> +> Consolidated Metadata is optional. If present, then readers should use the +> consolidated metadata. When not present, readers should use the +> non-consolidated metadata located in the Store to load the data. +> +> The `kind` field indicates that consolidated metadata is stored inline in +> the root `zarr.json` object. At this time, `'inline'` is the only supported +> value for `kind`. Future versions of the specification may allow for +> consolidated metadata in other locations. + +Zarr-Python's implementation of this schema, stated as requirements: + +| Field | Writer | Reader | +|-------------------|-----------------------------------------|--------| +| `kind` | MUST write `"inline"`. | MUST reject any value other than `"inline"` (`ValueError`). | +| `must_understand` | MUST write `false`. | Not checked. | +| `metadata` | MUST write a JSON object (possibly empty). | MUST reject a non-object (`TypeError`). Each value MUST be a JSON object (`TypeError` otherwise). | + +### Values of `metadata` + +Each value is a complete node metadata document, i.e. exactly what would be +found in that node's own `zarr.json`, with the following rules: + +* The document MUST contain `zarr_format`. Readers discriminate on this first; + Zarr-Python accepts `2` or `3` **per entry**, so a consolidated Zarr format 3 + group MAY in principle contain Zarr format 2 children. Writers SHOULD NOT + rely on this. +* For `zarr_format: 3`, the document MUST contain `node_type`, one of `"group"` + or `"array"`, and is parsed as `GroupMetadata` or `ArrayV3Metadata` + respectively. +* For `zarr_format: 2`, the presence of a `shape` key marks an array; + otherwise the entry is a group. +* Array documents are written by `ArrayV3Metadata.to_dict()` and therefore + contain every key Zarr-Python normally writes (including + `storage_transformers: []`, and `attributes: {}` when empty). + +### Child groups carry an empty marker + +Every group entry in `metadata` MUST itself carry a `consolidated_metadata` +key whose value is the empty marker: + +```json +{"kind": "inline", "must_understand": false, "metadata": {}} +``` + +This is the one place the nested in-memory form shows through. The marker +does **not** mean the child group has no children: the child's descendants are +still listed in the flat mapping at the consolidating group. It exists so +that, after a reader nests the flat mapping, a group with no children is +distinguishable from a group whose children are unknown (`consolidated_metadata` +absent / `null`). Zarr-Python's writer inserts the marker for every child +group ([`consolidate_metadata`][zarr.api.asynchronous.consolidate_metadata] +and `ConsolidatedMetadata.flattened_metadata`). + +Readers SHOULD tolerate a child group entry that omits the marker; Zarr-Python +does, and normalises such entries to the empty marker when nesting. + +### Key order + +`metadata` is a JSON object and therefore unordered in principle. Since +Zarr-Python 3.1.1 the writer sorts keys deterministically so that the same +hierarchy always serialises to byte-identical output: + +1. Primarily by **depth**, ascending, where depth is the number of `/` + characters in the path. +2. Secondarily by the path string after Unicode NFKC normalisation and + case-folding (`unicodedata.normalize("NFKC", key).casefold()`), ascending. + +Readers MUST NOT depend on key order. In particular, keys sharing a parent are +not guaranteed to be adjacent, and Zarr-Python's reader groups by parent +rather than assuming adjacency. + +### Root-level fields + +Zarr-Python writes `consolidated_metadata` at the top level of the group +document alongside `zarr_format`, `node_type`, and `attributes`, not inside +`attributes`. When a group has no consolidated metadata, the key is omitted +entirely (never written as `null`). + +### Complete example + +Consolidating the hierarchy from [Paths](#paths) (with `x` an `int32` array of +shape `(10,)` chunked by `(5,)`, `y` a `float64` array of shape `(2, 2)`, and +default codecs) yields the following `zarr.json` at `A`. This output was +produced by Zarr-Python and is reformatted only for indentation. + +```json +{ + "attributes": {"title": "example"}, + "zarr_format": 3, + "consolidated_metadata": { + "kind": "inline", + "must_understand": false, + "metadata": { + "B": { + "attributes": {"kind": "child"}, + "zarr_format": 3, + "consolidated_metadata": { + "kind": "inline", + "must_understand": false, + "metadata": {} + }, + "node_type": "group" + }, + "x": { + "shape": [10], + "data_type": "int32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [5]}}, + "chunk_key_encoding": {"name": "default", "configuration": {"separator": "/"}}, + "fill_value": 0, + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "zstd", "configuration": {"level": 0, "checksum": false}} + ], + "attributes": {}, + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + }, + "B/C": { + "attributes": {}, + "zarr_format": 3, + "consolidated_metadata": { + "kind": "inline", + "must_understand": false, + "metadata": {} + }, + "node_type": "group" + }, + "B/y": { + "shape": [2, 2], + "data_type": "float64", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [2, 2]}}, + "chunk_key_encoding": {"name": "default", "configuration": {"separator": "/"}}, + "fill_value": 0.0, + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "zstd", "configuration": {"level": 0, "checksum": false}} + ], + "attributes": {}, + "zarr_format": 3, + "node_type": "array", + "storage_transformers": [] + } + } + }, + "node_type": "group" +} +``` + +Observe the key order: depth 0 (`B`, `x`), then depth 1 (`B/C`, `B/y`), with +case-folded lexical order within each depth. + +## Zarr format 2 + +### Location + +Consolidated metadata is stored in a separate document at the key +`.zmetadata` in the consolidating group's directory, alongside `.zgroup` and +`.zattrs`. A reader MAY be configured to look for a different key +([`zarr.open_group`][] accepts a string for `use_consolidated`); Zarr-Python +only ever writes `.zmetadata`. + +### Schema + +```json +{ + "zarr_consolidated_format": 1, + "metadata": { "/.zarray": {...}, "/.zattrs": {...}, "/.zgroup": {...}, ... } +} +``` + +| Field | Writer | Reader | +|----------------------------|--------------------------|--------| +| `zarr_consolidated_format` | MUST write `1`. | Not checked. | +| `metadata` | MUST write an object. | MUST be an object; each key MUST end in `/.zarray`, `/.zattrs`, or `/.zgroup` (`ValueError` otherwise), except for the bare root keys below. | + +### Keys of `metadata` + +Unlike Zarr format 3, each *node* contributes up to **two** entries, one per +underlying metadata document, keyed by `/`: + +* Arrays: `/.zarray` and `/.zattrs`. +* Groups: `/.zgroup` and `/.zattrs`. + +Zarr-Python always writes the `.zattrs` entry, as `{}` when there are no +attributes. + +The consolidating group itself **is** included, under the bare keys `.zgroup` +and `.zattrs` (no path prefix). This is a difference from Zarr format 3 and +is inherited from Zarr-Python 2.x. On read, Zarr-Python **ignores** these two +root entries and uses the real `.zgroup` and `.zattrs` documents, which it +fetches in the same request batch. + +### Values of `metadata` + +`.zarray` and `.zgroup` values are the documents that would be found at the +corresponding store key, with one deviation: + +!!! warning "Deviation from Zarr-Python 2.x" + Zarr-Python 3 writes every child `.zgroup` entry as + + ```json + { + "zarr_format": 2, + "consolidated_metadata": {"metadata": {}, "must_understand": false, "kind": "inline"} + } + ``` + + rather than the `{"zarr_format": 2}` that Zarr-Python 2.x wrote. This is + the Zarr format 3 [empty child marker](#child-groups-carry-an-empty-marker) + leaking into the Zarr format 2 encoding. Zarr-Python's reader ignores + unknown keys in `.zgroup` entries, and Zarr-Python 2.x likewise ignored + them, but other readers that validate `.zgroup` strictly may reject this. + Readers SHOULD ignore the key; writers targeting maximum compatibility + with older Zarr format 2 tooling may wish to omit it. + +### Key order + +Entries are emitted in the same order as the Zarr format 3 case (depth, then +case-folded path), with each node's `.zattrs` entry immediately preceding its +`.zarray`/`.zgroup` entry. Readers MUST NOT depend on this order. + +### Serialisation + +`.zmetadata` is written **compactly** (no indentation), regardless of the +`json_indent` configuration setting that applies to other metadata documents. + +### Complete example + +The same hierarchy as above, written as Zarr format 2 with default +compressor. Output produced by Zarr-Python, reformatted for indentation. + +```json +{ + "metadata": { + ".zgroup": {"zarr_format": 2}, + ".zattrs": {"title": "example"}, + "B/.zattrs": {"kind": "child"}, + "B/.zgroup": { + "zarr_format": 2, + "consolidated_metadata": {"metadata": {}, "must_understand": false, "kind": "inline"} + }, + "x/.zattrs": {}, + "x/.zarray": { + "shape": [10], + "chunks": [5], + "dtype": " Like Zarr v2, consolidated metadata introduces the possibility of +> "inconsistent" metadata between the consolidated and non-consolidated +> forms. Should the spec take any stance on how to handle this? I've currently +> worded things to say that readers should always use the consolidated +> metadata if it's present. + +Zarr-Python follows that wording: when consolidated metadata is present and +not explicitly disabled, it is authoritative for reads, and no attempt is made +to detect drift from the per-node documents. Writers are responsible for +re-consolidating after modifying a hierarchy. Zarr-Python does **not** +re-consolidate automatically when nodes are created or have their attributes +updated. The one exception is deletion: deleting a member of a group that was +opened with consolidated metadata (`del group[name]`) removes that member's +entry from the consolidated copy and re-writes the group's metadata document. +See +[Synchronization and concurrency](consolidated_metadata.md#synchronization-and-concurrency). + +## Known limitations and open questions + +These are noted in zarr-specs#309 and remain open: + +* **Root document size.** Quoting the proposal: "For *very* large + hierarchies, this will bloat the size of the root `zarr.json`, slowing down + operations that just want to open the metadata for the root." The `kind` + field is reserved so that a future value could point at an external + document instead. +* **Overlap with child listing.** The proposal notes overlap with + [zarr-specs#284](https://github.com/zarr-developers/zarr-specs/issues/284), + which would store child *paths* only, allowing a reader to list the + hierarchy in one request and then fetch node metadata concurrently. +* **`must_understand`.** The field is always `false` and Zarr-Python's reader + does not inspect it. It exists so that a conforming Zarr format 3 reader + that does not implement consolidated metadata can safely ignore the key + per the core specification's rules for unknown metadata fields. + +## Attribution + +Quoted passages in this document are from +[zarr-developers/zarr-specs#309, "Added consolidated metadata to spec"](https://github.com/zarr-developers/zarr-specs/pull/309) +by Tom Augspurger, licensed under +[CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/). The Zarr format 3 +implementation in Zarr-Python was contributed in +[zarr-python#2113](https://github.com/zarr-developers/zarr-python/pull/2113), +also by Tom Augspurger. The Zarr format 2 `.zmetadata` layout originates in +Zarr-Python 2.x. diff --git a/docs/user-guide/extending.md b/docs/user-guide/extending.md index f852f9105e..507afedea7 100644 --- a/docs/user-guide/extending.md +++ b/docs/user-guide/extending.md @@ -63,6 +63,13 @@ New codecs need to have their own unique identifier. To avoid naming collisions, strongly recommended to prefix the codec identifier with a unique name. For example, the codecs from `numcodecs` are prefixed with `numcodecs.`, e.g. `numcodecs.delta`. +If someone opens an array that uses your codec without your package installed, Zarr raises +[`zarr.errors.UnknownCodecError`][] explaining how to register an implementation. Zarr also +keeps a small table of codec names and the published packages that provide them, and names +those packages in that error. Once your package is on PyPI, please open a pull request adding +it to the codec-package tables in `src/zarr/registry.py`, so that users get a message telling them +exactly what to install. + !!! note Note that the extension mechanism for the Zarr format 3 is still under development. Requirements for custom codecs including the choice of codec identifiers might diff --git a/mkdocs.yml b/mkdocs.yml index 4d06701a87..80922ad4ce 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -25,6 +25,7 @@ nav: - user-guide/extending.md - user-guide/gpu.md - user-guide/consolidated_metadata.md + - user-guide/consolidated_metadata_format.md - user-guide/experimental.md - user-guide/v3_migration.md - user-guide/glossary.md diff --git a/src/zarr/codecs/gzip.py b/src/zarr/codecs/gzip.py index 7f21872034..e5b54639f9 100644 --- a/src/zarr/codecs/gzip.py +++ b/src/zarr/codecs/gzip.py @@ -29,6 +29,13 @@ def parse_gzip_level(data: JSON) -> int: def _gzip_streams_equal_except_mtime(a: bytes, b: bytes) -> bool: + """Compare two gzip streams, ignoring the MTIME field of the header. + + Per RFC 1952 the gzip header is [magic(2)][CM(1)][FLG(1)][MTIME(4)][XFL(1)][OS(1)], + so bytes 4-8 are MTIME. The fixed offsets assume the standard 10-byte header + with no FNAME/FEXTRA/FCOMMENT flags set, which holds here because numcodecs' + ``GZip.encode`` wraps ``gzip.GzipFile`` without a filename. + """ if len(a) != len(b): return False diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index fc47f8fc95..988d2d369c 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -36,7 +36,7 @@ from zarr.core.dtype.common import check_dtype_spec_v3 from zarr.core.json_parse import parse_field, validate_json_value from zarr.core.metadata.common import parse_attributes -from zarr.errors import MetadataValidationError, NodeTypeValidationError, UnknownCodecError +from zarr.errors import MetadataValidationError, NodeTypeValidationError from zarr.registry import get_codec_class if TYPE_CHECKING: @@ -74,10 +74,17 @@ def parse_codecs(data: object) -> tuple[Codec, ...]: else: name_parsed, _ = parse_named_configuration(c, require_configuration=False) + codec_cls = get_codec_class(name_parsed) try: - out += (get_codec_class(name_parsed).from_dict(c),) + out += (codec_cls.from_dict(c),) except KeyError as e: - raise UnknownCodecError(f"Unknown codec: {e.args[0]!r}") from e + # A codec's `from_dict` may index its configuration directly, so a malformed + # configuration surfaces as a KeyError. Convert it: a bare KeyError escaping + # metadata parsing is swallowed by the array-then-group fallback in + # `zarr.api.asynchronous.open`, which then reports an unrelated group error. + raise MetadataValidationError( + f"Invalid configuration for codec {name_parsed!r}: missing key {e.args[0]!r}." + ) from e return out diff --git a/src/zarr/errors.py b/src/zarr/errors.py index 781bebe534..3e445de2e9 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -12,6 +12,7 @@ "MetadataValidationError", "NegativeStepError", "NodeTypeValidationError", + "UnknownCodecError", "UnstableSpecificationWarning", "VindexInvalidSelectionError", "ZarrDeprecationWarning", diff --git a/src/zarr/metadata/migrate_v3.py b/src/zarr/metadata/migrate_v3.py index 370af75a6d..ad177e19fb 100644 --- a/src/zarr/metadata/migrate_v3.py +++ b/src/zarr/metadata/migrate_v3.py @@ -29,6 +29,7 @@ from zarr.core.metadata.v2 import ArrayV2Metadata from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGridMetadata from zarr.core.sync import sync +from zarr.errors import UnknownCodecError from zarr.registry import get_codec_class from zarr.storage import StorePath from zarr.types import AnyArray @@ -273,7 +274,7 @@ def _find_numcodecs_zarr3(numcodecs_codec: numcodecs.abc.Codec) -> Codec: try: codec_v3 = get_codec_class(numcodec_name) - except KeyError as exc: + except UnknownCodecError as exc: raise ValueError( f"Couldn't find corresponding zarr.codecs.numcodecs codec for {numcodecs_codec.codec_id}" ) from exc diff --git a/src/zarr/registry.py b/src/zarr/registry.py index c2c0eb2921..a7537e5023 100644 --- a/src/zarr/registry.py +++ b/src/zarr/registry.py @@ -2,12 +2,13 @@ import warnings from collections import defaultdict +from collections.abc import Mapping from importlib.metadata import entry_points as get_entry_points from typing import TYPE_CHECKING, Any from zarr.core.config import BadConfigError, config from zarr.core.dtype import data_type_registry -from zarr.errors import ZarrUserWarning +from zarr.errors import UnknownCodecError, ZarrUserWarning if TYPE_CHECKING: from importlib.metadata import EntryPoint @@ -23,7 +24,7 @@ from zarr.abc.numcodec import Numcodec from zarr.core.buffer import Buffer, NDBuffer from zarr.core.chunk_key_encodings import ChunkKeyEncoding - from zarr.core.common import JSON + from zarr.core.common import JSON, ZarrFormat __all__ = [ "Registry", @@ -39,6 +40,131 @@ "register_pipeline", ] +_ZARR_CODEC_DOCS_URL = "https://zarr.readthedocs.io/en/stable/user-guide/extending/#custom-codecs" +_NUMCODECS_CODEC_DOCS_URL = ( + "https://numcodecs.readthedocs.io/en/stable/registry.html#numcodecs.registry.register_codec" +) + +# Codecs zarr-python does not implement, mapped to the names of Python packages that do. +# These tables exist purely to make the "no implementation for this codec" error actionable; +# nothing here affects which codecs zarr can actually read or write. Values are what you would +# pass to `pip install`. Only add an entry you have verified against the package's declared +# entry points, and only for a package that is actually published. +# +# The two Zarr formats resolve codecs through different registries, so they get different +# tables: a name can mean one thing as a Zarr format 3 codec name and another as a Zarr +# format 2 codec id. `imagecodecs_*` is exactly that -- `virtual-tiff` declares 15 of those +# names under `zarr.codecs`, while `imagecodecs-numcodecs` declares all 81 under +# `numcodecs.codecs`, so the format 2 side can use a prefix and the format 3 side cannot. + +# Zarr format 3 codec names (entry point group "zarr.codecs"). +_CODEC_PACKAGES: dict[str, tuple[str, ...]] = { + "gribberish": ("gribberish",), + # `virtual-tiff` declares these 15 `imagecodecs_*` names, out of the 81 that exist as + # numcodecs ids. They are listed exactly rather than by prefix so that the other 66 get no + # hint instead of a hint pointing at a package that does not provide them. + "imagecodecs_deflate": ("virtual-tiff",), + "imagecodecs_delta": ("virtual-tiff",), + "imagecodecs_floatpred": ("virtual-tiff",), + "imagecodecs_jetraw": ("virtual-tiff",), + "imagecodecs_jpeg": ("virtual-tiff",), + "imagecodecs_jpeg2k": ("virtual-tiff",), + "imagecodecs_jpeg8": ("virtual-tiff",), + "imagecodecs_jpegxl": ("virtual-tiff",), + "imagecodecs_jpegxr": ("virtual-tiff",), + "imagecodecs_lerc": ("virtual-tiff",), + "imagecodecs_lzw": ("virtual-tiff",), + "imagecodecs_packbits": ("virtual-tiff",), + "imagecodecs_png": ("virtual-tiff",), + "imagecodecs_webp": ("virtual-tiff",), + "imagecodecs_zstd": ("virtual-tiff",), + "n5_default": ("zarr-n5",), +} + +# As `_CODEC_PACKAGES`, but each key is matched against the start of the codec name. Packages +# that provide many codecs namespace them behind a shared prefix, so one entry covers them all. +_CODEC_PACKAGE_PREFIXES: dict[str, tuple[str, ...]] = { + "any-numcodecs.": ("zarr-any-numcodecs",), + "omfiles.": ("omfiles",), + "virtual_tiff.": ("virtual-tiff",), +} + +# Zarr format 2 codec ids (entry point group "numcodecs.codecs"). `numcodecs` itself gates +# several of its own codecs behind optional dependencies, so the package to install for those +# is an extra of numcodecs rather than a third-party distribution. +_NUMCODEC_PACKAGES: dict[str, tuple[str, ...]] = { + "FITSAscii": ("kerchunk",), + "FITSVarBintable": ("kerchunk",), + "crc32c": ("numcodecs[crc32c]",), + "fill_hdf_strings": ("kerchunk",), + "grib": ("kerchunk",), + "msgpack2": ("numcodecs[msgpack]",), + "pcodec": ("numcodecs[pcodec]",), + "rawgrib": ("gribscan",), + "record_member": ("kerchunk",), + "vc-delta3d": ("vc-delta3d",), + "wavpack": ("wavpack-numcodecs",), + "zfpy": ("numcodecs[zfpy]",), +} + +# As `_NUMCODEC_PACKAGES`, but matched against the start of the codec id. +_NUMCODEC_PACKAGE_PREFIXES: dict[str, tuple[str, ...]] = { + "gribscan.": ("gribscan",), + "imagecodecs_": ("imagecodecs-numcodecs",), +} + + +def _packages_for_codec(name: str, *, zarr_format: ZarrFormat) -> tuple[str, ...]: + """ + Names of Python packages known to provide an implementation of the codec ``name``. + + Returns an empty tuple if we don't know of any. + + Parameters + ---------- + name : str + The codec name (Zarr format 3) or codec id (Zarr format 2) we failed to resolve. + zarr_format : ZarrFormat + Which registry the codec was looked up in. + """ + if zarr_format == 2: + exact, prefixes = _NUMCODEC_PACKAGES, _NUMCODEC_PACKAGE_PREFIXES + else: + exact, prefixes = _CODEC_PACKAGES, _CODEC_PACKAGE_PREFIXES + if name in exact: + return exact[name] + for prefix, packages in prefixes.items(): + if name.startswith(prefix): + return packages + return () + + +def _missing_codec_message(name: str, *, zarr_format: ZarrFormat) -> str: + """ + Build the error message raised when no implementation of the codec ``name`` is available. + + Parameters + ---------- + name : str + The codec name (Zarr format 3) or codec id (Zarr format 2) we failed to resolve. + zarr_format : ZarrFormat + Which registry the codec was looked up in. Zarr format 2 codecs are resolved through + numcodecs, so that case points at the numcodecs registry rather than at zarr's. + """ + if zarr_format == 2: + docs_url, registry = _NUMCODECS_CODEC_DOCS_URL, "numcodecs" + else: + docs_url, registry = _ZARR_CODEC_DOCS_URL, "zarr" + msg = ( + f"An implementation for codec {name!r} is not available. Register one explicitly " + f"using the codec registry (see {docs_url}), or install a Python package that " + f"registers a codec implementation with {registry}." + ) + packages = _packages_for_codec(name, zarr_format=zarr_format) + if packages: + msg += f" Known packages supporting this codec: {', '.join(packages)}." + return msg + class Registry[T](dict[str, type[T]]): def __init__(self) -> None: @@ -168,7 +294,7 @@ def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]: codec_classes = _codec_registries[key] if not codec_classes: - raise KeyError(key) + raise UnknownCodecError(_missing_codec_message(key, zarr_format=3)) config_entry = config.get("codecs", {}).get(key) if config_entry is None: if len(codec_classes) == 1: @@ -179,11 +305,17 @@ def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]: category=ZarrUserWarning, ) return list(codec_classes.values())[-1] - selected_codec_cls = codec_classes[config_entry] - - if selected_codec_cls: - return selected_codec_cls - raise KeyError(key) + selected_codec_cls = codec_classes.get(config_entry) + if selected_codec_cls is None: + # Not UnknownCodecError: the codec is known, the implementation named in the config is + # not registered. That is a configuration problem, which is what the sibling getters in + # this module raise BadConfigError for. + raise BadConfigError( + f"Codec {key!r} is configured to use the implementation {config_entry!r}, which is " + f"not registered. Registered implementations of this codec: " + f"{sorted(codec_classes)}." + ) + return selected_codec_cls def _resolve_codec(data: dict[str, JSON]) -> Codec: @@ -321,6 +453,13 @@ def get_numcodec(data: CodecJSON_V2[str]) -> Numcodec: ------- codec : Numcodec + Raises + ------ + UnknownCodecError + If ``data`` carries a string ``"id"`` that is not registered with numcodecs. Any other + failure, including a registered codec rejecting its configuration and a ``data`` that is + not a mapping, propagates from numcodecs unchanged. + Examples -------- ```python @@ -331,6 +470,19 @@ def get_numcodec(data: CodecJSON_V2[str]) -> Numcodec: ``` """ - from numcodecs.registry import get_codec - + from numcodecs.registry import codec_registry, entries, get_codec + + # Check whether numcodecs can resolve the id *before* handing off, rather than catching what + # `get_codec` raises. Catching cannot tell "this id is unregistered" from "a registered codec + # rejected its configuration" or from "a wrapper codec failed to resolve an inner codec", and + # relabelling either of those with this id would attach a package hint that is simply wrong. + # This mirrors the two lookups `get_codec` performs (it then tests the result for + # truthiness rather than membership, which only differs for a falsy registry value). + # Widened to `object` deliberately: `data` is annotated as a TypedDict, but this is a public + # function and callers pass whatever they like. numcodecs coerces with `dict(config)` and + # raises for anything that is not a mapping, which is the behaviour to preserve. + raw: object = data + codec_id = raw.get("id") if isinstance(raw, Mapping) else None + if isinstance(codec_id, str) and codec_id not in codec_registry and codec_id not in entries: + raise UnknownCodecError(_missing_codec_message(codec_id, zarr_format=2)) return get_codec(data) # type: ignore[no-any-return] diff --git a/tests/test_codecs/test_numcodecs.py b/tests/test_codecs/test_numcodecs.py index 99cd89492f..d78b73f34c 100644 --- a/tests/test_codecs/test_numcodecs.py +++ b/tests/test_codecs/test_numcodecs.py @@ -8,15 +8,10 @@ import pytest from numcodecs import GZip -try: - from numcodecs.errors import UnknownCodecError -except ImportError: - # Older versions of numcodecs don't have a separate errors module - UnknownCodecError = ValueError - from zarr import config, create_array, open_array from zarr.abc.numcodec import _is_numcodec, _is_numcodec_cls from zarr.codecs import numcodecs as _numcodecs +from zarr.errors import UnknownCodecError from zarr.registry import get_codec_class, get_numcodec if TYPE_CHECKING: @@ -297,12 +292,7 @@ def test_generic_checksum(codec_class: type[_numcodecs._NumcodecsBytesBytesCodec def test_generic_bytes_codec(codec_class: type[_numcodecs._NumcodecsArrayBytesCodec]) -> None: try: codec_class()._codec # noqa: B018 - except ValueError as e: # pragma: no cover - if "codec not available" in str(e): - pytest.xfail(f"{codec_class.codec_name} is not available: {e}") # type: ignore[misc] - else: - raise - except ImportError as e: # pragma: no cover + except (UnknownCodecError, ImportError) as e: # pragma: no cover pytest.xfail(f"{codec_class.codec_name} is not available: {e}") # type: ignore[misc] data = np.arange(0, 256, dtype="float32").reshape((16, 16)) diff --git a/tests/test_config.py b/tests/test_config.py index 47f71a798e..c22aac7603 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -24,7 +24,7 @@ from zarr.core.codec_pipeline import BatchedCodecPipeline from zarr.core.config import BadConfigError, config from zarr.core.indexing import SelectorTuple -from zarr.errors import ChunkNotFoundError, ZarrUserWarning +from zarr.errors import ChunkNotFoundError, UnknownCodecError, ZarrUserWarning from zarr.registry import ( fully_qualified_name, get_buffer_class, @@ -334,7 +334,7 @@ class NewCodec2(BytesCodec): pass # error if codec is not registered - with pytest.raises(KeyError): + with pytest.raises(UnknownCodecError): get_codec_class("missing_codec") # no warning if only one implementation is available diff --git a/tests/test_fused_pipeline.py b/tests/test_fused_pipeline.py index 87bfaf71b4..6025ae8caf 100644 --- a/tests/test_fused_pipeline.py +++ b/tests/test_fused_pipeline.py @@ -5,11 +5,13 @@ import asyncio from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any +from unittest.mock import patch import numpy as np import pytest import zarr +from tests.conftest import Expect from zarr.abc.codec import ( ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, @@ -21,8 +23,13 @@ from zarr.codecs.gzip import GzipCodec, _gzip_streams_equal_except_mtime from zarr.codecs.transpose import TransposeCodec from zarr.codecs.zstd import ZstdCodec -from zarr.core.codec_pipeline import FusedCodecPipeline +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer import BufferPrototype, default_buffer_prototype +from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer +from zarr.core.chunk_utils import ChunkTransform, evolve_codecs +from zarr.core.codec_pipeline import AsyncChunkTransform, FusedCodecPipeline from zarr.core.config import config as zarr_config +from zarr.core.dtype import Float32, Float64, Int32, UInt8, ZDType from zarr.registry import register_codec from zarr.storage import MemoryStore, StorePath, WrapperStore from zarr.storage._utils import _normalize_byte_range_index @@ -32,8 +39,28 @@ from collections.abc import AsyncIterator, Callable, Iterable from zarr.abc.store import ByteRequest - from zarr.core.array_spec import ArraySpec - from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer + from zarr.core.buffer import Buffer, NDBuffer + + +_FLOAT64 = Float64() + + +def _make_spec( + shape: tuple[int, ...], + zdtype: ZDType[Any, Any] = _FLOAT64, + fill_value: float = 0, + *, + write_empty_chunks: bool = True, + prototype: BufferPrototype | None = None, +) -> ArraySpec: + """An ArraySpec with the C-order defaults shared by the tests in this file.""" + return ArraySpec( + shape=shape, + dtype=zdtype, + fill_value=zdtype.cast_scalar(fill_value), + config=ArrayConfig(order="C", write_empty_chunks=write_empty_chunks), + prototype=prototype if prototype is not None else default_buffer_prototype(), + ) @pytest.mark.parametrize( @@ -62,10 +89,6 @@ def test_sync_api_compute_off_event_loop(monkeypatch: pytest.MonkeyPatch) -> Non and the penalty grows with codec cost (observed as "fused pipeline is slower for zstd data" under dask-style multi-threaded single-chunk reads). """ - import asyncio - - from zarr.core.chunk_utils import ChunkTransform - compute_on_loop = {"decode": False, "encode": False} calls = {"decode": 0, "encode": 0} real_decode = ChunkTransform.decode_chunk @@ -116,22 +139,10 @@ def traced_encode(self: ChunkTransform, chunk_array: Any, chunk_spec: Any) -> An def test_evolve_from_array_spec() -> None: """evolve_from_array_spec creates a sync transform.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.dtype import get_data_type_from_native_dtype - pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) assert pipeline.sync_transform is None - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - spec = ArraySpec( - shape=(100,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) - evolved = pipeline.evolve_from_array_spec(spec) + evolved = pipeline.evolve_from_array_spec(_make_spec((100,))) assert evolved.sync_transform is not None @@ -147,31 +158,20 @@ def test_evolve_from_array_spec() -> None: @pytest.mark.parametrize( - ("dtype", "shape"), + ("zdtype", "shape"), [ - ("float64", (100,)), - ("float32", (50,)), - ("int32", (200,)), - ("float64", (10, 10)), + (Float64(), (100,)), + (Float32(), (50,)), + (Int32(), (200,)), + (Float64(), (10, 10)), ], ids=["f64-1d", "f32-1d", "i32-1d", "f64-2d"], ) -def test_read_write_sync_roundtrip(dtype: str, shape: tuple[int, ...]) -> None: +def test_read_write_sync_roundtrip(zdtype: ZDType[Any, Any], shape: tuple[int, ...]) -> None: """Data written via write_sync can be read back via read_sync.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype - store = MemoryStore() - zdtype = get_data_type_from_native_dtype(np.dtype(dtype)) - spec = ArraySpec( - shape=shape, - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + dtype = zdtype.to_native_dtype() + spec = _make_spec(shape, zdtype) pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) pipeline = pipeline.evolve_from_array_spec(spec) @@ -200,20 +200,8 @@ def test_read_write_sync_roundtrip(dtype: str, shape: tuple[int, ...]) -> None: def test_read_sync_missing_chunk_fills() -> None: """Sync read of a missing chunk fills with the fill value.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype - store = MemoryStore() - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - spec = ArraySpec( - shape=(10,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(42.0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((10,), fill_value=42.0) pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) pipeline = pipeline.evolve_from_array_spec(spec) @@ -232,21 +220,10 @@ def test_read_sync_missing_chunk_fills() -> None: def test_sync_write_async_read_roundtrip() -> None: """Data written via write_sync can be read back via async read.""" - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype from zarr.core.sync import sync store = MemoryStore() - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - spec = ArraySpec( - shape=(100,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((100,)) pipeline = FusedCodecPipeline.from_codecs((BytesCodec(),)) pipeline = pipeline.evolve_from_array_spec(spec) @@ -271,17 +248,14 @@ def test_sync_write_async_read_roundtrip() -> None: ) ) + np.testing.assert_array_equal(data, out.as_numpy_array()) + def test_chunk_transform_uses_runtime_prototype() -> None: """ChunkTransform must pass each codec the prototype from the runtime chunk_spec, not one captured at evolve time. Constructs ChunkTransform directly (a Fused-internal data structure with no BatchedCodecPipeline equivalent). """ - from zarr.abc.codec import BytesBytesCodec - from zarr.core.array_spec import ArrayConfig, ArraySpec - 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 class _PrototypeRecordingCodec(BytesBytesCodec): # type: ignore[misc,unused-ignore] """A no-op BB codec that records the prototype it was called with.""" @@ -319,16 +293,8 @@ async def _decode_single(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> Bu recording = _PrototypeRecordingCodec() transform = ChunkTransform(codecs=(BytesCodec(), recording)) - zdtype = get_data_type_from_native_dtype(np.dtype("float64")) - def _spec(prototype: BufferPrototype) -> ArraySpec: - return ArraySpec( - shape=(10,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0.0), - config=ArrayConfig(order="C", write_empty_chunks=False), - prototype=prototype, - ) + return _make_spec((10,), write_empty_chunks=False, prototype=prototype) proto_default = default_buffer_prototype() # A distinct BufferPrototype instance with the same buffer/nd_buffer types -- @@ -371,8 +337,6 @@ def test_read_write_with_thread_pool() -> None: returning 1) would silently degrade all the pool tests into re-testing the sequential branch while staying green. """ - from unittest.mock import patch - import zarr.core.codec_pipeline as cp_mod with zarr_config.set(_FUSED_POOL_CONFIG): @@ -404,8 +368,6 @@ def test_read_write_with_thread_pool() -> None: def test_thread_pool_write_worker_exception_propagates() -> None: """A store error raised inside a pool worker during write_sync surfaces to the caller (write_sync consumes pool.map, so worker exceptions re-raise).""" - from unittest.mock import patch - with zarr_config.set(_FUSED_POOL_CONFIG): store = MemoryStore() arr = zarr.create_array( @@ -426,8 +388,6 @@ def test_thread_pool_write_worker_exception_propagates() -> None: def test_thread_pool_read_worker_exception_propagates() -> None: """A store error raised inside a pool worker during read_sync surfaces to the caller (read_sync consumes pool.map into a tuple).""" - from unittest.mock import patch - with zarr_config.set(_FUSED_POOL_CONFIG): store = MemoryStore() arr = zarr.create_array( @@ -474,12 +434,7 @@ async def test_encode_and_write_as_completed_cancels_stray_writes_on_failure() - never retrieved (an unraisable "Task exception was never retrieved" warning if it later fails). """ - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.chunk_utils import ChunkTransform from zarr.core.codec_pipeline import _encode_and_write_as_completed - from zarr.core.dtype import get_data_type_from_native_dtype write_started = asyncio.Event() write_finished = False @@ -517,14 +472,7 @@ async def delete(self) -> None: async def set_if_not_exists(self, default: Buffer) -> None: pass - zdtype = get_data_type_from_native_dtype(np.dtype("uint8")) - chunk_spec = ArraySpec( - shape=(1,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + chunk_spec = _make_spec((1,), UInt8()) chunk_array = CPUNDBuffer.from_numpy_array(np.zeros(1, dtype="uint8")) transform = ChunkTransform(codecs=(BytesCodec(),)) @@ -598,28 +546,12 @@ def test_shared_transform_decode_alternating_specs() -> None: underpins that guarantee. (The concurrent counterpart is `test_concurrent_reads_shared_transform_with_pool`.) """ - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.chunk_utils import ChunkTransform - from zarr.core.dtype import get_data_type_from_native_dtype - - def _spec(shape: tuple[int, ...]) -> ArraySpec: - zdtype = get_data_type_from_native_dtype(np.dtype("int32")) - return ArraySpec( - shape=shape, - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) - transform = ChunkTransform(codecs=(TransposeCodec(order=(1, 0)), BytesCodec())) # two distinct specs (different shapes) sharing the one transform + cache slot cases = [] for shape in [(5, 7), (3, 11)]: - spec = _spec(shape) + spec = _make_spec(shape, Int32()) arr = np.arange(int(np.prod(shape)), dtype="int32").reshape(shape) encoded = transform.encode_chunk(CPUNDBuffer.from_numpy_array(arr), spec) assert encoded is not None @@ -643,11 +575,6 @@ def test_sharded_fallback_inner_chunks_avoid_async_transform() -> None: dict lookup plus an async per-chunk transform — measured at 1.5x (raw) to 3.6x (gzip) of sharded fallback read time. """ - from unittest.mock import patch - - from zarr.core.codec_pipeline import AsyncChunkTransform - from zarr.testing.store import LatencyStore - calls = {"decode": 0, "encode": 0} orig_decode = AsyncChunkTransform.decode_chunk orig_encode = AsyncChunkTransform.encode_chunk @@ -716,22 +643,9 @@ def test_write_over_sync_byte_setter_takes_sync_path() -> None: directly: without it, this write degrades to the async fallback (one coroutine per inner chunk for an in-memory dict store). """ - import asyncio - from unittest.mock import patch - from zarr.codecs.sharding import _ShardingByteSetter - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.dtype import get_data_type_from_native_dtype - zdtype = get_data_type_from_native_dtype(np.dtype("uint8")) - spec = ArraySpec( - shape=(10,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((10,), UInt8()) pipeline = FusedCodecPipeline.from_codecs([BytesCodec()]).evolve_from_array_spec(spec) assert pipeline.sync_transform is not None @@ -839,8 +753,6 @@ def test_sharded_roundtrip_with_async_only_inner_codec() -> None: # The stored bytes are valid for the default pipeline too: read them back # under BatchedCodecPipeline (default codec_pipeline.path). Opening from # metadata needs the codec name in the registry. - from zarr.registry import register_codec - register_codec("test-async-only-noop", _AsyncOnlyNoopCodec) reread = zarr.open_array(store=store, mode="r") np.testing.assert_array_equal(reread[:], data) @@ -849,51 +761,65 @@ def test_sharded_roundtrip_with_async_only_inner_codec() -> None: # --------------------------------------------------------------------------- # AsyncChunkTransform: the async per-chunk codec chain used on the async # fallback path. It is the async mirror of ChunkTransform, so it must produce -# identical bytes/arrays. The default (Fused, sync-store) path never uses it; +# identical bytes/arrays — except for gzip's embedded 4-byte MTIME header +# field, which the gzip case's comparator deliberately ignores. +# The default (Fused, sync-store) path never uses it; # these tests drive it directly over multi-codec chains so the aa/bb loops and # the all-fill drop branch are exercised. # --------------------------------------------------------------------------- -@pytest.mark.parametrize( - "codecs", - [ - (BytesCodec(),), - (BytesCodec(), GzipCodec(level=1)), - (TransposeCodec(order=(1, 0)), BytesCodec()), - (TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec(level=1)), - ], - ids=["bytes-only", "bb", "aa", "aa+ab+bb"], -) -def test_async_chunk_transform_matches_sync(codecs: tuple[Any, ...]) -> None: +def _bytes_identical(a: bytes, b: bytes) -> bool: + return a == b + + +# Each case pairs a codec chain with the expected relationship between the +# async and sync encoders' output: byte-identical, except for a chain +# containing gzip, whose embedded 4-byte MTIME header field varies with the +# encode time. +_ASYNC_SYNC_PARITY_CASES: list[Expect[tuple[Any, ...], Callable[[bytes, bytes], bool]]] = [ + Expect( + input=(BytesCodec(),), + output=_bytes_identical, + id="bytes-only", + ), + Expect( + input=(BytesCodec(), GzipCodec(level=1)), + output=_gzip_streams_equal_except_mtime, + id="bb", + ), + Expect( + input=(TransposeCodec(order=(1, 0)), BytesCodec()), + output=_bytes_identical, + id="aa", + ), + Expect( + input=(TransposeCodec(order=(1, 0)), BytesCodec(), ZstdCodec(level=1)), + output=_bytes_identical, + id="aa+ab+bb", + ), +] + + +@pytest.mark.parametrize("case", _ASYNC_SYNC_PARITY_CASES, ids=lambda c: c.id) +def test_async_chunk_transform_matches_sync( + case: Expect[tuple[Any, ...], Callable[[bytes, bytes], bool]], +) -> None: """`AsyncChunkTransform.decode_chunk`/`encode_chunk` must round-trip and - produce exactly what the synchronous `ChunkTransform` produces, across - array->array, array->bytes, and bytes->bytes codec combinations. + produce what the synchronous `ChunkTransform` produces, across + array->array, array->bytes, and bytes->bytes codec combinations. Each + case's expected output is a byte-comparison function: exact equality, + except for gzip, whose embedded 4-byte MTIME header field is deliberately + ignored. This is the async mirror of the codecs the default pipeline runs synchronously; a divergence here corrupts data only on the async fallback path (remote stores), which no end-to-end test of the default pipeline touches. """ - import asyncio - - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.chunk_utils import ChunkTransform, evolve_codecs - from zarr.core.codec_pipeline import AsyncChunkTransform - from zarr.core.dtype import get_data_type_from_native_dtype - shape = (4, 4) - zdtype = get_data_type_from_native_dtype(np.dtype("int32")) - spec = ArraySpec( - shape=shape, - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) - evolved = evolve_codecs(codecs, spec) + spec = _make_spec(shape, Int32()) + evolved = evolve_codecs(case.input, spec) sync_t = ChunkTransform(codecs=evolved) async_t = AsyncChunkTransform(codecs=evolved) @@ -904,19 +830,7 @@ def test_async_chunk_transform_matches_sync(codecs: tuple[Any, ...]) -> None: async_bytes = asyncio.run(async_t.encode_chunk(value, spec)) assert sync_bytes is not None assert async_bytes is not None - - has_timestamp_codec = any(isinstance(c, GzipCodec) for c in evolved) - - if has_timestamp_codec: - assert _gzip_streams_equal_except_mtime( - async_bytes.to_bytes(), - sync_bytes.to_bytes(), - ) - else: - np.testing.assert_array_equal( - async_bytes.to_bytes(), - sync_bytes.to_bytes(), - ) + assert case.output(async_bytes.to_bytes(), sync_bytes.to_bytes()) sync_arr = sync_t.decode_chunk(async_bytes, spec) async_arr = asyncio.run(async_t.decode_chunk(async_bytes, spec)) @@ -928,21 +842,7 @@ def test_async_decode_encode_passes_through_none_chunks() -> None: """`FusedCodecPipeline.decode`/`encode` (the async batch entry points used on the fallback path) map a None chunk to None and leave real chunks untouched — pins the None-passthrough branch the default sync path skips.""" - import asyncio - - from zarr.core.array_spec import ArrayConfig, ArraySpec - from zarr.core.buffer import default_buffer_prototype - from zarr.core.buffer.cpu import NDBuffer as CPUNDBuffer - from zarr.core.dtype import get_data_type_from_native_dtype - - zdtype = get_data_type_from_native_dtype(np.dtype("int32")) - spec = ArraySpec( - shape=(4,), - dtype=zdtype, - fill_value=zdtype.cast_scalar(0), - config=ArrayConfig(order="C", write_empty_chunks=True), - prototype=default_buffer_prototype(), - ) + spec = _make_spec((4,), Int32()) pipeline = FusedCodecPipeline.from_codecs([BytesCodec()]).evolve_from_array_spec(spec) data = np.arange(4, dtype="int32") @@ -1199,8 +1099,6 @@ def test_sync_io_capability_gates_fused_paths( (read-modify-write) writes, and all-fill (delete) writes — a store with a partial sync surface must get a clean async fallback, never a mid-batch error.""" - from unittest.mock import patch - store = store_factory() assert _store_supports_sync_io(store) is expect_sync_path diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 0000000000..767d3e0439 --- /dev/null +++ b/tests/test_registry.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +import json + +import numcodecs.registry +import pytest + +import zarr +import zarr.registry +from zarr.core.buffer import default_buffer_prototype +from zarr.core.config import BadConfigError, config +from zarr.core.metadata.v3 import parse_codecs +from zarr.errors import UnknownCodecError +from zarr.registry import ( + _CODEC_PACKAGE_PREFIXES, + _CODEC_PACKAGES, + _NUMCODEC_PACKAGE_PREFIXES, + _NUMCODECS_CODEC_DOCS_URL, + _ZARR_CODEC_DOCS_URL, + _missing_codec_message, + _packages_for_codec, + get_codec_class, + get_numcodec, +) +from zarr.storage import MemoryStore + + +@pytest.fixture +def unregistered_v3_codec(monkeypatch: pytest.MonkeyPatch) -> str: + """Guarantee a Zarr format 3 codec name resolves to nothing. + + The registry is entry-point driven, so a name the hint table advertises resolves for real + in any environment where the advertised package happens to be installed. + """ + name = "n5_default" + monkeypatch.setitem(zarr.registry._codec_registries, name, zarr.registry.Registry()) + return name + + +@pytest.fixture +def unregistered_v2_codec(monkeypatch: pytest.MonkeyPatch) -> str: + """As above, for a numcodecs codec id.""" + codec_id = "wavpack" + monkeypatch.delitem(numcodecs.registry.codec_registry, codec_id, raising=False) + monkeypatch.delitem(numcodecs.registry.entries, codec_id, raising=False) + return codec_id + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("n5_default", ("zarr-n5",)), + ("gribberish", ("gribberish",)), + ("imagecodecs_jpeg2k", ("virtual-tiff",)), + ("omfiles.pfor", ("omfiles",)), + ("any-numcodecs.array-array", ("zarr-any-numcodecs",)), + ("totally-made-up", ()), + ], +) +def test_packages_for_codec_v3(name: str, expected: tuple[str, ...]) -> None: + """Exact names and prefixes both resolve; unknown names resolve to nothing.""" + assert _packages_for_codec(name, zarr_format=3) == expected + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("wavpack", ("wavpack-numcodecs",)), + ("grib", ("kerchunk",)), + ("zfpy", ("numcodecs[zfpy]",)), + ("crc32c", ("numcodecs[crc32c]",)), + ("gribscan.rawgrib", ("gribscan",)), + ("imagecodecs_jpeg2k", ("imagecodecs-numcodecs",)), + ("totally-made-up", ()), + ], +) +def test_packages_for_numcodec_v2(name: str, expected: tuple[str, ...]) -> None: + """Zarr format 2 codec ids resolve against the numcodecs table.""" + assert _packages_for_codec(name, zarr_format=2) == expected + + +def test_packages_for_codec_is_format_specific() -> None: + """The same name can mean different packages in each format's registry.""" + assert _packages_for_codec("imagecodecs_jpeg2k", zarr_format=3) == ("virtual-tiff",) + assert _packages_for_codec("imagecodecs_jpeg2k", zarr_format=2) == ("imagecodecs-numcodecs",) + # `crc32c` is a codec zarr implements in format 3, so only format 2 gets a hint for it. + assert _packages_for_codec("crc32c", zarr_format=3) == () + + +def test_missing_codec_message_with_known_packages() -> None: + """The message names the codec, the docs page, and every known package.""" + msg = _missing_codec_message("n5_default", zarr_format=3) + assert msg == ( + "An implementation for codec 'n5_default' is not available. Register one explicitly " + f"using the codec registry (see {_ZARR_CODEC_DOCS_URL}), or install a Python package " + "that registers a codec implementation with zarr. Known packages supporting this " + "codec: zarr-n5." + ) + + +def test_missing_codec_message_without_known_packages() -> None: + """With no known package we still explain how to register one by hand.""" + msg = _missing_codec_message("totally-made-up", zarr_format=3) + assert msg == ( + "An implementation for codec 'totally-made-up' is not available. Register one " + f"explicitly using the codec registry (see {_ZARR_CODEC_DOCS_URL}), or install a " + "Python package that registers a codec implementation with zarr." + ) + assert "Known packages" not in msg + + +def test_missing_codec_message_for_zarr_format_2() -> None: + """Format 2 codecs live in the numcodecs registry, so the message must say so.""" + msg = _missing_codec_message("wavpack", zarr_format=2) + assert msg == ( + "An implementation for codec 'wavpack' is not available. Register one explicitly " + f"using the codec registry (see {_NUMCODECS_CODEC_DOCS_URL}), or install a Python " + "package that registers a codec implementation with numcodecs. Known packages " + "supporting this codec: wavpack-numcodecs." + ) + + +def test_no_prefix_shadows_another_prefix() -> None: + """First-match prefix lookup is only deterministic while no prefix contains another.""" + for table in (_CODEC_PACKAGE_PREFIXES, _NUMCODEC_PACKAGE_PREFIXES): + for a in table: + for b in table: + assert a == b or not a.startswith(b) + + +def test_mapping_does_not_shadow_builtin_codecs() -> None: + """A codec zarr implements itself must never appear in the format 3 hint table.""" + import zarr.codecs # noqa: F401 (importing registers the built-in codecs) + from zarr.registry import _codec_registries + + # Select on the implementing class's module rather than on "is this registry non-empty": + # a third-party entry-point codec lazy-loaded by an earlier test would otherwise show up + # here and fail the assertion even though it shadows nothing. + implemented = { + name + for name, reg in _codec_registries.items() + if any(cls.__module__.startswith("zarr.") for cls in reg.values()) + } + assert implemented, "expected importing zarr.codecs to populate the registry" + assert not (implemented & set(_CODEC_PACKAGES)) + for prefix in _CODEC_PACKAGE_PREFIXES: + assert not any(name.startswith(prefix) for name in implemented) + + +def test_get_codec_class_unknown_raises_with_package_hint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An unregistered codec with a known package names that package in the error.""" + from collections import defaultdict + + monkeypatch.setattr(zarr.registry, "_codec_registries", defaultdict(zarr.registry.Registry)) + with pytest.raises(UnknownCodecError, match="Known packages supporting this codec: zarr-n5"): + get_codec_class("n5_default") + + +def test_get_codec_class_unknown_raises_without_package_hint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An unregistered codec we know nothing about still explains manual registration.""" + from collections import defaultdict + + monkeypatch.setattr(zarr.registry, "_codec_registries", defaultdict(zarr.registry.Registry)) + with pytest.raises(UnknownCodecError, match="An implementation for codec 'nope' is not"): + get_codec_class("nope") + + +def test_unknown_codec_error_is_exported() -> None: + """UnknownCodecError is public API now that we tell users to catch it.""" + import zarr.errors + + assert "UnknownCodecError" in zarr.errors.__all__ + + +def test_get_numcodec_unknown_raises_with_package_hint(unregistered_v2_codec: str) -> None: + """A Zarr format 2 codec id we don't have names the package that provides it.""" + with pytest.raises( + UnknownCodecError, match="Known packages supporting this codec: wavpack-numcodecs" + ): + get_numcodec({"id": unregistered_v2_codec}) + + +def test_get_numcodec_unknown_points_at_numcodecs_registry() -> None: + """The Zarr format 2 message links the numcodecs registry, not zarr's extending guide.""" + with pytest.raises(UnknownCodecError) as excinfo: + get_numcodec({"id": "definitely-not-a-real-codec"}) + assert _NUMCODECS_CODEC_DOCS_URL in str(excinfo.value) + assert _ZARR_CODEC_DOCS_URL not in str(excinfo.value) + + +def test_get_numcodec_known_codec_still_works() -> None: + """The happy path is untouched.""" + from numcodecs import GZip + + assert get_numcodec({"id": "gzip", "level": 2}) == GZip(level=2) # type: ignore[typeddict-unknown-key] + + +def test_get_numcodec_without_an_id_keeps_the_numcodecs_error() -> None: + """With no codec id there is nothing to look up, so numcodecs' own error stands.""" + # Not asserting on numcodecs' exception class: it only gained a dedicated + # `numcodecs.errors.UnknownCodecError` in 0.15.1, and zarr supports numcodecs >= 0.14. + with pytest.raises(ValueError) as excinfo: + get_numcodec({"level": 2}) # type: ignore[typeddict-item,typeddict-unknown-key] + assert not isinstance(excinfo.value, UnknownCodecError) + + +def test_get_numcodec_does_not_relabel_a_bad_configuration() -> None: + """A registered codec rejecting its config is not a missing codec. + + Telling the user to install a package they already have would be the same misleading + error this module exists to remove, pointing the other way. + """ + with pytest.raises(ValueError) as excinfo: + get_numcodec({"id": "bitround", "keepbits": -1}) # type: ignore[typeddict-unknown-key] + assert not isinstance(excinfo.value, UnknownCodecError) + assert str(excinfo.value) == "keepbits must be zero or positive" + + +@pytest.mark.parametrize("via", ["zarr", "numcodecs"]) +def test_get_numcodec_does_not_relabel_a_missing_inner_codec(via: str) -> None: + """A wrapper codec failing on an inner codec must not be blamed on the outer id. + + The outer id is registered, so relabelling it would advertise a package the user has + already installed. Covers both routes a real wrapper takes to resolve its inner codec: + zarr's ``get_numcodec`` and numcodecs' own ``get_codec``. + """ + import numcodecs.registry + from numcodecs.abc import Codec + + resolve = get_numcodec if via == "zarr" else numcodecs.registry.get_codec + + class WrapperCodec(Codec): # type: ignore[misc] + codec_id = "test_wrapper" + + @classmethod + def from_config(cls, config: dict[str, object]) -> WrapperCodec: + resolve({"id": "some-missing-inner-codec"}) + return cls() + + def encode(self, buf: object) -> object: + return buf + + def decode(self, buf: object, out: object = None) -> object: + return buf + + numcodecs.registry.register_codec(WrapperCodec, codec_id="test_wrapper") + try: + with pytest.raises(ValueError) as excinfo: + get_numcodec({"id": "test_wrapper"}) + assert "some-missing-inner-codec" in str(excinfo.value) + assert "test_wrapper" not in str(excinfo.value) + finally: + numcodecs.registry.codec_registry.pop("test_wrapper", None) + + +@pytest.mark.parametrize("data", ["abc", ["ab"]]) +def test_get_numcodec_non_mapping_input_still_raises_value_error(data: object) -> None: + """Input that is not a mapping must reach numcodecs rather than raising AttributeError. + + Reading ``id`` off the input unconditionally turned numcodecs' ValueError into an + AttributeError, which is not a ValueError and so escaped handlers that caught it. Both + params end in a ValueError from numcodecs, by different routes: ``"abc"`` fails + ``dict(config)`` coercion, while ``["ab"]`` coerces to ``{"a": "b"}`` and then has no id. + """ + with pytest.raises(ValueError): + get_numcodec(data) # type: ignore[arg-type] + + +def test_imagecodecs_prefix_does_not_over_match_in_zarr_format_3() -> None: + """virtual-tiff provides 15 of the 81 `imagecodecs_*` names; the rest must get no hint. + + Recommending virtual-tiff for a name it does not provide is worse than saying nothing. + """ + assert _packages_for_codec("imagecodecs_jpeg2k", zarr_format=3) == ("virtual-tiff",) + for name in ("imagecodecs_jpegls", "imagecodecs_avif", "imagecodecs_blosc"): + assert _packages_for_codec(name, zarr_format=3) == () + assert _packages_for_codec(name, zarr_format=2) == ("imagecodecs-numcodecs",) + + +def test_resolve_codec_reports_missing_codec() -> None: + """The other public entry point into `get_codec_class` gets the message too.""" + import zarr + + with pytest.raises(UnknownCodecError, match="Known packages supporting this codec: zarr-n5"): + zarr.create_array({}, shape=(4,), dtype="uint8", compressors=[{"name": "n5_default"}]) + + +async def test_open_array_with_missing_v3_codec_reports_package( + unregistered_v3_codec: str, +) -> None: + """Opening a Zarr format 3 array naming a codec we lack points at the package.""" + store = MemoryStore() + metadata = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "data_type": "float64", + "fill_value": 0.0, + "codecs": [{"name": "bytes"}, {"name": unregistered_v3_codec}], + "attributes": {}, + } + await store.set( + "zarr.json", + default_buffer_prototype().buffer.from_bytes(json.dumps(metadata).encode()), + ) + with pytest.raises(UnknownCodecError, match="Known packages supporting this codec: zarr-n5"): + zarr.open_array(store=store, mode="r") + + +async def test_open_array_with_missing_v2_codec_reports_package( + unregistered_v2_codec: str, +) -> None: + """Same, for a Zarr format 2 array whose compressor id we lack.""" + store = MemoryStore() + metadata = { + "zarr_format": 2, + "shape": [4], + "chunks": [4], + "dtype": " None: + """A config pinning an implementation that isn't registered is a config error.""" + with config.set({"codecs.bytes": "some.package.RemovedBytesCodec"}): + with pytest.raises(BadConfigError, match="RemovedBytesCodec"): + get_codec_class("bytes") + + +def test_parse_codecs_with_unregistered_config_pin_raises_bad_config_error() -> None: + """The same, through the array-open path, and never as a bare KeyError.""" + with config.set({"codecs.bytes": "some.package.RemovedBytesCodec"}): + with pytest.raises(BadConfigError, match="RemovedBytesCodec"): + parse_codecs([{"name": "bytes"}]) + + +def test_parse_codecs_converts_keyerror_from_from_dict(monkeypatch: pytest.MonkeyPatch) -> None: + """A codec whose from_dict indexes a malformed config must not leak a KeyError. + + A bare KeyError out of metadata parsing is caught by the array-then-group fallback in + `zarr.api.asynchronous.open`, which then reports an unrelated group error. + """ + from zarr.codecs import BytesCodec + from zarr.errors import MetadataValidationError + from zarr.registry import register_codec + + class PickyCodec(BytesCodec): + @classmethod + def from_dict(cls, data: object) -> PickyCodec: + data["configuration"]["required_option"] # type: ignore[index] + return cls() + + monkeypatch.setitem(zarr.registry._codec_registries, "test_picky", zarr.registry.Registry()) + register_codec("test_picky", PickyCodec) + with pytest.raises(MetadataValidationError, match="test_picky.*required_option"): + parse_codecs([{"name": "test_picky", "configuration": {}}])