Repository: dask/dask

Issue:
da.compress() error with zarr array
I'm running into a problem using `da.compress()` with a 2D zarr array. Minimal example is below. I think what happens is that dask implements compress by trying to index (`__getitem__`) the input array with a list of integers. Zarr arrays don't support this, only contiguous slices, and so an error gets raised by zarr. Would it be possible to implement `da.compress()` without requiring the underling storage to support this? 

I sort of naively expected that dask would just load each full chunk required from the input array, via a contiguous slice, into a numpy array, then call numpy.compress on this chunk with the appropriate subset of the condition array, if that makes any sense. And so the input array only needs to support contiguous slices in `__getitem__`.

```
>>> import numpy as np
>>> import zarr
>>> zarr.__version__
'2.0.1'
>>> import dask
>>> dask.__version__
'0.11.0'
>>> x = zarr.array(np.random.randint(0, 4, size=(10000, 100)))
>>> f = np.random.randint(0, 2, size=10000, dtype=bool)
>>> import dask.array as da
>>> d = da.from_array(x, chunks=x.chunks)
>>> d
dask.array<array-6..., shape=(10000, 100), dtype=int64, chunksize=(1250, 13)>
>>> y = da.compress(f, d, axis=0)
>>> y
dask.array<getitem..., shape=(4991, 100), dtype=int64, chunksize=(598, 13)>
>>> y.compute()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/dask/base.py", line 86, in compute
    return compute(self, **kwargs)[0]
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/dask/base.py", line 179, in compute
    results = get(dsk, keys, **kwargs)
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/dask/threaded.py", line 58, in get
    **kwargs)
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/dask/async.py", line 481, in get_async
    raise(remote_exception(res, tb))
dask.async.TypeError: expected integer or slice, found: [3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 23, 25, 27, 30, 32, 33, 36, 37, 38, 39, 40, 41, 45, 46, 47, 48, 50, 52, 55, 56, 57, 61, 63, 65, 68, 70, 71, 75, 76, 78, 79, 80, 81, 84, 86, 89, 90, 92, 93, 94, 96, 97, 98, 99, 102, 105, 106, 108, 109, 110, 113, 115, 117, 118, 119, 120, 121, 123, 125, 126, 128, 129, 131, 132, 134, 135, 136, 137, 138, 139, 141, 142, 143, 144, 149, 150, 151, 152, 154, 155, 157, 158, 164, 165, 166, 167, 168, 169, 171, 172, 174, 176, 178, 182, 183, 184, 187, 188, 190, 194, 195, 198, 202, 203, 204, 208, 209, 210, 211, 212, 213, 216, 217, 218, 219, 220, 222, 223, 224, 228, 229, 232, 234, 235, 238, 240, 242, 247, 252, 253, 254, 258, 263, 264, 267, 268, 269, 271, 272, 273, 282, 286, 287, 289, 290, 291, 292, 296, 297, 299, 300, 304, 305, 311, 312, 313, 315, 319, 321, 325, 328, 329, 331, 334, 338, 339, 341, 343, 344, 345, 348, 352, 354, 355, 364, 365, 370, 372, 373, 377, 378, 381, 383, 384, 389, 396, 398, 403, 404, 405, 409, 411, 412, 416, 420, 421, 422, 424, 427, 428, 430, 433, 434, 437, 439, 441, 448, 450, 451, 452, 458, 459, 461, 462, 463, 465, 472, 474, 475, 477, 479, 480, 482, 483, 484, 485, 486, 487, 488, 489, 490, 493, 496, 498, 501, 502, 505, 506, 507, 510, 512, 514, 515, 520, 521, 525, 528, 535, 536, 537, 538, 540, 544, 548, 549, 550, 559, 560, 562, 564, 565, 566, 568, 570, 571, 577, 578, 579, 580, 581, 584, 585, 591, 592, 593, 594, 595, 599, 600, 602, 604, 607, 608, 609, 613, 615, 616, 618, 624, 626, 629, 631, 632, 634, 636, 637, 642, 644, 645, 646, 648, 650, 651, 652, 653, 657, 660, 661, 662, 663, 666, 668, 670, 672, 675, 679, 681, 682, 687, 688, 692, 694, 696, 697, 699, 701, 702, 704, 706, 708, 710, 711, 713, 715, 717, 718, 720, 724, 726, 727, 730, 731, 733, 735, 737, 738, 739, 740, 741, 742, 747, 748, 750, 753, 758, 761, 762, 768, 769, 771, 776, 778, 779, 780, 781, 783, 785, 788, 791, 792, 793, 794, 795, 797, 803, 807, 810, 814, 818, 821, 826, 827, 830, 832, 839, 841, 843, 844, 845, 846, 848, 850, 851, 852, 855, 857, 858, 861, 862, 864, 866, 868, 869, 873, 877, 878, 880, 882, 884, 887, 889, 892, 893, 894, 895, 899, 900, 902, 903, 904, 905, 909, 911, 913, 914, 917, 920, 921, 926, 932, 935, 936, 938, 940, 942, 946, 947, 948, 949, 952, 953, 956, 957, 959, 962, 963, 964, 966, 969, 971, 972, 973, 974, 976, 979, 980, 990, 991, 996, 998, 999, 1001, 1002, 1003, 1004, 1006, 1007, 1009, 1010, 1011, 1013, 1014, 1015, 1016, 1017, 1018, 1020, 1023, 1024, 1025, 1030, 1031, 1033, 1034, 1035, 1036, 1037, 1043, 1048, 1050, 1054, 1056, 1059, 1061, 1062, 1069, 1070, 1075, 1076, 1077, 1080, 1082, 1084, 1087, 1089, 1092, 1093, 1094, 1099, 1103, 1107, 1108, 1111, 1112, 1114, 1115, 1116, 1117, 1118, 1120, 1125, 1130, 1131, 1132, 1133, 1135, 1139, 1141, 1142, 1143, 1145, 1146, 1147, 1148, 1149, 1152, 1153, 1154, 1156, 1157, 1160, 1161, 1162, 1163, 1165, 1167, 1168, 1170, 1173, 1176, 1177, 1178, 1181, 1186, 1188, 1189, 1192, 1194, 1195, 1196, 1197, 1198, 1200, 1203, 1208, 1209, 1214, 1215, 1217, 1219, 1221, 1222, 1223, 1224, 1227, 1229, 1231, 1232, 1234, 1235, 1236, 1240, 1241, 1242, 1246, 1248]

Traceback
---------
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/dask/async.py", line 263, in execute_task
    result = _execute_task(task, data)
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/dask/async.py", line 245, in _execute_task
    return func(*args2)
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/dask/array/core.py", line 49, in getarray
    c = a[b]
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/zarr/core.py", line 370, in __getitem__
    selection = normalize_array_selection(item, self._shape)
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/zarr/util.py", line 183, in normalize_array_selection
    for i, l in zip(item, shape))
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/zarr/util.py", line 183, in <genexpr>
    for i, l in zip(item, shape))
  File "/home/aliman/miniconda3/envs/biipy/lib/python3.5/site-packages/zarr/util.py", line 162, in normalize_axis_selection
    raise TypeError('expected integer or slice, found: %r' % item)
```



Discussion:
I agree. By default, dask.array shouldn't presume anything beyond contiguous slicing. But if supported by the backend (e.g., a memory-mapped numpy array), list based indexing can be much more efficient.

It seems that dask.array needs some way to annotate the capabilities / indexing mode of input arrays, both internally (e.g., to distinguish between numpy and non-numpy arrays) and as part of the external API (e.g., for tools like xarray, that has it's own non-public array classes that we pass to `da.from_array`).

This has to do with our optimization passes, which fuse repeated slices together. If you try computing without optimization then everything works:

``` python
In [99]: y.compute(optimize_graph=False)
Out[99]:
array([[3, 1, 0, ..., 3, 2, 3],
       [3, 1, 1, ..., 3, 3, 1],
       [2, 1, 2, ..., 2, 1, 1],
       ...,
       [3, 0, 1, ..., 0, 3, 2],
       [3, 2, 0, ..., 1, 0, 3],
       [2, 2, 0, ..., 2, 1, 2]])
```

Perhaps we shouldn't be fusing non-slices if the base store isn't a numpy array. This is a fairly easy fix, if it doesn't cause performance degradations for other backends. Thoughts?


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -55,6 +55,15 @@ def getarray(a, b, lock=None):
     return c
 
 
+def getarray_nofancy(a, b, lock=None):
+    """ A simple wrapper around ``getarray``.
+
+    Used to indicate to the optimization passes that the backend doesn't
+    support "fancy indexing"
+    """
+    return getarray(a, b, lock=lock)
+
+
 from .optimization import optimize
 
 
@@ -76,7 +85,7 @@ def slices_from_chunks(chunks):
                 for start, shape in zip(starts, shapes)]
 
 
-def getem(arr, chunks, shape=None, out_name=None):
+def getem(arr, chunks, shape=None, out_name=None, fancy=True, lock=False):
     """ Dask getting various chunks from an array-like
 
     >>> getem('X', chunks=(2, 3), shape=(4, 6))  # doctest: +SKIP
@@ -95,8 +104,13 @@ def getem(arr, chunks, shape=None, out_name=None):
     chunks = normalize_chunks(chunks, shape)
 
     keys = list(product([out_name], *[range(len(bds)) for bds in chunks]))
+    slices = slices_from_chunks(chunks)
+    getter = getarray if fancy else getarray_nofancy
 
-    values = [(getarray, arr, x) for x in slices_from_chunks(chunks)]
+    if lock:
+        values = [(getter, arr, x, lock) for x in slices]
+    else:
+        values = [(getter, arr, x) for x in slices]
 
     return dict(zip(keys, values))
 
@@ -1504,17 +1518,28 @@ def normalize_chunks(chunks, shape=None):
     return tuple(map(tuple, chunks))
 
 
-def from_array(x, chunks, name=None, lock=False):
+def from_array(x, chunks, name=None, lock=False, fancy=True):
     """ Create dask array from something that looks like an array
 
     Input must have a ``.shape`` and support numpy-style slicing.
 
-    The ``chunks`` argument must be one of the following forms:
-
-    -   a blocksize like 1000
-    -   a blockshape like (1000, 1000)
-    -   explicit sizes of all blocks along all dimensions
-        like ((1000, 1000, 500), (400, 400)).
+    Parameters
+    ----------
+    x : array_like
+    chunks : int, tuple
+        How to chunk the array. Must be one of the following forms:
+        - A blocksize like 1000.
+        - A blockshape like (1000, 1000).
+        - Explicit sizes of all blocks along all dimensions
+          like ((1000, 1000, 500), (400, 400)).
+    name : str, optional
+        The key name to use for the array. Defaults to a hash of ``x``.
+    lock : bool or Lock, optional
+        If ``x`` doesn't support concurrent reads then provide a lock here, or
+        pass in True to have dask.array create one for you.
+    fancy : bool, optional
+        If ``x`` doesn't support fancy indexing (e.g. indexing with lists or
+        arrays) then set to False. Default is True.
 
     Examples
     --------
@@ -1538,11 +1563,9 @@ def from_array(x, chunks, name=None, lock=False):
     token = tokenize(x, chunks)
     original_name = (name or 'array-') + 'original-' + token
     name = name or 'array-' + token
-    dsk = getem(original_name, chunks, out_name=name)
     if lock is True:
         lock = Lock()
-    if lock:
-        dsk = dict((k, v + (lock,)) for k, v in dsk.items())
+    dsk = getem(original_name, chunks, out_name=name, fancy=fancy, lock=lock)
     return Array(merge({original_name: x}, dsk), name, chunks, dtype=x.dtype)
 
 
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from .core import getarray
+from .core import getarray, getarray_nofancy
 from ..core import flatten
 from ..optimize import cull, fuse, inline_functions
 
@@ -18,7 +18,7 @@ def optimize(dsk, keys, **kwargs):
     """
     keys = list(flatten(keys))
     fast_functions = kwargs.get('fast_functions',
-                                set([getarray, np.transpose]))
+                                set([getarray, getarray_nofancy, np.transpose]))
     dsk2, dependencies = cull(dsk, keys)
     dsk4, dependencies = fuse(dsk2, keys, dependencies)
     dsk5 = optimize_slices(dsk4)
@@ -36,48 +36,53 @@ def optimize_slices(dsk):
     See also:
         fuse_slice_dict
     """
-    getters = (getarray, getitem)
+    fancy_ind_types = (list, np.ndarray)
+    getters = (getarray_nofancy, getarray, getitem)
     dsk = dsk.copy()
     for k, v in dsk.items():
-        if type(v) is tuple:
-            if v[0] in getters:
-                try:
-                    func, a, a_index = v
-                    use_getarray = func is getarray
-                except ValueError:  # has four elements, includes a lock
-                    continue
-                while type(a) is tuple and a[0] in getters:
-                    try:
-                        f2, b, b_index = a
-                        use_getarray |= f2 is getarray
-                    except ValueError:  # has four elements, includes a lock
-                        break
-                    if (type(a_index) is tuple) != (type(b_index) is tuple):
-                        break
-                    if ((type(a_index) is tuple) and
-                            (len(a_index) != len(b_index)) and
-                            any(i is None for i in b_index + a_index)):
+        if type(v) is tuple and v[0] in getters and len(v) == 3:
+            f, a, a_index = v
+            getter = f
+            while type(a) is tuple and a[0] in getters and len(a) == 3:
+                f2, b, b_index = a
+                if (type(a_index) is tuple) != (type(b_index) is tuple):
+                    break
+                if type(a_index) is tuple:
+                    indices = b_index + a_index
+                    if (len(a_index) != len(b_index) and
+                            any(i is None for i in indices)):
                         break
-                    try:
-                        c_index = fuse_slice(b_index, a_index)
-                    except NotImplementedError:
+                    if (f2 is getarray_nofancy and
+                            any(isinstance(i, fancy_ind_types) for i in indices)):
                         break
-                    (a, a_index) = (b, c_index)
-                if use_getarray:
-                    dsk[k] = (getarray, a, a_index)
-                elif (type(a_index) is slice and
-                      not a_index.start and
-                      a_index.stop is None and
-                      a_index.step is None):
-                    dsk[k] = a
-                elif type(a_index) is tuple and all(type(s) is slice and
-                                                    not s.start and
-                                                    s.stop is None and
-                                                    s.step is None
-                                                    for s in a_index):
-                    dsk[k] = a
-                else:
-                    dsk[k] = (getitem, a, a_index)
+                elif (f2 is getarray_nofancy and
+                        (type(a_index) in fancy_ind_types or
+                         type(b_index) in fancy_ind_types)):
+                    break
+                try:
+                    c_index = fuse_slice(b_index, a_index)
+                    # rely on fact that nested gets never decrease in
+                    # strictness e.g. `(getarray, (getitem, ...))` never
+                    # happens
+                    getter = f2
+                except NotImplementedError:
+                    break
+                a, a_index = b, c_index
+            if getter is not getitem:
+                dsk[k] = (getter, a, a_index)
+            elif (type(a_index) is slice and
+                    not a_index.start and
+                    a_index.stop is None and
+                    a_index.step is None):
+                dsk[k] = a
+            elif type(a_index) is tuple and all(type(s) is slice and
+                                                not s.start and
+                                                s.stop is None and
+                                                s.step is None
+                                                for s in a_index):
+                dsk[k] = a
+            else:
+                dsk[k] = (getitem, a, a_index)
     return dsk
 
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -21,10 +21,10 @@
 from dask.utils_test import inc
 
 from dask.array import chunk
-from dask.array.core import (getem, getarray, top, dotmany, concatenate3,
-                             broadcast_dimensions, Array, stack, concatenate,
-                             from_array, take, elemwise, isnull, notnull,
-                             broadcast_shapes, partial_by_order, exp,
+from dask.array.core import (getem, getarray, getarray_nofancy, top, dotmany,
+                             concatenate3, broadcast_dimensions, Array, stack,
+                             concatenate, from_array, take, elemwise, isnull,
+                             notnull, broadcast_shapes, partial_by_order, exp,
                              tensordot, choose, where, coarsen, insert,
                              broadcast_to, reshape, fromfunction,
                              blockdims_from_blockshape, store, optimize,
@@ -51,11 +51,12 @@ def key(k):
 
 
 def test_getem():
-    assert getem('X', (2, 3), shape=(4, 6)) == \
-    {('X', 0, 0): (getarray, 'X', (slice(0, 2), slice(0, 3))),
-     ('X', 1, 0): (getarray, 'X', (slice(2, 4), slice(0, 3))),
-     ('X', 1, 1): (getarray, 'X', (slice(2, 4), slice(3, 6))),
-     ('X', 0, 1): (getarray, 'X', (slice(0, 2), slice(3, 6)))}
+    for fancy, getter in [(True, getarray), (False, getarray_nofancy)]:
+        sol = {('X', 0, 0): (getter, 'X', (slice(0, 2), slice(0, 3))),
+               ('X', 1, 0): (getter, 'X', (slice(2, 4), slice(0, 3))),
+               ('X', 1, 1): (getter, 'X', (slice(2, 4), slice(3, 6))),
+               ('X', 0, 1): (getter, 'X', (slice(0, 2), slice(3, 6)))}
+        assert getem('X', (2, 3), shape=(4, 6), fancy=fancy) == sol
 
 
 def test_top():
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -7,7 +7,7 @@
         fuse_slice)
 
 from dask.utils import raises
-from dask.array.core import getarray
+from dask.array.core import getarray, getarray_nofancy
 
 
 def test_fuse_getitem():
@@ -18,6 +18,10 @@ def test_fuse_getitem():
                         (slice(15, 20), slice(50, 60))),
               (getarray, 'x', (slice(1015, 1020), slice(150, 160)))),
 
+             ((getitem, (getarray_nofancy, 'x', (slice(1000, 2000), slice(100, 200))),
+                        (slice(15, 20), slice(50, 60))),
+              (getarray_nofancy, 'x', (slice(1015, 1020), slice(150, 160)))),
+
              ((getarray, (getarray, 'x', slice(1000, 2000)), 10),
               (getarray, 'x', 1010)),
 
@@ -25,6 +29,10 @@ def test_fuse_getitem():
                         (slice(15, 20),)),
               (getarray, 'x', (slice(1015, 1020), 10))),
 
+             ((getitem, (getarray_nofancy, 'x', (slice(1000, 2000), 10)),
+                        (slice(15, 20),)),
+              (getarray_nofancy, 'x', (slice(1015, 1020), 10))),
+
              ((getarray, (getarray, 'x', (10, slice(1000, 2000))),
                         (slice(15, 20),)),
               (getarray, 'x', (10, slice(1015, 1020)))),
@@ -117,3 +125,12 @@ def test_dont_fuse_different_slices():
     y = x.rechunk((1, 10))
     dsk = optimize(y.dask, y._keys())
     assert len(dsk) > 100
+
+
+def test_dont_fuse_fancy_indexing_in_getarray_nofancy():
+    dsk = {'a': (getitem, (getarray_nofancy, 'x', (slice(10, 20, None), slice(100, 200, None))),
+                 ([1, 3], slice(50, 60, None)))}
+    assert optimize_slices(dsk) == dsk
+
+    dsk = {'a': (getitem, (getarray_nofancy, 'x', [1, 2, 3]), 0)}
+    assert optimize_slices(dsk) == dsk