Add test and make max gap and max coalesce size config options

Author: aldenks

src/zarr/codecs/sharding.py

@@ -569,9 +569,10 @@ async def _load_partial_shard_maybe(
     def _coalesce_chunks(
         self,
         chunks: list[_ChunkCoordsByteSlice],
-        max_gap_bytes: int = 2**20,  # 1MiB
-        coalesce_max_bytes: int = 100 * 2**20,  # 100MiB
     ) -> list[list[_ChunkCoordsByteSlice]]:
+        max_gap_bytes = config.get("sharding.read.coalesce_max_gap_bytes")
+        coalesce_max_bytes = config.get("sharding.read.coalesce_max_bytes")
+
         sorted_chunks = sorted(chunks, key=lambda c: c.byte_slice.start)
 
         groups = []
@@ -590,15 +591,6 @@ def _coalesce_chunks(
 
         groups.append(current_group)
 
-        from pprint import pprint
-
-        pprint(
-            [
-                f"{len(g)} chunks, {(g[-1].byte_slice.stop - g[0].byte_slice.start) / 1e6:.1f}MB"
-                for g in groups
-            ]
-        )
-
         return groups
 
     async def _get_group_bytes(

src/zarr/core/config.py

@@ -108,6 +108,12 @@ def enable_gpu(self) -> ConfigSet:
             },
             "async": {"concurrency": 10, "timeout": None},
             "threading": {"max_workers": None},
+            "sharding": {
+                "read": {
+                    "coalesce_max_bytes": 100 * 2**20,  # 100MiB
+                    "coalesce_max_gap_bytes": 2**20,  # 1MiB
+                }
+            },
             "json_indent": 2,
             "codec_pipeline": {
                 "path": "zarr.core.codec_pipeline.BatchedCodecPipeline",

tests/test_codecs/test_sharding.py

@@ -243,10 +243,39 @@ def test_partial_shard_read_performance(store: Store) -> None:
             }
         )
 
-    with open("zarr-python-partial-shard-read-performance-no-coalesce.json", "w") as f:
+    with open("zarr-python-partial-shard-read-performance.json", "w") as f:
         json.dump(experiments, f)
 
 
+@pytest.mark.parametrize("index_location", ["start", "end"])
+@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"])
+def test_sharding_multiple_chunks_partial_shard_read(
+    store: Store, index_location: ShardingCodecIndexLocation
+) -> None:
+    array_shape = (8, 64)
+    shard_shape = (4, 32)
+    chunk_shape = (2, 4)
+
+    data = np.arange(np.prod(array_shape), dtype="float32").reshape(array_shape)
+
+    a = zarr.create_array(
+        StorePath(store),
+        shape=data.shape,
+        chunks=chunk_shape,
+        shards={"shape": shard_shape, "index_location": index_location},
+        compressors=BloscCodec(cname="lz4"),
+        dtype=data.dtype,
+        fill_value=1,
+    )
+    a[:] = data
+
+    # Reads 2.5 (3 full, one partial) chunks each from 2 shards (a subset of both shards)
+    assert np.allclose(a[0, 22:42], np.arange(22, 42, dtype="float32"))
+
+    # Reads 2 chunks from both shards along dimension 0
+    assert np.allclose(a[:, 0], np.arange(0, data.size, array_shape[1], dtype="float32"))
+
+
 @pytest.mark.parametrize(
     "array_fixture",
     [