Add the offset chunk key encoding and Array.prepend

changes/4516.feature.md

@@ -0,0 +1,3 @@
+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.

docs/user-guide/arrays.md

@@ -154,6 +154,22 @@ z.append(np.vstack([a, a]), axis=1)
 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:
+
+```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
+z.prepend(np.full((6, 4), 2, dtype='int32'))
+print(f"Shape after prepend: {z.shape}, offset: {z.metadata.chunk_key_encoding.offset}")
+```
+
 ## Runtime configuration
 
 Zarr arrays are parametrized with a configuration that determines certain aspects of array behavior.

src/zarr/core/array.py

@@ -54,6 +54,7 @@
     ChunkKeyEncoding,
     ChunkKeyEncodingLike,
     DefaultChunkKeyEncoding,
+    OffsetChunkKeyEncoding,
     V2ChunkKeyEncoding,
     parse_chunk_key_encoding,
 )
@@ -1801,6 +1802,29 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]:
         """
         return await _append(self, data, axis)
 
+    async def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]:
+        """Prepend `data` to `axis`.
+
+        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`, 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)
+
     async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
         """
         Asynchronously update the array's attributes.
@@ -4025,6 +4049,29 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]:
         """
         return sync(self.async_array.append(data, axis=axis))
 
+    def prepend(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]:
+        """Prepend `data` to `axis`.
+
+        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`, 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))
+
     def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
         """
         Update the array's attributes.
@@ -6066,6 +6113,86 @@ async def _append(
     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 size 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)
+
+    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)
+    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]
+    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 = 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)
+
+    return new_shape
+
+
 async def _update_attributes(
     array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata],
     new_attributes: dict[str, JSON],

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,64 @@ 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``, 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.
@@ -113,8 +173,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 +192,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)

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)