feat: recursive resolve_index for nested sharding with leaf_transform

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

Author: d-v-b

src/zarr/core/codec_pipeline.py

@@ -1346,49 +1346,209 @@ def pack_blob(
 
         return template.combine(buffers)
 
+    def _get_leaf_transform(self) -> ChunkTransform:
+        """Get the innermost (leaf) transform, traversing nested ShardingCodecs."""
+        from zarr.codecs.sharding import ShardingCodec
+
+        transform = self.inner_transform
+        while isinstance(transform._ab_codec, ShardingCodec):
+            inner_sc = transform._ab_codec
+            inner_spec = inner_sc._get_chunk_spec(transform.array_spec)
+            inner_evolved = tuple(
+                c.evolve_from_array_spec(array_spec=inner_spec) for c in inner_sc.codecs
+            )
+            transform = ChunkTransform(codecs=inner_evolved, array_spec=inner_spec)
+        return transform
+
+    def _fetch_index_from_blob(self, blob: Buffer) -> Any:
+        """Parse the shard index from an in-memory blob."""
+        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 :]
+        return self._decode_index(index_bytes)
+
     # -- Phase 1: resolve index --
 
     def resolve_index(self, byte_getter: Any, key: str, chunk_selection: SelectorTuple | None = None) -> ShardIndex:
         from zarr.abc.store import RangeByteRequest
+        from zarr.codecs.sharding import ShardingCodec
 
         shard_index = self._fetch_index_sync(byte_getter)
         if shard_index is None:
-            return ShardIndex(key=key, leaf_transform=self.inner_transform, is_sharded=True)
+            return ShardIndex(key=key, leaf_transform=self._get_leaf_transform(), is_sharded=True)
 
         if chunk_selection is not None:
             needed = self.needed_coords(chunk_selection)
         else:
             needed = set(np.ndindex(self.chunks_per_shard))
 
-        chunks: dict[tuple[int, ...], RangeByteRequest | None] = {}
-        for coord in needed:  # type: ignore[union-attr]
-            chunk_slice = shard_index.get_chunk_slice(coord)
-            if chunk_slice is not None:
-                chunks[coord] = RangeByteRequest(chunk_slice[0], chunk_slice[1])
-            else:
-                chunks[coord] = None
-        return ShardIndex(key=key, chunks=chunks, leaf_transform=self.inner_transform, is_sharded=True)
+        inner_ab = self.inner_transform._ab_codec
+        if not isinstance(inner_ab, ShardingCodec):
+            # Non-nested: same as before
+            chunks: dict[tuple[int, ...], RangeByteRequest | None] = {}
+            for coord in needed:  # type: ignore[union-attr]
+                chunk_slice = shard_index.get_chunk_slice(coord)
+                if chunk_slice is not None:
+                    chunks[coord] = RangeByteRequest(chunk_slice[0], chunk_slice[1])
+                else:
+                    chunks[coord] = None
+            return ShardIndex(key=key, chunks=chunks, leaf_transform=self.inner_transform, is_sharded=True)
+
+        # NESTED sharding
+        from zarr.core.buffer import default_buffer_prototype
+        from zarr.core.chunk_grids import ChunkGrid as _ChunkGrid
+        from zarr.core.indexing import get_indexer
+
+        leaf_transform = self._get_leaf_transform()
+
+        # Build inner layout for the nested ShardingCodec
+        inner_spec = self.inner_transform.array_spec
+        inner_layout = ShardedChunkLayout.from_sharding_codec(inner_ab, inner_spec)
+
+        # Build inner indexer to determine which inner shards overlap selection
+        sel = chunk_selection if chunk_selection is not None else tuple(
+            slice(0, s) for s in self.chunk_shape
+        )
+        inner_indexer = get_indexer(
+            sel,
+            shape=self.chunk_shape,
+            chunk_grid=_ChunkGrid.from_sizes(self.chunk_shape, self.inner_chunk_shape),
+        )
+
+        flat: dict[tuple[int, ...], RangeByteRequest | None] = {}
+        for inner_coords, inner_sel, _, _ in inner_indexer:
+            chunk_slice = shard_index.get_chunk_slice(inner_coords)
+            if chunk_slice is None:
+                continue
+            start, end = chunk_slice
+
+            # Fetch the inner shard blob
+            inner_blob = byte_getter.get_sync(
+                prototype=default_buffer_prototype(),
+                byte_range=RangeByteRequest(start, end),
+            )
+            if inner_blob is None:
+                continue
+
+            # Parse inner shard index
+            inner_index = inner_layout._fetch_index_from_blob(inner_blob)
+            if inner_index is None:
+                continue
+
+            # Determine which leaf chunks within this inner shard are needed
+            inner_needed = inner_layout.needed_coords(inner_sel)
+            if inner_needed is None:
+                inner_needed = set(np.ndindex(inner_layout.chunks_per_shard))
+
+            # Translate coords and byte ranges
+            for leaf_coord in inner_needed:
+                leaf_slice = inner_index.get_chunk_slice(leaf_coord)
+                global_coord = tuple(
+                    ic * cps + lc
+                    for ic, cps, lc in zip(
+                        inner_coords, inner_layout.chunks_per_shard, leaf_coord, strict=True
+                    )
+                )
+                if leaf_slice is not None:
+                    abs_start = start + leaf_slice[0]
+                    abs_end = start + leaf_slice[1]
+                    flat[global_coord] = RangeByteRequest(abs_start, abs_end)
+                else:
+                    flat[global_coord] = None
+
+        return ShardIndex(key=key, chunks=flat, leaf_transform=leaf_transform, is_sharded=True)
 
     async def resolve_index_async(self, byte_getter: Any, key: str, chunk_selection: SelectorTuple | None = None) -> ShardIndex:
         from zarr.abc.store import RangeByteRequest
+        from zarr.codecs.sharding import ShardingCodec
 
         shard_index = await self._fetch_index(byte_getter)
         if shard_index is None:
-            return ShardIndex(key=key, leaf_transform=self.inner_transform, is_sharded=True)
+            return ShardIndex(key=key, leaf_transform=self._get_leaf_transform(), is_sharded=True)
 
         if chunk_selection is not None:
             needed = self.needed_coords(chunk_selection)
         else:
             needed = set(np.ndindex(self.chunks_per_shard))
 
-        chunks: dict[tuple[int, ...], RangeByteRequest | None] = {}
-        for coord in needed:  # type: ignore[union-attr]
-            chunk_slice = shard_index.get_chunk_slice(coord)
-            if chunk_slice is not None:
-                chunks[coord] = RangeByteRequest(chunk_slice[0], chunk_slice[1])
-            else:
-                chunks[coord] = None
-        return ShardIndex(key=key, chunks=chunks, leaf_transform=self.inner_transform, is_sharded=True)
+        inner_ab = self.inner_transform._ab_codec
+        if not isinstance(inner_ab, ShardingCodec):
+            # Non-nested: same as before
+            chunks: dict[tuple[int, ...], RangeByteRequest | None] = {}
+            for coord in needed:  # type: ignore[union-attr]
+                chunk_slice = shard_index.get_chunk_slice(coord)
+                if chunk_slice is not None:
+                    chunks[coord] = RangeByteRequest(chunk_slice[0], chunk_slice[1])
+                else:
+                    chunks[coord] = None
+            return ShardIndex(key=key, chunks=chunks, leaf_transform=self.inner_transform, is_sharded=True)
+
+        # NESTED sharding
+        from zarr.core.buffer import default_buffer_prototype
+        from zarr.core.chunk_grids import ChunkGrid as _ChunkGrid
+        from zarr.core.indexing import get_indexer
+
+        leaf_transform = self._get_leaf_transform()
+
+        # Build inner layout for the nested ShardingCodec
+        inner_spec = self.inner_transform.array_spec
+        inner_layout = ShardedChunkLayout.from_sharding_codec(inner_ab, inner_spec)
+
+        # Build inner indexer to determine which inner shards overlap selection
+        sel = chunk_selection if chunk_selection is not None else tuple(
+            slice(0, s) for s in self.chunk_shape
+        )
+        inner_indexer = get_indexer(
+            sel,
+            shape=self.chunk_shape,
+            chunk_grid=_ChunkGrid.from_sizes(self.chunk_shape, self.inner_chunk_shape),
+        )
+
+        flat: dict[tuple[int, ...], RangeByteRequest | None] = {}
+        for inner_coords, inner_sel, _, _ in inner_indexer:
+            chunk_slice = shard_index.get_chunk_slice(inner_coords)
+            if chunk_slice is None:
+                continue
+            start, end = chunk_slice
+
+            # Fetch the inner shard blob
+            inner_blob = await byte_getter.get(
+                prototype=default_buffer_prototype(),
+                byte_range=RangeByteRequest(start, end),
+            )
+            if inner_blob is None:
+                continue
+
+            # Parse inner shard index
+            inner_index = inner_layout._fetch_index_from_blob(inner_blob)
+            if inner_index is None:
+                continue
+
+            # Determine which leaf chunks within this inner shard are needed
+            inner_needed = inner_layout.needed_coords(inner_sel)
+            if inner_needed is None:
+                inner_needed = set(np.ndindex(inner_layout.chunks_per_shard))
+
+            # Translate coords and byte ranges
+            for leaf_coord in inner_needed:
+                leaf_slice = inner_index.get_chunk_slice(leaf_coord)
+                global_coord = tuple(
+                    ic * cps + lc
+                    for ic, cps, lc in zip(
+                        inner_coords, inner_layout.chunks_per_shard, leaf_coord, strict=True
+                    )
+                )
+                if leaf_slice is not None:
+                    abs_start = start + leaf_slice[0]
+                    abs_end = start + leaf_slice[1]
+                    flat[global_coord] = RangeByteRequest(abs_start, abs_end)
+                else:
+                    flat[global_coord] = None
+
+        return ShardIndex(key=key, chunks=flat, leaf_transform=leaf_transform, is_sharded=True)
 
     # -- Phase 2: fetch chunk data --
 

tests/test_codec_pipeline.py

@@ -366,3 +366,28 @@ async def test_read_missing_chunks_true_fills(pipeline_class: str) -> None:
     )
     # Don't write anything
     np.testing.assert_array_equal(arr[:], np.full(20, -999.0))
+
+
+async def test_nested_sharding_roundtrip(pipeline_class: str) -> None:
+    """Nested sharding: data survives write/read roundtrip."""
+    from zarr.codecs.bytes import BytesCodec
+    from zarr.codecs.sharding import ShardingCodec
+
+    inner_sharding = ShardingCodec(chunk_shape=(10,), codecs=[BytesCodec()])
+    outer_sharding = ShardingCodec(chunk_shape=(50,), codecs=[inner_sharding])
+
+    store = MemoryStore()
+    arr = zarr.create_array(
+        store=store,
+        shape=(100,),
+        dtype="uint8",
+        chunks=(100,),
+        compressors=None,
+        fill_value=0,
+        serializer=outer_sharding,
+    )
+    data = np.arange(100, dtype="uint8")
+    arr[:] = data
+    np.testing.assert_array_equal(arr[:], data)
+    # Partial read
+    np.testing.assert_array_equal(arr[40:60], data[40:60])