Repository: dask/dask

Issue:
dask should not assume indexing an input array creates a numpy.ndarray
Consider:

``` python
import dask.array as da
import dask
import numpy as np

class Wrapper(object):
    def __init__(self, array):
        self.array = array
        self.shape = array.shape
        self.dtype = array.dtype

    def __getitem__(self, key):
        return Wrapper(self.array[key])

    def __array__(self, dtype=None):
        return np.array(self.array, dtype=dtype)

x = da.from_array(Wrapper(np.arange(5)), blockshape=(5,))
with dask.set_options(get=dask.get):
    (x + 1).compute()
```

```
TypeError: unsupported operand type(s) for +: 'Wrapper' and 'int'
```

Many wrappers classes created numpy arrays when indexed (it's often the simplest thing to do), but this is by no means always the case (e.g., consider wrapping an alternate implementation of dask arrays).

Thus, for the purposes of arithmetic, at least, input should be coerced to numpy arrays using `np.asarray`. It is probably a good idea to do this on all arrays as soon as the computation gets rolling, though this might have some unfortunate complications for fused indexing (#103).



Discussion:

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 from operator import add, getitem
+import operator
 import inspect
 from collections import Iterable
 from bisect import bisect
@@ -18,7 +19,6 @@
 from ..utils import deepmap, ignoring, repr_long_list
 from ..async import inline_functions
 from ..optimize import cull, inline
-from .optimization import optimize
 from ..compatibility import unicode
 from .. import threaded, core
 from ..context import _globals
@@ -27,20 +27,34 @@
 names = ('x_%d' % i for i in count(1))
 
 
+def getarray(a, b):
+    """ Mimics getitem but includes call to np.asarray
+
+    >>> getarray([1, 2, 3, 4, 5], slice(1, 4))
+    array([2, 3, 4])
+    """
+    c = a[b]
+    if type(c) != np.ndarray:
+        c = np.asarray(c)
+    return c
+
+
+from .optimization import optimize
+
 def getem(arr, blockdims=None, blockshape=None, shape=None):
     """ Dask getting various chunks from an array-like
 
     >>> getem('X', blockshape=(2, 3), shape=(4, 6))  # doctest: +SKIP
-    {('X', 0, 0): (getitem, 'X', (slice(0, 2), slice(0, 3))),
-     ('X', 1, 0): (getitem, 'X', (slice(2, 4), slice(0, 3))),
-     ('X', 1, 1): (getitem, 'X', (slice(2, 4), slice(3, 6))),
-     ('X', 0, 1): (getitem, 'X', (slice(0, 2), slice(3, 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)))}
 
     >>> getem('X', blockdims=((2, 2), (3, 3)))  # doctest: +SKIP
-    {('X', 0, 0): (getitem, 'X', (slice(0, 2), slice(0, 3))),
-     ('X', 1, 0): (getitem, 'X', (slice(2, 4), slice(0, 3))),
-     ('X', 1, 1): (getitem, 'X', (slice(2, 4), slice(3, 6))),
-     ('X', 0, 1): (getitem, 'X', (slice(0, 2), slice(3, 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)))}
     """
     if not blockdims:
         blockdims = blockdims_from_blockshape(shape, blockshape)
@@ -51,7 +65,7 @@ def getem(arr, blockdims=None, blockshape=None, shape=None):
     shapes = product(*blockdims)
     starts = product(*cumdims)
 
-    values = ((getitem, arr) + (tuple(slice(s, s+dim)
+    values = ((getarray, arr) + (tuple(slice(s, s+dim)
                                  for s, dim in zip(start, shape)),)
                 for start, shape in zip(starts, shapes))
 
@@ -656,7 +670,7 @@ def __getitem__(self, index):
                 dt = np.dtype([(name, self._dtype[name]) for name in index])
             else:
                 dt = None
-            return elemwise(getitem, self, index, dtype=dt)
+            return elemwise(getarray, self, index, dtype=dt)
 
         # Slicing
         out = next(names)
@@ -948,7 +962,7 @@ def stack(seq, axis=0):
     names = [a.name for a in seq]
     inputs = [(names[key[axis+1]],) + key[1:axis + 1] + key[axis + 2:]
                 for key in keys]
-    values = [(getitem, inp, (slice(None, None, None),) * axis
+    values = [(getarray, inp, (slice(None, None, None),) * axis
                            + (None,)
                            + (slice(None, None, None),) * (ndim - axis))
                 for inp in inputs]
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -2,9 +2,10 @@
 from ..core import flatten
 from ..async import inline_functions
 from ..optimize import dealias
+from .core import getarray
 from operator import getitem
 from dask.rewrite import RuleSet, RewriteRule
-from toolz import valmap
+from toolz import valmap, partial
 import numpy as np
 
 
@@ -16,7 +17,7 @@ def optimize(dsk, keys, **kwargs):
     3.  Inline fast functions like getitem and np.transpose
     """
     fast_functions=kwargs.get('fast_functions',
-                             set([getitem, np.transpose]))
+                             set([getarray, getitem, np.transpose]))
     dsk2 = cull(dsk, list(flatten(keys)))
     dsk3 = remove_full_slices(dsk2)
     dsk4 = fuse(dsk3)
@@ -36,7 +37,7 @@ def is_full_slice(task):
     False
     """
     return (isinstance(task, tuple) and
-            task[0] == getitem and
+            (task[0] in (getitem, getarray)) and
             (task[2] == slice(None, None, None) or
              isinstance(task[2], tuple) and
              all(ind == slice(None, None, None) for ind in task[2])))
@@ -72,21 +73,30 @@ def remove_full_slices(dsk):
 a, b, x = '~a', '~b', '~x'
 
 
-def fuse_slice_dict(match):
+def fuse_slice_dict(match, getter=getarray):
     # Making new dimensions?  Need two getitems
     # TODO: well, we could still optimize the non-None axes
-
     try:
         c = fuse_slice(match[a], match[b])
-        return (getitem, match[x], c)
+        return (getter, match[x], c)
     except NotImplementedError:
-        return (getitem, (getitem, match[x], match[a]),
+        return (getitem, (getter, match[x], match[a]),
                          match[b])
 
 
-rewrite_rules = RuleSet(RewriteRule((getitem, (getitem, x, a), b),
-                                    fuse_slice_dict,
-                                    (a, b, x)))
+rewrite_rules = RuleSet(
+        RewriteRule((getitem, (getitem, x, a), b),
+                    partial(fuse_slice_dict, getter=getitem),
+                    (a, b, x)),
+        RewriteRule((getarray, (getitem, x, a), b),
+                    fuse_slice_dict,
+                    (a, b, x)),
+        RewriteRule((getarray, (getarray, x, a), b),
+                    fuse_slice_dict,
+                    (a, b, x)),
+        RewriteRule((getitem, (getarray, x, a), b),
+                    fuse_slice_dict,
+                    (a, b, x)))
 
 
 def normalize_slice(s):


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
@@ -13,10 +13,10 @@
 
 def test_getem():
     assert getem('X', blockshape=(2, 3), shape=(4, 6)) == \
-    {('X', 0, 0): (getitem, 'X', (slice(0, 2), slice(0, 3))),
-     ('X', 1, 0): (getitem, 'X', (slice(2, 4), slice(0, 3))),
-     ('X', 1, 1): (getitem, 'X', (slice(2, 4), slice(3, 6))),
-     ('X', 0, 1): (getitem, 'X', (slice(0, 2), slice(3, 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)))}
 
 
 def test_top():
@@ -209,25 +209,25 @@ def test_stack():
 
     assert s.shape == (3, 4, 6)
     assert s.blockdims == ((1, 1, 1), (2, 2), (3, 3))
-    assert s.dask[(s.name, 0, 1, 0)] == (getitem, ('A', 1, 0),
+    assert s.dask[(s.name, 0, 1, 0)] == (getarray, ('A', 1, 0),
                                           (None, colon, colon))
-    assert s.dask[(s.name, 2, 1, 0)] == (getitem, ('C', 1, 0),
+    assert s.dask[(s.name, 2, 1, 0)] == (getarray, ('C', 1, 0),
                                           (None, colon, colon))
 
     s2 = stack([a, b, c], axis=1)
     assert s2.shape == (4, 3, 6)
     assert s2.blockdims == ((2, 2), (1, 1, 1), (3, 3))
-    assert s2.dask[(s2.name, 0, 1, 0)] == (getitem, ('B', 0, 0),
+    assert s2.dask[(s2.name, 0, 1, 0)] == (getarray, ('B', 0, 0),
                                             (colon, None, colon))
-    assert s2.dask[(s2.name, 1, 1, 0)] == (getitem, ('B', 1, 0),
+    assert s2.dask[(s2.name, 1, 1, 0)] == (getarray, ('B', 1, 0),
                                             (colon, None, colon))
 
     s2 = stack([a, b, c], axis=2)
     assert s2.shape == (4, 6, 3)
     assert s2.blockdims == ((2, 2), (3, 3), (1, 1, 1))
-    assert s2.dask[(s2.name, 0, 1, 0)] == (getitem, ('A', 0, 1),
+    assert s2.dask[(s2.name, 0, 1, 0)] == (getarray, ('A', 0, 1),
                                             (colon, colon, None))
-    assert s2.dask[(s2.name, 1, 1, 2)] == (getitem, ('C', 1, 1),
+    assert s2.dask[(s2.name, 1, 1, 2)] == (getarray, ('C', 1, 1),
                                             (colon, colon, None))
 
     assert raises(ValueError, lambda: stack([a, b, c], axis=3))
@@ -839,6 +839,38 @@ def test_optimize():
     assert all(key in result for key in expr._keys())
 
 
+def test_slicing_with_non_ndarrays():
+    class ARangeSlice(object):
+        def __init__(self, start, stop):
+            self.start = start
+            self.stop = stop
+
+        def __array__(self):
+            return np.arange(self.start, self.stop)
+
+    class ARangeSlicable(object):
+        dtype = 'i8'
+
+        def __init__(self, n):
+            self.n = n
+
+        @property
+        def shape(self):
+            return (self.n,)
+
+        def __getitem__(self, key):
+            return ARangeSlice(key[0].start, key[0].stop)
+
+
+    x = da.from_array(ARangeSlicable(10), blockshape=(4,))
+
+    assert eq((x + 1).sum(), (np.arange(10) + 1).sum())
+
+
+def test_getarray():
+    assert type(getarray(np.matrix([[1]]), 0)) == np.ndarray
+
+
 def test_squeeze():
     x = da.ones((10, 1), blockshape=(3, 1))
 
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
@@ -1,38 +1,39 @@
 from dask.array.optimization import (getitem, rewrite_rules, optimize,
         remove_full_slices, fuse_slice)
 from dask.utils import raises
+from dask.array.core import getarray
 
 
 def test_fuse_getitem():
-    pairs = [((getitem, (getitem, 'x', slice(1000, 2000)), slice(15, 20)),
-              (getitem, 'x', slice(1015, 1020))),
+    pairs = [((getarray, (getarray, 'x', slice(1000, 2000)), slice(15, 20)),
+              (getarray, 'x', slice(1015, 1020))),
 
-             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(100, 200))),
+             ((getitem, (getarray, 'x', (slice(1000, 2000), slice(100, 200))),
                         (slice(15, 20), slice(50, 60))),
-              (getitem, 'x', (slice(1015, 1020), slice(150, 160)))),
+              (getarray, 'x', (slice(1015, 1020), slice(150, 160)))),
 
-             ((getitem, (getitem, 'x', slice(1000, 2000)), 10),
-              (getitem, 'x', 1010)),
+             ((getarray, (getarray, 'x', slice(1000, 2000)), 10),
+              (getarray, 'x', 1010)),
 
-             ((getitem, (getitem, 'x', (slice(1000, 2000), 10)),
+             ((getitem, (getarray, 'x', (slice(1000, 2000), 10)),
                         (slice(15, 20),)),
-              (getitem, 'x', (slice(1015, 1020), 10))),
+              (getarray, 'x', (slice(1015, 1020), 10))),
 
-             ((getitem, (getitem, 'x', (10, slice(1000, 2000))),
+             ((getarray, (getarray, 'x', (10, slice(1000, 2000))),
                         (slice(15, 20),)),
-              (getitem, 'x', (10, slice(1015, 1020)))),
+              (getarray, 'x', (10, slice(1015, 1020)))),
 
-             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(100, 200))),
+             ((getarray, (getarray, 'x', (slice(1000, 2000), slice(100, 200))),
                         (slice(None, None), slice(50, 60))),
-              (getitem, 'x', (slice(1000, 2000), slice(150, 160)))),
+              (getarray, 'x', (slice(1000, 2000), slice(150, 160)))),
 
-             ((getitem, (getitem, 'x', (None, slice(None, None))),
+             ((getarray, (getarray, 'x', (None, slice(None, None))),
                         (slice(None, None), 5)),
-              (getitem, 'x', (None, 5))),
+              (getarray, 'x', (None, 5))),
 
-             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(10, 20))),
+             ((getarray, (getarray, 'x', (slice(1000, 2000), slice(10, 20))),
                         (slice(5, 10),)),
-              (getitem, 'x', (slice(1005, 1010), slice(10, 20)))),
+              (getarray, 'x', (slice(1005, 1010), slice(10, 20)))),
 
              ((getitem, (getitem, 'x', (slice(1000, 2000),)),
                         (slice(5, 10), slice(10, 20))),
@@ -47,23 +48,23 @@ def test_fuse_getitem():
 
 def test_optimize_with_getitem_fusion():
     dsk = {'a': 'some-array',
-           'b': (getitem, 'a', (slice(10, 20), slice(100, 200))),
-           'c': (getitem, 'b', (5, slice(50, 60)))}
+           'b': (getarray, 'a', (slice(10, 20), slice(100, 200))),
+           'c': (getarray, 'b', (5, slice(50, 60)))}
 
     result = optimize(dsk, ['c'])
-    expected = {'c': (getitem, 'some-array', (15, slice(150, 160)))}
+    expected = {'c': (getarray, 'some-array', (15, slice(150, 160)))}
     assert result == expected
 
 
 def test_optimize_slicing():
     dsk = {'a': (range, 10),
-           'b': (getitem, 'a', (slice(None, None, None),)),
-           'c': (getitem, 'b', (slice(None, None, None),)),
-           'd': (getitem, 'c', (slice(None, 5, None),)),
-           'e': (getitem, 'd', (slice(None, None, None),))}
+           'b': (getarray, 'a', (slice(None, None, None),)),
+           'c': (getarray, 'b', (slice(None, None, None),)),
+           'd': (getarray, 'c', (slice(None, 5, None),)),
+           'e': (getarray, 'd', (slice(None, None, None),))}
 
     expected = {'a': (range, 10),
-                'e': (getitem, 'a', (slice(None, 5, None),))}
+                'e': (getarray, 'a', (slice(None, 5, None),))}
 
     assert remove_full_slices(dsk) == expected
 
@@ -94,6 +95,6 @@ def test_fuse_slice_with_lists():
 
 
 def test_hard_fuse_slice_cases():
-    term = (getitem, (getitem, 'x', (None, slice(None, None))),
+    term = (getarray, (getarray, 'x', (None, slice(None, None))),
                      (slice(None, None), 5))
-    assert rewrite_rules.rewrite(term) == (getitem, 'x', (None, 5))
+    assert rewrite_rules.rewrite(term) == (getarray, 'x', (None, 5))