fix: nested sharding write + leaf_transform property + remove dead unpack_blob

- Add leaf_transform property to ChunkLayout base class (returns
  inner_transform) and override on ShardedChunkLayout (traverses
  nested ShardingCodecs to find innermost codec chain)
- Fix write path complete-overwrite to use layout.leaf_transform
  instead of layout.inner_transform (was using wrong transform for
  nested sharding)
- Fix decode_chunks_from_index to use index.is_sharded instead of
  fragile shape-based is_simple heuristic
- Add _pack_nested to ShardedChunkLayout: groups flat leaf chunks
  by inner shard, packs each group into an inner shard blob, then
  packs into outer shard — produces correct nested shard structure
- Remove dead unpack_blob from all layout classes

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: d-v-b

src/zarr/core/codec_pipeline.py

@@ -844,13 +844,7 @@ def decode_chunks_from_index(
         order=chunk_spec.order,
     )
 
-    # Non-sharded: the transform covers the entire chunk (or more, for rectilinear edges).
-    # Sharded: the transform's inner shape is smaller than chunk_spec in at least one dim.
-    is_simple = all(
-        t >= c for t, c in zip(transform.array_spec.shape, chunk_spec.shape, strict=True)
-    )
-
-    if is_simple:
+    if not index.is_sharded:
         assert len(raw_chunks) == 1
         raw = next(iter(raw_chunks.values()))
         if raw is None:
@@ -1048,6 +1042,16 @@ class ChunkLayout:
     def is_sharded(self) -> bool:
         return False
 
+    @property
+    def leaf_transform(self) -> ChunkTransform:
+        """The codec chain that decodes individual leaf chunks.
+
+        For non-sharded layouts, this is the full transform.
+        For sharded layouts, this traverses nested ShardingCodecs to
+        find the innermost codec chain.
+        """
+        return self.inner_transform
+
     def needed_coords(self, chunk_selection: SelectorTuple) -> set[tuple[int, ...]] | None:
         return None
 
@@ -1069,9 +1073,6 @@ async def pack_and_store_async(self, byte_setter: Any, encoded_chunks: dict[tupl
 
     # -- Low-level helpers --
 
-    def unpack_blob(self, blob: Buffer) -> dict[tuple[int, ...], Buffer | None]:
-        raise NotImplementedError
-
     def pack_blob(
         self, chunk_dict: dict[tuple[int, ...], Buffer | None], prototype: BufferPrototype
     ) -> Buffer | None:
@@ -1116,10 +1117,6 @@ async def pack_and_store_async(self, byte_setter: Any, encoded_chunks: dict[tupl
 
     # -- Low-level --
 
-    def unpack_blob(self, blob: Buffer) -> dict[tuple[int, ...], Buffer | None]:
-        key = (0,) * len(self.chunks_per_shard)
-        return {key: blob}
-
     def pack_blob(
         self, chunk_dict: dict[tuple[int, ...], Buffer | None], prototype: BufferPrototype
     ) -> Buffer | None:
@@ -1172,6 +1169,10 @@ def supports_partial_write(self) -> bool:
         """True when inner codecs are fixed-size, enabling byte-range writes."""
         return self._fixed_size
 
+    @property
+    def leaf_transform(self) -> ChunkTransform:
+        return self._get_leaf_transform()
+
     def _decode_index(self, index_bytes: Buffer) -> Any:
         from zarr.codecs.sharding import _ShardIndex
 
@@ -1186,24 +1187,6 @@ def _encode_index(self, index: Any) -> Buffer:
         assert result is not None
         return result
 
-    def unpack_blob(self, blob: Buffer) -> dict[tuple[int, ...], Buffer | None]:
-        from zarr.codecs.sharding import ShardingCodecIndexLocation
-
-        if self._index_location == ShardingCodecIndexLocation.start:
-            index_bytes = blob[: self._index_size]
-        else:
-            index_bytes = blob[-self._index_size :]
-
-        index = self._decode_index(index_bytes)
-        result: dict[tuple[int, ...], Buffer | None] = {}
-        for chunk_coords in np.ndindex(self.chunks_per_shard):
-            chunk_slice = index.get_chunk_slice(chunk_coords)
-            if chunk_slice is not None:
-                result[chunk_coords] = blob[chunk_slice[0] : chunk_slice[1]]
-            else:
-                result[chunk_coords] = None
-        return result
-
     def pack_blob(
         self, chunk_dict: dict[tuple[int, ...], Buffer | None], prototype: BufferPrototype
     ) -> Buffer | None:
@@ -1444,25 +1427,82 @@ async def resolve_index_async(self, byte_getter: Any, key: str, chunk_selection:
 
     # -- pack and store --
 
+    def _pack_nested(
+        self,
+        encoded_chunks: dict[tuple[int, ...], Buffer | None],
+    ) -> Buffer | None:
+        """Pack flat leaf chunks into a nested shard blob.
+
+        Groups leaf chunks by inner shard, packs each group into an
+        inner shard blob, then packs inner shard blobs into the outer
+        shard blob.
+        """
+        from zarr.codecs.sharding import ShardingCodec
+        from zarr.core.buffer import default_buffer_prototype
+
+        inner_ab = self.inner_transform._ab_codec
+        assert isinstance(inner_ab, ShardingCodec)
+
+        inner_spec = self.inner_transform.array_spec
+        inner_layout = ShardedChunkLayout.from_sharding_codec(inner_ab, inner_spec)
+        inner_cps = inner_layout.chunks_per_shard
+
+        # Group leaf coords by inner shard
+        groups: dict[tuple[int, ...], dict[tuple[int, ...], Buffer | None]] = {}
+        for global_coord, chunk_bytes in encoded_chunks.items():
+            inner_shard_coord = tuple(gc // cps for gc, cps in zip(global_coord, inner_cps, strict=True))
+            leaf_coord = tuple(gc % cps for gc, cps in zip(global_coord, inner_cps, strict=True))
+            if inner_shard_coord not in groups:
+                groups[inner_shard_coord] = {}
+            groups[inner_shard_coord][leaf_coord] = chunk_bytes
+
+        # Pack each group into an inner shard blob
+        proto = default_buffer_prototype()
+        inner_shard_blobs: dict[tuple[int, ...], Buffer | None] = {}
+        for inner_shard_coord in np.ndindex(self.chunks_per_shard):
+            group = groups.get(inner_shard_coord, {})
+            # Fill missing leaf coords with None
+            for lc in np.ndindex(inner_cps):
+                if lc not in group:
+                    group[lc] = None
+            inner_blob = inner_layout.pack_blob(group, proto)
+            inner_shard_blobs[inner_shard_coord] = inner_blob
+
+        # Pack inner shard blobs into outer shard
+        return self.pack_blob(inner_shard_blobs, proto)
+
     def pack_and_store_sync(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        from zarr.codecs.sharding import ShardingCodec
         from zarr.core.buffer import default_buffer_prototype
 
         if all(v is None for v in encoded_chunks.values()):
             byte_setter.delete_sync()  # type: ignore[attr-defined]
             return
-        blob = self.pack_blob(encoded_chunks, default_buffer_prototype())
+
+        # Check for nested sharding
+        if isinstance(self.inner_transform._ab_codec, ShardingCodec):
+            blob = self._pack_nested(encoded_chunks)
+        else:
+            blob = self.pack_blob(encoded_chunks, default_buffer_prototype())
+
         if blob is None:
             byte_setter.delete_sync()  # type: ignore[attr-defined]
         else:
             byte_setter.set_sync(blob)  # type: ignore[attr-defined]
 
     async def pack_and_store_async(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        from zarr.codecs.sharding import ShardingCodec
         from zarr.core.buffer import default_buffer_prototype
 
         if all(v is None for v in encoded_chunks.values()):
             await byte_setter.delete()
             return
-        blob = self.pack_blob(encoded_chunks, default_buffer_prototype())
+
+        if isinstance(self.inner_transform._ab_codec, ShardingCodec):
+            blob = self._pack_nested(encoded_chunks)
+        else:
+            blob = self.pack_blob(encoded_chunks, default_buffer_prototype())
+
         if blob is None:
             await byte_setter.delete()
         else:
@@ -1852,7 +1892,7 @@ async def _process_chunk(
 
             # Phase 1: resolve index (IO)
             if is_complete:
-                index = ShardIndex(key=key, leaf_transform=layout.inner_transform, is_sharded=layout.is_sharded)
+                index = ShardIndex(key=key, leaf_transform=layout.leaf_transform, is_sharded=layout.is_sharded)
             elif layout.is_sharded:
                 async with sem:
                     index = await layout.resolve_index_async(byte_setter, key, chunk_selection=None)  # ALL coords
@@ -1974,7 +2014,7 @@ def write_sync(
 
             # Phase 1: resolve index
             if is_complete:
-                index = ShardIndex(key=key, leaf_transform=layout.inner_transform, is_sharded=layout.is_sharded)
+                index = ShardIndex(key=key, leaf_transform=layout.leaf_transform, is_sharded=layout.is_sharded)
             elif layout.is_sharded:
                 index = layout.resolve_index(bs, key, chunk_selection=None)  # ALL coords
             else: