diff --git a/changes/4286.bugfix.md b/changes/4286.bugfix.md new file mode 100644 index 0000000000..ee0a60a82c --- /dev/null +++ b/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`. diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index a1b050cb7b..b791a9b6c9 100644 --- a/src/zarr/core/indexing.py +++ b/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 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 04fbdad8c6..02150ca2e5 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -937,6 +937,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: