Fix orthogonal indexing with scalar. (#1947)

Co-authored-by: David Stansby <[email protected]>

Author: dcherian

docs/release.rst

@@ -18,6 +18,17 @@ Release notes
    See `GH1777 <https://github.com/zarr-developers/zarr-python/issues/1777>`_ for more details on the upcoming
    3.0 release.
 
+.. _release_2.18.3:
+
+2.18.3
+------
+
+Maintenance
+~~~~~~~~~~~
+* Fix a regression when using orthogonal indexing with a scalar.
+  By :user:`Deepak Cherian <dcherian>` :issue:`1931`
+
+
 .. _release_2.18.2:
 
 Enhancements

zarr/core.py

@@ -2053,7 +2053,12 @@ def _process_chunk(
                     if isinstance(cdata, UncompressedPartialReadBufferV3):
                         cdata = cdata.read_full()
                     chunk = ensure_ndarray_like(cdata).view(self._dtype)
-                    chunk = chunk.reshape(self._chunks, order=self._order)
+                    # dest.shape is not self._chunks when a dimensions is squeezed out
+                    # For example, assume self._chunks = (5, 5, 1)
+                    # and the selection is [:, :, 0]
+                    # Then out_selection is (slice(5), slice(5))
+                    # See https://github.com/zarr-developers/zarr-python/issues/1931
+                    chunk = chunk.reshape(dest.shape, order=self._order)
                     np.copyto(dest, chunk)
                 return
 

zarr/tests/test_core.py

@@ -3206,3 +3206,20 @@ def test_object_array_indexing():
     elem = [1, 3]
     arr[1] = elem
     assert arr[1] == elem
+
+
+@pytest.mark.parametrize("shape", ((1, 1, 1), (5, 5, 1), (1, 5, 5)))
+def test_scalar_orthogonal_indexing(shape):
+    # regression test for https://github.com/zarr-developers/zarr-python/issues/1931
+    store = zarr.MemoryStore({})
+    data = np.random.randint(0, 255, shape)
+    arr = zarr.zeros(
+        shape=shape, chunks=shape[:-1] + (1,), compressor=None, store=store, dtype="u1"
+    )
+    arr[:, :, :] = data
+    store.close()
+
+    zf = zarr.open(store, "r")
+    assert_array_equal(zf[0, :, :], data[0, :, :])
+    assert_array_equal(zf[:, 0, :], data[:, 0, :])
+    assert_array_equal(zf[:, :, 0], data[:, :, 0])