Repository: dask/dask
Pull request #10854 (2024-01-23 12:57:03)

numpy 2.0: fix slicing by uint64 array
- Part of #10593

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     15 files  ± 0       15 suites  ±0   3h 26m 55s :stopwatch: + 2m 42s
 12 975 tests + 7   12 044 :white_check_mark: + 8     929 :zzz: ±0  2 :x:  - 1 
160 051 runs  +91  143 471 :white_check_mark: +98  16 578 :zzz:  - 6  2 :x:  - 1 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/20775475190).

Results for commit f4ba6580. ± Comparison against base commit 62e5bed4.

<details>
  <summary>This pull request <b>removes</b> 2 and <b>adds</b> 9 tests. <i>Note that renamed tests count towards both.</i></summary>

```
dask.array.tests.test_array_core ‑ test_index_with_integer_types
dask.array.tests.test_array_core ‑ test_slice_with_integer_types
```

```
dask.array.tests.test_array_core ‑ test_index_with_integer_types[int32]
dask.array.tests.test_array_core ‑ test_index_with_integer_types[int64]
dask.array.tests.test_array_core ‑ test_index_with_integer_types[int]
dask.array.tests.test_array_core ‑ test_index_with_integer_types[uint32]
dask.array.tests.test_array_core ‑ test_index_with_integer_types[uint64]
dask.array.tests.test_array_core ‑ test_slice_with_integer_types[int32]
dask.array.tests.test_array_core ‑ test_slice_with_integer_types[int64]
dask.array.tests.test_array_core ‑ test_slice_with_integer_types[uint32]
dask.array.tests.test_array_core ‑ test_slice_with_integer_types[uint64]
```
</details>

[test-results]:data:application/gzip;base64,H4sIAH/Nr2UC/03MSw7CIBSF4a00jB3cXi4U3IyhUBJiH4bCyLh3EZV2+H8nOU/mwzzt7Nr14tKxPYd0lMvRpLCtn0aqVNZUd9RD69uera0IRAfew6OgRt3ImzAXwgZTjFssAkViXuuzBBD9r9szcRoO/D73Ugzqb6fr2udnuy1LSCWYp9FIocAYJTlqsFyMDgGd5aNCSaDU5Dxw9noDTJ8giBkBAAA=


Patch:
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index e2a6b34a0..20df6fe02 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -559,9 +559,14 @@ def slicing_plan(chunks, index):
 
     if not is_arraylike(index):
         index = np.asanyarray(index)
-    cum_chunks = cached_cumsum(chunks)
 
-    cum_chunks = asarray_safe(cum_chunks, like=index)
+    cum_chunks_tup = cached_cumsum(chunks)
+    cum_chunks = asarray_safe(cum_chunks_tup, like=index)
+    if cum_chunks.dtype.kind != "f":  # Don't cast NaN chunks to int
+        # This is important when index.dtype=uint64 (or uint32 on 32-bit hosts) to
+        # prevent accidental automatic casting during `index - cum_chunks` below
+        cum_chunks = cum_chunks.astype(index.dtype)
+
     # this dispactches to the array library
     chunk_locations = np.searchsorted(cum_chunks, index, side="right")
 
@@ -634,8 +639,6 @@ def take(outname, inname, chunks, index, itemsize, axis=0):
             PerformanceWarning,
             stacklevel=6,
         )
-    if not is_arraylike(index):
-        index = np.asarray(index)
 
     # Check for chunks from the plan that would violate the user's
     # configured chunk size.


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 9f9e78a40..f4b190f2f 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -3112,25 +3112,19 @@ def test_slice_with_floats():
         d[[1, 1.5]]
 
 
-def test_slice_with_integer_types():
+@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.uint32, np.uint64])
+def test_slice_with_integer_types(dtype):
     x = np.arange(10)
     dx = da.from_array(x, chunks=5)
-    inds = np.array([0, 3, 6], dtype="u8")
+    inds = np.array([0, 3, 6], dtype=dtype)
     assert_eq(dx[inds], x[inds])
-    assert_eq(dx[inds.astype("u4")], x[inds.astype("u4")])
-
-    inds = np.array([0, 3, 6], dtype=np.int64)
-    assert_eq(dx[inds], x[inds])
-    assert_eq(dx[inds.astype("u4")], x[inds.astype("u4")])
 
 
-def test_index_with_integer_types():
+@pytest.mark.parametrize("cls", [int, np.int32, np.int64, np.uint32, np.uint64])
+def test_index_with_integer_types(cls):
     x = np.arange(10)
     dx = da.from_array(x, chunks=5)
-    inds = int(3)
-    assert_eq(dx[inds], x[inds])
-
-    inds = np.int64(3)
+    inds = cls(3)
     assert_eq(dx[inds], x[inds])