Repository: dask/dask
Pull request #11267 (2024-08-01 14:46:20)

Use the shuffle algorithm for take
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


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 33m 4s :stopwatch: + 4m 47s
 13 135 tests + 6   12 067 :white_check_mark: + 6   1 068 :zzz: ±0  0 :x: ±0 
162 873 runs  +78  140 084 :white_check_mark: +81  22 789 :zzz:  - 3  0 :x: ±0 

Results for commit e4860ebe. ± Comparison against base commit 8dcf306e.

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

```
dask.array.tests.test_slicing ‑ test_slicing_plan[chunks0-index0-expected0]
dask.array.tests.test_slicing ‑ test_slicing_plan[chunks1-index1-expected1]
dask.array.tests.test_slicing ‑ test_slicing_plan[chunks2-index2-expected2]
```

```
dask.array.tests.test_shuffle ‑ test_incompatible_indexer
dask.array.tests.test_shuffle ‑ test_oob_axis
dask.array.tests.test_shuffle ‑ test_oob_indexer
dask.array.tests.test_shuffle ‑ test_shuffle[indexer0-chunks0]
dask.array.tests.test_shuffle ‑ test_shuffle[indexer1-chunks1]
dask.array.tests.test_shuffle ‑ test_shuffle[indexer2-chunks2]
dask.array.tests.test_shuffle ‑ test_shuffle[indexer3-chunks3]
dask.array.tests.test_shuffle ‑ test_shuffle_larger_array
dask.array.tests.test_shuffle ‑ test_unknown_chunk_sizes
```
</details>

[test-results]:data:application/gzip;base64,H4sIAHunq2YC/02OSw6DIBBAr2JYdzEI8ullGoQxIfXTIKya3r2Iim4meW8mL/Mlgx9xJc+Gdo+GrMnHi1wKJvpl3riVimeVt7HsGWXdya81WVuOQMhLvv1nkyBUdYPxY3ZQBYawhMOENJe0aJVkB9c0BygP7HJPt/kpfbpbuvC9bJdp8jEDQa4EYI+WAgjHtZJaO0FdDwbRKIV5MtZx8vsDzWDwMRoBAAA=


Patch:
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index 1ed630a25..f9add043a 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -11,7 +11,7 @@ from operator import itemgetter
 import numpy as np
 from tlz import concat, memoize, merge, pluck
 
-from dask import config, core, utils
+from dask import core
 from dask.array.chunk import getitem
 from dask.base import is_dask_collection, tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -173,13 +173,13 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
     index += (slice(None, None, None),) * missing
 
     # Pass down to next function
-    dsk_out, bd_out = slice_with_newaxes(out_name, in_name, blockdims, index, itemsize)
+    dsk_out, bd_out = slice_with_newaxes(out_name, in_name, blockdims, index)
 
     bd_out = tuple(map(tuple, bd_out))
     return dsk_out, bd_out
 
 
-def slice_with_newaxes(out_name, in_name, blockdims, index, itemsize):
+def slice_with_newaxes(out_name, in_name, blockdims, index):
     """
     Handle indexing with Nones
 
@@ -195,7 +195,7 @@ def slice_with_newaxes(out_name, in_name, blockdims, index, itemsize):
             where_none[i] -= n
 
     # Pass down and do work
-    dsk, blockdims2 = slice_wrap_lists(out_name, in_name, blockdims, index2, itemsize)
+    dsk, blockdims2 = slice_wrap_lists(out_name, in_name, blockdims, index2)
 
     if where_none:
         expand = expander(where_none)
@@ -220,7 +220,7 @@ def slice_with_newaxes(out_name, in_name, blockdims, index, itemsize):
         return dsk, blockdims2
 
 
-def slice_wrap_lists(out_name, in_name, blockdims, index, itemsize):
+def slice_wrap_lists(out_name, in_name, blockdims, index):
     """
     Fancy indexing along blocked array dasks
 
@@ -262,7 +262,7 @@ def slice_wrap_lists(out_name, in_name, blockdims, index, itemsize):
     if all(is_arraylike(i) or i == slice(None, None, None) for i in index):
         axis = where_list[0]
         blockdims2, dsk3 = take(
-            out_name, in_name, blockdims, index[where_list[0]], itemsize, axis=axis
+            out_name, in_name, blockdims, index[where_list[0]], axis=axis
         )
     # Mixed case. Both slices/integers and lists. slice/integer then take
     else:
@@ -279,7 +279,7 @@ def slice_wrap_lists(out_name, in_name, blockdims, index, itemsize):
         )
 
         # Do work
-        blockdims2, dsk2 = take(out_name, tmp, blockdims2, index[axis], 8, axis=axis2)
+        blockdims2, dsk2 = take(out_name, tmp, blockdims2, index[axis], axis=axis2)
         dsk3 = merge(dsk, dsk2)
 
     return dsk3, blockdims2
@@ -540,56 +540,7 @@ def issorted(seq):
     return np.all(seq[:-1] <= seq[1:])
 
 
-def slicing_plan(chunks, index):
-    """Construct a plan to slice chunks with the given index
-
-    Parameters
-    ----------
-    chunks : Tuple[int]
-        One dimensions worth of chunking information
-    index : np.ndarray[int]
-        The index passed to slice on that dimension
-
-    Returns
-    -------
-    out : List[Tuple[int, np.ndarray]]
-        A list of chunk/sub-index pairs corresponding to each output chunk
-    """
-    from dask.array.utils import asarray_safe
-
-    if not is_arraylike(index):
-        index = np.asanyarray(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")
-
-    # but we need chunk_locations as python ints for getitem calls downstream
-    chunk_locations = chunk_locations.tolist()
-    where = np.where(np.diff(chunk_locations))[0] + 1
-
-    extra = asarray_safe([0], like=where)
-    c_loc = asarray_safe([len(chunk_locations)], like=where)
-    where = np.concatenate([extra, where, c_loc])
-
-    out = []
-    for i in range(len(where) - 1):
-        sub_index = index[where[i] : where[i + 1]]
-        chunk = chunk_locations[where[i]]
-        if chunk > 0:
-            sub_index = sub_index - cum_chunks[chunk - 1]
-        out.append((chunk, sub_index))
-
-    return out
-
-
-def take(outname, inname, chunks, index, itemsize, axis=0):
+def take(outname, inname, chunks, index, axis=0):
     """Index array with an iterable of index
 
     Handles a single index by a single list
@@ -597,115 +548,44 @@ def take(outname, inname, chunks, index, itemsize, axis=0):
     Mimics ``np.take``
 
     >>> from pprint import pprint
-    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], 8, axis=0)
+    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
     >>> chunks
-    ((2, 1, 1),)
-    >>> pprint(dsk)   # doctest: +ELLIPSIS
-    {('y', 0): (<function getitem at ...>, ('x', 0), (array([5, 1]),)),
-     ('y', 1): (<function getitem at ...>, ('x', 2), (array([7]),)),
-     ('y', 2): (<function getitem at ...>, ('x', 0), (array([3]),))}
+    ((4,),)
 
-    When list is sorted we retain original block structure
+    When list is sorted we still try to preserve properly sized chunks.
 
-    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], 8, axis=0)
+    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
     >>> chunks
-    ((3, 1),)
-    >>> pprint(dsk)     # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
-    {('y', 0): (<function getitem at ...>,
-                ('x', 0),
-                (array([1, 3, 5]),)),
-     ('y', 1): (<function getitem at ...>, ('x', 2), (array([7]),))}
+    ((4,),)
 
     When any indexed blocks would otherwise grow larger than
-    dask.config.array.chunk-size, we might split them,
-    depending on the value of ``dask.config.slicing.split-large-chunks``.
-
-    >>> import dask
-    >>> with dask.config.set({"array.slicing.split-large-chunks": True}):
-    ...      chunks, dsk = take('y', 'x', [(1, 1, 1), (2000, 2000), (2000, 2000)],
-    ...                        [0] + [1] * 6 + [2], axis=0, itemsize=8)
-    >>> chunks
-    ((1, 3, 3, 1), (2000, 2000), (2000, 2000))
+    dask.config.array.chunk-size, we will split them to avoid
+    growing chunksizes.
     """
-    from dask.array.core import PerformanceWarning
 
-    plan = slicing_plan(chunks[axis], index)
-    if len(plan) >= len(chunks[axis]) * 10:
-        factor = math.ceil(len(plan) / len(chunks[axis]))
+    if not np.isnan(chunks[axis]).any():
+        from dask.array._shuffle import _shuffle
 
-        warnings.warn(
-            "Slicing with an out-of-order index is generating %d "
-            "times more chunks" % factor,
-            PerformanceWarning,
-            stacklevel=6,
-        )
+        average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))
 
-    # Check for chunks from the plan that would violate the user's
-    # configured chunk size.
-    nbytes = utils.parse_bytes(config.get("array.chunk-size"))
-    other_chunks = [chunks[i] for i in range(len(chunks)) if i != axis]
-    other_numel = math.prod(max(x) for x in other_chunks)
+        indexer = []
+        index = np.asarray(index)
+        for i in range(0, len(index), average_chunk_size):
+            indexer.append(index[i : i + average_chunk_size].tolist())
 
-    if math.isnan(other_numel) or other_numel == 0:
-        warnsize = maxsize = math.inf
+        token = (
+            outname.split("-")[-1]
+            if "-" in outname
+            else tokenize(outname, chunks, index, axis)
+        )
+        chunks, graph = _shuffle(chunks, indexer, axis, inname, outname, token)
+        return chunks, graph
     else:
-        maxsize = math.ceil(nbytes / (other_numel * itemsize))
-        warnsize = maxsize * 5
-
-    split = config.get("array.slicing.split-large-chunks", None)
-
-    # Warn only when the default is not specified.
-    warned = split is not None
-
-    for _, index_list in plan:
-        if not warned and len(index_list) > warnsize:
-            msg = (
-                "Slicing is producing a large chunk. To accept the large\n"
-                "chunk and silence this warning, set the option\n"
-                "    >>> with dask.config.set(**{'array.slicing.split_large_chunks': False}):\n"
-                "    ...     array[indexer]\n\n"
-                "To avoid creating the large chunks, set the option\n"
-                "    >>> with dask.config.set(**{'array.slicing.split_large_chunks': True}):\n"
-                "    ...     array[indexer]"
-            )
-            warnings.warn(msg, PerformanceWarning, stacklevel=6)
-            warned = True
-
-    where_index = []
-    index_lists = []
-    for where_idx, index_list in plan:
-        index_length = len(index_list)
-        if split and index_length > maxsize:
-            index_sublist = np.array_split(
-                index_list, math.ceil(index_length / maxsize)
-            )
-            index_lists.extend(index_sublist)
-            where_index.extend([where_idx] * len(index_sublist))
-        else:
-            if not is_arraylike(index_list):
-                index_list = np.array(index_list)
-            index_lists.append(index_list)
-            where_index.append(where_idx)
-
-    dims = [range(len(bd)) for bd in chunks]
-
-    indims = list(dims)
-    indims[axis] = list(range(len(where_index)))
-    keys = list(product([outname], *indims))
-
-    outdims = list(dims)
-    outdims[axis] = where_index
-    slices = [[colon] * len(bd) for bd in chunks]
-    slices[axis] = index_lists
-    slices = list(product(*slices))
-    inkeys = list(product([inname], *outdims))
-    values = [(getitem, inkey, slc) for inkey, slc in zip(inkeys, slices)]
-
-    chunks2 = list(chunks)
-    chunks2[axis] = tuple(map(len, index_lists))
-    dsk = dict(zip(keys, values))
-
-    return tuple(chunks2), dsk
+        from dask.array.core import unknown_chunk_message
+
+        raise ValueError(
+            f"Array chunk size or shape is unknown. {unknown_chunk_message}"
+        )
 
 
 def posify_index(shape, ind):


Test patch:
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index 8c72fc25c..ed8a99e0c 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -4,7 +4,8 @@ import itertools
 import warnings
 
 import pytest
-from tlz import merge
+
+from dask.array._shuffle import concatenate_arrays
 
 np = pytest.importorskip("numpy")
 
@@ -21,7 +22,6 @@ from dask.array.slicing import (
     sanitize_index,
     shuffle_slice,
     slice_array,
-    slicing_plan,
     take,
 )
 from dask.array.utils import assert_eq, same_keys
@@ -315,62 +315,111 @@ def test_slicing_with_newaxis():
 
 
 def test_take():
-    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], [5, 1, 47, 3], itemsize=8, axis=0)
+    chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
     expected = {
-        ("y", 0): (getitem, ("x", 0), (np.array([5, 1]),)),
-        ("y", 1): (getitem, ("x", 2), (np.array([7]),)),
-        ("y", 2): (getitem, ("x", 0), (np.array([3]),)),
+        ("y-y", 0): (
+            concatenate_arrays,
+            [
+                ("shuffle-split-y", 0),
+                ("shuffle-split-y", 1),
+            ],
+            np.array([1, 3, 0, 2]),
+            0,
+        ),
+        ("shuffle-split-y", 0): (
+            getitem,
+            ("x", 0),
+            (np.array([1, 3, 5]),),
+        ),
+        ("shuffle-split-y", 1): (
+            getitem,
+            ("x", 2),
+            (np.array([7]),),
+        ),
     }
     np.testing.assert_equal(sorted(dsk.items()), sorted(expected.items()))
-    assert chunks == ((2, 1, 1),)
+    assert chunks == ((4,),)
 
-    chunks, dsk = take(
-        "y", "x", [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], itemsize=8, axis=0
-    )
+    chunks, dsk = take("y-y", "x", [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], axis=0)
     expected = {
-        ("y", 0, 0): (
+        ("y-y", 0, 0): (
+            concatenate_arrays,
+            [
+                ("shuffle-split-y", 0),
+                ("shuffle-split-y", 1),
+            ],
+            np.array([1, 3, 0, 2]),
+            0,
+        ),
+        ("y-y", 0, 1): (
+            concatenate_arrays,
+            [
+                ("shuffle-split-y", 2),
+                ("shuffle-split-y", 3),
+            ],
+            np.array([1, 3, 0, 2]),
+            0,
+        ),
+        ("shuffle-split-y", 0): (
             getitem,
             ("x", 0, 0),
-            (np.array([5, 1]), slice(None, None, None)),
+            (np.array([1, 3, 5]), slice(None)),
         ),
-        ("y", 0, 1): (
+        ("shuffle-split-y", 1): (
+            getitem,
+            ("x", 2, 0),
+            (np.array([7]), slice(None)),
+        ),
+        ("shuffle-split-y", 2): (
             getitem,
             ("x", 0, 1),
-            (np.array([5, 1]), slice(None, None, None)),
+            (np.array([1, 3, 5]), slice(None)),
+        ),
+        ("shuffle-split-y", 3): (
+            getitem,
+            ("x", 2, 1),
+            (np.array([7]), slice(None)),
         ),
-        ("y", 1, 0): (getitem, ("x", 2, 0), (np.array([7]), slice(None, None, None))),
-        ("y", 1, 1): (getitem, ("x", 2, 1), (np.array([7]), slice(None, None, None))),
-        ("y", 2, 0): (getitem, ("x", 0, 0), (np.array([3]), slice(None, None, None))),
-        ("y", 2, 1): (getitem, ("x", 0, 1), (np.array([3]), slice(None, None, None))),
     }
     np.testing.assert_equal(sorted(dsk.items()), sorted(expected.items()))
-    assert chunks == ((2, 1, 1), (20, 20))
+    assert chunks == ((4,), (20, 20))
 
 
 def test_take_sorted():
-    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], itemsize=8, axis=0)
+    chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
     expected = {
-        ("y", 0): (getitem, ("x", 0), ([1, 3, 5],)),
-        ("y", 1): (getitem, ("x", 2), ([7],)),
+        ("y-y", 0): (
+            concatenate_arrays,
+            [
+                ("shuffle-split-y", 0),
+                ("shuffle-split-y", 1),
+            ],
+            np.array([0, 1, 2, 3]),
+            0,
+        ),
+        ("shuffle-split-y", 0): (
+            getitem,
+            ("x", 0),
+            (np.array([1, 3, 5]),),
+        ),
+        ("shuffle-split-y", 1): (
+            getitem,
+            ("x", 2),
+            (np.array([7]),),
+        ),
     }
     np.testing.assert_equal(dsk, expected)
-    assert chunks == ((3, 1),)
+    assert chunks == ((4,),)
 
-    chunks, dsk = take(
-        "y", "x", [(20, 20, 20, 20), (20, 20)], [1, 3, 5, 37], itemsize=8, axis=1
-    )
-    expected = merge(
-        {
-            ("y", i, 0): (getitem, ("x", i, 0), (slice(None, None, None), [1, 3, 5]))
-            for i in range(4)
-        },
-        {
-            ("y", i, 1): (getitem, ("x", i, 1), (slice(None, None, None), [17]))
-            for i in range(4)
-        },
-    )
+    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
+    expected = {
+        ("y", 0): (getitem, ("x", 0), (np.arange(0, 20),)),
+        ("y", 1): (getitem, ("x", 1), (np.arange(0, 20),)),
+        ("y", 2): (getitem, ("x", 2), (np.arange(0, 20),)),
+        ("y", 3): (getitem, ("x", 3), (np.arange(0, 20),)),
+    }
     np.testing.assert_equal(dsk, expected)
-    assert chunks == ((20, 20, 20, 20), (3, 1))
+    assert chunks == ((20, 20, 20, 20),)
 
 
 def test_slicing_chunks():
@@ -392,14 +441,14 @@ def test_slicing_chunks():
 
 def test_slicing_with_numpy_arrays():
     a, bd1 = slice_array(
-        "y",
+        "y-y",
         "x",
         ((3, 3, 3, 1), (3, 3, 3, 1)),
         (np.array([1, 2, 9]), slice(None, None, None)),
         itemsize=8,
     )
     b, bd2 = slice_array(
-        "y",
+        "y-y",
         "x",
         ((3, 3, 3, 1), (3, 3, 3, 1)),
         (np.array([1, 2, 9]), slice(None, None, None)),
@@ -412,7 +461,7 @@ def test_slicing_with_numpy_arrays():
     i = [False, True, True, False, False, False, False, False, False, True]
     index = (i, slice(None, None, None))
     index = normalize_index(index, (10, 10))
-    c, bd3 = slice_array("y", "x", ((3, 3, 3, 1), (3, 3, 3, 1)), index, itemsize=8)
+    c, bd3 = slice_array("y-y", "x", ((3, 3, 3, 1), (3, 3, 3, 1)), index, itemsize=8)
     assert bd1 == bd3
     np.testing.assert_equal(a, c)
 
@@ -860,27 +909,6 @@ def test_take_semi_sorted():
     assert y.chunks == ((5, 5, 5),)
 
 
-@pytest.mark.parametrize(
-    "chunks,index,expected",
-    [
-        ((5, 5, 5), np.arange(5, 15) % 10, [(1, np.arange(5)), (0, np.arange(5))]),
-        (
-            (5, 5, 5, 5),
-            np.arange(20) // 2,
-            [(0, np.arange(10) // 2), (1, np.arange(10) // 2)],
-        ),
-        ((10, 10), [15, 2, 3, 15], [(1, [5]), (0, [2, 3]), (1, [5])]),
-    ],
-)
-def test_slicing_plan(chunks, index, expected):
-    plan = slicing_plan(chunks, index=index)
-    assert len(plan) == len(expected)
-    for (i, x), (j, y) in zip(plan, expected):
-        assert i == j
-        assert len(x) == len(y)
-        assert (x == y).all()
-
-
 def test_getitem_avoids_large_chunks():
     with dask.config.set({"array.chunk-size": "0.1Mb"}):
         a = np.arange(2 * 128 * 128, dtype="int64").reshape(2, 128, 128)
@@ -895,12 +923,9 @@ def test_getitem_avoids_large_chunks():
         arr = da.from_array(a, chunks=(1, 128, 128))  # large chunks
         expected = a[indexer]
 
-        # By default, we warn
-        with pytest.warns(da.PerformanceWarning):
-            result = arr[indexer]
-
+        result = arr[indexer]
         assert_eq(result, expected)
-        assert result.chunks == ((1, 11), (128,), (128,))
+        assert result.chunks == ((1,) * 12, (128,), (128,))
 
         # Users can silence the warning
         with dask.config.set({"array.slicing.split-large-chunks": False}):
@@ -936,38 +961,36 @@ def test_take_avoids_large_chunks():
     # unit test for https://github.com/dask/dask/issues/6270
     with dask.config.set({"array.slicing.split-large-chunks": True}):
         chunks = ((1, 1, 1, 1), (500,), (500,))
-        itemsize = 8
         index = np.array([0, 1] + [2] * 101 + [3])
-        chunks2, dsk = take("a", "b", chunks, index, itemsize)
-        assert chunks2 == ((1, 1, 51, 50, 1), (500,), (500,))
-        assert len(dsk) == 5
+        chunks2, dsk = take("a", "b", chunks, index)
+        assert chunks2 == ((1,) * 104, (500,), (500,))
+        assert len(dsk) == 104
 
         index = np.array([0] * 101 + [1, 2, 3])
-        chunks2, dsk = take("a", "b", chunks, index, itemsize)
-        assert chunks2 == ((51, 50, 1, 1, 1), (500,), (500,))
-        assert len(dsk) == 5
+        chunks2, dsk = take("a", "b", chunks, index)
+        assert chunks2 == ((1,) * 104, (500,), (500,))
+        assert len(dsk) == 104
 
         index = np.array([0, 1, 2] + [3] * 101)
-        chunks2, dsk = take("a", "b", chunks, index, itemsize)
-        assert chunks2 == ((1, 1, 1, 51, 50), (500,), (500,))
-        assert len(dsk) == 5
+        chunks2, dsk = take("a", "b", chunks, index)
+        assert chunks2 == ((1,) * 104, (500,), (500,))
+        assert len(dsk) == 104
 
         chunks = ((500,), (1, 1, 1, 1), (500,))
         index = np.array([0, 1, 2] + [3] * 101)
-        chunks2, dsk = take("a", "b", chunks, index, itemsize, axis=1)
-        assert chunks2 == ((500,), (1, 1, 1, 51, 50), (500,))
-        assert len(dsk) == 5
+        chunks2, dsk = take("a", "b", chunks, index, axis=1)
+        assert chunks2 == ((500,), (1,) * 104, (500,))
+        assert len(dsk) == 104
 
 
 def test_take_uses_config():
     with dask.config.set({"array.slicing.split-large-chunks": True}):
         chunks = ((1, 1, 1, 1), (500,), (500,))
         index = np.array([0, 1] + [2] * 101 + [3])
-        itemsize = 8
         with config.set({"array.chunk-size": "10GB"}):
-            chunks2, dsk = take("a", "b", chunks, index, itemsize)
-        assert chunks2 == ((1, 1, 101, 1), (500,), (500,))
-        assert len(dsk) == 4
+            chunks2, dsk = take("a", "b", chunks, index)
+        assert chunks2 == ((1,) * 104, (500,), (500,))
+        assert len(dsk) == 104
 
 
 def test_pathological_unsorted_slicing():
@@ -976,11 +999,7 @@ def test_pathological_unsorted_slicing():
     # [0, 10, 20, ... 90, 1, 11, 21, ... 91, ...]
     index = np.arange(100).reshape(10, 10).ravel(order="F")
 
-    with pytest.warns(da.PerformanceWarning) as info:
-        x[index]
-
-    assert "10" in str(info.list[0])
-    assert "out-of-order" in str(info.list[0])
+    assert_eq(x[index], x.compute()[index])
 
 
 @pytest.mark.parametrize("params", [(2, 2, 1), (5, 3, 2)])