Fix is_total_slice for size-1 dimensions (#1800)

Closes #1730

Co-authored-by: Ryan Abernathey <[email protected]>
Co-authored-by: Joe Hamman <[email protected]>

Author: 3 people

docs/release.rst

@@ -20,6 +20,8 @@ Unreleased
 
 Enhancements
 ~~~~~~~~~~~~
+* Performance improvement for reading and writing chunks if any of the dimensions is size 1. :issue:`1730`
+  By :user:`Deepak Cherian <dcherian>`.
 
 
 Docs

zarr/tests/test_util.py

@@ -89,6 +89,15 @@ def test_is_total_slice():
     assert not is_total_slice((slice(0, 50), slice(0, 50)), (100, 100))
     assert not is_total_slice((slice(0, 100, 2), slice(0, 100)), (100, 100))
 
+    # size-1 dimension edge-case
+    # https://github.com/zarr-developers/zarr-python/issues/1730
+    assert is_total_slice((slice(0, 1),), (1,))
+    # this is an equivalent selection (without a slice)
+    assert is_total_slice((0,), (1,))
+    # same for multidimensional selection
+    assert is_total_slice((slice(0, 1), slice(0, 10)), (1, 10))
+    assert is_total_slice((0, slice(0, 10)), (1, 10))
+
     with pytest.raises(TypeError):
         is_total_slice("foo", (100,))
 

zarr/util.py

@@ -234,8 +234,17 @@ def is_total_slice(item, shape: Tuple[int]) -> bool:
     if isinstance(item, tuple):
         return all(
             (
-                isinstance(it, slice)
-                and ((it == slice(None)) or ((it.stop - it.start == sh) and (it.step in [1, None])))
+                (
+                    isinstance(it, slice)
+                    and (
+                        (it == slice(None))
+                        or ((it.stop - it.start == sh) and (it.step in [1, None]))
+                    )
+                )
+                # The only scalar edge case, indexing with int 0 along a size-1 dimension
+                # is identical to a total slice
+                # https://github.com/zarr-developers/zarr-python/issues/1730
+                or (isinstance(it, int) and it == 0 and sh == 1)
             )
             for it, sh in zip(item, shape)
         )