diff --git a/changes/2782.bugfix.rst b/changes/2782.bugfix.rst new file mode 100644 index 0000000000..baea970f51 --- /dev/null +++ b/changes/2782.bugfix.rst @@ -0,0 +1 @@ +Optimize full chunk writes diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index f1226821ba..733b2464ac 100644 --- a/src/zarr/core/indexing.py +++ b/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 --git a/tests/test_array.py b/tests/test_array.py index abc048423c..80ff8444fc 100644 --- a/tests/test_array.py +++ b/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 --git a/tests/test_indexing.py b/tests/test_indexing.py index 30d0d75f22..932c32f1ae 100644 --- a/tests/test_indexing.py +++ b/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))