fix(sharding): address adversarial-review findings on clipped grids

Fixes from a multi-agent adversarial review of the semi-regular grid
change:

- nchunks_initialized used floor division for chunks-per-shard, which
  undercounts on clipped grids and reports 0 when chunk_shape exceeds
  the shard shape; use ceiling division to match the inner grid
- the scalar-broadcast write memo in _encode_partial_sync gated on
  ArraySpec object identity, which never holds because _get_chunk_spec
  is uncached (#3054); gate on shape equality so the memo fires again
  (one encode instead of one per complete chunk)
- the rectilinear clipping test used shard edges [[6, 4]], which
  normalize to a regular grid; use [[4, 6]] with the rectilinear config
  flag so RectilinearChunkGridMetadata is actually exercised

Adds regression tests for the first two.

Assisted-by: ClaudeCode:claude-fable-5

Author: d-v-b

src/zarr/codecs/sharding.py

@@ -763,7 +763,6 @@ def _encode_partial_sync(
         """
         shard_shape = shard_spec.shape
         chunks_per_shard = self._get_chunks_per_shard(shard_spec)
-        nominal_chunk_spec = self._get_chunk_spec(shard_spec)
         get_chunk_spec = self._make_chunk_spec_getter(shard_spec)
         inner_transform = self._get_inner_chunk_transform(shard_spec)
 
@@ -818,7 +817,7 @@ def _encode_partial_sync(
 
         for chunk_coords, chunk_sel, out_sel, is_complete_chunk in indexer:
             chunk_spec = get_chunk_spec(chunk_coords)
-            if is_scalar and is_complete_chunk and chunk_spec is nominal_chunk_spec:
+            if is_scalar and is_complete_chunk and chunk_spec.shape == self.chunk_shape:
                 if scalar_complete_result is _sentinel:
                     scalar_complete_result = merge_and_encode_chunk(
                         None,
@@ -1510,9 +1509,12 @@ def _make_chunk_spec_getter(
 
         Inner chunks clipped by the shard boundary are encoded at their clipped
         shape, so their spec differs from the nominal `chunk_shape` spec.
-        Chunks with the nominal shape all share one spec object, which keeps
-        per-spec caches downstream (e.g. `ChunkTransform._resolve_specs`)
-        effective and makes `spec is nominal_spec` a valid uniformity check.
+        Within one getter, chunks with the nominal shape all share one spec
+        object, which keeps per-spec caches downstream (e.g.
+        `ChunkTransform._resolve_specs`) effective. Note that `_get_chunk_spec`
+        is not cached (see #3054), so the nominal spec here is NOT the same
+        object as one obtained from a separate `_get_chunk_spec` call — callers
+        must compare specs by shape, not identity.
         """
         nominal = self._get_chunk_spec(shard_spec)
         if self._is_evenly_divided(shard_spec.shape):

src/zarr/core/array.py

@@ -5345,8 +5345,10 @@ async def _nchunks_initialized(
     if array.shards is None:
         chunks_per_shard = 1
     else:
+        # Inner chunks are clipped by the shard shape, so the count per shard
+        # is the ceiling division of the shard shape by the chunk shape.
         chunks_per_shard = product(
-            tuple(a // b for a, b in zip(array.shards, array.chunks, strict=True))
+            tuple(ceildiv(a, b) for a, b in zip(array.shards, array.chunks, strict=True))
         )
     return (await _nshards_initialized(array)) * chunks_per_shard
 

tests/test_codecs/test_sharding.py

@@ -887,17 +887,27 @@ def test_sharding_clipped_inner_chunks_byte_layout() -> None:
 
 def test_sharding_clipped_inner_chunks_rectilinear() -> None:
     """A rectilinear outer grid whose shard edges are not divisible by the inner
-    chunk shape round-trips: the inner grid is clipped per shard."""
-    arr = zarr.create_array(
-        {}, shape=(10,), shards=[[6, 4]], chunks=(4,), dtype="uint8", fill_value=0
-    )
+    chunk shape round-trips: the inner grid is clipped per shard.
+
+    Shard edges `[4, 6]` cannot be normalized to a regular grid (the larger edge
+    comes last), so this exercises `RectilinearChunkGridMetadata` for real —
+    unlike e.g. `[6, 4]`, which coincides with a clipped regular grid of size 6.
+    """
+    from zarr.core.metadata.v3 import RectilinearChunkGridMetadata
+
+    with zarr.config.set({"array.rectilinear_chunks": True}):
+        arr = zarr.create_array(
+            {}, shape=(10,), shards=[[4, 6]], chunks=(4,), dtype="uint8", fill_value=0
+        )
+    assert isinstance(arr.metadata, ArrayV3Metadata)
+    assert isinstance(arr.metadata.chunk_grid, RectilinearChunkGridMetadata)
     data = np.arange(10, dtype="uint8")
     arr[...] = data
     assert np.array_equal(arr[...], data)
-    assert arr.read_chunk_sizes == ((4, 2, 4),)
+    assert arr.read_chunk_sizes == ((4, 4, 2),)
 
-    arr[5:9] = 42
-    data[5:9] = 42
+    arr[3:9] = 42
+    data[3:9] = 42
     assert np.array_equal(arr[...], data)
 
 
@@ -935,6 +945,55 @@ def test_sharding_clipped_inner_chunk_sizes() -> None:
     assert arr.nchunks == 6
 
 
+def test_sharding_clipped_nchunks_initialized() -> None:
+    """`nchunks_initialized` counts ceiling-division inner chunks per shard.
+
+    Regression test: it previously used floor division, which undercounts for
+    clipped grids and reports 0 when the inner chunk shape exceeds the shard
+    shape.
+    """
+    arr = zarr.create_array({}, shape=(16,), shards=(8,), chunks=(3,), dtype="uint8", fill_value=0)
+    arr[...] = np.arange(16, dtype="uint8")
+    # 2 shards, each holding ceil(8/3) = 3 inner chunks
+    assert arr.nchunks_initialized == 6
+
+    arr2 = zarr.create_array({}, shape=(11,), shards=(4,), chunks=(9,), dtype="uint8", fill_value=0)
+    arr2[...] = np.arange(11, dtype="uint8")
+    # 3 shards, each holding one inner chunk clipped to the shard shape
+    assert arr2.nchunks_initialized == 3
+
+
+def test_sharding_scalar_write_memoizes_complete_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A scalar broadcast partial write encodes the repeated complete inner
+    chunk once, not once per chunk.
+
+    Regression test: the memo gate must compare chunk specs by shape — an
+    object-identity gate never fires because `_get_chunk_spec` is uncached.
+    """
+    from zarr.core import chunk_utils
+
+    calls = 0
+    real_encode = chunk_utils.ChunkTransform.encode_chunk
+
+    def counting_encode(self: Any, *args: Any, **kwargs: Any) -> Any:
+        nonlocal calls
+        calls += 1
+        return real_encode(self, *args, **kwargs)
+
+    with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}):
+        arr = zarr.create_array(
+            {}, shape=(100,), shards=(100,), chunks=(10,), dtype="uint8", fill_value=0
+        )
+        monkeypatch.setattr(chunk_utils.ChunkTransform, "encode_chunk", counting_encode)
+        arr[5:] = 7
+        expected = np.full(100, 7, dtype="uint8")
+        expected[:5] = 0
+        assert np.array_equal(arr[...], expected)
+    # one merge for the partial edge chunk, one memoized encode shared by the
+    # nine complete chunks, one shard-index encode
+    assert calls <= 4, f"scalar write encoded {calls} times; memo is not firing"
+
+
 @pytest.mark.parametrize("store", ["local"], indirect=["store"])
 def test_sharding_mixed_integer_list_indexing(store: Store) -> None:
     """Regression test for https://github.com/zarr-developers/zarr-python/issues/3691.