Repository: dask/dask

Issue:
dask scheduler struggles with concatenate, following by indexing, followed by mean
In xray, this looks something like the following:

``` python
# concatenate a bunch of arrays along the "time" dimension
ds = xray.open_mfdataset('/Users/shoyer/data/era-interim/2t/*.nc', engine='scipy')
# this is 6-hourly data from 35 years; select out only data for January and take the mean 
ds.sel(time=ds['time.month'] == 1).mean()
```

The scheduler eats up all my memory when I try to compute the data. In contrast, `ds.mean()` (no intermediate indexing) works fine.



Discussion:
Ah, ok, so when we index with a list or array of booleans

```
y = x[[5, 1, 2, 4]]
```

We don't have any knowledge about the size of blocks that we should produce.  Should we have one block of four? two blocks of two?  It's unclear.  Currently we follow the one-giant-block approach, hoping that either the list of items is not very large or that the blocking along the other dimensions manages the memory blowup.

This may be a fail case.  We have a very large list of elements to grab, so a very long block along one dimension.  Presumably this data is large enough so that cutting in the other dimensions doesn't save us.

To resolve this problem we probably need a heuristic to determine good output block sizes when given large indexing lists / array-of-bools.

I wonder if we can do something like the current slicing logic: simply preserve the existing blocks, if possible (dropping unused blocks). That would work as long as indexing doesn't do anything pathological like accessing values out of order.

I view indexing with lists as likely to span different blocks.  Indexing with arrays of booleans though might be a different story.  Perhaps in the list case we do one-giant-block while with arrays of booleans we do all of your original blocks.

Hmm. With xray we currently canonicalize all indexing to use integers instead of booleans. I'm not sure it's a good idea to make behavior that relies on that distinction...

Here's another choice that avoids heuristics: use the size of the largest block as the blockshape along the indexed dimension.

That could work.  We'll still run into issues when the list is out-of-order but hopefully large out-of-order lists are rare.

@mrocklin I don't follow you on the issue for out of order large lists.

To be more explicit, I was thinking: if `array.blockdims[axis] == (5, 8, 10)`, then the result of any array based indexing `result = array[key]` satisfies `max(result.blockdims[axis]) <= 10`.

Yeah, nm, it's just a performance concern.  We might still have problems when someone indexes with 

```
x[[1000000, 1, 50000, 3, ...]]
```

All of the blocks would be smallish, as you say, but the scheduler is likely to hold on to things longer than we'd like

Hrm, The way we handle out-of-order lists I'm not sure there is much to be gained here.  We're going to end up pulling a lot of the data in to memory regardless.

Instead, I suggest a short-cut route in the case where the list is sorted (this is fast to check) in which case we stay within the original block structure.

Sounds good to me. I think the sorted case will be quite common.


Resolving patch:
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -178,9 +178,8 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
     if all(isinstance(i, list) or i == slice(None, None, None)
             for i in index2):
         axis = where_list[0]
-        dsk3 = take(out_name, in_name, blockdims, index2[where_list[0]],
-                    axis=axis)
-        blockdims2 = blockdims
+        blockdims2, dsk3 = take(out_name, in_name, blockdims,
+                                index2[where_list[0]], axis=axis)
     # Mixed case. Both slices/integers and lists. slice/integer then take
     else:
         # Do first pass without lists
@@ -193,15 +192,11 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
                            if i < axis and isinstance(ind, (int, long)))
 
         # Do work
-        dsk2 = take(out_name, tmp, blockdims2, index2[axis], axis=axis2)
+        blockdims2, dsk2 = take(out_name, tmp, blockdims2, index2[axis],
+                                axis=axis2)
         dsk3 = merge(dsk, dsk2)
 
-    # Replace blockdims of list entries with single block
-    index4 = [ind for ind in index2 if not isinstance(ind, (int, long))]
-    blockdims3 = tuple([bd if not isinstance(i, list) else (len(i),)
-                        for i, bd in zip(index4, blockdims2)])
-
-    return dsk3, blockdims3
+    return dsk3, blockdims2
 
 
 def slice_slices_and_integers(out_name, in_name, blockdims, index):
@@ -383,6 +378,73 @@ def partition_by_size(sizes, seq):
     return result
 
 
+def issorted(seq):
+    """ Is sequence sorted?
+
+    >>> issorted([1, 2, 3])
+    True
+    >>> issorted([3, 1, 2])
+    False
+    """
+    if not seq:
+        return True
+    x = seq[0]
+    for elem in seq[1:]:
+        if elem < x:
+            return False
+        x = elem
+    return True
+
+
+colon = slice(None, None, None)
+
+
+def take_sorted(outname, inname, blockdims, index, axis=0):
+    """ Index array with sorted list index
+
+    Forms a dask for the following case
+
+        x[:, [1, 3, 5, 10], ...]
+
+    where the index, ``[1, 3, 5, 10]`` is sorted in non-decreasing order.
+
+    >>> blockdims, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
+    >>> blockdims
+    ((3, 1),)
+    >>> dsk  # doctest: +SKIP
+    {('y', 0): (getitem, ('x', 0), ([1, 3, 5],)),
+     ('y', 1): (getitem, ('x', 2), ([7],))}
+
+    See also:
+        take - calls this function
+    """
+    n = len(blockdims)
+    sizes = blockdims[axis]  # the blocksizes on the axis that we care about
+
+    index_lists = partition_by_size(sizes, sorted(index))
+    where_index = [i for i, il in enumerate(index_lists) if il]
+    index_lists = [il for il in index_lists if il]
+
+    dims = [range(len(bd)) for bd in blockdims]
+
+    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 blockdims]
+    slices[axis] = index_lists
+    slices = list(product(*slices))
+    inkeys = list(product([inname], *outdims))
+    values = [(getitem, inkey, slc) for inkey, slc in zip(inkeys, slices)]
+
+    blockdims2 = list(blockdims)
+    blockdims2[axis] = tuple(map(len, index_lists))
+
+    return tuple(blockdims2), dict(zip(keys, values))
+
+
 def take(outname, inname, blockdims, index, axis=0):
     """ Index array with an iterable of index
 
@@ -390,12 +452,27 @@ def take(outname, inname, blockdims, index, axis=0):
 
     Mimics ``np.take``
 
-    >>> take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)  # doctest: +SKIP
+    >>> blockdims, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
+    >>> blockdims
+    ((4,),)
+    >>> dsk  # doctest: +SKIP
     {('y', 0): (getitem, (np.concatenate, [(getitem, ('x', 0), ([1, 3, 5],)),
                                            (getitem, ('x', 2), ([7],))],
                                           0),
                          (2, 0, 4, 1))}
+
+    When list is sorted we retain original block structure
+
+    >>> blockdims, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
+    >>> blockdims
+    ((3, 1),)
+    >>> dsk  # doctest: +SKIP
+    {('y', 0): (getitem, ('x', 0), ([1, 3, 5],)),
+     ('y', 2): (getitem, ('x', 2), ([7],))}
     """
+    if issorted(index):
+        return take_sorted(outname, inname, blockdims, index, axis)
+
     n = len(blockdims)
     sizes = blockdims[axis]  # the blocksizes on the axis that we care about
 
@@ -405,8 +482,6 @@ def take(outname, inname, blockdims, index, axis=0):
                 for i, bd in enumerate(blockdims)]
     keys = list(product([outname], *dims))
 
-    colon = slice(None, None, None)
-
     rev_index = list(map(sorted(index).index, index))
     vals = [(getitem, (np.concatenate,
                   (list, [(getitem, ((inname,) + d[:axis] + (i,) + d[axis+1:]),
@@ -417,7 +492,10 @@ def take(outname, inname, blockdims, index, axis=0):
                      ((colon,)*axis + (rev_index,) + (colon,)*(n-axis-1)))
             for d in product(*dims)]
 
-    return dict(zip(keys, vals))
+    blockdims2 = list(blockdims)
+    blockdims2[axis] = (len(index),)
+
+    return tuple(blockdims2), dict(zip(keys, vals))
 
 
 def posify_index(shape, ind):


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
@@ -4,6 +4,7 @@
 from dask.array.slicing import slice_array, _slice_1d, take
 from operator import getitem
 import numpy as np
+from toolz import merge
 
 
 def eq(a, b):
@@ -251,7 +252,7 @@ def test_slicing_with_newaxis():
 
 
 def test_take():
-    result = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
+    blockdims, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
     expected = {('y', 0):
             (getitem,
               (np.concatenate, (list,
@@ -259,9 +260,10 @@ def test_take():
                  (getitem, ('x', 2), ([7],))]),
                0),
              ([2, 0, 3, 1],))}
-    assert result == expected
+    assert dsk == expected
+    assert blockdims == ((4,),)
 
-    result = take('y', 'x', [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], axis=0)
+    blockdims, dsk = take('y', 'x', [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], axis=0)
     expected = dict((('y', 0, j),
             (getitem,
               (np.concatenate, (list,
@@ -270,9 +272,11 @@ def test_take():
                 0),
               ([2, 0, 3, 1], slice(None, None, None))))
             for j in range(2))
-    assert result == expected
+    assert dsk == expected
+    assert blockdims == ((4,), (20, 20))
 
-    result = take('y', 'x', [(20, 20, 20, 20), (20, 20)], [5, 1, 37, 3], axis=1)
+
+    blockdims, dsk = take('y', 'x', [(20, 20, 20, 20), (20, 20)], [5, 1, 37, 3], axis=1)
     expected = dict((('y', i, 0),
             (getitem,
               (np.concatenate, (list,
@@ -281,13 +285,32 @@ def test_take():
                 1),
              (slice(None, None, None), [2, 0, 3, 1])))
            for i in range(4))
+    assert dsk == expected
+    assert blockdims == ((20, 20, 20, 20), (4,))
+
+
+def test_take_sorted():
+    blockdims, dsk = take('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],))}
+    assert dsk == expected
+    assert blockdims == ((3, 1),)
 
-    assert result == expected
+    blockdims, dsk = take('y', 'x', [(20, 20, 20, 20), (20, 20)], [1, 3, 5, 37], axis=1)
+    expected = merge(
+            dict((('y', i, 0),
+                  (getitem, ('x', i, 0), (slice(None, None, None), [1, 3, 5])))
+                  for i in range(4)),
+            dict((('y', i, 1),
+                  (getitem, ('x', i, 1), (slice(None, None, None), [17])))
+                  for i in range(4)))
+    assert dsk == expected
+    assert blockdims == ((20, 20, 20, 20), (3, 1))
 
 
 def test_slice_lists():
     y, blockdims = slice_array('y', 'x', ((3, 3, 3, 1), (3, 3, 3, 1)),
-                                ([1, 2, 9], slice(None, None, None)))
+                                ([2, 1, 9], slice(None, None, None)))
     assert y == \
         dict((('y', 0, i), (getitem,
                        (np.concatenate,
@@ -300,7 +323,7 @@ def test_slice_lists():
                            ([0], slice(None, None, None))),
                            ]),
                         0),
-                       ([0, 1, 2], slice(None, None, None))))
+                       ([1, 0, 2], slice(None, None, None))))
                 for i in range(4))
 
     assert blockdims == ((3,), (3, 3, 3, 1))