From 8c6b986518b097e2df22eceae24f29d8688f5361 Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 21:34:09 +0000 Subject: [PATCH 1/7] Add the offset chunk key encoding and Array.prepend The `offset` encoding stores each chunk under the `default` key of its grid coordinate plus a fixed per-dimension integer offset, so that an array grows in the negative direction of an axis by whole chunks as a metadata-only change: `shape` grows, `offset` drops, and every existing chunk keeps its key. `Array.prepend_chunks` performs that edit and `Array.prepend` grows and writes in one call; both refuse an array on another encoding rather than rename its chunks. The offset must have one entry per dimension, checked when the metadata is built; the `ChunkKeyEncodingParams` shortcut carries `offset` through to the configuration instead of dropping it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01CSaEKF9nB55gU7PQgQ22X5 --- changes/+offset-chunk-key-encoding.feature.md | 3 + docs/user-guide/arrays.md | 24 ++ src/zarr/core/array.py | 227 ++++++++++++++++++ src/zarr/core/chunk_key_encodings.py | 75 +++++- src/zarr/core/metadata/v3.py | 8 + tests/test_array.py | 133 ++++++++++ tests/test_chunk_key_encodings.py | 82 ++++++- 7 files changed, 548 insertions(+), 4 deletions(-) create mode 100644 changes/+offset-chunk-key-encoding.feature.md diff --git a/changes/+offset-chunk-key-encoding.feature.md b/changes/+offset-chunk-key-encoding.feature.md new file mode 100644 index 0000000000..4d2587e33c --- /dev/null +++ b/changes/+offset-chunk-key-encoding.feature.md @@ -0,0 +1,3 @@ +Added the ``offset`` chunk key encoding (the ``default`` keys of ``chunk_coords + offset``) and +``Array.prepend_chunks`` / ``Array.prepend``, which grow an array before index 0 by whole chunks +as a metadata-only change: every existing chunk keeps its key. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index e8df79a102..06b64ae3e1 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -154,6 +154,30 @@ z.append(np.vstack([a, a]), axis=1) print(f"Shape after second append: {z.shape}") ``` +### Growing before index 0 + +`resize` and `append` only move the high edge of a dimension. Adding data *before* index +0 would shift every array index and, under the `default` chunk key encoding, change the key +of every chunk. The `offset` chunk key encoding stores each chunk under the `default` key of +its grid coordinate plus a fixed per-dimension offset, so that +[`zarr.Array.prepend_chunks`][] can grow a dimension by whole chunks as a metadata-only +change: the new leading chunks take new (negative) keys and every existing chunk keeps its +key. [`zarr.Array.prepend`][] grows and writes in one call. An array opts in at creation: + +```python exec="true" session="arrays" source="above" result="ansi" +z = zarr.create_array(store='data/example-5.zarr', shape=(6, 4), dtype='int32', chunks=(3, 4), + chunk_key_encoding={'name': 'offset', 'offset': (0, 0)}, overwrite=True) +z[:] = 1 +print(f"Original shape: {z.shape}, offset: {z.metadata.chunk_key_encoding.offset}") +z.prepend(np.full((6, 4), 2, dtype='int32')) +print(f"Shape after prepend: {z.shape}, offset: {z.metadata.chunk_key_encoding.offset}") +print(f"Chunks written: {z.nchunks_initialized}") +``` + +With an all-zero offset the keys are those of the `default` encoding, so an existing array +can adopt the encoding by rewriting its metadata, without moving a chunk. Readers that do not +implement the encoding refuse the array rather than misplace its chunks. + ## Runtime configuration Zarr arrays are parametrized with a configuration that determines certain aspects of array behavior. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 06f7338755..c476a4cc00 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -54,6 +54,7 @@ ChunkKeyEncoding, ChunkKeyEncodingLike, DefaultChunkKeyEncoding, + OffsetChunkKeyEncoding, V2ChunkKeyEncoding, parse_chunk_key_encoding, ) @@ -1801,6 +1802,57 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """ return await _append(self, data, axis) + async def prepend_chunks(self, n_chunks: int, axis: int = 0) -> tuple[int, ...]: + """Grow the array by `n_chunks` whole chunks before index 0 of `axis`. + + A metadata-only operation: the array must use the ``offset`` chunk key encoding, + under which the new leading chunks take new keys and every existing chunk keeps + its key, so no chunk is read, moved or rewritten. The new region reads as the fill + value until written. + + Parameters + ---------- + n_chunks : int + Number of whole chunks to add before index 0. + axis : int + Axis along which to grow. + + Returns + ------- + new_shape : tuple + + Raises + ------ + TypeError + If the array does not use the ``offset`` chunk key encoding, or its chunk grid + is not regular. + """ + return await _prepend_chunks(self, n_chunks, axis) + + async def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: + """Prepend `data` to `axis`: grow the array before index 0 and write `data` there. + + The array must use the ``offset`` chunk key encoding, and the extent of `data` + along `axis` must be a whole number of chunks; see `prepend_chunks`. + + Parameters + ---------- + data : array-like + Data to be prepended. + axis : int + Axis along which to prepend. + + Returns + ------- + new_shape : tuple + + Notes + ----- + The size of all dimensions other than `axis` must match between this + array and `data`. + """ + return await _prepend(self, data, axis) + async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ Asynchronously update the array's attributes. @@ -4025,6 +4077,64 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """ return sync(self.async_array.append(data, axis=axis)) + def prepend_chunks(self, n_chunks: int, axis: int = 0) -> tuple[int, ...]: + """Grow the array by `n_chunks` whole chunks before index 0 of `axis`. + + A metadata-only operation: the array must use the ``offset`` chunk key encoding, + under which the new leading chunks take new keys and every existing chunk keeps + its key, so no chunk is read, moved or rewritten. The new region reads as the fill + value until written. + + Parameters + ---------- + n_chunks : int + Number of whole chunks to add before index 0. + axis : int + Axis along which to grow. + + Returns + ------- + new_shape : tuple + + Examples + -------- + ```python + import zarr + z = zarr.create_array(store={}, shape=(6, 4), chunks=(3, 4), dtype="int32", + chunk_key_encoding={"name": "offset", "offset": (0, 0)}) + z[:] = 1 + z.prepend_chunks(2, axis=0) + z.shape + #> (12, 4) + z[:6] = 0 + ``` + """ + return sync(self.async_array.prepend_chunks(n_chunks, axis=axis)) + + def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: + """Prepend `data` to `axis`: grow the array before index 0 and write `data` there. + + The array must use the ``offset`` chunk key encoding, and the extent of `data` + along `axis` must be a whole number of chunks; see `prepend_chunks`. + + Parameters + ---------- + data : array-like + Data to be prepended. + axis : int + Axis along which to prepend. + + Returns + ------- + new_shape : tuple + + Notes + ----- + The size of all dimensions other than `axis` must match between this + array and `data`. + """ + return sync(self.async_array.prepend(data, axis=axis)) + def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ Update the array's attributes. @@ -6066,6 +6176,123 @@ async def _append( return new_shape +async def _prepend_chunks( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + n_chunks: int, + axis: int = 0, +) -> tuple[int, ...]: + """ + Grow an array by whole chunks before index 0 of `axis`, as a metadata-only change. + + Parameters + ---------- + array : AsyncArray + The array to grow. Must use the ``offset`` chunk key encoding and a regular chunk grid. + n_chunks : int + Number of whole chunks to add. + axis : int + Axis along which to grow. + + Returns + ------- + new_shape : tuple[int, ...] + The new shape of the array. + """ + metadata = array.metadata + if not isinstance(metadata, ArrayV3Metadata) or not isinstance( + metadata.chunk_key_encoding, OffsetChunkKeyEncoding + ): + raise TypeError( + "Growing an array before index 0 without moving its chunks needs the offset chunk " + "key encoding (chunk_key_encoding={'name': 'offset', 'offset': (0,) * ndim} at " + "creation); this array uses " + f"{metadata.chunk_key_encoding.to_dict() if isinstance(metadata, ArrayV3Metadata) else 'Zarr v2 keys'}." + ) + if not isinstance(metadata.chunk_grid, RegularChunkGridMetadata): + raise TypeError("prepend_chunks needs a regular chunk grid.") + ndim = len(metadata.shape) + if not -ndim <= axis < ndim: + raise ValueError(f"axis {axis} is out of range for a {ndim}-dimensional array.") + axis %= ndim + if n_chunks < 0: + raise ValueError(f"n_chunks must be non-negative, got {n_chunks}.") + if n_chunks == 0: + return metadata.shape + + chunk = metadata.chunk_grid.chunk_shape[axis] + new_shape = tuple( + s + n_chunks * chunk if i == axis else s for i, s in enumerate(metadata.shape) + ) + new_offset = tuple( + o - n_chunks if i == axis else o for i, o in enumerate(metadata.chunk_key_encoding.offset) + ) + new_metadata = replace( + metadata.update_shape(new_shape), + chunk_key_encoding=replace(metadata.chunk_key_encoding, offset=new_offset), + ) + new_chunk_grid = ChunkGrid.from_metadata(new_metadata) + + # Both edits land in one document write, so a reader never sees one without the other. + await store_documents(array.store_path, encode_documents(array.store_path, new_metadata)) + array._stored_document_replaced() + object.__setattr__(array, "metadata", new_metadata) + object.__setattr__(array, "_chunk_grid", new_chunk_grid) + return new_shape + + +async def _prepend( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + data: npt.ArrayLike, + axis: int = 0, +) -> tuple[int, ...]: + """ + Prepend data to an array along the specified axis. + + Parameters + ---------- + array : AsyncArray + The array to prepend to. Must use the ``offset`` chunk key encoding. + data : npt.ArrayLike + Data to be prepended. Its extent along `axis` must be a whole number of chunks. + axis : int + Axis along which to prepend. + + Returns + ------- + new_shape : tuple[int, ...] + The new shape of the array after prepending. + """ + if not hasattr(data, "shape"): + data = np.asanyarray(data) + + self_shape_preserved = tuple(s for i, s in enumerate(array.shape) if i != axis) + data_shape_preserved = tuple(s for i, s in enumerate(data.shape) if i != axis) + if self_shape_preserved != data_shape_preserved: + raise ValueError( + f"shape of data to prepend is not compatible with the array. " + f"The shape of the data is ({data_shape_preserved})" + f"and the shape of the array is ({self_shape_preserved})." + "All dimensions must match except for the dimension being " + "prepended." + ) + extent = data.shape[axis] + chunk = array.chunks[axis] + if extent % chunk: + raise ValueError( + f"prepend grows the array by whole chunks, but the data spans {extent} along axis " + f"{axis} and the chunk length there is {chunk}." + ) + + new_shape = await _prepend_chunks(array, extent // chunk, axis) + + prepend_selection = tuple( + slice(None) if i != axis else slice(0, extent) for i in range(len(new_shape)) + ) + await array.setitem(prepend_selection, data) + + return new_shape + + async def _update_attributes( array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], new_attributes: dict[str, JSON], diff --git a/src/zarr/core/chunk_key_encodings.py b/src/zarr/core/chunk_key_encodings.py index d871e279d2..a113b46ab7 100644 --- a/src/zarr/core/chunk_key_encodings.py +++ b/src/zarr/core/chunk_key_encodings.py @@ -1,6 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Iterable from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast @@ -24,8 +25,9 @@ def parse_separator(data: JSON) -> SeparatorLiteral: class ChunkKeyEncodingParams(TypedDict): - name: Literal["v2", "default"] + name: Literal["v2", "default", "offset"] separator: NotRequired[SeparatorLiteral] + offset: NotRequired[tuple[int, ...]] @dataclass(frozen=True) @@ -105,6 +107,70 @@ def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: return "0" if chunk_identifier == "" else chunk_identifier +def parse_offset(data: object) -> tuple[int, ...]: + if isinstance(data, (str, bytes)) or not isinstance(data, Iterable): + raise TypeError(f"offset must be a sequence of integers, got {data!r}") + offset = tuple(data) + if any(isinstance(o, bool) or not isinstance(o, int) for o in offset): + raise TypeError(f"offset must be a sequence of integers, got {data!r}") + return offset + + +@dataclass(frozen=True) +class OffsetChunkKeyEncoding(ChunkKeyEncoding): + """ + The ``default`` encoding applied to ``chunk_coords + offset``. + + ``offset`` holds one integer per array dimension and may be negative. Array indexing is + unchanged and zero-based; only the mapping from chunk coordinates to store keys moves, so + growing the array by ``k`` whole chunks before index 0 along an axis is two metadata edits + (``shape += k * chunk``, ``offset -= k``) under which every existing chunk keeps its key. + With an all-zero ``offset`` the keys are those of ``default`` with the same separator. + + See https://github.com/zarr-developers/zarr-extensions/tree/main/chunk-key-encodings/offset + """ + + name: ClassVar[Literal["offset"]] = "offset" + offset: tuple[int, ...] + separator: SeparatorLiteral = "/" + + def __post_init__(self) -> None: + object.__setattr__(self, "offset", parse_offset(self.offset)) + object.__setattr__(self, "separator", parse_separator(self.separator)) + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + _, config_parsed = parse_named_configuration(data, require_configuration=False) + if not config_parsed or "offset" not in config_parsed: + raise ValueError( + "The offset chunk key encoding requires an offset in its configuration." + ) + return cls(**cast("dict[str, Any]", config_parsed)) + + def to_dict(self) -> dict[str, JSON]: + return { + "name": self.name, + "configuration": {"separator": self.separator, "offset": list(self.offset)}, + } + + def _check_rank(self, chunk_coords: tuple[int, ...]) -> None: + if len(chunk_coords) != len(self.offset): + raise ValueError( + f"The offset has {len(self.offset)} dimension(s) but the chunk coordinates " + f"{chunk_coords!r} have {len(chunk_coords)}." + ) + + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: + shifted = DefaultChunkKeyEncoding(separator=self.separator).decode_chunk_key(chunk_key) + self._check_rank(shifted) + return tuple(c - o for c, o in zip(shifted, self.offset, strict=True)) + + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + self._check_rank(chunk_coords) + shifted = tuple(c + o for c, o in zip(chunk_coords, self.offset, strict=True)) + return self.separator.join(map(str, ("c",) + shifted)) + + def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding: """ Take an implicit specification of a chunk key encoding and parse it into a ChunkKeyEncoding object. @@ -113,8 +179,10 @@ def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding: return data # handle ChunkKeyEncodingParams - if "name" in data and "separator" in data: - data = {"name": data["name"], "configuration": {"separator": data["separator"]}} # type: ignore[typeddict-item] + if "name" in data and "configuration" not in data and ("separator" in data or "offset" in data): + params = cast("dict[str, JSON]", data) + configuration = {key: params[key] for key in ("separator", "offset") if key in params} + data = cast("dict[str, JSON]", {"name": params["name"], "configuration": configuration}) # Now must be a named config data = cast("dict[str, JSON]", data) @@ -130,3 +198,4 @@ def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding: register_chunk_key_encoding("default", DefaultChunkKeyEncoding) register_chunk_key_encoding("v2", V2ChunkKeyEncoding) +register_chunk_key_encoding("offset", OffsetChunkKeyEncoding) diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2d5d169e6b..8bcc0424ee 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -17,6 +17,7 @@ from zarr.core.chunk_key_encodings import ( ChunkKeyEncoding, ChunkKeyEncodingLike, + OffsetChunkKeyEncoding, parse_chunk_key_encoding, ) from zarr.core.common import ( @@ -544,6 +545,13 @@ def __init__( shape_parsed = parse_shapelike(shape) chunk_grid_parsed = parse_chunk_grid(chunk_grid) chunk_key_encoding_parsed = parse_chunk_key_encoding(chunk_key_encoding) + if isinstance(chunk_key_encoding_parsed, OffsetChunkKeyEncoding) and len( + chunk_key_encoding_parsed.offset + ) != len(shape_parsed): + raise ValueError( + f"The offset chunk key encoding has {len(chunk_key_encoding_parsed.offset)} " + f"offset(s) but the array has {len(shape_parsed)} dimension(s)." + ) dimension_names_parsed = parse_dimension_names(dimension_names) # Note: relying on a type method is numpy-specific fill_value_parsed = data_type.cast_scalar(fill_value) diff --git a/tests/test_array.py b/tests/test_array.py index 3e4835e686..c86478fdc7 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -6,6 +6,7 @@ import pickle import re import sys +from collections.abc import AsyncIterator from itertools import accumulate, starmap from typing import TYPE_CHECKING, Any, Literal from unittest import mock @@ -676,6 +677,138 @@ async def test_info_complete_async( ) +def _chunk_keys(store: Store) -> set[str]: + return {key for key in sync(_collect(store.list())) if not key.endswith("zarr.json")} + + +async def _collect(it: AsyncIterator[str]) -> list[str]: + return [key async for key in it] + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_chunks_is_a_metadata_only_edit(store: MemoryStore) -> None: + """Growing before index 0 touches zarr.json and nothing else: the chunk key set is + unchanged, the old data reads back at its new indices, and the new region reads as the + fill value until written.""" + z = zarr.create_array( + store=store, + shape=(6, 4), + chunks=(3, 2), + dtype="int32", + fill_value=-1, + chunk_key_encoding={"name": "offset", "offset": (0, 0)}, + ) + data = np.arange(24, dtype="int32").reshape(6, 4) + z[:] = data + keys_before = _chunk_keys(store) + assert keys_before == {"c/0/0", "c/0/1", "c/1/0", "c/1/1"} + + assert z.prepend_chunks(2, axis=0) == (12, 4) + + assert z.shape == (12, 4) + assert _chunk_keys(store) == keys_before + np.testing.assert_array_equal(z[6:], data) + np.testing.assert_array_equal(z[:6], -1) + assert z.metadata.to_dict()["chunk_key_encoding"] == { + "name": "offset", + "configuration": {"separator": "/", "offset": [-2, 0]}, + } + + z[:3] = 7 + assert _chunk_keys(store) == keys_before | {"c/-2/0", "c/-2/1"} + + reopened = zarr.open_array(store=store, mode="r") + assert reopened.shape == (12, 4) + np.testing.assert_array_equal(reopened[:], z[:]) + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_writes_the_data_in_front(store: MemoryStore) -> None: + z = zarr.create_array( + store=store, + shape=(4, 3), + chunks=(2, 3), + dtype="uint8", + chunk_key_encoding={"name": "offset", "offset": (0, 0)}, + ) + z[:] = 1 + front = np.full((4, 3), 9, dtype="uint8") + + assert z.prepend(front, axis=0) == (8, 3) + + np.testing.assert_array_equal(z[:4], 9) + np.testing.assert_array_equal(z[4:], 1) + assert _chunk_keys(store) == {"c/-2/0", "c/-1/0", "c/0/0", "c/1/0"} + + with pytest.raises(ValueError, match="whole chunks"): + z.prepend(np.zeros((3, 3), dtype="uint8"), axis=0) + with pytest.raises(ValueError, match="not compatible"): + z.prepend(np.zeros((2, 4), dtype="uint8"), axis=0) + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_chunks_along_a_later_axis_and_a_negative_axis(store: MemoryStore) -> None: + z = zarr.create_array( + store=store, + shape=(2, 6), + chunks=(2, 3), + dtype="int8", + chunk_key_encoding={"name": "offset", "separator": ".", "offset": (0, 0)}, + ) + z[:] = np.arange(12, dtype="int8").reshape(2, 6) + + assert z.prepend_chunks(1, axis=-1) == (2, 9) + + np.testing.assert_array_equal(z[:, 3:], np.arange(12, dtype="int8").reshape(2, 6)) + assert _chunk_keys(store) == {"c.0.0", "c.0.1"} + assert z.prepend_chunks(0, axis=1) == (2, 9) + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_prepend_chunks_refuses_an_array_without_the_offset_encoding( + store: MemoryStore, zarr_format: ZarrFormat +) -> None: + """A default-encoded array cannot grow in front without every chunk changing its key, + so the request is refused and names the encoding that allows it.""" + z = zarr.create_array( + store=store, shape=(4,), chunks=(2,), dtype="int8", zarr_format=zarr_format + ) + + with pytest.raises(TypeError, match="offset chunk key encoding"): + z.prepend_chunks(1) + with pytest.raises(TypeError, match="offset chunk key encoding"): + z.prepend(np.zeros(2, dtype="int8")) + assert z.shape == (4,) + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_chunks_rejects_a_bad_axis_or_count(store: MemoryStore) -> None: + z = zarr.create_array( + store=store, + shape=(4,), + chunks=(2,), + dtype="int8", + chunk_key_encoding={"name": "offset", "offset": (0,)}, + ) + with pytest.raises(ValueError, match="out of range"): + z.prepend_chunks(1, axis=1) + with pytest.raises(ValueError, match="non-negative"): + z.prepend_chunks(-1) + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_create_array_rejects_an_offset_of_the_wrong_rank(store: MemoryStore) -> None: + with pytest.raises(ValueError, match="2 dimension"): + zarr.create_array( + store=store, + shape=(4, 4), + chunks=(2, 2), + dtype="int8", + chunk_key_encoding={"name": "offset", "offset": (0,)}, + ) + + @pytest.mark.parametrize("store", ["memory"], indirect=True) def test_resize_wrong_ndim_raises(store: MemoryStore, zarr_format: ZarrFormat) -> None: """ diff --git a/tests/test_chunk_key_encodings.py b/tests/test_chunk_key_encodings.py index dcc93b9249..4b746aace0 100644 --- a/tests/test_chunk_key_encodings.py +++ b/tests/test_chunk_key_encodings.py @@ -2,7 +2,12 @@ import pytest -from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding, V2ChunkKeyEncoding +from zarr.core.chunk_key_encodings import ( + DefaultChunkKeyEncoding, + OffsetChunkKeyEncoding, + V2ChunkKeyEncoding, + parse_chunk_key_encoding, +) @pytest.mark.parametrize("separator", ["/", "."]) @@ -63,3 +68,78 @@ def test_default_encoding_rejects_key_using_the_other_separator() -> None: with pytest.raises(ValueError, match="Invalid chunk key"): encoding.decode_chunk_key("c.0.1") + + +@pytest.mark.parametrize("separator", ["/", "."]) +@pytest.mark.parametrize("coords", [(0,), (1, 2), (10, 0, 3)]) +def test_offset_encoding_with_zero_offset_is_the_default_encoding( + separator: str, coords: tuple[int, ...] +) -> None: + """An all-zero offset gives byte-for-byte the default encoding's keys, so an array can + switch to this encoding without any chunk moving.""" + default = DefaultChunkKeyEncoding(separator=separator) # type: ignore[arg-type] + offset = OffsetChunkKeyEncoding(separator=separator, offset=(0,) * len(coords)) # type: ignore[arg-type] + + assert offset.encode_chunk_key(coords) == default.encode_chunk_key(coords) + assert offset.decode_chunk_key(default.encode_chunk_key(coords)) == coords + + +@pytest.mark.parametrize("separator", ["/", "."]) +@pytest.mark.parametrize("coords", [(0, 0), (5, 4), (7, 0)]) +def test_offset_encoding_shifts_keys_and_round_trips( + separator: str, coords: tuple[int, ...] +) -> None: + encoding = OffsetChunkKeyEncoding(separator=separator, offset=(-2, 1)) # type: ignore[arg-type] + + key = encoding.encode_chunk_key(coords) + assert key == separator.join(["c", str(coords[0] - 2), str(coords[1] + 1)]) + assert encoding.decode_chunk_key(key) == coords + + +def test_offset_encoding_produces_negative_key_components() -> None: + """The chunk that holds array index 0 after two chunks were added in front of it is + stored under a negative key; the chunk that held index 0 before keeps `c/0/0`.""" + encoding = OffsetChunkKeyEncoding(offset=(-2, 0)) + + assert encoding.encode_chunk_key((0, 0)) == "c/-2/0" + assert encoding.encode_chunk_key((2, 0)) == "c/0/0" + assert encoding.decode_chunk_key("c/-2/0") == (0, 0) + + +def test_offset_encoding_rejects_a_rank_mismatch() -> None: + encoding = OffsetChunkKeyEncoding(offset=(0, 0)) + + with pytest.raises(ValueError, match="2 dimension"): + encoding.encode_chunk_key((1, 2, 3)) + with pytest.raises(ValueError, match="2 dimension"): + encoding.decode_chunk_key("c/1/2/3") + + +@pytest.mark.parametrize("offset", [(1.5,), ("1",), (True,), "12", 3]) +def test_offset_encoding_rejects_a_non_integer_offset(offset: object) -> None: + with pytest.raises(TypeError, match="sequence of integers"): + OffsetChunkKeyEncoding(offset=offset) # type: ignore[arg-type] + + +def test_offset_encoding_metadata_round_trips() -> None: + encoding = OffsetChunkKeyEncoding(offset=(-1, 0), separator=".") + + assert encoding.to_dict() == { + "name": "offset", + "configuration": {"separator": ".", "offset": [-1, 0]}, + } + assert parse_chunk_key_encoding(encoding.to_dict()) == encoding + assert ( + parse_chunk_key_encoding({"name": "offset", "offset": (-1, 0), "separator": "."}) + == encoding + ) + assert parse_chunk_key_encoding({"name": "offset", "offset": (3,)}) == OffsetChunkKeyEncoding( + offset=(3,) + ) + + +def test_offset_encoding_requires_an_offset() -> None: + with pytest.raises(ValueError, match="requires an offset"): + parse_chunk_key_encoding({"name": "offset"}) + with pytest.raises(ValueError, match="requires an offset"): + parse_chunk_key_encoding({"name": "offset", "configuration": {"separator": "/"}}) From 6b0cdf7ac0275c1bbf4d030b5c502e4d1db50294 Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 21:45:30 +0000 Subject: [PATCH 2/7] prepend: grow by whole shards when sharded, normalise a negative axis The store key addresses a shard, so the unit of growth is the chunk grid's shape, not the inner chunk; and a negative axis has to be normalised before the preserved-shape check and the write selection use it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01CSaEKF9nB55gU7PQgQ22X5 --- src/zarr/core/array.py | 35 ++++++++++++++++++++++------------ tests/test_array.py | 43 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 12 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index c476a4cc00..09c63cd06e 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -1803,7 +1803,8 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: return await _append(self, data, axis) async def prepend_chunks(self, n_chunks: int, axis: int = 0) -> tuple[int, ...]: - """Grow the array by `n_chunks` whole chunks before index 0 of `axis`. + """Grow the array by `n_chunks` whole chunks (shards, when sharded) before index 0 of + `axis`. A metadata-only operation: the array must use the ``offset`` chunk key encoding, under which the new leading chunks take new keys and every existing chunk keeps @@ -1833,7 +1834,8 @@ async def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """Prepend `data` to `axis`: grow the array before index 0 and write `data` there. The array must use the ``offset`` chunk key encoding, and the extent of `data` - along `axis` must be a whole number of chunks; see `prepend_chunks`. + along `axis` must be a whole number of chunks (shards, when sharded); see + `prepend_chunks`. Parameters ---------- @@ -4078,7 +4080,8 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: return sync(self.async_array.append(data, axis=axis)) def prepend_chunks(self, n_chunks: int, axis: int = 0) -> tuple[int, ...]: - """Grow the array by `n_chunks` whole chunks before index 0 of `axis`. + """Grow the array by `n_chunks` whole chunks (shards, when sharded) before index 0 of + `axis`. A metadata-only operation: the array must use the ``offset`` chunk key encoding, under which the new leading chunks take new keys and every existing chunk keeps @@ -4115,7 +4118,8 @@ def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """Prepend `data` to `axis`: grow the array before index 0 and write `data` there. The array must use the ``offset`` chunk key encoding, and the extent of `data` - along `axis` must be a whole number of chunks; see `prepend_chunks`. + along `axis` must be a whole number of chunks (shards, when sharded); see + `prepend_chunks`. Parameters ---------- @@ -6265,6 +6269,13 @@ async def _prepend( if not hasattr(data, "shape"): data = np.asanyarray(data) + # Validates the encoding, the chunk grid and the axis without changing anything. + await _prepend_chunks(array, 0, axis) + metadata = cast("ArrayV3Metadata", array.metadata) + chunk_grid = cast("RegularChunkGridMetadata", metadata.chunk_grid) + ndim = len(metadata.shape) + axis %= ndim + self_shape_preserved = tuple(s for i, s in enumerate(array.shape) if i != axis) data_shape_preserved = tuple(s for i, s in enumerate(data.shape) if i != axis) if self_shape_preserved != data_shape_preserved: @@ -6276,18 +6287,18 @@ async def _prepend( "prepended." ) extent = data.shape[axis] - chunk = array.chunks[axis] - if extent % chunk: + # The unit of growth is what a store key addresses: the shard when sharded, else the chunk. + unit = chunk_grid.chunk_shape[axis] + if extent % unit: + what = "shards" if array.shards is not None else "chunks" raise ValueError( - f"prepend grows the array by whole chunks, but the data spans {extent} along axis " - f"{axis} and the chunk length there is {chunk}." + f"prepend grows the array by whole {what}, but the data spans {extent} along axis " + f"{axis} and the {what[:-1]} length there is {unit}." ) - new_shape = await _prepend_chunks(array, extent // chunk, axis) + new_shape = await _prepend_chunks(array, extent // unit, axis) - prepend_selection = tuple( - slice(None) if i != axis else slice(0, extent) for i in range(len(new_shape)) - ) + prepend_selection = tuple(slice(None) if i != axis else slice(0, extent) for i in range(ndim)) await array.setitem(prepend_selection, data) return new_shape diff --git a/tests/test_array.py b/tests/test_array.py index c86478fdc7..11521b61da 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -746,6 +746,49 @@ def test_prepend_writes_the_data_in_front(store: MemoryStore) -> None: z.prepend(np.zeros((2, 4), dtype="uint8"), axis=0) +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_along_a_negative_axis_writes_in_front_of_that_axis(store: MemoryStore) -> None: + z = zarr.create_array( + store=store, + shape=(2, 4), + chunks=(2, 2), + dtype="uint8", + chunk_key_encoding={"name": "offset", "offset": (0, 0)}, + ) + z[:] = 1 + + assert z.prepend(np.full((2, 2), 5, dtype="uint8"), axis=-1) == (2, 6) + + np.testing.assert_array_equal(z[:, :2], 5) + np.testing.assert_array_equal(z[:, 2:], 1) + assert _chunk_keys(store) == {"c/0/-1", "c/0/0", "c/0/1"} + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_on_a_sharded_array_grows_by_whole_shards(store: MemoryStore) -> None: + """The store key addresses a shard, so the unit of growth is the shard, not the inner + chunk: four rows in front of a shard of four is one shard, one new key.""" + z = zarr.create_array( + store=store, + shape=(4, 2), + chunks=(2, 2), + shards=(4, 2), + dtype="int16", + chunk_key_encoding={"name": "offset", "offset": (0, 0)}, + ) + z[:] = 1 + assert _chunk_keys(store) == {"c/0/0"} + + assert z.prepend(np.full((4, 2), 3, dtype="int16"), axis=0) == (8, 2) + + assert z.shape == (8, 2) + np.testing.assert_array_equal(z[:4], 3) + np.testing.assert_array_equal(z[4:], 1) + assert _chunk_keys(store) == {"c/-1/0", "c/0/0"} + with pytest.raises(ValueError, match="whole shards"): + z.prepend(np.zeros((2, 2), dtype="int16"), axis=0) + + @pytest.mark.parametrize("store", ["memory"], indirect=True) def test_prepend_chunks_along_a_later_axis_and_a_negative_axis(store: MemoryStore) -> None: z = zarr.create_array( From fd613c3f4b9b949498bae19fa498d5904978e652 Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 22:32:21 +0000 Subject: [PATCH 3/7] prepend: keep the append mirror only, move the reasoning to the PR Drops the public prepend_chunks; Array.prepend grows the array by the data's extent (whole chunks, or shards when sharded) and writes it. Docstrings and the docs paragraph are cut to what the API needs. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01CSaEKF9nB55gU7PQgQ22X5 --- changes/+offset-chunk-key-encoding.feature.md | 4 +- docs/user-guide/arrays.md | 20 +- src/zarr/core/array.py | 183 ++++-------------- src/zarr/core/chunk_key_encodings.py | 10 +- tests/test_array.py | 95 +++------ 5 files changed, 72 insertions(+), 240 deletions(-) diff --git a/changes/+offset-chunk-key-encoding.feature.md b/changes/+offset-chunk-key-encoding.feature.md index 4d2587e33c..39f4820e9c 100644 --- a/changes/+offset-chunk-key-encoding.feature.md +++ b/changes/+offset-chunk-key-encoding.feature.md @@ -1,3 +1,3 @@ Added the ``offset`` chunk key encoding (the ``default`` keys of ``chunk_coords + offset``) and -``Array.prepend_chunks`` / ``Array.prepend``, which grow an array before index 0 by whole chunks -as a metadata-only change: every existing chunk keeps its key. +``Array.prepend``, which adds data before index 0 of an axis by whole chunks without changing any +existing chunk's key. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index 06b64ae3e1..1249ce93ed 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -154,30 +154,22 @@ z.append(np.vstack([a, a]), axis=1) print(f"Shape after second append: {z.shape}") ``` -### Growing before index 0 +### Prepending -`resize` and `append` only move the high edge of a dimension. Adding data *before* index -0 would shift every array index and, under the `default` chunk key encoding, change the key -of every chunk. The `offset` chunk key encoding stores each chunk under the `default` key of -its grid coordinate plus a fixed per-dimension offset, so that -[`zarr.Array.prepend_chunks`][] can grow a dimension by whole chunks as a metadata-only -change: the new leading chunks take new (negative) keys and every existing chunk keeps its -key. [`zarr.Array.prepend`][] grows and writes in one call. An array opts in at creation: +[`zarr.Array.prepend`][] adds data before index 0 of an axis. Under the `default` chunk key +encoding that would change every chunk's key, so it needs the `offset` encoding, which stores +each chunk under the `default` key of its grid coordinate plus a per-dimension offset. The +growth is then a metadata change and the existing chunks stay where they are; the data must +span a whole number of chunks along the axis. An array opts in at creation: ```python exec="true" session="arrays" source="above" result="ansi" z = zarr.create_array(store='data/example-5.zarr', shape=(6, 4), dtype='int32', chunks=(3, 4), chunk_key_encoding={'name': 'offset', 'offset': (0, 0)}, overwrite=True) z[:] = 1 -print(f"Original shape: {z.shape}, offset: {z.metadata.chunk_key_encoding.offset}") z.prepend(np.full((6, 4), 2, dtype='int32')) print(f"Shape after prepend: {z.shape}, offset: {z.metadata.chunk_key_encoding.offset}") -print(f"Chunks written: {z.nchunks_initialized}") ``` -With an all-zero offset the keys are those of the `default` encoding, so an existing array -can adopt the encoding by rewriting its metadata, without moving a chunk. Readers that do not -implement the encoding refuse the array rather than misplace its chunks. - ## Runtime configuration Zarr arrays are parametrized with a configuration that determines certain aspects of array behavior. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 09c63cd06e..13a5b589b1 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -1802,40 +1802,8 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """ return await _append(self, data, axis) - async def prepend_chunks(self, n_chunks: int, axis: int = 0) -> tuple[int, ...]: - """Grow the array by `n_chunks` whole chunks (shards, when sharded) before index 0 of - `axis`. - - A metadata-only operation: the array must use the ``offset`` chunk key encoding, - under which the new leading chunks take new keys and every existing chunk keeps - its key, so no chunk is read, moved or rewritten. The new region reads as the fill - value until written. - - Parameters - ---------- - n_chunks : int - Number of whole chunks to add before index 0. - axis : int - Axis along which to grow. - - Returns - ------- - new_shape : tuple - - Raises - ------ - TypeError - If the array does not use the ``offset`` chunk key encoding, or its chunk grid - is not regular. - """ - return await _prepend_chunks(self, n_chunks, axis) - async def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: - """Prepend `data` to `axis`: grow the array before index 0 and write `data` there. - - The array must use the ``offset`` chunk key encoding, and the extent of `data` - along `axis` must be a whole number of chunks (shards, when sharded); see - `prepend_chunks`. + """Prepend `data` to `axis`. Parameters ---------- @@ -1851,7 +1819,9 @@ async def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: Notes ----- The size of all dimensions other than `axis` must match between this - array and `data`. + array and `data`, and the size along `axis` must be a whole number of + chunks (shards, when sharded). The array must use the ``offset`` chunk + key encoding; existing chunks keep their keys. """ return await _prepend(self, data, axis) @@ -4079,47 +4049,8 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """ return sync(self.async_array.append(data, axis=axis)) - def prepend_chunks(self, n_chunks: int, axis: int = 0) -> tuple[int, ...]: - """Grow the array by `n_chunks` whole chunks (shards, when sharded) before index 0 of - `axis`. - - A metadata-only operation: the array must use the ``offset`` chunk key encoding, - under which the new leading chunks take new keys and every existing chunk keeps - its key, so no chunk is read, moved or rewritten. The new region reads as the fill - value until written. - - Parameters - ---------- - n_chunks : int - Number of whole chunks to add before index 0. - axis : int - Axis along which to grow. - - Returns - ------- - new_shape : tuple - - Examples - -------- - ```python - import zarr - z = zarr.create_array(store={}, shape=(6, 4), chunks=(3, 4), dtype="int32", - chunk_key_encoding={"name": "offset", "offset": (0, 0)}) - z[:] = 1 - z.prepend_chunks(2, axis=0) - z.shape - #> (12, 4) - z[:6] = 0 - ``` - """ - return sync(self.async_array.prepend_chunks(n_chunks, axis=axis)) - def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: - """Prepend `data` to `axis`: grow the array before index 0 and write `data` there. - - The array must use the ``offset`` chunk key encoding, and the extent of `data` - along `axis` must be a whole number of chunks (shards, when sharded); see - `prepend_chunks`. + """Prepend `data` to `axis`. Parameters ---------- @@ -4135,7 +4066,9 @@ def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: Notes ----- The size of all dimensions other than `axis` must match between this - array and `data`. + array and `data`, and the size along `axis` must be a whole number of + chunks (shards, when sharded). The array must use the ``offset`` chunk + key encoding; existing chunks keep their keys. """ return sync(self.async_array.prepend(data, axis=axis)) @@ -6180,70 +6113,6 @@ async def _append( return new_shape -async def _prepend_chunks( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], - n_chunks: int, - axis: int = 0, -) -> tuple[int, ...]: - """ - Grow an array by whole chunks before index 0 of `axis`, as a metadata-only change. - - Parameters - ---------- - array : AsyncArray - The array to grow. Must use the ``offset`` chunk key encoding and a regular chunk grid. - n_chunks : int - Number of whole chunks to add. - axis : int - Axis along which to grow. - - Returns - ------- - new_shape : tuple[int, ...] - The new shape of the array. - """ - metadata = array.metadata - if not isinstance(metadata, ArrayV3Metadata) or not isinstance( - metadata.chunk_key_encoding, OffsetChunkKeyEncoding - ): - raise TypeError( - "Growing an array before index 0 without moving its chunks needs the offset chunk " - "key encoding (chunk_key_encoding={'name': 'offset', 'offset': (0,) * ndim} at " - "creation); this array uses " - f"{metadata.chunk_key_encoding.to_dict() if isinstance(metadata, ArrayV3Metadata) else 'Zarr v2 keys'}." - ) - if not isinstance(metadata.chunk_grid, RegularChunkGridMetadata): - raise TypeError("prepend_chunks needs a regular chunk grid.") - ndim = len(metadata.shape) - if not -ndim <= axis < ndim: - raise ValueError(f"axis {axis} is out of range for a {ndim}-dimensional array.") - axis %= ndim - if n_chunks < 0: - raise ValueError(f"n_chunks must be non-negative, got {n_chunks}.") - if n_chunks == 0: - return metadata.shape - - chunk = metadata.chunk_grid.chunk_shape[axis] - new_shape = tuple( - s + n_chunks * chunk if i == axis else s for i, s in enumerate(metadata.shape) - ) - new_offset = tuple( - o - n_chunks if i == axis else o for i, o in enumerate(metadata.chunk_key_encoding.offset) - ) - new_metadata = replace( - metadata.update_shape(new_shape), - chunk_key_encoding=replace(metadata.chunk_key_encoding, offset=new_offset), - ) - new_chunk_grid = ChunkGrid.from_metadata(new_metadata) - - # Both edits land in one document write, so a reader never sees one without the other. - await store_documents(array.store_path, encode_documents(array.store_path, new_metadata)) - array._stored_document_replaced() - object.__setattr__(array, "metadata", new_metadata) - object.__setattr__(array, "_chunk_grid", new_chunk_grid) - return new_shape - - async def _prepend( array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], data: npt.ArrayLike, @@ -6257,7 +6126,7 @@ async def _prepend( array : AsyncArray The array to prepend to. Must use the ``offset`` chunk key encoding. data : npt.ArrayLike - Data to be prepended. Its extent along `axis` must be a whole number of chunks. + Data to be prepended. Its size along `axis` must be a whole number of chunks. axis : int Axis along which to prepend. @@ -6269,11 +6138,19 @@ async def _prepend( if not hasattr(data, "shape"): data = np.asanyarray(data) - # Validates the encoding, the chunk grid and the axis without changing anything. - await _prepend_chunks(array, 0, axis) - metadata = cast("ArrayV3Metadata", array.metadata) - chunk_grid = cast("RegularChunkGridMetadata", metadata.chunk_grid) + metadata = array.metadata + if not isinstance(metadata, ArrayV3Metadata) or not isinstance( + metadata.chunk_key_encoding, OffsetChunkKeyEncoding + ): + raise TypeError( + "prepend needs the offset chunk key encoding " + "(chunk_key_encoding={'name': 'offset', 'offset': (0,) * ndim} at creation)." + ) + if not isinstance(metadata.chunk_grid, RegularChunkGridMetadata): + raise TypeError("prepend needs a regular chunk grid.") ndim = len(metadata.shape) + if not -ndim <= axis < ndim: + raise ValueError(f"axis {axis} is out of range for a {ndim}-dimensional array.") axis %= ndim self_shape_preserved = tuple(s for i, s in enumerate(array.shape) if i != axis) @@ -6287,16 +6164,28 @@ async def _prepend( "prepended." ) extent = data.shape[axis] - # The unit of growth is what a store key addresses: the shard when sharded, else the chunk. - unit = chunk_grid.chunk_shape[axis] + unit = metadata.chunk_grid.chunk_shape[axis] if extent % unit: what = "shards" if array.shards is not None else "chunks" raise ValueError( f"prepend grows the array by whole {what}, but the data spans {extent} along axis " f"{axis} and the {what[:-1]} length there is {unit}." ) + n_units = extent // unit - new_shape = await _prepend_chunks(array, extent // unit, axis) + new_shape = tuple(s + n_units * unit if i == axis else s for i, s in enumerate(metadata.shape)) + new_offset = tuple( + o - n_units if i == axis else o for i, o in enumerate(metadata.chunk_key_encoding.offset) + ) + new_metadata = replace( + metadata.update_shape(new_shape), + chunk_key_encoding=replace(metadata.chunk_key_encoding, offset=new_offset), + ) + new_chunk_grid = ChunkGrid.from_metadata(new_metadata) + await store_documents(array.store_path, encode_documents(array.store_path, new_metadata)) + array._stored_document_replaced() + object.__setattr__(array, "metadata", new_metadata) + object.__setattr__(array, "_chunk_grid", new_chunk_grid) prepend_selection = tuple(slice(None) if i != axis else slice(0, extent) for i in range(ndim)) await array.setitem(prepend_selection, data) diff --git a/src/zarr/core/chunk_key_encodings.py b/src/zarr/core/chunk_key_encodings.py index a113b46ab7..c9649fd0db 100644 --- a/src/zarr/core/chunk_key_encodings.py +++ b/src/zarr/core/chunk_key_encodings.py @@ -119,15 +119,9 @@ def parse_offset(data: object) -> tuple[int, ...]: @dataclass(frozen=True) class OffsetChunkKeyEncoding(ChunkKeyEncoding): """ - The ``default`` encoding applied to ``chunk_coords + offset``. + The ``default`` encoding applied to ``chunk_coords + offset``, one integer per dimension. - ``offset`` holds one integer per array dimension and may be negative. Array indexing is - unchanged and zero-based; only the mapping from chunk coordinates to store keys moves, so - growing the array by ``k`` whole chunks before index 0 along an axis is two metadata edits - (``shape += k * chunk``, ``offset -= k``) under which every existing chunk keeps its key. - With an all-zero ``offset`` the keys are those of ``default`` with the same separator. - - See https://github.com/zarr-developers/zarr-extensions/tree/main/chunk-key-encodings/offset + https://github.com/zarr-developers/zarr-extensions/tree/main/chunk-key-encodings/offset """ name: ClassVar[Literal["offset"]] = "offset" diff --git a/tests/test_array.py b/tests/test_array.py index 11521b61da..eb3c3b2cb4 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -685,11 +685,18 @@ async def _collect(it: AsyncIterator[str]) -> list[str]: return [key async for key in it] +def _chunk_bytes(store: Store) -> dict[str, bytes]: + proto = default_buffer_prototype() + return { + key: sync(store.get(key, prototype=proto)).to_bytes() # type: ignore[union-attr] + for key in _chunk_keys(store) + } + + @pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_chunks_is_a_metadata_only_edit(store: MemoryStore) -> None: - """Growing before index 0 touches zarr.json and nothing else: the chunk key set is - unchanged, the old data reads back at its new indices, and the new region reads as the - fill value until written.""" +def test_prepend_writes_in_front_and_leaves_every_existing_chunk_untouched( + store: MemoryStore, +) -> None: z = zarr.create_array( store=store, shape=(6, 4), @@ -700,30 +707,27 @@ def test_prepend_chunks_is_a_metadata_only_edit(store: MemoryStore) -> None: ) data = np.arange(24, dtype="int32").reshape(6, 4) z[:] = data - keys_before = _chunk_keys(store) - assert keys_before == {"c/0/0", "c/0/1", "c/1/0", "c/1/1"} + before = _chunk_bytes(store) + assert set(before) == {"c/0/0", "c/0/1", "c/1/0", "c/1/1"} - assert z.prepend_chunks(2, axis=0) == (12, 4) + assert z.prepend(np.full((6, 4), 7, dtype="int32"), axis=0) == (12, 4) - assert z.shape == (12, 4) - assert _chunk_keys(store) == keys_before + after = _chunk_bytes(store) + assert {k: after[k] for k in before} == before + assert set(after) - set(before) == {"c/-2/0", "c/-2/1", "c/-1/0", "c/-1/1"} + np.testing.assert_array_equal(z[:6], 7) np.testing.assert_array_equal(z[6:], data) - np.testing.assert_array_equal(z[:6], -1) assert z.metadata.to_dict()["chunk_key_encoding"] == { "name": "offset", "configuration": {"separator": "/", "offset": [-2, 0]}, } - - z[:3] = 7 - assert _chunk_keys(store) == keys_before | {"c/-2/0", "c/-2/1"} - reopened = zarr.open_array(store=store, mode="r") assert reopened.shape == (12, 4) np.testing.assert_array_equal(reopened[:], z[:]) @pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_writes_the_data_in_front(store: MemoryStore) -> None: +def test_prepend_requires_whole_chunks_and_matching_other_dimensions(store: MemoryStore) -> None: z = zarr.create_array( store=store, shape=(4, 3), @@ -731,29 +735,23 @@ def test_prepend_writes_the_data_in_front(store: MemoryStore) -> None: dtype="uint8", chunk_key_encoding={"name": "offset", "offset": (0, 0)}, ) - z[:] = 1 - front = np.full((4, 3), 9, dtype="uint8") - - assert z.prepend(front, axis=0) == (8, 3) - - np.testing.assert_array_equal(z[:4], 9) - np.testing.assert_array_equal(z[4:], 1) - assert _chunk_keys(store) == {"c/-2/0", "c/-1/0", "c/0/0", "c/1/0"} - with pytest.raises(ValueError, match="whole chunks"): z.prepend(np.zeros((3, 3), dtype="uint8"), axis=0) with pytest.raises(ValueError, match="not compatible"): z.prepend(np.zeros((2, 4), dtype="uint8"), axis=0) + with pytest.raises(ValueError, match="out of range"): + z.prepend(np.zeros((2, 3), dtype="uint8"), axis=2) + assert z.shape == (4, 3) @pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_along_a_negative_axis_writes_in_front_of_that_axis(store: MemoryStore) -> None: +def test_prepend_along_a_negative_axis_with_a_dot_separator(store: MemoryStore) -> None: z = zarr.create_array( store=store, shape=(2, 4), chunks=(2, 2), dtype="uint8", - chunk_key_encoding={"name": "offset", "offset": (0, 0)}, + chunk_key_encoding={"name": "offset", "separator": ".", "offset": (0, 0)}, ) z[:] = 1 @@ -761,13 +759,11 @@ def test_prepend_along_a_negative_axis_writes_in_front_of_that_axis(store: Memor np.testing.assert_array_equal(z[:, :2], 5) np.testing.assert_array_equal(z[:, 2:], 1) - assert _chunk_keys(store) == {"c/0/-1", "c/0/0", "c/0/1"} + assert _chunk_keys(store) == {"c.0.-1", "c.0.0", "c.0.1"} @pytest.mark.parametrize("store", ["memory"], indirect=True) def test_prepend_on_a_sharded_array_grows_by_whole_shards(store: MemoryStore) -> None: - """The store key addresses a shard, so the unit of growth is the shard, not the inner - chunk: four rows in front of a shard of four is one shard, one new key.""" z = zarr.create_array( store=store, shape=(4, 2), @@ -781,7 +777,6 @@ def test_prepend_on_a_sharded_array_grows_by_whole_shards(store: MemoryStore) -> assert z.prepend(np.full((4, 2), 3, dtype="int16"), axis=0) == (8, 2) - assert z.shape == (8, 2) np.testing.assert_array_equal(z[:4], 3) np.testing.assert_array_equal(z[4:], 1) assert _chunk_keys(store) == {"c/-1/0", "c/0/0"} @@ -789,57 +784,19 @@ def test_prepend_on_a_sharded_array_grows_by_whole_shards(store: MemoryStore) -> z.prepend(np.zeros((2, 2), dtype="int16"), axis=0) -@pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_chunks_along_a_later_axis_and_a_negative_axis(store: MemoryStore) -> None: - z = zarr.create_array( - store=store, - shape=(2, 6), - chunks=(2, 3), - dtype="int8", - chunk_key_encoding={"name": "offset", "separator": ".", "offset": (0, 0)}, - ) - z[:] = np.arange(12, dtype="int8").reshape(2, 6) - - assert z.prepend_chunks(1, axis=-1) == (2, 9) - - np.testing.assert_array_equal(z[:, 3:], np.arange(12, dtype="int8").reshape(2, 6)) - assert _chunk_keys(store) == {"c.0.0", "c.0.1"} - assert z.prepend_chunks(0, axis=1) == (2, 9) - - @pytest.mark.parametrize("store", ["memory"], indirect=True) @pytest.mark.parametrize("zarr_format", [2, 3]) -def test_prepend_chunks_refuses_an_array_without_the_offset_encoding( +def test_prepend_refuses_an_array_without_the_offset_encoding( store: MemoryStore, zarr_format: ZarrFormat ) -> None: - """A default-encoded array cannot grow in front without every chunk changing its key, - so the request is refused and names the encoding that allows it.""" z = zarr.create_array( store=store, shape=(4,), chunks=(2,), dtype="int8", zarr_format=zarr_format ) - - with pytest.raises(TypeError, match="offset chunk key encoding"): - z.prepend_chunks(1) with pytest.raises(TypeError, match="offset chunk key encoding"): z.prepend(np.zeros(2, dtype="int8")) assert z.shape == (4,) -@pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_chunks_rejects_a_bad_axis_or_count(store: MemoryStore) -> None: - z = zarr.create_array( - store=store, - shape=(4,), - chunks=(2,), - dtype="int8", - chunk_key_encoding={"name": "offset", "offset": (0,)}, - ) - with pytest.raises(ValueError, match="out of range"): - z.prepend_chunks(1, axis=1) - with pytest.raises(ValueError, match="non-negative"): - z.prepend_chunks(-1) - - @pytest.mark.parametrize("store", ["memory"], indirect=True) def test_create_array_rejects_an_offset_of_the_wrong_rank(store: MemoryStore) -> None: with pytest.raises(ValueError, match="2 dimension"): From 5375df3ebc24c1fee1ce6719e16736b9b5a9d80d Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 22:46:22 +0000 Subject: [PATCH 4/7] changelog: name the fragment after the PR Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01CSaEKF9nB55gU7PQgQ22X5 --- .../{+offset-chunk-key-encoding.feature.md => 4516.feature.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{+offset-chunk-key-encoding.feature.md => 4516.feature.md} (100%) diff --git a/changes/+offset-chunk-key-encoding.feature.md b/changes/4516.feature.md similarity index 100% rename from changes/+offset-chunk-key-encoding.feature.md rename to changes/4516.feature.md From 86e0b74de3c53e28fd5197b6b05fa911fd35e33a Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Fri, 16 Oct 2026 00:45:22 +0000 Subject: [PATCH 5/7] tests: prepend from a list and on a rectilinear grid Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01CSaEKF9nB55gU7PQgQ22X5 --- tests/test_array.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/test_array.py b/tests/test_array.py index eb3c3b2cb4..33b7746d7e 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -797,6 +797,30 @@ def test_prepend_refuses_an_array_without_the_offset_encoding( assert z.shape == (4,) +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_accepts_a_plain_list_and_refuses_a_rectilinear_grid(store: MemoryStore) -> None: + z = zarr.create_array( + store=store, + shape=(2,), + chunks=(2,), + dtype="int8", + chunk_key_encoding={"name": "offset", "offset": (0,)}, + ) + z[:] = 1 + assert z.prepend([4, 5]) == (4,) + np.testing.assert_array_equal(z[:], [4, 5, 1, 1]) + + rectilinear = zarr.create_array( + store=MemoryStore(), + shape=(4,), + chunks=[[2, 2]], + dtype="int8", + chunk_key_encoding={"name": "offset", "offset": (0,)}, + ) + with pytest.raises(TypeError, match="regular chunk grid"): + rectilinear.prepend(np.zeros(2, dtype="int8")) + + @pytest.mark.parametrize("store", ["memory"], indirect=True) def test_create_array_rejects_an_offset_of_the_wrong_rank(store: MemoryStore) -> None: with pytest.raises(ValueError, match="2 dimension"): From 7501ff30565680342e266600df4bcf55fc0512e8 Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Fri, 16 Oct 2026 00:45:57 +0000 Subject: [PATCH 6/7] tests: enable rectilinear chunks for the refusal case Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01CSaEKF9nB55gU7PQgQ22X5 --- tests/test_array.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/tests/test_array.py b/tests/test_array.py index 33b7746d7e..77ded6f1f9 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -810,15 +810,16 @@ def test_prepend_accepts_a_plain_list_and_refuses_a_rectilinear_grid(store: Memo assert z.prepend([4, 5]) == (4,) np.testing.assert_array_equal(z[:], [4, 5, 1, 1]) - rectilinear = zarr.create_array( - store=MemoryStore(), - shape=(4,), - chunks=[[2, 2]], - dtype="int8", - chunk_key_encoding={"name": "offset", "offset": (0,)}, - ) - with pytest.raises(TypeError, match="regular chunk grid"): - rectilinear.prepend(np.zeros(2, dtype="int8")) + with zarr.config.set({"array.rectilinear_chunks": True}): + rectilinear = zarr.create_array( + store=MemoryStore(), + shape=(4,), + chunks=[[2, 2]], + dtype="int8", + chunk_key_encoding={"name": "offset", "offset": (0,)}, + ) + with pytest.raises(TypeError, match="regular chunk grid"): + rectilinear.prepend(np.zeros(2, dtype="int8")) @pytest.mark.parametrize("store", ["memory"], indirect=True) From fb48e955fde94719614a39129665c06319243f19 Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Fri, 16 Oct 2026 22:50:52 +0000 Subject: [PATCH 7/7] Prepend through a regular_bounded chunk grid instead of an offset key encoding Rework following zarr-extensions#79: the `offset` chunk key encoding is removed. `regular_bounded` is a regular chunk grid with an `origin` (the array coordinate where chunk 0 starts) and a `lower_bound` (the coordinate of the first element), so chunk grid indices can be negative and the first chunk can be partly filled, mirroring the last chunk. `Array.prepend` no longer needs whole chunks or an opt-in at creation: the first prepend switches a `regular` grid to `regular_bounded`, and existing chunks keep their keys and bytes. `Array.domain` reports each dimension's coordinate range; indexing stays zero-based. Internally a regular dimension gains a `shift` (where the first element sits in the first chunk), and `buffer_offset` gives where a chunk's data starts inside its buffer, which the indexers add when converting to chunk-local positions. Key decoding rejects `-0` components. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CSaEKF9nB55gU7PQgQ22X5 --- changes/4516.feature.md | 7 +- docs/user-guide/arrays.md | 22 +-- src/zarr/core/array.py | 78 ++++++---- src/zarr/core/chunk_grids.py | 50 +++++-- src/zarr/core/chunk_key_encodings.py | 80 ++-------- src/zarr/core/indexing.py | 16 +- src/zarr/core/metadata/v3.py | 98 +++++++++++- tests/test_array.py | 215 ++++++++++++++++----------- tests/test_chunk_key_encodings.py | 88 ++--------- 9 files changed, 359 insertions(+), 295 deletions(-) diff --git a/changes/4516.feature.md b/changes/4516.feature.md index 39f4820e9c..985f10f83b 100644 --- a/changes/4516.feature.md +++ b/changes/4516.feature.md @@ -1,3 +1,4 @@ -Added the ``offset`` chunk key encoding (the ``default`` keys of ``chunk_coords + offset``) and -``Array.prepend``, which adds data before index 0 of an axis by whole chunks without changing any -existing chunk's key. +Added ``Array.prepend``, which adds data before index 0 of an axis without moving existing chunks, +and the ``regular_bounded`` chunk grid it switches an array to: a regular grid with an ``origin`` +and a ``lower_bound``, under which chunk grid indices can be negative. ``Array.domain`` reports the +coordinate range of each dimension. Chunk keys with a ``-0`` component are now rejected. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index 1249ce93ed..f96d872f4b 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -156,18 +156,22 @@ print(f"Shape after second append: {z.shape}") ### Prepending -[`zarr.Array.prepend`][] adds data before index 0 of an axis. Under the `default` chunk key -encoding that would change every chunk's key, so it needs the `offset` encoding, which stores -each chunk under the `default` key of its grid coordinate plus a per-dimension offset. The -growth is then a metadata change and the existing chunks stay where they are; the data must -span a whole number of chunks along the axis. An array opts in at creation: +[`zarr.Array.prepend`][] adds data before index 0 of an axis, the way [`zarr.Array.append`][] +adds it after the end. The existing chunks keep their keys and their bytes. Data that does not +fill a whole chunk goes into a partly filled first chunk, like the partly filled last chunk of +any array. + +The first prepend switches the chunk grid from `regular` to `regular_bounded`, which records +the coordinate of the first element (`lower_bound`) and of the start of chunk 0 (`origin`), so +chunks in front of the old start get negative keys such as `c/-1/0`. Indexing stays zero-based, +and [`zarr.Array.domain`][] gives the coordinate range of each dimension. Readers that do not +implement `regular_bounded` refuse such an array rather than misread it. ```python exec="true" session="arrays" source="above" result="ansi" -z = zarr.create_array(store='data/example-5.zarr', shape=(6, 4), dtype='int32', chunks=(3, 4), - chunk_key_encoding={'name': 'offset', 'offset': (0, 0)}, overwrite=True) +z = zarr.create_array(store='data/example-5.zarr', shape=(6, 4), dtype='int32', chunks=(4, 4), overwrite=True) z[:] = 1 -z.prepend(np.full((6, 4), 2, dtype='int32')) -print(f"Shape after prepend: {z.shape}, offset: {z.metadata.chunk_key_encoding.offset}") +z.prepend(np.full((3, 4), 2, dtype='int32')) +print(f"Shape: {z.shape}, domain: {z.domain}") ``` ## Runtime configuration diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 13a5b589b1..252fa0062c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -54,7 +54,6 @@ ChunkKeyEncoding, ChunkKeyEncodingLike, DefaultChunkKeyEncoding, - OffsetChunkKeyEncoding, V2ChunkKeyEncoding, parse_chunk_key_encoding, ) @@ -139,6 +138,7 @@ from zarr.core.metadata.v3 import ( ChunkGridMetadata, RectilinearChunkGridMetadata, + RegularBoundedChunkGridMetadata, RegularChunkGridMetadata, create_chunk_grid_metadata, parse_node_type_array, @@ -950,6 +950,20 @@ def write_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: return self._chunk_grid.chunk_sizes + @property + def domain(self) -> tuple[tuple[int, int], ...]: + """The half-open coordinate range ``(start, stop)`` of each dimension. + + Index ``i`` along a dimension is coordinate ``start + i``. ``start`` is 0 + unless the array has been grown with `prepend`. + """ + grid = getattr(self.metadata, "chunk_grid", None) + if isinstance(grid, RegularBoundedChunkGridMetadata): + lower_bound = grid.lower_bound + else: + lower_bound = (0,) * len(self.shape) + return tuple((lb, lb + n) for lb, n in zip(lower_bound, self.shape, strict=True)) + @property def shards(self) -> tuple[int, ...] | None: """Returns the shard shape of the Array. @@ -1819,9 +1833,9 @@ async def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: Notes ----- The size of all dimensions other than `axis` must match between this - array and `data`, and the size along `axis` must be a whole number of - chunks (shards, when sharded). The array must use the ``offset`` chunk - key encoding; existing chunks keep their keys. + array and `data`. Existing chunks keep their keys. The first prepend + switches a ``regular`` chunk grid to ``regular_bounded``, which readers + without that chunk grid refuse. """ return await _prepend(self, data, axis) @@ -2218,6 +2232,15 @@ def write_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: """ return self.async_array.write_chunk_sizes + @property + def domain(self) -> tuple[tuple[int, int], ...]: + """The half-open coordinate range ``(start, stop)`` of each dimension. + + Index ``i`` along a dimension is coordinate ``start + i``. ``start`` is 0 + unless the array has been grown with `prepend`. + """ + return self.async_array.domain + @property def shards(self) -> tuple[int, ...] | None: """Returns a tuple of integers describing the length of each dimension of a shard of the array. @@ -4066,9 +4089,9 @@ def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: Notes ----- The size of all dimensions other than `axis` must match between this - array and `data`, and the size along `axis` must be a whole number of - chunks (shards, when sharded). The array must use the ``offset`` chunk - key encoding; existing chunks keep their keys. + array and `data`. Existing chunks keep their keys. The first prepend + switches a ``regular`` chunk grid to ``regular_bounded``, which readers + without that chunk grid refuse. """ return sync(self.async_array.prepend(data, axis=axis)) @@ -5612,7 +5635,7 @@ async def _nchunks_initialized( grid = array._chunk_grid total = 0 for key in await _shards_initialized(array): - spec = grid[meta.chunk_key_encoding.decode_chunk_key(key)] + spec = grid[meta.decode_chunk_key(key)] if spec is not None: total += product( tuple(s // c for s, c in zip(spec.codec_shape, inner_chunks, strict=True)) @@ -6124,9 +6147,9 @@ async def _prepend( Parameters ---------- array : AsyncArray - The array to prepend to. Must use the ``offset`` chunk key encoding. + The array to prepend to. Must be a Zarr v3 array with a regular chunk grid. data : npt.ArrayLike - Data to be prepended. Its size along `axis` must be a whole number of chunks. + Data to be prepended. axis : int Axis along which to prepend. @@ -6140,14 +6163,9 @@ async def _prepend( metadata = array.metadata if not isinstance(metadata, ArrayV3Metadata) or not isinstance( - metadata.chunk_key_encoding, OffsetChunkKeyEncoding + metadata.chunk_grid, RegularChunkGridMetadata ): - raise TypeError( - "prepend needs the offset chunk key encoding " - "(chunk_key_encoding={'name': 'offset', 'offset': (0,) * ndim} at creation)." - ) - if not isinstance(metadata.chunk_grid, RegularChunkGridMetadata): - raise TypeError("prepend needs a regular chunk grid.") + raise TypeError("prepend needs a Zarr v3 array with a regular chunk grid.") ndim = len(metadata.shape) if not -ndim <= axis < ndim: raise ValueError(f"axis {axis} is out of range for a {ndim}-dimensional array.") @@ -6164,23 +6182,19 @@ async def _prepend( "prepended." ) extent = data.shape[axis] - unit = metadata.chunk_grid.chunk_shape[axis] - if extent % unit: - what = "shards" if array.shards is not None else "chunks" - raise ValueError( - f"prepend grows the array by whole {what}, but the data spans {extent} along axis " - f"{axis} and the {what[:-1]} length there is {unit}." - ) - n_units = extent // unit + grid = metadata.chunk_grid + if isinstance(grid, RegularBoundedChunkGridMetadata): + origin, lower_bound = grid.origin, grid.lower_bound + else: + origin = lower_bound = (0,) * ndim - new_shape = tuple(s + n_units * unit if i == axis else s for i, s in enumerate(metadata.shape)) - new_offset = tuple( - o - n_units if i == axis else o for i, o in enumerate(metadata.chunk_key_encoding.offset) - ) - new_metadata = replace( - metadata.update_shape(new_shape), - chunk_key_encoding=replace(metadata.chunk_key_encoding, offset=new_offset), + new_shape = tuple(s + extent if i == axis else s for i, s in enumerate(metadata.shape)) + new_grid = RegularBoundedChunkGridMetadata( + chunk_shape=grid.chunk_shape, + origin=origin, + lower_bound=tuple(lb - extent if i == axis else lb for i, lb in enumerate(lower_bound)), ) + new_metadata = replace(metadata, shape=new_shape, chunk_grid=new_grid) new_chunk_grid = ChunkGrid.from_metadata(new_metadata) await store_documents(array.store_path, encode_documents(array.store_path, new_metadata)) array._stored_document_replaced() diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 76bf399a65..cbc66bd081 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -48,10 +48,12 @@ @dataclass(frozen=True) class FixedDimension: """Uniform chunk size. Boundary chunks contain less data but are - encoded at full size by the codec pipeline.""" + encoded at full size by the codec pipeline. With a non-zero ``shift`` the + first element sits ``shift`` positions into the first chunk's buffer.""" size: int # chunk edge length (>= 1) extent: int # array dimension length (>= 0) + shift: int = 0 # position of array index 0 inside the first chunk, in [0, size) nchunks: int = field(init=False, repr=False) ngridcells: int = field(init=False, repr=False) @@ -60,7 +62,9 @@ def __post_init__(self) -> None: raise ValueError(f"FixedDimension size must be >= 1, got {self.size}") if self.extent < 0: raise ValueError(f"FixedDimension extent must be >= 0, got {self.extent}") - n = ceildiv_int(self.extent, self.size) + if not 0 <= self.shift < self.size: + raise ValueError(f"FixedDimension shift must be in [0, {self.size}), got {self.shift}") + n = ceildiv_int(self.extent + self.shift, self.size) if self.extent else 0 object.__setattr__(self, "nchunks", n) object.__setattr__(self, "ngridcells", n) @@ -69,15 +73,19 @@ def index_to_chunk(self, idx: int) -> int: raise IndexError(f"Negative index {idx} is not allowed") if idx >= self.extent: raise IndexError(f"Index {idx} is out of bounds for extent {self.extent}") - return idx // self.size + return (idx + self.shift) // self.size def chunk_offset(self, chunk_ix: int) -> int: - """Byte-aligned start position of chunk *chunk_ix* in array coordinates. + """Start of the valid data region of chunk *chunk_ix* in array coordinates. Does not validate *chunk_ix* — callers must ensure it is in ``[0, nchunks)``. Use ``ChunkGrid.__getitem__`` for safe access. """ - return chunk_ix * self.size + return max(0, chunk_ix * self.size - self.shift) + + def buffer_offset(self, chunk_ix: int) -> int: + """Position of the valid data region inside the chunk's buffer.""" + return self.shift if chunk_ix == 0 else 0 def chunk_size(self, chunk_ix: int) -> int: """Buffer size for codec processing — always uniform. @@ -93,7 +101,8 @@ def data_size(self, chunk_ix: int) -> int: Does not validate *chunk_ix* — callers must ensure it is in ``[0, nchunks)``. Use ``ChunkGrid.__getitem__`` for safe access. """ - return max(0, min(self.size, self.extent - chunk_ix * self.size)) + start = self.chunk_offset(chunk_ix) + return max(0, min((chunk_ix + 1) * self.size - self.shift, self.extent) - start) @property def _unique_edge_lengths(self) -> Iterable[int]: @@ -106,7 +115,7 @@ def _unique_edge_lengths(self) -> Iterable[int]: return (self.size,) def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: - return indices // self.size + return (indices + self.shift) // self.size def with_extent(self, new_extent: int) -> FixedDimension: """Re-bind to *new_extent* without modifying edges. @@ -115,7 +124,7 @@ def with_extent(self, new_extent: int) -> FixedDimension: are already correct. Raises on ``VaryingDimension`` if edges don't cover the new extent. """ - return FixedDimension(size=self.size, extent=new_extent) + return FixedDimension(size=self.size, extent=new_extent, shift=self.shift) def resize(self, new_extent: int) -> FixedDimension: """Adapt for a user-initiated array resize, growing edges if needed. @@ -123,7 +132,7 @@ def resize(self, new_extent: int) -> FixedDimension: For ``FixedDimension`` this is identical to ``with_extent`` since regular grids don't store explicit edges. """ - return FixedDimension(size=self.size, extent=new_extent) + return FixedDimension(size=self.size, extent=new_extent, shift=self.shift) @property def _size_repr(self) -> str: @@ -205,6 +214,10 @@ def data_size(self, chunk_ix: int) -> int: offset = self.cumulative[chunk_ix - 1] if chunk_ix > 0 else 0 return max(0, min(self.edges[chunk_ix], self.extent - offset)) + def buffer_offset(self, chunk_ix: int) -> int: + """Position of the valid data region inside the chunk's buffer.""" + return 0 + @property def _unique_edge_lengths(self) -> Iterable[int]: """Distinct chunk edge lengths for this dimension (lazily deduplicated). @@ -271,6 +284,7 @@ def index_to_chunk(self, idx: int) -> int: ... def chunk_offset(self, chunk_ix: int) -> int: ... def chunk_size(self, chunk_ix: int) -> int: ... def data_size(self, chunk_ix: int) -> int: ... + def buffer_offset(self, chunk_ix: int) -> int: ... def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: ... @property def _unique_edge_lengths(self) -> Iterable[int]: ... @@ -416,11 +430,27 @@ def from_metadata(cls, metadata: ArrayMetadata) -> ChunkGrid: For v3 metadata, dispatches on the chunk grid type. """ from zarr.core.metadata import ArrayV2Metadata - from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata + from zarr.core.metadata.v3 import ( + RectilinearChunkGridMetadata, + RegularBoundedChunkGridMetadata, + RegularChunkGridMetadata, + ) if isinstance(metadata, ArrayV2Metadata): return cls.from_sizes(metadata.shape, tuple(metadata.chunks)) chunk_grid_meta = metadata.chunk_grid + if isinstance(chunk_grid_meta, RegularBoundedChunkGridMetadata): + return cls( + dimensions=tuple( + FixedDimension(size=c, extent=e, shift=s) + for c, e, s in zip( + chunk_grid_meta.chunk_shape, + metadata.shape, + chunk_grid_meta.first_chunk_offset, + strict=True, + ) + ) + ) if isinstance(chunk_grid_meta, RegularChunkGridMetadata): return cls.from_sizes(metadata.shape, tuple(chunk_grid_meta.chunk_shape)) elif isinstance(chunk_grid_meta, RectilinearChunkGridMetadata): diff --git a/src/zarr/core/chunk_key_encodings.py b/src/zarr/core/chunk_key_encodings.py index c9649fd0db..3dff3a91ef 100644 --- a/src/zarr/core/chunk_key_encodings.py +++ b/src/zarr/core/chunk_key_encodings.py @@ -1,7 +1,6 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Iterable from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast @@ -24,10 +23,16 @@ def parse_separator(data: JSON) -> SeparatorLiteral: return cast("SeparatorLiteral", parse_field(data, Literal[".", "/"], "separator")) +def _parse_key_component(part: str) -> int: + # A chunk grid index can be negative; "-0" would name the same chunk as "0". + if part.startswith("-0"): + raise ValueError(f"Invalid chunk key component: {part!r}") + return int(part) + + class ChunkKeyEncodingParams(TypedDict): - name: Literal["v2", "default", "offset"] + name: Literal["v2", "default"] separator: NotRequired[SeparatorLiteral] - offset: NotRequired[tuple[int, ...]] @dataclass(frozen=True) @@ -83,7 +88,7 @@ def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: # Strip the "c" prefix (e.g. "c/" or "c.") before splitting. prefix = "c" + self.separator if chunk_key.startswith(prefix): - return tuple(map(int, chunk_key[len(prefix) :].split(self.separator))) + return tuple(map(_parse_key_component, chunk_key[len(prefix) :].split(self.separator))) raise ValueError(f"Invalid chunk key for default encoding: {chunk_key!r}") def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: @@ -100,71 +105,13 @@ def __post_init__(self) -> None: object.__setattr__(self, "separator", separator_parsed) def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: - return tuple(map(int, chunk_key.split(self.separator))) + return tuple(map(_parse_key_component, chunk_key.split(self.separator))) def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: chunk_identifier = self.separator.join(map(str, chunk_coords)) return "0" if chunk_identifier == "" else chunk_identifier -def parse_offset(data: object) -> tuple[int, ...]: - if isinstance(data, (str, bytes)) or not isinstance(data, Iterable): - raise TypeError(f"offset must be a sequence of integers, got {data!r}") - offset = tuple(data) - if any(isinstance(o, bool) or not isinstance(o, int) for o in offset): - raise TypeError(f"offset must be a sequence of integers, got {data!r}") - return offset - - -@dataclass(frozen=True) -class OffsetChunkKeyEncoding(ChunkKeyEncoding): - """ - The ``default`` encoding applied to ``chunk_coords + offset``, one integer per dimension. - - https://github.com/zarr-developers/zarr-extensions/tree/main/chunk-key-encodings/offset - """ - - name: ClassVar[Literal["offset"]] = "offset" - offset: tuple[int, ...] - separator: SeparatorLiteral = "/" - - def __post_init__(self) -> None: - object.__setattr__(self, "offset", parse_offset(self.offset)) - object.__setattr__(self, "separator", parse_separator(self.separator)) - - @classmethod - def from_dict(cls, data: dict[str, JSON]) -> Self: - _, config_parsed = parse_named_configuration(data, require_configuration=False) - if not config_parsed or "offset" not in config_parsed: - raise ValueError( - "The offset chunk key encoding requires an offset in its configuration." - ) - return cls(**cast("dict[str, Any]", config_parsed)) - - def to_dict(self) -> dict[str, JSON]: - return { - "name": self.name, - "configuration": {"separator": self.separator, "offset": list(self.offset)}, - } - - def _check_rank(self, chunk_coords: tuple[int, ...]) -> None: - if len(chunk_coords) != len(self.offset): - raise ValueError( - f"The offset has {len(self.offset)} dimension(s) but the chunk coordinates " - f"{chunk_coords!r} have {len(chunk_coords)}." - ) - - def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: - shifted = DefaultChunkKeyEncoding(separator=self.separator).decode_chunk_key(chunk_key) - self._check_rank(shifted) - return tuple(c - o for c, o in zip(shifted, self.offset, strict=True)) - - def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: - self._check_rank(chunk_coords) - shifted = tuple(c + o for c, o in zip(chunk_coords, self.offset, strict=True)) - return self.separator.join(map(str, ("c",) + shifted)) - - def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding: """ Take an implicit specification of a chunk key encoding and parse it into a ChunkKeyEncoding object. @@ -173,10 +120,8 @@ def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding: return data # handle ChunkKeyEncodingParams - if "name" in data and "configuration" not in data and ("separator" in data or "offset" in data): - params = cast("dict[str, JSON]", data) - configuration = {key: params[key] for key in ("separator", "offset") if key in params} - data = cast("dict[str, JSON]", {"name": params["name"], "configuration": configuration}) + if "name" in data and "separator" in data: + data = {"name": data["name"], "configuration": {"separator": data["separator"]}} # type: ignore[typeddict-item] # Now must be a named config data = cast("dict[str, JSON]", data) @@ -192,4 +137,3 @@ def parse_chunk_key_encoding(data: ChunkKeyEncodingLike) -> ChunkKeyEncoding: register_chunk_key_encoding("default", DefaultChunkKeyEncoding) register_chunk_key_encoding("v2", V2ChunkKeyEncoding) -register_chunk_key_encoding("offset", OffsetChunkKeyEncoding) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 8ea59eb960..24dc570b5a 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -380,7 +380,7 @@ def __iter__(self) -> Iterator[ChunkDimProjection]: g = self.dim_grid dim_chunk_ix = g.index_to_chunk(self.dim_sel) dim_offset = g.chunk_offset(dim_chunk_ix) - dim_chunk_sel = self.dim_sel - dim_offset + dim_chunk_sel = self.dim_sel - dim_offset + g.buffer_offset(dim_chunk_ix) dim_out_sel = None is_complete_chunk = g.data_size(dim_chunk_ix) == 1 yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) @@ -453,7 +453,6 @@ def __iter__(self) -> Iterator[ChunkDimProjection]: # selection ends within current chunk dim_chunk_sel_stop = self.stop - dim_offset - dim_chunk_sel = slice(dim_chunk_sel_start, dim_chunk_sel_stop, self.step) dim_chunk_nitems = ceildiv_int((dim_chunk_sel_stop - dim_chunk_sel_start), self.step) # If there are no elements on the selection within this chunk, then skip @@ -464,6 +463,10 @@ def __iter__(self) -> Iterator[ChunkDimProjection]: is_complete_chunk = ( dim_chunk_sel_start == 0 and (self.stop >= dim_limit) and self.step in [1, None] ) + buffer_offset = g.buffer_offset(dim_chunk_ix) + dim_chunk_sel = slice( + dim_chunk_sel_start + buffer_offset, dim_chunk_sel_stop + buffer_offset, self.step + ) yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) @@ -686,8 +689,9 @@ def __iter__(self) -> Iterator[ChunkDimProjection]: # pad out if boundary chunk (codec buffer may be larger than valid data region) codec_size = g.chunk_size(dim_chunk_ix) if dim_chunk_sel.shape[0] < codec_size: + buffer_offset = g.buffer_offset(dim_chunk_ix) tmp = np.zeros(codec_size, dtype=bool) - tmp[: dim_chunk_sel.shape[0]] = dim_chunk_sel + tmp[buffer_offset : buffer_offset + dim_chunk_sel.shape[0]] = dim_chunk_sel dim_chunk_sel = tmp # find region in output @@ -877,7 +881,7 @@ def __iter__(self) -> Iterator[ChunkDimProjection]: dim_out_sel = self.dim_out_sel[start:stop] # find region in chunk - dim_offset = g.chunk_offset(dim_chunk_ix) + dim_offset = g.chunk_offset(dim_chunk_ix) - g.buffer_offset(dim_chunk_ix) dim_chunk_sel = self.dim_sel[start:stop] - dim_offset is_complete_chunk = False # TODO yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) @@ -1111,7 +1115,7 @@ def __init__( ) start = dim_grid.chunk_offset(dim_sel) - stop = start + dim_grid.chunk_size(dim_sel) + stop = start + dim_grid.data_size(dim_sel) slice_ = slice(start, stop) elif is_slice(dim_sel): @@ -1398,7 +1402,7 @@ def __iter__(self) -> Iterator[ChunkProjection]: out_selection = self.sel_sort[start:stop] chunk_offsets = tuple( - g.chunk_offset(dim_chunk_ix) + g.chunk_offset(dim_chunk_ix) - g.buffer_offset(dim_chunk_ix) for dim_chunk_ix, g in zip(chunk_coords, self.dim_grids, strict=True) ) chunk_selection = tuple( diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 8bcc0424ee..cb957869fb 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -17,7 +17,6 @@ from zarr.core.chunk_key_encodings import ( ChunkKeyEncoding, ChunkKeyEncodingLike, - OffsetChunkKeyEncoding, parse_chunk_key_encoding, ) from zarr.core.common import ( @@ -232,6 +231,12 @@ class RegularChunkGridMetadataConfig(TypedDict): chunk_shape: Sequence[int] +class RegularBoundedChunkGridMetadataConfig(TypedDict): + chunk_shape: Sequence[int] + origin: Sequence[int] + lower_bound: Sequence[int] + + class RectilinearChunkGridMetadataConfig(TypedDict): kind: Literal["inline"] chunk_shapes: Sequence[RectilinearDimSpecJSON] @@ -240,6 +245,9 @@ class RectilinearChunkGridMetadataConfig(TypedDict): RegularChunkGridMetadataJSON = NamedRequiredConfig[ Literal["regular"], RegularChunkGridMetadataConfig ] +RegularBoundedChunkGridMetadataJSON = NamedRequiredConfig[ + Literal["regular_bounded"], RegularBoundedChunkGridMetadataConfig +] RectilinearChunkGridMetadataJSON = NamedRequiredConfig[ Literal["rectilinear"], RectilinearChunkGridMetadataConfig ] @@ -297,6 +305,73 @@ def from_dict(cls, data: RegularChunkGridMetadataJSON) -> Self: # type: ignore[ return cls(chunk_shape=parse_chunk_shape(configuration["chunk_shape"])) +def _parse_int_tuple(data: object, name: str, ndim: int) -> tuple[int, ...]: + if isinstance(data, (str, bytes)) or not isinstance(data, Sequence): + raise TypeError(f"{name} must be a sequence of integers, got {data!r}") + parsed = tuple(data) + if any(isinstance(v, bool) or not isinstance(v, int) for v in parsed): + raise TypeError(f"{name} must be a sequence of integers, got {data!r}") + if len(parsed) != ndim: + raise ValueError(f"{name} has {len(parsed)} element(s), expected {ndim}") + return parsed + + +@dataclass(frozen=True, kw_only=True) +class RegularBoundedChunkGridMetadata(RegularChunkGridMetadata): + """A regular chunk grid whose chunk 0 starts at array coordinate ``origin`` and whose + first element is at array coordinate ``lower_bound``. + + Chunk grid indices may be negative. Array indexing in zarr-python stays zero-based: + index ``i`` is coordinate ``lower_bound + i``. + """ + + origin: tuple[int, ...] + lower_bound: tuple[int, ...] + + def __post_init__(self) -> None: + super().__post_init__() + object.__setattr__(self, "origin", _parse_int_tuple(self.origin, "origin", self.ndim)) + object.__setattr__( + self, "lower_bound", _parse_int_tuple(self.lower_bound, "lower_bound", self.ndim) + ) + + @property + def first_chunk(self) -> tuple[int, ...]: + """Grid index of the chunk holding the first element, per dimension.""" + return tuple( + (lb - o) // c + for lb, o, c in zip(self.lower_bound, self.origin, self.chunk_shape, strict=True) + ) + + @property + def first_chunk_offset(self) -> tuple[int, ...]: + """Position of the first element inside its chunk, per dimension.""" + return tuple( + (lb - o) % c + for lb, o, c in zip(self.lower_bound, self.origin, self.chunk_shape, strict=True) + ) + + def to_dict(self) -> RegularBoundedChunkGridMetadataJSON: # type: ignore[override] + return { + "name": "regular_bounded", + "configuration": { + "chunk_shape": self.chunk_shape, + "origin": self.origin, + "lower_bound": self.lower_bound, + }, + } + + @classmethod + def from_dict(cls, data: RegularBoundedChunkGridMetadataJSON) -> Self: # type: ignore[override] + parse_named_configuration(data, "regular_bounded") + configuration = data["configuration"] + return cls( + chunk_shape=parse_chunk_shape(configuration["chunk_shape"]), + origin=configuration["origin"], # type: ignore[arg-type] + lower_bound=configuration["lower_bound"], # type: ignore[arg-type] + ) + + @dataclass(frozen=True, kw_only=True) class RectilinearChunkGridMetadata(Metadata): """Metadata-only description of a rectilinear chunk grid. @@ -462,6 +537,8 @@ def parse_chunk_grid( name, _ = parse_named_configuration(data) if name == "regular": return RegularChunkGridMetadata.from_dict(data) # type: ignore[arg-type] + if name == "regular_bounded": + return RegularBoundedChunkGridMetadata.from_dict(data) # type: ignore[arg-type] if name == "rectilinear": return RectilinearChunkGridMetadata.from_dict(data) # type: ignore[arg-type] raise ValueError(f"Unknown chunk grid name: {name!r}") @@ -545,13 +622,6 @@ def __init__( shape_parsed = parse_shapelike(shape) chunk_grid_parsed = parse_chunk_grid(chunk_grid) chunk_key_encoding_parsed = parse_chunk_key_encoding(chunk_key_encoding) - if isinstance(chunk_key_encoding_parsed, OffsetChunkKeyEncoding) and len( - chunk_key_encoding_parsed.offset - ) != len(shape_parsed): - raise ValueError( - f"The offset chunk key encoding has {len(chunk_key_encoding_parsed.offset)} " - f"offset(s) but the array has {len(shape_parsed)} dimension(s)." - ) dimension_names_parsed = parse_dimension_names(dimension_names) # Note: relying on a type method is numpy-specific fill_value_parsed = data_type.cast_scalar(fill_value) @@ -665,8 +735,20 @@ def inner_codecs(self) -> tuple[Codec, ...]: return self.codecs def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + if isinstance(self.chunk_grid, RegularBoundedChunkGridMetadata): + chunk_coords = tuple( + c + f for c, f in zip(chunk_coords, self.chunk_grid.first_chunk, strict=True) + ) return self.chunk_key_encoding.encode_chunk_key(chunk_coords) + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: + chunk_coords = self.chunk_key_encoding.decode_chunk_key(chunk_key) + if isinstance(self.chunk_grid, RegularBoundedChunkGridMetadata): + chunk_coords = tuple( + c - f for c, f in zip(chunk_coords, self.chunk_grid.first_chunk, strict=True) + ) + return chunk_coords + def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: check_storable(self) indent = config.get("json_indent") diff --git a/tests/test_array.py b/tests/test_array.py index 77ded6f1f9..389e2b2f60 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -694,54 +694,140 @@ def _chunk_bytes(store: Store) -> dict[str, bytes]: @pytest.mark.parametrize("store", ["memory"], indirect=True) +@pytest.mark.parametrize("extent", [3, 4, 9]) def test_prepend_writes_in_front_and_leaves_every_existing_chunk_untouched( - store: MemoryStore, + store: MemoryStore, extent: int ) -> None: - z = zarr.create_array( - store=store, - shape=(6, 4), - chunks=(3, 2), - dtype="int32", - fill_value=-1, - chunk_key_encoding={"name": "offset", "offset": (0, 0)}, - ) + z = zarr.create_array(store=store, shape=(6, 4), chunks=(4, 2), dtype="int32", fill_value=-1) data = np.arange(24, dtype="int32").reshape(6, 4) z[:] = data before = _chunk_bytes(store) assert set(before) == {"c/0/0", "c/0/1", "c/1/0", "c/1/1"} - assert z.prepend(np.full((6, 4), 7, dtype="int32"), axis=0) == (12, 4) + front = np.full((extent, 4), 7, dtype="int32") + assert z.prepend(front, axis=0) == (6 + extent, 4) after = _chunk_bytes(store) assert {k: after[k] for k in before} == before - assert set(after) - set(before) == {"c/-2/0", "c/-2/1", "c/-1/0", "c/-1/1"} - np.testing.assert_array_equal(z[:6], 7) - np.testing.assert_array_equal(z[6:], data) - assert z.metadata.to_dict()["chunk_key_encoding"] == { - "name": "offset", - "configuration": {"separator": "/", "offset": [-2, 0]}, + first = -(-extent // 4) + assert set(after) - set(before) == {f"c/{-g}/{d}" for g in range(1, first + 1) for d in (0, 1)} + np.testing.assert_array_equal(z[:extent], front) + np.testing.assert_array_equal(z[extent:], data) + assert z.domain == ((-extent, 6), (0, 4)) + assert z.metadata.to_dict()["chunk_grid"] == { + "name": "regular_bounded", + "configuration": {"chunk_shape": (4, 2), "origin": (0, 0), "lower_bound": (-extent, 0)}, } reopened = zarr.open_array(store=store, mode="r") - assert reopened.shape == (12, 4) + assert reopened.domain == z.domain np.testing.assert_array_equal(reopened[:], z[:]) @pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_requires_whole_chunks_and_matching_other_dimensions(store: MemoryStore) -> None: +def test_prepend_twice_and_resize_keep_coordinates_stable(store: MemoryStore) -> None: + z = zarr.create_array(store=store, shape=(5,), chunks=(4,), dtype="int16", fill_value=0) + z[:] = np.arange(5, dtype="int16") + z.prepend(np.array([-2, -1], dtype="int16")) + z.prepend(np.array([-5, -4, -3], dtype="int16")) + assert z.domain == ((-5, 5),) + np.testing.assert_array_equal(z[:], np.arange(-5, 5)) + assert _chunk_keys(store) == {"c/-2", "c/-1", "c/0", "c/1"} + + z.resize((7,)) + assert z.domain == ((-5, 2),) + np.testing.assert_array_equal(z[:], np.arange(-5, 2)) + assert _chunk_keys(store) == {"c/-2", "c/-1", "c/0"} + + +def _bounded_reference() -> tuple[zarr.Array[Any], npt.NDArray[Any]]: + z = zarr.create_array(store={}, shape=(7, 9), chunks=(3, 4), dtype="int32", fill_value=-1) + ref = np.arange(63, dtype="int32").reshape(7, 9) + z[:] = ref + z.prepend(np.full((2, 9), 100, dtype="int32"), axis=0) + z.prepend(np.full((9, 3), 200, dtype="int32"), axis=1) + ref = np.concatenate([np.full((2, 9), 100, dtype="int32"), ref], axis=0) + ref = np.concatenate([np.full((9, 3), 200, dtype="int32"), ref], axis=1) + assert z.metadata.chunk_grid.first_chunk_offset == (1, 1) # type: ignore[union-attr] + return z, ref + + +@pytest.mark.parametrize( + "selection", + [ + (0, 0), + (slice(None), slice(None)), + (slice(0, 2), slice(1, 10)), + (slice(None, None, 2), slice(1, None, 3)), + (slice(1, 8), 5), + ], +) +def test_basic_selection_matches_numpy_on_a_partial_first_chunk(selection: Any) -> None: + z, ref = _bounded_reference() + np.testing.assert_array_equal(z[selection], ref[selection]) + z[selection] = 9 + ref[selection] = 9 + np.testing.assert_array_equal(z[:], ref) + + +def test_advanced_selections_match_numpy_on_a_partial_first_chunk() -> None: + z, ref = _bounded_reference() + rows = np.array([0, 2, 3, 8]) + cols = np.array([11, 0, 3, 4]) + mask = np.zeros(ref.shape, dtype=bool) + mask[0, 0] = mask[2, 3] = mask[8, 11] = True + row_mask = np.zeros(ref.shape[0], dtype=bool) + row_mask[[0, 1, 5]] = True + + np.testing.assert_array_equal(z.get_orthogonal_selection((rows, cols)), ref[np.ix_(rows, cols)]) + np.testing.assert_array_equal( + z.get_orthogonal_selection((row_mask, slice(0, 4))), ref[row_mask, :4] + ) + np.testing.assert_array_equal(z.get_coordinate_selection((rows, cols)), ref[rows, cols]) + np.testing.assert_array_equal(z.get_mask_selection(mask), ref[mask]) + np.testing.assert_array_equal(z.get_block_selection((0, 0)), ref[:2, :3]) + np.testing.assert_array_equal(z.get_block_selection((slice(0, 2), 1)), ref[:5, 3:7]) + + z.set_orthogonal_selection((rows, cols), 1) + ref[np.ix_(rows, cols)] = 1 + z.set_coordinate_selection((np.array([0, 8]), np.array([0, 11])), 2) + ref[[0, 8], [0, 11]] = 2 + z.set_mask_selection(mask, 3) + ref[mask] = 3 + z.set_orthogonal_selection((row_mask, slice(0, 2)), 4) + ref[row_mask, :2] = 4 + z.set_block_selection((0, 0), 5) + ref[:2, :3] = 5 + np.testing.assert_array_equal(z[:], ref) + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_on_a_sharded_array_by_part_of_a_shard(store: MemoryStore) -> None: z = zarr.create_array( - store=store, - shape=(4, 3), - chunks=(2, 3), - dtype="uint8", - chunk_key_encoding={"name": "offset", "offset": (0, 0)}, + store=store, shape=(4, 2), chunks=(2, 2), shards=(4, 2), dtype="int16", fill_value=0 ) - with pytest.raises(ValueError, match="whole chunks"): - z.prepend(np.zeros((3, 3), dtype="uint8"), axis=0) + z[:] = 1 + before = _chunk_bytes(store) + assert set(before) == {"c/0/0"} + + z.prepend(np.full((3, 2), 3, dtype="int16"), axis=0) + + np.testing.assert_array_equal(z[:3], 3) + np.testing.assert_array_equal(z[3:], 1) + after = _chunk_bytes(store) + assert after["c/0/0"] == before["c/0/0"] + assert set(after) == {"c/-1/0", "c/0/0"} + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +def test_prepend_validates_its_input(store: MemoryStore) -> None: + z = zarr.create_array(store=store, shape=(4, 3), chunks=(2, 3), dtype="uint8") with pytest.raises(ValueError, match="not compatible"): z.prepend(np.zeros((2, 4), dtype="uint8"), axis=0) with pytest.raises(ValueError, match="out of range"): z.prepend(np.zeros((2, 3), dtype="uint8"), axis=2) - assert z.shape == (4, 3) + assert z.prepend([[4, 5, 6]]) == (5, 3) + np.testing.assert_array_equal(z[0], [4, 5, 6]) + assert z.prepend(np.zeros((0, 3), dtype="uint8")) == (5, 3) @pytest.mark.parametrize("store", ["memory"], indirect=True) @@ -751,87 +837,46 @@ def test_prepend_along_a_negative_axis_with_a_dot_separator(store: MemoryStore) shape=(2, 4), chunks=(2, 2), dtype="uint8", - chunk_key_encoding={"name": "offset", "separator": ".", "offset": (0, 0)}, + chunk_key_encoding={"name": "default", "separator": "."}, ) z[:] = 1 - assert z.prepend(np.full((2, 2), 5, dtype="uint8"), axis=-1) == (2, 6) - np.testing.assert_array_equal(z[:, :2], 5) np.testing.assert_array_equal(z[:, 2:], 1) assert _chunk_keys(store) == {"c.0.-1", "c.0.0", "c.0.1"} -@pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_on_a_sharded_array_grows_by_whole_shards(store: MemoryStore) -> None: - z = zarr.create_array( - store=store, - shape=(4, 2), - chunks=(2, 2), - shards=(4, 2), - dtype="int16", - chunk_key_encoding={"name": "offset", "offset": (0, 0)}, - ) - z[:] = 1 - assert _chunk_keys(store) == {"c/0/0"} - - assert z.prepend(np.full((4, 2), 3, dtype="int16"), axis=0) == (8, 2) - - np.testing.assert_array_equal(z[:4], 3) - np.testing.assert_array_equal(z[4:], 1) - assert _chunk_keys(store) == {"c/-1/0", "c/0/0"} - with pytest.raises(ValueError, match="whole shards"): - z.prepend(np.zeros((2, 2), dtype="int16"), axis=0) - - @pytest.mark.parametrize("store", ["memory"], indirect=True) @pytest.mark.parametrize("zarr_format", [2, 3]) -def test_prepend_refuses_an_array_without_the_offset_encoding( +def test_prepend_refuses_v2_and_rectilinear_arrays( store: MemoryStore, zarr_format: ZarrFormat ) -> None: z = zarr.create_array( store=store, shape=(4,), chunks=(2,), dtype="int8", zarr_format=zarr_format ) - with pytest.raises(TypeError, match="offset chunk key encoding"): - z.prepend(np.zeros(2, dtype="int8")) - assert z.shape == (4,) - - -@pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_prepend_accepts_a_plain_list_and_refuses_a_rectilinear_grid(store: MemoryStore) -> None: - z = zarr.create_array( - store=store, - shape=(2,), - chunks=(2,), - dtype="int8", - chunk_key_encoding={"name": "offset", "offset": (0,)}, - ) - z[:] = 1 - assert z.prepend([4, 5]) == (4,) - np.testing.assert_array_equal(z[:], [4, 5, 1, 1]) - + if zarr_format == 2: + with pytest.raises(TypeError, match="regular chunk grid"): + z.prepend(np.zeros(2, dtype="int8")) + assert z.shape == (4,) with zarr.config.set({"array.rectilinear_chunks": True}): rectilinear = zarr.create_array( - store=MemoryStore(), - shape=(4,), - chunks=[[2, 2]], - dtype="int8", - chunk_key_encoding={"name": "offset", "offset": (0,)}, + store=MemoryStore(), shape=(4,), chunks=[[2, 2]], dtype="int8" ) with pytest.raises(TypeError, match="regular chunk grid"): rectilinear.prepend(np.zeros(2, dtype="int8")) -@pytest.mark.parametrize("store", ["memory"], indirect=True) -def test_create_array_rejects_an_offset_of_the_wrong_rank(store: MemoryStore) -> None: - with pytest.raises(ValueError, match="2 dimension"): - zarr.create_array( - store=store, - shape=(4, 4), - chunks=(2, 2), - dtype="int8", - chunk_key_encoding={"name": "offset", "offset": (0,)}, - ) +def test_regular_bounded_grid_round_trips_and_validates() -> None: + from zarr.core.metadata.v3 import RegularBoundedChunkGridMetadata, parse_chunk_grid + + grid = RegularBoundedChunkGridMetadata(chunk_shape=(4, 5), origin=(-2, 0), lower_bound=(-7, 3)) + assert parse_chunk_grid(grid.to_dict()) == grid + assert grid.first_chunk == (-2, 0) + assert grid.first_chunk_offset == (3, 3) + with pytest.raises(ValueError, match="origin has 1"): + RegularBoundedChunkGridMetadata(chunk_shape=(4, 5), origin=(0,), lower_bound=(0, 0)) + with pytest.raises(TypeError, match="lower_bound must be a sequence of integers"): + RegularBoundedChunkGridMetadata(chunk_shape=(4,), origin=(0,), lower_bound=(1.5,)) # type: ignore[arg-type] @pytest.mark.parametrize("store", ["memory"], indirect=True) diff --git a/tests/test_chunk_key_encodings.py b/tests/test_chunk_key_encodings.py index 4b746aace0..80df5054b0 100644 --- a/tests/test_chunk_key_encodings.py +++ b/tests/test_chunk_key_encodings.py @@ -2,12 +2,7 @@ import pytest -from zarr.core.chunk_key_encodings import ( - DefaultChunkKeyEncoding, - OffsetChunkKeyEncoding, - V2ChunkKeyEncoding, - parse_chunk_key_encoding, -) +from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding, V2ChunkKeyEncoding @pytest.mark.parametrize("separator", ["/", "."]) @@ -71,75 +66,20 @@ def test_default_encoding_rejects_key_using_the_other_separator() -> None: @pytest.mark.parametrize("separator", ["/", "."]) -@pytest.mark.parametrize("coords", [(0,), (1, 2), (10, 0, 3)]) -def test_offset_encoding_with_zero_offset_is_the_default_encoding( - separator: str, coords: tuple[int, ...] -) -> None: - """An all-zero offset gives byte-for-byte the default encoding's keys, so an array can - switch to this encoding without any chunk moving.""" - default = DefaultChunkKeyEncoding(separator=separator) # type: ignore[arg-type] - offset = OffsetChunkKeyEncoding(separator=separator, offset=(0,) * len(coords)) # type: ignore[arg-type] - - assert offset.encode_chunk_key(coords) == default.encode_chunk_key(coords) - assert offset.decode_chunk_key(default.encode_chunk_key(coords)) == coords - - -@pytest.mark.parametrize("separator", ["/", "."]) -@pytest.mark.parametrize("coords", [(0, 0), (5, 4), (7, 0)]) -def test_offset_encoding_shifts_keys_and_round_trips( - separator: str, coords: tuple[int, ...] -) -> None: - encoding = OffsetChunkKeyEncoding(separator=separator, offset=(-2, 1)) # type: ignore[arg-type] - - key = encoding.encode_chunk_key(coords) - assert key == separator.join(["c", str(coords[0] - 2), str(coords[1] + 1)]) - assert encoding.decode_chunk_key(key) == coords - - -def test_offset_encoding_produces_negative_key_components() -> None: - """The chunk that holds array index 0 after two chunks were added in front of it is - stored under a negative key; the chunk that held index 0 before keeps `c/0/0`.""" - encoding = OffsetChunkKeyEncoding(offset=(-2, 0)) - - assert encoding.encode_chunk_key((0, 0)) == "c/-2/0" - assert encoding.encode_chunk_key((2, 0)) == "c/0/0" - assert encoding.decode_chunk_key("c/-2/0") == (0, 0) - - -def test_offset_encoding_rejects_a_rank_mismatch() -> None: - encoding = OffsetChunkKeyEncoding(offset=(0, 0)) - - with pytest.raises(ValueError, match="2 dimension"): - encoding.encode_chunk_key((1, 2, 3)) - with pytest.raises(ValueError, match="2 dimension"): - encoding.decode_chunk_key("c/1/2/3") - - -@pytest.mark.parametrize("offset", [(1.5,), ("1",), (True,), "12", 3]) -def test_offset_encoding_rejects_a_non_integer_offset(offset: object) -> None: - with pytest.raises(TypeError, match="sequence of integers"): - OffsetChunkKeyEncoding(offset=offset) # type: ignore[arg-type] - +def test_default_encoding_round_trips_negative_coordinates(separator: str) -> None: + encoding = DefaultChunkKeyEncoding(separator=separator) # type: ignore[arg-type] + key = encoding.encode_chunk_key((-3, 0, 12)) + assert key == separator.join(["c", "-3", "0", "12"]) + assert encoding.decode_chunk_key(key) == (-3, 0, 12) -def test_offset_encoding_metadata_round_trips() -> None: - encoding = OffsetChunkKeyEncoding(offset=(-1, 0), separator=".") - assert encoding.to_dict() == { - "name": "offset", - "configuration": {"separator": ".", "offset": [-1, 0]}, - } - assert parse_chunk_key_encoding(encoding.to_dict()) == encoding - assert ( - parse_chunk_key_encoding({"name": "offset", "offset": (-1, 0), "separator": "."}) - == encoding - ) - assert parse_chunk_key_encoding({"name": "offset", "offset": (3,)}) == OffsetChunkKeyEncoding( - offset=(3,) - ) +@pytest.mark.parametrize("chunk_key", ["c/-0", "c/1/-0", "c/-01"]) +def test_default_encoding_rejects_negative_zero(chunk_key: str) -> None: + """`-0` would name the same chunk as `0`, so it is not a valid key component.""" + with pytest.raises(ValueError, match="Invalid chunk key component"): + DefaultChunkKeyEncoding(separator="/").decode_chunk_key(chunk_key) -def test_offset_encoding_requires_an_offset() -> None: - with pytest.raises(ValueError, match="requires an offset"): - parse_chunk_key_encoding({"name": "offset"}) - with pytest.raises(ValueError, match="requires an offset"): - parse_chunk_key_encoding({"name": "offset", "configuration": {"separator": "/"}}) +def test_v2_encoding_rejects_negative_zero() -> None: + with pytest.raises(ValueError, match="Invalid chunk key component"): + V2ChunkKeyEncoding(separator=".").decode_chunk_key("1.-0")