diff --git a/changes/4516.feature.md b/changes/4516.feature.md new file mode 100644 index 0000000000..985f10f83b --- /dev/null +++ b/changes/4516.feature.md @@ -0,0 +1,4 @@ +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 e8df79a102..f96d872f4b 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -154,6 +154,26 @@ 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, 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=(4, 4), overwrite=True) +z[:] = 1 +z.prepend(np.full((3, 4), 2, dtype='int32')) +print(f"Shape: {z.shape}, domain: {z.domain}") +``` + ## 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..252fa0062c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -138,6 +138,7 @@ from zarr.core.metadata.v3 import ( ChunkGridMetadata, RectilinearChunkGridMetadata, + RegularBoundedChunkGridMetadata, RegularChunkGridMetadata, create_chunk_grid_metadata, parse_node_type_array, @@ -949,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. @@ -1801,6 +1816,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`. 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) + async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ Asynchronously update the array's attributes. @@ -2194,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. @@ -4025,6 +4072,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`. 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)) + def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ Update the array's attributes. @@ -5565,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)) @@ -6066,6 +6136,77 @@ 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 be a Zarr v3 array with a regular chunk grid. + data : npt.ArrayLike + Data to be prepended. + 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_grid, RegularChunkGridMetadata + ): + 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.") + 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] + 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 + 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() + 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], 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 d871e279d2..3dff3a91ef 100644 --- a/src/zarr/core/chunk_key_encodings.py +++ b/src/zarr/core/chunk_key_encodings.py @@ -23,6 +23,13 @@ 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"] separator: NotRequired[SeparatorLiteral] @@ -81,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: @@ -98,7 +105,7 @@ 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)) 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 2d5d169e6b..cb957869fb 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -231,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] @@ -239,6 +245,9 @@ class RectilinearChunkGridMetadataConfig(TypedDict): RegularChunkGridMetadataJSON = NamedRequiredConfig[ Literal["regular"], RegularChunkGridMetadataConfig ] +RegularBoundedChunkGridMetadataJSON = NamedRequiredConfig[ + Literal["regular_bounded"], RegularBoundedChunkGridMetadataConfig +] RectilinearChunkGridMetadataJSON = NamedRequiredConfig[ Literal["rectilinear"], RectilinearChunkGridMetadataConfig ] @@ -296,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. @@ -461,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}") @@ -657,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 3e4835e686..389e2b2f60 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,208 @@ 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] + + +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) +@pytest.mark.parametrize("extent", [3, 4, 9]) +def test_prepend_writes_in_front_and_leaves_every_existing_chunk_untouched( + store: MemoryStore, extent: int +) -> None: + 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"} + + 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 + 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.domain == z.domain + np.testing.assert_array_equal(reopened[:], z[:]) + + +@pytest.mark.parametrize("store", ["memory"], indirect=True) +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, 2), chunks=(2, 2), shards=(4, 2), dtype="int16", fill_value=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.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) +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": "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) +@pytest.mark.parametrize("zarr_format", [2, 3]) +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 + ) + 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" + ) + with pytest.raises(TypeError, match="regular chunk grid"): + rectilinear.prepend(np.zeros(2, dtype="int8")) + + +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) 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..80df5054b0 100644 --- a/tests/test_chunk_key_encodings.py +++ b/tests/test_chunk_key_encodings.py @@ -63,3 +63,23 @@ 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", ["/", "."]) +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) + + +@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_v2_encoding_rejects_negative_zero() -> None: + with pytest.raises(ValueError, match="Invalid chunk key component"): + V2ChunkKeyEncoding(separator=".").decode_chunk_key("1.-0")