fix: unsigned index arrays raise IndexError because np.diff wraps and uint64 promotes (#4286)

* fix:  array comparison and add test for unsorted indexing with unsigned dtypes

* rewrite comment

* fix: resolve IndexError in integer array indexing with unsigned dtypes

---------

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

Author: selmanozleyen

changes/4286.bugfix.md

@@ -0,0 +1,6 @@
+Fixed integer array indexing with unsigned index dtypes. An unsorted index such as
+`np.array([3, 0], dtype="uint8")` spanning more than one chunk raised `IndexError`, because
+the order check used `np.diff`, which wraps on unsigned dtypes and misclassified a
+descending selection as increasing. Separately, a `uint64` index raised `IndexError` on both
+`array[...]` and `array.vindex[...]` — sorted or not — because `uint64` promotes to
+`float64` against a signed chunk offset. Index arrays are now cast to `intp`.

src/zarr/core/indexing.py

@@ -715,8 +715,8 @@ class Order(Enum):
 
     @staticmethod
     def check(a: npt.NDArray[Any]) -> Order:
-        diff = np.diff(a)
-        diff_positive = diff >= 0
+        # compare, don't subtract: np.diff wraps on unsigned dtypes
+        diff_positive = a[1:] >= a[:-1]
         n_diff_positive = np.count_nonzero(diff_positive)
         all_increasing = n_diff_positive == len(diff_positive)
         any_increasing = n_diff_positive > 0
@@ -769,6 +769,8 @@ def __init__(
         dim_sel = np.asanyarray(dim_sel)
         if not is_integer_array(dim_sel, 1):
             raise IndexError("integer arrays in an orthogonal selection must be 1-dimensional only")
+        # uint64 promotes to float against the signed chunk offset
+        dim_sel = dim_sel.astype(np.intp, copy=False)
 
         nitems = len(dim_sel)
         g = dim_grid
@@ -1207,6 +1209,11 @@ def __init__(
                 "(coordinate) array per dimension of the target array, "
                 f"got {selection!r}"
             )
+        # keep indices integral: uint64 against a signed offset promotes to float
+        selection_normalized = cast(
+            "CoordinateSelectionNormalized",
+            tuple(np.asarray(s, dtype=np.intp) for s in selection_normalized),
+        )
 
         # Optimization for a single sorted, in-bounds, 1-D integer coordinate array over a
         # regular (fixed-size) chunk grid. The general path below makes several full passes over

tests/test_indexing.py

@@ -939,6 +939,16 @@ def test_orthogonal_indexing_edge_cases(store: StorePath) -> None:
     assert_array_equal(expect, actual)
 
 
+@pytest.mark.parametrize("dtype", ["int8", "int64", "uint8", "uint16", "uint32", "uint64"])
+def test_unsorted_index_unsigned_dtype(store: StorePath, dtype: str) -> None:
+    a = np.arange(8).reshape(4, 2)
+    z = zarr_array_from_numpy_array(store, a, chunk_shape=(2, 1))
+    rows = np.array([3, 0], dtype=dtype)
+
+    assert_array_equal(a[[3, 0], :], z[rows, :])
+    assert_array_equal(a[[3, 0], [1, 0]], z.vindex[rows, np.array([1, 0], dtype=dtype)])
+
+
 def _test_set_orthogonal_selection(
     v: npt.NDArray[np.int_], a: npt.NDArray[Any], z: Array, selection: OrthogonalSelection
 ) -> None: