Optimize is_total_slice for chunksize == 1 (#2782)

Author: dcherian

Diff for: changes/2782.bugfix.rst

@@ -0,0 +1 @@
+Optimize full chunk writes

Diff for: src/zarr/core/indexing.py

@@ -1373,10 +1373,13 @@ def is_total_slice(item: Selection, shape: ChunkCoords) -> bool:
         item = (item,)
     if isinstance(item, tuple):
         return all(
-            isinstance(dim_sel, slice)
-            and (
-                (dim_sel == slice(None))
-                or ((dim_sel.stop - dim_sel.start == dim_len) and (dim_sel.step in [1, None]))
+            (isinstance(dim_sel, int) and dim_len == 1)
+            or (
+                isinstance(dim_sel, slice)
+                and (
+                    (dim_sel == slice(None))
+                    or ((dim_sel.stop - dim_sel.start == dim_len) and (dim_sel.step in [1, None]))
+                )
             )
             for dim_sel, dim_len in zip(item, shape, strict=False)
         )

Diff for: tests/test_array.py

@@ -5,6 +5,7 @@
 import re
 from itertools import accumulate
 from typing import TYPE_CHECKING, Any, Literal
+from unittest import mock
 
 import numcodecs
 import numpy as np
@@ -1328,3 +1329,11 @@ async def test_scalar_array() -> None:
     assert arr[...] == 1.5
     assert arr[()] == 1.5
     assert arr.shape == ()
+
+
+async def test_orthogonal_set_total_slice() -> None:
+    """Ensure that a whole chunk overwrite does not read chunks"""
+    store = MemoryStore()
+    array = zarr.create_array(store, shape=(20, 20), chunks=(1, 2), dtype=int, fill_value=-1)
+    with mock.patch("zarr.storage.MemoryStore.get", side_effect=ValueError):
+        array[0, slice(4, 10)] = np.arange(6)

Diff for: tests/test_indexing.py

@@ -19,6 +19,7 @@
     OrthogonalSelection,
     Selection,
     _iter_grid,
+    is_total_slice,
     make_slice_selection,
     normalize_integer_selection,
     oindex,
@@ -1953,3 +1954,8 @@ def test_vectorized_indexing_incompatible_shape(store) -> None:
     )
     with pytest.raises(ValueError, match="Attempting to set"):
         arr[np.array([1, 2]), np.array([1, 2])] = np.array([[-1, -2], [-3, -4]])
+
+
+def test_is_total_slice():
+    assert is_total_slice((0, slice(4, 6)), (1, 2))
+    assert is_total_slice((slice(0, 1, None), slice(4, 6)), (1, 2))