Merge branch 'main' into glossary

Author: d-v-b

docs/_static/favicon-96x96.png

@@ -85,6 +85,7 @@ theme:
   name: material
   custom_dir: docs/overrides
   logo: _static/logo_bw.png
+  favicon: _static/favicon-96x96.png
 
   palette:
     # Light mode

mkdocs.yml

@@ -1235,24 +1235,28 @@ def basename(self) -> str:
     @property
     def cdata_shape(self) -> tuple[int, ...]:
         """
-        The shape of the chunk grid for this array.
+        The number of chunks along each dimension.
+
+        When sharding is used, this counts inner chunks (not shards) per dimension.
 
         Returns
         -------
         tuple[int, ...]
-            The shape of the chunk grid for this array.
+            The number of chunks along each dimension.
         """
         return self._chunk_grid_shape
 
     @property
     def _chunk_grid_shape(self) -> tuple[int, ...]:
         """
-        The shape of the chunk grid for this array.
+        The number of chunks along each dimension.
+
+        When sharding is used, this counts inner chunks (not shards) per dimension.
 
         Returns
         -------
         tuple[int, ...]
-            The shape of the chunk grid for this array.
+            The number of chunks along each dimension.
         """
         return tuple(starmap(ceildiv, zip(self.shape, self.chunks, strict=True)))
 
@@ -2399,14 +2403,23 @@ def compressors(self) -> tuple[Numcodec, ...] | tuple[BytesBytesCodec, ...]:
     @property
     def cdata_shape(self) -> tuple[int, ...]:
         """
-        The shape of the chunk grid for this array.
+        The number of chunks along each dimension.
+
+        When sharding is used, this counts inner chunks (not shards) per dimension.
         """
         return self.async_array._chunk_grid_shape
 
     @property
     def _chunk_grid_shape(self) -> tuple[int, ...]:
         """
-        The shape of the chunk grid for this array.
+        The number of chunks along each dimension.
+
+        When sharding is used, this counts inner chunks (not shards) per dimension.
+
+        Returns
+        -------
+        tuple[int, ...]
+            The number of chunks along each dimension.
         """
         return self.async_array._chunk_grid_shape
 

src/zarr/core/array.py

@@ -126,11 +126,14 @@ def normalize_chunks(chunks: Any, shape: tuple[int, ...], typesize: int) -> tupl
         chunks = tuple(int(chunks) for _ in shape)
 
     # handle dask-style chunks (iterable of iterables)
-    if all(isinstance(c, (tuple | list)) for c in chunks):
-        # take first chunk size for each dimension
-        chunks = tuple(
-            c[0] for c in chunks
-        )  # TODO: check/error/warn for irregular chunks (e.g. if c[0] != c[1:-1])
+    if all(isinstance(c, (tuple, list)) for c in chunks):
+        for i, c in enumerate(chunks):
+            if any(x != y for x, y in itertools.pairwise(c[:-1])) or (len(c) > 1 and c[-1] > c[0]):
+                raise ValueError(
+                    f"Irregular chunk sizes in dimension {i}: {tuple(c)}. "
+                    "Only uniform chunks (with an optional smaller final chunk) are supported."
+                )
+        chunks = tuple(c[0] for c in chunks)
 
     # handle bad dimensionality
     if len(chunks) > len(shape):

src/zarr/core/chunk_grids.py

@@ -263,6 +263,8 @@ async def read_batch(
                 chunk_array_batch, batch_info, strict=False
             ):
                 if chunk_array is not None:
+                    if drop_axes:
+                        chunk_array = chunk_array.squeeze(axis=drop_axes)
                     out[out_selection] = chunk_array
                 else:
                     out[out_selection] = fill_value_or_default(chunk_spec)
@@ -285,7 +287,7 @@ async def read_batch(
             ):
                 if chunk_array is not None:
                     tmp = chunk_array[chunk_selection]
-                    if drop_axes != ():
+                    if drop_axes:
                         tmp = tmp.squeeze(axis=drop_axes)
                     out[out_selection] = tmp
                 else:
@@ -324,7 +326,7 @@ def _merge_chunk_array(
         else:
             chunk_value = value[out_selection]
             # handle missing singleton dimensions
-            if drop_axes != ():
+            if drop_axes:
                 item = tuple(
                     None  # equivalent to np.newaxis
                     if idx in drop_axes

src/zarr/core/codec_pipeline.py

@@ -35,6 +35,10 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None:
         ((30, None, None), (100, 20, 10), 1, (30, 20, 10)),
         ((30, 20, None), (100, 20, 10), 1, (30, 20, 10)),
         ((30, 20, 10), (100, 20, 10), 1, (30, 20, 10)),
+        # dask-style chunks (uniform with optional smaller final chunk)
+        (((100, 100, 100), (50, 50)), (300, 100), 1, (100, 50)),
+        (((100, 100, 50),), (250,), 1, (100,)),
+        (((100,),), (100,), 1, (100,)),
         # auto chunking
         (None, (100,), 1, (100,)),
         (-1, (100,), 1, (100,)),
@@ -52,3 +56,8 @@ def test_normalize_chunks_errors() -> None:
         normalize_chunks("foo", (100,), 1)
     with pytest.raises(ValueError):
         normalize_chunks((100, 10), (100,), 1)
+    # dask-style irregular chunks should raise
+    with pytest.raises(ValueError, match="Irregular chunk sizes"):
+        normalize_chunks(((10, 20, 30),), (60,), 1)
+    with pytest.raises(ValueError, match="Irregular chunk sizes"):
+        normalize_chunks(((100, 100), (10, 20)), (200, 30), 1)

tests/test_chunk_grids.py

@@ -490,7 +490,8 @@ def test_invalid_shard_shape() -> None:
     with pytest.raises(
         ValueError,
         match=re.escape(
-            "The array's `chunk_shape` (got (16, 16)) needs to be divisible by the shard's inner `chunk_shape` (got (9,))."
+            "The array's `chunk_shape` (got (16, 16)) needs to be divisible "
+            "by the shard's inner `chunk_shape` (got (9,))."
         ),
     ):
         zarr.create_array(
@@ -501,3 +502,56 @@ def test_invalid_shard_shape() -> None:
             dtype=np.dtype("uint8"),
             fill_value=0,
         )
+
+
+@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.
+
+    Mixed integer/list indexing on sharded arrays should return the same
+    shape and data as on equivalent chunked arrays.
+    """
+    import numpy as np
+
+    data = np.arange(200 * 100 * 10, dtype=np.uint8).reshape(200, 100, 10)
+
+    chunked = zarr.create_array(
+        store,
+        name="chunked",
+        shape=(200, 100, 10),
+        dtype=np.uint8,
+        chunks=(200, 100, 1),
+        overwrite=True,
+    )
+    chunked[:, :, :] = data
+
+    sharded = zarr.create_array(
+        store,
+        name="sharded",
+        shape=(200, 100, 10),
+        dtype=np.uint8,
+        chunks=(200, 100, 1),
+        shards=(200, 100, 10),
+        overwrite=True,
+    )
+    sharded[:, :, :] = data
+
+    # Mixed integer + list indexing
+    c = chunked[0:10, 0, [0, 1]]  # type: ignore[index]
+    s = sharded[0:10, 0, [0, 1]]  # type: ignore[index]
+    assert c.shape == s.shape == (10, 2), (  # type: ignore[union-attr]
+        f"Expected (10, 2), got chunked={c.shape}, sharded={s.shape}"  # type: ignore[union-attr]
+    )
+    np.testing.assert_array_equal(c, s)
+
+    # Multiple integer axes
+    c2 = chunked[0, 0, [0, 1, 2]]  # type: ignore[index]
+    s2 = sharded[0, 0, [0, 1, 2]]  # type: ignore[index]
+    assert c2.shape == s2.shape == (3,)  # type: ignore[union-attr]
+    np.testing.assert_array_equal(c2, s2)
+
+    # Slice + integer + slice
+    c3 = chunked[0:5, 1, 0:3]
+    s3 = sharded[0:5, 1, 0:3]
+    assert c3.shape == s3.shape == (5, 3)  # type: ignore[union-attr]
+    np.testing.assert_array_equal(c3, s3)