Repository: dask/dask

Issue:
optimization of array.concatenate depends strongly on endianness
I have encountered a dask optimization issue that I think is at the core of xgcm/xmitgcm#73.

Basically, I am constructing a big dask array by concatenating many numpy memmaps, each created within a `from_delayed` function. Then I want to get back out a single value from this big array. In theory this should go very fast and use very little memory, as if I had accessed the original memmap. And indeed it does...unless the dtype is big endian!

Here is how to reproduce this issue. First create some test data
```python
import numpy as np
import dask
import dask.array

# create some example binary data (8GB)
# large dataset is necessary to see timing differences
shape = (1, 50, 1000, 2000)
nfiles = 10
dtype = np.dtype('f8') # note: this doesn't matter at this stage
data = np.zeros(shape, dtype)

filenames = ['data.{:04d}'.format(n) for n in range(nfiles)]
for fname in filenames:
    data.tofile(fname)
```

Now define some functions to read it into dask arrays
```python
def read_as_memmap(fname, dtype):
    return np.memmap(fname, dtype=dtype, shape=shape)
    
def lazy_memmap(fname, dtype):
    return dask.array.from_delayed(
                dask.delayed(read_as_memmap)(fname, dtype),
                shape, dtype)

def read_all_data(dtype):
    return dask.array.concatenate(
            [lazy_memmap(fname, dtype) for fname in filenames])
```

Now perform a timing test on reading back a single value with default (little endian) datatype
```python
dtype = np.dtype('f8')
all_data = read_all_data(dtype)
%timeit lazy_memmap(filenames[-1], dtype)[0, 0, 0].compute()
%timeit all_data[-1, 0, 0, 0].compute()
```
On my machine I get
```
100 loops, best of 3: 2.54 ms per loop
100 loops, best of 3: 2.3 ms per loop
```
basically identical and very fast, as we expect. The dask graph for `all_data[-1, 0, 0, 0]` looks like this:
![image](https://user-images.githubusercontent.com/1197350/37015504-abb9adee-20d5-11e8-8506-14fafe84dba0.png)


Instead, if I repeat the test with a big-endian dtype
```python
dtype = np.dtype('>f8')
all_data = read_all_data(dtype)
%timeit lazy_memmap(filenames[-1], dtype)[0, 0, 0].compute()
%timeit all_data[-1, 0, 0, 0].compute()
```
I get this:
```
100 loops, best of 3: 2.57 ms per loop
1 loop, best of 3: 929 ms per loop
```
The `ResourceProfiler` diagnostics also indicate much higher memory usage. Now the dask graph looks like this
![image](https://user-images.githubusercontent.com/1197350/37015518-bc1160ce-20d5-11e8-84d6-0f5a9ace8dd8.png)

There appears to be an extra call to `astype` which is interfering with the optimization somehow.

I'm using dask version 0.17.1.




Discussion:
Yeah, dask array optimizations aren't very serious, we do some obvious
things, but we're not a fully-fledged query optimization engine (this would
need to be higher in the stack and doesn't currently exist).  We merge
getitem calls into each other, which means that you often don't have to
read out full chunks and then get small bits, you can just get the small
bits directly.  The presence of an astype call interferes with this
optimization.  Dask.array doesn't know which operations easily commute with
getitems and so it leaves this alone.

On Tue, Mar 6, 2018 at 12:34 AM, Ryan Abernathey <notifications@github.com>
wrote:

> I have encountered a dask optimization issue that I think is at the core
> of xgcm/xmitgcm#73 <https://github.com/xgcm/xmitgcm/issues/73>.
>
> Basically, I am constructing a big dask array by concatenating many numpy
> memmaps, each created within a from_delayed function. Then I want to get
> back out a single value from this big array. In theory this should go very
> fast and use very little memory, as if I had accessed the original memmap.
> And indeed it does...unless the dtype is big endian!
>
> Here is how to reproduce this issue. First create some test data
>
> import numpy as npimport daskimport dask.array
> # create some example binary data (8GB)# large dataset is necessary to see timing differences
> shape = (1, 50, 1000, 2000)
> nfiles = 10
> dtype = np.dtype('f8') # note: this doesn't matter at this stage
> data = np.zeros(shape, dtype)
>
> filenames = ['data.{:04d}'.format(n) for n in range(nfiles)]for fname in filenames:
>     data.tofile(fname)
>
> Now define some functions to read it into dask arrays
>
> def read_as_memmap(fname, dtype):
>     return np.memmap(fname, dtype=dtype, shape=shape)
>     def lazy_memmap(fname, dtype):
>     return dask.array.from_delayed(
>                 dask.delayed(read_as_memmap)(fname, dtype),
>                 shape, dtype)
> def read_all_data(dtype):
>     return dask.array.concatenate(
>             [lazy_memmap(fname, dtype) for fname in filenames])
>
> Now perform a timing test on reading back a single value with default
> (little endian) datatype
>
> dtype = np.dtype('f8')
> all_data = read_all_data(dtype)%timeit lazy_memmap(filenames[-1], dtype)[0, 0, 0].compute()%timeit all_data[-1, 0, 0, 0].compute()
>
> On my machine I get
>
> 100 loops, best of 3: 2.54 ms per loop
> 100 loops, best of 3: 2.3 ms per loop
>
> basically identical and very fast, as we expect. The dask graph for all_data[-1,
> 0, 0, 0] looks like this:
> [image: image]
> <https://user-images.githubusercontent.com/1197350/37015504-abb9adee-20d5-11e8-8506-14fafe84dba0.png>
>
> Instead, if I repeat the test with a big-endian dtype
>
> dtype = np.dtype('>f8')
> all_data = read_all_data(dtype)%timeit lazy_memmap(filenames[-1], dtype)[0, 0, 0].compute()%timeit all_data[-1, 0, 0, 0].compute()
>
> I get this:
>
> 100 loops, best of 3: 2.57 ms per loop
> 1 loop, best of 3: 929 ms per loop
>
> The ResourceProfiler diagnostics also indicate much higher memory usage.
> Now the dask graph looks like this
> [image: image]
> <https://user-images.githubusercontent.com/1197350/37015518-bc1160ce-20d5-11e8-84d6-0f5a9ace8dd8.png>
>
> There appears to be an extra call to astype which is interfering with the
> optimization somehow.
>
> I'm using dask version 0.17.1.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/3248>, or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszJDewALSGjVoM6zFzQJNmY0hgrh8ks5tbh_LgaJpZM4SeIZL>
> .
>

Why does astype need to get called in the first place? We are just selecting data. Why can't it just remain a big endian dtype?
Not sure.  I'd have to investigate.

On Fri, Mar 9, 2018 at 10:50 AM, Ryan Abernathey <notifications@github.com>
wrote:

> Why does astype need to get called in the first place? We are just
> selecting data. Why can't it just remain a big endian dtype?
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/3248#issuecomment-371851713>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszKkoeontVZa87Idlg9x2eI0qkwXHks5tcqTTgaJpZM4SeIZL>
> .
>

What steps could I and @dhruvbalwada take to help debug and fix this issue? It is a severe bottleneck on our work. Dask internals are somewhat intimidating to us, and we don't know where to start.
One approach would be to determine why the astype call is there in the first place and if it can be removed.  Presumably this would involve looking at the functions that you call and tracing the code path to find the location of the astype call.  
> Presumably this would involve looking at the functions that you call and tracing the code path to find the location of the astype call.

To clarify: we do not call `astype` directly in the user code. So you mean tracing the code path *within dask*?
Yes
I believe this may be because of 

```
numpy.promote_types('>f8', '>f8') == '<f8'
```

This is use in `concatenate` to determine the output dtype, and the result being different from the input dtype is what I think results in the `astype` call. (see https://github.com/dask/dask/blob/master/dask/array/core.py#L2614 )
The following would avoid this situation

```diff
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2611,8 +2611,12 @@ def concatenate(seq, axis=0, allow_unknown_chunksizes=False):

     cum_dims = [0] + list(accumulate(add, [len(a.chunks[axis]) for a in seq]))

-    dt = reduce(np.promote_types, [a.dtype for a in seq])
-    seq = [x.astype(dt) for x in seq]
+    seq_dtypes = [a.dtype for a in seq]
+    if len(set(seq_dtypes)) > 1:
+        dt = reduce(np.promote_types, seq_dtypes)
+        seq = [x.astype(dt) for x in seq]
+    else:
+        dt = seq_dtypes[0]

     names = [a.name for a in seq]
```
@martindurant care to submit a PR with this fix and a small test?

On Mon, Mar 19, 2018 at 11:48 AM, Martin Durant <notifications@github.com>
wrote:

> The following would avoid this situation
>
> --- a/dask/array/core.py+++ b/dask/array/core.py@@ -2611,8 +2611,12 @@ def concatenate(seq, axis=0, allow_unknown_chunksizes=False):
>
>      cum_dims = [0] + list(accumulate(add, [len(a.chunks[axis]) for a in seq]))
> -    dt = reduce(np.promote_types, [a.dtype for a in seq])-    seq = [x.astype(dt) for x in seq]+    seq_dtypes = [a.dtype for a in seq]+    if len(set(seq_dtypes)) > 1:+        dt = reduce(np.promote_types, seq_dtypes)+        seq = [x.astype(dt) for x in seq]+    else:+        dt = seq_dtypes[0]
>
>      names = [a.name for a in seq]
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/3248#issuecomment-374260413>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszGveOmDV5ukr7OCIkAWS_2T_pXAsks5tf9M7gaJpZM4SeIZL>
> .
>

Yes, will do - posted here first because I wasn't certain this was the right thing to do. Also, it will not be foolproof: a set of dtypes like ('>f8', '>f8', '>f4') will become '<f8' and use `astype` for every chunk.

--edit-- 
right, it will prefer little-endian, or perhaps this is determined by the OS/build.

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2611,8 +2611,12 @@ def concatenate(seq, axis=0, allow_unknown_chunksizes=False):
 
     cum_dims = [0] + list(accumulate(add, [len(a.chunks[axis]) for a in seq]))
 
-    dt = reduce(np.promote_types, [a.dtype for a in seq])
-    seq = [x.astype(dt) for x in seq]
+    seq_dtypes = [a.dtype for a in seq]
+    if len(set(seq_dtypes)) > 1:
+        dt = reduce(np.promote_types, seq_dtypes)
+        seq = [x.astype(dt) for x in seq]
+    else:
+        dt = seq_dtypes[0]
 
     names = [a.name for a in seq]
 


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
@@ -335,6 +335,17 @@ def test_concatenate():
     pytest.raises(ValueError, lambda: concatenate([a, b, c], axis=2))
 
 
+@pytest.mark.parametrize('dtypes', [(('>f8', '>f8'), '>f8'),
+                                    (('<f4', '<f8'), '<f8')])
+def test_concatenate_types(dtypes):
+    dts_in, dt_out = dtypes
+    arrs = [np.zeros(4, dtype=dt) for dt in dts_in]
+    darrs = [from_array(arr, chunks=(2,)) for arr in arrs]
+
+    x = concatenate(darrs, axis=0)
+    assert x.dtype == dt_out
+
+
 def test_concatenate_unknown_axes():
     dd = pytest.importorskip('dask.dataframe')
     pd = pytest.importorskip('pandas')