diff --git a/changes/4174.bugfix.md b/changes/4174.bugfix.md new file mode 100644 index 0000000000..2e4fcb2eec --- /dev/null +++ b/changes/4174.bugfix.md @@ -0,0 +1,11 @@ +Array creation is now O(1) in the number of chunks per dimension. Chunk +normalization returns a `ChunkGrid` whose uniform dimensions are stored as a +size + extent pair (`FixedDimension`) instead of being expanded to one entry +per chunk, so creating arrays like +`zarr.create_array(store, shape=(2**62,), chunks=(1,), dtype='int32')` succeeds +instantly instead of raising `ValueError` or allocating gigabytes of memory. +The intermediate `ChunksTuple` representation was removed in the process, and +`ChunksLike` now admits per-dimension specs that mix a bare int (uniform chunk +size) with explicit edge-length sequences, matching what the normalizer and +the rectilinear grid spec already accepted. +This is the creation-time counterpart of the indexing fix in #4172. diff --git a/changes/4272.bugfix.md b/changes/4272.bugfix.md new file mode 100644 index 0000000000..417128d092 --- /dev/null +++ b/changes/4272.bugfix.md @@ -0,0 +1,22 @@ +Explicit per-chunk size lists now always produce a rectilinear chunk grid, +even when the sizes happen to describe a regular grid (all equal, or all equal +with a smaller trailing chunk). Previously such input was silently collapsed to +a regular grid, which changed resize semantics: a regular grid grows by +extending the uniform pattern, while a rectilinear grid appends a new edge +chunk — the behavior an append-oriented layout like `(168,) * 13 + (24,)` +relies on. The grid kind now follows the input syntax, matching 3.2.x: +scalar chunk sizes (including numpy integers and the `-1` sentinel) produce a +regular grid, nested sequences produce a rectilinear grid. Rectilinear grids +remain gated behind `zarr.config.set({"array.rectilinear_chunks": True})`. +See #4174 for the accompanying O(1) chunk normalization change. + +`zarr.from_array` with the default `chunks="keep"` / `shards="keep"` now +reproduces the source's stored grid exactly: a rectilinear grid is passed +through in O(number of dimensions), with uniform dimensions keeping their +bare-int shorthand; sharding under a rectilinear shard grid is preserved +instead of being silently dropped; and the default `write_data=True` copy +works for every grid kind. `Array.chunks` is now defined for any sharded array +(the inner chunks of a shard are always regular), and for sharded arrays with +a rectilinear shard grid `Array.info` no longer raises — it reports the shard +shape as `` — while `Array.nchunks_initialized` counts the chunks of +each initialized shard individually instead of raising. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index a192845f9e..6e31b2e9f9 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -674,8 +674,11 @@ z_regular = zarr.create_array( print(z_regular.write_chunk_sizes) ``` -Note that the `.chunks` property is only available for regular chunk grids. For -rectilinear arrays, use `.write_chunk_sizes` (or `.read_chunk_sizes`) instead. +Note that the `.chunks` property is not available for non-sharded rectilinear +arrays, since there is no single uniform chunk shape — use `.write_chunk_sizes` +(or `.read_chunk_sizes`) instead. Sharded arrays always have `.chunks`: it +returns the inner chunk shape, which is regular even when the shard grid is +rectilinear (see [Rectilinear shard boundaries](#rectilinear-shard-boundaries)). ### Resizing and appending @@ -746,6 +749,17 @@ print(z[50:70, 40:60]) Note that rectilinear inner chunks with sharding are not supported — only the shard boundaries can be rectilinear. +For such arrays, `.chunks` returns the (regular) inner chunk shape, while +`.shards` raises `NotImplementedError` since there is no single uniform shard +shape — use `.write_chunk_sizes` for the per-dimension shard sizes. `.info` +reports the shard shape as ``: + +```python exec="true" session="arrays" source="above" result="ansi" +print(f"chunks={z.chunks}") +print(f"shard sizes={z.write_chunk_sizes}") +print(z.info) +``` + ### Metadata format Rectilinear chunk grid metadata uses run-length encoding (RLE) for compact diff --git a/src/zarr/core/_info.py b/src/zarr/core/_info.py index 1503f05b26..58dcbdefb3 100644 --- a/src/zarr/core/_info.py +++ b/src/zarr/core/_info.py @@ -82,7 +82,9 @@ class ArrayInfo: _data_type: ZDType[TBaseDType, TBaseScalar] _fill_value: object _shape: tuple[int, ...] - _shard_shape: tuple[int, ...] | None = None + # "" marks a sharded array whose shard grid is rectilinear, so + # there is no uniform shard shape; None means the array is not sharded. + _shard_shape: tuple[int, ...] | Literal[""] | None = None _chunk_shape: tuple[int, ...] | None = None _order: Literal["C", "F"] _read_only: bool diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 903f7b8f26..80b356ad81 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -43,7 +43,6 @@ SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, _is_rectilinear_chunks, - as_regular_shape, guess_chunks, normalize_chunks_nd, resolve_outer_and_inner_chunks, @@ -101,7 +100,6 @@ Selection, VIndex, _iter_grid, - _iter_regions, check_fields, check_no_multi_fields, is_pure_fancy_indexing, @@ -126,6 +124,7 @@ ) from zarr.core.metadata.v3 import ( ChunkGridMetadata, + RectilinearChunkGridMetadata, create_chunk_grid_metadata, parse_node_type_array, ) @@ -154,7 +153,7 @@ from zarr.abc.codec import CodecPipeline from zarr.abc.store import Store - from zarr.codecs.sharding import IndexLocation + from zarr.codecs.sharding import IndexLocation, ShardingCodec from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -523,7 +522,7 @@ async def _create( outer_chunks = guess_chunks(shape, item_size) else: outer_chunks = normalize_chunks_nd(_raw, shape) - _chunks = as_regular_shape(outer_chunks) + _chunks = outer_chunks.chunk_shape if order is None: order_parsed = config_parsed.order @@ -848,10 +847,12 @@ def shape(self) -> tuple[int, ...]: @property def chunks(self) -> tuple[int, ...]: """Returns the chunk shape of the Array. - If sharding is used the inner chunk shape is returned. + If sharding is used the inner chunk shape is returned, which is defined + for any chunk grid (inner chunks are always regular). - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Otherwise, only defined for arrays using a regular chunk grid: for a + non-sharded array with a rectilinear chunk grid, `NotImplementedError` + is raised. Use `read_chunk_sizes` for general access. Returns ------- @@ -923,8 +924,9 @@ def shards(self) -> tuple[int, ...] | None: """Returns the shard shape of the Array. Returns None if sharding is not used. - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Only defined when the shard grid is regular: for a sharded array with a + rectilinear chunk grid, `NotImplementedError` is raised. Use + `write_chunk_sizes` for general access. Returns ------- @@ -1121,14 +1123,9 @@ def _chunk_grid_shape(self) -> tuple[int, ...]: tuple[int, ...] The number of chunks along each dimension. """ - # TODO: refactor — extract a sharding_codec property on ArrayV3Metadata - # to replace the repeated `len == 1 and isinstance` pattern. - from zarr.codecs.sharding import ShardingCodec - - codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) - if len(codecs) == 1 and isinstance(codecs[0], ShardingCodec): + if (sharding_codec := _sharding_codec(self.metadata)) is not None: # When sharding, count inner chunks across the whole array - chunk_shape = codecs[0].chunk_shape + chunk_shape = sharding_codec.chunk_shape return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True))) return self._chunk_grid.grid_shape @@ -1144,11 +1141,9 @@ def _shard_grid_shape(self) -> tuple[int, ...]: tuple[int, ...] The shape of the shard grid for this array. """ - if self.shards is None: - shard_shape = self.chunks - else: - shard_shape = self.shards - return tuple(starmap(ceildiv, zip(self.shape, shard_shape, strict=True))) + # The stored chunk grid is the shard grid when sharding is used, the + # chunk grid otherwise. Works for regular and rectilinear grids alike. + return self._chunk_grid.grid_shape @property def nchunks(self) -> int: @@ -1212,10 +1207,10 @@ async def nchunks_initialized(self) -> int: """ Calculate the number of chunks that have been initialized in storage. - This value is calculated as the product of the number of initialized shards and the number - of chunks per shard. For arrays that do not use sharding, the number of chunks per shard is - effectively 1, and in that case the number of chunks initialized is the same as the number - of stored objects associated with an array. + This value is calculated as the sum of the number of chunks in every initialized shard + (shard sizes can vary when the shard grid is rectilinear). For arrays that do not use + sharding, each stored object holds one chunk, so the number of chunks initialized is the + same as the number of stored objects associated with an array. Returns ------- @@ -1822,14 +1817,27 @@ async def info_complete(self) -> Any: def _info( self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None ) -> Any: - chunk_shape = self.chunks if self._chunk_grid.is_regular else None + rectilinear_grid = _stored_rectilinear_grid(self.metadata) + sharded = _sharding_codec(self.metadata) is not None + # `.chunks` (the inner chunk shape when sharded) is undefined only for a + # non-sharded rectilinear grid, which ArrayInfo renders as ""; + # `.shards` is undefined for a rectilinear shard grid, where the + # "" sentinel keeps the array rendered as sharded. + chunk_shape = self.chunks if (rectilinear_grid is None or sharded) else None + shard_shape: tuple[int, ...] | Literal[""] | None + if rectilinear_grid is None: + shard_shape = self.shards + elif sharded: + shard_shape = "" + else: + shard_shape = None return ArrayInfo( _zarr_format=self.metadata.zarr_format, _data_type=self._zdtype, _fill_value=self.metadata.fill_value, _shape=self.shape, _order=self.order, - _shard_shape=self.shards, + _shard_shape=shard_shape, _chunk_shape=chunk_shape, _read_only=self.read_only, _compressors=self.compressors, @@ -2023,10 +2031,12 @@ def shape(self, value: tuple[int, ...]) -> None: @property def chunks(self) -> tuple[int, ...]: """Returns a tuple of integers describing the length of each dimension of a chunk of the array. - If sharding is used the inner chunk shape is returned. + If sharding is used the inner chunk shape is returned, which is defined + for any chunk grid (inner chunks are always regular). - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Otherwise, only defined for arrays using a regular chunk grid: for a + non-sharded array with a rectilinear chunk grid, `NotImplementedError` + is raised. Use `read_chunk_sizes` for general access. Returns ------- @@ -2090,8 +2100,9 @@ def shards(self) -> tuple[int, ...] | None: """Returns a tuple of integers describing the length of each dimension of a shard of the array. Returns None if sharding is not used. - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Only defined when the shard grid is regular: for a sharded array with a + rectilinear chunk grid, `NotImplementedError` is raised. Use + `write_chunk_sizes` for general access. Returns ------- @@ -2297,10 +2308,10 @@ def nchunks_initialized(self) -> int: """ Calculate the number of chunks that have been initialized in storage. - This value is calculated as the product of the number of initialized shards and the number of - chunks per shard. For arrays that do not use sharding, the number of chunks per shard is effectively 1, - and in that case the number of chunks initialized is the same as the number of stored objects associated with an - array. For a direct count of the number of initialized stored objects, see `nshards_initialized`. + This value is calculated as the sum of the number of chunks in every initialized shard + (shard sizes can vary when the shard grid is rectilinear). For arrays that do not use sharding, + each stored object holds one chunk, so the number of chunks initialized is the same as the number + of stored objects associated with an array. For a direct count of the number of initialized stored objects, see `nshards_initialized`. Returns ------- @@ -4062,7 +4073,9 @@ class ShardsConfigParam(TypedDict): index_location: IndexLocation | None -type ShardsLike = tuple[int, ...] | Sequence[Sequence[int]] | ShardsConfigParam | Literal["auto"] +type ShardsLike = ( + tuple[int, ...] | Sequence[int | Sequence[int]] | ShardsConfigParam | Literal["auto"] +) async def from_array( @@ -4472,7 +4485,7 @@ async def init_array( "chunks=(inner_size, ...), shards=[[shard_sizes], ...]" ) - # Normalize the user's chunks into canonical ChunksTuple form + # Normalize the user's chunks into a canonical ChunkGrid if chunks == "auto": max_bytes = None if shards is None else SHARDED_INNER_CHUNK_MAX_BYTES @@ -4515,7 +4528,7 @@ async def init_array( meta = AsyncArray._create_metadata_v2( shape=shape_parsed, dtype=zdtype, - chunks=as_regular_shape(outer_chunks), + chunks=outer_chunks.chunk_shape, dimension_separator=chunk_key_encoding_parsed.separator, fill_value=fill_value, order=order_parsed, @@ -4534,7 +4547,7 @@ async def init_array( grid = create_chunk_grid_metadata(outer_chunks) codecs_out: tuple[Codec, ...] if inner is not None: - inner_chunks_flat = as_regular_shape(inner.outer_chunks) + inner_chunks_flat = inner.outer_chunks.chunk_shape index_location: IndexLocation = "end" if isinstance(shards, dict): index_location = cast("IndexLocation", shards.get("index_location", "end")) @@ -4765,6 +4778,33 @@ async def create_array( ) +def _sharding_codec(metadata: ArrayMetadata) -> ShardingCodec | None: + """The array's sharding codec, or None if the array is not sharded. + + Zarr format 2 arrays are never sharded. + """ + if isinstance(metadata, ArrayV3Metadata): + return metadata.sharding_codec + return None + + +def _stored_rectilinear_grid(metadata: ArrayMetadata) -> RectilinearChunkGridMetadata | None: + """The *stored* rectilinear chunk grid, or None if the stored grid is regular + (in which case `.chunks` and `.shards` are defined). + + Dispatches on the stored metadata, not the runtime ``ChunkGrid``: the + runtime grid collapses a rectilinear dimension whose edges happen to be + uniform to a ``FixedDimension`` as an optimization, so it can report regular + for an array whose stored metadata — and therefore `.chunks` — is + rectilinear. Zarr format 2 grids are always regular. + """ + if isinstance(metadata, ArrayV3Metadata) and isinstance( + metadata.chunk_grid, RectilinearChunkGridMetadata + ): + return metadata.chunk_grid + return None + + def _parse_keep_array_attr( data: AnyArray | npt.ArrayLike, chunks: ChunksLike | Literal["auto", "keep"], @@ -4792,13 +4832,27 @@ def _parse_keep_array_attr( dict[str, JSON] | None, ]: if isinstance(data, Array): + rectilinear_grid = _stored_rectilinear_grid(data.metadata) + sharded = _sharding_codec(data.metadata) is not None if chunks == "keep": - if data._chunk_grid.is_regular: + if rectilinear_grid is None or sharded: + # `.chunks` is the inner chunk shape when sharding is used, and + # inner chunks are regular whatever the shape of the shard grid. chunks = data.chunks else: - chunks = data.write_chunk_sizes + # Pass the stored spec through as-is: it is O(ndim), not + # O(nchunks), and keeps the bare-int shorthand of uniform + # dimensions (which also preserves resize semantics). + chunks = rectilinear_grid.chunk_shapes if shards == "keep": - shards = data.shards if data._chunk_grid.is_regular else None + if rectilinear_grid is None: + shards = data.shards + elif sharded: + # The stored grid is the shard grid; its spec round-trips + # through the `shards=` parameter of array creation. + shards = rectilinear_grid.chunk_shapes + else: + shards = None if zarr_format is None: zarr_format = data.metadata.zarr_format if filters == "keep": @@ -5271,14 +5325,10 @@ def _iter_shard_regions( A tuple of slice objects representing the region spanned by each shard in the selection or chunk when no shards are present. """ - if array.shards is None: - shard_shape = array.chunks - else: - shard_shape = array.shards - - return _iter_regions( - array.shape, shard_shape, origin=origin, selection_shape=selection_shape, trim_excess=True - ) + # The stored chunk grid always describes the write regions: the shard grid + # when sharding is used, the chunk grid otherwise. Iterating it directly + # works for regular and rectilinear grids alike. + return array._chunk_grid.iter_chunk_regions(origin=origin, selection_shape=selection_shape) def _iter_chunk_regions( @@ -5316,10 +5366,10 @@ async def _nchunks_initialized( """ Calculate the number of chunks that have been initialized in storage. - This value is calculated as the product of the number of initialized shards and the number - of chunks per shard. For arrays that do not use sharding, the number of chunks per shard is - effectively 1, and in that case the number of chunks initialized is the same as the number - of stored objects associated with an array. + This value is calculated as the sum of the number of chunks in every initialized shard + (shard sizes can vary when the shard grid is rectilinear). For arrays that do not use + sharding, each stored object holds one chunk, so the number of chunks initialized is the + same as the number of stored objects associated with an array. Parameters ---------- @@ -5331,13 +5381,22 @@ async def _nchunks_initialized( nchunks_initialized : int The number of chunks that have been initialized. """ - if array.shards is None: - chunks_per_shard = 1 - else: - chunks_per_shard = product( - tuple(a // b for a, b in zip(array.shards, array.chunks, strict=True)) - ) - return (await _nshards_initialized(array)) * chunks_per_shard + if _sharding_codec(array.metadata) is None: + return await _nshards_initialized(array) + # Count the inner chunks of each initialized shard individually: shard + # sizes vary when the shard grid is rectilinear. `.chunks` is the inner + # chunk shape, which is defined for any sharded array. + inner_chunks = array.chunks + grid = array._chunk_grid + initialized = set(await _shards_initialized(array)) + total = 0 + for coords in grid.all_chunk_coords(): + spec = grid[coords] + if spec is not None and array.metadata.encode_chunk_key(coords) in initialized: + total += product( + tuple(s // c for s, c in zip(spec.codec_shape, inner_chunks, strict=True)) + ) + return total async def _nshards_initialized( diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 584829bc6c..f23011cbdb 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -12,7 +12,6 @@ TYPE_CHECKING, Any, NamedTuple, - NewType, Protocol, TypeGuard, cast, @@ -43,31 +42,6 @@ is not `None`. Explicit chunk sizes are not affected by this value. """ -ChunksTuple = NewType("ChunksTuple", tuple[np.ndarray[tuple[int], np.dtype[np.int64]], ...]) -"""Normalized chunk specification: one 1D int64 array of chunk sizes per dimension. - -Produced exclusively by `normalize_chunks_nd` and `guess_chunks`. -Consumers should use this type to ensure they receive validated, -canonical chunk specifications rather than raw user input. -""" - - -class ChunkLayout(NamedTuple): - """Result of resolving user `chunks`/`shards` into grid metadata inputs. - - outer_chunks - Chunk sizes for the chunk grid metadata. When sharding is active - these are the shard sizes; otherwise they are the user's chunk sizes. - inner - Recursive sub-structure inside each chunk. `None` means the chunk is - opaque (no sharding). When present, `inner.outer_chunks` gives the - sub-chunk sizes passed to `ShardingCodec`, and `inner.inner` gives - the next level of nesting (for nested sharding), or `None`. - """ - - outer_chunks: ChunksTuple - inner: ChunkLayout | None = None - @dataclass(frozen=True) class FixedDimension: @@ -360,9 +334,7 @@ def _is_rectilinear_chunks(chunks: Any) -> TypeGuard[Sequence[Sequence[int]]]: return False -def is_regular_1d( - dim_chunks: Sequence[int] | np.ndarray[tuple[int], np.dtype[np.int64]], -) -> bool: +def is_regular_1d(dim_chunks: Sequence[int]) -> bool: """Check if a single dimension's chunk sizes represent a regular grid. A regular dimension has either all chunks the same size, or all @@ -372,9 +344,6 @@ def is_regular_1d( if len(dim_chunks) <= 1: return True first = dim_chunks[0] - if isinstance(dim_chunks, np.ndarray): - # Vectorized comparison avoids per-element Python iteration over int64 arrays. - return bool((dim_chunks[1:-1] == first).all() and dim_chunks[-1] <= first) for c in dim_chunks[1:-1]: if c != first: return False @@ -382,19 +351,11 @@ def is_regular_1d( return dim_chunks[-1] <= first -def is_regular_nd( - chunks: Iterable[Sequence[int] | np.ndarray[tuple[int], np.dtype[np.int64]]], -) -> bool: +def is_regular_nd(chunks: Iterable[Sequence[int]]) -> bool: """Check if an N-dimensional chunk specification represents a regular grid.""" return all(is_regular_1d(d) for d in chunks) -def as_regular_shape(chunks: ChunksTuple) -> tuple[int, ...]: - """Flatten a regular ChunksTuple to one int per dimension.""" - assert is_regular_nd(chunks), f"expected regular chunks, got {chunks}" - return tuple(int(dim[0]) for dim in chunks) - - @dataclass(frozen=True) class ChunkGrid: """ @@ -489,6 +450,11 @@ def from_sizes( # -- Properties -- + @property + def dimensions(self) -> tuple[DimensionGrid, ...]: + """The per-dimension grids (`FixedDimension` or `VaryingDimension`).""" + return self._dimensions + @property def ndim(self) -> int: return len(self._dimensions) @@ -641,6 +607,23 @@ def update_shape(self, new_shape: tuple[int, ...]) -> ChunkGrid: return ChunkGrid(dimensions=dims) +class ChunkLayout(NamedTuple): + """Result of resolving user `chunks`/`shards` into grid metadata inputs. + + outer_chunks + Chunk grid for the chunk grid metadata. When sharding is active + this holds the shard sizes; otherwise it holds the user's chunk sizes. + inner + Recursive sub-structure inside each chunk. `None` means the chunk is + opaque (no sharding). When present, `inner.outer_chunks` gives the + sub-chunk sizes passed to `ShardingCodec`, and `inner.inner` gives + the next level of nesting (for nested sharding), or `None`. + """ + + outer_chunks: ChunkGrid + inner: ChunkLayout | None = None + + def _guess_regular_chunks( shape: tuple[int, ...] | int, typesize: int, @@ -717,30 +700,33 @@ def _guess_regular_chunks( return tuple(int(x) for x in chunks) -def normalize_chunks_1d( - chunks: int | Iterable[object], span: int -) -> np.ndarray[tuple[int], np.dtype[np.int64]]: +def normalize_chunks_1d(chunks: int | Iterable[object], span: int) -> DimensionGrid: """ - Normalize a one-dimensional chunk specification into a 1D int64 array of - chunk sizes that cover the span. + Normalize a one-dimensional chunk specification into a dimension grid: + `FixedDimension` for scalar chunk sizes, `VaryingDimension` for explicit + per-chunk size lists. Both variants bind the chunk sizes to + the span, and the uniform form is O(1) in the number of chunks — a + dimension with `2**62` chunks must not materialize one entry per chunk. `-1` means "one chunk covering the entire span." - For an integer chunk size, all chunks are uniform — the last chunk may - overhang the span. The actual data extent of each chunk is determined - by the chunk grid at runtime, not by this function. + Explicit chunk size lists must sum to the span exactly and always produce + `VaryingDimension`, even when the sizes happen to be uniform: the input + syntax declares the grid kind, so a per-chunk list is preserved as a + rectilinear dimension rather than silently collapsed to a regular one, + which would change how the dimension grows on resize. For scalar sizes + the last chunk may overhang the span. """ # `numbers.Integral` rather than `int` so that numpy integer scalars (which are not # `int` subclasses) take the uniform-chunk path instead of being treated as a sequence. + # The `-1` sentinel check lives inside this branch so that numpy-array chunk + # specifications never hit an ambiguous-truth-value error on `chunks == -1`. if isinstance(chunks, numbers.Integral): chunk_size = int(chunks) if chunk_size == -1: - return np.array([span], dtype=np.int64) + return FixedDimension(size=span, extent=span) if chunk_size <= 0: raise ValueError(f"Chunk size must be positive, got {chunk_size}") - if span == 0: - return np.array([chunk_size], dtype=np.int64) - n = ceildiv(span, chunk_size) - return np.full(n, chunk_size, dtype=np.int64) + return FixedDimension(size=chunk_size, extent=span) else: try: chunk_list = list(chunks) # type: ignore[arg-type] @@ -766,23 +752,28 @@ def normalize_chunks_1d( raise ValueError(f"All chunk sizes must be positive, got {ints}") if sum(ints) != span: raise ValueError(f"Chunk sizes {ints} do not sum to span {span}") - return np.asarray(ints, dtype=np.int64) + return VaryingDimension(ints, extent=span) def normalize_chunks_nd( chunks: Any, shape: tuple[int, ...], -) -> ChunksTuple: +) -> ChunkGrid: """ - Normalize a chunk specification into a `ChunksTuple`. + Normalize a chunk specification into a `ChunkGrid`. This is a mechanical transformation — no heuristics, no guessing. Handles `False` ("all data in one chunk"), scalar ints, `-1` sentinels (one chunk per dimension covering the full span), and explicit per-dimension lists of chunk sizes (regular or rectilinear). + This is the strict parser for user-supplied chunk specifications; use + `ChunkGrid.from_sizes` / `ChunkGrid.from_metadata` for stored metadata, + which is validated under more tolerant rules (e.g. trailing edges beyond + the array extent). + For auto-chunking, use `guess_chunks` which returns a - `ChunksTuple` directly. `chunks=None` and `chunks=True` are rejected + `ChunkGrid` directly. `chunks=None` and `chunks=True` are rejected here — the caller is responsible for choosing between explicit sizes and auto-chunking. """ @@ -793,7 +784,9 @@ def normalize_chunks_nd( # handle no chunking if chunks is False: - return ChunksTuple(tuple(np.array([s], dtype=np.int64) for s in shape)) + return ChunkGrid( + dimensions=tuple(FixedDimension(size=int(s), extent=int(s)) for s in shape) + ) # handle 1D convenience form. bool is excluded above so this only catches actual ints. if isinstance(chunks, numbers.Integral): @@ -805,20 +798,20 @@ def normalize_chunks_nd( f"chunks has {len(chunks)} dimensions but shape has {len(shape)} dimensions" ) - return ChunksTuple( - tuple(normalize_chunks_1d(c, span=s) for c, s in zip(chunks, shape, strict=True)) + return ChunkGrid( + dimensions=tuple(normalize_chunks_1d(c, span=s) for c, s in zip(chunks, shape, strict=True)) ) def guess_chunks( shape: tuple[int, ...], typesize: int, *, max_bytes: int | None = None -) -> ChunksTuple: +) -> ChunkGrid: """ Heuristically determine chunk sizes for an array. This is the policy function — it makes opinionated choices about chunk sizes based on array shape and element size, and returns a - normalized `ChunksTuple`. + normalized `ChunkGrid`. Parameters ---------- @@ -877,7 +870,7 @@ def _guess_num_chunks_per_axis_shard( def resolve_outer_and_inner_chunks( *, array_shape: tuple[int, ...], - chunks: ChunksTuple, + chunks: ChunkGrid, shard_shape: ShardsLike | None, item_size: int, ) -> ChunkLayout: @@ -888,7 +881,7 @@ def resolve_outer_and_inner_chunks( array_shape The array shape. chunks - Normalized chunk specification (the user's `chunks=`). + Normalized chunk grid (the user's `chunks=`). shard_shape Raw shard specification (the user's `shards=`). `None` means no sharding, `"auto"` triggers heuristic inference, @@ -900,7 +893,7 @@ def resolve_outer_and_inner_chunks( Returns ------- ChunkLayout - `outer_chunks` is the `ChunksTuple` for chunk grid + `outer_chunks` is the `ChunkGrid` for chunk grid metadata. `inner` holds the sub-chunk structure for `ShardingCodec`, or is `None` when sharding is not active. """ @@ -912,8 +905,8 @@ def resolve_outer_and_inner_chunks( outer = normalize_chunks_nd(shard_shape, array_shape) return ChunkLayout(outer_chunks=outer, inner=ChunkLayout(outer_chunks=chunks)) - # Extract the flat chunk shape (first size per dimension) for arithmetic. - chunk_shape_flat = as_regular_shape(chunks) + # Extract the flat chunk shape (uniform size per dimension) for arithmetic. + chunk_shape_flat = chunks.chunk_shape if shard_shape == "auto": warnings.warn( diff --git a/src/zarr/core/common.py b/src/zarr/core/common.py index 3da5c108b6..3008b1fef7 100644 --- a/src/zarr/core/common.py +++ b/src/zarr/core/common.py @@ -35,7 +35,9 @@ BytesLike = bytes | bytearray | memoryview ShapeLike = Iterable[int | np.integer[Any]] | int | np.integer[Any] -ChunksLike = ShapeLike | Iterable[Iterable[int]] +# Per-dimension chunk specs may mix a bare int (uniform chunk size, the +# rectilinear spec's step-size shorthand) with explicit edge-length sequences. +ChunksLike = ShapeLike | Iterable[int | Iterable[int]] # For backwards compatibility ChunkCoords = tuple[int, ...] ZarrFormat = Literal[2, 3] diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index fc47f8fc95..2ab3ff56fb 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -12,7 +12,7 @@ from zarr.core._json import json_to_buffer from zarr.core.array_spec import ArrayConfig, ArraySpec from zarr.core.buffer.core import default_buffer_prototype -from zarr.core.chunk_grids import is_regular_nd +from zarr.core.chunk_grids import FixedDimension, VaryingDimension from zarr.core.chunk_key_encodings import ( ChunkKeyEncoding, ChunkKeyEncodingLike, @@ -42,8 +42,9 @@ if TYPE_CHECKING: from typing import Self + from zarr.codecs.sharding import ShardingCodec from zarr.core.buffer import Buffer, BufferPrototype - from zarr.core.chunk_grids import ChunksTuple + from zarr.core.chunk_grids import ChunkGrid from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar @@ -373,32 +374,36 @@ def from_dict(cls, data: RectilinearChunkGridMetadataJSON) -> Self: # type: ign def create_chunk_grid_metadata( - chunks: ChunksTuple, + chunks: ChunkGrid, ) -> ChunkGridMetadata: - """Construct a chunk grid metadata object from a normalized `ChunksTuple`. + """Construct a chunk grid metadata object from a normalized `ChunkGrid`. - Regular chunks produce a `RegularChunkGridMetadata`. - Rectilinear chunks produce a `RectilinearChunkGridMetadata`. + Regular grids produce a `RegularChunkGridMetadata`. + Rectilinear grids produce a `RectilinearChunkGridMetadata`. Parameters ---------- - chunks : ChunksTuple - Normalized chunk specification, as returned by + chunks : ChunkGrid + Normalized chunk grid, as returned by `normalize_chunks_nd` or `guess_chunks`. See Also -------- parse_chunk_grid : Deserialize a chunk grid from stored JSON metadata. """ - if is_regular_nd(chunks): - # If we know the chunks specification is regular, then we can take the first - # chunk size for each dimension as the chunk shape. - chunk_shape = tuple(int(dim_chunks[0]) for dim_chunks in chunks) - return RegularChunkGridMetadata(chunk_shape=chunk_shape) - else: - return RectilinearChunkGridMetadata( - chunk_shapes=tuple(tuple(int(x) for x in d) for d in chunks) - ) + if chunks.is_regular: + return RegularChunkGridMetadata(chunk_shape=chunks.chunk_shape) + # Uniform dimensions stay bare ints — the rectilinear grid spec treats + # a bare int as a step size repeating to cover the axis. + chunk_shapes: list[int | tuple[int, ...]] = [] + for dim in chunks.dimensions: + if isinstance(dim, FixedDimension): + chunk_shapes.append(dim.size) + elif isinstance(dim, VaryingDimension): + chunk_shapes.append(dim.edges) + else: + raise TypeError(f"Unknown dimension grid type: {type(dim)}") + return RectilinearChunkGridMetadata(chunk_shapes=tuple(chunk_shapes)) def parse_chunk_grid( @@ -565,26 +570,32 @@ def dtype(self) -> ZDType[TBaseDType, TBaseScalar]: # TODO: move these properties to the Array class. # They require knowledge of codecs (ShardingCodec) and don't belong on a metadata DTO. + @property + def sharding_codec(self) -> ShardingCodec | None: + """The array's sharding codec, or None if the array is not sharded.""" + from zarr.codecs.sharding import ShardingCodec + + if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): + return self.codecs[0] + return None + @property def chunks(self) -> tuple[int, ...]: + if (sharding_codec := self.sharding_codec) is not None: + # Inner chunks are always regular, whatever the shape of the outer + # (shard) grid. + return sharding_codec.chunk_shape if not isinstance(self.chunk_grid, RegularChunkGridMetadata): msg = ( "The `chunks` attribute is only defined for arrays using regular chunk grids. " "This array has a rectilinear chunk grid. Use `read_chunk_sizes` for general access." ) raise NotImplementedError(msg) - - from zarr.codecs.sharding import ShardingCodec - - if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): - return self.codecs[0].chunk_shape return self.chunk_grid.chunk_shape @property def shards(self) -> tuple[int, ...] | None: - from zarr.codecs.sharding import ShardingCodec - - if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): + if self.sharding_codec is not None: if not isinstance(self.chunk_grid, RegularChunkGridMetadata): msg = ( "The `shards` attribute is only defined for arrays using regular chunk grids. " @@ -596,10 +607,8 @@ def shards(self) -> tuple[int, ...] | None: @property def inner_codecs(self) -> tuple[Codec, ...]: - from zarr.codecs.sharding import ShardingCodec - - if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): - return self.codecs[0].codecs + if (sharding_codec := self.sharding_codec) is not None: + return sharding_codec.codecs return self.codecs def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: diff --git a/src/zarr/testing/stateful.py b/src/zarr/testing/stateful.py index 652e90e60a..59c5b72e62 100644 --- a/src/zarr/testing/stateful.py +++ b/src/zarr/testing/stateful.py @@ -152,16 +152,12 @@ def add_array(self, data: DataObject, name: str) -> None: ) note(f"Adding array: path='{path}' shape={a.shape} chunks={a.metadata.chunk_grid}") - # Recreate the same array in the store under test. - # The data is copied here rather than by `write_data=True`, - # whose shard-wise copy does not support rectilinear chunk grids. - arr = zarr.from_array( + # Recreate the same array, including its data, in the store under test. + zarr.from_array( self.store, data=a, name=path, - write_data=False, ) - arr[:] = a[:] self.all_arrays.add(path) @rule() diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 6679dbcee4..db14d5f29b 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -312,17 +312,14 @@ def arrays( # - RegularChunkGridMetadata -> flat tuple of ints # - RectilinearChunkGridMetadata -> nested list of ints (triggers rectilinear path) # - v2 -> flat tuple of ints - chunks_param: tuple[int, ...] | list[list[int]] + chunks_param: tuple[int, ...] | list[int | list[int]] shard_shape = None dim_names = None if zarr_format == 3: chunk_grid_meta = draw(st.none() | chunk_grids(shape=nparray.shape), label="chunk grid") dim_names = draw(dimension_names(ndim=nparray.ndim), label="dimension names") if isinstance(chunk_grid_meta, RectilinearChunkGridMetadata): - chunks_param = [ - list(dim) if isinstance(dim, tuple) else [dim] - for dim in chunk_grid_meta.chunk_shapes - ] + chunks_param = chunks_param_from_rectilinear(chunk_grid_meta) elif isinstance(chunk_grid_meta, RegularChunkGridMetadata): chunks_param = chunk_grid_meta.chunk_shape else: @@ -404,6 +401,19 @@ def simple_arrays( ) +def chunks_param_from_rectilinear( + meta: RectilinearChunkGridMetadata, +) -> list[int | list[int]]: + """Convert rectilinear chunk grid metadata into a `chunks=` argument. + + Explicit edge tuples become lists. Bare ints — the spec's step-size + shorthand meaning "repeat to cover the axis" — pass through unchanged; + wrapping one in a single-element list would instead declare exactly one + chunk, which fails normalization whenever the axis needs more than one. + """ + return [list(dim) if isinstance(dim, tuple) else dim for dim in meta.chunk_shapes] + + @st.composite def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[int]]: """Generate valid rectilinear chunk shapes for a given array shape. diff --git a/tests/conftest.py b/tests/conftest.py index 7ccf9958e7..6fc96237af 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -25,7 +25,6 @@ ) from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, - as_regular_shape, guess_chunks, normalize_chunks_nd, resolve_outer_and_inner_chunks, @@ -393,7 +392,7 @@ def create_array_metadata( return ArrayV2Metadata( shape=shape_parsed, dtype=dtype_parsed, - chunks=as_regular_shape(outer_chunks), + chunks=outer_chunks.chunk_shape, order=order_parsed, dimension_separator=chunk_key_encoding_parsed.separator, fill_value=fill_value, @@ -412,7 +411,7 @@ def create_array_metadata( sub_codecs: tuple[Codec, ...] = (*array_array, array_bytes, *bytes_bytes) codecs_out: tuple[Codec, ...] if inner is not None: - inner_chunks_flat = as_regular_shape(inner.outer_chunks) + inner_chunks_flat = inner.outer_chunks.chunk_shape index_location: IndexLocation = "end" if isinstance(shards, dict): index_location = cast("IndexLocation", shards.get("index_location", "end")) diff --git a/tests/test_array.py b/tests/test_array.py index e18ce3cfb1..46890244ec 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -1091,7 +1091,7 @@ def test_auto_partition_auto_shards( shard_shape="auto", item_size=dtype.itemsize, ) - auto_shards = tuple(dim[0] for dim in outer_chunks) + auto_shards = outer_chunks.chunk_shape assert auto_shards == expected_shards @@ -1106,7 +1106,7 @@ def test_auto_partition_auto_shards_with_auto_chunks_should_be_close_to_1MiB() - chunks_normalized = guess_chunks( array_shape, item_size, max_bytes=SHARDED_INNER_CHUNK_MAX_BYTES ) - chunk_shape = tuple(dim[0] for dim in chunks_normalized) + chunk_shape = chunks_normalized.chunk_shape chunk_bytes = np.prod(chunk_shape) * item_size assert chunk_bytes <= SHARDED_INNER_CHUNK_MAX_BYTES assert chunk_bytes > SHARDED_INNER_CHUNK_MAX_BYTES // 4 # should be in the right ballpark @@ -1123,7 +1123,7 @@ def test_auto_partition_auto_shards_with_auto_chunks_should_be_close_to_1MiB() - item_size=item_size, ) assert inner is not None - shard_shape = tuple(dim[0] for dim in outer_chunks) + shard_shape = outer_chunks.chunk_shape # Shard dimensions must be multiples of chunk dimensions assert all(s % c == 0 for s, c in zip(shard_shape, chunk_shape, strict=True)) @@ -2449,3 +2449,15 @@ async def test_create_array_chunks_3d( shape = (10, 12, 15) arr = await create_array(store={}, shape=shape, chunks=chunk_input, dtype="float64") assert arr.write_chunk_sizes == expected + + +async def test_create_array_huge_chunk_count() -> None: + """Array creation must be O(1) in the number of chunks per dimension. + + With `shape=(2**62,)` and `chunks=(1,)` this dimension has 2**62 chunks; + materializing one entry per chunk would raise ("array is too big"). + Companion to the indexing-time fix from gh-4174. + """ + arr = await create_array(store={}, shape=(2**62,), chunks=(1,), dtype="int32") + assert arr.shape == (2**62,) + assert arr.chunks == (1,) diff --git a/tests/test_chunk_grids.py b/tests/test_chunk_grids.py index 4640c43d1c..90ec67fe9d 100644 --- a/tests/test_chunk_grids.py +++ b/tests/test_chunk_grids.py @@ -5,7 +5,10 @@ from tests.conftest import Expect, ExpectFail from zarr.core.chunk_grids import ( + ChunkGrid, ChunkLayout, + FixedDimension, + VaryingDimension, _guess_regular_chunks, normalize_chunks_1d, normalize_chunks_nd, @@ -14,15 +17,30 @@ def _assert_chunks_equal( - actual: tuple[Any, ...], - expected: tuple[tuple[int, ...], ...], + actual: ChunkGrid, + expected: tuple[int | tuple[int, ...], ...], + shape: tuple[int, ...], ) -> None: - """Compare a ChunksTuple (tuple of np.int64 arrays) against a tuple of int tuples.""" - assert len(actual) == len(expected), f"axis count mismatch: {len(actual)} vs {len(expected)}" - for axis, (a, e) in enumerate(zip(actual, expected, strict=True)): - assert np.array_equal(a, np.asarray(e, dtype=np.int64)), ( - f"axis {axis}: {list(a)} != {list(e)}" - ) + """Compare a normalized ChunkGrid against per-dimension expectations. + + An expected bare `int` requires a `FixedDimension` with that uniform size; + an expected tuple requires a `VaryingDimension` with those exact edges. + Every dimension must carry the corresponding extent from `shape`. + """ + dims = actual.dimensions + assert len(dims) == len(expected) == len(shape), ( + f"axis count mismatch: {len(dims)} vs {len(expected)} vs {len(shape)}" + ) + for axis, (a, e, span) in enumerate(zip(dims, expected, shape, strict=True)): + if isinstance(e, int): + assert isinstance(a, FixedDimension), f"axis {axis}: expected FixedDimension, got {a!r}" + assert a.size == e, f"axis {axis}: size {a.size} != {e}" + else: + assert isinstance(a, VaryingDimension), ( + f"axis {axis}: expected VaryingDimension, got {a!r}" + ) + assert a.edges == tuple(e), f"axis {axis}: edges {a.edges} != {tuple(e)}" + assert a.extent == span, f"axis {axis}: extent {a.extent} != {span}" @pytest.mark.parametrize( @@ -42,79 +60,83 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None: @pytest.mark.parametrize( ("chunks", "shape", "expected"), [ - # 1D cases - ((10,), (100,), ((10,) * 10,)), - ([10], (100,), ((10,) * 10,)), - (10, (100,), ((10,) * 10,)), + # 1D cases (uniform sizes stay bare ints) + ((10,), (100,), (10,)), + ([10], (100,), (10,)), + (10, (100,), (10,)), # 2D cases - ((10, 10), (100, 10), ((10,) * 10, (10,))), - (10, (100, 10), ((10,) * 10, (10,))), - ((10, -1), (100, 10), ((10,) * 10, (10,))), + ((10, 10), (100, 10), (10, 10)), + (10, (100, 10), (10, 10)), + ((10, -1), (100, 10), (10, 10)), # 3D cases - (30, (100, 20, 10), ((30, 30, 30, 30), (30,), (30,))), - ((30, -1, -1), (100, 20, 10), ((30, 30, 30, 30), (20,), (10,))), - ((30, 20, -1), (100, 20, 10), ((30, 30, 30, 30), (20,), (10,))), - ((30, 20, 10), (100, 20, 10), ((30, 30, 30, 30), (20,), (10,))), - # dask-style chunks (explicit per-chunk sizes) + (30, (100, 20, 10), (30, 30, 30)), + ((30, -1, -1), (100, 20, 10), (30, 20, 10)), + ((30, 20, -1), (100, 20, 10), (30, 20, 10)), + ((30, 20, 10), (100, 20, 10), (30, 20, 10)), + # dask-style explicit lists always keep the per-chunk rectilinear form, + # even when the sizes describe a regular grid (gh-4272) (((100, 100, 100), (50, 50)), (300, 100), ((100, 100, 100), (50, 50))), (((100, 100, 50),), (250,), ((100, 100, 50),)), (((100,),), (100,), ((100,),)), + (((10, 20, 70), (50, 50)), (100, 100), ((10, 20, 70), (50, 50))), # no chunking (False means each dimension is one chunk spanning the full extent) - (False, (100,), ((100,),)), - (False, (100, 50), ((100,), (50,))), + (False, (100,), (100,)), + (False, (100, 50), (100, 50)), # sentinel values - (-1, (100,), ((100,),)), + (-1, (100,), (100,)), # zero-length dimensions preserve the declared chunk size - (10, (0,), ((10,),)), - ((5, 10), (0, 100), ((5,), (10,) * 10)), - ((5, 10), (20, 0), ((5, 5, 5, 5), (10,))), + (10, (0,), (10,)), + ((5, 10), (0, 100), (5, 10)), + ((5, 10), (20, 0), (5, 10)), # numpy integers are accepted anywhere a python int is, whether as the scalar # convenience form, as per-dimension entries, or as the `-1` sentinel. - (np.int64(10), (100,), ((10,) * 10,)), - ((np.int64(2), np.int64(2)), (4, 4), ((2, 2), (2, 2))), - ((1, 3, np.int64(16), np.int64(16)), (1, 3, 32, 32), ((1,), (3,), (16, 16), (16, 16))), - ((np.int32(30), np.int64(-1)), (100, 20), ((30, 30, 30, 30), (20,))), - (np.array([10, 10]), (100, 100), ((10,) * 10, (10,) * 10)), + (np.int64(10), (100,), (10,)), + ((np.int64(2), np.int64(2)), (4, 4), (2, 2)), + ((1, 3, np.int64(16), np.int64(16)), (1, 3, 32, 32), (1, 3, 16, 16)), + ((np.int32(30), np.int64(-1)), (100, 20), (30, 20)), + (np.array([10, 10]), (100, 100), (10, 10)), # rectilinear chunks given as numpy arrays ((np.array([60, 40]), np.array([50, 50])), (100, 100), ((60, 40), (50, 50))), ], ) def test_normalize_chunks( - chunks: Any, shape: tuple[int, ...], expected: tuple[tuple[int, ...], ...] + chunks: Any, shape: tuple[int, ...], expected: tuple[int | tuple[int, ...], ...] ) -> None: - _assert_chunks_equal(normalize_chunks_nd(chunks, shape), expected) + _assert_chunks_equal(normalize_chunks_nd(chunks, shape), expected, shape) @pytest.mark.parametrize( ("array_shape", "chunks_input", "shard_shape", "expected_outer", "expected_inner_outer"), [ # no sharding: outer = chunks, inner = None - ((100,), (10,), None, ((10,) * 10,), None), + ((100,), (10,), None, (10,), None), # explicit regular shards - ((100,), (10,), (50,), ((50, 50),), ((10,) * 10,)), - # rectilinear shards - ((100,), (10,), ((60, 40),), ((60, 40),), ((10,) * 10,)), + ((100,), (10,), (50,), (50,), (10,)), + # rectilinear shards keep the per-chunk form even when the sizes + # describe a regular-with-boundary grid (gh-4272) + ((100,), (10,), ((60, 40),), ((60, 40),), (10,)), + ((100,), (10,), ((30, 60, 10),), ((30, 60, 10),), (10,)), # dict-style shards - ((100, 100), (10, 10), {"shape": (50, 50)}, ((50, 50), (50, 50)), ((10,) * 10, (10,) * 10)), + ((100, 100), (10, 10), {"shape": (50, 50)}, (50, 50), (10, 10)), ], ) def test_resolve_outer_and_inner_chunks( array_shape: tuple[int, ...], chunks_input: tuple[int, ...], shard_shape: Any, - expected_outer: tuple[tuple[int, ...], ...], - expected_inner_outer: tuple[tuple[int, ...], ...] | None, + expected_outer: tuple[int | tuple[int, ...], ...], + expected_inner_outer: tuple[int | tuple[int, ...], ...] | None, ) -> None: chunks = normalize_chunks_nd(chunks_input, array_shape) outer_chunks, inner = resolve_outer_and_inner_chunks( array_shape=array_shape, chunks=chunks, shard_shape=shard_shape, item_size=1 ) - _assert_chunks_equal(outer_chunks, expected_outer) + _assert_chunks_equal(outer_chunks, expected_outer, array_shape) if expected_inner_outer is None: assert inner is None else: assert inner is not None - _assert_chunks_equal(inner.outer_chunks, expected_inner_outer) + _assert_chunks_equal(inner.outer_chunks, expected_inner_outer, array_shape) assert inner.inner is None @@ -128,11 +150,11 @@ def test_chunk_layout_nested() -> None: top = ChunkLayout(outer_chunks=normalize_chunks_nd((50, 50), (100, 100)), inner=mid) # Three levels: top -> mid -> leaf - _assert_chunks_equal(top.outer_chunks, ((50, 50), (50, 50))) + _assert_chunks_equal(top.outer_chunks, (50, 50), (100, 100)) assert top.inner is not None - _assert_chunks_equal(top.inner.outer_chunks, ((25,) * 4, (25,) * 4)) + _assert_chunks_equal(top.inner.outer_chunks, (25, 25), (100, 100)) assert top.inner.inner is not None - _assert_chunks_equal(top.inner.inner.outer_chunks, ((5,) * 20, (5,) * 20)) + _assert_chunks_equal(top.inner.inner.outer_chunks, (5, 5), (100, 100)) assert top.inner.inner.inner is None @@ -258,22 +280,42 @@ def test_normalize_chunks_nd_errors(case: ExpectFail[tuple[Any, tuple[int, ...]] @pytest.mark.parametrize( "case", [ - # uniform-chunks branch: one int → broadcast across span via np.full. - Expect(input=(1000, 100_000), output=[1000] * 100, id="uniform"), - # explicit-per-chunk branch. - Expect(input=([10, 20, 30, 40], 100), output=[10, 20, 30, 40], id="explicit-list"), + # uniform-chunks branch: O(1) size+extent record, never one entry per chunk. + Expect( + input=(1000, 100_000), output=FixedDimension(size=1000, extent=100_000), id="uniform" + ), + # uniform chunks on a span too large to expand per-chunk (creation-time + # counterpart of the gh-4174 indexing fix). + Expect(input=(1, 2**62), output=FixedDimension(size=1, extent=2**62), id="uniform-huge"), # -1 sentinel branch: one chunk covering the full span. - Expect(input=(-1, 100), output=[100], id="full-span-sentinel"), + Expect(input=(-1, 100), output=FixedDimension(size=100, extent=100), id="full-span"), + # zero-length span preserves the declared chunk size. + Expect(input=(10, 0), output=FixedDimension(size=10, extent=0), id="uniform-zero-span"), + # explicit lists always keep the per-chunk form, even when the sizes + # describe a regular grid (gh-4272). + Expect( + input=([10, 10, 10], 30), + output=VaryingDimension([10, 10, 10], extent=30), + id="explicit-regular", + ), + Expect( + input=([10, 10, 5], 25), + output=VaryingDimension([10, 10, 5], extent=25), + id="explicit-boundary", + ), + Expect( + input=([10, 20, 70], 100), + output=VaryingDimension([10, 20, 70], extent=100), + id="explicit-irregular", + ), ], ids=lambda c: c.id, ) -def test_normalize_chunks_1d_returns_int64_array( - case: Expect[tuple[Any, int], list[int]], +def test_normalize_chunks_1d( + case: Expect[tuple[Any, int], FixedDimension | VaryingDimension], ) -> None: - """Every branch of normalize_chunks_1d must produce a 1D int64 array.""" + """Both output variants bind chunk sizes to the span: scalar specs become + `FixedDimension` (O(1) regardless of chunk count), explicit per-chunk + lists become `VaryingDimension`.""" chunks, span = case.input - result = normalize_chunks_1d(chunks, span) - assert isinstance(result, np.ndarray) - assert result.dtype == np.int64 - assert result.ndim == 1 - assert result.tolist() == case.output + assert normalize_chunks_1d(chunks, span) == case.output diff --git a/tests/test_metadata/test_v3.py b/tests/test_metadata/test_v3.py index d1e156e500..9f78ed7b70 100644 --- a/tests/test_metadata/test_v3.py +++ b/tests/test_metadata/test_v3.py @@ -5,13 +5,12 @@ import json from typing import TYPE_CHECKING -import numpy as np import pytest from tests.conftest import Expect, ExpectFail from tests.test_metadata.conftest import minimal_metadata_dict_v3 from zarr.core.buffer import default_buffer_prototype -from zarr.core.chunk_grids import is_regular_1d, is_regular_nd +from zarr.core.chunk_grids import ChunkGrid, is_regular_1d, is_regular_nd from zarr.core.config import config from zarr.core.dtype import Float64, UInt8 from zarr.core.group import GroupMetadata, parse_node_type @@ -20,6 +19,7 @@ ARRAY_METADATA_KEYS, ArrayMetadataJSON_V3, ArrayV3Metadata, + create_chunk_grid_metadata, parse_codecs, parse_dimension_names, parse_node_type_array, @@ -110,9 +110,6 @@ def test_parse_codecs_unknown_raises(monkeypatch: pytest.MonkeyPatch) -> None: # Chunk-grid regularity helpers # --------------------------------------------------------------------------- -# Cases used for both list/tuple (Python-sequence path) and ndarray (vectorized -# path) of `is_regular_1d`. Parametrizing the input form ensures both branches -# are exercised by the same suite of edge cases. _REGULAR_1D_CASES: list[Expect[list[int], bool]] = [ Expect(input=[], output=True, id="empty"), Expect(input=[10], output=True, id="single-chunk"), @@ -129,19 +126,11 @@ def test_parse_codecs_unknown_raises(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.parametrize("case", _REGULAR_1D_CASES, ids=lambda c: c.id) def test_is_regular_1d_sequence(case: Expect[list[int], bool]) -> None: - """`is_regular_1d` accepts plain Python sequences and uses the iterative path.""" - # list and tuple both go through the non-ndarray branch. + """`is_regular_1d` accepts plain Python sequences.""" assert is_regular_1d(case.input) is case.output assert is_regular_1d(tuple(case.input)) is case.output -@pytest.mark.parametrize("case", _REGULAR_1D_CASES, ids=lambda c: c.id) -def test_is_regular_1d_ndarray(case: Expect[list[int], bool]) -> None: - """`is_regular_1d` accepts int64 ndarrays and uses the vectorized path.""" - arr = np.asarray(case.input, dtype=np.int64) - assert is_regular_1d(arr) is case.output - - @pytest.mark.parametrize( "case", [ @@ -156,8 +145,13 @@ def test_is_regular_1d_ndarray(case: Expect[list[int], bool]) -> None: def test_is_regular_nd_sequence(case: Expect[list[list[int]], bool]) -> None: """`is_regular_nd` returns True iff every per-dim spec is regular.""" assert is_regular_nd(case.input) is case.output - # Same result via ndarray inputs. - assert is_regular_nd([np.asarray(d, dtype=np.int64) for d in case.input]) is case.output + + +def test_create_chunk_grid_metadata_unknown_dimension_type() -> None: + """`create_chunk_grid_metadata` rejects dimension grids it does not recognize.""" + grid = ChunkGrid(dimensions=(object(),)) # type: ignore[arg-type] + with pytest.raises(TypeError, match="Unknown dimension grid type"): + create_chunk_grid_metadata(grid) # --------------------------------------------------------------------------- diff --git a/tests/test_properties.py b/tests/test_properties.py index 33888bfd4e..651c61e264 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -448,3 +448,27 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N assert serialized_complex_float_is_valid(asdict_dict["fill_value"]) elif dtype_native.kind in ("M", "m") and np.isnat(meta.fill_value): assert asdict_dict["fill_value"] == -9223372036854775808 + + +def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: + """Bare-int dims in rectilinear metadata (the spec's step-size shorthand, + produced by a scalar dimension of a mixed chunk spec) must pass + through the `chunks=` conversion unchanged. Wrapping one in a + single-element list turns "repeat to cover the axis" into "exactly one + chunk" and re-creation fails the sum-to-span check.""" + from zarr.core.metadata.v3 import RectilinearChunkGridMetadata + from zarr.storage import MemoryStore + from zarr.testing.strategies import chunks_param_from_rectilinear + + with zarr.config.set({"array.rectilinear_chunks": True}): + src = zarr.create_array(MemoryStore(), shape=(3, 3), chunks=([1, 2], 1), dtype="uint8") + grid = src.metadata.chunk_grid # type: ignore[union-attr] + assert isinstance(grid, RectilinearChunkGridMetadata) + assert grid.chunk_shapes == ((1, 2), 1) + dst = zarr.create_array( + MemoryStore(), + shape=src.shape, + chunks=chunks_param_from_rectilinear(grid), + dtype="uint8", + ) + assert dst.metadata.chunk_grid == grid # type: ignore[union-attr] diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index f0b54519ab..963bd95973 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -1027,6 +1027,197 @@ def test_e2e_chunk_grid_name_regular_from_dict(tmp_path: Path) -> None: assert chunk_grid_dict["name"] == "regular" +# --------------------------------------------------------------------------- +# Input syntax determines grid kind (gh-4272) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("shape", "chunks", "expected_type"), + [ + # scalar specs (including numpy integers and the -1 sentinel) produce a regular grid + ((30,), (10,), RegularChunkGridMetadata), + ((30,), 10, RegularChunkGridMetadata), + ((30,), (np.int64(10),), RegularChunkGridMetadata), + ((30,), -1, RegularChunkGridMetadata), + # explicit per-chunk lists stay rectilinear even when the sizes are uniform + ((30,), [[10, 10, 10]], RectilinearChunkGridMetadata), + # ... or uniform with a short tail (the append-oriented time-series layout of gh-4272) + ((168 * 13 + 24,), [[168] * 13 + [24]], RectilinearChunkGridMetadata), + # genuinely varying edges are rectilinear too + ((60,), [[10, 20, 30]], RectilinearChunkGridMetadata), + ], + ids=[ + "scalar-tuple", + "scalar-int", + "scalar-numpy-int", + "full-span-sentinel", + "explicit-uniform", + "explicit-uniform-short-tail", + "explicit-varying", + ], +) +def test_chunk_grid_kind_follows_input_syntax( + tmp_path: Path, shape: tuple[int, ...], chunks: Any, expected_type: type +) -> None: + """The stored grid kind is decided by the input syntax, not the chunk sizes: + scalar specs produce a regular grid, while explicit per-chunk lists produce + a rectilinear grid even when the sizes happen to describe a regular one + (gh-4272).""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + arr = zarr.create_array(store=tmp_path / "arr.zarr", shape=shape, chunks=chunks, dtype="int32") + assert isinstance(arr.metadata, ArrayV3Metadata) + assert isinstance(arr.metadata.chunk_grid, expected_type) + + +def test_mixed_scalar_and_list_dims_keep_shorthand(tmp_path: Path) -> None: + """In a mixed spec, a scalar dimension serializes as the bare-int step-size + shorthand while an explicit list keeps its per-chunk edges.""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + arr = zarr.create_array( + store=tmp_path / "arr.zarr", shape=(30, 100), chunks=(5, [10, 20, 70]), dtype="int32" + ) + assert isinstance(arr.metadata, ArrayV3Metadata) + grid = arr.metadata.chunk_grid + assert isinstance(grid, RectilinearChunkGridMetadata) + assert grid.chunk_shapes == (5, (10, 20, 70)) + + +# The chunk/shard spec matrix for `from_array` round-trip tests. Every kind of +# spec `create_array` accepts should appear here, so that any change to how +# specs are normalized or stored is automatically checked against the "keep" +# path of `from_array` as well. +FROM_ARRAY_KEEP_CASES = [ + pytest.param((30, 40), (10, 20), None, id="regular"), + pytest.param((40, 40), (5, 10), (10, 20), id="regular-sharded"), + pytest.param((60, 100), [[10, 20, 30], [50, 50]], None, id="rectilinear"), + pytest.param((9, 30), (3, [10, 20]), None, id="rectilinear-mixed-bare-int"), + pytest.param((3,), [[3]], None, id="rectilinear-uniform-list"), + pytest.param((100,), (10,), [[50, 50]], id="rectilinear-sharded"), +] + + +@pytest.mark.parametrize("write_data", [True, False]) +@pytest.mark.parametrize(("shape", "chunks", "shards"), FROM_ARRAY_KEEP_CASES) +def test_from_array_keep_roundtrips_chunk_grid( + shape: tuple[int, ...], chunks: Any, shards: Any, write_data: bool +) -> None: + """For every chunk/shard spec accepted by `create_array`, `from_array` with + the default "keep" parameters reproduces the source's stored chunk grid and + codecs exactly: the grid kind is preserved (gh-4272), uniform dimensions + keep their bare-int shorthand instead of being expanded per chunk, and + sharding survives — including under a rectilinear shard grid. `.info` is + defined for every kind of source, and the default data copy works for + every kind of grid.""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + src = zarr.create_array( + MemoryStore(), shape=shape, chunks=chunks, shards=shards, dtype="uint16" + ) + src[:] = np.arange(np.prod(shape), dtype="uint16").reshape(shape) + assert src.info is not None + # The array is fully written, so every chunk is initialized, whatever the + # grid kind and sharding layout. + assert src.nchunks_initialized == src.nchunks + dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=write_data) + assert isinstance(src.metadata, ArrayV3Metadata) + assert isinstance(dst.metadata, ArrayV3Metadata) + assert dst.metadata.chunk_grid == src.metadata.chunk_grid + assert dst.metadata.codecs == src.metadata.codecs + if write_data: + np.testing.assert_array_equal(dst[:], src[:]) + + +def test_info_reports_variable_shard_shape_for_rectilinear_shard_grid() -> None: + """A sharded array whose shard grid is rectilinear has no uniform shard + shape. `.info` must still present it as sharded — shard shape rendered as + ``, and the initialized count labeled as shards, not chunks.""" + arr = zarr.create_array( + MemoryStore(), shape=(100,), chunks=(10,), shards=[[50, 50]], dtype="int32" + ) + arr[:] = 1 + assert "Shard shape : " in repr(arr.info) + info = arr.info_complete() + assert "Shards Initialized : 2" in repr(info) + assert arr.nchunks_initialized == 10 + + +def test_nchunks_initialized_counts_per_shard_for_rectilinear_shard_grid() -> None: + """With a rectilinear shard grid the chunk count per shard varies, so the + initialized-chunk count is summed per stored shard rather than derived + from a uniform multiplier.""" + arr = zarr.create_array( + MemoryStore(), shape=(90,), chunks=(10,), shards=[[60, 30]], dtype="int32" + ) + assert arr.nchunks_initialized == 0 + arr[60:90] = 1 # initializes only the second shard, which holds 3 chunks + assert arr.nchunks_initialized == 3 + arr[0:10] = 1 # initializes the first shard, which holds 6 chunks + assert arr.nchunks_initialized == 9 + + +def test_chunks_raises_for_nonsharded_rectilinear_grid() -> None: + """`.chunks` has no uniform value for a non-sharded rectilinear grid.""" + arr = zarr.create_array(MemoryStore(), shape=(30,), chunks=[[10, 20]], dtype="int32") + with pytest.raises(NotImplementedError, match="regular chunk grids"): + _ = arr.chunks + + +def test_shards_raises_for_rectilinear_shard_grid() -> None: + """`.shards` has no uniform value when the shard grid is rectilinear.""" + arr = zarr.create_array( + MemoryStore(), shape=(100,), chunks=(10,), shards=[[50, 50]], dtype="int32" + ) + with pytest.raises(NotImplementedError, match="regular chunk grids"): + _ = arr.shards + + +def test_from_array_keep_is_o1_in_chunk_count() -> None: + """`chunks="keep"` passes uniform dimensions through as bare-int shorthand + rather than expanding one entry per chunk, so copying array metadata is + O(ndim), not O(nchunks): this completes instantly despite ~2**60 chunks + along the first dimension (and would hang if the shorthand were expanded).""" + src = zarr.create_array(MemoryStore(), shape=(2**62, 30), chunks=(3, [10, 20]), dtype="uint8") + dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=False) + grid = dst.metadata.chunk_grid + assert isinstance(grid, RectilinearChunkGridMetadata) + assert grid.chunk_shapes == (3, (10, 20)) + + +def test_resize_uniform_rectilinear_appends_edge() -> None: + """Growing an explicitly rectilinear array whose edges look regular appends + a new edge chunk, while the same sizes declared as a scalar extend the + uniform pattern instead (gh-4272), so an append-only workload writing the + grown region touches exactly one new chunk.""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + rect_store: dict[str, Any] = {} + rect = zarr.create_array(store=rect_store, shape=(30,), chunks=[[10, 10, 10]], dtype="int32") + rect.resize((45,)) + assert isinstance(rect.metadata, ArrayV3Metadata) + rect_grid = rect.metadata.chunk_grid + assert isinstance(rect_grid, RectilinearChunkGridMetadata) + assert rect_grid.chunk_shapes == ((10, 10, 10, 15),) + assert rect.nchunks == 4 + # the appended region is exactly the new edge chunk + rect[30:45] = 1 + assert sorted(k for k in rect_store if k.startswith("c/")) == ["c/3"] + + reg_store: dict[str, Any] = {} + reg = zarr.create_array(store=reg_store, shape=(30,), chunks=(10,), dtype="int32") + reg.resize((45,)) + assert isinstance(reg.metadata, ArrayV3Metadata) + reg_grid = reg.metadata.chunk_grid + assert isinstance(reg_grid, RegularChunkGridMetadata) + assert reg_grid.chunk_shape == (10,) + assert reg.nchunks == 5 + # the same write straddles two chunks of the extended uniform pattern + reg[30:45] = 1 + assert sorted(k for k in reg_store if k.startswith("c/")) == ["c/3", "c/4"] + + # --------------------------------------------------------------------------- # Sharding compatibility tests # ---------------------------------------------------------------------------