test: update read/write plan tests for RangeByteRequest and nested sharding

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

Author: d-v-b

tests/test_read_plan.py

@@ -4,7 +4,7 @@
 
 The model:
 - A shard (or non-sharded chunk) is a flat key-value space:
-  ``coords → ByteRange`` within one store key.
+  ``coords → RangeByteRequest`` within one store key.
 - Index resolution (possibly recursive for nested sharding) produces
   this flat mapping.
 - The pipeline then filters to needed coords, fetches those byte ranges,
@@ -17,47 +17,30 @@
 from typing import Any
 
 import numpy as np
+import pytest
 
 import zarr
+from zarr.abc.store import RangeByteRequest
 
 # ---------------------------------------------------------------------------
 # Data model
 # ---------------------------------------------------------------------------
 
 
-@dataclass(frozen=True)
-class ByteRange:
-    """A contiguous byte range within a store value."""
-
-    offset: int
-    length: int
-
-
 @dataclass(frozen=True)
 class ShardIndex:
     """Flat mapping from inner chunk coordinates to byte ranges.
 
-    Produced by resolving the shard index (and any nested indexes).
-    For non-sharded chunks, contains a single entry mapping ``(0,)``
-    to ``None`` (meaning: read the full value).
-
-    Parameters
-    ----------
-    key : str
-        The store key for this shard/chunk.
-    chunks : dict
-        Mapping from inner chunk coords to their byte range within
-        the blob at ``key``. A value of ``None`` means the chunk
-        is absent (fill value).
+    Uses ``RangeByteRequest`` from ``zarr.abc.store`` for byte ranges.
     """
 
     key: str
-    chunks: dict[tuple[int, ...], ByteRange | None] = field(default_factory=dict)
+    chunks: dict[tuple[int, ...], RangeByteRequest | None] = field(default_factory=dict)
 
     @property
     def nbytes_data(self) -> int:
         """Total data bytes across all present chunks."""
-        return sum(r.length for r in self.chunks.values() if r is not None)
+        return sum(r.end - r.start for r in self.chunks.values() if r is not None)
 
     def filter(self, needed: set[tuple[int, ...]] | None = None) -> ShardIndex:
         """Return a new ShardIndex with only the needed coords."""
@@ -176,10 +159,10 @@ def test_single_inner_chunk(self) -> None:
         idx = indices[0]
         # Only the needed inner chunk
         assert len(idx.chunks) == 1
-        coords = list(idx.chunks.keys())[0]
+        coords = next(iter(idx.chunks.keys()))
         byte_range = idx.chunks[coords]
         assert byte_range is not None
-        assert byte_range.length == 10
+        assert byte_range.end - byte_range.start == 10
 
     def test_two_inner_chunks(self) -> None:
         arr, _store = _create_and_fill(shape=(100,), chunks=(10,), shards=(100,))
@@ -224,7 +207,7 @@ def test_single_inner_chunk_compressed(self) -> None:
         assert len(idx.chunks) == 1
         byte_range = list(idx.chunks.values())[0]
         assert byte_range is not None
-        assert byte_range.length > 0
+        assert byte_range.end - byte_range.start > 0
 
 
 # ---------------------------------------------------------------------------
@@ -234,7 +217,7 @@ def test_single_inner_chunk_compressed(self) -> None:
 
 class TestNestedShardedIndex:
     """For nested sharding, index resolution recurses through levels
-    but produces the same flat coords → ByteRange mapping.
+    but produces the same flat coords → RangeByteRequest mapping.
     """
 
     @staticmethod
@@ -263,7 +246,7 @@ def test_single_leaf_chunk(self) -> None:
         assert len(idx.chunks) == 1
         byte_range = list(idx.chunks.values())[0]
         assert byte_range is not None
-        assert byte_range.length == 10
+        assert byte_range.end - byte_range.start == 10
 
     def test_full_inner_shard(self) -> None:
         """One full inner shard (50 bytes = 5 leaf chunks)."""
@@ -299,109 +282,14 @@ def test_all_leaf_chunks(self) -> None:
 # ---------------------------------------------------------------------------
 
 
-def _resolve_shard_index(
-    layout: Any,
-    chunk_selection: Any,
-    shard_blob: Any | None,
-    base_offset: int = 0,
-) -> dict[tuple[int, ...], ByteRange | None]:
-    """Recursively resolve a flat coords → ByteRange mapping for a shard.
-
-    For fixed-size codecs, byte ranges are computed from coordinates alone
-    (shard_blob can be None). For variable-size codecs, the shard blob
-    is needed to read the index. For nested sharding, recurses into
-    inner shards.
-    """
-    from zarr.codecs.sharding import ShardingCodec
-    from zarr.core.codec_pipeline import ShardedChunkLayout
-
-    needed_coords = layout.needed_coords(chunk_selection)
-    if needed_coords is None:
-        return {}
-
-    if layout._fixed_size:
-        chunk_spec = layout.inner_transform.array_spec
-        chunk_byte_length = layout.inner_chunk_byte_length(chunk_spec)
-        return {
-            coords: ByteRange(
-                offset=base_offset + layout.chunk_byte_offset(coords, chunk_byte_length),
-                length=chunk_byte_length,
-            )
-            for coords in needed_coords
-        }
-
-    # Variable-size: need the blob to read the index
-    assert shard_blob is not None
-    chunk_dict = layout.unpack_blob(shard_blob)
-
-    # Check for nested sharding
-    inner_ab = layout.inner_transform._ab_codec
-    is_nested = isinstance(inner_ab, ShardingCodec)
-
-    if not is_nested:
-        # Leaf level: read byte ranges from the index
-        from zarr.codecs.sharding import ShardingCodecIndexLocation
-
-        if layout._index_location == ShardingCodecIndexLocation.start:
-            index_bytes = shard_blob[: layout._index_size]
-        else:
-            index_bytes = shard_blob[-layout._index_size :]
-        index = layout._decode_index(index_bytes)
-
-        result: dict[tuple[int, ...], ByteRange | None] = {}
-        for coords in needed_coords:
-            chunk_slice = index.get_chunk_slice(coords)
-            if chunk_slice is not None:
-                start, end = chunk_slice
-                result[coords] = ByteRange(offset=base_offset + start, length=end - start)
-            else:
-                result[coords] = None
-        return result
-
-    # Nested: resolve inner shard indexes and flatten
-    from zarr.codecs.sharding import ShardingCodecIndexLocation
-    from zarr.core.chunk_grids import ChunkGrid as _ChunkGrid
-    from zarr.core.indexing import get_indexer
-
-    if layout._index_location == ShardingCodecIndexLocation.start:
-        index_bytes = shard_blob[: layout._index_size]
-    else:
-        index_bytes = shard_blob[-layout._index_size :]
-    outer_index = layout._decode_index(index_bytes)
-
-    inner_indexer = get_indexer(
-        chunk_selection,
-        shape=layout.chunk_shape,
-        chunk_grid=_ChunkGrid.from_sizes(layout.chunk_shape, layout.inner_chunk_shape),
-    )
-
-    inner_spec = layout.inner_transform.array_spec
-    inner_layout = ShardedChunkLayout.from_sharding_codec(inner_ab, inner_spec)
-
-    flat: dict[tuple[int, ...], ByteRange | None] = {}
-    for inner_coords, inner_sel, _, _ in inner_indexer:
-        chunk_slice = outer_index.get_chunk_slice(inner_coords)
-        if chunk_slice is None:
-            continue
-        start, end = chunk_slice
-        inner_blob = shard_blob[start:end]
-        inner_chunks = _resolve_shard_index(
-            inner_layout, inner_sel, inner_blob, base_offset=base_offset + start
-        )
-        # Prefix leaf coords with the outer coords to make them globally unique
-        for leaf_coords, byte_range in inner_chunks.items():
-            flat[inner_coords + leaf_coords] = byte_range
-
-    return flat
-
-
 def _resolve_indices(arr: zarr.Array, selection: Any) -> list[ShardIndex]:
     """Given an array and a selection, resolve ShardIndex for each chunk/shard.
 
-    Each ShardIndex is a flat mapping from inner chunk coords to byte ranges,
-    regardless of how many levels of nesting exist.
+    Uses the pipeline's ``layout.resolve_index`` to get the flat
+    coords → RangeByteRequest mapping, then wraps in the test's
+    ShardIndex (which has extra helper methods).
     """
-    from zarr.core.codec_pipeline import PhasedCodecPipeline, ShardedChunkLayout
+    from zarr.core.codec_pipeline import PhasedCodecPipeline
     from zarr.core.indexing import BasicIndexer
 
     aa = arr._async_array
@@ -423,27 +311,13 @@ def _resolve_indices(arr: zarr.Array, selection: Any) -> list[ShardIndex]:
             continue
 
         layout = pipeline.layout
+        store_path = aa.store_path / key
 
-        if not layout.is_sharded:
-            indices.append(ShardIndex(key=key, chunks={(0,) * len(chunk_coords): None}))
-            continue
-
-        assert isinstance(layout, ShardedChunkLayout)
-
-        if layout._fixed_size:
-            # No blob needed
-            chunks = _resolve_shard_index(layout, chunk_selection, shard_blob=None)
-        else:
-            # Need the blob to read indexes
-            from zarr.core.buffer import default_buffer_prototype
-
-            store_path = aa.store_path / key
-            shard_blob = store_path.get_sync(prototype=default_buffer_prototype())
-            if shard_blob is None:
-                indices.append(ShardIndex(key=key))
-                continue
-            chunks = _resolve_shard_index(layout, chunk_selection, shard_blob)
+        # Use the pipeline's resolve_index — it handles all cases
+        # (non-sharded, fixed-size, variable-size)
+        pipeline_index = layout.resolve_index(store_path, key, chunk_selection=chunk_selection)
 
-        indices.append(ShardIndex(key=key, chunks=chunks))
+        # Convert pipeline ShardIndex to test ShardIndex
+        indices.append(ShardIndex(key=pipeline_index.key, chunks=dict(pipeline_index.chunks)))
 
     return indices