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/src/zarr/core/array.py b/src/zarr/core/array.py index 2b31eefcd4..32121c7b4e 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, @@ -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 @@ -4467,7 +4466,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 @@ -4510,7 +4509,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, @@ -4529,7 +4528,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")) diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 2cb9762775..81f75f28ed 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,27 +700,26 @@ 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 uniform chunk sizes, `VaryingDimension` for explicit + per-chunk sizes that genuinely vary. 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; lists that describe + a regular grid (all sizes equal, or equal with a smaller boundary chunk) + collapse to `FixedDimension`. For uniform sizes the last chunk may overhang + the span. """ if chunks == -1: - return np.array([span], dtype=np.int64) + return FixedDimension(size=span, extent=span) if isinstance(chunks, int): if chunks <= 0: raise ValueError(f"Chunk size must be positive, got {chunks}") - if span == 0: - return np.array([chunks], dtype=np.int64) - n = ceildiv(span, chunks) - return np.full(n, chunks, dtype=np.int64) + return FixedDimension(size=chunks, extent=span) else: chunk_list = list(chunks) if not chunk_list: @@ -757,23 +739,30 @@ 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) + if is_regular_1d(ints): + return FixedDimension(size=ints[0], extent=span) + 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. """ @@ -784,7 +773,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): @@ -796,20 +787,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 ---------- @@ -868,7 +859,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: @@ -879,7 +870,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, @@ -891,7 +882,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. """ @@ -903,8 +894,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 1541683b09..271c56bd34 100644 --- a/src/zarr/core/common.py +++ b/src/zarr/core/common.py @@ -34,7 +34,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 9eaccc5076..1a769039d3 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,7 +42,7 @@ from typing import Self 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 @@ -372,32 +372,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( diff --git a/src/zarr/testing/stateful.py b/src/zarr/testing/stateful.py index 9105b8234e..f36de766ec 100644 --- a/src/zarr/testing/stateful.py +++ b/src/zarr/testing/stateful.py @@ -154,13 +154,12 @@ def add_array(self, data: DataObject, name: str) -> None: # Recreate the same array in the store under test from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata + from zarr.testing.strategies import chunks_param_from_rectilinear chunk_grid = a.metadata.chunk_grid - chunks_param: tuple[int, ...] | list[list[int]] + chunks_param: tuple[int, ...] | list[int | list[int]] if isinstance(chunk_grid, RectilinearChunkGridMetadata): - chunks_param = [ - list(dim) if isinstance(dim, tuple) else [dim] for dim in chunk_grid.chunk_shapes - ] + chunks_param = chunks_param_from_rectilinear(chunk_grid) elif isinstance(chunk_grid, RegularChunkGridMetadata): chunks_param = chunk_grid.chunk_shape else: diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 99e81b0389..344ce57c46 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 b1a7a3c0f2..54f723c64e 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)) @@ -2386,3 +2386,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 b730a43901..49a44ff7c2 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,70 +60,74 @@ 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) - (((100, 100, 100), (50, 50)), (300, 100), ((100, 100, 100), (50, 50))), - (((100, 100, 50),), (250,), ((100, 100, 50),)), - (((100,),), (100,), ((100,),)), + (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 chunks describing a regular grid collapse to the uniform form + (((100, 100, 100), (50, 50)), (300, 100), (100, 50)), + (((100, 100, 50),), (250,), (100,)), + (((100,),), (100,), (100,)), + # genuinely irregular explicit chunks keep the per-chunk form + (((10, 20, 70), (50, 50)), (100, 100), ((10, 20, 70), 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)), ], ) 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 describing a regular-with-boundary grid collapse + ((100,), (10,), ((60, 40),), (60,), (10,)), + # genuinely irregular rectilinear shards keep the per-chunk form + ((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 @@ -119,11 +141,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 @@ -234,22 +256,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 that describe a regular grid collapse to the uniform form. + Expect( + input=([10, 10, 10], 30), + output=FixedDimension(size=10, extent=30), + id="explicit-regular", + ), + Expect( + input=([10, 10, 5], 25), + output=FixedDimension(size=10, extent=25), + id="explicit-boundary", + ), + # genuinely irregular edges keep the explicit per-chunk form. + 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: uniform specs become + `FixedDimension` (O(1) regardless of chunk count), irregular explicit + 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..a73a6d8a7c 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -448,3 +448,29 @@ 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 when a uniform dimension of a mixed grid is collapsed) 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, 1, 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]