More extensive orthogonal indexing in get/setitem (#1333)

* More extensive orthogonal indexing in get/setitem

Added pass-through to orthogonal indexing for the following cases:
* index is iterable of integers
* index is iterable of length ndim, with each element being a slice, integer, or list. Maximum one list.

* Add test cases for indexing with single integer iterable

---------

Co-authored-by: Josh Moore <josh@openmicroscopy.org>

Author: AndreasAlbertQC

docs/release.rst

@@ -6,15 +6,18 @@ Release notes
     # to document your changes. On releases it will be
     # re-indented so that it does not show up in the notes.
 
-    .. _unreleased:
+.. _unreleased:
 
-    Unreleased
-    ----------
+Unreleased
+----------
 
 ..
     # .. warning::
     #    Pre-release! Use :command:`pip install --pre zarr` to evaluate this release.
 
+* Implement more extensive fallback of getitem/setitem for orthogonal indexing.
+  By :user:`Andreas Albert <AndreasAlbertQC>` :issue:`1029`.
+
 .. _release_2.14.2:
 
 2.14.2

docs/tutorial.rst

@@ -634,6 +634,13 @@ For convenience, the orthogonal indexing functionality is also available via the
 Any combination of integer, slice, 1D integer array and/or 1D Boolean array can
 be used for orthogonal indexing.
 
+If the index contains at most one iterable, and otherwise contains only slices and integers,
+orthogonal indexing is also available directly on the array:
+
+    >>> z = zarr.array(np.arange(15).reshape(3, 5))
+    >>> all(z.oindex[[0, 2], :] == z[[0, 2], :])
+    True
+
 Indexing fields in structured arrays
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

zarr/core.py

@@ -28,6 +28,7 @@
     err_too_many_indices,
     is_contiguous_selection,
     is_pure_fancy_indexing,
+    is_pure_orthogonal_indexing,
     is_scalar,
     pop_fields,
 )
@@ -817,6 +818,8 @@ def __getitem__(self, selection):
         fields, pure_selection = pop_fields(selection)
         if is_pure_fancy_indexing(pure_selection, self.ndim):
             result = self.vindex[selection]
+        elif is_pure_orthogonal_indexing(pure_selection, self.ndim):
+            result = self.get_orthogonal_selection(pure_selection, fields=fields)
         else:
             result = self.get_basic_selection(pure_selection, fields=fields)
         return result
@@ -1387,6 +1390,8 @@ def __setitem__(self, selection, value):
         fields, pure_selection = pop_fields(selection)
         if is_pure_fancy_indexing(pure_selection, self.ndim):
             self.vindex[selection] = value
+        elif is_pure_orthogonal_indexing(pure_selection, self.ndim):
+            self.set_orthogonal_selection(pure_selection, value, fields=fields)
         else:
             self.set_basic_selection(pure_selection, value, fields=fields)
 

zarr/indexing.py

@@ -101,6 +101,26 @@ def is_pure_fancy_indexing(selection, ndim):
     )
 
 
+def is_pure_orthogonal_indexing(selection, ndim):
+    if not ndim:
+        return False
+
+    # Case 1: Selection is a single iterable of integers
+    if is_integer_list(selection) or is_integer_array(selection, ndim=1):
+        return True
+
+    # Case two: selection contains either zero or one integer iterables.
+    # All other selection elements are slices or integers
+    return (
+            isinstance(selection, tuple) and len(selection) == ndim and
+            sum(is_integer_list(elem) or is_integer_array(elem) for elem in selection) <= 1 and
+            all(
+                is_integer_list(elem) or is_integer_array(elem)
+                or isinstance(elem, slice) or isinstance(elem, int) for
+                elem in selection)
+    )
+
+
 def normalize_integer_selection(dim_sel, dim_len):
 
     # normalize type to int

zarr/tests/test_indexing.py

@@ -283,8 +283,6 @@ def test_get_basic_selection_2d():
     for selection in bad_selections:
         with pytest.raises(IndexError):
             z.get_basic_selection(selection)
-        with pytest.raises(IndexError):
-            z[selection]
     # check fallback on fancy indexing
     fancy_selection = ([0, 1], [0, 1])
     np.testing.assert_array_equal(z[fancy_selection], [0, 11])
@@ -317,14 +315,179 @@ def test_fancy_indexing_fallback_on_get_setitem():
     )
 
 
-def test_fancy_indexing_doesnt_mix_with_slicing():
-    z = zarr.zeros((20, 20))
-    with pytest.raises(IndexError):
-        z[[1, 2, 3], :] = 2
-    with pytest.raises(IndexError):
-        np.testing.assert_array_equal(
-            z[[1, 2, 3], :], 0
+@pytest.mark.parametrize("index,expected_result",
+                         [
+                             # Single iterable of integers
+                             (
+                                [0, 1],
+                                [[0, 1, 2],
+                                 [3, 4, 5]]
+                             ),
+                             # List first, then slice
+                             (
+                                     ([0, 1], slice(None)),
+                                     [[0, 1, 2],
+                                      [3, 4, 5]]
+                             ),
+                             # List first, then slice
+                             (
+                                     ([0, 1], slice(1, None)),
+                                     [[1, 2],
+                                      [4, 5]]
+                             ),
+                             # Slice first, then list
+                             (
+                                     (slice(0, 2), [0, 2]),
+                                     [[0, 2],
+                                      [3, 5]]
+                             ),
+                             # Slices only
+                             (
+                                     (slice(0, 2), slice(0, 2)),
+                                     [[0, 1],
+                                      [3, 4]]
+                             ),
+                             # List with repeated index
+                             (
+                                     ([1, 0, 1], slice(1, None)),
+                                     [[4, 5],
+                                      [1, 2],
+                                      [4, 5]]
+                             ),
+                             # 1D indexing
+                             (
+                                     ([1, 0, 1]),
+                                     [
+                                         [3, 4, 5],
+                                         [0, 1, 2],
+                                         [3, 4, 5]
+                                     ]
+                             )
+
+                         ])
+def test_orthogonal_indexing_fallback_on_getitem_2d(index, expected_result):
+    """
+    Tests the orthogonal indexing fallback on __getitem__ for a 2D matrix.
+
+    In addition to checking expected behavior, all indexing
+    is also checked against numpy.
+    """
+    # [0, 1, 2],
+    # [3, 4, 5],
+    # [6, 7, 8]
+    a = np.arange(9).reshape(3, 3)
+    z = zarr.array(a)
+
+    np.testing.assert_array_equal(z[index], a[index], err_msg="Indexing disagrees with numpy")
+    np.testing.assert_array_equal(z[index], expected_result)
+
+
+@pytest.mark.parametrize("index,expected_result",
+                         [
+                             # Single iterable of integers
+                             (
+                                [0, 1],
+                                [[[0,  1,  2],
+                                  [3,  4,  5],
+                                  [6,  7,  8]],
+                                 [[9, 10, 11],
+                                  [12, 13, 14],
+                                  [15, 16, 17]]]
+                             ),
+                             # One slice, two integers
+                             (
+                                     (slice(0, 2), 1, 1),
+                                     [4, 13]
+                             ),
+                             # One integer, two slices
+                             (
+                                     (slice(0, 2), 1, slice(0, 2)),
+                                     [[3, 4], [12, 13]]
+                             ),
+                             # Two slices and a list
+                             (
+                                     (slice(0, 2), [1, 2], slice(0, 2)),
+                                     [[[3, 4], [6, 7]], [[12, 13], [15, 16]]]
+                             ),
+                         ])
+def test_orthogonal_indexing_fallback_on_getitem_3d(index, expected_result):
+    """
+    Tests the orthogonal indexing fallback on __getitem__ for a 3D matrix.
+
+    In addition to checking expected behavior, all indexing
+    is also checked against numpy.
+    """
+    #   [[[ 0,  1,  2],
+    #     [ 3,  4,  5],
+    #     [ 6,  7,  8]],
+
+    #    [[ 9, 10, 11],
+    #     [12, 13, 14],
+    #     [15, 16, 17]],
+
+    #    [[18, 19, 20],
+    #     [21, 22, 23],
+    #     [24, 25, 26]]]
+    a = np.arange(27).reshape(3, 3, 3)
+    z = zarr.array(a)
+
+    np.testing.assert_array_equal(z[index], a[index], err_msg="Indexing disagrees with numpy")
+    np.testing.assert_array_equal(z[index], expected_result)
+
+
+@pytest.mark.parametrize(
+    "index,expected_result",
+    [
+        # Single iterable of integers
+        (
+            [0, 1],
+            [
+                [1, 1, 1],
+                [1, 1, 1],
+                [0, 0, 0]
+            ]
+        ),
+        # List and slice combined
+        (
+                ([0, 1], slice(1, 3)),
+                [[0, 1, 1],
+                 [0, 1, 1],
+                 [0, 0, 0]]
+        ),
+        # Index repetition is ignored on setitem
+        (
+                ([0, 1, 1, 1, 1, 1, 1], slice(1, 3)),
+                [[0, 1, 1],
+                 [0, 1, 1],
+                 [0, 0, 0]]
+        ),
+        # Slice with step
+        (
+                ([0, 2], slice(None, None, 2)),
+                [[1, 0, 1],
+                 [0, 0, 0],
+                 [1, 0, 1]]
         )
+    ]
+)
+def test_orthogonal_indexing_fallback_on_setitem_2d(index, expected_result):
+    """
+    Tests the orthogonal indexing fallback on __setitem__ for a 3D matrix.
+
+    In addition to checking expected behavior, all indexing
+    is also checked against numpy.
+    """
+    # Slice + fancy index
+    a = np.zeros((3, 3))
+    z = zarr.array(a)
+    z[index] = 1
+    a[index] = 1
+    np.testing.assert_array_equal(
+        z, expected_result
+    )
+    np.testing.assert_array_equal(
+        z, a, err_msg="Indexing disagrees with numpy"
+    )
 
 
 def test_fancy_indexing_doesnt_mix_with_implicit_slicing():
@@ -335,12 +498,6 @@ def test_fancy_indexing_doesnt_mix_with_implicit_slicing():
         np.testing.assert_array_equal(
             z2[[1, 2, 3], [1, 2, 3]], 0
         )
-    with pytest.raises(IndexError):
-        z2[[1, 2, 3]] = 2
-    with pytest.raises(IndexError):
-        np.testing.assert_array_equal(
-            z2[[1, 2, 3]], 0
-        )
     with pytest.raises(IndexError):
         z2[..., [1, 2, 3]] = 2
     with pytest.raises(IndexError):
@@ -770,6 +927,33 @@ def test_set_orthogonal_selection_3d():
         _test_set_orthogonal_selection_3d(v, a, z, ix0, ix1, ix2)
 
 
+def test_orthogonal_indexing_fallback_on_get_setitem():
+    z = zarr.zeros((20, 20))
+    z[[1, 2, 3], [1, 2, 3]] = 1
+    np.testing.assert_array_equal(
+        z[:4, :4],
+        [
+            [0, 0, 0, 0],
+            [0, 1, 0, 0],
+            [0, 0, 1, 0],
+            [0, 0, 0, 1],
+        ],
+    )
+    np.testing.assert_array_equal(
+        z[[1, 2, 3], [1, 2, 3]], 1
+    )
+    # test broadcasting
+    np.testing.assert_array_equal(
+        z[1, [1, 2, 3]], [1, 0, 0]
+    )
+    # test 1D fancy indexing
+    z2 = zarr.zeros(5)
+    z2[[1, 2, 3]] = 1
+    np.testing.assert_array_equal(
+        z2, [0, 1, 1, 1, 0]
+    )
+
+
 def _test_get_coordinate_selection(a, z, selection):
     expect = a[selection]
     actual = z.get_coordinate_selection(selection)