DimitriPapadopoulos
diff --git a/‎changes/4001.misc.md‎
Lines changed: 6 additions & 1 deletion b/‎changes/4001.misc.md‎
Lines changed: 6 additions & 1 deletion
diff --git a/‎src/zarr/codecs/sharding.py‎
Lines changed: 44 additions & 29 deletions b/‎src/zarr/codecs/sharding.py‎
Lines changed: 44 additions & 29 deletions
diff --git a/‎src/zarr/core/indexing.py‎
Lines changed: 55 additions & 15 deletions b/‎src/zarr/core/indexing.py‎
Lines changed: 55 additions & 15 deletions
@@ -1 +1,6 @@
-Consolidated the array indexing test suite (`tests/test_indexing.py`): the loop-and-`np.random` based selection tests were rewritten as deterministic, parametrized `Expect`/`ExpectFail` cases on small arrays, error paths were split into their own named tests, and the two divergent `Expect` test-case dataclass pairs were unified onto the canonical one in `tests/conftest.py` (whose `ExpectFail` now has an optional regex `msg` and a `raises()` helper). Test-only change with no effect on the public API.
+Restore sharding write performance for shards with many inner chunks. The
+`subchunk_write_order` feature inadvertently rebuilt the per-shard chunk
+coordinate grid (up to tens of thousands of coordinate tuples) on every shard
+write. These coordinates are now computed once per shard shape and cached, so
+repeated writes to same-shaped shards reuse them, restoring write throughput to
+its previous level.
@@ -46,9 +46,11 @@
     BasicIndexer,
     ChunkProjection,
     SelectorTuple,
-    c_order_iter,
+    _lexicographic_order,
+    colexicographic_order_coords,
     get_indexer,
-    morton_order_iter,
+    lexicographic_order_coords,
+    morton_order_coords,
 )
 from zarr.core.metadata.v3 import (
     ChunkGridMetadata,
@@ -261,31 +263,42 @@ def __len__(self) -> int:
         return int(self.index.offsets_and_lengths.size / 2)
 
     def __iter__(self) -> Iterator[tuple[int, ...]]:
-        return c_order_iter(self.index.chunks_per_shard)
+        return iter(lexicographic_order_coords(self.index.chunks_per_shard))
 
-    def to_dict_vectorized(
-        self,
-        chunk_coords_array: npt.NDArray[np.integer[Any]],
-    ) -> dict[tuple[int, ...], Buffer | None]:
+    def to_dict_vectorized(self) -> dict[tuple[int, ...], Buffer | None]:
         """Build a dict of chunk coordinates to buffers using vectorized lookup.
 
-        Parameters
-        ----------
-        chunk_coords_array : ndarray of shape (n_chunks, n_dims)
-            Array of chunk coordinates for vectorized index lookup.
+        The full per-shard chunk coordinate grid (both the array used for the
+        vectorized index lookup and the plain tuples used as dict keys) is
+        cached on `chunks_per_shard`, so neither is rebuilt on every call. For a
+        shard with tens of thousands of chunks this avoids reconstructing that
+        many tuples on every partial write.
 
         Returns
         -------
         dict mapping chunk coordinate tuples to Buffer or None
         """
+        chunks_per_shard = self.index.chunks_per_shard
+        # The same chunk-grid coordinates are needed in two forms, and neither can
+        # stand in for the other:
+        #   - `chunk_coords_array`: an (n_chunks, n_dims) numpy array, fed to the
+        #     vectorized index lookup, which does modulo + advanced indexing on it.
+        #     A list of tuples can't be used for that without first being arrayified.
+        #   - `chunk_coords_keys`: the same coordinates as hashable Python tuples,
+        #     used as the result dict's keys. numpy array rows are unhashable
+        #     (mutable), so they can't key a dict.
+        # Both are cached per shape (see indexing.py), so neither is rebuilt here;
+        # row i of the array and key i refer to the same chunk.
+        chunk_coords_array = _lexicographic_order(chunks_per_shard)
+        chunk_coords_keys = lexicographic_order_coords(chunks_per_shard)
         starts, ends, valid = self.index.get_chunk_slices_vectorized(chunk_coords_array)
 
         result: dict[tuple[int, ...], Buffer | None] = {}
-        for i, coords in enumerate(chunk_coords_array):
+        for i, coords in enumerate(chunk_coords_keys):
             if valid[i]:
-                result[tuple(coords.ravel())] = self.buf[int(starts[i]) : int(ends[i])]
+                result[coords] = self.buf[int(starts[i]) : int(ends[i])]
             else:
-                result[tuple(coords.ravel())] = None
+                result[coords] = None
 
         return result
 
@@ -533,13 +546,14 @@ async def _decode_partial_single(
     def _subchunk_order_iter(
         self, chunks_per_shard: tuple[int, ...], subchunk_write_order: SubchunkWriteOrder
     ) -> Iterable[tuple[int, ...]]:
+        subchunk_iter: Iterable[tuple[int, ...]]
         match subchunk_write_order:
             case "morton":
-                subchunk_iter = morton_order_iter(chunks_per_shard)
+                subchunk_iter = morton_order_coords(chunks_per_shard)
             case "lexicographic":
-                subchunk_iter = np.ndindex(chunks_per_shard)
+                subchunk_iter = lexicographic_order_coords(chunks_per_shard)
             case "colexicographic":
-                subchunk_iter = (c[::-1] for c in np.ndindex(chunks_per_shard[::-1]))
+                subchunk_iter = colexicographic_order_coords(chunks_per_shard)
             case "unordered":
                 # "unordered" promises no particular layout; today it happens to be
                 # lexicographic, but callers must not rely on that.
@@ -565,9 +579,7 @@ async def _encode_single(
                 chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape),
             )
         )
-        # The key order of this intermediate dict is immaterial; the physical layout is
-        # decided later by the `subchunk_write_order` loop in `_encode_shard_dict`.
-        shard_builder = dict.fromkeys(np.ndindex(chunks_per_shard))
+        shard_builder = dict.fromkeys(lexicographic_order_coords(chunks_per_shard))
 
         await self.codec_pipeline.write(
             [
@@ -610,19 +622,18 @@ async def _encode_partial_single(
         )
 
         if self._is_complete_shard_write(indexer, chunks_per_shard):
-            # Intermediate key order is immaterial (see `_encode_single`).
-            shard_dict = dict.fromkeys(np.ndindex(chunks_per_shard))
+            shard_dict = dict.fromkeys(lexicographic_order_coords(chunks_per_shard))
         else:
             shard_reader = await self._load_full_shard_maybe(
                 byte_getter=byte_setter,
                 prototype=chunk_spec.prototype,
                 chunks_per_shard=chunks_per_shard,
             )
             shard_reader = shard_reader or _ShardReader.create_empty(chunks_per_shard)
-            # Use vectorized lookup for better performance
-            shard_dict = shard_reader.to_dict_vectorized(
-                np.array(list(np.ndindex(chunks_per_shard)))
-            )
+            # Use vectorized lookup for better performance. The lexicographic
+            # coordinate array and keys are cached, so neither is rebuilt on
+            # every write.
+            shard_dict = shard_reader.to_dict_vectorized()
 
         await self.codec_pipeline.write(
             [
@@ -692,9 +703,13 @@ async def _encode_shard_dict(
     def _is_total_shard(
         self, all_chunk_coords: set[tuple[int, ...]], chunks_per_shard: tuple[int, ...]
     ) -> bool:
-        return len(all_chunk_coords) == product(chunks_per_shard) and all(
-            chunk_coords in all_chunk_coords for chunk_coords in c_order_iter(chunks_per_shard)
-        )
+        # `all_chunk_coords` comes from an indexer over this shard's chunk grid, so
+        # it is always a subset of that grid (`validate` requires the shard shape to
+        # be divisible by the inner chunk shape, so the indexer cannot produce an
+        # out-of-grid coordinate). A subset whose size equals the grid's is the
+        # whole grid, so the count check alone proves totality — no need to build
+        # and membership-test the full coordinate set on this hot path.
+        return len(all_chunk_coords) == product(chunks_per_shard)
 
     def _is_complete_shard_write(
         self,
 
@@ -1521,19 +1521,19 @@ def decode_morton_vectorized(
 
 
 @lru_cache(maxsize=16)
-def _morton_order(chunk_shape: tuple[int, ...]) -> npt.NDArray[np.intp]:
-    n_total = product(chunk_shape)
-    n_dims = len(chunk_shape)
+def _morton_order(shape: tuple[int, ...]) -> npt.NDArray[np.intp]:
+    n_total = product(shape)
+    n_dims = len(shape)
     if n_total == 0:
         out = np.empty((0, n_dims), dtype=np.intp)
         out.flags.writeable = False
         return out
 
     # Ceiling hypercube: smallest power-of-2 hypercube whose Morton codes span
-    # all valid coordinates in chunk_shape. (c-1).bit_length() gives the number
+    # all valid coordinates in shape. (c-1).bit_length() gives the number
     # of bits needed to index c values (0 for singleton dims). n_z = 2**total_bits
     # is the size of this hypercube.
-    total_bits = sum((c - 1).bit_length() for c in chunk_shape)
+    total_bits = sum((c - 1).bit_length() for c in shape)
     n_z = 1 << total_bits if total_bits > 0 else 1
 
     # Decode all Morton codes in the ceiling hypercube, then filter to valid coords.
@@ -1544,8 +1544,8 @@ def _morton_order(chunk_shape: tuple[int, ...]) -> npt.NDArray[np.intp]:
         # Ceiling strategy: decode all n_z codes vectorized, filter in-bounds.
         # Works well when the overgeneration ratio n_z/n_total is small (≤4).
         z_values = np.arange(n_z, dtype=np.intp)
-        all_coords = decode_morton_vectorized(z_values, chunk_shape)
-        shape_arr = np.array(chunk_shape, dtype=np.intp)
+        all_coords = decode_morton_vectorized(z_values, shape)
+        shape_arr = np.array(shape, dtype=np.intp)
         valid_mask = np.all(all_coords < shape_arr, axis=1)
         order = all_coords[valid_mask]
     else:
@@ -1554,11 +1554,11 @@ def _morton_order(chunk_shape: tuple[int, ...]) -> npt.NDArray[np.intp]:
         # larger overgeneration penalty for near-miss shapes like (33,33,33).
         # Cost: O(n_total * bits) encode + O(n_total log n_total) sort,
         # vs O(n_z * bits) = O(8 * n_total * bits) for ceiling.
-        grids = np.meshgrid(*[np.arange(c, dtype=np.intp) for c in chunk_shape], indexing="ij")
+        grids = np.meshgrid(*[np.arange(c, dtype=np.intp) for c in shape], indexing="ij")
         all_coords = np.stack([g.ravel() for g in grids], axis=1)
 
         # Encode all coordinates to Morton codes (vectorized).
-        bits_per_dim = tuple((c - 1).bit_length() for c in chunk_shape)
+        bits_per_dim = tuple((c - 1).bit_length() for c in shape)
         max_coord_bits = max(bits_per_dim)
         z_codes = np.zeros(n_total, dtype=np.intp)
         output_bit = 0
@@ -1576,16 +1576,56 @@ def _morton_order(chunk_shape: tuple[int, ...]) -> npt.NDArray[np.intp]:
 
 
 @lru_cache(maxsize=16)
-def _morton_order_keys(chunk_shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
-    return tuple(tuple(int(x) for x in row) for row in _morton_order(chunk_shape))
+def morton_order_coords(shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+    # The grid coordinates in Morton (Z) order, as a cached sequence. The
+    # coordinate set of a finite grid has a known length and is reused in full on
+    # every shard write, so it is built once (vectorized, via `_morton_order`) and
+    # cached per shape rather than recomputed. Indexable and `len`-able; iterate it
+    # directly where an iterator is needed.
+    #
+    # `.tolist()` converts the whole array to native Python ints in one C-level
+    # call; building the tuples row-by-row with `int(x)` is ~9x slower.
+    return tuple(map(tuple, _morton_order(shape).tolist()))
 
 
-def morton_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
-    return iter(_morton_order_keys(tuple(chunk_shape)))
+@lru_cache(maxsize=16)
+def _lexicographic_order(shape: tuple[int, ...]) -> npt.NDArray[np.intp]:
+    # Lexicographic (C-order) coordinates, computed vectorized and cached so that
+    # the sharding codec's per-shard chunk grid is not rebuilt on every call.
+    # Equivalent to `np.array(list(np.ndindex(shape)))` but without the
+    # Python-level iteration over every coordinate.
+    n_dims = len(shape)
+    if n_dims == 0:
+        # A 0-d shard holds a single chunk addressed by the empty coordinate, so
+        # the coordinate array has one row and zero columns. np.indices(()) cannot
+        # express this, so build it directly. Matches list(np.ndindex(())) == [()].
+        order = np.empty((1, 0), dtype=np.intp)
+    else:
+        order = np.indices(shape, dtype=np.intp).reshape(n_dims, -1).T
+    order.flags.writeable = False
+    return order
 
 
-def c_order_iter(chunks_per_shard: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
-    return itertools.product(*(range(x) for x in chunks_per_shard))
+@lru_cache(maxsize=16)
+def lexicographic_order_coords(shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+    # The grid coordinates in lexicographic (row-major / C) order, as a cached
+    # sequence. The coordinate set of a finite grid has a known length and is
+    # reused in full on every shard write, so it is built once (vectorized, via
+    # `_lexicographic_order`) and cached per shape. Indexable and `len`-able;
+    # iterate it directly where an iterator is needed.
+    #
+    # `.tolist()` converts the whole array to native Python ints in one C-level
+    # call; building the tuples row-by-row with `int(x)` is ~9x slower.
+    return tuple(map(tuple, _lexicographic_order(shape).tolist()))
+
+
+@lru_cache(maxsize=16)
+def colexicographic_order_coords(shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+    # The grid coordinates in colexicographic (column-major / F) order, as a cached
+    # sequence: the first axis varies fastest. Equivalent to reversing each axis,
+    # taking lexicographic order, and reversing the coordinates back. Cached per
+    # shape like its siblings so shard writes don't rebuild it.
+    return tuple(c[::-1] for c in lexicographic_order_coords(shape[::-1]))
 
 
 def get_indexer(