Repository: dask/dask

Issue:
Array.__getitem__ should consider `array.chunk-size` for concrete indexers?
Reported in https://github.com/pydata/xarray/issues/4112, where xarray calling `Array.__getitem__` resulted in an unbalanced Array.

```python
import dask.array

arr = dask.array.from_array([0, 1, 2, 3], chunks=(1,))
print(arr.chunks)  # ((1, 1, 1, 1),)
# align calls reindex which indexes with something like this
indexer = [0, 1, 2, 3, ] + [-1,] * 111
print(arr[indexer].chunks)  # ((1, 1, 1, 112),)
```

That last chunk is large, since all of the `-1`s are in the same input chunk.

To be explicit, the user could chunk the `indexer`:

```
# maybe something like this is a solution
lazy_indexer = dask.array.from_array(indexer, chunks=arr.chunks[0][0], name="idx")
print(arr[lazy_indexer].chunks) # ((1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1),)
```

If the indexer isn't `chunked`, perhaps dask could examine the input array and indexer and chunk the indexer if we detect that the output chunksize would exceed `dask.config.get('array.chunk-size')`?


Discussion:
> If the indexer isn't `chunked`, perhaps dask could examine the input array and indexer and chunk the indexer if we detect that the output chunksize would exceed `dask.config.get('array.chunk-size')`?

I think this strategy could work well for chunking `arr[indexer]` if any chunks would exceed some maximum chunk sizes.

However, I wouldn't implement this via chunking `indexer`. This removes information about which source chunks different parts of the indexer refer to, which makes indexing require a quadratic number of tasks.

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1539,7 +1539,6 @@ def __getitem__(self, index):
         )
 
         index2 = normalize_index(index, self.shape)
-
         dependencies = {self.name}
         for i in index2:
             if isinstance(i, Array):
@@ -1554,7 +1553,7 @@ def __getitem__(self, index):
             return self
 
         out = "getitem-" + tokenize(self, index2)
-        dsk, chunks = slice_array(out, self.name, self.chunks, index2)
+        dsk, chunks = slice_array(out, self.name, self.chunks, index2, self.itemsize)
 
         graph = HighLevelGraph.from_collections(out, dsk, dependencies=[self])
 
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -10,6 +10,8 @@
 from tlz import memoize, merge, pluck, concat, accumulate
 
 from .. import core
+from .. import config
+from .. import utils
 from ..highlevelgraph import HighLevelGraph
 from ..base import tokenize, is_dask_collection
 
@@ -83,7 +85,7 @@ def sanitize_index(ind):
         raise TypeError("Invalid index type", type(ind), ind)
 
 
-def slice_array(out_name, in_name, blockdims, index):
+def slice_array(out_name, in_name, blockdims, index, itemsize):
     """
     Master function for array slicing
 
@@ -104,6 +106,8 @@ def slice_array(out_name, in_name, blockdims, index):
       This is the dask variable output name
     blockshape - iterable of integers
     index - iterable of integers, slices, lists, or None
+    itemsize : int
+        The number of bytes required for each element of the array.
 
     Returns
     -------
@@ -155,13 +159,13 @@ def slice_array(out_name, in_name, blockdims, index):
     index += (slice(None, None, None),) * missing
 
     # Pass down to next function
-    dsk_out, bd_out = slice_with_newaxes(out_name, in_name, blockdims, index)
+    dsk_out, bd_out = slice_with_newaxes(out_name, in_name, blockdims, index, itemsize)
 
     bd_out = tuple(map(tuple, bd_out))
     return dsk_out, bd_out
 
 
-def slice_with_newaxes(out_name, in_name, blockdims, index):
+def slice_with_newaxes(out_name, in_name, blockdims, index, itemsize):
     """
     Handle indexing with Nones
 
@@ -177,7 +181,7 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
             where_none[i] -= n
 
     # Pass down and do work
-    dsk, blockdims2 = slice_wrap_lists(out_name, in_name, blockdims, index2)
+    dsk, blockdims2 = slice_wrap_lists(out_name, in_name, blockdims, index2, itemsize)
 
     if where_none:
         expand = expander(where_none)
@@ -202,7 +206,7 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
         return dsk, blockdims2
 
 
-def slice_wrap_lists(out_name, in_name, blockdims, index):
+def slice_wrap_lists(out_name, in_name, blockdims, index, itemsize):
     """
     Fancy indexing along blocked array dasks
 
@@ -244,7 +248,7 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
     if all(isinstance(i, np.ndarray) 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]], axis=axis
+            out_name, in_name, blockdims, index[where_list[0]], itemsize, axis=axis
         )
     # Mixed case. Both slices/integers and lists. slice/integer then take
     else:
@@ -261,7 +265,7 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
         )
 
         # Do work
-        blockdims2, dsk2 = take(out_name, tmp, blockdims2, index[axis], axis=axis2)
+        blockdims2, dsk2 = take(out_name, tmp, blockdims2, index[axis], 8, axis=axis2)
         dsk3 = merge(dsk, dsk2)
 
     return dsk3, blockdims2
@@ -554,14 +558,14 @@ def slicing_plan(chunks, index):
     return out
 
 
-def take(outname, inname, chunks, index, axis=0):
+def take(outname, inname, chunks, index, itemsize, axis=0):
     """Index array with an iterable of index
 
     Handles a single index by a single list
 
     Mimics ``np.take``
 
-    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
+    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], 8, axis=0)
     >>> chunks
     ((2, 1, 1),)
     >>> dsk  # doctest: +SKIP
@@ -572,12 +576,20 @@ def take(outname, inname, chunks, index, axis=0):
 
     When list is sorted we retain original block structure
 
-    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
+    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], 8, axis=0)
     >>> chunks
     ((3, 1),)
     >>> dsk  # doctest: +SKIP
     {('y', 0): (getitem, ('x', 0), ([1, 3, 5],)),
      ('y', 2): (getitem, ('x', 2), ([7],))}
+
+    When any indexed blocks would otherwise grow larger than
+    dask.config.array.chunk-size, we split them.
+
+    >>> chunks, dsk = take('y', 'x', [(1, 1, 1), (1000, 1000), (1000, 1000)],
+    ...                    [0] + [1] * 6 + [2], axis=0, itemsize=8)
+    >>> chunks
+    ((1, 3, 3, 1), (1000, 1000), (1000, 1000))
     """
     plan = slicing_plan(chunks[axis], index)
     if len(plan) >= len(chunks[axis]) * 10:
@@ -590,9 +602,29 @@ def take(outname, inname, chunks, index, axis=0):
             PerformanceWarning,
             stacklevel=6,
         )
-
-    index_lists = [idx for _, idx in plan]
-    where_index = [i for i, _ in plan]
+    index = np.asarray(index)
+
+    # 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 = np.prod([sum(x) for x in other_chunks])
+
+    maxsize = nbytes / (other_numel * itemsize)
+
+    where_index = []
+    index_lists = []
+    for where_idx, index_list in plan:
+        index_length = len(index_list)
+        if 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:
+            index_lists.append(np.array(index_list))
+            where_index.append(where_idx)
 
     dims = [range(len(bd)) for bd in chunks]
 


Test patch:
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -7,6 +7,7 @@
 np = pytest.importorskip("numpy")
 
 import dask
+from dask import config, utils
 import dask.array as da
 from dask.array.slicing import (
     _sanitize_index_element,
@@ -191,7 +192,7 @@ def test_slice_array_1d():
         ("y", 2): (getitem, ("x", 2), (slice(0, 25, 2),)),
         ("y", 3): (getitem, ("x", 3), (slice(1, 25, 2),)),
     }
-    result, chunks = slice_array("y", "x", [[25] * 4], [slice(24, None, 2)])
+    result, chunks = slice_array("y", "x", [[25] * 4], [slice(24, None, 2)], 8)
 
     assert expected == result
 
@@ -202,7 +203,7 @@ def test_slice_array_1d():
         ("y", 2): (getitem, ("x", 3), (slice(1, 25, 2),)),
     }
 
-    result, chunks = slice_array("y", "x", [[25] * 4], [slice(26, None, 2)])
+    result, chunks = slice_array("y", "x", [[25] * 4], [slice(26, None, 2)], 8)
     assert expected == result
 
     # x[24::2]
@@ -212,7 +213,7 @@ def test_slice_array_1d():
         ("y", 2): (getitem, ("x", 2), (slice(0, 25, 2),)),
         ("y", 3): (getitem, ("x", 3), (slice(1, 25, 2),)),
     }
-    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(24, None, 2),))
+    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(24, None, 2),), 8)
 
     assert expected == result
 
@@ -223,7 +224,7 @@ def test_slice_array_1d():
         ("y", 2): (getitem, ("x", 3), (slice(1, 25, 2),)),
     }
 
-    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(26, None, 2),))
+    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(26, None, 2),), 8)
     assert expected == result
 
 
@@ -244,7 +245,11 @@ def test_slice_array_2d():
     }
 
     result, chunks = slice_array(
-        "y", "x", [[20], [20, 20, 5]], [slice(13, None, 2), slice(10, None, 1)]
+        "y",
+        "x",
+        [[20], [20, 20, 5]],
+        [slice(13, None, 2), slice(10, None, 1)],
+        itemsize=8,
     )
 
     assert expected == result
@@ -256,7 +261,9 @@ def test_slice_array_2d():
         ("y", 2): (getitem, ("x", 0, 2), (5, slice(None, None, None))),
     }
 
-    result, chunks = slice_array("y", "x", ([20], [20, 20, 5]), [5, slice(10, None, 1)])
+    result, chunks = slice_array(
+        "y", "x", ([20], [20, 20, 5]), [5, slice(10, None, 1)], 8
+    )
 
     assert expected == result
 
@@ -264,7 +271,7 @@ def test_slice_array_2d():
 def test_slice_optimizations():
     # bar[:]
     expected = {("foo", 0): ("bar", 0)}
-    result, chunks = slice_array("foo", "bar", [[100]], (slice(None, None, None),))
+    result, chunks = slice_array("foo", "bar", [[100]], (slice(None, None, None),), 8)
     assert expected == result
 
     # bar[:,:,:]
@@ -274,12 +281,15 @@ def test_slice_optimizations():
         "bar",
         [(100, 1000, 10000)],
         (slice(None, None, None), slice(None, None, None), slice(None, None, None)),
+        itemsize=8,
     )
     assert expected == result
 
 
 def test_slicing_with_singleton_indices():
-    result, chunks = slice_array("y", "x", ([5, 5], [5, 5]), (slice(0, 5), 8))
+    result, chunks = slice_array(
+        "y", "x", ([5, 5], [5, 5]), (slice(0, 5), 8), itemsize=8
+    )
 
     expected = {("y", 0): (getitem, ("x", 0, 1), (slice(None, None, None), 3))}
 
@@ -288,7 +298,11 @@ def test_slicing_with_singleton_indices():
 
 def test_slicing_with_newaxis():
     result, chunks = slice_array(
-        "y", "x", ([5, 5], [5, 5]), (slice(0, 3), None, slice(None, None, None))
+        "y",
+        "x",
+        ([5, 5], [5, 5]),
+        (slice(0, 3), None, slice(None, None, None)),
+        itemsize=8,
     )
 
     expected = {
@@ -309,7 +323,7 @@ def test_slicing_with_newaxis():
 
 
 def test_take():
-    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
+    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], [5, 1, 47, 3], itemsize=8, axis=0)
     expected = {
         ("y", 0): (getitem, ("x", 0), (np.array([5, 1]),)),
         ("y", 1): (getitem, ("x", 2), (np.array([7]),)),
@@ -318,7 +332,9 @@ def test_take():
     np.testing.assert_equal(sorted(dsk.items()), sorted(expected.items()))
     assert chunks == ((2, 1, 1),)
 
-    chunks, dsk = take("y", "x", [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], axis=0)
+    chunks, dsk = take(
+        "y", "x", [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], itemsize=8, axis=0
+    )
     expected = {
         ("y", 0, 0): (
             getitem,
@@ -340,7 +356,7 @@ def test_take():
 
 
 def test_take_sorted():
-    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
+    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], itemsize=8, axis=0)
     expected = {
         ("y", 0): (getitem, ("x", 0), ([1, 3, 5],)),
         ("y", 1): (getitem, ("x", 2), ([7],)),
@@ -348,7 +364,9 @@ def test_take_sorted():
     np.testing.assert_equal(dsk, expected)
     assert chunks == ((3, 1),)
 
-    chunks, dsk = take("y", "x", [(20, 20, 20, 20), (20, 20)], [1, 3, 5, 37], axis=1)
+    chunks, dsk = take(
+        "y", "x", [(20, 20, 20, 20), (20, 20)], [1, 3, 5, 37], itemsize=8, axis=1
+    )
     expected = merge(
         dict(
             (("y", i, 0), (getitem, ("x", i, 0), (slice(None, None, None), [1, 3, 5])))
@@ -364,15 +382,19 @@ def test_take_sorted():
 
 
 def test_slicing_chunks():
-    result, chunks = slice_array("y", "x", ([5, 5], [5, 5]), (1, np.array([2, 0, 3])))
+    result, chunks = slice_array(
+        "y", "x", ([5, 5], [5, 5]), (1, np.array([2, 0, 3])), itemsize=8
+    )
     assert chunks == ((3,),)
 
     result, chunks = slice_array(
-        "y", "x", ([5, 5], [5, 5]), (slice(0, 7), np.array([2, 0, 3]))
+        "y", "x", ([5, 5], [5, 5]), (slice(0, 7), np.array([2, 0, 3])), itemsize=8
     )
     assert chunks == ((5, 2), (3,))
 
-    result, chunks = slice_array("y", "x", ([5, 5], [5, 5]), (slice(0, 7), 1))
+    result, chunks = slice_array(
+        "y", "x", ([5, 5], [5, 5]), (slice(0, 7), 1), itemsize=8
+    )
     assert chunks == ((5, 2),)
 
 
@@ -382,12 +404,14 @@ def test_slicing_with_numpy_arrays():
         "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",
         "x",
         ((3, 3, 3, 1), (3, 3, 3, 1)),
         (np.array([1, 2, 9]), slice(None, None, None)),
+        itemsize=8,
     )
 
     assert bd1 == bd2
@@ -396,7 +420,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)
+    c, bd3 = slice_array("y", "x", ((3, 3, 3, 1), (3, 3, 3, 1)), index, itemsize=8)
     assert bd1 == bd3
     np.testing.assert_equal(a, c)
 
@@ -575,7 +599,7 @@ def test_sanitize_index():
 def test_uneven_blockdims():
     blockdims = ((31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30), (100,))
     index = (slice(240, 270), slice(None))
-    dsk_out, bd_out = slice_array("in", "out", blockdims, index)
+    dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
     sol = {
         ("in", 0, 0): (getitem, ("out", 7, 0), (slice(28, 31, 1), slice(None))),
         ("in", 1, 0): (getitem, ("out", 8, 0), (slice(0, 27, 1), slice(None))),
@@ -585,7 +609,7 @@ def test_uneven_blockdims():
 
     blockdims = ((31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30),) * 2
     index = (slice(240, 270), slice(180, 230))
-    dsk_out, bd_out = slice_array("in", "out", blockdims, index)
+    dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
     sol = {
         ("in", 0, 0): (getitem, ("out", 7, 5), (slice(28, 31, 1), slice(29, 30, 1))),
         ("in", 0, 1): (getitem, ("out", 7, 6), (slice(28, 31, 1), slice(None))),
@@ -849,7 +873,7 @@ def test_take_semi_sorted():
     ],
 )
 def test_slicing_plan(chunks, index, expected):
-    plan = slicing_plan(chunks, index)
+    plan = slicing_plan(chunks, index=index)
     assert len(plan) == len(expected)
     for (i, x), (j, y) in zip(plan, expected):
         assert i == j
@@ -857,6 +881,76 @@ def test_slicing_plan(chunks, index, expected):
         assert (x == y).all()
 
 
+def test_getitem_avoids_large_chunks():
+    a = np.arange(4 * 500 * 500).reshape(4, 500, 500)
+    arr = da.from_array(a, chunks=(1, 500, 500))
+    indexer = [0, 1] + [2] * 100 + [3]
+    result = arr[indexer]
+    chunk_size = utils.parse_bytes(config.get("array.chunk-size"))
+    assert all(x.nbytes < chunk_size for x in result.blocks)
+    expected = a[indexer]
+
+    assert_eq(result, expected)
+
+
+@pytest.mark.parametrize(
+    "chunks",
+    [
+        ((1, 1, 1, 1), (np.nan,), (np.nan,)),
+        pytest.param(
+            ((np.nan, np.nan, np.nan, np.nan), (500,), (500,)),
+            marks=pytest.mark.xfail(reason="https://github.com/dask/dask/issues/6586"),
+        ),
+    ],
+)
+def test_getitem_avoids_large_chunks_missing(chunks):
+    # We cannot apply the "avoid large chunks" optimization when
+    # the chunks have unknown sizes.
+    a = np.arange(4 * 500 * 500).reshape(4, 500, 500)
+    arr = da.from_array(a, chunks=(1, 500, 500))
+    arr._chunks = chunks
+    indexer = [0, 1] + [2] * 100 + [3]
+    expected = a[indexer]
+    result = arr[indexer]
+    assert_eq(result, expected)
+
+
+def test_take_avoids_large_chunks():
+    # unit test for https://github.com/dask/dask/issues/6270
+    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
+
+    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
+
+    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
+
+    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
+
+
+def test_take_uses_config():
+    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
+
+
 def test_pathological_unsorted_slicing():
     x = da.ones(100, chunks=10)