Repository: dask/dask
Pull request #11335 (2024-08-20 17:54:41)

Add ``crick`` back to Python 3.11+ CI builds
`numpy=1.24+` has been added, so we can include `crick` in these CI builds again 

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._

     13 files  ±0       13 suites  ±0   3h 1m 34s :stopwatch: - 1m 51s
 13 162 tests ±0   12 106 :white_check_mark: ±  0   1 056 :zzz: ±  0  0 :x: ±0 
137 151 runs  ±0  118 145 :white_check_mark: +111  19 006 :zzz:  - 111  0 :x: ±0 

Results for commit bb35ce9d. ± Comparison against base commit 28a38328.

[test-results]:data:application/gzip;base64,H4sIAMPfxGYC/02Nyw6DIBBFf8Ww7oLhIUx/pkGEhNRHg7Ay/fciWnV5zs2cWYkPg1vIswH+aMiSQ7qoz9GkME8bU42iqLKmfYeW/fm1ZGs3yYC2l3yHT72Ul/MmDMXRU7gY53iYmKc9rUDCwWcaNAh5yiONtP6r7paufC/beRxDKkC6jkvrsGdeA6JhyKkRIFvvtJKWWc1RKUaRfH8iizBqGgEAAA==


Patch:
diff --git a/continuous_integration/environment-3.11.yaml b/continuous_integration/environment-3.11.yaml
index 5806843d9..438f71e5f 100644
--- a/continuous_integration/environment-3.11.yaml
+++ b/continuous_integration/environment-3.11.yaml
@@ -49,9 +49,7 @@ dependencies:
   - httpretty
   - aiohttp
   - s3fs>=2021.9.0
-  # Need a new `crick` release with support for `numpy=1.24+`
-  # https://github.com/dask/crick/issues/25
-  # - crick
+  - crick
   - cytoolz
   - distributed
   - ipython
diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
index 341fa3722..e0e8f432f 100644
--- a/continuous_integration/environment-3.12.yaml
+++ b/continuous_integration/environment-3.12.yaml
@@ -49,9 +49,7 @@ dependencies:
   - httpretty
   - aiohttp
   - s3fs>=2021.9.0
-  # Need a new `crick` release with support for `numpy=1.24+`
-  # https://github.com/dask/crick/issues/25
-  # - crick
+  - crick
   - cytoolz
   - distributed
   - ipython
diff --git a/dask/array/core.py b/dask/array/core.py
index a35c05f76..75515993a 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -11,6 +11,7 @@ import traceback
 import uuid
 import warnings
 from bisect import bisect
+from collections import defaultdict
 from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence
 from functools import partial, reduce, wraps
 from itertools import product, zip_longest
@@ -5540,20 +5541,38 @@ def _vindex_array(x, dict_indexes):
     token = tokenize(x, flat_indexes)
     out_name = "vindex-merge-" + token
 
+    max_chunk_point_dimensions = reduce(
+        mul, map(max, [c for i, c in zip(flat_indexes, x.chunks) if i is not None])
+    )
+
     points = list()
     for i, idx in enumerate(zip(*[i for i in flat_indexes if i is not None])):
         block_idx = [bisect(b, ind) - 1 for b, ind in zip(bounds2, idx)]
         inblock_idx = [
             ind - bounds2[k][j] for k, (ind, j) in enumerate(zip(idx, block_idx))
         ]
-        points.append((i, tuple(block_idx), tuple(inblock_idx)))
+        points.append(
+            (
+                divmod(i, max_chunk_point_dimensions)[1],
+                tuple(block_idx),
+                tuple(inblock_idx),
+                (i // max_chunk_point_dimensions,) + tuple(block_idx),
+            )
+        )
 
     chunks = [c for i, c in zip(flat_indexes, x.chunks) if i is None]
-    chunks.insert(0, (len(points),) if points else (0,))
+    n_chunks, remainder = divmod(len(points), max_chunk_point_dimensions)
+    chunks.insert(
+        0,
+        (max_chunk_point_dimensions,) * n_chunks
+        + ((remainder,) if remainder > 0 else ())
+        if points
+        else (0,),
+    )
     chunks = tuple(chunks)
 
     if points:
-        per_block = groupby(1, points)
+        per_block = groupby(3, points)
         per_block = {k: v for k, v in per_block.items() if v}
 
         other_blocks = list(
@@ -5571,23 +5590,28 @@ def _vindex_array(x, dict_indexes):
         vindex_merge_name = "vindex-merge-" + token
         dsk = {}
         for okey in other_blocks:
+            merge_inputs = defaultdict(list)
+            merge_indexer = defaultdict(list)
             for i, key in enumerate(per_block):
                 dsk[keyname(name, i, okey)] = (
                     _vindex_transpose,
                     (
                         _vindex_slice,
-                        (x.name,) + interleave_none(okey, key),
+                        (x.name,) + interleave_none(okey, tuple(list(key)[1:])),
                         interleave_none(
                             full_slices, list(zip(*pluck(2, per_block[key])))
                         ),
                     ),
                     axis,
                 )
-            dsk[keyname(vindex_merge_name, 0, okey)] = (
-                _vindex_merge,
-                [list(pluck(0, per_block[key])) for key in per_block],
-                [keyname(name, i, okey) for i in range(len(per_block))],
-            )
+                merge_inputs[key[0]].append(keyname(name, i, okey))
+                merge_indexer[key[0]].append(list(pluck(0, per_block[key])))
+            for i in sorted(merge_inputs.keys()):
+                dsk[keyname(vindex_merge_name, i, okey)] = (
+                    _vindex_merge,
+                    merge_indexer[i],
+                    merge_inputs[i],
+                )
 
         result_1d = Array(
             HighLevelGraph.from_collections(out_name, dsk, dependencies=[x]),


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index d975cfced..144671518 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -3088,6 +3088,21 @@ def test_vindex_nd():
     assert_eq(result, x.T)
 
 
+@pytest.mark.parametrize("size", [0, 1])
+def test_vindex_preserve_chunksize(size):
+    np_arr = np.random.rand(10_000 * 40).reshape(100, 100, 40)
+    arr = da.from_array(np_arr, chunks=(50, 50, 20))
+    indices_2d = np.random.choice(np.arange(100), size=(10000 + size, 2))
+    idx1 = indices_2d[:, 0]
+    idx2 = indices_2d[:, 0]
+    result = arr.vindex[idx1, idx2, slice(None)]
+    assert result.chunks == (
+        (2500, 2500, 2500, 2500) + ((1,) if size else ()),
+        (20, 20),
+    )
+    assert_eq(result, np_arr[idx1, idx2, :])
+
+
 def test_vindex_negative():
     x = np.arange(10)
     d = da.from_array(x, chunks=(5, 5))